[mattteochen]

Motivation

With MTP on DSA models (GLM-5.x, DeepSeek-V3.2), every decode step pays costs that scale with --context-length and the vocab rather than with the live work:

• The replay metadata kernels launch one program per 128 columns of the context-wide page table: 8193 per row at the default 1M context.
• The MTP draft-decode replay precomputes into a full-width scratch table, then copies it into every step backend.
• Each greedy draft step all-gathers rows x vocab logits only to take an argmax.

Modifications

• [DSA] Bound the replay metadata kernels by the live KV length.
  • Adds bounded variants of the decode, target-verify and draft-extend replay metadata kernels (kernels/ops/attention/dsa_metadata_bounded.py).
  • Each request gets at most min(budget / rows, 256) programs, which stride over its live column blocks.
  • The grid depends only on (rows, width, budget), so it stays CUDA-graph safe.
  • Outputs are bit-identical to the unbounded kernels.
  • Used on CUDA with kpool 1 and without KPool metadata fusion.
• [DSA] Write the MTP draft-decode replay metadata in one bounded launch.
  • DeepseekSparseAttnMultiStepBackend writes each step backend's graph metadata straight from req_to_token, in one bounded launch for up to 4 steps. This replaces the precompute + full-width copies.
  • Draft steps seeded with the draft-extend top-k no longer run their DeepGEMM schedule refresh, which they never used.
  • Enabled per captured bs, only when the step backends' tables match.
• [Spec] Greedy EAGLE drafts: exchange vocab-shard argmax partials instead of gathering logits.
  • For topk=1 drafts, each TP rank reduces its LM-head shard to per-split argmax candidates, 4 exact fp32 lanes each.
  • The TP sum all-reduce exchanges the candidates. A one-warp finalize then picks the token with the same tie and NaN rules as the gathered path.
  • Tokens, topk_p and positions are bit-identical, inside the draft-decode and draft-extend CUDA graphs.
  • Draft LogitsProcessors opt in through enable_draft_argmax_exchange().
  • The gather is kept for non-vocab-parallel heads, softcapping, DP attention, and when the all-reduce would run on NCCL.

Flags (kill switches, both default on):

• SGLANG_DSA_BOUNDED_REPLAY_METADATA=0 turns off the first two commits.
• SGLANG_SPEC_DRAFT_ARGMAX_EXCHANGE=0 turns off the third.

Related:

• [Perf] DSA: Share one metadata buffer set across MTP draft steps #44358 shares one metadata buffer set across MTP draft steps; if it lands, the second commit reduces to a single-destination bounded write.
• perf(spec): add vocab-parallel MTP local argmax #37106 and [qwen 3.8 next] Fuse exact greedy chain verification with argmax finalization #41172 take other approaches to the greedy-draft argmax.
• [Perf] speculative: opt-in step CUDA graph and fused verify / relay paths for EAGLE v2 chains #46211: its _fused_draft_extend_tail_eligible reads next_token_logits, which is None when the exchange is engaged, so it should decline when next_token_argmax_partials is set.

Accuracy Tests

GSM8K on GB300 with GLM-5.2 NVFP4: full test set (1319 questions), 5-shot, greedy, 512 new tokens.

• Harness: few-shot /generate, with the same prompts and answer extraction as sglang.test.few_shot_gsm8k. One question is sent first, so both arms read the same 5-shot prefix KV from the radix cache.
• Runs per arm: 2 orderings x 2 server instances.

	TP4, EAGLE 4/1/5	TP8, EAGLE 5/1/6
main	94.48% (accept length 3.737)	94.14% (accept length 4.025)
this PR	94.33% (accept length 3.740)	94.05% (accept length 4.026)

Speed Tests and Profiling

Setup: GB300, GLM-5.2 NVFP4. 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.

End to end:

• Workload: 1024 in / 1024 out, prompts from ShareGPT text (--dataset-name random), greedy, real acceptance, default 1M context.
• Both arms ran in the same allocation, interleaved main / PR / PR / main, with 2 runs per server instance (n = 4 per point).
• tok/s/user = 1000 / median TPOT. Spread = (max - min) / median over the 4 runs. The run ranges of the two arms do not overlap at any point.

TP4 (one tray), EAGLE 4/1/5:

concurrency	tok/s/user main	this PR	gain	spread main / PR	output tok/s main → PR	median TTFT ms main → PR	accept length
1	528.8	533.0	+0.8%	0.03% / 0.06%	471 → 475	72 → 74	4.56
2	428.5	432.1	+0.8%	0.11% / 0.13%	749 → 753	75 → 76	4.56
4	328.7	333.3	+1.4%	0.91% / 0.69%	1178 → 1192	93 → 103	4.75
8	251.0	255.7	+1.9%	1.10% / 0.84%	1722 → 1755	124 → 124	4.68
16	190.8	195.1	+2.2%	0.32% / 0.28%	2550 → 2605	178 → 156	4.54 / 4.50
32	138.4	142.9	+3.2%	0.75% / 0.59%	3977 → 4105	151 → 151	4.67

TP8 (two trays, one NVL72 rack), EAGLE 5/1/6:

concurrency	tok/s/user main	this PR	gain	spread main / PR	output tok/s main → PR	median TTFT ms main → PR	accept length
1	680.9	695.2	+2.1%	0.06% / 0.02%	547 → 557	72 → 73	4.94
4	461.9	473.7	+2.6%	0.31% / 0.33%	1580 → 1615	79 → 80	5.58
16	282.3	293.4	+3.9%	2.37% / 2.20%	3717 → 3782	139 → 154	5.29

Kernels. Times are CUDA-graph replays in us. The metadata kernels are timed with a cold L2.

Replay metadata, 1M-token table, 5 verify rows per request. Unbounded → bounded:

op	live KV	bs 1	bs 4	bs 16
decode	1k	5.1 → 1.8	17.4 → 2.2	67.2 → 2.8
decode	64k	5.5 → 2.8	17.7 → 3.5	67.3 → 5.1
target verify	1k	21.9 → 2.3	83.6 → 2.6	332.1 → 2.9
target verify	64k	22.2 → 3.1	84.0 → 3.9	331.7 → 7.9
draft extend	1k	6.4 → 2.4	18.8 → 2.7	68.3 → 2.9
draft extend	64k	6.5 → 3.2	19.1 → 4.4	68.6 → 9.6

MTP draft-decode metadata. Precompute + copies (plus, at 5/1/6, the step-3 schedule refresh) → one launch:

step backends	live KV	bs 1	bs 4	bs 16
3 (EAGLE 4/1/5)	1k	11.9 → 2.4	33.6 → 2.4	123.1 → 2.7
3 (EAGLE 4/1/5)	64k	11.7 → 3.1	33.6 → 3.9	123.6 → 6.5
4 (EAGLE 5/1/6)	1k	21.0 → 2.3	47.2 → 2.6	168.6 → 3.1
4 (EAGLE 5/1/6)	64k	20.9 → 3.0	47.2 → 3.9	169.0 → 6.9

Draft step (LM head + top-1). Logits gather → partials exchange; vocab 154880, hidden 6144, FlashInfer all-reduce, max over ranks:

rows	1	2	4	8	16
TP4	91.4 → 82.1	92.0 → 82.1	93.4 → 83.0	96.4 → 83.5	103.9 → 83.8
TP8 (two trays)	87.2 → 65.7	95.8 → 63.3	89.1 → 60.5	95.9 → 59.0	112.3 → 59.6

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 \
  --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 32 --cuda-graph-max-bs-decode 32
# TP8 (two trays, MNNVL), EAGLE 5/1/6: --tp-size 8 --nnodes 2 --node-rank {0,1} --dist-init-addr <node0>:5757 \
#   --speculative-num-steps 5 --speculative-num-draft-tokens 6 --mem-fraction-static 0.8 \
#   --max-running-requests 16 --cuda-graph-max-bs-decode 16
# Baseline on this branch: SGLANG_DSA_BOUNDED_REPLAY_METADATA=0 SGLANG_SPEC_DRAFT_ARGMAX_EXCHANGE=0

# Warm-up: for every c in 1..max-running-requests, send c concurrent 1k-token requests with 48 new tokens
# (the Triton metadata kernels compile once per batch size). Then, for each concurrency C:
python -m sglang.benchmark.serving --backend sglang --dataset-name random \
  --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
  --random-input-len 1024 --random-output-len 1024 --random-range-ratio 1 \
  --num-prompts $(( C * 6 > 16 ? C * 6 : 16 )) --max-concurrency $C --seed 1 --flush-cache --warmup-requests 1

# Kernels
python test/registered/kernels/benchmark/attention/bench_dsa_metadata_bounded.py
python test/registered/kernels/benchmark/speculative/bench_draft_argmax_exchange.py --num-gpu 4
torchrun --nnodes 2 --nproc-per-node 4 --node-rank {0,1} --master-addr <node0> --master-port 29511 \
  test/registered/kernels/benchmark/speculative/bench_draft_argmax_exchange.py

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.

New tests:

• test/registered/kernels/ops/attention/test_dsa_metadata_bounded.py and test_dsa_draft_decode_metadata.py: bit-identity against the unbounded kernels and against precompute + copies at a 1M-token table; CUDA-graph replay with new lengths; backend selection; the one-launch gate; the schedule-refresh steps.
• test/registered/kernels/ops/speculative/test_draft_topk1_partials.py: partials + sum all-reduce against the gathered logits, covering ties, NaN, inf and padding, with TP 2-8 emulated on one GPU.
• test/registered/kernels/ops/speculative/test_draft_argmax_exchange_tp.py: a real TP group, covering an exact all-reduce, LogitsProcessor forward and the CUDA graph.
• test/registered/unit/spec/test_draft_argmax_exchange.py: the opt-in protocol.

Results on GB300:

• All new tests pass.
• The existing DSA metadata, speculative and logits-processor tests (72 GPU and 613 unit cases) pass on both main and this branch.
• The only change to an existing test adds next_token_argmax_partials=None to the fake LogitsProcessorOutput in test_eagle_draft_sampling.py.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #40379582163
Latest PR Test (Extra): ❌ Run #40379582127
Latest PR Test (AMD ROCm 10): ❌ Run #40379581904