[Perf] dsa: Capture DSA prefill attention and the NextN draft prefill in the prefill CUDA graph

docs/docs/references/environment_variables.mdx

@@ -591,6 +591,16 @@ SGLang supports various environment variables that can be used to configure its
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>When the maximum kv len in current prefill batch exceeds this value, the sparse mla kernel will be applied, else it falls back to dense MHA implementation. Default to the index topk of model (2048 for DeepSeek V3.2).</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>2048</code></td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_ENABLE_DSA_PREFILL_ATTN_IN_CUDA_GRAPH</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Record DSA prefill attention (TRT-LLM sparse MLA with FP8 KV and the fused top-k) inside the breakable prefill CUDA graph instead of an eager graph break. Buckets below the smallest one whose rows all exceed half the SM count keep the eager break, so outputs stay bitwise equal to it.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_ENABLE_DSA_PREFILL_INDEXER_PREP_IN_CUDA_GRAPH</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>With <code>SGLANG_ENABLE_DSA_PREFILL_ATTN_IN_CUDA_GRAPH</code>, also record the DSA indexer's Q/K preparation and index-K store in the prefill CUDA graph; only its logits and top-k stay an eager break.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
+    </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_DSA_TOPK_BROADCAST</code></td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Experimental. When enabled, broadcast the finalized NSA/DSA indexer top-k result from attention TP rank 0 to the other attention TP ranks. This can mitigate top-k mismatches in TP attention runs at the cost of some speed.</td>
@@ -1809,6 +1819,11 @@ SGLang supports various environment variables that can be used to configure its
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Skip draft_extend while adaptive spec is at steps=0; saves a draft forward but the draft KV goes stale.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_ENABLE_NEXTN_DRAFT_PREFILL_CUDA_GRAPH</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DeepSeek-family NextN (MTP) drafts: replay the breakable prefill CUDA graph for the draft extend that follows each target prefill chunk. Not used with DP attention, pipeline parallelism or prefill context parallelism.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
+    </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_NGRAM_FORCE_GREEDY_VERIFY</code></td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Force greedy verification for the n-gram speculative path.</td>

python/sglang/srt/environ.py

@@ -1734,6 +1734,9 @@ class Envs:
     # than the graph saves (e.g. DeepEP MoE workspace captured at full dispatch
     # capacity).
     SGLANG_DISABLE_DRAFT_EXTEND_CUDA_GRAPH = EnvBool(False)
+    # DeepSeek-family NextN (MTP) drafts: the draft extend after each target
+    # prefill chunk replays the breakable prefill CUDA graph instead of eager.
+    SGLANG_ENABLE_NEXTN_DRAFT_PREFILL_CUDA_GRAPH = EnvBool(True)
     # Use the split-KV (flash-decode) kernel for EAGLE target-verify on the
     # Triton backend (ROCm). Only active at speculative topk == 1; falls back to
     # extend_attention_fwd for unsupported cases or when set false (e.g. for
@@ -2183,6 +2186,12 @@ class Envs:
     SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC = EnvBool(False)
     SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK = EnvStr(None)
     SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD = EnvInt(2048)
+    # Capture DSA prefill attention (k_b bmm, RoPE + FP8, KV write, TRT-LLM sparse
+    # MLA) in the breakable prefill CUDA graph instead of an eager break.
+    SGLANG_ENABLE_DSA_PREFILL_ATTN_IN_CUDA_GRAPH = EnvBool(True)
+    # With the above: also capture the indexer's Q/K prep and index-K store;
+    # only its logits + top-k stay an eager break.
+    SGLANG_ENABLE_DSA_PREFILL_INDEXER_PREP_IN_CUDA_GRAPH = EnvBool(True)
     SGLANG_DSA_HIP_DISABLE_PRESHUFFLE = EnvBool(False)
     SGLANG_DSA_MQA_LOGITS_FREE_MEM_FRACTION = EnvFloat(0.2)
     SGLANG_ENABLE_PCG_DSV2_DUAL_STREAM = EnvBool(False)

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -31,6 +31,7 @@
     GRAPH_WEIGHTS_PROJ_LORA_ERROR,
     _is_in_piecewise_or_breakable_cuda_graph,
     bcg_dsa_indexer_prefill_split,
