[ROCm][Perf][DSv4.1] Mono decode layer for DeepSeek-V4.1-Flash on MI355X

[Fangzhou-Ai]

Purpose

Opt-in (VLLM_ROCM_MONO_DECODE=1, meant as the one switch for every model's mono decode kernels) fast path for DeepSeek-V4.1-Flash decode on MI355X (gfx950). Every backbone layer of a decode step runs on persistent FlyDSL launches:

Launch	Layers	Stages
K1	30 standard layers	the mHC attention seam, then the attention front: wqkv, q/kv norms, the KV insert, wq_b
K2	30 standard layers	the attention back (split attention, combine, wo_a, wo_b), the attention TP all-reduce, the FFN seam, the MoE (router, top-6, shared expert, MXFP4 experts) and its TP all-reduce
FFN	the other 10	after vLLM's attention seam and attention: the attention TP all-reduce (wo_b's unreduced output pushed to every rank), then K2's FFN seam and MoE

All TP all-reduces run in-kernel over symmetric peer memory. K1 + K2 replace the ~40 kernels this path launches per standard layer today; the FFN launch replaces the attention all-reduce, the FFN seam and the MoE's kernels on the others.

Scope

• Layers: K1 + K2 on the 30 backbone layers with the standard mHC seam and no compressor or indexer of their own: 3–7, 9–13, 15–19, 21–23, 25–27, 29–31, 33–35, 37–39. The FFN launch on the other 10: layer 0 (embedding-broadcast seam), the Engram layer 1, and the compressor/indexer layers 2, 8, 14 (also Engram), 20, 24, 28, 32, 36.
• Steps: decode only, ≤ 48 rows (8 requests × 1 + 5 DSpark drafts), TP2 or TP4, eager or a FULL CUDA graph. K1 + K2 also need causal SWA windows and fp8_ds_mla KV records.
• Everything else runs as before: the DSpark draft layers, and prefill, mixed or larger steps. The paths take and return the same tensors at the layer boundary, so they interleave freely.

Layout

vllm/models/deepseek_v41/amd/mono_decode.py   eligibility per layer / per step, vLLM tensors -> kernel inputs
vllm/models/deepseek_v41/amd/mono/
    runner.py     DSV41MonoLayer: scratch, epoch, peer memory, the launches
    layer.py      K1 / K2 and the FFN launch, composed from the stages below
    attention/    attention stages
    stages/       seam.py (mHC seams, attention reduce), moe.py, moe_shape.py, moe_plan.py, router.py, gemv.py, dims.py
    common/       tagged mailbox, device ops, MX helpers, build keys, peer memory

stages/ and common/ are adapted from ROCm/ATOM's DeepSeek-V4.1 mono decode (MIT License, commit 97359d6). Each adapted file names its origin and keeps AMD's copyright notice. The changes from ATOM:

• the seams are specialized to these layers;
• vLLM's numerics: f32 router logits, the shared expert reading vLLM's row-major FP8 weights, vLLM's mxfp8_e4m3_quantize and silu_and_mul_with_clamp;
• hand-offs carry epoch tags instead of per-step mailbox clears, so no host step hook and the launches are graph-safe. ATOM's clears also wiped what another step width had laid out differently; without them each width gets a scratch of its own;
• ATOM's engine-side modules are dropped.

Design

