[ROCm][Perf][DSv4.1] Mono decode for the DSpark draft layers by Fangzhou-Ai · Pull Request #62131 · vllm-project/vllm

Fangzhou-Ai · 2026-10-17T03:57:30Z

Overview

DSv4.1-Flash mono decode (MI355X, --monokernel): the DSpark draft's three layers now run on the mono K1 / K2 kernels too, as ATOM's DraftMonoRunner does, instead of vLLM's per-op path. Follow-up to #60397 and #61082; they follow --monokernel like the target's layers, with no new switch. Overlaps files with #60731 (gfx942 port), so whichever merges second needs a rebase.

Claims

The draft's three layers cost about 227 us of kernel time per step on mono, in place of about 640 us on the original path (graph profile, TP4 c1): about 414 us less GPU time a step.
AgentX at c1, per-request ITL p50 is 0.030 ms (1.8%) lower than InferenceX's run with --monokernel, at TP4 (1.646 against 1.678 ms) and at TP2 (1.970 against 2.000 ms).
Draft quality is unchanged: DSpark mean acceptance length 3.889 / 3.935 (TP4) and 3.917 (TP2) against 3.899 / 3.896 without, and GSM8K 0.975 / 0.972.
Target layers' kernels are untouched: their ISA is identical to main's.

Validation

All on MI355X, DeepSeek-V4.1-Flash, the InferenceX #3690 recipe, main at ce49c47cbb plus this PR unless a row says otherwise. "Off" is the same build with the draft layers on vLLM's path (a throwaway env switch for the A/B, not in this PR).

