diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 58d3be72a..74650d342 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -293,6 +293,20 @@ class CommonKVManager(BaseKVManager): f"Unsupported DisaggregationMode: {self.disaggregation_mode}" ) + def _should_skip_cp_replicated_state_transfer(self) -> bool: + """Whether this prefill rank should omit CP-replicated state. + + Prefill CP materializes global token order before writing state pools, so + every CP rank holds the same state. When all CP ranks transfer their KV + shards, only rank 0 needs to send that state. Cache layer split is the + exception because each CP rank owns different state layers. + """ + return ( + self.attn_cp_size > 1 + and self.attn_cp_rank != 0 + and not get_parallel().enable_dsa_cache_layer_split + ) + def requires_dcp_relayout(self, dst_dcp_size: int, dst_dcp_rank: int) -> bool: if self.dcp_size == dst_dcp_size: if self.dcp_rank != dst_dcp_rank: diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 47bec0d76..d8e47bf73 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -62,7 +62,7 @@ from sglang.srt.observability.trace import ( TraceReqContext, trace_set_thread_info, ) -from sglang.srt.runtime_context import get_memory, get_parallel, get_schedule +from sglang.srt.runtime_context import get_memory, get_schedule from sglang.srt.server_args import ServerArgs from sglang.srt.utils.network import NetworkAddress @@ -1186,11 +1186,7 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): # 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 get_parallel().enable_dsa_cache_layer_split - ): + if self._should_skip_cp_replicated_state_transfer(): skip_state = True if not self.is_hybrid_mla_backend: diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py index 7954877f7..33996f7a7 100644 --- a/python/sglang/srt/disaggregation/mori/conn.py +++ b/python/sglang/srt/disaggregation/mori/conn.py @@ -1425,12 +1425,17 @@ class MoriKVSender(CommonKVSender): if should_skip: return + transfer_state_indices = ( + None + if self.kv_mgr._should_skip_cp_replicated_state_transfer() + else state_indices + ) normalized_state = ( - _normalize_state_indices_per_component(state_indices) + _normalize_state_indices_per_component(transfer_state_indices) if is_last_chunk else None ) - self._record_transfer_indices(kv_indices, state_indices) + self._record_transfer_indices(kv_indices, transfer_state_indices) wait_event = getattr(self, "_early_send_wait_event", None) self._early_send_wait_event = None self.kv_mgr.enqueue_transfer( diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 800c427f1..e65364289 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -8,6 +8,7 @@ import numpy as np import torch from sglang.srt.disaggregation.base.conn import KVArgs, StateType +from sglang.srt.disaggregation.common.conn import CommonKVManager from sglang.srt.disaggregation.common.staging_handler import ( handle_staging_req, ) @@ -138,6 +139,53 @@ class TestDisaggregationWire(unittest.TestCase): self.assertEqual(unpack_list_of_buffers(pack_list_of_buffers(bufs)), bufs) +class TestCPReplicatedStateTransfer(unittest.TestCase): + def test_only_nonzero_cp_ranks_without_layer_split_skip_state(self): + cases = [ + (1, 0, False, False), + (8, 0, False, False), + (8, 1, False, True), + (8, 7, False, True), + (8, 1, True, False), + ] + + for cp_size, cp_rank, layer_split, expected in cases: + with self.subTest( + cp_size=cp_size, + cp_rank=cp_rank, + layer_split=layer_split, + ): + manager = object.__new__(CommonKVManager) + manager.attn_cp_size = cp_size + manager.attn_cp_rank = cp_rank + parallel = SimpleNamespace( + enable_dsa_cache_layer_split=layer_split, + ) + with patch( + "sglang.srt.disaggregation.common.conn.get_parallel", + return_value=parallel, + ): + self.assertEqual( + manager._should_skip_cp_replicated_state_transfer(), + expected, + ) + + def test_mooncake_uses_common_cp_state_policy(self): + manager = object.__new__(MooncakeKVManager) + manager.attn_cp_size = 8 + manager.attn_cp_rank = 3 + manager.is_hybrid_mla_backend = False + + with patch( + "sglang.srt.disaggregation.common.conn.get_parallel", + return_value=SimpleNamespace(enable_dsa_cache_layer_split=False), + ): + self.assertEqual( + manager._get_dsa_cache_transfer_skip_flags(None), + (False, True), + ) + + class TestGroupConcurrentContiguous(unittest.TestCase): @staticmethod def _arr(values):