[ROCm][Perf][Qwen3.8] Mono decode follow-up: MXFP6 projections, attention front kernel, router through o_proj by maeehart · Pull Request #61271 · vllm-project/vllm

maeehart · 2026-10-12T20:17:57Z

Summary

This draft follows #60833. On amd/Qwen3.8-2.4T-A95B-Quark-MXFP4 with TP8 on MI355X, it moves the GDN and attention projections of the specialized decode from MXFP4 to MXFP6 E2M3, which keeps the gsm8k answers as short as with bf16 projections, and adds an attention front kernel, router logits through o_proj and a better task placement. A first measurement on the older layout put the decode step 0.94 ms faster at 1M context (median ITL 9.39 ms to 8.45 ms, with the caveats under Validation).

The branch is stacked on #60833 at 63036c1291, the --monokernel version in the layout of RFC #60904, so the first commit is #60833 and the last 6 are new. #60833 rewrote the kernels in that version, so these commits were ported by hand, not rebased. The earlier version of this PR on 1301a7cfcf had 6 commits with the same intent.

What this PR does

A slot per pick at 2-3 rows, and load order in the layer block. At 2 and 3 rows, the layer block gives every expert pick its own slot instead of building the union of the picked experts. A router CTA issues its router weights after its o_proj tasks, and when one batch covers the union, the down rounds run without a prefetch.
MXFP6 E2M3 projections. --additional-config '{"monokernel_projections": "mxfp6"}' with --monokernel rounds the bf16 in_proj_qkvz, out_proj, qkv_proj and o_proj to MXFP6 after loading, so prefill multiplies the same values that decode reads, and packs MXFP6 copies of them. The GDN and layer blocks multiply those copies with the bf16 rows (W6A16). These builds serve up to 8 rows, and wider steps run vLLM's decode. The layer block's router also reads the AR sums and sums of squares as tagged pairs, so the norm tasks' drain is not on its path.
Attention front kernel. One launch runs the residual add, input_layernorm, the MXFP6 qkv_proj, the q and k norms, RoPE and the FP8 KV cache write of a full-attention layer. The attention then runs with q only, so vLLM does not write the cache again. Any other backend, cache dtype or layout keeps vLLM's cache write, and the reason is logged once.
Router logits through o_proj in the layer block. Up to 5 rows, a folded bf16 weight, the router rows times the norm's (1 + w) times this rank's o_proj rows, gives the router logits straight from the core rows. So every CTA routes before the all-reduce and the norm have built the x rows.
GDN block task placement for the MXFP6 in_proj. The q, k and v rows, which the conv of every GDN task waits for, run on CTAs with one task, and the z rows on CTAs with two tasks. The outputs stay the same bit for bit.
README and build tests.

A configuration that does not fit is refused at startup with every reason, as in #60833. An unknown monokernel_projections value is refused with the config reasons, and projections that are not bf16 with the loaded-model reasons. No environment variable is added.

Numerics

Unlike #60833, two of these changes do not keep the stock decode numerics bit for bit. MXFP6 is a different weight format, and the folded router weight rounds the logits differently, so a near tie between two experts can go the other way. On the older layout, gsm8k and the needle prompts below stayed within noise.

Validation

The numbers below are from the earlier version of these commits on #60833 at 6a1c13b2fc, which had the same kernels in the older layout. They have not been measured again on this port yet.

amd/Qwen3.8-2.4T-A95B-Quark-MXFP4, TP8 on MI355X, vllm/vllm-openai-rocm:nightly-8cbd5d03006c33185f402249ff2b448efd594986 with the branch on top. The #60833 arm serves the projections in online MXFP4. This PR's arm serves them in bf16 with --quantization-config.targets '{"*shared_expert*": "mxfp4"}' and the MXFP6 projections.

	#60833 at `6a1c13b2fc`	This PR on it
Median ITL, 2 requests of 1M input and 1000 output tokens at concurrency 2	9.39 ms	8.45 ms
gsm8k 5-shot, flexible-extract, `max_gen_toks` 2048	0.9477 ± 0.0061	0.9469 ± 0.0062
Mean gsm8k answer length	560 characters	449 characters
Passphrase found in a needle prompt of 2k, 128k and 1M tokens	yes	yes

Both arms also ran the FlyDSL paged-attention decode of ROCm/aiter#5949 in ROCM_AITER_FA, so the difference between them is the decode kernels.

