From 983e4aa18d7e8746d3ef357ad9f8c8d757a3de5b Mon Sep 17 00:00:00 2001 From: hunhokim <74707298+hunhokim@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:58:48 +0900 Subject: [PATCH] Eliminate redundant DSA state transfers (Mooncake) (#32620) Co-authored-by: Hun-ho Kim --- .../srt/disaggregation/mooncake/conn.py | 22 +++++++++++++------ 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 3a92fd2a5..d4489fd03 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -974,6 +974,21 @@ class MooncakeKVManager(CommonKVManager): ) -> Tuple[bool, bool]: skip_kv = 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: return skip_kv, skip_state @@ -986,13 +1001,6 @@ class MooncakeKVManager(CommonKVManager): # Hybrid-MLA KV is replicated across these source ranks, but # 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 def _is_generic_kvcache_state_type(self, st: StateType) -> bool: