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,