[Bugfix][ROCm] Fix MiniMax-M3 fused MXFP8 block scale (#53674)
Signed-off-by: maithilijoshi20 <97733343+maithilijoshi20@users.noreply.github.com>
maithilijoshi20 committed
Sep 16, 2026 at 11:59 UTC
22bb158c21fcba630f267a12ebab4e86296ccc42
2 files changed
+28
-4
tests/kernels/test_minimax_m3_amd_ops.py
+22
-2
@@ -42,6 +42,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( # noqa:
42
from vllm.models.minimax_m3.amd.ops import ( # noqa: E402
43
gemma_fused_add_rmsnorm,
44
gemma_rmsnorm,
45
+ swiglu_oai_quantize_mxfp8,
46
swiglu_oai_split,
47
)
48
from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import _num_warps # noqa: E402
@@ -167,6 +168,25 @@ def test_swiglu_oai_split(m, inter, limit, dtype):
168
assert _relerr(got, ref) < 5e-3
169
170
171
+@torch.inference_mode()
172
+def test_swiglu_oai_quantize_mxfp8_uses_e4m3_range_for_scale():
173
+ # Keep a small value in two MX blocks next to their maxima. The scale must
174
+ # use E4M3's finite range (448), otherwise that value underflows to zero.
175
+ # Keep the third block empty to cover the reference's finite tiny clamp.
176
+ m, inter = 3, 96
177
+ gate = torch.full((m, inter), 2.0, device=DEVICE, dtype=torch.float16)
178
+ up = torch.full((m, inter), 2**-10, device=DEVICE, dtype=torch.float16)
179
+ up[:, :64:32] = 1.0
180
+ up[:, 64:] = 0.0
181
+ gate_up = torch.cat((gate, up), dim=-1)
182
+
183
+ got_q, got_s = swiglu_oai_quantize_mxfp8(gate_up, alpha=0.0, beta=0.0, limit=None)
184
+ ref_q, ref_s = _mxfp8_e4m3_quantize_torch(up, is_sf_swizzled_layout=False)
185
+
186
+ assert torch.equal(got_s, ref_s)
187
+ assert torch.equal(got_q, ref_q)
188
+
189
+
190
# --------------------------------------------------------------------------- #
191
# Fused MXFP8 activation quant (Triton vs torch reference)
192
# --------------------------------------------------------------------------- #
@@ -179,8 +199,8 @@ def test_mxfp8_quant_triton_matches_torch(shape, dtype):
199
xq_t, s_t = _mxfp8_e4m3_quantize_torch(x, is_sf_swizzled_layout=False)
200
xq_k, s_k = _mxfp8_e4m3_quantize_triton(x)
201
assert s_k.shape == s_t.shape == (shape[0], shape[1] // 32)
182
- # E8M0 block exponents share the floor(log2(amax))+127 algorithm; allow at
183
- # most a 1-step difference at exact powers of two.
202
+ # Both paths use the E4M3-aware scale calculation; allow a 1-step
203
+ # difference at exact powers of two due to floating-point rounding.
204
assert (s_k.int() - s_t.int()).abs().max().item() <= 1
205
# Dequantized values agree to fp8 granularity.
206
deq_t = dequant_mxfp8_to_bf16(xq_t, s_t)
vllm/models/minimax_m3/amd/ops/swiglu_oai.py
+6
-2
@@ -83,6 +83,8 @@ def _swiglu_oai_quant_kernel(
83
limit,
84
HAS_LIMIT: tl.constexpr,
85
BLOCK_M: tl.constexpr,
86
+ FP8_MAX: tl.constexpr,
87
+ TINY: tl.constexpr,
88
):
89
"""SwiGLU-OAI (split layout) fused with per-32-block MXFP8 (E4M3 + E8M0)
90
quant. Each program handles ``[BLOCK_M, 32]`` of the ``[M, I]`` output (one
@@ -109,8 +111,8 @@ def _swiglu_oai_quant_kernel(
111
gate = tl.minimum(gate, limit)
112
up = tl.minimum(tl.maximum(up, -limit), limit)
113
act = gate * tl.sigmoid(alpha * gate) * (up + beta) # [BLOCK_M, 32] fp32
112
- amax = tl.maximum(tl.max(tl.abs(act), axis=1), 1e-30) # [BLOCK_M]
113
- sb = tl.minimum(tl.maximum(tl.floor(tl.log2(amax)) + 127.0, 0.0), 254.0)
114
+ amax = tl.maximum(tl.max(tl.abs(act), axis=1), TINY) # [BLOCK_M]
115
+ sb = tl.minimum(tl.maximum(tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0, 0.0), 254.0)
116
descale = tl.exp2(sb - 127.0)
117
aq = (act / descale[:, None]).to(aq_ptr.dtype.element_ty)
118
tl.store(
@@ -173,6 +175,8 @@ def swiglu_oai_quantize_mxfp8(
175
0.0 if limit is None else float(limit),
176
HAS_LIMIT=limit is not None,
177
BLOCK_M=block_m,
178
+ FP8_MAX=float(torch.finfo(MXFP8_VALUE_DTYPE).max),
179
+ TINY=float(torch.finfo(torch.float32).tiny),
180
num_warps=4,
181
)
182
return aq, asc