[RuibinCheung]

Overview

On ROCm, a MiniMax-M3 decoder layer whose qkv_proj is MXFP8 (e.g. MiniMax-M3.1-Flash-Quark-MXFP4, whose attention projections are 1x32 MXFP8) runs the fused add + Gemma RMSNorm, writes the normed bf16 rows, and then RocmDotScaledMxfp8LinearKernel reads them back and quantizes them with a separate mxfp8_e4m3_quantize launch. This PR quantizes the normed row per 32-element block in the epilogue of the norm kernel and hands qkv_proj a QuantizedActivation, so each layer drops one launch and the bf16 round trip. The fp8 values and E8M0 scales are bitwise equal to norm -> mxfp8_e4m3_quantize, so the model output does not change.

Claims

• Norm + activation quant: 1.41x-2.06x per call (hidden 6144, M = 1-4096, bf16, MI355X); norm + quant + qkv_proj GEMM: 1.07x-1.19x.
• Bitwise equal fp8 values, scales and residual vs. norm -> mxfp8_e4m3_quantize, for bf16 and fp16.
• MiniMax-M3.1-Flash-Quark-MXFP4, TP4, decode: the separate activation quant for qkv_proj disappears in 59 of 60 layers (kernel trace), the norm kernel time is unchanged, and the removed launches are 1.9% of decode GPU kernel time. AgentX P90 interactivity is 1.018x at concurrency 1 and 1.017x-1.045x at concurrency 16 (two runs, order swapped); c4, c8 and c32 are within run-to-run noise; GSM8K unchanged within stderr.
• No change for layers whose qkv_proj is not served by RocmDotScaledMxfp8LinearKernel with fp8 weights, under VLLM_BATCH_INVARIANT, with LoRA on qkv_proj, or with mono decode.

Validation

