[Bugfix][ROCM] Fix the MXFP8 block scale exponent (#53110)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>
stefankoncarevic committed
Aug 22, 2026 at 03:26 UTC
7f4a1b7e246c5df89077c65831f349f698aa2b73
1 file changed
+9
-2
vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
+9
-2
@@ -107,6 +107,8 @@ def _mxfp8_quant_triton_kernel():
107
ssm,
108
ssk,
109
BLOCK_M: tl.constexpr,
110
+ FP8_MAX: tl.constexpr,
111
+ TINY: tl.constexpr,
112
):
113
pid_m = tl.program_id(0)
114
pid_b = tl.program_id(1) # which 32-element block along K
@@ -118,8 +120,11 @@ def _mxfp8_quant_triton_kernel():
120
mask=m_mask[:, None],
121
other=0.0,
122
).to(tl.float32)
121
- amax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-30) # [BLOCK_M]
122
- sb = tl.floor(tl.log2(amax)) + 127.0
123
+ # Mirror _mxfp8_e4m3_quantize_torch: the scale has to put the block amax
124
+ # at the top of the e4m3 range rather than at 1.0, or small elements of
125
+ # the block end up in the subnormals.
126
+ amax = tl.maximum(tl.max(tl.abs(x), axis=1), TINY) # [BLOCK_M]
127
+ sb = tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0
128
sb = tl.minimum(tl.maximum(sb, 0.0), 254.0)
129
descale = tl.exp2(sb - 127.0)
130
xq = (x / descale[:, None]).to(xq_ptr.dtype.element_ty)
@@ -167,6 +172,8 @@ def _mxfp8_e4m3_quantize_triton(
172
scales.stride(0),
173
scales.stride(1),
174
BLOCK_M=BLOCK_M,
175
+ FP8_MAX=float(torch.finfo(MXFP8_VALUE_DTYPE).max),
176
+ TINY=float(torch.finfo(torch.float32).tiny),
177
)
178
return xq, scales
179