From f20ef7ace4e731b9ea7c5f38adce891f095e1a1b Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Thu, 7 May 2026 06:37:50 +0800 Subject: [PATCH] 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. --- .../attention/nsa/cp_shared_kv_prefetch.py | 10 ++++++-- .../srt/layers/attention/nsa_backend.py | 25 +++++++++++++------ .../attention_forward_methods/forward_mla.py | 23 ++++++++++++++--- 3 files changed, 45 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py index c945d7f0c..613487d2d 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py +++ b/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: diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index e333031bf..696253ae9 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/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, ) ) diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 9ee93cbaa..2c06599c3 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/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,