On this port so far: every commit passes pre-commit (ruff and mypy) on its own, the CPU tests in tests/models/qwen3_5/amd_specialized_decode/ give 29 passed and 30 skipped (the kernel tests need gfx950), and every build that the executor compiles for the MXFP6 projections at 1 to 8 rows compiles for gfx950 without register spills or scratch.

Commands

VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MHA=1 \
VLLM_ROCM_USE_AITER_FP4_ASM_GEMM=1 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1 \
VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4 VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB=512 \
VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 HSA_NO_SCRATCH_RECLAIM=1 \
vllm serve amd/Qwen3.8-2.4T-A95B-Quark-MXFP4 -tp 8 --trust-remote-code \
  --monokernel --additional-config '{"monokernel_projections": "mxfp6"}' \
  --async-scheduling --attention-backend ROCM_AITER_FA --block-size 64 --kv-cache-dtype fp8 \
  --max-model-len 1010000 --hf-overrides '{"max_position_embeddings": 1010000}' \
  --max-num-batched-tokens 32768 --max-num-seqs 256 --no-enable-prefix-caching \
  --gpu-memory-utilization 0.9 --gdn-prefill-backend aiter_flydsl --reasoning-parser qwen3 \
  --generation-config vllm \
  --quantization-config.targets '{"*shared_expert*": "mxfp4"}'

vllm bench serve --model amd/Qwen3.8-2.4T-A95B-Quark-MXFP4 \
  --dataset-name random --random-input-len 1000000 --random-output-len 1000 \
  --num-prompts 2 --max-concurrency 2

The #60833 arm drops --additional-config and uses --quantization-config.targets '{"*shared_expert*": "mxfp4", "*linear_attn.in_proj_qkv": "mxfp4", "*linear_attn.in_proj_z": "mxfp4", "*linear_attn.out_proj": "mxfp4"}'.

Notes

The MXFP6 projections need the bf16 weights, which prefill still reads, and add packed MXFP6 copies, about 3.9 GB a GPU on the older layout. A W6A16 GEMM for prefill would let the bf16 weights go.
The earlier version also issued both ug tasks' loads of a CTA before the first task's MFMAs. On [ROCm][Perf] Qwen3.8-2.4T specialized decode on MI355X (--monokernel) #60833's ug tasks that run in two halves, that spills about 70 VGPRs at 1 to 5 rows, so it is left out.
[ROCm][Perf] Qwen3.8-2.4T specialized decode on MI355X (--monokernel) #60833 at 63036c1291 refuses speculative decoding, so the MTP verify steps that the attention front kernel could serve do not reach it here.

Test plan

pre-commit on every commit
CPU tests in tests/models/qwen3_5/amd_specialized_decode/
gfx950 compile of every MXFP6 build, no spills
Kernel tests on 8 MI355X GPUs: test_kernels.py (GDN block with q6 and qp, layer block with o6 and rm, attention front with and without the KV cache write)
1M ITL, gsm8k and the needle prompts against [ROCm][Perf] Qwen3.8-2.4T specialized decode on MI355X (--monokernel) #60833 at 63036c1291

mergify · 2026-10-12T20:20:04Z

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

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

fxmarty-amd · 2026-10-12T21:39:37Z

+        # VLLM_QWEN38_FAKE_MXFP6_PROJ=1 rounds the bf16 projections to MXFP6
+        # here, after loading and before the mono kernels or a CUDA graph
+        # read them.
+        from vllm.models.qwen3_5.amd import mxfp6
+
+        if mxfp6.ROUND_PROJ:
+            mxfp6.round_projections(self)
+        # The MXFP6 copies for the mono kernels are packed from the weights
+        # after the rounding, so that they hold the same values.
+        from vllm.models.qwen3_5.amd import mono_decode
+
+        if mono_decode.FP6:
+            mono_decode.pack_fp6(self)
+        return loaded


Please refer to https://docs.vllm.ai/en/stable/features/quantization/online/

VLLM_QWEN38_FAKE_MXFP6_PROJ and VLLM_QWEN38_MONO_FP6 env vars should be removed, please use standard --quantization-config API.

At 2 and 3 rows, K2 gives every expert pick its own slot (`up`) instead of building the union of the picked experts. An expert that two tokens pick gets two slots, and each token's contributions are summed in top-k order, so the outputs are those of the union in expert order. At 2 rows, the fused shared expert's ug also starts before the route lands (`ps`). K2 also issues some loads in a better order. The comment next to each change explains why. - A router CTA issues its router weights after run_oproj. Before, they delayed its o_proj, which every norm task waits for. - A CTA with at most two ug tasks issues both tasks' loads before the first task's MFMAs, with no loop that carries the tiles. - When one batch covers the union, the down rounds run without a prefetch, and the second row group's tiles go out behind the INTER polls. A prefetch past the last batch reads no memory. This is the first commit of vllm-project#61271 on the fixed schedule of vllm-project#60833. The W4A16 option for MXFP4 projections (`wa`) is left out, because the MXFP6 projections of the next commit replace it and W4A16 does not match the stock decode. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

