Avoid redundant CP KV rebuild on shared-KV MLA path

Shared-KV prefill already persists each rank's MLA shard and reconstructs the dense attention KV from the shared pool before attention. Keeping the legacy CP rebuild all-gather after direct write duplicated communication on the hot MQA-to-attention path. The rebuild remains enabled only for current-only reuse, where the backend intentionally consumes full current KV tensors instead of materializing history from the pool. Index and MLA prefetch now share one FIFO CUDA stream so their CP collectives preserve local launch order.

Constraint: CP shared-KV materialize is the authoritative KV source for prefix/cache-hit MLA attention.

Rejected: Gate prefetch by prefix owner-lane coverage | owner skew is not the root cause and would add an extra collective plus CPU sync.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce rebuild_cp_kv_cache for shared-KV direct-write unless the backend consumes k_nope/k_pe directly.

Tested: git diff --check; py_compile for cp_shared_kv_prefetch.py, nsa_backend.py, forward_mla.py; remote container py_compile after scp sync.

Not-tested: Full multi-node GLM5 performance run after this exact commit.

python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py

@@ -84,6 +84,7 @@ class CpSharedKVMlaPrefetcher:
         slot_logical_pages: torch.Tensor,
         page_inverse: torch.Tensor,
         dense_num_pages: int,
+        stream: Optional[torch.cuda.Stream] = None,
     ) -> None:
         self.layout = layout
         self.page_size = page_size
@@ -92,7 +93,7 @@ class CpSharedKVMlaPrefetcher:
         self.page_inverse = page_inverse
         self.dense_num_pages = dense_num_pages
         self.total_slots = int(slot_logical_pages.numel())
-        self.stream = torch.cuda.Stream()
+        self.stream = stream if stream is not None else torch.cuda.Stream()
         self.handles: dict[int, CpSharedKVMlaPrefetchHandle] = {}
         self.pending_attention_handle: Optional[CpSharedKVMlaPrefetchHandle] = None
         self.disabled = False
@@ -104,6 +105,7 @@ class CpSharedKVMlaPrefetcher:
         forward_batch: Any,
         metadata: Any,
         topk_transform_is_paged: bool,
+        stream: Optional[torch.cuda.Stream] = None,
     ) -> Optional["CpSharedKVMlaPrefetcher"]:
         if not cp_shared_kv_mla_prefetch_enabled():
             return None
@@ -226,6 +228,7 @@ class CpSharedKVMlaPrefetcher:
             slot_logical_pages=remap.slot_logical_pages,
             page_inverse=remap.page_inverse,
             dense_num_pages=remap.dense_num_pages,
+            stream=stream,
         )
 
     def _layer_in_pool(self, token_to_kv_pool: Any, layer_id: int) -> bool:
@@ -459,6 +462,7 @@ class CpSharedKVIndexPrefetcher:
         slot_logical_pages: torch.Tensor,
         page_inverse: torch.Tensor,
         dense_num_pages: int,
+        stream: Optional[torch.cuda.Stream] = None,
     ) -> None:
         self.layout = layout
         self.prefix_pages = prefix_pages
@@ -466,7 +470,7 @@ class CpSharedKVIndexPrefetcher:
         self.page_inverse = page_inverse
         self.dense_num_pages = dense_num_pages
         self.total_slots = int(slot_logical_pages.numel())
-        self.stream = torch.cuda.Stream()
+        self.stream = stream if stream is not None else torch.cuda.Stream()
         self.handles: dict[int, CpSharedKVIndexPrefetchHandle] = {}
         self.pending_attention_handle: Optional[CpSharedKVIndexPrefetchHandle] = None
         self.disabled = False
@@ -478,6 +482,7 @@ class CpSharedKVIndexPrefetcher:
         forward_batch: Any,
         metadata: Any,
         topk_transform_is_paged: bool,
+        stream: Optional[torch.cuda.Stream] = None,
     ) -> Optional["CpSharedKVIndexPrefetcher"]:
         if not cp_shared_kv_mla_prefetch_enabled():
             return None
@@ -649,6 +654,7 @@ class CpSharedKVIndexPrefetcher:
             slot_logical_pages=remap.slot_logical_pages,
             page_inverse=remap.page_inverse,
             dense_num_pages=remap.dense_num_pages,
+            stream=stream,
         )
 
     def _layer_in_pool(self, token_to_kv_pool: Any, layer_id: int) -> bool:

python/sglang/srt/layers/attention/nsa_backend.py

@@ -878,14 +878,22 @@ class NativeSparseAttnBackend(
             token_to_batch_idx=token_to_batch_idx,
         )
         self.forward_metadata = metadata
-        forward_batch.cp_shared_kv_mla_prefetcher = (
-            CpSharedKVMlaPrefetcher.maybe_create(
-                forward_batch=forward_batch,
-                metadata=metadata,
-                topk_transform_is_paged=(
-                    topk_transform_method == TopkTransformMethod.PAGED
-                ),
-            )
+        mla_prefetcher = CpSharedKVMlaPrefetcher.maybe_create(
+            forward_batch=forward_batch,
+            metadata=metadata,
+            topk_transform_is_paged=(
+                topk_transform_method == TopkTransformMethod.PAGED
+            ),
+        )
+        forward_batch.cp_shared_kv_mla_prefetcher = mla_prefetcher
+        # Use one FIFO stream for index and MLA prefix prefetch.  Both paths
+        # enqueue CP collectives; independent streams can let one rank advance
+        # to the next prefetch collective while another rank is still launching
+        # the previous one.
+        shared_prefetch_stream = (
+            getattr(mla_prefetcher, "stream", None)
+            if mla_prefetcher is not None
+            else None
         )
         forward_batch.cp_shared_kv_index_prefetcher = (
             CpSharedKVIndexPrefetcher.maybe_create(
@@ -894,6 +902,7 @@ class NativeSparseAttnBackend(
                 topk_transform_is_paged=(
                     topk_transform_method == TopkTransformMethod.PAGED
                 ),
+                stream=shared_prefetch_stream,
             )
         )
 

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py

@@ -13,6 +13,8 @@ from sglang.srt.layers.attention.nsa.utils import (
     nsa_use_prefill_cp,
 )
 from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
+    cp_shared_kv_current_reuse_enabled,
+    is_current_only_extend_batch,
     try_tai_fused_mla_store,
 )
 from sglang.srt.layers.communicator import get_attn_tp_context
@@ -317,10 +319,25 @@ class DeepseekMLAForwardMixin:
             forward_batch.cp_shared_mla_direct_write_done = (
                 shared_mla_direct_write_done
             )
-            # support allgather+rerrange
-            k_nope, k_pe = self.rebuild_cp_kv_cache(
-                latent_cache, forward_batch, k_nope, k_pe
+            shared_kv_materialize_will_read_pool = (
+                shared_mla_direct_write_done
+                and getattr(forward_batch, "uses_cp_shared_kv", False)
             )
+            current_reuse_needs_full_current_kv = (
+                cp_shared_kv_current_reuse_enabled()
+                and is_current_only_extend_batch(forward_batch)
+            )
+            if (
+                not shared_kv_materialize_will_read_pool
+                or current_reuse_needs_full_current_kv
+            ):
+                # Legacy CP path needs full KV here.  CP shared KV normally
+                # reconstructs the attention KV from the persistent pool inside
+                # the backend, so this all-gather would duplicate the later
+                # materialize.  Keep it only for the current-only reuse fast path.
+                k_nope, k_pe = self.rebuild_cp_kv_cache(
+                    latent_cache, forward_batch, k_nope, k_pe
+                )
 
         return (
             q_pe,