[abmfy]

Overview

Lowers the step time of the DeepSeek-V4.1-Flash latency_specialized decode at TP4, concurrency 1, with DSpark on GB300. All changes are unconditional: no new env vars or config options, and they also apply above concurrency 1 (see Limitations).

Stacked on #61869: one commit, ebcc07f, on its head bc6b0f2, which has merged main 5415568 (with #63921). Review that commit. The two Model Runner V2 changes that apply beyond this executor (the FULL graph's sliced output copy and the skipped draft-token gather) moved to #64701. This PR no longer changes the shared runner or speculators: #61869 removed the post-forward CUDA graphs on review, and the changes here that extended them went with them. Outside the executor it adds skip_topk_buffer_clear (default False) to the two shared sparse indexers, optional arguments to query_quant.py and all_reduce_mhc.py that default to the old behaviour, and a package-data line in setup.py.

Claims

• AgentX P90 interactivity 891.9 → 964.0 tok/s (+8.1%), median 921.3 → 999.3 tok/s (+8.5%), ITL p50 1.085 → 1.001 ms: the mean of two interleaved 1500 s runs per arm on one GB300 tray. Every base/PR pair falls within +7.6% to +8.6% (P90) and +8.2% to +8.7% (median).
• The new top-k kernel (topk_bf16_small) is not bitwise equal to the code it replaces (the same set up to ties at the k-th value, in another order). GSM8K reads 0.45 pt lower with this PR (0.9050 → 0.9005, four runs per arm, every PR run below every base run). The evidence for a real loss is mixed: a per-question sign-flip test gives p = 0.15, a permutation test on the run totals p = 0.029 (0.083 exchanging runs only within jobs), and the 95% interval on the difference is about −1.0 to +0.1 pt. The other changes that touch a kernel's arithmetic are bitwise equal to the code they replace; the rest moves work between streams, changes launch timing or the routed GEMMs' CTA count, or adds a read-only L2 prefetch.

Validation

Setup: 1 GB300 tray (4 GPUs), TP4, --max-num-seqs 1, --numa-bind, DSpark k=5, synthetic acceptance length 3.51. vLLM nightly image dd4e401 (FlashInfer 0.7.2.post1) with each arm's Python package mounted over it, both without #62172's Python hunks (the image predates its kernels). AgentX 1.0.12, 1500 s, seed 42.

Arm	P90 tok/s	Median tok/s	ITL p50 (ms)
Base: #61869 (bc6b0f2), run 1	894.2	923.2	1.0832
Base, run 2	889.6	919.5	1.0876
This PR (ebcc07f), run 1	966.1	999.0	1.0010
This PR, run 2	961.8	999.5	1.0005

One job on one tray, in the order PR, base, PR, base; scored on the 193 requests with an inter-token latency that the base runs completed, all of which every run also completed. Mean ± sd of two runs per arm: P90 891.9 ± 3.3 → 964.0 ± 3.0 tok/s (t ≈ 23), median 921.3 ± 2.6 → 999.3 ± 0.3 tok/s (t ≈ 42). AgentX runs at temperature 1 with synthetic rejection. Apart from #62172's Python hunks, each arm's package has the same vllm/ tree as its commit. The specialized decode re-tunes its FlashInfer GEMMs at every start, which can move the step time by up to about 1%; both arms loaded the same pinned tactics before tuning through a benchmark-side hook (not part of this PR).

On the previous stack (#63921 + #61869 with its post-forward graphs, nightly 3966924), an earlier revision of this PR that also extended those graphs gave 907.5 → 986.5 and 986.1 / 936.0 → 1017.0 and 1017.4. Both arms read about 1.6–2.3% lower on this base; that comparison crosses trays, nightlies and the removal of the post-forward graphs, and is not attributed further.

Accuracy (GSM8K, real DSpark acceptance, 5-shot, 1319 questions, temperature 0; one run's standard error about 0.8 pt):

Build	Run 1	Run 2	Median acceptance length
Base: #61869	0.9067	0.9030	3.56
This PR	0.8999	0.8992	3.56–3.57

Paired by question, the base and this PR differ in which questions they get right on 59–75 questions per run pair (only the base right / only this PR right: 34/25, 35/25, 34/30, 40/35), about as many as two runs of one build (base 33/28, this PR 35/34). McNemar exact p = 0.25 to 0.71.

Two more runs per arm, in earlier jobs (trays 11, 15 and 17), on an earlier merge of main and #61869 with this PR's changes (the same diff apart from one AGENTS.md line), read 0.9075, 0.9030 → 0.9014, 0.9014. Over all eight runs every PR run reads below every base run, 0.9050 → 0.9005 (−0.45 pt). A sign-flip test on each question's right counts gives p = 0.15 (24 fewer right answers out of 5276; 133 questions have unequal counts; paired standard error 0.30 pt, 95% interval −1.0 to +0.1 pt). A permutation test on the eight run totals gives p = 0.029 (2 of 70 splits; 0.083 exchanging runs only within jobs). The data are inconclusive: they do not establish a loss, and they do not rule out one of up to about 1 pt. Earlier revisions of this PR on other stacks read 0.896 and 0.898, with no base run on those stacks to pair them with.

topk_bf16_small is the only change here with its own selection rule: on layers 24, 28, 32 and 36 it replaces DeepSelect's kernel, and among columns equal to the k-th value it keeps the lowest-numbered ones, which need not be the ones DeepSelect keeps. decode_topk on layer 20 keeps its rule, but it breaks ties by CTA arrival order in both arms, so this PR's launch changes there can change which tied columns it keeps. Real indexer logits hold few distinct BF16 values, so the k-th value may often be tied; neither that rate nor either effect was measured. GSM8K is the only eval run on this PR.

Tests (GB300):

# 225 passed
pytest tests/models/deepseek_v41/nvidia/specialized_decode/test_decode_topk.py \
  tests/models/deepseek_v41/nvidia/specialized_decode/test_sparse_mla_cute.py \
  tests/kernels/core/test_fused_q_kv_rmsnorm.py
# 4 GPUs: 200 passed
pytest tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py \
  tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py \
  tests/models/deepseek_v41/nvidia/specialized_decode/test_all_gather.py \
  tests/models/deepseek_v41/nvidia/specialized_decode/test_layers.py

Run on ebcc07f with nightly dd4e401's image, the branch's Python package mounted over it, without #62172's Python hunks (the image predates #62172's kernels, so it cannot run them). test_numerics.py ran with a stand-in for ray, which the image lacks and which it only imports. test_shared_query_quant_preserves_projection fails on main on GB300 (24 cases, AttributeError: params_dtype: #64111 made the kernel read an attribute the test's hand-built layer lacks), so this PR sets it there.

test_layers runs 5 layers with random weights through the executor and compares each with the native path within tolerance. The new tests compare the fused query RoPE, the KV insert in the norm and the local-index attention bitwise with the kernels they replace, and topk_bf16_small exactly with its specified output (torch.topk's k-th value, the selected columns ascending, ties toward the lowest column).

CI. No CI agent has four GB300s. The nightly 4xB200 step runs the kernel, numerics, top-k and sparse-MLA tests, and test_layers skips there; tests/kernels/core runs on H200, where the new KV-insert test skips. The GB300 runs above cover that gap.

Serve and benchmark commands

VLLM_USE_V2_MODEL_RUNNER=1 VLLM_USE_RUST_FRONTEND=1 vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
  --tensor-parallel-size 4 --max-num-seqs 1 --numa-bind \
  --language-model-only \
  --tokenizer-mode deepseek_v41 --tool-call-parser deepseek_v41 --enable-auto-tool-choice --reasoning-parser deepseek_v41 \
  --engram-config '{"cpu_offload":true,"use_thp":true}' \
  --additional-config '{"decode_backend":"latency_specialized"}' \
  --kernel-config '{"enable_flashinfer_autotune":true}' \
  --attention-config '{"backend":"FLASHINFER_MLA_SPARSE_DSV41","indexer_kv_dtype":"mxfp4","indexer_sparse_logits":true}' \
  --kv-cache-dtype fp8 \
  --speculative-config '{"method":"dspark","num_speculative_tokens":5,"draft_sample_method":"probabilistic","rejection_sample_method":"synthetic","synthetic_acceptance_length":3.51,"enable_adaptive_verification":false}' \
  --max-model-len 1048576 --prefill-schedule-interval 1 \
  --compilation-config '{"mode":"VLLM_COMPILE","cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[5,6,12,18,24,30,36,48,60,72,96,120,144,192,240,288,384,480,576,768,1020,1536,2046,3072,4092]}' \
  --max-cudagraph-capture-size 4092 --max-num-batched-tokens 4096

agentx profile --url http://localhost:8000 \
  --model deepseek-ai/DeepSeek-V4.1-Flash --tokenizer <model dir> \
  --concurrency 1 --benchmark-duration 1500 --random-seed 42 \
  --extra-inputs '{"chat_template_kwargs":{"reasoning_effort":"high","thinking":true},"temperature":1}' \
  --tokenizer-trust-remote-code --public-dataset semianalysis_cc_traces_weka_062126

GSM8K used "rejection_sample_method":"block" (real acceptance) and tests/evals/gsm8k/gsm8k_eval.py-style prompting (5-shot, 256 tokens; the client sends 4 at once and --max-num-seqs 1 runs them one at a time).

Details

Where the time went. An nsys trace of an earlier stack (#63921 + #61869) at 3.73 ms per step showed these costs:

• Lamport all-reduce + mHC time exposed between sublayers: 365 µs/step.
• Idle gaps in the PDL chain, where a kernel waited for its predecessor's tail.
• The indexer top-k kernels: about 90 µs/step.
• The MLA launch chain on plain layers: KV insert, query RoPE and quantization between wq_b and the attention.
• Host and graph gaps outside the target layers (not addressed here: the changes that targeted them extended the post-forward graphs).

Per-layer medians on this base, from one trace per arm:

Median	Base (µs)	This PR (µs)
Plain-layer attention half	31.7	27.4
Plain-layer MoE half	38.6	34.0
Indexer-layer attention half	62.4	54.8
Layers 1–39	3015.8	2695.6

Layers 2, 8 and 14 keep the stock indexer's top-k (attention half 68–78 → 61–77 µs); layer 20's block selection moves to the side stream (L20 70.4 → 61.7 µs) and layers 24, 28, 32 and 36 take topk_bf16_small (about 56 → 48 µs). Layer 1 reads slower in the PR's trace (95.3 → 117.9 µs); the total is net of that. The MoE half moves between traces of identical code (31.6 and 34.4 µs per plain layer in two earlier traces), probably with the number of experts the window hits.

The 320 µs saved across the layers (one trace per arm; identical code has moved this sum by up to about 110 µs, in the MoE half) is in line with the roughly 300 µs per step behind the AgentX ITL (0.085 ms per token at acceptance length 3.51), and a fixed 6000-token decode without the tracer reads 1.02–1.07 → 0.95–0.97 ms per token in the same runs. The traced whole step moves less (3757 → 3579 µs): in the PR's trace rank 0's GPU catches up with its host in the metadata chain before the next step, so it runs no kernel for about 216 µs per step after the last layer (59 in the base; gaps over 2 µs, copies not counted as busy), and the other ranks wait for it in the step's first all-reduce (about 170 µs, against 5–13 in the base). Outside the trace window the same request's chunk gaps read 3.76 and 3.71 ms in the base and 3.42 and 3.36 ms in the PR (inside it, 3.75 and 3.59), so that floor binds only under the tracer; whether the untraced step has a lower host floor of its own was not measured.

The changes and what each bought (measured during development, when each was still switchable; single starts without pinned tactics, so differences under about 40 µs/step are within start-to-start tuning noise)

Area	Change	Measured
Boundary	The fused all-reduce + mHC consumers trigger their PDL dependents early: the MXFP8 consumer at its start, the plain one after its poll, the low-batch finalize-publish after its wait. The dependents read only weights before their own wait.	−89 µs/step (probe)
Boundary	Capture a boundary's main-stream successor before the side-stream kernels forked at it; the boundary before the MoE writes the MoE input as MXFP8 for the shared gate_up	exposed AR 287 → 191 µs/step (nsys)
MoE	Persistent routed GEMMs leave 16 SMs to the shared expert	−66 µs/step (probe)
MoE	GEMM1 triggers GEMM2 after 21 k tiles	GEMM2 14.5 → 7.5 µs; −51 to −96 µs/step (probe)
Attention	Query RoPE and FP8 quantization inside the CuTe sparse-MLA kernel (bitwise equal to the CUDA kernel's rounding)	−45 µs/step (probe)
Attention	KV cache write inside the Q/KV norm-quant kernel, so the attention follows wq_b via PDL	plain attention half 28.9 → 27.1 µs; AgentX median +18
Attention	Sparse-MLA loads graph-invariant inputs before its grid wait; side-stream L2 prefetch of the KV rows a later layer reads	−0.7 µs per plain layer (nsys)
Attention	An index source's MLA reads its request-local top-k; no table build in between	top-k end → MLA end 10.85 → 9.18 µs
Indexer	Candidate consumers' top-k and remap in one kernel (topk_bf16_small)	−17 to −20 µs/step (nsys)
Indexer	Candidate-block selection on the candidate side stream, in 2-row launches after the token top-k; early top-k reset	L20 attention half 74.5 → 62.7 µs

Probe = per-step median of a fixed 6000-token decode; nsys = one trace.

Third-party code. topk_bf16_small is adapted from SGLang (Apache-2.0), whose kernel adapts DeepSelect's init-window select (MIT); the file carries both notices. It is JIT-built with torch.utils.cpp_extension at warm-up (on the first start, then cached), like vllm/models/deepseek_v41/amd/ops/topk.py but with no fallback: without nvcc and a host compiler the executor's warm-up fails (its AGENTS.md states the requirement). An earlier revision of this PR also vendored a bf16 cluster top-k for the dense indexers' token top-k; it is out, to keep the vendored tree out of this PR, and those layers stay on decode_topk or the stock indexer as in the base. On the previous stack, without #62172, it was worth about 1.5% P90 / 1.2% median in one run against two; it would need measuring again. The ahead-of-time DeepSelect build vLLM pins does not fit as is: its gather signals completion barriers in the wrong CTA (deepseek-ai/DeepSelect#16), and a row holding a NaN gets a marker the sparse-index table build does not mask.

Limitations

• The changes also apply above concurrency 1 (for example, the routed GEMMs leave 16 SMs to the shared expert at every step up to 192 rows). Only concurrency 1 was benchmarked; test_layers covers correctness up to 12 requests.
• The fused query RoPE and the norm's KV write match nvcc's FMA contraction of the CUDA kernels they replace; test_sparse_mla_cute.py (the RoPE) and tests/kernels/core/test_fused_q_kv_rmsnorm.py (the KV write) catch a compiler change that breaks that.
• Both arms ran without [Perf][Kernel] Let cooperative_topk read the BF16 indexer logits #62172's Python hunks, which the branch has, because the image predates its kernels. Their layers 2, 8 and 14 therefore upcast the logits and run the FP32 cooperative top-k (about 38 µs/step in the traces), which [Perf][Kernel] Let cooperative_topk read the BF16 indexer logits #62172 replaces with a BF16 read in both arms; this PR does not touch that path.

Not a duplicate. Related open PRs make other changes, some in the same files: #64146 (TP8 on GB300), #64144 (the fused all-reduce + mHC without NVLink multicast) and #64893 (H200). Searches for latency_specialized, "specialized decode", "DSv4.1 TP4" and "DSv4.1 top-k DeepSelect" found no open PR with these changes.

AI assistance. Claude was used to profile, design, implement and benchmark these changes and to write this description.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @abmfy.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[zhou9402]

Hi! We've been looking at the same draft-token gather overhead and found that this change in model_runner.py is generic: it lives in the shared MRV2 runner and benefits every spec-decode model, not just specialized decode. We'd like to avoid duplicate PRs, so would you mind splitting this part into its own small PR? It'd be easier to review and could land independently of this one.

When you rebase it, note that #64458 changed broadcast_drafts to take a pre-gathered tensor. PP still needs its own full-width, batch-ordered copy, because the broadcast runs on a side stream and the next step can overwrite the source before it's sent. So the slice shouldn't go to PP:

broadcast = self.pp_handler is not None and num_spec_tokens > 0
if broadcast or input_batch.has_structured_output_reqs:
    draft_tokens = self.req_states.draft_tokens[input_batch.idx_mapping]
else:
    draft_tokens = self.req_states.draft_tokens[: input_batch.num_reqs]  # width only
self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens[:, :num_spec_tokens])
if broadcast:
    self.pp_handler.broadcast_drafts(draft_tokens)

A couple of tests would help: check that the placeholder width matches K (0, partial, max), and that PP still sends the full width when K < max.

Happy to help review. Thanks!