+    bcg_dsa_indexer_prefill_topk_split,
     pcg_dsa_indexer_prefill_split,
 )
 from sglang.srt.layers.attention.dsa.paged_mqa_logits_backend import (
@@ -848,6 +849,26 @@ def _gfx950_fused_decode(
             rows=rows,
         )
 
+    def _prefill_prep_in_graph(
+        self, forward_batch: ForwardBatch, num_rows: int
+    ) -> bool:
+        # Same gate as the captured prefill attention, whose padded-row
+        # conventions (zero-padded out_cache_loc, dummy slot 0) this relies on.
+        from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend
+
+        if not (
+            self.use_dsa_indexer_fusion
+            and envs.SGLANG_ENABLE_DSA_PREFILL_INDEXER_PREP_IN_CUDA_GRAPH.get()
+            and not self._aiter_fused_fp8_active(forward_batch)
+            and not self.dsa_enable_prefill_cp
+            and not hasattr(get_token_to_kv_pool(), "_is_layer_owned")
+        ):
+            return False
+        backend = get_attn_backend()
+        return isinstance(
+            backend, DeepseekSparseAttnBackend
+        ) and backend.can_run_prefill_attention_in_graph(forward_batch, num_rows)
+
     def _should_skip_logits_computation(self, forward_batch: ForwardBatch) -> bool:
         # topk_transform selects every valid page slot when kv_len <= index_topk;
         # logits are unnecessary in that case.
@@ -2075,22 +2096,47 @@ def forward_cuda(
                 topk_result = torch.empty(
                     (0, self.index_topk), device=x_meta.device, dtype=torch.int32
                 )
-            graph_dispatch_fn = (
-                bcg_dsa_indexer_prefill_split
-                if is_in_breakable_cuda_graph()
-                else pcg_dsa_indexer_prefill_split
-            )
-            # ROCm aiter fused RMSNorm+quant hands over (x, x_s[, x_bf16]).
-            x_parts = x if isinstance(x, tuple) else (x,)
-            graph_dispatch_fn(
-                layer_id=layer_id,
-                x=x_parts[0],
-                q_lora=q_lora,
-                positions=positions,
-                topk_result=topk_result,
-                x_s=x_parts[1] if len(x_parts) > 1 else None,
-                x_bf16=x_parts[2] if len(x_parts) > 2 else None,
-            )
+            if (
+                return_indices
+                and not isinstance(x, tuple)
+                and self._prefill_prep_in_graph(forward_batch, num_rows=x.shape[0])
+            ):
+                # Row-wise Q/K prep and index-K store over every bucket row
+                # (padded rows store to the dummy slot); only logits + top-k break.
+                # CUDA hands over plain bf16 x; the ROCm (x, x_s[, x_bf16]) tuple
+                # takes the whole-indexer split below.
+                q_fp8, weights = self._fused_q_prepare_and_store(
+                    x,
+                    q_lora,
+                    positions,
+                    forward_batch,
+                    layer_id,
+                    act_quant,
+                    enable_dual_stream=False,
+                )
+                bcg_dsa_indexer_prefill_topk_split(
+                    layer_id=layer_id,
+                    q_fp8=q_fp8,
+                    weights=weights,
+                    topk_result=topk_result,
+                )
+            else:
+                graph_dispatch_fn = (
+                    bcg_dsa_indexer_prefill_split
+                    if is_in_breakable_cuda_graph()
+                    else pcg_dsa_indexer_prefill_split
+                )
+                # ROCm aiter fused RMSNorm+quant hands over (x, x_s[, x_bf16]).
+                x_parts = x if isinstance(x, tuple) else (x,)
+                graph_dispatch_fn(
+                    layer_id=layer_id,
+                    x=x_parts[0],
+                    q_lora=q_lora,
+                    positions=positions,
+                    topk_result=topk_result,
+                    x_s=x_parts[1] if len(x_parts) > 1 else None,
+                    x_bf16=x_parts[2] if len(x_parts) > 2 else None,
+                )
             result = _broadcast_indexer_topk_from_rank0(
                 topk_result if return_indices else None
             )

python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py

@@ -172,3 +172,47 @@ def pcg_dsa_indexer_prefill_split(
 
 
 bcg_dsa_indexer_prefill_split = eager_on_graph(True)(pcg_dsa_indexer_prefill_split)
+
+
+def dsa_indexer_prefill_topk_split(
+    layer_id: int,
+    q_fp8: torch.Tensor,
+    weights: torch.Tensor,
+    topk_result: torch.Tensor,
+) -> None:
+    # Eager top-k of the BCG prefill indexer whose Q/K prep and index-K store
+    # were captured over every bucket row (q_fp8 / weights padded to the
+    # bucket). Logits are sized by the batch's ragged KV, so this stays a break;
+    # it fills the real rows of the static topk_result.
+    forward_context = get_tc_piecewise_forward_context()
+    forward_batch = forward_context.forward_batch
+    indexer = forward_context.dsa_indexers[layer_id]
+    metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch)
+    if indexer._should_skip_logits_computation(forward_batch):
+        # Every key is selected (kv_len <= top-k): the top-k kernel's fast path
+        # turns zero logits into the identity selection.
+        seq_lens_expanded = metadata.get_seqlens_expanded()
+        dummy_logits = torch.zeros(
+            seq_lens_expanded.shape[0],
+            indexer.index_topk,
+            dtype=torch.float32,
+            device=q_fp8.device,
+        )
+        raw_topk_result = metadata.topk_transform(dummy_logits, indexer.index_topk)
+        topk_result[: raw_topk_result.shape[0]] = raw_topk_result
+        return
+    num_tokens = forward_batch.extend_num_tokens
+    indexer._get_topk_ragged(
+        False,
+        forward_batch,
+        layer_id,
+        q_fp8[:num_tokens],
+        weights[:num_tokens],
+        metadata,
+        topk_result,
+    )
+
+
+bcg_dsa_indexer_prefill_topk_split = eager_on_graph(True)(
+    dsa_indexer_prefill_topk_split
+)