[mattteochen]

Motivation

With the breakable prefill CUDA graph, every DSA layer's sparse attention and indexer still run as eager graph breaks, and the NextN (MTP) draft extend that follows each prefill chunk runs eagerly. Short prefills and prefix-cache hits are then bound by host launch work rather than GPU time. This records the DSA prefill attention, the indexer's Q/K prep and the NextN draft extend inside the prefill graph, which cuts TTFT on GLM-5.2 NVFP4 (GB300) by up to 44% for short fresh prompts and by up to 37% for cached-prefix prefill under concurrency, with outputs bitwise equal to the eager breaks.

Modifications

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

• DSA prefill attention (k_b bmm, RoPE + FP8, KV write, TRT-LLM sparse MLA) runs inside the breakable prefill graph. Per-row lengths come from a static buffer that init_forward_metadata refreshes; padded rows attend to the dummy slot. Buckets below the smallest one whose rows all exceed SMs / 2 (96 on GB300) keep the eager break, so every prefill stays bitwise equal to it.
• With it, the indexer's Q/K prep and index-K store are captured too; only logits + top-k stay a break.
• DeepseekModelNextN exposes its decoder as a one-layer layers view (no DP attention, PP or prefill CP), so the draft extend replays the prefill graph. The DSA index-share seed is published from a graph break at replay; multimodal batches stay eager.
• Off for DCP, DP attention, prefill CP, hisparse, k-pool, TBO, page-sharded KV caches ([Feat] KVShard: Add attention and indexer support for page-sharded DSA caches #44280) and per-layer KV loaders (FlexKV, external cache linker, radix cache plug-ins, host-pool retraction backup). HiCache load-backs are waited for in full before the forward.
• The in-graph attention writes and reads KV through the typed KV views (kv_layers.layer(id).store_mla, kv_layers.read(id).key(), [Refactor] attention: Use typed KV views and shared layer routing #46990), like _forward_trtllm.
• Tests: CPU unit tests for the bucket floor, the length refresh, the HiCache wait and the gates (test/registered/unit/layers/attention/test_dsa_prefill_cuda_graph.py, test/registered/unit/models/test_deepseek_nextn_draft_prefill_graph.py); a GB300 nightly e2e test (test/registered/e2e/models/test_dsa_glm52_nvfp4_prefill_graph.py) comparing captured vs eager servers bitwise, with and without HiCache load-back. Env var docs updated.

Kill switches (default on): SGLANG_ENABLE_DSA_PREFILL_ATTN_IN_CUDA_GRAPH, SGLANG_ENABLE_DSA_PREFILL_INDEXER_PREP_IN_CUDA_GRAPH, SGLANG_ENABLE_NEXTN_DRAFT_PREFILL_CUDA_GRAPH.

Accuracy Tests

GSM8K on GB300 with GLM-5.2 NVFP4: full test set (1319 questions), 5-shot, greedy, real acceptance. Same harness as #46630 (few-shot /generate, prompts and answer extraction of sglang.test.few_shot_gsm8k, one question sent first so both arms read the same 5-shot prefix KV). Runs per arm: 2 orderings x 2 server instances.

	TP4, EAGLE 4/1/5	TP8, EAGLE 5/1/6
main	94.81% (accept length 3.743)	94.03% (accept length 4.020)
this PR	94.73% (accept length 3.737)	94.33% (accept length 4.012)

The run-to-run range per arm is 0.3-1.3 pp.

Speed Tests and Profiling

Setup: GB300, GLM-5.2 NVFP4, HiCache on, chunked prefill 8192, default prefill graph capture size. The 2026-10-26 SGLang nightly image (lmsysorg/sglang:nightly-dev-cu13-20261026-8252e3af, SGLang 8252e3af63) is the main arm; the other arm is this branch on top of it. TP4 is one tray (EAGLE 4/1/5, max 8 running requests); TP8 is two trays in one NVL72 rack (EAGLE 5/1/6, max 16).

• Both arms ran in the same allocation, interleaved main / PR / PR / main / main / PR (3 server instances per arm).
• Prompts are random tokens; "fresh" prompts have no cache hit, "cached" ones reuse a warmed 32k / 128k prefix and add 1024 or 4096 new tokens. 64-token answers, 16 requests at concurrency 1 and 8 per lane above it. Acceptance is pinned to the recipe's accept length (3.33 / 3.61) with real draft tokens, so both arms do the same decode work.
• Median TTFT pooled over the 3 instances. Bold: at least 3% and the per-instance medians of the two arms do not overlap.

prompt (new tokens)	conc.	TP4 TTFT ms main → PR	TP4	TP8 TTFT ms main → PR	TP8
fresh, 512	1	71 → 49	-30%	74 → 41	-44%
fresh, 512	4	104 → 78	-25%	87 → 72	-17%
fresh, 1024	1	71 → 62	-13%	74 → 52	-30%
fresh, 1024	4	95 → 94	-2%	98 → 96	-2%
fresh, 2048	1	91 → 90	-1%	78 → 76	-2%
fresh, 2048	4	213 → 161	-24%	182 → 183	+0.5%
fresh, 8192	1	296 → 293	-1%	264 → 262	-0.9%
fresh, 8192	4	581 → 389	-33%	515 → 508	-1%
32k cached + 1024	1	85 → 83	-2%	88 → 73	-17%
32k cached + 1024	4	230 → 225	-2%	216 → 210	-3%
32k cached + 1024	16	612 → 430	-30%	531 → 502	-5%
32k cached + 4096	1	209 → 208	-0.8%	184 → 184	-0.3%
32k cached + 4096	4	483 → 307	-37%	430 → 385	-10%
32k cached + 4096	16	1347 → 1072	-20%	1216 → 972	-20%
128k cached + 1024	1	125 → 124	-0.6%	116 → 115	-0.8%
128k cached + 1024	4	304 → 297	-2%	287 → 285	-0.8%
128k cached + 1024	16	706 → 595	-16%	686 → 538	-22%
128k cached + 4096	1	302 → 300	-0.5%	278 → 279	+0.2%
128k cached + 4096	4	707 → 570	-19%	626 → 467	-25%
128k cached + 4096	16	2108 → 1678	-20%	1895 → 1543	-19%

Gains come from fewer host-bound graph breaks: at concurrency 1 they are largest where the extend is short (fresh 512-1k; at TP8 also 32k + 1k). Under concurrency, shorter prefill steps cut the queueing of the other lanes.

Steady decode (fresh 1000 in / 512 out, median TPOT ms, main → PR):

	concurrency 1	concurrency 4
TP4	2.87 → 2.86	4.61 → 4.61
TP8	2.44 → 2.45	3.75 → 3.71

64-token answers. TPOT is not a per-step measure here, and two effects show up in it:

• At concurrency 1, TPOT of fresh 512 / 1k prompts is +1.7-2.2%: most likely because the first token now leaves before the GPU tail of the prefill step (draft extend), so ~2-3 ms move from TTFT into the first inter-token gap. E2E still falls (TP4 -8% / -2%, TP8 -15% / -9%).
• At concurrency 4-16, TPOT rises (up to 2x) because faster prefills let the lanes drift out of lockstep, and the other lanes' prefills then land in a request's decode window instead of before its first token. E2E moves by -3.6% to +1.0% at these points, except at TP4 fresh 512 and 2048 at concurrency 4. There, E2E is +9% and +5% with the client's default lane offset, which depends on each arm's latency. With the same fixed offset for both arms, it is -4% and +3%. At 2048, the four prefills on main batch into one 8192-token chunk and the lanes stay in lockstep. At TP8, the same two points are -2.7% and -0.5%.

Graph memory (rank 0): target prefill graph 4.04 → 3.61 GB (TP4) and 3.79 → 3.37 GB (TP8), plus 0.25 GB for the draft prefill graph; the KV pool size is unchanged.

Nov-4 re-check on lmsysorg/sglang:nightly-cu134-20261104-27aedf7 (Nov-4 SGLang nightly image, current main), this branch rebased onto it. TP4, one server instance per arm, real acceptance; the e2e test TestGLM52NVFP4DsaPrefillCudaGraph passes (4 tests). Median TTFT ms, main → PR:

prompt (new tokens)	concurrency 1	concurrency 4
fresh, 512	51.6 → 41.9 (-19%)	101.7 → 58.8 (-42%)
fresh, 1024	56.3 → 55.0 (-2%)	104.5 → 68.0 (-35%)
32k cached + 4096	207.0 → 205.7 (-0.7%)	325.9 → 230.8 (-29%)

Main is faster at concurrency 1 than on the Oct-26 image, so the concurrency-1 gains are smaller. GSM8K (1319, 5-shot, parallel 4): main 94.69% (accept length 3.757), PR 94.47% (3.762).

Reproduce

# Server, TP4 (one GB300 tray), EAGLE 4/1/5
python -m sglang.launch_server --model-path nvidia/GLM-5.2-NVFP4 --trust-remote-code --tp-size 4 \
  --quantization modelopt_fp4 --kv-cache-dtype fp8_e4m3 --bf16-gemm-backend cutedsl \
  --speculative-algorithm EAGLE --speculative-eagle-topk 1 --speculative-num-steps 4 --speculative-num-draft-tokens 5 \
  --chunked-prefill-size 8192 --max-prefill-tokens 8192 --mem-fraction-static 0.83 \
  --max-running-requests 8 --cuda-graph-max-bs-decode 8 \
  --enable-hierarchical-cache --hicache-size 100 --hicache-write-policy write_back --hicache-io-backend direct \
  --hicache-mem-layout page_first_direct
# TP8 (two trays, MNNVL), EAGLE 5/1/6: --tp-size 8 --nnodes 2 --node-rank {0,1} --dist-init-addr <node0>:5757 \
#   --fp4-gemm-backend flashinfer_trtllm --speculative-num-steps 5 --speculative-num-draft-tokens 6 \
#   --mem-fraction-static 0.8 --max-running-requests 16 --cuda-graph-max-bs-decode 16
# Pinned acceptance for the speed runs: SGLANG_SIMULATE_ACC_LEN=3.33 (TP8: 3.61) SGLANG_SIMULATE_ACC_METHOD=match-expected \
#   SGLANG_SIMULATE_ACC_TOKEN_MODE=real-draft-token
# Baseline on this branch: SGLANG_ENABLE_DSA_PREFILL_ATTN_IN_CUDA_GRAPH=0 SGLANG_ENABLE_NEXTN_DRAFT_PREFILL_CUDA_GRAPH=0

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Tests on GB300:

• New CPU unit tests (15) pass. The new e2e class passes at TP4 (one tray) and TP8 (two trays): captured vs eager servers and HiCache load-back vs HiCache eager give bitwise-equal input logprobs and greedy ids on 13 requests, the captured paths engage, and the MTP accept length stays within 0.1.
• Existing prefill-graph, DSA, NextN and speculative unit tests: 287 passed / 6 skipped on the CPU suite and 292 passed on GPU, the same as main. The one GPU failure (test_deepseek_v41_bcg_multimodal.py) fails on main too.

Related:

• [Spec] Capture the prefill CUDA graph for DeepSeek-style NextN drafts #43982 also captures the NextN draft prefill; this version gates DP attention / PP / prefill CP off, publishes the DSA index-share seed at replay and keeps multimodal batches eager.
• [Perf] dsa: Prefill indexer TP row split and per-request page-table rows (opt-in) #42006 (prefill indexer TP row split) is complementary: it speeds up the logits + top-k break that this PR leaves eager. An own row split was dropped from this PR in favour of it.
• With page-sharded KV caches ([Feat] KVShard: Add attention and indexer support for page-sharded DSA caches #44280) the captured path stays off; it is not validated there.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #42039509434
Latest PR Test (Extra): ❌ Run #42039509462
Latest PR Test (AMD ROCm 10): ❌ Run #42039509342

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

[mintlify]

Preview approval required

This pull request comes from a fork, so Mintlify won't build a preview until someone with write access approves commit 81a9bd4.

Approve preview

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[nvpohanh]

[by Claude Code] This review is done by agents. Please feel free to disagree if they do not make sense.