[yaoyaoding]

Overview

This PR makes the latency_specialized decode from #61868 faster for DeepSeek-V4.1 on NVIDIA GPUs. Most of the work is in the decode step's kernels: larger batches now stay on the low-latency paths, more of each step overlaps, and the indexers split their work across the TP ranks. "What changed" below goes through each part.

On one GB300 tray with DSpark, the decode step takes 5.9–17.0% less time than on main at batch 1 to 32, and AgentX P90 interactivity is 4–26% higher at concurrency 1 to 32, with slightly higher throughput.

Left: three ABBA sets at 0cdfc566, 12 measurements per arm and batch. Right: AgentX at 0cdfc566, both arms of each concurrency on one tray.

Validation

This branch includes #63914 (merged into this branch), which keeps the specialized decode's FlashInfer paths on main's FlashInfer 0.7.2.post1; this PR's two further paths on FlashInfer internals follow the same rule. The runs below are from before this branch merged main: main at 76518259 on that commit's wheel with FlashInfer 0.7.2, against this PR at 0cdfc566, except compute-sanitizer, which ran at 5ab5ef9e, the commit before it (0cdfc566 adds only a start-up requirement, no runtime weight updates, a test option, a comment and documentation).

After the merge, on main's wheel at ed1a7307 with FlashInfer 0.7.2.post1, four GB300: these GPU tests pass at b7bb822b (test_kernels.py, test_all_gather.py, test_decode_topk.py, test_sparse_mla_cute.py, test_numerics.py, test_layers.py, tests/kernels/test_engram.py, tests/distributed/test_custom_all_reduce.py; all but test_kernels.py ran at d7fcb755, which differs only in that file), and a TP4 server with decode_backend: latency_specialized starts with the specialized decode on and no fallback warning, and answers a completion and a chat request correctly. FlashInfer 0.7.2.post1's MoE finalize adds the top-k products in a different order at small batches, so b7bb822b lets one test's finalize comparison allow FP32's reordering error.

Correctness

Task	main	This PR
GSM8K, strict match (1,319, c8, real verification)	0.9742	0.9735
GPQA-Diamond (198 × 16, c32)	0.789	0.791
AIME25 (30 × 16, c32)	0.983	0.985
RULER NIAH (single_2, multikey_2, multikey_3; 32K, 128K, 256K; 40 samples each), concurrency 8	1.000 in all nine cells (40 of 40 each)	1.000 in all nine cells (40 of 40 each)
RULER NIAH, same, concurrency 32	1.000 in all nine cells (40 of 40 each)	1.000 in all nine cells (40 of 40 each)
GSM8K with VLLM_GPU_SYNC_CHECK=error		0.9719, no GPU↔CPU sync reported
GSM8K with decode_backend unset (stock decode)		0.9704
All-gather stress: CUDA graphs of three back-to-back gathers of different sizes, some with the ranks skewed, replayed 2,000,000 times in a seeded random order on four GPUs; every word encodes its rank, gather and replay, checked on the GPU (test_all_gather.py as a module; --custom-all-gather replays the same graphs on vLLM's custom all-gather)	custom all-gather: 392 stale words	0 bad words (6,000,000 gathers)

Tests

# One SM100-family GPU: the new kernels against FlashInfer, cuBLAS or torch references
.venv/bin/python -m pytest tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py 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/test_engram.py -v
# Four GPUs: the all-gather, AllReduceMHC's LL and two-shot protocols, TP4 numerics
.venv/bin/python -m pytest tests/models/deepseek_v41/nvidia/specialized_decode/test_all_gather.py tests/distributed/test_custom_all_reduce.py tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py -v
# The all-gather stress run (four GPUs), then the same replays on vLLM's custom all-gather
.venv/bin/python -m tests.models.deepseek_v41.nvidia.specialized_decode.test_all_gather --replays 2000000
.venv/bin/python -m tests.models.deepseek_v41.nvidia.specialized_decode.test_all_gather --replays 2000000 --custom-all-gather
# Four GB300: five-layer model at TP4 with DSpark, switch on vs off
.venv/bin/python -m pytest tests/models/deepseek_v41/nvidia/specialized_decode/test_layers.py -v

Check	What it covers	Result
Kernel tests	Routed experts against FlashInfer up to 192 rows, eagerly, in graphs and over 500 replays with alternating inputs (bitwise); the front against the stock chain; pinned and wide low-latency MXFP8 GEMMs; the replicated embedding lookup; the metadata graph against eager builds, with slow clamping builders	299 passed (with tests/kernels/test_engram.py), 1 skipped
Decode top-k (CuTe DSL)	An exact top-k of the BF16 logits against torch.topk, k 512 and 2048, 1–64 rows: sampled and direct paths, crowded, constant and rising values, NaN and infinities, a sample that misses the top values, the wide 16-CTA rows, CUDA graphs; block mode against the stock select_candidate_blocks	77 passed
Sparse MLA (CuTe DSL)	Against an FP32 reference: widths of 1–5 chunks (128 to 640 columns), masked columns, a token with no active column, as many tokens as the GPU holds clusters for, with and without sinks, CUDA graph replay	33 passed
All-gather	Against NCCL, bitwise: int32, BF16 and FP32 inputs of 1 pack to the full 4 MiB, growing and shrinking from call to call, one rank late each call; the sentinel word; CUDA graphs of three gathers of different sizes, some with the ranks skewed, replayed in a random order with new contents each replay, 4 GPUs	Passed; the stress run: 0 bad words in 6,000,000 gathers
AllReduceMHC	Shrinking and growing batches in one mailbox with LL only, two-shot only (to 200 rows) and both in turn; the two protocols give the same bits, 4 GPUs	61 passed (the whole file)
Numerics, layers (4 GB300)	TP4 collectives through the new paths, including the logits and the row-sharded indexer at 1, 2, 5, 8 and 32 requests through the new all-gather; with the switch on the executor runs, embeds from the replicated table and replays its graphs	Both passed
compute-sanitizer	memcheck and synccheck over the new kernels' single-GPU tests	memcheck and synccheck: 0 errors each (199 tests: 110 in the two new test files, 89 selected from test_kernels.py; at 5ab5ef9e)
pre-commit	All hooks	Passes

CI. The new test files (the all-gather's among them), the AllReduceMHC tests (test_custom_all_reduce.py -k all_reduce_mhc) and test_engram.py join the nightly 4×B200 DeepSeek-V4.1 step. Three tests need hardware that step does not have, and skip there: test_pinned_mxfp8_mm and test_wide_low_latency_mxfp8_mm (the tactics are pinned for GB300, SM103) and the layers test (four GB300).

Performance

All runs: DeepSeek-V4.1-Flash on one GB300 tray (TP4), DSpark (5 draft tokens), Model Runner V2, FULL CUDA graphs for decode. The decode step and AgentX use DSpark's synthetic acceptance (length 3.51, as InferenceX runs AgentX); the accuracy runs above verify for real.

Server command and clients

Both arms ran the same command, on main's wheel with either main's or this PR's source:

VLLM_USE_V2_MODEL_RUNNER=1 VLLM_HTTP_TIMEOUT_KEEP_ALIVE=900 vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
  --tensor-parallel-size 4 --kv-cache-dtype fp8 \
  --attention-config '{"backend":"FLASHINFER_MLA_SPARSE_DSV41","indexer_kv_dtype":"mxfp4","indexer_sparse_logits":true}' \
  --speculative-config '{"method":"dspark","num_speculative_tokens":5,"draft_sample_method":"probabilistic","enable_adaptive_verification":false}' \
  --engram-config '{"cpu_offload":true,"use_thp":true}' \
  --kernel-config '{"enable_flashinfer_autotune":true}' \
  --additional-config '{"decode_backend":"latency_specialized"}' \
  --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 --max-num-seqs 64 \
  --max-model-len 1048576 --language-model-only \
  --tokenizer-mode deepseek_v41 --tool-call-parser deepseek_v41 --reasoning-parser deepseek_v41 \
  --enable-auto-tool-choice

• Decode step and AgentX: the speculative config adds "rejection_sample_method":"synthetic","synthetic_acceptance_length":3.51.
• GSM8K: InferenceX's lm-eval GSM8K (5-shot, temperature 0, all 1,319 questions), 8 requests in flight, one run per cell.
• GPQA-Diamond and AIME25: zero-shot through lm-eval, reasoning effort high, temperature 0.6, 16 samples of every question, 32 requests in flight, as for [Perf][DSv4.1] DeepSeek-V4.1-Flash specialized decode on GB300 #61868. GPQA-Diamond is InferenceX's task (the last standalone A–D letter, 16,384 output tokens); AIME25 is lm-eval's with DeepSeek-R1's math prompt (65,536 output tokens).
• RULER NIAH: lm-eval's RULER generator and scorer for niah_single_2, niah_multikey_2 and niah_multikey_3 at 32K, 128K and 256K tokens, 40 seeded samples per task and length (upstream lm-eval draws 500), 128 output tokens, thinking off, one run per cell.
• AgentX: InferenceX's AgentX scenario (AIPerf, the 393-trace corpus), closed loop at a fixed concurrency for 3,600 s.
• Decode step: one server per block. A block runs batch 1 (a 524K-token context), batch 8 and batch 32 (84K–648K tokens a request) twice each. The requests are prefilled one at a time, then sent together for 3,000 tokens each, so they decode from the prefix cache. The step is the client-side window in which every request decodes, divided by its chunk count.

AgentX. The InferenceX AgentX trace corpus (393 traces), 3,600 s per concurrency, closed loop, DSpark at the synthetic acceptance length 3.51 as InferenceX runs it, at 0cdfc566. Both arms of a concurrency ran on the same tray, one after the other (the first arm alternates with the concurrency). Each cell reads main → this PR.

Concurrency	P90 interactivity (tok/s/user)	Change	Total tok/s/GPU	Change	TPOT p50 (ms)	TTFT p50 (s)	E2E p50 (s)	Requests
1	817.7 → 852.0	+4.2%	7,856 → 8,016	+2.0%	1.110 → 1.050	0.984 → 0.783	2.10 → 1.52	341 → 395
4	666.3 → 808.7	+21.4%	10,336 → 10,477	+1.4%	1.110 → 1.040	0.434 → 0.431	1.13 → 1.09	764 → 769
8	539.6 → 678.7	+25.8%	16,813 → 17,198	+2.3%	1.280 → 1.110	0.414 → 0.434	1.26 → 1.16	1,443 → 1,456
16	431.1 → 536.2	+24.4%	29,316 → 29,922	+2.1%	1.520 → 1.230	0.496 → 0.448	1.30 → 1.11	2,715 → 2,748
32	225.0 → 258.2	+14.8%	61,858 → 63,312	+2.3%	2.210 → 1.640	0.829 → 0.671	2.30 → 1.80	4,869 → 4,966

Some trace turns ask for one output token. When that token ends the reasoning, the response has no content, and the client records the request as an error: 0–1 such requests a run for main and 0–1 for this PR.

TTFT: API server and engine

AgentX TTFT is mostly the API server's single-threaded work before the engine sees the request: chat-template rendering and tokenization of the whole conversation. aiperf's per-request request_ack_ns separates that work from the engine's (acknowledgement to first token). This PR changes only worker-side code, and the engine part stays within 18 ms of main at every concurrency; the TTFT differences in the table are in the API server:

Concurrency	API server p50 (s)	Engine p50 (s)
1	0.693 → 0.479	0.267 → 0.249
4	0.215 → 0.219	0.177 → 0.167
8	0.237 → 0.244	0.150 → 0.153
16	0.284 → 0.239	0.158 → 0.162
32	0.503 → 0.354	0.214 → 0.213

Decode top-k. decode_topk against vLLM's opt-in DeepSelect and the stock path (FP32 upcast + cooperative_topk), k = 512, µs per call:

Random BF16 logits; "valid" is the share of each row's columns before its end. One GB300, a CUDA graph of 20 back-to-back calls, median of 20 replays. Each cell: decode_topk / DeepSelect / stock.

With row sharding, a rank selects 12 rows at batch 8 and 48 rows at batch 32 (6 rows a request, over 4 ranks).

• Up to 12 rows, decode_topk is faster than DeepSelect at every shape but two: 65,536 columns with 1/8 valid (7.3–7.7 against 4.2–4.3 µs), and 1,048,576 columns with 1/8 valid at 12 rows (17.1 against 16.5).
• From 24 rows, DeepSelect is faster when 1/8 of the columns are valid (65,536 columns at 48 rows: 17.1 against 4.3 µs) and at 65,536 columns from 48 rows and 524,288 all-valid columns from 48 rows (44.2–44.1 against 46.1–44.7 µs). decode_topk stays faster with all 262,144 columns valid (11.8–22.5 against 26.1–26.2 µs).
• Above 12 rows, decode_topk does not take 1,048,576-column rows, so they use the stock path (81.7–289.3 µs against DeepSelect's 16.6–76.6).
• In block mode decode_topk is faster at every size (5.9–12.7 against 9.3–80.0 µs).
• Against the stock path decode_topk is faster at every shape it takes, except 65,536 columns with 1/8 valid (7.3–17.4 against 6.7–12.9 µs; equal at 24 rows).

All shapes

Columns, valid	1 row	6 rows	12 rows	24 rows	48 rows	64 rows
65,536, 1/8	7.3 / 4.2 / 6.7	7.4 / 4.2 / 7.0	7.7 / 4.3 / 7.4	8.9 / 4.3 / 8.9	17.1 / 4.3 / 11.3	17.4 / 4.3 / 12.9
65,536, all	9.0 / 10.8 / 11.5	9.3 / 10.7 / 12.2	9.4 / 11.0 / 14.1	10.3 / 11.0 / 15.7	17.9 / 11.0 / 21.4	20.2 / 11.0 / 23.3
262,144, 1/8	7.4 / 8.2 / 11.1	7.9 / 8.2 / 13.2	8.0 / 8.2 / 17.0	9.9 / 8.3 / 22.3	19.3 / 8.3 / 35.9	19.6 / 8.4 / 54.6
262,144, all	9.5 / 25.8 / 14.6	11.3 / 25.9 / 18.6	10.1 / 26.0 / 29.7	11.8 / 26.1 / 34.9	20.8 / 26.2 / 62.0	22.5 / 26.2 / 80.7
524,288, 1/8	8.7 / 11.0 / 11.8	9.0 / 11.1 / 17.2	9.3 / 10.9 / 23.8	15.5 / 11.0 / 36.3	28.5 / 11.0 / 81.4	35.0 / 11.1 / 102.6
524,288, all	12.4 / 16.8 / 18.0	11.5 / 17.0 / 29.7	28.9 / 43.9 / 49.3	37.0 / 44.0 / 63.2	46.1 / 44.2 / 130.0	44.7 / 44.1 / 151.9
1,048,576, 1/8	9.4 / 11.8 / 14.3	10.1 / 12.0 / 23.9	17.1 / 16.5 / 39.4	– / 16.6 / 81.7	– / 16.7 / 149.9	– / 16.8 / 190.7
1,048,576, all	12.1 / 20.9 / 27.0	12.4 / 21.1 / 49.1	27.3 / 76.1 / 90.1	– / 76.2 / 129.9	– / 76.4 / 246.8	– / 76.6 / 289.3

Block mode (the candidate source: top 2,048 blocks of 8 columns; DeepSelect over the block maxima), 1/8 valid, decode_topk / DeepSelect:

Columns	1 row	6 rows	12 rows
524,288	5.9 / 9.3	6.2 / 24.5	6.3 / 42.7
1,048,576	6.3 / 13.3	6.5 / 43.6	12.7 / 80.0

Memory and start-up. The replicated embedding table takes 1.32 GB per GPU and the all-gather's mailbox 48 MiB, which the KV cache gives up (106,281,697 → 105,286,588 KV-cache tokens on the tray, −0.9%). Engine start: the engine's init (profile, KV cache, warm-up) took 98.8–103.0 s (main) and 98.8–108.9 s (this PR) with warm compile caches, three starts each; a tray's first start, with cold caches, took 680–748 s in either arm.

Details

Duplicate-work check. No open PR extends the DeepSeek-V4.1 specialized decode. Related open PRs:

• [Perf][Kernel][DSv4.1] Swap-AB CuTe DSL decode attention for small TP2/TP4 steps #62633 adds a swap-AB CuTe DSL decode attention to the FlashMLA mega-attention layer. This PR's sparse-MLA kernel replaces the TRTLLM-gen decode of the FLASHINFER_MLA_SPARSE_DSV41 backend, which the specialized decode requires.
• [Perf][Kernel] Let cooperative_topk read the BF16 indexer logits #62172 lets cooperative_topk read BF16 logits, and [Perf][DSA] Add FlashInfer gvr_2 as a sparse-indexer top-k backend #57109 adds FlashInfer's sampled gvr_2 as an opt-in indexer top-k. decode_topk serves only this model's decode indexers. It also selects the candidate source's blocks, and it covers the 1M-wide rows on 16 CTAs. The comparison with DeepSelect is above.
• [Perf][Kernel][DSv4.1] Fuse MoE top-k finalize with shared-expert add #62635 fuses the finalize with the shared-expert add where the fused all-reduce + mHC does not take it. This PR's MoE finalize stays deferred into that fused kernel.

What changed, and what each part saved when it was added

Each row was measured against the tree before it, mostly as decode-step ABBA blocks on one tray (batch B = B requests × 6 rows). They were measured one at a time, on a moving base that also held DSpark-draft changes not in this PR, so they do not add up exactly.

Change	Why	Measured when added
Pinned CuTe DSL MXFP8 tactics for the attention projections above 8 rows	The autotuner sometimes picked an unclustered wo_b tactic (15.3 against 5.9 µs at 48 rows), about 310 µs a step at batch 8, depending on the server start	Batch 8 −7.3%
The next layer's mHC coefficients on the side stream at every executor step (was up to 48 rows)	They ran inline above 48 rows	Batch 32 −3.9%
Low-batch routed experts up to 64, then 192 rows: a fused front, a routing-record prefix, the router's padding mask	Above 8 rows the stock FlashInfer chain ran (routing, sort, GEMMs, finalize)	Batch 2 −8.6%, batch 4 −7.0%; 12 cold layers at 80–128 rows 10–13% under FlashInfer
Low-latency MXFP8 projections up to 48 rows	FlashInfer's block-scaled GEMMs, 13.5 µs for the three projections at 12 rows against 4.3 µs of HBM time	Attention sublayer at batch 2 50.4 → 48.5 µs
Dense indexers row-sharded from two requests; Engram WKV run ahead; candidate expansion overlapped	Every rank scored every row of the indexers	Batch 2–21 −0.6% to −4.7% (all separated)
Two-shot all-reduce + mHC from 64 rows, adapted from FlashInfer's MNNVL two-shot all-reduce	Each rank read four contributions per value; the owner now sums them once	Batch 21 −1.0%, batch 28 −0.9%; 192 rows 15.7 → 12.1 µs per call
The dense indexer planned before the index-key wait	6.6–7.2 µs of reset, schedule and cast waited for the keys without reading them	About −25 µs a step at batch 2, −50 at batch 8 (profile)
Embedding from a replicated table into the hyper-connection streams	Sharded lookup + all-reduce + copy	Not measured on its own: it was added in a set of four changes, which together gave batch 4 −1.6% and batch 21 −1.4%. From the work it removes, about −13 µs a step is expected at batch 2 and −23 at batch 32
Indexer metadata builds on branches of their own	13 nodes ran in series on one block-table branch	Draft end → next target −7 to −10 µs
Decode top-k from BF16 logits (CuTe DSL)	cooperative_topk plus the FP32 upcast took 182 µs a step on the slowest rank at batch 8	Batch 2 −1.5%, batch 8 −1.5% (two trays)
Sparse MLA up to 8 requests (CuTe DSL): a cluster per token, one CTA per 128 columns, merged over distributed shared memory	TRTLLM-gen's kernel at these shapes	Batch 2 −5.3%, batch 4 −2.7%, batch 8 −4.9%
GEMM2 starts per token tile	GEMM2 waited for GEMM1's whole grid	MoE pair at batch 2 39.2 → 35.4 µs a layer; batch 2 −1.6%, batch 8 −1.3%
The front's top-k over eight CTAs; its stores after the cluster release	The top-k was issue-bound on one SM (0.18 µs a row)	Batch 2 −1.1%, batch 4 −1.6% (five sets, three trays)
Layer 20's top-k on 16 CTAs a row	Its 1M-wide rows fell back to the upcast and the stock top-k	Layer 20 at batch 2 38.5 → 24.7 µs a step; batch 2 −1.1%
The executor's own all-gather for the indexer rows and the logits (CuTe DSL)	main gathers them with the custom all-gather, whose Lamport kernel never waits for the rank's own multicast copy, so a late copy can corrupt a later gather of another size (test_all_gather.py --custom-all-gather reproduces it)	Not measured on its own
The decode Q/KV norm and the fused wo_a trigger their dependents early (≤ 48 rows)	wq_b and wo_b load their weights before their own wait, but were launched late	Attention per layer at batch 2 34.9 → 33.4 µs, about −60 µs a step

---

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.

BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing (anything written below this line will be removed by GitHub Actions)

[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.

[mergify]

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

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

[yaoyaoding]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99406 for commit 6330fc83b0a6.

ℹ️ This PR is 21 commits behind upstream main. Merged upstream main into this branch as 6330fc83b0a6 before starting CI. Pull the branch before pushing more changes.

[mergify]

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

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

[yaoyaoding]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99611 for commit 3bd96601ad21.

ℹ️ This PR is 5 commits behind upstream main. Merged upstream main into this branch as 3bd96601ad21 before starting CI. Pull the branch before pushing more changes.

[mergify]

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

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

[khluu]

stamped for CI

[yaoyaoding]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99839 for commit fe843bccff3f.

ℹ️ This PR is 5 commits behind upstream main. Merged upstream main into this branch as fe843bccff3f before starting CI. Pull the branch before pushing more changes.