[Fangzhou-Ai]

Purpose

On ROCm, MiniMax-M3.1 decode runs one separate Triton _kernel launch (the MXFP8 activation quant, _mxfp8_quant_triton_kernel) in front of several MXFP8 linears. #63637 fused it into input_layernorm for qkv_proj. Two more per-layer launches remain on the FFN side:

1. post-attention add+RMSNorm -> quant -> gate_up_proj (dense MLP, or the MoE shared expert), every layer, every N.
2. swiglu -> quant -> down_proj, every layer when the GEMM does not already quantize in-kernel (decode batches with more than 2 requests, i.e. M = 8N > 16). Dense layers skip the separate quant in the same way when the in-GEMM quant applies.

This PR (two commits) turns both into epilogues of kernels that already run:

• Commit 1: _gemma_fused_add_rmsnorm_kernel gets an ALSO_BF16 mode: it writes the MXFP8 values/scales and the bf16 normed rows (the router and routed experts need bf16). New allreduce_gemma_rms_norm_mxfp8 (all-reduce + that norm) returns None when a fused all-reduce+norm kernel applies, so that path is untouched. The MoE passes the quantized rows to the shared expert through set_prequant (checked by data pointer + shape, consumed once).
• Commit 2: the dense MLP / shared expert call the existing swiglu_oai_quantize_mxfp8 (fp32 swiglu + quant in one pass) and give down_proj a QuantizedActivation; skipped where the GEMM already quantizes in-kernel (M <= 16, K <= 1024).

Related: #63637 (same idea for input_layernorm -> qkv_proj); the fp32 swiglu kernel is from the MiniMax-M3 AMD ops. Not covered by any open PR that I could find.

What changes per decode step (main model, 60 layers)

Trace evidence (torch profiler, TP4, DSpark 7-token drafts so M = 8 x requests; rank-0 profile windows at N=1, 2, 4 and 16): _kernel.kd per step = 145 at N=1 and N=2, 202 at N>=4 (profiled durations ~1.6x inflated).

Launch removed	Layers	N=1 (M=8)	N=2 (M=16)	N>=3 (M>=24)
post-attn quant for gate_up (commit 1)	60	no (aiter fused AR+norm runs)	yes	yes
swiglu quant for down_proj (commit 2), shared expert	57	no (in-GEMM quant)	no (in-GEMM quant)	yes
swiglu quant for down_proj (commit 2), dense layers	3	yes	yes	yes

Still unfused after this PR: the o_proj input quant (K=2048; in-GEMM quant did not win above M=16 in my sweep), layer 0's input norm, and the DSpark draft layers.

Kernel-level measurements (MI355X, gfx950, 1 GPU, CUDA-graph replay of 200 back-to-back launches; us per launch)

Benchmark script below (collapsed), MiniMax-M3.1 TP4 shapes: hidden 6144, shared-expert intermediate 768.

M	norm + separate quant	fused norm+quant+bf16	saved	swiglu + separate quant	fused swiglu+quant (block_m=16)	saved
8	4.27	2.71	1.56	3.82	1.83	1.99
16	4.32	2.77	1.55	3.96	1.93	2.03
32	4.40	2.80	1.60	3.97	1.95	2.02
64	4.79	2.96	1.83	4.00	1.94	2.06
128	5.55	3.29	2.26	4.06	1.90	2.16
192	6.94	3.65	3.29	4.04	1.95	2.09

(second run on the production code of this branch, GPU 6; an earlier run on GPU 7 agreed within 0.1 us.)

Microbenchmark script (CUDA-graph replay of 200 back-to-back launches)