Layer level. The draft layers against the original layers on real eager decode steps (TP4 c1 and TP2 c1, 5- and 10-row steps), same inputs: x output cosine >= 0.99975 (relative L2 up to 2.3%, the MoE's top-3 routing near-ties), residual, post / comb / pre mixes >= 0.9999 cosine.

Acceptance and accuracy (real block rejection, GSM8K as in InferenceX's DSv4.1-Flash eval: custom gsm8k task, chat, 5-shot multiturn, max_tokens 12288, 64 requests in flight, full 1,319):

	draft on mono: GSM8K strict / flexible; mean acceptance length	draft off
TP4	0.9750 / 0.9742; 3.889 (3.935 on the final rebase)	0.9750 / 0.9742; 3.899
TP2	0.9719 / 0.9712; 3.917	0.9727 / 0.9719; 3.896

AgentX, c1, 3,600 s (InferenceX's infx.bench agentic client, the #3690 recipe with --monokernel, DSpark with synthetic acceptance 3.51, the Rust frontend). Reference: InferenceX's run 38377181342 (same recipe, nightly 5555a2d, its own node). TP4 and TP2 ran at the same time on one node. The confidence intervals are bootstrap over the per-request ITL of the profiling phase (about 330 requests an arm).

	TP4 PR	TP4 ref	TP2 PR	TP2 ref
ITL p50 (ms), from the summary	1.650	1.680	1.970	2.000
per-request ITL p50 (ms), diff [95% CI]	1.646	1.678: -0.032 [-0.034, -0.027]	1.970	2.000: -0.030 [-0.034, -0.026]
per-request ITL p90 (ms), diff [95% CI]	1.690	1.757: -0.068 [-0.189, +0.012]	2.053	2.178: -0.125 [-0.292, +0.196]
interactivity p50 / p90 (tok/s/user)	607.6 / 591.8	596.1 / 569.1	507.6 / 487.1	499.9 / 459.1
TTFT p50 / p90 (s)	0.170 / 0.247	0.245 / 0.336	0.183 / 0.293	0.210 / 0.383
E2E p50 (s)	1.85	1.93	2.14	2.18
requests completed	335	331	318	316

The p50 gain is outside the noise at both TP sizes; the p90 interactivity gain (+4.0% and +6.1%) is not (its interval spans zero), and the TTFT difference is the node, not this PR (prefill is untouched). The reference is not a same-node baseline, so the p50 difference also includes any difference of node, image and commit; an earlier run on the same node of main plus only #61934's change (a different kernel change) gave ITL p50 1.670 / 2.000 ms (per-request 1.675 / 1.998).

Kernel time (graph-mode torch profile, TP4 c1, draft on mono): K1 about 20 us and K2 about 56 us a layer for the three draft layers (227 us a step); the non-target kernels of a step total 1.76 ms against 2.17 ms with the draft off (excluding all-reduce spin-waits).

Tests. tests/models/test_dsv41_mono_decode.py: 41 pass, 3 new (the draft layer takes the mono path on its own runner, its first layer takes the broadcast embedding, a layer takes only its own SWA row width). The K1 / K2 ISA of the target builds at TP4, M=6 is identical to main's (FLYDSL_DUMP_IR diff: 0 lines).

Accuracy (RFC #59665 protocol, as in #60397: thinking on, reasoning_effort high, temperature 0.6, top-p 0.95, 8 requests in flight a server, DSpark with block rejection, --monokernel so the draft layers run on mono; pass@1, ± is the standard error across questions). The baselines are #60397's numbers; no reference arm of this build was rerun.

	this PR	#60397
TP2 GPQA-Diamond (8 samples a question)	84.8 ± 2.2	85.3 ± 2.2 reference, 85.1 ± 2.2 mono
TP2 AIME25 (16 samples a problem)	98.3 ± 0.9	99.4 ± 0.3 reference, 98.3 ± 0.8 mono
TP4 GPQA-Diamond (4 samples a question)	84.5 ± 2.3	85.2 ± 2.2 mono
TP4 AIME25 (8 samples a problem)	98.8 ± 0.7	100.0 mono (240/240)

AIME25's wrong samples: 8 of 480 at TP2 (3 empty, out of the 64K budget; 4 on one hard problem) and 3 of 240 at TP4 (1 empty).

RULER NIAH (lm-eval, prefix-free chat copies, thinking on, greedy, 40 samples a cell): niah_single_1, niah_multikey_2 and niah_multikey_3 at 32K, 128K and 256K score 40/40 in every cell at both TP2 and TP4, as in #60397.

Details

The draft layers are DeepseekV4DecoderLayers with window-only attention (compress ratio 0, no indexer), so the kernels need no new stage. What differs is carried by the build and the step's integers:

Routing: 128 routed experts, top-3 (dspark_n_routed_experts / dspark_num_experts_per_tok) and one FP4 group per ug task. MoeBuild already took these; MonoBuild now carries experts, topk, ug_groups and swa_width into the MoE stage and the scratch layout.
SWA rows: the draft's metadata is non-causal, with rows get_dspark_swa_index_width(window, num_speculative_tokens) wide (192 for window 128 and 5 drafts). K1's key-table stage read the row width as the constant 128; it reads it from the build now.
First layer: the draft hands it the embedding already repeated to the hc copies, [rows, hc, dim]; K1 takes it as the settled residual (fold=False), as for the target's layer 0.
Its own runner: the scratch layout depends on the number of experts, and a mailbox word must never be a tag to one launch and data to another, so the draft's layers share a second DSV41MonoLayer (epoch, scratch, peer buffers) as in ATOM. Every rank creates it at the same layer of the same eager step, like the target's.

Target builds are unchanged: every new field defaults to the target's value, the kernel symbols gain no suffix for them, and the ISA diff above is empty.

Limitations: MI355X (CDNA4), TP2 / TP4, steps of at most 48 rows (the draft's are 5 per request, so up to 8 requests); the layer comparison is on eager steps only. The draft's remaining per-step cost is outside its layers: _prepare_dflash_inputs (3 launches of 113 us), the six skinny head GEMMs (279 us), and about 75 small elementwise and copy kernels. I found no open PR doing this. Developed and measured with Claude Code; the submitter's line-by-line review is still pending.

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

The DSpark draft's three layers are DeepseekV4DecoderLayers with window-only attention, so the mono K1 / K2 serve them with four differences, all in the build or the step's integers: - their routing: 128 routed experts, top-3 (the MoE stage already takes ``MoeBuild(experts, topk, ug_groups)``); - their SWA index rows: the block's non-causal rows, 192 wide on V4.1-Flash (the key table reads the row width from the build); - the first layer takes the embedding already broadcast to the hc copies; - a runner of their own (its epoch, scratch and peer buffers), since the scratch layout depends on the experts. They follow ``--monokernel``, as the target's layers do. Target builds are unchanged: their kernels' ISA is identical to main's. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com>

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.

Fangzhou-Ai · 2026-10-17T19:59:55Z

/ci run

github-actions · 2026-10-17T20:00:14Z

✅ Triggered Buildkite CI #95619 for commit 980dc3696b8e.

ChuanLi1101

LGTM. The draft's three layers drop from about 640 us to 227 us a step. On the same build, GSM8K and the acceptance length stay flat.

The AgentX ITL p50 is against a different node and nightly, so I am not counting that 0.030 ms as this PR's gain.

shen-shanshan

Reviewed by @ChuanLi1101

main now has vllm-project#62131, which runs the DSpark draft layers whole on gfx950 on a runner of their own. It conflicted with this branch in mono_decode.py, mono/layer.py, mono/runner.py and the model test. Both sides are kept: - On gfx950 the draft layers run as vllm-project#62131 runs them. - On gfx942 the draft layers keep the FFN launch after vLLM's attention, on the target's runner. The draft routing comes from the config, as in vllm-project#62131, and the test config now carries it. - run_ffn reads the target's scratch layout in the FFN launch, because that launch runs on the target's runner. K2 reads the layout of its own runner. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

…groups The DSpark draft's runner (vllm-project#62131) builds its MoE with one FP4 group an ug task (DRAFT_UG_GROUPS), but its scratch was laid out for the default two. With one group a task the ug flags take twice the room. So the flags of the upper half of the tasks were written past their region, over the start of the normed rows that the MoE reads as its input, at TP 2 and 4 and for every step width. scratch_layout and scratch_bytes now take the runner's ug_groups, and K1, K2 and the runner pass it. The target's layout does not change. run_ffn checks that each MoE region of a build fits the layout, and that check stopped the build of the draft runner's K2. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

…de on the latest ROCm nightly / [AMD][MI355X] DSv4.1 Flash vLLM：在最新 ROCm nightly 上以草稿模型 mono decode 重新测试 (#4319) * dsv41flash mi355x vllm: move to the e68efc68 ROCm nightly with draft mono decode Update dsv41flash-fp4-mi355x-vllm-agentic-dspark to vllm/vllm-openai-rocm:nightly-e68efc685ab1e6f1eda004ced7157e3ca411da2f (pinned by digest), the 2026-10-19 nightly. ROCm 10.0 is now the default nightly build, so the tag drops the rocm100 prefix. The image carries vllm-project/vllm#62131, which runs the DSpark draft layers on the mono decode kernels under --monokernel. Recipe and search space are unchanged. 将 dsv41flash-fp4-mi355x-vllm-agentic-dspark 更新到 vllm/vllm-openai-rocm:nightly-e68efc685ab1e6f1eda004ced7157e3ca411da2f （以 digest 固定），即 2026-10-19 的 nightly。ROCm 10.0 现为默认 nightly 构建，因此标签不再带 rocm100 前缀。该镜像包含 vllm-project/vllm#62131，在 --monokernel 下让 DSpark 草稿层也走 mono decode 内核。配方与搜索空间不变。 Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * perf-changelog: link #4319 perf-changelog：关联 #4319 Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * perf-changelog: restore other entries' placeholder pr-links perf-changelog：恢复其他条目的占位 pr-link Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: Chun Fang <chun.fang@amd.com> Co-authored-by: Wenyao Gao <wenyao.gao28@gmail.com>

…the stock front/head runtime/draft.py: DraftDecode, the MTP drafter's own executor, built and bound the same way the target's is (--monokernel; no switch of its own). This commit's draft step runs the stock embedding branch (plain torch: pre_fc_norm_embedding/hidden, fc_embedding/fc_hidden, the residual add) and the proposer's own greedy-token path; only the decoder layer's hyper- connection fronts and blocks are on the kernels. Reuses Mikko Tukiainen's generic launches for the draft decoder layer -- front.hc_front for both hyper-connections' mixes, blocks.oproj_block, blocks.moe_block and the final mixer's front -- since the draft layer's shapes match the target's exactly; "QSA attention stays stock" (unmodified, its own test suite covers it). Its own epoch, health, scratch and peer buffers (DSv4.1's drafter precedent, vllm-project#62131): a mailbox word must never be a tag to one launch and data to another, so the drafter's mailbox is never the target's. Default-deny config refusals (draft_unsupported) share the target's parallelism / feature / platform checks and add MTP-only and the draft model's own shape; model-stage refusals (model_reasons) check the loaded draft layer's modules and weight layouts, same as the target's. vllm/models/qwen4_exp/amd/mtp.py: wires DraftDecode into Qwen4ExpMTP -- init at __init__ (the config check, before any weight loads), build/refusal-check in process_weights_after_loading (the loader's generic post-load hook, so no core change is needed), run in forward, and the fail-stop watch in compute_logits. get_top_tokens (the fused draft head) is not added yet. Tests: - tests/models/qwen4_exp/amd_specialized_decode/test_runtime.py (CPU): the drafter's config refusals and widths, SimpleNamespace configs, no FlyDSL import (draft.py keeps mono.* / runtime.layers imports local to DraftDecode's own methods for exactly this reason). - tests/.../test_kernels.py (GPU, two ranks, S in 1/3/8): the stock-front draft step's whole launch sequence -- the attention HC front, a fixed random (core, gate) standing in for QSA (out of scope here), the o_proj block, the MLP HC front, the MoE block, the final mixer's front (its down zero-padded from LORA to layout.DOWN_ROWS rows) -- each stage checked against the same stock-ops references the individual block tests use (GatedResidual.mix, the HC combine math, the MoE emulation), a port of md-int's test_amd_mono_draft.py::test_draft_step_matches_vllm_ops onto the split package's materialized-state API (no scratch-buffer peeking needed: oproj_block / moe_block hand back the next state directly). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Tres Popp <tres.popp@amd.com>

mergify Bot added deepseek Related to DeepSeek models DSv4.1 Related to DeepSeek-V4.1 models rocm Related to AMD ROCm dflash labels Oct 17, 2026

github-project-automation Bot added this to Sprint - DFlash and AMD Oct 17, 2026

github-project-automation Bot moved this to Todo in AMD Oct 17, 2026

github-project-automation Bot moved this to Backlog in Sprint - DFlash Oct 17, 2026

Fangzhou-Ai assigned shen-shanshan Oct 17, 2026

Fangzhou-Ai marked this pull request as ready for review October 17, 2026 19:59

Fangzhou-Ai requested review from AndreasKaratzas, DarkLight1337, dllehr-amd, hongxiayang, tjtanaa and ywang96 as code owners October 17, 2026 19:59

claude Bot reviewed Oct 17, 2026

View reviewed changes

Merge branch 'main' into mono-draft-layers

980dc36

Fangzhou-Ai added the ready ONLY add when PR is ready to merge/full CI is needed label Oct 17, 2026

ChuanLi1101 approved these changes Oct 17, 2026

View reviewed changes

shen-shanshan approved these changes Oct 19, 2026

View reviewed changes

shen-shanshan merged commit e719b41 into vllm-project:main Oct 19, 2026
107 checks passed

github-project-automation Bot moved this from Backlog to Done in Sprint - DFlash Oct 19, 2026

github-project-automation Bot moved this from Todo to Done in AMD Oct 19, 2026

maeehart mentioned this pull request Oct 19, 2026

[ROCm][Perf][DSv4.1] Mono decode on MI300X and MI325X (gfx942) #60731

Merged

6 tasks done

Fangzhou-Ai mentioned this pull request Oct 19, 2026

[AMD][MI355X] DSv4.1 Flash vLLM: re-sweep with DSpark draft mono decode on the latest ROCm nightly / [AMD][MI355X] DSv4.1 Flash vLLM：在最新 ROCm nightly 上以草稿模型 mono decode 重新测试 SemiAnalysisAI/InferenceX#4319

Merged

zkawfanx mentioned this pull request Oct 22, 2026

[Start here] Classify the pinned DSv4.1 timeline and prepare chronological backport batches zkawfanx/vllm#18

Open

3 of 5 tasks

AndreasKaratzas mentioned this pull request Oct 24, 2026

[CI][ROCm] Add a monokernels test area with a group per AMD monokernel #63486

Draft

Navigation Menu

Repository navigation

[ROCm][Perf][DSv4.1] Mono decode for the DSpark draft layers - #62131

[ROCm][Perf][DSv4.1] Mono decode for the DSpark draft layers#62131
shen-shanshan merged 2 commits into
vllm-project:mainfrom
Fangzhou-Ai:mono-draft-layers

Fangzhou-Ai commented Oct 17, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

Fangzhou-Ai commented Oct 17, 2026

Uh oh!

github-actions Bot commented Oct 17, 2026

Uh oh!

ChuanLi1101 left a comment

Uh oh!

shen-shanshan left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

Fangzhou-Ai commented Oct 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Claims

Validation

Details

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

Fangzhou-Ai commented Oct 17, 2026

Uh oh!

github-actions Bot commented Oct 17, 2026

Uh oh!

ChuanLi1101 left a comment

Choose a reason for hiding this comment

Uh oh!

shen-shanshan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Fangzhou-Ai commented Oct 17, 2026 •

edited

Loading