mergify · 2026-10-20T02:29:36Z

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

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

Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Mikko Tukiainen <Mikko.Tukiainen@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

…r block load order At 2 and 3 rows, the layer block gives every expert pick its own slot (`up`) instead of building the union of the picked experts. An expert that two tokens pick gets two slots, and each token's contributions are summed in top-k order, so the outputs are those of the union in expert order. At 2 rows, the fused shared expert's ug also starts before the route lands (`ps`). The layer block also issues some loads in a better order. The comment next to each change explains why. - A router CTA issues its router weights after run_oproj. Before, they delayed its o_proj, which every norm task waits for. - When one batch covers the union (up to 3 rows), the down rounds run without a prefetch, and the second row group's tiles go out behind the INTER polls. A prefetch past the last batch reads no memory. The earlier version of this commit also issued both ug tasks' loads of a CTA before the first task's MFMAs. On the ug tasks that run in two halves, that needs about 70 VGPRs more than the CTA has at 1 to 5 rows, so it is left out. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>

The online MXFP4 in_proj_qkvz, out_proj and o_proj keep the gsm8k score within noise, but they make the answers longer. In gsm8k 5-shot with max_gen_toks 2048, the mean answer is 564 characters with MXFP4 projections against 433 with bf16, and 27 answers run into the token limit against 8. MXFP6 E2M3 brings that back to 452 characters and 9 answers, and the kernels take about the same time as with MXFP4. --additional-config '{"monokernel_projections": "mxfp6"}' together with --monokernel selects them. After the weights load, runtime/mxfp6.py rounds the bf16 in_proj_qkvz, out_proj, qkv_proj and o_proj to MXFP6 in place, so the prefill GEMMs multiply the same values that decode reads. It then packs MXFP6 copies of in_proj_qkvz, out_proj and o_proj for the kernels. A copy holds 32 codes in 24 bytes and one e8m0 scale per 32 weights along K, in tiles of 16 rows and 128 columns, so that one v_cvt_scalef32_pk32_bf16_fp6 turns 32 codes into bf16 exactly. The GDN block multiplies in_proj_qkvz in MXFP6 (`q6`) and the layer block out_proj or o_proj (`o6`), both with the bf16 rows (W6A16). The bf16 x rows sit in LDS, so these builds serve up to 8 rows, and wider steps run vLLM's decode. The setting needs the projections in bf16, for example with `--quantization-config.targets '{"*shared_expert*": "mxfp4"}'`. The loaded-model check refuses any other format of them, with the other reasons. An unknown value of monokernel_projections is refused at startup. The copies come on top of the bf16 weights, which prefill still reads. The layer block's router tasks now read the norm tasks' bf16 AR sums and sums of squares as tagged pairs (XAR and MSQ), so the norm tasks' drain of their xbuf stores is not on the router's path. The other consumers read xbuf after MSQ2, which a norm task puts after the drain. The norm task also loads its residual before it polls the partials. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>

…V cache write Between the layer block of one layer and the attention of a full-attention layer, the step ran the residual add with input_layernorm, qkv_proj, fused_qk_rmsnorm_rope_gate and reshape_and_cache_flash as separate launches. afront.py runs all of them in one launch for up to 8 rows when qkv_proj has an MXFP6 copy, and the MXFP6 projections now also pack qkv_proj when it has no bias. Every CTA builds the x rows itself, so no CTA waits for another before its qkv_proj tasks. The q and k tasks hand their rows through a mailbox to the CTAs that run the q and k norms and RoPE, and the CTAs with two tasks run gate rows, because nothing in the kernel waits for them. The mailbox is the bounded one of the other two blocks, so a wait that gives up records itself in the rank's health words. The kernel also writes k and v into the FP8 KV cache as reshape_and_cache_flash does. The attention then runs with q only (_attention_q_only), so vLLM does not write them again. _af_kv keeps vLLM's cache write for any other attention backend, KV cache dtype or layout, for a shared KV cache, a quantized query or a slot_mapping without the pad rows, and logs the reason once. Every row has its own position and KV cache slot, so the kernel would also serve the full-attention layers of MTP verify steps. Against the model's kernels, only the order of the fp32 sums differs, and whether the compiler fuses the RoPE products into FMAs. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>

