[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 (
-            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:

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(