[micah-wil]

NOTE: We should wait to merge this one after this AITER PR makes it into a release ROCm/aiter#6060

Overview

Adds AiterMxfp8LinearKernel, an MXFP8 linear kernel for gfx950 that runs AITER's tuned FlyDSL GEMM (gemm_a8w8_blockscale_bpreshuffle with 1x32 E8M0 scales). It is placed ahead of RocmDotScaledMxfp8LinearKernel (Triton tl.dot_scaled) in the ROCm MXFP8 kernel list. It only claims (N, K) shapes that AITER has tuned rows for; all other shapes fall through to Triton unchanged.
AITER configs: ROCm/aiter#6060 and ROCm/aiter#7411.

Claims

• 1.2-1.9x faster per GEMM than the Triton dot_scaled kernel on the MiniMax-M3.1-Flash TP4 MXFP8 shapes at decode/verify sizes (M = 32-512). The rotating-weight micro-benchmark is below.
• In a DSpark3 conc-32 decode profile, the MXFP8 linear kernels drop from 7.46 ms to 4.28 ms per step (261 launches in both).
• No accuracy change: GSM8K is within run-to-run noise.

Validation

Kernel selection. I logged the kernel chosen for every MXFP8 linear while loading MiniMax-M3.1-Flash-Quark-MXFP4 with TP4 and DSpark:

Layer	N x K (per rank)	Count	Kernel
self_attn.qkv_proj	2560 x 6144	60	AITER
self_attn.o_proj	6144 x 2048	60	AITER
shared_experts.gate_up_proj	1536 x 6144	57	AITER
shared_experts.down_proj	6144 x 768	57	AITER
dense mlp.gate_up_proj (main 3, draft 5)	6144 x 6144	8	AITER
dense mlp.down_proj (main 3, draft 5)	6144 x 3072	8	AITER
DSpark qkv_proj	2304 x 6144	5	AITER
DSpark o_proj	6144 x 2048	5	AITER
DSpark fused context K/V projection	1280 x 6144	1	AITER with ROCm/aiter#7411, Triton without it

With VLLM_ROCM_MINIMAX_M3_MONO=1, every layer selects RocmDotScaledMxfp8LinearKernel.

Per-GEMM latency (µs), Triton dot_scaled -> AITER FlyDSL.

• Setup: MI355X, MXFP8 weights and bf16 activations. The timing covers the full apply_weights path, so the activation MXFP8 quantization is included.

N x K	M=1	M=32	M=128	M=512	M=2048	M=8192
2560 x 6144 (QKV)	8.5 -> 8.1 (1.04x)	12.5 -> 8.9 (1.41x)	23.1 -> 13.5 (1.72x)	34.7 -> 26.8 (1.30x)	84.9 -> 53.3 (1.59x)	228.5 -> 173.7 (1.32x)
6144 x 2048 (o_proj)	5.8 -> 6.2 (0.94x)	10.8 -> 6.7 (1.61x)	11.6 -> 9.2 (1.26x)	24.1 -> 15.5 (1.56x)	50.7 -> 37.2 (1.36x)	189.6 -> 127.6 (1.49x)
1536 x 6144 (shared gate/up)	8.3 -> 7.3 (1.14x)	10.2 -> 7.8 (1.30x)	20.0 -> 11.3 (1.76x)	34.7 -> 20.9 (1.66x)	46.0 -> 47.6 (0.97x)	148.7 -> 116.9 (1.27x)
6144 x 768 (shared down)	3.9 -> 5.2 (0.76x)	6.2 -> 5.4 (1.15x)	7.1 -> 6.4 (1.12x)	10.8 -> 9.8 (1.11x)	28.8 -> 20.3 (1.42x)	101.7 -> 61.7 (1.65x)
6144 x 6144 (dense gate/up)	13.5 -> 10.8 (1.26x)	21.5 -> 12.5 (1.71x)	24.4 -> 19.5 (1.25x)	64.9 -> 36.1 (1.80x)	132.7 -> 98.1 (1.35x)	523.8 -> 358.4 (1.46x)
6144 x 3072 (dense down)	8.9 -> 7.4 (1.20x)	14.0 -> 8.3 (1.68x)	14.6 -> 12.4 (1.18x)	35.1 -> 21.1 (1.66x)	70.3 -> 51.6 (1.36x)	269.6 -> 176.8 (1.52x)
2304 x 6144 (DSpark QKV)	10.0 -> 8.0 (1.24x)	12.2 -> 8.8 (1.38x)	23.1 -> 13.4 (1.72x)	34.7 -> 26.1 (1.33x)	84.2 -> 51.4 (1.64x)	221.4 -> 160.4 (1.38x)
1280 x 6144 (DSpark ctx K/V)	9.8 -> 6.9 (1.42x)	9.9 -> 7.5 (1.32x)	19.0 -> 11.3 (1.69x)	34.8 -> 18.3 (1.90x)	44.6 -> 37.3 (1.19x)	143.2 -> 101.7 (1.41x)