Unit tests (tests/kernels/test_minimax_m3_amd_ops.py, already run by the AMD kernels CI step): 81 passed on this branch (MI355X, main nightly image with this branch's Python tree), 81 passed + 1 skipped with this PR applied on #63232. tests/kernels/quantization/test_rocm_mxfp8_linear.py and tests/v1/determinism/test_mx_linear_batch_invariant.py (the existing consumer tests of the changed kernel): 97 passed. New tests:

• test_gemma_fused_add_rmsnorm_mxfp8_bitwise_eq_unfused: torch.equal against gemma_fused_add_rmsnorm -> mxfp8_e4m3_quantize for values, scales and residual, n in {2048, 3072, 6144}, bf16 and fp16. Every row has a block with a scale above 1, an all-zero block and a tiny block (1 + w = 2**-8), plus all-zero, tiny, single-block-dominated and partly-zero rows.
• test_fused_input_norm_mxfp8_qkv_proj: a real QuarkOCP_MX MXFP8 layer (per-row scales, 32x32 block scales, and K % 128 != 0, where the weight is dequantized to bf16 at load), M in {1, 37}. Checks the producer gate (fuses iff the kernel reads the fp8 activation as is) and that qkv_proj(QuantizedActivation) is bitwise equal to qkv_proj(norm output).
• test_fused_input_norm_mxfp8_off_under_batch_invariance

Microbenchmark (method as in #59782: one HIP graph of 128 back-to-back calls cycling through 16 input sets, best of 5; MI355X; hidden 6144, bf16; qkv_proj is 2560 x 6144 at TP4). "separate" is gemma_fused_add_rmsnorm + mxfp8_e4m3_quantize (current behaviour), "fused" is gemma_fused_add_rmsnorm_mxfp8. Speedup of fused over separate:

M	1	2	4	8	16	32	64	128	256	1024	4096
norm + quant	1.66x	1.58x	1.58x	1.57x	1.58x	1.59x	1.64x	1.70x	2.06x	1.41x	1.49x
norm + quant + qkv GEMM	1.18x	1.16x	1.15x	1.16x	1.19x	1.15x	1.09x	1.10x	1.10x	1.07x	1.11x

The plain gemma_fused_add_rmsnorm (the kernel gains a QUANT_MXFP8 constexpr) is unchanged: the separate path times within 1% of the pre-PR kernel at every M.

End to end: MiniMax-M3.1-Flash-Quark-MXFP4, TP4, MI355X, mono decode off, FULL_DECODE_ONLY CUDA graphs, on #63232 (8d0f05a0c, image vllm/vllm-openai-rocm:minimax-m31-flash) with this PR applied on top; baseline is the same image unchanged. AgentX trace replay (InferenceX), 600 s per point, max-num-seqs = 2 x concurrency, the arm that runs first alternates per concurrency. opt / baseline (> 1 is better):

Concurrency	P90 interactivity
1	1.018x
4	0.990x
8	1.106x
16	1.017x
32	1.064x

GSM8K (lm-eval, 5-shot, full test set of 1319, greedy, concurrency 16): baseline 96.51 / opt 96.44 strict-match (stderr 0.51 / 0.51), flexible-extract 96.44 / 96.36. The difference is one question; vLLM is not batch invariant at concurrency 16, so greedy outputs vary between runs even with bit-identical kernels.

Details

• Producer and consumer. RocmDotScaledMxfp8LinearKernel declares kMxfp8Dynamic as its input quant key and accepts a QuantizedActivation on the dot_scaled and block32 GEMMs; QuarkOCP_MX exposes the key for MXFP8 layers (ModelOpt MXFP8 already does). The QuantizedActivation contract itself is unchanged.
• Scale layout. kMxfp8Dynamic does not encode the scale layout: the FlashInfer MXFP8 kernels read swizzled scales under the same key, this kernel reads row-major [M, K // 32]. So the producer (_can_fuse_input_norm_mxfp8) also requires the consumer to be RocmDotScaledMxfp8LinearKernel with fp8 weights, the same way DeepSeek V4.1's can_fuse_query_quant checks for the FlashInfer kernels. If a QuantizedActivation reaches the kernel for a weight it dequantized to bf16 at load (K % 128 != 0), apply_weights dequantizes the activation instead of failing; the producer does not fuse in that case.
• Other producers of kMxfp8Dynamic. On main, no other producer fuses into a Quark MXFP8 linear: DeepSeek V4.1's query-quant fusion is CUDA-only and requires the FlashInfer kernels, and maybe_fused_act_quant has no MXFP8 entry. ModelOpt MXFP8 MiniMax-M3 checkpoints on ROCm also take this path (same kernel, same bitwise result).
• First layer. Layer 0 has no residual and uses the non-add gemma_rmsnorm; it keeps the separate quant (1/60 of the calls on M3.1-Flash). A non-add variant is a small follow-up if wanted.
• Batch invariance, LoRA. Off under VLLM_BATCH_INVARIANT, like the other producer-side quant fusions ([ROCm][Kimi-K3] Fuse AttnRes output with per-token FP8 input quantization #59069), so that mode keeps exactly the path tests/v1/determinism covers. With LoRA on qkv_proj, get_input_quant_key returns None and nothing changes.
• Mono decode. The VLLM_ROCM_MINIMAX_M3_MONO executor calls input_layernorm and self_attn directly, so it is unaffected; the end-to-end numbers are with mono off.
• No new env var. The output is bitwise equal, so there is no accuracy reason to turn it off.

---

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.

---

^{Need help on this PR? Tag @codesmith-bot with what you need.}

---

^{Need help on this PR? Tag @codesmith-bot with what you need.}

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

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[shen-shanshan]

The e2e perf results show that there could be regression at conc=4, could you please offer more detailed benchmark results? e.g., high conc part, together with output token throughput, TTFT, TPOT...

[shen-shanshan]

I'm not sure should we guard this within rocm platform?

[RuibinCheung]

Added. ACK. Thx!

[RuibinCheung]

Thanks for looking at this. The single P90-interactivity column in the description was too thin, so here is the full A/B for the same runs. All numbers are opt / baseline, so > 1 is better everywhere.

Concurrency	Output tput/GPU	P90 TTFT	P90 TPOT
1	1.001x	1.015x	1.018x
4	1.006x	1.002x	0.990x
8	1.029x	1.056x	1.106x
16	1.007x	1.053x	1.017x
32	1.065x	1.309x	1.064x

I think the perf regression of c4 is caused by benchmark noise. It's not a real perf regression.

[shen-shanshan]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97534 for commit a1d568916471.

[shen-shanshan]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97851 for commit f0610ae15d86.

[RuibinCheung]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97893 for commit a2066a5defa4.

ℹ️ This PR is 3 commits behind upstream main. Merged upstream main into this branch as a2066a5defa4 before starting CI. Pull the branch before pushing more changes.

[peymanr]

PR Pundit review

(testing PR Pundit Code Review agent. Not all comments may be relevant. Please reach out to @peymanr for suggestions and feedback.)

#		Finding	Where
1	🔵	Fusion decided pre-load from unprocessed layer state	vllm/models/minimax_m3/amd/model.py:1386
2	🔵	Model couples to concrete kernel class via attribute duck-typing	vllm/models/minimax_m3/amd/model.py:1295
3	🔵	Pre-quantized activation path skips scale shape/layout validation	vllm/model_executor/kernels/linear/mxfp8/rocm_native.py:314
4	🔵	Third inline copy of MXFP8 block quant math	vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py:103
5	🔵	Forward fusion branch not exercised by any test	vllm/models/minimax_m3/amd/model.py:1401
6	🔵	Module docstring stale after adding third entry point	vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py:19
7	🔵	Fusion skipped when residual is None; MTP log overstates	vllm/models/minimax_m3/amd/model.py:1398
8	🔵	ModelOpt kernel fallback branch untested and likely unreachable	vllm/models/minimax_m3/amd/model.py:1307

1. 🔵 Nit: Fusion decided pre-load from unprocessed layer state

vllm/models/minimax_m3/amd/model.py:1386

_can_fuse_input_norm_mxfp8 runs in MiniMaxM3DecoderLayer.__init__, before process_weights_after_loading. Pre-load weight_scale is [N, K//32], so the weight_scale.shape[0] != weight.shape[0] branch of supports_quantized_input is dead there and the decision reduces to K % 128 == 0. This is conservative and never fuses wrongly: both post-load fp8 paths (block32 and dot_scaled) consume the pre-quantized input when K % 128 == 0. The only effect is that fusion is skipped for 32x32 block-scale checkpoints with K % 32 == 0 but K % 128 != 0, which rocm_mxfp8_block32_gemm could serve. Consider noting the pre-load evaluation in a comment, or deciding after load.

Who is affected: Only hypothetical block-scaled MXFP8 checkpoints whose hidden size is a multiple of 32 but not 128 on gfx95x lose the fusion; the default MiniMax-M3 hidden_size of 6144 always fuses.

2. 🔵 Nit: Model couples to concrete kernel class via attribute duck-typing

vllm/models/minimax_m3/amd/model.py:1295

_can_fuse_input_norm_mxfp8 imports the concrete RocmDotScaledMxfp8LinearKernel and finds it by attribute name (scheme.ocp_mx_linear for Quark, quant_method.kernel for ModelOpt). If either attribute is renamed, the fused input-norm MXFP8 path turns off without warning; output stays correct and only the speedup is lost. The Quark path is covered by test_fused_input_norm_mxfp8_qkv_proj, which asserts the fusion is on for K % 128 == 0; the ModelOpt lookup is not tested. The class check is a documented workaround because kMxfp8Dynamic does not encode scale layout (row-major vs FlashInfer swizzled), already tracked as a TODO in quant_activation.py. A small accessor such as a kernel getter, or a layout field in the quant key, would remove the duck-typing.

Who is affected: MiniMax-M3 AMD users with MXFP8 (Quark/ModelOpt) checkpoints on gfx950, only if those attributes are later renamed; the result is a silent performance regression, not a correctness bug.

Related past reviews: #50605 on vllm/model_executor/kernels/linear/mxfp8/rocm_native.py

3. 🔵 Nit: Pre-quantized activation path skips scale shape/layout validation

vllm/model_executor/kernels/linear/mxfp8/rocm_native.py:314

The QuantizedActivation branch of RocmDotScaledMxfp8LinearKernel.apply_weights passes qa.scale to _mxfp8_dot_scaled_gemm / rocm_mxfp8_block32_gemm without checking its shape ([M, K//32]), dtype (uint8) or row count against qa.data. kMxfp8Dynamic does not encode scale layout, so a future producer handing in a mismatched 2-D scale would get wrong output rather than an error. No current producer can do this: the only one routed to this kernel is MiniMax-M3's fused norm, which checks the kernel type and writes contiguous row-major scales. The K alignment is already enforced by supports_quantized_input. Cheap defensive asserts would still help.

Who is affected: No current configuration; only a future producer that sends non-row-major kMxfp8Dynamic scales to the ROCm dot_scaled kernel.

4. 🔵 Nit: Third inline copy of MXFP8 block quant math

vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py:103

The QUANT_MXFP8 epilogue in _gemma_fused_add_rmsnorm_kernel repeats the block-quant math of _mxfp8_quant_triton_kernel: the amax/TINY clamp, the sb formula, the [0,254] clamp and the exp2(127 - sb) reciprocal multiply. _swiglu_oai_quant_kernel already had a copy, so this is a third. The docstring's 'bit-identical to mxfp8_e4m3_quantize' refers to the ROCm Triton path, which is what mxfp8_e4m3_quantize dispatches to for these 2D row-major inputs. The new test checks this bitwise with torch.equal, so any future drift would fail the test. Consider factoring the block quant into a shared @triton.jit helper, or naming _mxfp8_quant_triton_kernel as the reference in the comment.

Who is affected: Maintainers changing the MXFP8 quant math on ROCm; there is no runtime impact today.

Verified by running it.

Related past reviews: #46117 on vllm/model_executor/kernels/linear/mxfp8/rocm_native.py

5. 🔵 Nit: Forward fusion branch not exercised by any test

vllm/models/minimax_m3/amd/model.py:1401

No test drives the new fused branch in MiniMaxM3DecoderLayer.forward. test_fused_input_norm_mxfp8_qkv_proj builds the QuantizedActivation by hand and calls scheme.apply_weights directly. The end-to-end amd_specialized_decode/test_model.py excludes self_attn from Quark quantization, so fuse_input_norm_mxfp8 is never True there. A regression in the residual handoff, orig_shape/dtype, or norm weight/eps in this hunk would go unnoticed. The hunk currently reads correctly; consider a small layer-level test, or quantizing qkv_proj in the mini model.

Who is affected: Only maintainers changing this code: it lacks regression coverage on gfx950 with an MXFP8-quantized qkv_proj; no runtime misbehavior was found.

Verified by running it.

6. 🔵 Nit: Module docstring stale after adding third entry point

vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py:19

The module docstring now lists three entry points (gemma_rmsnorm, gemma_fused_add_rmsnorm, gemma_fused_add_rmsnorm_mxfp8), but the next paragraph still says "Both normalize over the last dim ...". That wording is stale and should cover all three. The MXFP8 variant's flattened [M, N] output and N % 32 == 0 requirement are already documented in its own docstring and enforced by an assert. At most, the module paragraph could point to that docstring.

Who is affected: Readers of vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py; documentation only, no runtime effect.

Verified by running it.

Related past reviews: #9138 on csrc/layernorm_quant_kernels.cu

7. 🔵 Nit: Fusion skipped when residual is None; MTP log overstates

vllm/models/minimax_m3/amd/model.py:1398

The MXFP8 fused path only runs when residual is not None, and only gemma_fused_add_rmsnorm_mxfp8 has the MXFP8 epilogue. So the model's first decoder layer on the first PP rank and the MTP transformer_layer (called with residual=None in mtp.py) always run the unfused input_layernorm -> qkv_proj quant path. Non-first PP ranks receive residual from intermediate_tensors, so their first layer does fuse. fuse_input_norm_mxfp8 is still computed as True for the MTP layer even though it can never take that branch. Consider a comment, or a non-residual gemma_rmsnorm MXFP8 variant.

Who is affected: MiniMax-M3 on ROCm with an MXFP8 qkv_proj using RocmDotScaledMxfp8LinearKernel: layer 0 on the first PP rank and the MTP draft layer keep the separate quant (performance only).

8. 🔵 Nit: ModelOpt kernel fallback branch untested and likely unreachable

vllm/models/minimax_m3/amd/model.py:1307

The ModelOpt fallback in _can_fuse_input_norm_mxfp8 (reading quant_method.kernel) is reachable but untested. ModelOptLinearMethod.create_weights calls expose_input_quant_key, and on gfx95x init_mxfp8_linear_kernel selects RocmDotScaledMxfp8LinearKernel, whose key is kMxfp8Dynamic. The new tests only build Quark layers (_quark_mxfp8_linear). Consider adding a ModelOpt MXFP8 qkv_proj case to test_fused_input_norm_mxfp8_qkv_proj.

Who is affected: ModelOpt MXFP8 MiniMax-M3 checkpoints on gfx95x ROCm, where the fusion decision for qkv_proj goes through the untested quant_method.kernel lookup.

Verified by running it.

_{Shared from a pilot of PR Pundit, an agentic code reviewer that builds and tests the change and reads the project's past reviews. Reviewed at a2066a5d; findings checked and posted by @peymanr; points already raised on this PR are left out.}