[mjkvaak-amd]

Overview

Opt-in --monokernel decode for Qwen3.8-2.4T (amd/Qwen3.8-2.4T-A95B-Quark-MXFP4, Qwen3_5MoeForCausalLM) at TP8 on MI355X. Laid out per the specialized-decode RFC #60904 under vllm/models/qwen3_5/amd/specialized_decode/, as Qwen3.8-Flash-Next's (#63032). A decode step of 1–16 rows runs two persistent FlyDSL launches a layer; everything else runs vLLM's forward. No core changes beyond the --monokernel docstring.

Claims

• Decode TPOT at 1–16 concurrent requests on MI355X TP8: −33% at 1 (8.36 → 5.61 ms, 8k/1k), −18% at 8, −9% at 16; above 16 rows steps run stock (within 1%).
• Accuracy within noise of stock: GSM8K 0.976 vs 0.973, GPQA Diamond (thinking, 198 questions) 0.919 vs 0.934 (±0.019 per run). GSM8K under VLLM_GPU_SYNC_CHECK=error (Triton GDN prefill, since AITER's FlyDSL GDN prefill metadata syncs the host on prefills) completes with 0.968: no host syncs in the decode path. RULER NIAH at 64k / 128k: being measured.
• Follows RFC [RFC]: A standard for Specialized (MonoKernels) Decodes in vllm/models #60904: --monokernel switch, default-deny refusals in three stages (each listing every reason, TP ranks agreeing), one plan() per step, MRV2 only, model-owned code, weights and caches read in place.
• A rank whose cross-rank wait gives up (10 s) marks itself failed; the TP ranks raise together within 16 steps, so the engine stops rather than hangs (tested with a stalled rank).
• Kernel numerics tests (TP8, MI355X) and CPU refusal/routing tests, with a Buildkite MI355 (8 GPU) step.

Validation

Serving (MI355X TP8, vllm/vllm-openai-rocm:nightly-5d35ce3934a7804b298559d2d99eec33792e1eb9 with this branch's vllm/models/qwen3_5/amd/ and model hook overlaid; both arms VLLM_ROCM_USE_AITER=1, AITER fused shared experts, FULL CUDA graphs, FP8 KV cache, prefix caching off, --max-num-seqs 256; stock adds nothing, the candidate adds --monokernel; vllm bench serve random dataset, about 4× concurrency prompts (at least 8), as many warmups as concurrency; mean TPOT, ms):

concurrency	8k/1k stock	8k/1k --monokernel	1k/1k stock	1k/1k --monokernel
1	8.36	5.61 (−33%)	8.33	5.59 (−33%)
2	8.82	6.45 (−27%)	8.81	6.46 (−27%)
4	10.47	8.38 (−20%)	10.35	8.15 (−21%)
8	13.78	11.30 (−18%)	12.80	10.39 (−19%)
12	18.51	16.98 (−8%)	16.34	14.83 (−9%)
16	19.23	17.55 (−9%)	16.51	15.00 (−9%)
24			19.25	19.28
32			21.99	22.14
48			26.51	26.71
64			29.55	29.76

Per launch (µs, TP8 on MI355X, each launch in a CUDA graph of 46, weights rotated over enough copies that the 256 MB MALL does not hold a layer's weights across launches, random routing):

rows	GDN block	layer block (out_proj MXFP4)	layer block (o_proj BF16, gated)
1	13.3	32.2	35.9
2	15.3	38.2	42.9
4	19.8	54.8	59.0
8	25.4	74.1	78.8
12	30.6	99.0	104.7
16	34.7	120.8	127.2

From 8 rows up the layer block's routed experts stream about 6 TB/s of MXFP4 weights; the rest is the all-reduce, norm and route hand-offs (about 3 µs each between CTAs).

Tests (MI355X, vllm/vllm-openai-rocm:nightly-5d35ce3934a7804b298559d2d99eec33792e1eb9 with this branch's vllm/models/qwen3_5/amd/ overlaid, VLLM_ROCM_USE_AITER=1 pytest -v -s tests/models/qwen3_5/amd_specialized_decode): all 41 tests pass.

• test_kernels.py: the GDN block against the Gemma residual add + norm, in_proj, AITER's conv update and fused gated delta rule and RMSNormGated at widths 1, 5, 8, 9, 16 (pad rows included, first layer included), on either MXFP4 GEMM layout; the layer block against o_proj / out_proj, the all-reduce, fused_add_rms_norm, the router's softmax top-10, the MXFP4 experts with the fused shared expert and the MoE all-reduce at widths 1, 3, 8, 9, 15, 16 (each twice: both epoch parities), on all 8 ranks with rank agreement and the health words left clear. A stalled-rank test launches the layer block on rank 0 only: its waits give up after ~10 s, the failed word and rank 0's flag are set, the given-up waits name only mailbox regions of this kernel, and the next launch returns at once.
• test_runtime.py (host side, no kernels launched; skipped off ROCm like test_kernels.py): off without the flag, every-reason refusal messages, parallel/feature refusals, the PIECEWISE warning, and plan() routing.
• Serving: --monokernel compiles 64 builds (92 layers, widths 1–16) at startup, captures FULL graphs, and takes specialized steps.

Details

Launches. Per layer: (1) on the GDN layers, the GDN block: residual add + input_layernorm, in_proj_qkvz (MXFP4) and in_proj_ba, the conv update, the gated delta rule on the cached states, RMSNormGated; on the full attention layers vLLM's input norm and attention run instead, up to o_proj's input. (2) The layer block: out_proj (MXFP4) / o_proj (BF16, with the output gate), the TP all-reduce, the residual add + post_attention_layernorm, the router's softmax top-10 and the shared-expert gate, the MXFP4 experts with the fused shared expert (expert 512), the TP all-reduce. Embeddings, final norm, head and sampling stay vLLM's.

x without an LDS tile. At H = 8192 a 16-row bf16 x is 262 KB, more than a CU's LDS, so neither launch holds it. The norm CTAs put their columns' sums of squares and every consumer builds what it reads of x from global memory (in the GDN block at up to 3 rows every CTA computes the rows' rstd itself, cheaper than the hand-off); from 4 (GDN) / 5 (layer) rows, the norm CTAs instead quantize their columns once into a shared MXFP4 image that in_proj_qkvz and the experts copy, as every CTA rebuilding all of x would move 256 × 2 [S, H] bf16 through the L2s each launch.

MoE. At decode widths stock runs AITER's AITER_MXFP4_MXFP4 path with x and silu(g)·u through dynamic MXFP4 (1×32 blocks, RCEIL scales); the layer block does the same, the intermediate quantized where it is produced, and feeds the codes to the scaled fp4 MFMA. From 4 rows the fused shared expert's gate/up runs before the route is known; from 8 rows down starts on the union slots whose gate/up finished in the full rounds, under the last, partial round.

MXFP4 layouts. The GDN projections are read in the layout of whichever AITER FP4 GEMM vLLM picked for the shape (runtime/guards.py, mxfp4_shuffled), quantizing their inputs with the matching rounding (Triton's vs ASM's RCEIL). Both layouts are covered by test_kernels.py. No weights are duplicated.

Refusals. Config (model __init__): an architecture other than Qwen3_5MoeForCausalLM, TP≠8, PP, DP, EP, EPLB, DCP, PCP, DBO, sequence-parallel MoE, LoRA, KV connectors, routed-expert capture, speculative decoding, sleep mode, CPU weight offload, weight updates, GDN prefix caching, other GDN cache dtypes, an unfused shared expert, MRV1, STOCK_TORCH_COMPILE, other shapes, unnormalized top-k weights, layer scale, non-gfx950 or partitioned GPUs. Model (after load): every layer's modules, quant methods, weight layouts and shapes. Caches (every bind): GDN state dtypes, shapes and slot layout. Without --monokernel none of this runs.

Steps. plan() returns a Step for a pure decode of 1–16 rows with int32 GDN state indices, else None (stock forward). Every width is compiled when the executor is built (FlyDSL COMPILE_ONLY), so no JIT on the first step of a width. PIECEWISE steps run stock.

Shared helpers. The launch ABI and health words, the tagged mailbox with its bounded waits, the build keys, the GEMV stage, peer memory and FailStop are adapted from the Kimi-K3 decode (#60683, by @hyukjlee, after ROCm/ATOM), as in #63032, and trimmed to what these kernels use. The package imports nothing from other models; the kernel cache key covers every source under its mono/, and the Buildkite step triggers only on this model's directories.

Limitations / follow-ups.

• Speculative decoding is refused; MTP verify steps are a follow-up.
• At 4–8 rows the CTAs that run the routed gate/up's last, partial round still own two down row groups each and set the launch's end; moving row groups off them is a follow-up (~3 µs at those widths).
• Failure: every cross-rank wait is bounded (10 s). A rank that gives up records the wait, flags the output rank and sets its sticky failed word; FailStop reads that word after each step without blocking and the TP ranks agree every 16 steps, then raise together. There is no fallback to the stock forward after a failure.
• Startup does not check the FlyDSL version.

---

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.

[mergify]

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

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

[maeehart]

@mjkvaak-amd, I think that the GDN projections in this PR should be MXFP6 instead of MXFP4. With the online MXFP4 targets for in_proj_qkv, in_proj_z and out_proj, the gsm8k score stays within noise, but the answers get about 30% longer and about three times as many of them run into the token cap before the answer line. MXFP6 E2M3 brings the answers back to the bf16 length, and K2 reads it as fast as MXFP4.

Accuracy. lm_eval gsm8k, 5-shot completions, max_gen_toks=2048, 8 concurrent, all 1,319 questions. amd/Qwen3.8-2.4T-A95B-Quark-MXFP4 TP8 on MI355X, vllm/vllm-openai-rocm:nightly-8cbd5d03006c33185f402249ff2b448efd594986, the shared expert in online MXFP4 in every row. The standard error of the score is about 0.006, so the answer length shows the difference more clearly than the score.

GDN projections	Mono decode	gsm8k flexible	Mean answer length (characters)	Answers at the 2048-token cap
bf16	off	0.9492	433	8
online MXFP8	off	0.9560	458	9
online MXFP4	off	0.9424	564	27
online MXFP4	on, 6a1c13b2fc	0.9477	560	21
MXFP6 E2M3	on, K1 and K2 read MXFP6	0.9500	452	9

The MXFP6 row also has the attention qkv_proj and o_proj in MXFP6, so it quantizes more weights than the MXFP4 rows and still keeps the bf16 length.

How to make the MXFP6 weights.

1. Serve the GDN projections in bf16, so the targets keep only the shared expert: --quantization-config.targets '{"*shared_expert*": "mxfp4"}'. in_proj_ba stays bf16 as before.
2. After loading, round every bf16 weight of in_proj_qkvz and out_proj (and the attention qkv_proj and o_proj) to its MXFP6 E2M3 value in place, one block of 32 along K at a time. The block scale follows AITER's dynamic_mxfp4_quant with the E2M3 maximum of 7.5. The block maximum is rounded to 3 mantissa bits, and the scale is 2 to the power of its exponent minus 2, so no value in the block is above 7.5 times the scale. Each value then rounds to the nearest E2M3 value, ties to even. An E2M3 value times a power of two fits in bf16 exactly, so prefill, the steps wider than 8 rows and K1 and K2 all multiply the same values.
3. Pack a copy for the kernels with the OCP E2M3 codes (sign at bit 5, then 2 exponent and 3 mantissa bits) and one e8m0 byte a block. 32 codes are 24 bytes, with code i at bits 6i to 6i+5 in little-endian order, which is the order in which v_cvt_scalef32_pk32_bf16_fp6 reads them. One such instruction turns a lane's 32 codes times 2^(e8m0 - 127) into 32 bf16 values exactly, and they go into the bf16 MFMA with the bf16 activations (W6A16). We store a tile of 16 rows and 128 K as 64 lanes of 24 bytes, with the first 16 bytes of every lane together and the last 8 together, so a lane loads one dwordx4 and one dwordx2. The converted dwords then come out in the order in which K2 reads an MXFP4 tile, so the MFMA loop does not change.

The rounding of step 2:

def round_e2m3(x: torch.Tensor) -> torch.Tensor:
    v = x.float().reshape(-1, x.shape[-1] // 32, 32)
    amax = v.abs().amax(dim=-1, keepdim=True)
    # Adding 2**19 to the fp32 bits of the block maximum rounds its mantissa
    # at the third bit, and a carry moves into the exponent. The mask keeps
    # only the exponent field.
    exp = (((amax.view(torch.int32) + 0x80000) & 0x7F800000) >> 23) - 127 - 2
    # An all-zero block keeps a normal scale, because a scale of zero would
    # turn 0 / 0 into NaN.
    scale = torch.exp2(exp.clamp(-126, 127).float())
    a = v.abs() / scale
    q = torch.where(
        a < 2.0,
        torch.round(a * 8.0) / 8.0,
        torch.where(a < 4.0, torch.round(a * 4.0) / 4.0, torch.round(a * 2.0) / 2.0),
    )
    return (torch.sign(v) * q.clamp(max=7.5) * scale).reshape(x.shape).to(x.dtype)

Cost. In our build on top of 6a1c13b2fc, K2 with an MXFP6 out_proj or o_proj takes 34.7 µs a launch at 2 rows, as fast as with MXFP4. K1 with an MXFP6 in_proj_qkvz takes 14.41 µs at 2 rows, against 13.37 µs with MXFP4 and 28.54 µs with bf16, so about 0.07 ms a decode step over the 69 GDN layers. The bf16 weights stay for prefill and for the steps wider than 8 rows, and the packed copies of in_proj_qkvz, out_proj and o_proj add 3.24 GB a GPU. At --gpu-memory-utilization 0.9 the KV cache goes from 81.44 GiB (7.33M tokens) with the MXFP4 targets to 72.57 GiB (6.53M tokens). A W6A16 GEMM for the stock path would let the bf16 weights go.

[mergify]

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

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

[mjkvaak-amd]

@mjkvaak-amd, I think that the GDN projections in this PR should be MXFP6 instead of MXFP4. With the online MXFP4 targets for in_proj_qkv, in_proj_z and out_proj, the gsm8k score stays within noise, but the answers get about 30% longer and about three times as many of them run into the token cap before the answer line. MXFP6 E2M3 brings the answers back to the bf16 length, and K2 reads it as fast as MXFP4.

Accuracy. lm_eval gsm8k, 5-shot completions, max_gen_toks=2048, 8 concurrent, all 1,319 questions. amd/Qwen3.8-2.4T-A95B-Quark-MXFP4 TP8 on MI355X, vllm/vllm-openai-rocm:nightly-8cbd5d03006c33185f402249ff2b448efd594986, the shared expert in online MXFP4 in every row. The standard error of the score is about 0.006, so the answer length shows the difference more clearly than the score.

GDN projections Mono decode gsm8k flexible Mean answer length (characters) Answers at the 2048-token cap
bf16 off 0.9492 433 8
online MXFP8 off 0.9560 458 9
online MXFP4 off 0.9424 564 27
online MXFP4 on, 6a1c13b2fc 0.9477 560 21
MXFP6 E2M3 on, K1 and K2 read MXFP6 0.9500 452 9
The MXFP6 row also has the attention qkv_proj and o_proj in MXFP6, so it quantizes more weights than the MXFP4 rows and still keeps the bf16 length.

How to make the MXFP6 weights.

1. Serve the GDN projections in bf16, so the targets keep only the shared expert: --quantization-config.targets '{"*shared_expert*": "mxfp4"}'. in_proj_ba stays bf16 as before.
2. After loading, round every bf16 weight of in_proj_qkvz and out_proj (and the attention qkv_proj and o_proj) to its MXFP6 E2M3 value in place, one block of 32 along K at a time. The block scale follows AITER's dynamic_mxfp4_quant with the E2M3 maximum of 7.5. The block maximum is rounded to 3 mantissa bits, and the scale is 2 to the power of its exponent minus 2, so no value in the block is above 7.5 times the scale. Each value then rounds to the nearest E2M3 value, ties to even. An E2M3 value times a power of two fits in bf16 exactly, so prefill, the steps wider than 8 rows and K1 and K2 all multiply the same values.
3. Pack a copy for the kernels with the OCP E2M3 codes (sign at bit 5, then 2 exponent and 3 mantissa bits) and one e8m0 byte a block. 32 codes are 24 bytes, with code i at bits 6i to 6i+5 in little-endian order, which is the order in which v_cvt_scalef32_pk32_bf16_fp6 reads them. One such instruction turns a lane's 32 codes times 2^(e8m0 - 127) into 32 bf16 values exactly, and they go into the bf16 MFMA with the bf16 activations (W6A16). We store a tile of 16 rows and 128 K as 64 lanes of 24 bytes, with the first 16 bytes of every lane together and the last 8 together, so a lane loads one dwordx4 and one dwordx2. The converted dwords then come out in the order in which K2 reads an MXFP4 tile, so the MFMA loop does not change.

The rounding of step 2:

def round_e2m3(x: torch.Tensor) -> torch.Tensor:
    v = x.float().reshape(-1, x.shape[-1] // 32, 32)
    amax = v.abs().amax(dim=-1, keepdim=True)
    # Adding 2**19 to the fp32 bits of the block maximum rounds its mantissa
    # at the third bit, and a carry moves into the exponent. The mask keeps
    # only the exponent field.
    exp = (((amax.view(torch.int32) + 0x80000) & 0x7F800000) >> 23) - 127 - 2
    # An all-zero block keeps a normal scale, because a scale of zero would
    # turn 0 / 0 into NaN.
    scale = torch.exp2(exp.clamp(-126, 127).float())
    a = v.abs() / scale
    q = torch.where(
        a < 2.0,
        torch.round(a * 8.0) / 8.0,
        torch.where(a < 4.0, torch.round(a * 4.0) / 4.0, torch.round(a * 2.0) / 2.0),
    )
    return (torch.sign(v) * q.clamp(max=7.5) * scale).reshape(x.shape).to(x.dtype)

Cost. In our build on top of 6a1c13b2fc, K2 with an MXFP6 out_proj or o_proj takes 34.7 µs a launch at 2 rows, as fast as with MXFP4. K1 with an MXFP6 in_proj_qkvz takes 14.41 µs at 2 rows, against 13.37 µs with MXFP4 and 28.54 µs with bf16, so about 0.07 ms a decode step over the 69 GDN layers. The bf16 weights stay for prefill and for the steps wider than 8 rows, and the packed copies of in_proj_qkvz, out_proj and o_proj add 3.24 GB a GPU. At --gpu-memory-utilization 0.9 the KV cache goes from 81.44 GiB (7.33M tokens) with the MXFP4 targets to 72.57 GiB (6.53M tokens). A W6A16 GEMM for the stock path would let the bf16 weights go.

Thanks for pointing these up! 💯 The issues with the long MXFP4 thinking chains have been fixed in #62750. Alone the Hadamard rotations from #62750 hurt the throughput, but there's another stacked PR #63272 that returns the performance by fusing the rotations with preceding quant kernels of the MXFP4 GEMM. Please have a look and let me know what you think.

[mergify]

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

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