A high-throughput and memory-efficient inference and serving engine for LLMs
Fix the E8M0 scale computation in the MXFP4 (W4A4) MOE CUTLASS kernel (#43557)
Signed-off-by: Xin He <xin3.he@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Xin He committed
Jun 15, 2026 at 21:04 UTC
5ed15f42b93d4ea7b40f4dfc2dd7f12a44c99d75
3 files changed
+253
-17
csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh
+33
-16
@@ -237,21 +237,30 @@ __device__ __forceinline__ fp4_packed_t cvt_warp_fp16_to_fp4(
237
// Get the final absolute maximum values.
238
float vecMax = float(__hmax(localMax.x, localMax.y));
239
240
- // Get the SF (max value of the vector / max value of e2m1).
241
- // maximum value of e2m1 = 6.0.
242
- // TODO: use half as compute data type.
243
- float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
240
// 8 bits representation of the SF.
241
+ float SFValue;
242
uint8_t fp8SFVal;
246
- // Write the SF to global memory (STG.8).
243
+
244
if constexpr (UE8M0_SF) {
248
- // Extract the 8 exponent bits from float32.
249
- // float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits.
250
- uint32_t tmp = reinterpret_cast<uint32_t&>(SFValue) >> 23;
251
- fp8SFVal = tmp & 0xff;
252
- // Convert back to fp32.
253
- reinterpret_cast<uint32_t&>(SFValue) = tmp << 23;
245
+ // OCP MX spec E8M0 scale computation (MXFP4 path):
246
+ // scale_exp = biased_exponent(round_up(vecMax)) - 2
247
+ // -2 because max E2M1 value is 6.0 ≈ 2^2.58; we use 2^2=4 as the
248
+ // safe divisor so that max_val / scale <= 6.0 for values near 2^n.
249
+ uint32_t max_bits = __float_as_uint(vecMax);
250
+ // Add rounding bias at mantissa bit 21 (equivalent to bf16 val_to_add=32
251
+ // at bit 5). Threshold: values with mantissa >= 0.75 (i.e. >= 1.75*2^n)
252
+ // round up to the next power of 2.
253
+ uint32_t rounded_bits = (max_bits + (1u << 21)) & 0xFF800000u;
254
+ uint32_t biased_exp = (rounded_bits >> 23) & 0xFFu;
255
+ uint32_t scale_exp = (biased_exp > 2u) ? (biased_exp - 2u) : 0u;
256
+ scale_exp = min(scale_exp, 254u);
257
+ fp8SFVal = static_cast<uint8_t>(scale_exp);
258
+ // Reconstruct scale as float32: scale = 2^(scale_exp - 127)
259
+ uint32_t sf_bits = scale_exp << 23;
260
+ SFValue = __uint_as_float(sf_bits);
261
} else {
262
+ // NVFP4 path: scale = max / 6.0, stored as E4M3.
263
+ SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
264
// Here SFValue is always positive, so E4M3 is the same as UE4M3.
265
__nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
266
reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp;
@@ -262,13 +271,21 @@ __device__ __forceinline__ fp4_packed_t cvt_warp_fp16_to_fp4(
271
// Write the SF to global memory (STG.8).
272
if (SFout) *SFout = fp8SFVal;
273
265
- // Get the output scale.
266
- // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
267
- // reciprocal(SFScaleVal))
268
- float outputScale =
269
- SFValue != 0.0f ? reciprocal_approximate_ftz(
274
+ // Get the output scale (= 1 / SFValue for the MXFP4/UE8M0 path where
275
+ // SFScaleVal=1). Use exact division for UE8M0 to ensure bit-exact scaling
276
+ // that matches the reference QDQ implementation (dividing by a power-of-2
277
+ // scale is exact in IEEE 754).
278
+ float outputScale;
279
+ if constexpr (UE8M0_SF) {
280
+ // SFValue is always a power of 2 for UE8M0, so 1/SFValue is exact.
281
+ outputScale = SFValue != 0.0f ? (1.0f / SFValue) : 0.0f;
282
+ } else {
283
+ // NVFP4 path: use fast approximate reciprocal (original behavior).
284
+ outputScale = SFValue != 0.0f
285
+ ? reciprocal_approximate_ftz(
286
SFValue * reciprocal_approximate_ftz(SFScaleVal))
287
: 0.0f;
288
+ }
289
290
// Convert the input to float.
291
float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
tests/kernels/moe/test_mxfp4_moe.py
+219
@@ -244,5 +244,224 @@ def test_mxfp4_experts_quant_basic():
244
print("PASSED")
245
246
247
+def untile_cutlass_scale(scale_raw: torch.Tensor, rows: int, K: int) -> torch.Tensor:
248
+ """Convert CUTLASS tiled scale back to flat [M, K//32] layout.
249
+
250
+ CUTLASS tiled layout: [numMTiles, numKTiles, 32(outerM), 4(innerM), 4(innerK)]
251
+ Produced by: padded.reshape(numMTiles, 4, 32, numKTiles, 4).permute(0,3,2,1,4)
252
+ To undo: tiled.permute(0, 3, 2, 1, 4).reshape(padded_M, padded_sK)
253
+ """
254
+ num_scale_cols = K // MXFP4_BLOCK_SIZE
255
+ num_m_tiles = (rows + 127) // 128
256
+ num_k_tiles = (num_scale_cols + 3) // 4
257
+ padded_M = num_m_tiles * 128
258
+ padded_sK = num_k_tiles * 4
259
+
260
+ scale_bytes = scale_raw.view(torch.uint8).flatten()
261
+ total_bytes = padded_M * padded_sK
262
+ tiled = scale_bytes[:total_bytes].reshape(num_m_tiles, num_k_tiles, 32, 4, 4)
263
+ undone = tiled.permute(0, 3, 2, 1, 4).contiguous()
264
+ return undone.reshape(padded_M, padded_sK)[:rows, :num_scale_cols]
265
+
266
+
267
+def compute_reference_e8m0_scale(block_max: float) -> int:
268
+ """Compute the expected OCP MX spec E8M0 scale for a given block max.
269
+
270
+ The CUTLASS kernel uses round-to-nearest on the mantissa:
271
+ rounded_bits = (float_bits + (1 << 21)) & 0xFF800000
272
+ biased_exp = (rounded_bits >> 23) & 0xFF
273
+ scale_exp = max(biased_exp - 2, 0)
274
+
275
+ This ensures max_val / scale <= 6.0 for most inputs.
276
+ """
277
+ import struct
278
+
279
+ if block_max <= 0:
280
+ return 0
281
+ # Replicate the kernel's rounding logic in Python
282
+ float_bytes = struct.pack("f", block_max)
283
+ max_bits = struct.unpack("I", float_bytes)[0]
284
+ rounded_bits = (max_bits + (1 << 21)) & 0xFF800000
285
+ biased_exp = (rounded_bits >> 23) & 0xFF
286
+ scale_exp = max(int(biased_exp) - 2, 0)
287
+ scale_exp = min(scale_exp, 254)
288
+ return scale_exp
289
+
290
+
291
+@pytest.mark.skipif(
292
+ not is_sm100_supported(),
293
+ reason="mxfp4_experts_quant requires CUDA SM100",
294
+)
295
+@pytest.mark.parametrize("k", [256, 7168])
296
+@pytest.mark.parametrize("m", [16, 64])
297
+def test_mxfp4_experts_quant_e8m0_scale_correctness(m, k):
298
+ """
299
+ Test that mxfp4_experts_quant computes E8M0 block scales correctly
300
+ per OCP MX spec (not the NVFP4 formula).
301
+
302
+ The old buggy kernel used: floor(log2(max/6)) + 127
303
+ The fixed kernel uses: round_nearest_exp(max) - 2
304
+
305
+ This test verifies:
306
+ 1. Scales match the expected OCP MX formula for all blocks
307
+ 2. No block max exceeds the representable range (no unexpected saturation)
308
+ 3. Reconstruction error is within expected bounds for MXFP4
309
+ """
310
+ device = "cuda"
311
+
312
+ # Generate input with controlled range
313
+ input_tensor = torch.randn(m, k, device=device, dtype=torch.bfloat16) * 0.5
314
+
315
+ # Quantize
316
+ num_experts = 1
317
+ expert_offsets = torch.tensor([0, m], device=device, dtype=torch.int32)
318
+ num_k_tiles = (k // MXFP4_BLOCK_SIZE + 3) // 4
319
+ blockscale_offsets = torch.tensor(
320
+ [0, align(m, 128) * num_k_tiles], device=device, dtype=torch.int32
321
+ )
322
+
323
+ output_fp4, output_sf = ops.mxfp4_experts_quant(
324
+ input_tensor, expert_offsets, blockscale_offsets, num_experts, topk=1
325
+ )
326
+
327
+ # Untile scale to flat layout for verification
328
+ scale_flat = untile_cutlass_scale(output_sf, m, k)
329
+ assert scale_flat.shape == (m, k // MXFP4_BLOCK_SIZE)
330
+
331
+ # Verify each block's scale matches the OCP MX spec formula
332
+ num_blocks = k // MXFP4_BLOCK_SIZE
333
+ mismatches = 0
334
+ buggy_pattern = 0 # count blocks where scale is 1-2 lower than expected
335
+
336
+ for row in range(m):
337
+ for blk in range(num_blocks):
338
+ block_start = blk * MXFP4_BLOCK_SIZE
339
+ block_end = block_start + MXFP4_BLOCK_SIZE
340
+ block_max = (
341
+ input_tensor[row, block_start:block_end].float().abs().max().item()
342
+ )
343
+
344
+ actual_scale = scale_flat[row, blk].item()
345
+ expected_scale = compute_reference_e8m0_scale(block_max)
346
+
347
+ if actual_scale != expected_scale:
348
+ mismatches += 1
349
+ if actual_scale < expected_scale:
350
+ buggy_pattern += 1
351
+
352
+ total_blocks = m * num_blocks
353
+ match_rate = (total_blocks - mismatches) / total_blocks
354
+
355
+ print(
356
+ f" m={m}, k={k}: scale match rate = {match_rate * 100:.2f}% "
357
+ f"({mismatches}/{total_blocks} mismatches)"
358
+ )
359
+
360
+ # The fixed kernel should match the reference formula exactly
361
+ assert match_rate > 0.99, (
362
+ f"E8M0 scale match rate too low: {match_rate * 100:.2f}%. "
363
+ f"Buggy pattern (scale too low): {buggy_pattern}/{mismatches}. "
364
+ f"This suggests the NVFP4 formula bug is present."
365
+ )
366
+
367
+ # Extra check: if most mismatches show scale < expected, it's the old bug
368
+ if mismatches > 0:
369
+ assert buggy_pattern / mismatches < 0.5, (
370
+ f"Most scale mismatches show scale too LOW ({buggy_pattern}/{mismatches}). "
371
+ "This is the signature of the NVFP4 formula bug in nvfp4_utils.cuh."
372
+ )
373
+
374
+ # Verify reconstruction error is within MXFP4 expected bounds
375
+ # Dequantize and check cosine similarity
376
+ fp4_lut = torch.tensor(
377
+ [0, 0.5, 1, 1.5, 2, 3, 4, 6, 0, -0.5, -1, -1.5, -2, -3, -4, -6],
378
+ device=device,
379
+ dtype=torch.float32,
380
+ )
381
+ lo = (output_fp4 & 0x0F).long()
382
+ hi = ((output_fp4 >> 4) & 0x0F).long()
383
+ unpacked = torch.stack([lo, hi], dim=-1).reshape(m, k)
384
+ fp4_vals = fp4_lut[unpacked]
385
+
386
+ scales_expanded = 2.0 ** (scale_flat.float() - 127.0)
387
+ scales_expanded = scales_expanded.unsqueeze(-1).expand(-1, -1, MXFP4_BLOCK_SIZE)
388
+ scales_expanded = scales_expanded.reshape(m, k)
389
+ recon = (fp4_vals * scales_expanded).bfloat16()
390
+
391
+ # Cosine similarity should be > 0.99 for well-behaved MXFP4 quantization
392
+ cos_sim = torch.nn.functional.cosine_similarity(
393
+ recon.float().flatten().unsqueeze(0),
394
+ input_tensor.float().flatten().unsqueeze(0),
395
+ ).item()
396
+ max_abs_diff = (recon.float() - input_tensor.float()).abs().max().item()
397
+
398
+ print(
399
+ f" Reconstruction: cosine_sim={cos_sim:.6f}, max_abs_diff={max_abs_diff:.4f}"
400
+ )
401
+
402
+ assert cos_sim > 0.99, (
403
+ f"Reconstruction cosine similarity too low: {cos_sim:.6f}. "
404
+ f"Expected > 0.99 for correct MXFP4 quantization."
405
+ )
406
+ # With correct E8M0, max abs diff should be bounded by scale * 6
407
+ # (worst case: value just below threshold rounds to wrong FP4 code)
408
+ assert max_abs_diff < 1.0, (
409
+ f"Max reconstruction error too large: {max_abs_diff:.4f}. "
410
+ "Likely caused by incorrect E8M0 scale (values saturating to ±6)."
411
+ )
412
+
413
+
414
+@pytest.mark.skipif(
415
+ not is_sm100_supported(),
416
+ reason="mxfp4_experts_quant requires CUDA SM100",
417
+)
418
+def test_mxfp4_experts_quant_no_saturation():
419
+ """
420
+ Test that the E8M0 scale is large enough to avoid unexpected saturation.
421
+
422
+ With the buggy NVFP4 formula, the scale was too small causing most values
423
+ to saturate to ±6 in FP4. The fixed OCP MX formula should ensure that
424
+ block_max / scale <= 6.0 (the max E2M1 value) in almost all cases.
425
+ """
426
+ device = "cuda"
427
+
428
+ m, k = 128, 1024
429
+ # Use inputs with known range to make saturation detectable
430
+ input_tensor = torch.randn(m, k, device=device, dtype=torch.bfloat16) * 0.5
431
+
432
+ num_experts = 1
433
+ expert_offsets = torch.tensor([0, m], device=device, dtype=torch.int32)
434
+ num_k_tiles = (k // MXFP4_BLOCK_SIZE + 3) // 4
435
+ blockscale_offsets = torch.tensor(
436
+ [0, align(m, 128) * num_k_tiles], device=device, dtype=torch.int32
437
+ )
438
+
439
+ output_fp4, output_sf = ops.mxfp4_experts_quant(
440
+ input_tensor, expert_offsets, blockscale_offsets, num_experts, topk=1
441
+ )
442
+
443
+ # Check saturation rate: count FP4 values that are ±6 (codes 7 and 15)
444
+ lo = output_fp4 & 0x0F
445
+ hi = (output_fp4 >> 4) & 0x0F
446
+ # Code 7 = +6.0, code 15 = -6.0
447
+ saturated = ((lo == 7) | (lo == 15) | (hi == 7) | (hi == 15)).sum().item()
448
+ total_values = m * k
449
+ saturation_rate = saturated / total_values
450
+
451
+ print(
452
+ f" Saturation rate: {saturation_rate * 100:.2f}% "
453
+ f"({saturated}/{total_values} values at ±6)"
454
+ )
455
+
456
+ # For Gaussian input with std=0.5, saturation should be very rare
457
+ # (±6 * scale is far from the typical range).
458
+ # The buggy kernel had ~30-50% saturation; fixed should be < 5%.
459
+ assert saturation_rate < 0.05, (
460
+ f"FP4 saturation rate too high: {saturation_rate * 100:.2f}%. "
461
+ "This suggests the E8M0 scale is too small (NVFP4 formula bug). "
462
+ "Expected < 5% for Gaussian(0, 0.5) input with correct OCP MX scale."
463
+ )
464
+
465
+
466
if __name__ == "__main__":
467
pytest.main([__file__, "-v", "-s"])
vllm/model_executor/kernels/linear/mxfp4/flashinfer.py
+1
-1
@@ -56,7 +56,7 @@ class FlashInferMxFp4LinearKernel(MxFp4LinearKernel):
56
out_shape = x.shape[:-1] + (layer.output_size_per_partition,)
57
x_2d = x.reshape(-1, x.shape[-1])
58
59
- x_fp4, x_scale = flashinfer_mxfp4_quantize(x_2d)
59
+ x_fp4, x_scale = flashinfer_mxfp4_quantize(x_2d.contiguous())
60
out = flashinfer_scaled_fp4_mm(
61
x_fp4,
62
weight,