Eliminate redundant DSA state transfers (Mooncake) (#32620)
Co-authored-by: Hun-ho Kim <hunho.kim@samsung.com>
This commit is contained in:
@@ -974,6 +974,21 @@ class MooncakeKVManager(CommonKVManager):
|
|||||||
) -> Tuple[bool, bool]:
|
) -> Tuple[bool, bool]:
|
||||||
skip_kv = False
|
skip_kv = False
|
||||||
skip_state = False
|
skip_state = False
|
||||||
|
|
||||||
|
# Must be checked before the non-hybrid early return below, or every CP
|
||||||
|
# rank re-sends the same state and we transfer it cp_size times over.
|
||||||
|
# Prefill CP all-gathers before writing the pool, so every CP rank holds
|
||||||
|
# the full state regardless of whether the pool is hybrid. We assume no
|
||||||
|
# structure about the state rows, so we don't split them across CP ranks
|
||||||
|
# -- just let rank 0 send the whole thing (unless layer split already
|
||||||
|
# shards it per rank).
|
||||||
|
if (
|
||||||
|
self.attn_cp_size > 1
|
||||||
|
and self.attn_cp_rank != 0
|
||||||
|
and not self.server_args.enable_dsa_cache_layer_split
|
||||||
|
):
|
||||||
|
skip_state = True
|
||||||
|
|
||||||
if not self.is_hybrid_mla_backend:
|
if not self.is_hybrid_mla_backend:
|
||||||
return skip_kv, skip_state
|
return skip_kv, skip_state
|
||||||
|
|
||||||
@@ -986,13 +1001,6 @@ class MooncakeKVManager(CommonKVManager):
|
|||||||
# Hybrid-MLA KV is replicated across these source ranks, but
|
# Hybrid-MLA KV is replicated across these source ranks, but
|
||||||
# TP-sharded state needs every rank for the aggregation path.
|
# TP-sharded state needs every rank for the aggregation path.
|
||||||
|
|
||||||
if (
|
|
||||||
self.attn_cp_size > 1
|
|
||||||
and self.attn_cp_rank != 0
|
|
||||||
and not self.server_args.enable_dsa_cache_layer_split
|
|
||||||
):
|
|
||||||
skip_state = True
|
|
||||||
|
|
||||||
return skip_kv, skip_state
|
return skip_kv, skip_state
|
||||||
|
|
||||||
def _is_generic_kvcache_state_type(self, st: StateType) -> bool:
|
def _is_generic_kvcache_state_type(self, st: StateType) -> bool:
|
||||||
|
|||||||
Reference in New Issue
Block a user