From 20abaee26cd8ac1387584775babf910c517c3388 Mon Sep 17 00:00:00 2001 From: Ziang Li <106564213+zianglih@users.noreply.github.com> Date: Fri, 9 Jan 2026 21:06:44 -0800 Subject: [PATCH] [DSv32] Overlap indexer weights_proj during dual_stream decode (#16637) Co-authored-by: Ziang Li --- .../srt/layers/attention/nsa/nsa_indexer.py | 41 ++++++++++++----- python/sglang/srt/models/deepseek_v2.py | 45 ++++++++++++++----- 2 files changed, 64 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index d4ce2271c..9ed629967 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -205,6 +205,12 @@ class Indexer(MultiPlatformOp): else: yield + @torch.compile(dynamic=True) + def _project_and_scale_head_gates(self, x: torch.Tensor): + weights, _ = self.weights_proj(x.float()) + weights = weights * self.n_heads**-0.5 + return weights + @torch.compile(dynamic=True) def _get_logits_head_gate(self, x: torch.Tensor, q_scale: torch.Tensor): weights, _ = self.weights_proj(x.float()) @@ -856,21 +862,36 @@ class Indexer(MultiPlatformOp): return_indices, ) - query, key = self._get_q_k_bf16( - q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch - ) - - if enable_dual_stream: + if enable_dual_stream and forward_batch.forward_mode.is_decode_or_idle(): current_stream = torch.cuda.current_stream() self.alt_stream.wait_stream(current_stream) - - q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) + weights = self._project_and_scale_head_gates(x) with torch.cuda.stream(self.alt_stream): + query, key = self._get_q_k_bf16( + q_lora, x, positions, False, forward_batch=forward_batch + ) + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt) current_stream.wait_stream(self.alt_stream) + weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale else: - q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) - k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt) + query, key = self._get_q_k_bf16( + q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch + ) + + if enable_dual_stream: + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) + with torch.cuda.stream(self.alt_stream): + k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt) + current_stream.wait_stream(self.alt_stream) + else: + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) + k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt) + + weights = self._get_logits_head_gate(x, q_scale) # k_fp8: (seq_len, head_dim) fp8_e4m3fn # k_buffer: (num_total_tokens + page_size, head_dim) fp8_e4m3fn @@ -885,8 +906,6 @@ class Indexer(MultiPlatformOp): index_k_scale=k_scale, ) - weights = self._get_logits_head_gate(x, q_scale) - if is_cuda(): assert forward_batch.seq_lens_cpu is not None if len(forward_batch.seq_lens_cpu) == 0: diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 96bb3812d..748912af6 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1666,6 +1666,7 @@ class DeepseekV2AttentionMLA(nn.Module): from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode q_lora = None + topk_indices = None if self.q_lora_rank is not None: q, latent_cache = ( get_attn_tp_context() @@ -1722,8 +1723,39 @@ class DeepseekV2AttentionMLA(nn.Module): if self.use_nsa: q_lora = q - k_nope = k_nope.unsqueeze(1) - q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) + # overlap q_b_proj and indexer during decode + if ( + self.alt_stream is not None + and get_is_capture_mode() + and forward_batch.forward_mode.is_decode_or_idle() + and q_lora is not None + ): + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + with torch.cuda.stream(self.alt_stream): + k_nope = k_nope.unsqueeze(1) + q = self.q_b_proj(q)[0].view( + -1, self.num_local_heads, self.qk_head_dim + ) + topk_indices = self.indexer( + x=hidden_states, + q_lora=q_lora, + positions=positions, + forward_batch=forward_batch, + layer_id=self.layer_id, + ) + current_stream.wait_stream(self.alt_stream) + else: + k_nope = k_nope.unsqueeze(1) + q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) + if q_lora is not None: + topk_indices = self.indexer( + x=hidden_states, + q_lora=q_lora, + positions=positions, + forward_batch=forward_batch, + layer_id=self.layer_id, + ) else: q = self.q_proj(hidden_states)[0].view( -1, self.num_local_heads, self.qk_head_dim @@ -1818,15 +1850,6 @@ class DeepseekV2AttentionMLA(nn.Module): k_nope, k_pe = self.rebuild_cp_kv_cache( latent_cache, forward_batch, k_nope, k_pe ) - topk_indices = None - if q_lora is not None: - topk_indices = self.indexer( - x=hidden_states, - q_lora=q_lora, - positions=positions, - forward_batch=forward_batch, - layer_id=self.layer_id, - ) return ( q_pe,