Decode profile. I compared torch-profiler traces of MiniMax-M3.1-Flash (DSpark3, conc 32, TP4) with and without this kernel. The traces also contain my other MiniMax-M3 ROCm changes, which don't touch the MXFP8 linears.

	Triton dot_scaled	AITER FlyDSL
MXFP8 linear time per decode step	7.46 ms	4.28 ms
MXFP8 linear launches per step	261	261

Accuracy. GSM8K 5-shot strict-match, measured with the full MiniMax-M3 ROCm branch (including this kernel) against main:

	main	with branch
conc 64, DSpark3	0.9621	0.9568
conc 8, DSpark7	0.9591	0.9644

e2e results

Conc	Metric	Nightly	Nightly + PR	Change
32	Throughput per GPU (tok/s)	49,084	49,523	+0.9%
32	ITL p50 (ms)	6.76	6.53	−3.4%
32	Interactivity p50 (tok/s/user)	148.0	153.2	+3.5%
64	Throughput per GPU (tok/s)	57,527	59,563	+3.5%
64	ITL p50 (ms)	13.29	12.88	−3.1%
64	Interactivity p50 (tok/s/user)	75.2	77.6	+3.2%

Details

Opt-out: Set VLLM_DISABLED_KERNELS=AiterMxfp8LinearKernel.

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

[micah-wil]

I know we shouldn't have os.getenv here, it's because MiniMax-M3.1 mono added its own envs file https://github.com/vllm-project/vllm/blob/main/vllm/models/minimax_m3/amd/mono/envs.py, I am going to make a separate PR to clean that up

[Fangzhou-Ai]

why we need to wait for AITER PR 7411? From that PR seems that untuned aiter gemm is still faster than vllm tl.dot kernel.

[Fangzhou-Ai]

Is this new gemm path guarded by VLLM_ROCM_USE_AITER etc env?

[micah-wil]

why we need to wait for AITER PR 7411? From that PR seems that untuned aiter gemm is still faster than vllm tl.dot kernel.

Sorry that note was supposed to link the other AITER PR ROCm/aiter#6060, I fixed the description. The second PR ROCm/aiter#7411 is not strictly require for this one since the N=1280 layer will just fall back to Triton.

[micah-wil]

Is this new gemm path guarded by VLLM_ROCM_USE_AITER etc env?

Yep it's guarded under rocm_aiter_ops.is_linear_enabled() (which requires VLLM_ROCM_USE_AITER=1) in is_supported.

[Fangzhou-Ai]

Thanks! LGTM, let's wait for aiter v0.1.28

[Fangzhou-Ai]

Thanks for this, the per-GEMM numbers look good. One compatibility question with two open MiniMax-M3.1 PRs that feed this GEMM a pre-quantized activation: #65478 (fuses the MXFP8 quant into the post-attention norm and into swiglu) and #65508 (the NVFP4 sparse decode kernel emits the MXFP8 input for o_proj).

Both pass a QuantizedActivation (row-major [M, K/32] E8M0 scales, kMxfp8Dynamic) to the linear, and only when the layer's kernel is RocmDotScaledMxfp8LinearKernel. With AiterMxfp8LinearKernel selected first for these shapes, two things happen:

1. apply_weights always calls mxfp8_e4m3_quantize(x), so it cannot take a QuantizedActivation.
2. The model-side gate (_can_fuse_input_norm_mxfp8) fails, so the fusions silently turn off for qkv, o_proj and the shared/dense gate_up/down in every layer. Each linear then goes back to a separate quant launch.

Could AiterMxfp8LinearKernel accept a QuantizedActivation the way RocmDotScaledMxfp8LinearKernel does? Roughly:

• input_quant_key() returning kMxfp8Dynamic, and a supports_quantized_input(layer) that returns True when the weight is 1 byte per element.
• In apply_weights, qa = as_quantized_activation(x, self.input_quant_key()). If qa is not None, use qa.data / qa.scale directly, and take the output dtype and shape from qa.orig_dtype / qa.orig_shape; otherwise quantize as today.

The scale layout needs no change: the row-major E8M0 scales those fusions write are what gemm_a8w8_blockscale_bpreshuffle reads here (viewed as float8_e8m0fnu). I have this as an untested sketch and can share it if useful. On our side I'll make the model-side gate accept either kernel once the aiter kernel declares it can take quantized input, so the two PRs don't depend on landing order.

[peymanr]

PR Pundit review

(Testing PR Pundit Code Review agent. Not all comments may be relevant. This is a one-time automated pass: please address what is useful, no need to request a re-review. Please reply to each point, whether it applies or not: we harvest these replies, and both kinds are enormously helpful for tuning the agent. Reach out to @peymanr for suggestions.)

#		Finding	Where
1	🟠	Engine startup fails with 'Failed to find a kernel' when --linear-backend=aiter and shape untuned: this PR registers a shape-gated kernel under 'aiter'	vllm/model_executor/kernels/linear/__init__.py:364
2	🟠	Model load aborts with AITER exception instead of falling back: this PR's tuned probe catches only ImportError	vllm/_aiter_ops.py:4432
3	🟠	Weight-cache daemon and IPC loader fail with UnsupportedQuantForIPCError on tuned shapes: this PR precedes RocmDotScaled without setting supports_pre_processed_weights=True	vllm/model_executor/kernels/linear/mxfp8/aiter.py:20
4	🟠	New default-first ROCm MXFP8 kernel has no selection-gating or output-parity tests	vllm/model_executor/kernels/linear/mxfp8/aiter.py:1

1. 🟠 Should fix: Engine startup fails with 'Failed to find a kernel' when --linear-backend=aiter and shape untuned: this PR registers a shape-gated kernel under 'aiter'

vllm/model_executor/kernels/linear/__init__.py:364

Adding AiterMxfp8LinearKernel to _LINEAR_BACKEND_KERNEL_MAP['aiter'] makes _resolve_backend_kernels narrow the ROCm MXFP8 list to [AiterMxfp8LinearKernel] whenever --linear-backend=aiter (or linear_backend_per_quant={'mxfp8': 'aiter'}) is set. Because the filtered list is non-empty, the fall-back-to-automatic-selection path no longer runs. Any MXFP8 layer this kernel rejects then makes init_mxfp8_linear_kernel raise ValueError('Failed to find a kernel...') instead of using RocmDotScaledMxfp8LinearKernel or emulation. The kernel rejects a layer when its (N, K) is untuned, the default dtype is not bf16, VLLM_ROCM_MINIMAX_M3_MONO=1 is set, or VLLM_BATCH_INVARIANT is set. The default auto backend is unaffected.

Who is affected: ROCm (gfx950) users serving MXFP8 models with --linear-backend=aiter or an mxfp8: aiter per-quant override where some layer shapes are untuned or another gate fails.

Checked against the code by an independent pass.

Related past reviews: #48735 on vllm/model_executor/kernels/linear/__init__.py, #52231 on vllm/model_executor/kernels/linear/__init__.py, #39538 on vllm/model_executor/kernels/linear/__init__.py

2. 🟠 Should fix: Model load aborts with AITER exception instead of falling back: this PR's tuned probe catches only ImportError