"""Correctness + graph-replay microbench of the FFN quant fusions. 1 GPU.
Run: PYTHONPATH=<overlay dir> python bench_fuse.py  (overlay dir contains vllm/)"""
import itertools, sys, torch
import vllm
print("vllm from", vllm.__file__)
from vllm.models.minimax_m3.amd.ops import (gemma_fused_add_rmsnorm, gemma_fused_add_rmsnorm_mxfp8,
    gemma_fused_add_rmsnorm_mxfp8_and_bf16, swiglu_oai_split, swiglu_oai_quantize_mxfp8)
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import mxfp8_e4m3_quantize
from vllm.model_executor.kernels.linear.mxfp8 import rocm_native as rn
dev = "cuda"; torch.manual_seed(0)

def graph_time(fn, iters=200, reps=20):
    fn(); torch.cuda.synchronize()
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        for _ in range(iters): fn()
    for _ in range(3): g.replay()
    torch.cuda.synchronize()
    ts = []
    for _ in range(reps):
        s = torch.cuda.Event(enable_timing=True); e = torch.cuda.Event(enable_timing=True)
        s.record(); g.replay(); e.record(); torch.cuda.synchronize()
        ts.append(s.elapsed_time(e) * 1000 / iters)
    ts.sort(); return ts[len(ts)//2]

H = 6144; I = 768
w = (torch.randn(H, device=dev) * 0.1).bfloat16()
for M in (8, 16, 32, 64, 128, 192):
    x = (torch.randn(M, H, device=dev) * 3).bfloat16(); r = torch.randn(M, H, device=dev).bfloat16()
    # A: norm+quant+bf16
    xq, xs, nb, ro = gemma_fused_add_rmsnorm_mxfp8_and_bf16(x, r, w, 1e-6)
    ref_n, ref_r = gemma_fused_add_rmsnorm(x, r, w, 1e-6)
    rq, rs = mxfp8_e4m3_quantize(ref_n)
    okA = (nb == ref_n).all().item() and (xq.view(torch.uint8) == rq.view(torch.uint8)).all().item() \
        and (xs.view(torch.uint8) == rs.view(torch.uint8)).all().item() and (ro == ref_r).all().item()
    # unfused vs fused timing
    t_un = graph_time(lambda: (lambda n, _: mxfp8_e4m3_quantize(n))(*gemma_fused_add_rmsnorm(x, r, w, 1e-6)))
    t_norm = graph_time(lambda: gemma_fused_add_rmsnorm(x, r, w, 1e-6))
    t_f = graph_time(lambda: gemma_fused_add_rmsnorm_mxfp8_and_bf16(x, r, w, 1e-6))
    # B: swiglu+quant
    gu = (torch.randn(M, 2 * I, device=dev) * 8).bfloat16()
    aq, asc = swiglu_oai_quantize_mxfp8(gu, 1.702, 1.0, 7.0)
    act = swiglu_oai_split(gu, 1.702, 1.0, 7.0)
    bq, bs = mxfp8_e4m3_quantize(act)
    nq = (aq.view(torch.uint8) != bq.view(torch.uint8)).float().mean().item()
    ns = (asc.view(torch.uint8) != bs.view(torch.uint8)).float().mean().item()
    t_su = graph_time(lambda: mxfp8_e4m3_quantize(swiglu_oai_split(gu, 1.702, 1.0, 7.0)))
    t_sw = graph_time(lambda: swiglu_oai_split(gu, 1.702, 1.0, 7.0))
    t_sf = {bm: graph_time(lambda: swiglu_oai_quantize_mxfp8(gu, 1.702, 1.0, 7.0, block_m=bm)) for bm in (4, 8, 16, 64)}
    print(f"M={M:4d} A bitexact={okA} norm={t_norm:.2f} norm+quant={t_un:.2f} fused={t_f:.2f} us | "
          f"B mismatch q={nq:.2e} s={ns:.2e} swiglu={t_sw:.2f} swiglu+quant={t_su:.2f} fused(bm)={ {k: round(v,2) for k,v in t_sf.items()} } us", flush=True)

PROJECTION made before the step-level runs (NOT a measurement, and about 2x too high compared with the A+B-only arm below, which averaged -0.7% over N and about -0.1 to -0.4 ms at N>=3): A+B remove 0.22-0.32 ms/step at N>=3 (~1.8-2.3%), nothing at N=1.

Test Plan / Correctness

• Fused norm output (values, scales, bf16 normed, residual): bit-identical to unfused norm -> mxfp8_e4m3_quantize at all M in the table above (edge-case rows: all-zero, tiny, dominated and partly-zero blocks in the unit test).
• Swiglu+quant: scales identical in the microbench (0 differing blocks at M=8..192); 1.9-2.4% of E4M3 bytes differ (M=8: 2.36%, 16: 1.93%, 32: 2.22%, 64: 2.26%, 128: 2.17%, 192: 2.17%) because the fused kernel skips the bf16 rounding of the activation before quantizing; the test asserts it is at least as close to the fp32 activation as the unfused chain.
• Unit tests added to tests/kernels/test_minimax_m3_amd_ops.py (norm variant bitwise, allreduce_gemma_rms_norm_mxfp8 for TP 1/4 and deferral to a fused AR kernel, prequant hand-off, gate_up/down_proj on the fused outputs, swiglu+quant vs unfused). Run on MI355X (gfx950): pytest tests/kernels/test_minimax_m3_amd_ops.py -> 128 passed, 6 failed. The 6 failures are test_mxfp8_native_moe / test_mxfp8_grouped_gemm_native with AttributeError: '_moe_C' has no attribute 'moe_align_block_size' (the installed _moe_C extension of the test machine lacks the op); they fail identically on the unmodified base commit c1d8fad.

Launch counts and step-level measurements (MI355X, TP4, DSpark, 150K context)

All server settings as the AgentX recipe with DBO off and VLLM_ROCM_USE_SKINNY_GEMM=0; decode_bench.py (3 repeats per arm, p50 decode-step time; run-to-run noise about +-0.1-0.2 ms, ~2%). Torch-profiler window of 60 steps, rank 0, all streams.

Exact launch-count reduction (profile window, N=4 and N=16 identical): _kernel.kd (the MXFP8 activation quant) 202 -> 82 per decode step (-120 = 60 layers x 2). Total kernels per step in the A+B-only arm: 1348 -> 1228 on the main stream (-120). What remains of _kernel.kd: the o_proj input quant (60), layer 0's input norm quant and the draft layers (22).

A+B only (this PR's two commits on top of the installed tree), decode-step p50, ms:

N	installed (3 reps)	A+B (3 reps)	delta
1	9.22 9.16 9.28	9.22 9.20 9.19	-0.02 (-0.2%)
2	10.18 10.21 10.08	10.20 10.24 10.26	+0.06 (+0.6%)
3	11.86 11.56 11.76	11.73 11.39 11.34	-0.37 (-3.1%)
4	11.94 11.88 11.87	11.87 11.76 11.90	-0.01 (-0.1%)
5	13.21 12.92 12.82	12.81 12.78 12.79	-0.13 (-1.0%)
6	13.25 13.70 13.24	13.27 13.27 13.30	+0.02 (+0.2%)
8	14.65 14.65 14.34	14.74 14.70 14.77	+0.09 (+0.6%)
12	16.47 16.47 16.35	16.28 16.35 16.32	-0.15 (-0.9%)
16	18.58 18.55 18.41	18.19 18.12 18.07	-0.43 (-2.3%)

Average over N=1..16: about -0.7%. Only N=3 (-3.1%) and N=16 (-2.3%) are outside the noise band; the rest is within it, and the sign is not monotonic in N. I therefore do not claim a specific end-to-end gain for this PR; what is established is the exact launch removal and the per-launch kernel savings above (1.6-3.3 us for the norm epilogue, ~2.0 us for the swiglu fusion). At N=1 and N=2 the PR is expected to do little (see Limitations).

Full stack (this PR + two other, separate changes: folding the routed scaling into the AITER top-k, and launching the shared expert after the routed chain in the HIP graph). Decode-step p50 improved at every N (-0.4% to -4.6%, mean -1.7%; N>=3 mean -0.17 ms); this is not attributable to this PR alone and is listed only for context. A repeat with 5 repetitions at N=1 and N=2 confirmed the stack's N=1 gain (9.22 -> 8.96 ms, -2.8%) and N=2 (10.20 -> 9.93, -2.6%); the A+B-only arm shows nothing at N=1/N=2 (-0.2%/+0.6%), and the scale-fold-only arm accounts for about -2.1% / -1.0% there. The full-stack N=1 gain is not additive across the single-change arms and I could not explain it.

Quality

Full stack (not this PR alone), GSM8K, all 1319 test questions, chat mode with reasoning allowed, greedy, real (not synthetic) draft acceptance: strict match 0.9363 (installed) vs 0.9424 (stack), flexible 0.9356 vs 0.9416; binomial SE about 0.008, i.e. within the noise. Outputs are not bit-reproducible even within one tree (spec decode, batching), so only accuracy statistics are compared. "What is 17*23?" -> 391 on both. The fused norm output is bit-identical to norm -> quant (unit tests and microbench); the swiglu fusion changes 1.9-2.4% of FP8 bytes (see above), which this GSM8K run cannot isolate.

AgentX replay: not run.

Limitations

• No gain at N=1 for the post-attention fusion (the aiter fused all-reduce+norm kernel has no quant output); a quant-emitting variant of that kernel could plug into allreduce_gemma_rms_norm_mxfp8.
• Relies on the RocmDotScaledMxfp8LinearKernel consuming row-major E8M0 scales (same condition as [ROCm][MiniMax-M3.1] Fuse MXFP8 activation quant into the Gemma input_layernorm for qkv_proj #63637).
• Dispatch assumption: the MoE runner hands the shared expert the layer's own input tensor; otherwise the hand-off is ignored and the linear quantizes itself (correct, just not fused).

AI assistance: this change was developed with Claude Code; the tests and benchmarks above were run on MI355X (gfx950).

🤖 Generated with Claude Code

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

[selimfedakar]

Thanks for this — it's an unusually well-documented PR, and the honesty in the
perf section (separating the solid launch-count/per-kernel numbers from the
end-to-end deltas that sit in the noise, and explicitly not claiming an e2e
gain) makes it much easier to review. I went through the diff against main
and traced the surrounding code; overall this looks correct and I'm in favor of
merging. No blocking issues. A few questions and nits below, plus the things I
specifically checked and found safe.

Context: we've been working the MXFP8 activation scheme invariant on the
linear side (the native MXFP8 kernel must honor the activation scheme:
self-quantizing kernels advertise kMxfp8Dynamic with row-major E8M0 scales,
weight-only stays BF16). This PR is a producer for exactly that contract, so I
paid most attention to whether the fused epilogues emit activations the consumer
actually expects.

What I checked and found safe

• Norm epilogue is bit-identical to the unfused path. The new ALSO_BF16
  branch in gemma_rmsnorm.py stores the normed row before the QUANT_MXFP8
  block and still re-casts through out.to(x_ptr.dtype.element_ty), so the FP8
  values and E8M0 scales are computed from the same input-dtype-rounded row as
  norm -> mxfp8_e4m3_quantize. The bitwise-eq test covers bf16/fp16 and N in
  {2048, 6144}.
• Scale layout / scheme contract matches the consumer. The fused ops emit
  row-major [M, N//32] E8M0 scales (not swizzled) under kMxfp8Dynamic, which
  is exactly what RocmDotScaledMxfp8LinearKernel consumes (row-major, not
  swizzled). The gate _can_fuse_input_norm_mxfp8 requires that kernel +
  supports_quantized_input(...), so a FlashInfer-style consumer that wants
  swizzled scales under the same key can never receive these — the same
  layout-vs-key distinction matters here and it's handled correctly.
• down_proj in-GEMM-quant carve-out reproduces the kernel's own decision
  (0 < M <= _FUSED_QUANT_MAX_M and _select_fused_quant_cfg(...) is not None),
  so no double-quant / no missed fusion; the test pins in_gemm == (m<=16 and k<=1024).
• Prequant hand-off is fail-safe: forward only substitutes when the
  incoming x is a tensor with matching data_ptr() and shape, consumes it
  once, else falls back to self-quantization. DP/EP paths fall back cleanly.
• Batch-invariant mode is fully disabled (_can_fuse_input_norm_mxfp8
  returns False under VLLM_BATCH_INVARIANT), so the down_proj path not
  re-checking it is harmless.
• None-deferral keeps the fused-AR / N=1 path byte-for-byte unchanged.
• Perf claims appropriately scoped (per-launch savings + exact launch-count
  reduction are the credible core; e2e correctly not claimed).

Questions (non-blocking)

1. The fused swiglu path is numerically different (not bit-identical) by
   design — it quantizes the fp32 activation directly and skips the bf16
   round-trip the unfused chain does (~1.9-2.4% of E4M3 bytes differ, and your
   test only bounds it as at least as close to the fp32 ref). The op documents
   this, but the substitution at the MiniMaxM3MLP.forward call site is silent —
   a one-line comment there would save the next reader the trip.
2. The 6 failing kernel tests (moe_align_block_size missing) reproduce on the
   base commit too, i.e. a test-machine extension gap, not a regression — agreed,
   but worth confirming they pass on a real gfx950 CI image before merge.
3. return_normed=self.is_moe_layer: confirming the dense path has no consumer
   of the post-attn normed output other than gate_up_proj (which takes MXFP8
   directly), so discarding normed for dense is safe.

Nits

• _gemma_fused_add_rmsnorm's return-type annotation was dropped when
  also_bf16 was added (now a 3- or 4-tuple) — worth restoring.
• _SWIGLU_QUANT_BLOCK_M = 16 — note it was tuned on gfx950 specifically (the
  MoE default is 64).
• PR body says swiglu fusion applies to dense "always"; the code correctly also
  skips it when _down_proj_quantizes_in_gemm — minor wording.

Correct, well-tested, and the activation-scheme/scale-layout contract lines up
with the native MXFP8 consumer. The guards all fail safe toward the unfused
path. I'd approve once the swiglu-not-bit-identical comment (Q1) is added and
the two extension-gap failures are confirmed green on real gfx950 CI. Nice work.

[Fangzhou-Ai]

Thanks for the review, addressed in 6c4a595:

• Return annotation on _gemma_fused_add_rmsnorm: restored, using @overload on also_bf16 (Literal[True] -> 4-tuple, Literal[False] -> 3-tuple) so the existing 3- and 4-value unpacks type-check. The repo's mypy hook passes.
• Dense layers / normed output: checked. Dense decoder layers return ffn(qa), residual and never read the normed bf16 output, so skipping it there is safe. MoE layers return ffn(normed, shared_prequant=qa), residual and do use it.
• swiglu numerics: added a comment in MiniMaxM3MLP.forward. The fused path is intentionally not bit-identical: it quantizes the fp32 swiglu result directly and skips the bf16 rounding the unfused path applies first. About 2% of the E4M3 bytes differ, and the result is at least as close to the fp32 activation.
• _SWIGLU_QUANT_BLOCK_M = 16: commented as tuned on gfx950 (MI355X) only.
• PR body "always" wording: you're right, the dense path skips the separate quant when the in-GEMM quant applies. Body updated.
• Failing tests: the 6 failures on my machine are _moe_C.moe_align_block_size extension gaps in this local build and fail identically on the base commit; CI on a gfx950 image should confirm.

[Fangzhou-Ai]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99726 for commit beb2048674b2.

[ChuanLi1101]

LGTM. Fusing the FFN-side MXFP8 quant into the post-attn norm / swiglu epilogues is the right producer for the native MXFP8 linears, and the dense vs MoE bf16-norm handling looks correct. Thanks for keeping the e2e claim honest. Selim’s asks look addressed.