…he layer block The layer block's router could only start after the all-reduce of o_proj and the post-attention norm had built the x rows. With `rm`, at up to 5 rows, the router logits come from the core rows instead, so every CTA routes as soon as its rstd is in. rm_pack folds the router rows, the norm's (1 + w) and this rank's o_proj rows into one bf16 weight a layer, and appends the scaled router rows for this rank's slice of the residual. Over the 8 ranks, that weight times [core rows | residual slice] sums to the router logits before the rstd, which the route applies when the norm sums are in. The MXFP6 projections make that weight for every layer from the MXFP6 o_proj values, which the layer block multiplies. An MXFP4 o_proj runs without rm. The folded weight is rounded to bf16, so the logits differ from the stock router's in rounding, and a near tie between two experts can go the other way. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>

…MXFP6 in_proj With an MXFP6 in_proj_qkvz (`q6`), the GDN block now places its in_proj tasks by what waits for them (`qp`). The q, k and v rows, which the conv of every GDN task polls, come from CTAs with one in_proj task and no front. The CTAs with two in_proj tasks run z rows, which only the gated norm at the end reads, and they issue their second task's weights after the x loads. When the GDN tasks fill one round, the owner tasks of the key heads move to CTAs 64 to 95, whose in_proj tasks are v rows, so the front CTAs and the CTAs with two tasks run no owner task. The outputs are those of the build without `qp`, bit for bit, which test_gdn_block_mxfp6_matches_vllm_ops checks. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>

…FP6 paths The README describes the MXFP6 projections and the build options that come with them (q6 with qp, o6, rm, the attention front kernel) and the slot per pick at 2 and 3 rows. test_builds checks on the CPU which builds the executor compiles: with the MXFP6 projections the q6 and qp GDN builds, the o6 layer builds with rm up to 5 rows, and both attention front builds of every width. Without them, the builds stay those of the stock formats. The kernel tests of the new options came with their commits: test_gdn_block_mxfp6_matches_vllm_ops (q6, and qp bit for bit against the build without it), test_layer_tp8_matches_vllm_ops (o6 with and without the output gate, rm with the route checked against rm's own logits, and the slot per pick at 3 rows), test_attention_front_matches_vllm_ops (with and without the KV cache write) and test_mxfp6.py (the rounding and the packed layout, on the CPU). Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>

mergify Bot added ci/build new-model Requests to new models qwen Related to Qwen models kimi k3 rocm Related to AMD ROCm labels Oct 12, 2026

github-project-automation Bot added this to AMD Oct 12, 2026

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

mergify Bot added the needs-rebase label Oct 12, 2026

fxmarty-amd reviewed Oct 12, 2026

View reviewed changes

maeehart force-pushed the qwen38-mono-decode-followup branch from d22391e to 2cb5450 Compare October 14, 2026 19:02

mergify Bot removed the needs-rebase label Oct 14, 2026

mergify Bot added the needs-rebase label Oct 20, 2026

mjkvaak-amd and others added 7 commits October 28, 2026 16:24

[ROCm][Perf] Qwen3.8-2.4T specialized decode on MI355X (--monokernel)

63036c1

Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Mikko Tukiainen <Mikko.Tukiainen@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

maeehart force-pushed the qwen38-mono-decode-followup branch from 2cb5450 to 5c63c21 Compare November 10, 2026 20:42

Navigation Menu

Repository navigation

[ROCm][Perf][Qwen3.8] Mono decode follow-up: MXFP6 projections, attention front kernel, router through o_proj - #61271

[ROCm][Perf][Qwen3.8] Mono decode follow-up: MXFP6 projections, attention front kernel, router through o_proj#61271
maeehart wants to merge 7 commits into
vllm-project:mainfrom
maeehart:qwen38-mono-decode-followup

maeehart commented Oct 12, 2026 •

edited

Loading

Uh oh!

mergify Bot commented Oct 12, 2026

Uh oh!

fxmarty-amd Oct 12, 2026

Uh oh!

mergify Bot commented Oct 20, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

maeehart commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What this PR does

Numerics

Validation

Notes

Test plan

Uh oh!

mergify Bot commented Oct 12, 2026

Uh oh!

fxmarty-amd Oct 12, 2026

Choose a reason for hiding this comment

Uh oh!

mergify Bot commented Oct 20, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

maeehart commented Oct 12, 2026 •

edited

Loading