[AMD][MORI] Deduplicate CP-replicated state transfers (#36025)

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:

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 (
+        if self._should_skip_cp_replicated_state_transfer():
-            self.attn_cp_size > 1
-            and self.attn_cp_rank != 0
-            and not get_parallel().enable_dsa_cache_layer_split
-        ):
             skip_state = True
 
         if not self.is_hybrid_mla_backend:

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(

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):