• Weights are read in place, with no second copy. The attention linears stay vLLM's row-major MXFP8 (32×32 E8M0 scales). The routed experts stay in AITER's shuffled MXFP4 layout.
• Path selection. The model is not torch.compiled, so the per-step decision is plain Python. It runs when a FULL graph is captured (uniform decode) and on every eager step. PIECEWISE (breakable) captures are always declined, because they are replayed for mixed batches.
• wo_b's reduce on the FFN-launch layers. Their wo_b is built with reduce_results=False, as the sequence-parallel path already does. A decode step hands the partial to the FFN launch; any other step reduces it in the decoder layer with the same tensor_model_parallel_all_reduce. ROCm's preshuffled-FP8 wo_b path now honors reduce_results as the linear does. This also fixes a bug on main: with sequence parallelism and an FP8-block checkpoint, wo_b was all-reduced and then reduce-scattered.
• TP ranks stay in lockstep. Every rank must take the same path at every layer, since the kernels wait on each other's pushes. The conditions read only metadata that all ranks hold alike.
• Failure modes. An explicit opt-in that cannot run raises: a CDNA version outside the kernels' (get_cdna_version() in CDNA_VERSIONS, CDNA4 / gfx950 today, so a later generation is one entry once ported), TP other than 2 or 4, FlyDSL missing, expert / data parallelism or EPLB, VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4, a routed-expert backend other than AITER_MXFP4_BF16, or fewer than 256 compute units (a partitioned GPU: every CTA of a launch must stay resident). The MoE and shared-expert tensors are checked for the kernels' shapes and dtypes at the first step, as the attention ones are. A layer outside the kernels' scope is logged once.
• All-reduces are exact: fp32 sum in rank order, then bf16. At ≤ 48 rows vLLM's own decode all-reduces are below the INT4 quick-reduce threshold, so the current path runs them exact too.
• A scratch per step width. Hand-offs between CTAs are tagged mailbox pairs in a scratch laid out per step width. Each width has its own scratch, allocated at its first (eager) step: 3–16 MB, and at most 8 widths under DSpark. So a word one width polls as a tag only ever holds tags. The MoE's task counters outlive a width and live in the epoch buffer.
• Dependencies: the flydsl wheel, plus two small AITER helpers that path already depends on: aiter.ops.flydsl.kernels.buffer_ops (FlyDSL's wheel no longer ships it) and UncachedIpcHeap for peer memory.

Duplicate check

• No open PR targets DeepSeek-V4.1 mono decode. Searched "dsv41 mono", "DeepSeek V4.1 mono", "mono decode layer" and "dsv41_mono".
• The sibling mono-decode PRs are for other models: [ROCm][Perf][MiniMax-M3] Fused sparse-layer decode (mono decode kernel) #59705 (MiniMax-M3) and [ROCm][Perf][GLM-5] Fused sparse-layer decode via AITER MonoKernel #60216 (GLM-5).
• This PR hooks per decoder layer rather than replacing the model forward, because only 30 of DSv4.1's 40 layers are eligible.

Test Plan

• pytest tests/models/test_dsv41_mono_decode.py (CPU, wired into the MI355 Basic Models job): a decode step takes the mono path and hands the kernels [blocks, block, 584] cache views. A padded FULL-graph step takes it too (the kernels skip pad rows, whose slot is -1). PIECEWISE steps, prefill or mixed steps, more decode tokens than rows, non-causal windows, missing metadata, more than 48 rows and the first layer all take the original path. Deployments the kernels cannot run raise.
• pytest tests/models/test_dsv41_mono_numerics.py (4 × MI355, its own CI step). One layer with random weights in the checkpoint's formats, at TP2 and TP4: the kernels against the ops vLLM runs for the same decode step (the mHC seams, the ROCm attention chain, the all-reduce, DeepseekV4MoE). Each stage is checked on its own input, by cosine and relative error, for the whole layer and the FFN launch, at 1×6 and 8×6 rows. Then both widths' launches are captured in one graph and replayed 300 times, bit for bit against the first run.
• pytest tests/models/test_dsv41_mono_layers.py (same CI step). A five-layer DeepSeek-V4.1-Flash at TP2 and TP4 holding both layer kinds the mono path takes: the whole layer at a ratio-2 and a ratio-1 consumer, and the FFN launch at the sliding-window and KV / index-source layers. It is served with mono on and off on the same random weights, and every decoder layer's outputs at every decode token must match. It runs DeepseekV4DecoderLayer.forward as the engine does, so a change to the layer that the kernels do not follow fails it.
• Race stress: the whole layer and the FFN launch in one HIP graph, 48-row and 6-row steps alternating, every intermediate compared bit for bit with the first run. 100K replays each on layers 3 and 9 (ratio 2) and 21 (ratio 1).
• Compile-only build (COMPILE_ONLY=1) of the in-tree kernels: K1 and K2 at TP2 and TP4, M = 6 and 48, both layer types.
• No regression of K1 + K2: K2's FFN half is factored out and shared with the FFN launch. K2 disassembles instruction-identical to the first commit's (TP2 M=6 ratio 1, TP4 M=48 ratio 2), and the whole-layer accuracy check reproduces its earlier numbers exactly.
• FFN launch, kernel level: against vLLM's FFN half (the attention all-reduce, MHCPreDelayedOp, DeepseekV4MoE with its all-reduce) on the real weights of the 10 layers, with the attention partial from a real ROCm attention chain. 48 points: TP2/TP4 × concurrency 1/2/4/8 × 1–6 tokens per request; the 10 layers chained in one HIP graph, cold weights.
• Kernel level (MI355X): the layer against vLLM's own modules (DeepseekV4MoE, the mHC seams, the ROCm attention chain, the all-reduce) in a TP process group, with real checkpoint weights. 96 points: TP2/TP4 × concurrency 1–8 × 1–6 tokens per request × the two layer types. Timing: 8 layers chained in one HIP graph, cold weights, vLLM's path captured under graph_capture.
• GSM8K: InferenceX's task, 5-shot, full 1,319, temperature 0, concurrency 8, DSpark block rejection; TP2 for K1 + K2 alone, TP2 and TP4 with the FFN launch.
• AgentX e2e with the InferenceX#3690 recipe (this tree's image, a8w4 MoE, INT4 quick reduce), golden DSpark acceptance 3.51, 3600 s per point, c1/c2/c4/c8 at TP2 and TP4.
• Accuracy, per RFC [RFC]: Fast-Track Merging for Model Optimization PRs #59665: GPQA-Diamond and AIME25 at TP2, thinking on with reasoning_effort high, 8 requests in flight a server (every decode step within 48 rows), against the same build with VLLM_ROCM_MONO_DECODE=0. Long context: RULER NIAH (single, multikey 2 and 3) at 32K, 128K and 256K.

Test Result

Unit: 33 passed. Compile-only: all K1 / K2 / FFN builds compile.

Numerics test (GPU): passed at TP2 and TP4: every token's output cosine ≥ 0.9993 against vLLM's chain, relative errors 0.3–0.5% (bounds 1–1.5%; a routed scale off by 3% fails), 300 replays bit-exact. Race stress: no mismatch in 100K replays on each layer. Before the per-width scratch, layer 3 diverged at replay 13 on every run. Model-level test: passed at TP2 and TP4 (9 of 9 local runs of the final version; 40 of 41 of the TP2-only one before it, the failure not recurring and its log not kept). A routed scale off by 3% fails it at both sizes.

Accuracy, TP2, per RFC #59665 (details):

• pass@1, against this build with VLLM_ROCM_MONO_DECODE=0;
• thinking on, reasoning_effort high, temperature 0.6, 8 requests in flight a server;
• ± is the standard error across questions; Δ is paired by question.

	Reference	Mono
GPQA-Diamond (8 samples a question)	85.3 ± 2.2	85.1 ± 2.2 (Δ −0.2 ± 0.6)
AIME25 (16 samples a problem)	99.4 ± 0.3	98.3 ± 0.8 (Δ −1.0 ± 0.7)

RULER NIAH, long context (details): single, multikey 2 and multikey 3 at 32K, 128K and 256K, 40 samples a cell. Every cell is 40/40 for both mono and the reference.

FFN launch, kernel level, µs per layer (vLLM's FFN half → the FFN launch, at 6 tokens per request). Every point is faster or at parity: TP2 1.01–2.14×, TP4 1.12–2.32×. The residual and the MoE input match vLLM's bit for bit at every point; the layer output's cosine is ≥ 0.99999, and every token's ≥ 0.9999.

TP	c1	c2	c4	c8
2	136→81 (1.69×)	181→128 (1.41×)	241→215 (1.12×)	345→342 (1.01×)
4	111→57 (1.96×)	147→85 (1.74×)	176→140 (1.26×)	234→209 (1.12×)

Kernel level, µs per layer. Every one of 96 points is faster, 1.10–2.18×. Accuracy: residual cosine ≥ 0.99999. MoE output: 1,234 of 1,260 tokens at cosine ≥ 0.999. Each of the other 26 is a top-6 routing near-tie: the reference's own 6th/7th expert score gap is ≤ 3.2e-3, against medians near 1e-2. On the same FFN input the MoE matches vLLM's at cosine 0.99999.

Layers	TP	Speedup range	DSpark decode points c1 · c2 · c4 · c8 (current → mono µs)
ratio 1	2	1.11–2.04×	209→122 · 262→179 · 336→280 · 474→427
ratio 2	2	1.10–2.05×	210→123 · 255→172 · 325→269 · 446→406
ratio 1	4	1.15–2.18×	180→94 · 218→123 · 256→192 · 324→282
ratio 2	4	1.17–2.18×	180→95 · 216→121 · 248→185 · 316→270

GSM8K (strict / flexible; DSpark mean acceptance length):

	TP2	TP4
current path	0.9742 ± 0.0044 / 0.9735; 3.949	—
K1 + K2 (30 layers)	0.9750 ± 0.0043 / 0.9742; 3.911	0.9719 ± 0.0045 / 0.9712; 4.003
K1 + K2 + the FFN launch (40 layers)	0.9735 ± 0.0044 / 0.9727; 3.928	0.9719 ± 0.0045 / 0.9712; 3.989

AgentX with the FFN launch (all 40 layers) vs K1 + K2 alone (30 layers). Same recipe and protocol. AgentX replays a fixed request sequence, so the runs' requests match one to one; the last column is the median ratio of a request's inter-token latency (FFN launch / K1 + K2).

TP	conc	p90 interactivity, tok/s/user	total throughput, tok/s	decode, same requests
4	c1	523.7 → 571.4 (1.09×)	26,796 → 28,225 (1.05×)	0.915
4	c2	471.1 → 510.0 (1.08×)	27,217 → 28,241 (1.04×)	0.916
4	c4	445.1 → 502.5 (1.13×)	36,791 → 39,640 (1.08×)	0.911
4	c8	358.6 → 367.6 (1.02×)	62,621 → 62,877 (1.00×)	0.915
2	c1	429.1 → 471.5 (1.10×)	25,039 → 25,599 (1.02×)	0.928
2	c2	381.7 → 417.2 (1.09×)	25,695 → 25,962 (1.01×)	0.925
2	c4	332.4 → 394.4 (1.19×)	35,885 → 36,761 (1.02×)	0.927
2	c8	245.3 → 265.7 (1.08×)	60,501 → 61,180 (1.01×)	0.932

Decode on the same requests is 7–9% faster at every point. TP2 is a re-run: the first TP2 round hit node-wide engine stalls of 20–40 s, on all four servers at the same moments, that cost its throughput. Against InferenceX's run, the FFN launch reaches 1.56× / 1.46× / 1.47× / 1.30× p90 interactivity and 1.15× / 1.14× / 1.11× / 1.05× total throughput at TP2 c1 / c2 / c4 / c8, and 1.66× / 1.56× / 1.64× / 1.38× and 1.19× / 1.20× / 1.14× / 1.04× at TP4. The tables below are K1 + K2 alone.

AgentX, TP2. Against the same-node mono-off baseline, run the same way (four points at once on one node):

conc	p90 interactivity, tok/s/user	total throughput, tok/s	TTFT p90, s
c1	318.1 → 429.1 (1.35×)	21,573 → 25,039 (1.16×)	1.18 → 1.28
c2	299.3 → 381.7 (1.28×)	22,735 → 25,695 (1.13×)	1.10 → 1.08
c4	288.0 → 332.4 (1.15×)	33,043 → 35,885 (1.09×)	1.27 → 1.26
c8	214.6 → 245.3 (1.14×)	58,294 → 60,501 (1.04×)	1.01 → 1.08

AgentX, TP4. Against InferenceX's own run of the same image and recipe (run 37510490434). The TP2 same-node baseline above matches InferenceX's TP2 run within a few percent, so InferenceX serves as the TP4 reference.

conc	p90 interactivity, tok/s/user	total throughput, tok/s
c1	343.7 → 523.7 (1.52×)	23,691 → 26,796 (1.13×)
c2	327.4 → 471.1 (1.44×)	23,579 → 27,217 (1.15×)
c4	307.2 → 445.1 (1.45×)	34,831 → 36,791 (1.06×)
c8	265.8 → 358.6 (1.35×)	60,266 → 62,621 (1.04×)

Provenance of the numbers. The kernel, GSM8K and AgentX numbers come from these kernels while they were hosted in AITER during development (ROCm/aiter#6196). Moving them here changed only import paths, the merge of a few small helper modules, and formatting. The kernel-level accuracy check re-run against this in-tree copy (TP2 layer 21 M=48, TP2 layer 3 M=6, TP4 layer 21 M=48) reproduces the earlier results exactly, to six decimals of cosine. The TP4 GSM8K and every FFN-launch number come from the in-tree code. The kernel, GSM8K and AgentX numbers predate the per-width scratch. It leaves kernel time unchanged: 24 TP2/TP4 points, mean ratio 1.000, range 0.996–1.009.

AI assistance

This change was written with AI assistance (Claude Code). It is opened as a draft for the submitter to review every changed line before it is marked ready.

🤖 Generated with Claude Code

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

[AndreasKaratzas]

very prelim review

[Fangzhou-Ai]

/ci run

[github-actions]

✅ Triggered Buildkite CI #93417 for commit d18f7697a6bc.

[benchislett]

Given the scope of this change, I don't think GSM8k only is sufficient for correctness testing. Can you please ensure that at least GPQA-D and a hard eval like AIME25 are run to ensure there is no numerical leakage?

I haven't formalized this proposal yet but I outline some eval guidance in this RFC, which may be helpful.

[Fangzhou-Ai]

/ci run

[github-actions]

✅ Triggered Buildkite CI #93467 for commit 524bc76eeb3e.

[dllehr-amd]

one small question, I know we don't want to add more to this PR if we don't have to. But can we test the full DeepseekV4DecoderLayer, for a step? The DSV41MonoLayer helps run what a single gpu will see, but I'm a little nervous we don't run the full layer, we may miss something if a change is made to DeepseekV4DecoderLayer after your change.

[Fangzhou-Ai]

Good call. Added in 7c69067 and c25b7e9: tests/models/test_dsv41_mono_layers.py, in the same MI355 CI step (4 GPUs), at TP2 and TP4.

It serves a five-layer DeepSeek-V4.1-Flash (hf_overrides) holding both layer kinds this PR changes:

• layers 2 and 4, a ratio-2 and a ratio-1 consumer, run whole on the mono kernels;
• layers 0, 1 and 3 (the sliding-window layer and the KV / index sources) keep vLLM's attention and run the FFN launch.

The model runs twice on the same random weights, with VLLM_ROCM_MONO_DECODE on and off. Every DeepseekV4DecoderLayer output, x and the residual stream, must match at every decode token taken in a decode-only step. So it exercises the real forward: the hook, the per-step path choice, weights() off the layer, the metadata, and the top-k from the real indexer layers. It also checks that every layer took its mono path on every step. At TP4 it adds that size's own kernel builds, the padded routed intermediate (576 → 640) and the hand-offs across four ranks.

Details that made it reliable:

• Weights: plain dummy weights leave the E8M0 scales zero, so both paths agree trivially. A test-only loader fills checkpoint-range weights before post-load processing (with the routed experts' padding zero, as the loader keeps it). They're small enough that the random stack isn't chaotic: at larger scales, atomics' order alone moves vLLM's own logits run to run.
• Row matching: a request's row and step differ between engines, so rows are matched by prompt length and position.
• Bounds: layer 0 at a request's first decode token sees identical inputs on both paths and gets the op-level test's 1.5%. Deeper layers inherit small differences, which random weights amplify to 2–4%, so they get 8%.
• Routing near-ties: a token whose top-6 routing is a near-tie on vLLM's path may pick another expert. Once it does, it carries the flip in its residual stream, so its later layers aren't compared. Every other token is checked at every layer.
• Sensitivity: a routed scale off by 3% fails it at layer 0, at both TP sizes. A layer the kernels no longer follow would be off on every token.

It takes about 1.5 minutes at TP2 and 3 at TP4.

[WoosukKwon]

Excited to learn more about the kernel! 🚀

[Fangzhou-Ai]

/ci run

[github-actions]

✅ Triggered Buildkite CI #93923 for commit 03705674ac4c.

[Fangzhou-Ai]

/ci run

[github-actions]

❌ This PR is 4 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[Fangzhou-Ai]

/ci run

[github-actions]

✅ Triggered Buildkite CI #93936 for commit f6b54daade30.

[dllehr-amd]

Approving. We look well guarded and this will be a continual work in progress. Let's get the first piece in so future PRs have something to aim against. #60904 will be updated as we solidify the monokernel approach