[DSv32] Overlap indexer weights_proj during dual_stream decode (#16637)

Co-authored-by: Ziang Li <ziangli@humansand.ai>

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(
+        if enable_dual_stream and forward_batch.forward_mode.is_decode_or_idle():
-            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)
-
+            weights = self._project_and_scale_head_gates(x)
-            q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
             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)
+            query, key = self._get_q_k_bf16(
-            k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt)
+                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:

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)
+            # overlap q_b_proj and indexer during decode
-            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+            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,