vllm/_aiter_ops.py:4432

is_mxfp8_bpreshuffle_gemm_tuned catches only ImportError, so any other error from AITER's get_mxscale_bpreshuffle_config escapes can_implement and init_mxfp8_linear_kernel raises instead of falling through to RocmDotScaledMxfp8LinearKernel. Examples are a malformed or overridden tuned CSV (pandas/KeyError), an unknown arch, or an AITER with a different signature. On a supported host this is unlikely: is_supported() already requires a gfx95x GPU with AITER linear enabled, so get_gfx/get_cu_num succeed, and the default tuned CSV ships with AITER. The RuntimeError in the evidence only happens without a GPU. Catching Exception (or at least (AssertionError, ImportError, KeyError, RuntimeError)) and returning False would make the fallback robust; note the sibling is_asm_fp4_gemm_tuned has the same narrow catch.

Who is affected: ROCm gfx95x users with AITER linear enabled whose AITER install has a broken or overridden mxscale bpreshuffle tuned CSV, or a mismatched AITER version.

Checked against the code by an independent pass.

Related past reviews: #46184 on vllm/model_executor/layers/fused_moe/experts/flydsl_mxfp8_moe.py

3. 🟠 Should fix: Weight-cache daemon and IPC loader fail with UnsupportedQuantForIPCError on tuned shapes: this PR precedes RocmDotScaled without setting supports_pre_processed_weights=True

vllm/model_executor/kernels/linear/mxfp8/aiter.py:20

AiterMxfp8LinearKernel leaves supports_pre_processed_weights at the base default False, while the RocmDotScaledMxfp8LinearKernel it now precedes sets it to True. For ModelOpt MXFP8, the quant method's supports_pre_processed_weights defers to the kernel's flag. As a result check_ipc_quant_support raises UnsupportedQuantForIPCError at daemon and loader startup as soon as one layer has an AITER-tuned (N, K). That error is re-raised even with fallback enabled, so the weight cache stops working where it worked before. The kernel's process_weights_after_loading only replaces weight/weight_scale, which the daemon exports already preshuffled, so it can most likely just declare supports_pre_processed_weights = True.

Who is affected: Users of ModelOpt MXFP8 checkpoints with --load-format ipc_cache (weight cache daemon) on gfx950 with AITER linear enabled.

Checked against the code by an independent pass.

4. 🟠 Should fix: New default-first ROCm MXFP8 kernel has no selection-gating or output-parity tests

vllm/model_executor/kernels/linear/mxfp8/aiter.py:1

The PR adds AiterMxfp8LinearKernel and puts it first in the ROCm MXFP8 kernel list, but no test references it. The sibling RocmDotScaledMxfp8LinearKernel has dequant-reference parity tests in tests/kernels/quantization/test_rocm_mxfp8_linear.py. Neither the new kernel's selection gating nor its output parity is tested. Untested selection gating: tuned shapes pick AITER, untuned shapes fall through to RocmDotScaledMxfp8LinearKernel, and VLLM_ROCM_MINIMAX_M3_MONO, VLLM_BATCH_INVARIANT, non-bf16 and bmm_batch_size reject it. Untested output parity: process_weights_after_loading and apply_weights (shuffle layout, scale trim and E8M0 view) against a dequantized-BF16 reference. Add a CPU selection test with is_mxfp8_bpreshuffle_gemm_tuned monkeypatched and a gfx950 parity test.

Who is affected: ROCm gfx950 users serving MXFP8 checkpoints (e.g. MiniMax-M3.1) whose (N, K) shapes AITER has tuned; a regression in this path would not be caught by CI.

Checked against the code by an independent pass.

Related past reviews: #15433 on tests/quantization/test_compressed_tensors.py, #52016 on vllm/model_executor/kernels/linear/__init__.py, #43645 on vllm/model_executor/kernels/linear/__init__.py, #52204 on vllm/model_executor/kernels/linear/__init__.py

_{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 69481e58; findings checked and posted by @peymanr; points already raised on this PR are left out.}