[iseeyuan]

Overview

GDN layers with the flat qkvz layout (Qwen3.5 family) fall back to the unfused speculative-decode path on ROCm because fused_gdn_decode_post_conv_mtp (#51674) only builds for CUDA. This adds a Triton port with the same signature and uses it on ROCm: a state-update kernel tiled over the value dimension like fused_sigmoid_gating_delta_rule_update (tile size chosen from the grid size) plus a small gated-RMSNorm epilogue. It covers every model built on QwenGatedDeltaNetAttention with the flat layout: the Qwen3.5 family, Qwen3.8-27B and Qwen3.8-Flash-Next (table below). Qwen3-Next (interleaved layout) stays ineligible, as on CUDA.

Claims

• Qwen3.5-35B-A3B, MTP k=3, 1x MI350X, one request at a time: decode 2.39 -> 2.08 ms/token (-13%), E2E 544 -> 488 ms (-10%), on vLLM main (nightly 77154ca) with aiter main. The first version of this PR (one program per value head, no V tiling) gave 2.19 ms/token in the same side-by-side run.
• Per decode step on MI350X: copy/fill/cat kernels 1.56 -> 0.31 ms, kernel launches 1008 -> 768, GPU kernel time 7.74 -> 6.71 ms.
• Faster at higher concurrency too: per-step time (median ITL) -12% at 1, -6% at 8 and -6% at 32 concurrent requests, output throughput +14% / +16% / +10% (vllm bench serve, 2K/256).
• Numerics match the unfused path: the existing MTP kernel test, parametrized over CUDA op / Triton port, passes for the Triton port on MI350X and MI300X (22/22 each); states are bit-identical to the first version at every tile size, and outputs are too or differ only in the rounding of the norm's sum; MTP acceptance unchanged (2.85-2.87 accepted tokens per step).
• Accuracy unchanged: GPQA Diamond 0.821 (main) vs 0.826 (this PR), needle-in-a-haystack 100% for both at 4K-64K tokens.
• Long context (64K in / 1K out): decode step (median ITL) -12% / -6% / -3% at 1 / 8 / 32 concurrent requests; prefill unchanged.
• MI300X (gfx942): per call, 49-60% faster than the unfused path at 1 request and 25-44% at 32 (table below).
• GSM8K unchanged within run-to-run variation (base 0.854-0.861 over 4 runs, this PR 0.848-0.862 over 3 runs), above the 0.84 threshold in tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-DEP2.yaml.
• CUDA unchanged: the CUDA op is still used wherever it is built.

Validation

Kernel test (MI350X and MI300X, ROCm 7.2.3; the CUDA-op cases skip because the op is not built on ROCm):

pytest tests/kernels/fla/test_fused_gdn_post_conv.py -k mtp
22 passed, 22 skipped

End to end: Qwen/Qwen3.5-35B-A3B (BF16), 1x MI350X, --speculative-config '{"method":"mtp","num_speculative_tokens":3}' --attention-backend ROCM_AITER_FA --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE"}' --async-scheduling, VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1. Workload: 104 image-to-video prompt-rewriting requests (one image + 9.5K-token system prompt, ~180 output tokens, temperature 0.6 / top_p 0.95 / top_k 3), one at a time, first pass on a fresh server.

vLLM	GDN MTP decode	Decode ms/token	Accepted tokens/step	TTFT median ms	E2E mean ms
main (nightly 77154ca) + aiter main	unfused fallback	2.39	2.86	107	544
main (nightly 77154ca) + aiter main	first version of this PR (no V tiling)	2.19	2.85	112	510
main (nightly 77154ca) + aiter main	this PR	2.08	2.87	111	488

The three servers ran side by side on three GPUs of the same host. The I2V, profile, GSM8K, GPQA and NIAH runs predate the tile-size rule; at one request it picks the same 32-row tiles, and it leaves states and outputs unchanged in the checks above.

aiter main is used because aiter <= v0.1.25 sends the multi-token verify batch to an unsplit gfx950 attention kernel (fixed by ROCm/aiter#6336, not yet released); without it attention dominates the step and hides this change.

Per decode step GPU time (torch profiler, 39 steps, same servers):

	unfused fallback	first version	this PR
GDN kernels	0.48 ms	0.78 ms	0.62 ms (includes the gated RMSNorm)
copies / fills / cat	1.56 ms	0.31 ms	0.31 ms
other elementwise	0.48 ms	0.53 ms	0.53 ms
kernel launches per step	1008	738	768
total GPU kernel time	7.74 ms	6.86 ms	6.71 ms

vllm bench serve (base and PR servers side by side; random dataset, 2048 input / 256 output tokens, --ignore-eos, 4x concurrency prompts, min 32). Acceptance differs between runs because the outputs differ. With MTP each streamed chunk is one decode step, so median ITL is the per-step time independent of acceptance; TPOT x acceptance is shown too:

Concurrency	Base TPOT ms / acceptance	PR TPOT ms / acceptance	Median ITL (step)	TPOT x acceptance	Output tok/s
1	2.90 / 2.32	2.49 / 2.38	6.75 -> 5.93 ms (-12%)	-12%	314.0 -> 359.2
8	6.75 / 1.89	5.66 / 2.25	10.92 -> 10.31 ms (-6%)	-0.2%	971.7 -> 1126.8
32	12.46 / 2.24	11.48 / 2.26	18.74 -> 17.55 ms (-6%)	-7%	2150.1 -> 2354.9

Long context, vllm bench serve with 64K input / 1K output tokens (random dataset, --ignore-eos, 8 / 32 / 64 prompts at concurrency 1 / 8 / 32, both servers with --max-model-len 69632 and prompts neither had cached):

Concurrency	Median TTFT	Median ITL (step)	Base TPOT ms / acceptance	PR TPOT ms / acceptance	Output tok/s
1	1834 -> 1833 ms	7.26 -> 6.42 ms (-12%)	2.49 / 2.89	2.33 / 2.71	233.9 -> 242.3
8	1917 -> 1886 ms	14.65 -> 13.79 ms (-6%)	12.37 / 2.89	13.38 / 2.49	524.2 -> 512.5
32	2746 -> 3073 ms	31.43 -> 30.64 ms (-3%)	34.51 / 2.78	34.18 / 2.75	782.1 -> 801.9

The 64K prefill dominates these runs, and acceptance on random prompts varied from 2.49 to 2.89, which moves throughput more than this change does. The decode step itself is 3-12% faster and prefill is unchanged; at 32 concurrent requests, time to first token mostly reflects queueing behind other prefills and varied by about 30% between passes. A pass with another prompt seed gave -12% / -7% / -6% median ITL.

GSM8K (tests/evals/gsm8k/gsm8k_eval.py, 1319 questions, 5-shot, greedy, --max-concurrency 16, so MTP verify batches of up to 16 requests go through the new kernel):

	Run 1	Run 2	Run 3	Run 4
base	0.861	0.861	0.854	0.858
this PR	0.855	0.862	0.848

GPQA Diamond and needle-in-a-haystack, main vs this PR side by side (both servers with --max-model-len 69632):

	GPQA Diamond accuracy	per repeat	mean output tokens	NIAH 4K / 8K / 16K / 32K / 64K
main	0.821	0.808 / 0.833	8105	33/33 each
this PR	0.826	0.828 / 0.823	8087	33/33 each

GPQA Diamond uses the OpenAI simple-evals prompt and answer pattern: 198 questions x 2 repeats with shuffled choices, thinking on, temperature 0.6 / top_p 0.95 / top_k 20, up to 24K output tokens, 64 concurrent requests. NIAH follows RULER's niah_single_1: a 7-digit number hidden at 11 depths (0-100%) in the noise haystack, 3 samples per depth, thinking off, greedy.

Reproduce (on a ROCm image, from the vLLM repo root):

pytest tests/kernels/fla/test_fused_gdn_post_conv.py -k mtp
VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1 vllm serve Qwen/Qwen3.5-35B-A3B --port 8100 \
  --max-model-len 32768 --attention-backend ROCM_AITER_FA --async-scheduling \
  --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE"}' \
  --speculative-config '{"method":"mtp","num_speculative_tokens":3}'
python tests/evals/gsm8k/gsm8k_eval.py --port 8100 --max-concurrency 16
vllm bench serve --model Qwen/Qwen3.5-35B-A3B --port 8100 --dataset-name random --random-input-len 2048 \
  --random-output-len 256 --ignore-eos --num-prompts 32 --max-concurrency 1

Details

Models that take the fused path on ROCm with this change (eligibility is unchanged from the CUDA op: flat qkvz layout, K = V = 128, value/key head ratio 1/2/3/4/8, BF16 activations, at most 8 tokens per verify):

Model	vLLM model code	GDN K/V heads (ratio)	GDN layers	Validated
Qwen3.5-397B-A17B	qwen3_5.py	16/64 (4)	45	kernel test
Qwen3.5-122B-A10B	qwen3_5.py	16/64 (4)	36	kernel test
Qwen3.5-35B-A3B	qwen3_5.py	16/32 (2)	30	end to end (this PR)
Qwen3.5-27B / 9B / 2B	qwen3_5.py	16/48 (3) / 16/32 (2) / 16/16 (1)	48 / 24 / 18	kernel test
Qwen3.8-27B	qwen3_5.py (Qwen3_5ForConditionalGeneration)	16/48 (3)	48	kernel test
Qwen3.8-Flash-Next (incl. FP8)	vllm/models/qwen4_exp (same GDN layer, flat layout)	16/48 (3)	36	kernel test
Qwen3-Next-80B-A3B	qwen3_next.py (interleaved layout)	16/32 (2)	-	not eligible (unchanged)

• Root cause. On ROCm the layer logs Falling back to the Triton GDN spec decode path: torch.ops._C.fused_gdn_decode_post_conv_mtp is not built. The AITER fused GDN path only covers non-speculative decode, so every GDN layer's verify step runs ~8 small kernels around the recurrent kernel: contiguous copies of a/b, three copies plus a cat to unpack q/k/v, a zero-fill, an output copy, and the gated RMSNorm.

• Kernels. Two launches. The update kernel runs one program per (request, value head, tile of the value dimension). It reads q/k/v and a/b in place from the projections and runs the gated delta-rule update over the request's tokens with the same state reads and writes as fused_sigmoid_gating_delta_rule_update: start from the state at the last accepted token, store the state after every token, slot index <= 0 means no state. Each tile writes its rows of the output and their sum of squares per token. The epilogue, one program per (request, value head, token), adds the partial sums and applies the gated RMSNorm, the only step that needs the whole value dimension, then writes the out_proj input. Recurrent state layout is [slot, HV, V, K], BF16 or FP32.

• Tile size. 32 rows per program while there are fewer (request, value head) pairs than CUs, 64 up to 4x the CUs, else one tile per value head. A sweep of 32/64/128 rows on MI300X and MI350X (both head shapes, both state dtypes, 1-32 requests) put this rule within 8% of the best tile size everywhere; a fixed 32 was up to 36% slower at 16-32 requests on MI300X.

• Wiring. Same keyword signature as the CUDA op; the layer picks the Triton port on ROCm. forward_hip routes speculative-decode configs to the packed fused-norm path (forward_cuda); non-speculative ROCm decode keeps the AITER path. Opt out with the existing VLLM_GDN_DECODE_KERNEL=triton.

• Per call. Qwen3.5-35B-A3B shape (16/32 K/V heads), 4 tokens per request, CUDA-graph replay, current ROCm nightly image; "first version" is this PR before the review (one program per value head):

  GPU	State	Requests	unfused	first version	this PR
  MI350X	FP32	1	30.0 us	18.7 us	11.5 us
  MI350X	FP32	8	46.7 us	26.8 us	21.8 us
  MI350X	FP32	32	91.4 us	71.9 us	64.3 us
  MI350X	BF16	1	28.3 us	12.5 us	10.2 us
  MI350X	BF16	32	68.1 us	36.5 us	39.6 us
  MI300X	FP32	1	30.2 us	21.7 us	13.3 us
  MI300X	FP32	8	56.7 us	38.5 us	33.3 us
  MI300X	FP32	32	145.7 us	96.8 us	88.8 us
  MI300X	BF16	1	28.5 us	16.3 us	11.3 us
  MI300X	BF16	32	117.0 us	66.5 us	65.2 us

  With FP32 state, the Qwen3.5 and Qwen3.8 default (mamba_ssm_dtype: float32), this PR is faster than the first version at every batch size on both GPUs. With BF16 state at 32 requests on MI350X it is 8-9% slower than the first version (one tile per head plus the epilogue launch), still well ahead of the unfused path. The 16/48 head shape behaves the same way.

• Coverage. End to end this was run on Qwen3.5-35B-A3B only. The other models in the table are covered by the kernel test, which runs every eligible head ratio (1/2/3/4/8) at TP 1-16, but were not run end to end.

• CI. The parametrized test runs the Triton port on any GPU (NVIDIA too), so existing CUDA CI covers its numerics; AMD CI covers the ROCm path where it runs tests/kernels/fla.

• Related open PRs (not duplicates).

  • [Perf][ROCm][Kernel] Fuse Qwen3.8 MTP verify conv, GDN, and RMSNorm #61675 adds an opt-in (VLLM_ROCM_QWEN_MTP_GDN), shape-specialized ROCm kernel that also fuses the conv. It accepts only 16/48 K/V heads (Qwen3.8-27B), TP=1 and verify widths up to 4, and raises for any other shape, including Qwen3.5-35B-A3B (16/32). This PR makes the default ROCm path fused under the same contract and eligibility as the upstream CUDA op (head ratios 1/2/3/4/8, any TP, up to 8 tokens). The two can coexist, with [Perf][ROCm][Kernel] Fuse Qwen3.8 MTP verify conv, GDN, and RMSNorm #61675 taking over when it is enabled.
  • [Perf][Spec Decode] Defer GDN spec decode state stores to the next step #61434 defers state stores inside the unfused Triton spec-decode kernel (FP32 state, mamba_cache_mode="none"). On ROCm this PR replaces that path whenever the fused path is eligible; the deferred-store idea could be applied to this kernel in a follow-up.
  • [Spec Decode][Qwen4Exp] Run GDN RecoverSSM verify in the fused CUDA MTP kernel #59366 adds a separate CUDA replay op for RecoverSSM and leaves the op ported here unchanged.
  • [ROCm][Perf] Qwen3.8-2.4T specialized decode on MI355X (--monokernel) #60833 / [ROCm][Perf][Qwen3.8-Flash-Next] Mono decode layer for Qwen3.8-Flash-Next on MI355X (TP2) #61325 (Qwen3.8 mono decode on MI355X) run vLLM's GDN core for verify steps, so this change applies there too.

• Provenance. The first version of the kernel was written and validated in Meta's internal vLLM fork; the V tiling was added in this PR after review.

• AI assistance. Drafted with Claude Code (port, V tiling and tile-size rule after review, test parametrization, benchmarks); reviewed and run by the submitter.

---

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.

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

[tpopp]

This eliminates the copy/cat/fill overhead, but drops the BV=32 V-tiling that fused_sigmoid_gating_delta_rule_update_kernel uses.

Since only the final RMSNorm needs the full V reduced, would tiling the state update over V like the existing kernel, with the norm done as a small epilogue over the accumulated sum-of-squares, be faster?

[iseeyuan]

Good call, measured it: on MI350X at the Qwen3.5-35B-A3B shape (fp32 state), V-tiling with BV=32 plus a small norm epilogue is 33% faster per call at 1 request and 15% at 32, with bit-identical outputs and states. I'll switch the PR to that layout.

[simondanielsson]

Thanks for the fix, this looks good

Suggestion: Can we please attach

1. gpqa and niah accuracy results vs main
2. vllm bench serve with 64k/1k ISL/OSL

Suggestion: Can we also check the perf of this on gfx942, or add a 950 gate?

[iseeyuan]

Thanks, good call. Switched to that layout in 66eb7ec: the state update is tiled over V like fused_sigmoid_gating_delta_rule_update, and the gated RMSNorm runs as a small epilogue over the per-tile sums of squares. One addition: a fixed 32-row tile was up to 36% slower at 16-32 requests on MI300X, so the tile size now follows the grid (32 rows while there are fewer (request, value head) pairs than CUs, 64 up to 4x the CUs, else one tile per head); a 32/64/128 sweep on MI300X and MI350X puts that within 8% of the best tile everywhere. States are bit-identical to the previous version.

Per call (Qwen3.5-35B-A3B shape, FP32 state, 4 tokens per request), previous -> now: MI350X 18.7 -> 11.5 us at 1 request and 71.9 -> 64.3 us at 32; MI300X 21.7 -> 13.3 us and 96.8 -> 88.8 us. End to end on MI350X at batch 1, decode went from 2.19 to 2.08 ms/token (unfused fallback: 2.39). The description has the refreshed numbers.

[iseeyuan]

@simondanielsson Thanks for the suggestions; here are the results (methodology and full tables in the updated description).

• gfx942 (MI300X): the kernel test passes (22/22). Per call (Qwen3.5-35B-A3B shape, 4 tokens per request), this PR is 49-60% faster than the unfused path at 1 request and 25-44% at 32. The tile-size sweep there showed a fixed 32-row V tile was up to 36% slower at 16-32 requests, so the tile size now follows the grid size; with that, FP32 state (the Qwen3.5 default) is faster than the PR's first version at every batch size on MI300X and MI350X.
• GPQA Diamond (simple-evals prompt, 198 questions x 2, thinking on): main 0.821, this PR 0.826.
• NIAH (RULER niah_single_1, 4K-64K tokens, 11 depths): 100% for both.
• vllm bench serve 64K/1K: the decode step (median ITL) is 12% / 6% / 3% faster at concurrency 1 / 8 / 32, and prefill is unchanged (median TTFT 1834 vs 1833 ms at concurrency 1). Output throughput moved -2% to +4%, mostly with acceptance, which varied from 2.49 to 2.89 on the random prompts.

GPQA, NIAH and 64K/1K ran on MI350X with vLLM nightly 77154ca + aiter main and shared-expert fusion, main vs this PR side by side.

[github-actions]

✅ @iseeyuan, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[simondanielsson]

Thanks for the changes! LGTM

[simondanielsson]

/ci run

[github-actions]

✅ Triggered Buildkite CI #96005 for commit b2ec6890e848.

[simondanielsson]

/ci run

[github-actions]

❌ This PR is 61 commits behind upstream main. The bot could not merge upstream main into this branch (API returned 422: expected head sha didn’t match current head ref.). No new CI build was started. Either merge or rebase onto the latest main yourself and rerun /ci run, or use /ci run --allow-stale to test this branch as is.

[simondanielsson]

/ci run

[github-actions]

✅ Triggered Buildkite CI #96663 for commit 983d611112e7.

[iseeyuan]

@simondanielsson since all tests pass, could we merge the pr to main to have the improvement landed, and to unblock the deployment by default?

[fxmarty-amd]

@iseeyuan you can join vllm slack => ping maintainers in #pr-merge-requests as this is already approved by @simondanielsson (who likely does not have merge rights)