[ROCm][DSv4.1][Perf] Dequantize the MXFP8 weight once when dot_scaled cannot be used (#56560)
Signed-off-by: JohnQinAMD <yanyuan.qin@amd.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
JohnQinAMD committed
Sep 14, 2026 at 18:07 UTC
ef5f7cd119c60cdc75643beca232c159c9ad0930
2 files changed
+76
-5
tests/kernels/test_minimax_m3_amd_ops.py
+63
@@ -424,6 +424,69 @@ def test_mxfp8_linear_emulation_bf16_at_load(
424
assert _relerr(out.float(), out_ref.float()) < 2e-2
425
426
427
+# --------------------------------------------------------------------------- #
428
+# Native MXFP8 linear: the K % 128 != 0 weight is dequantized once at load
429
+# --------------------------------------------------------------------------- #
430
+@requires_gfx950
431
+# (5120, 576) and (5120, 288): DSv4.1 shared-expert down_proj at TP4 and TP8.
432
+@pytest.mark.parametrize("shape", [(5120, 576), (5120, 288), (256, 2048)])
433
+@torch.inference_mode()
434
+def test_mxfp8_rocm_native_unaligned_k_dequantizes_at_load(shape):
435
+ """``dot_scaled`` tiles K by 128, so an unaligned weight must not reach it.
436
+
437
+ ``process_weights_after_loading`` converts those to BF16 once; a weight that
438
+ never went through it has to fall back per step rather than take the
439
+ invalid path, which is what keying the dispatch on the weight dtype alone
440
+ would do.
441
+ """
442
+ from vllm.model_executor.kernels.linear.mxfp8 import rocm_native
443
+ from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
444
+ Mxfp8LinearLayerConfig,
445
+ )
446
+
447
+ N, K = shape
448
+ aligned = K % 128 == 0
449
+ torch.manual_seed(0)
450
+ w_bf16 = torch.randn(N, K, device=DEVICE, dtype=torch.bfloat16) * 0.1
451
+ w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False)
452
+ x = torch.randn(7, K, device=DEVICE, dtype=torch.bfloat16) * 0.5
453
+ out_ref = torch.nn.functional.linear(x, dequant_mxfp8_to_bf16(w_fp8, w_scale))
454
+ # dot_scaled re-quantizes the activation, the dequant paths do not.
455
+ tol = 5e-2 if aligned else 2e-3
456
+
457
+ def make_layer():
458
+ layer = torch.nn.Module()
459
+ layer.weight = torch.nn.Parameter(w_fp8.clone(), requires_grad=False)
460
+ layer.weight_scale = torch.nn.Parameter(w_scale.clone(), requires_grad=False)
461
+ return layer
462
+
463
+ kernel = rocm_native.RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig())
464
+ layer = make_layer()
465
+ kernel.process_weights_after_loading(layer)
466
+ assert (layer.weight.element_size() >= 2) is not aligned
467
+ assert _relerr(kernel.apply_weights(layer, x).float(), out_ref.float()) < tol
468
+
469
+ # An unaligned weight that is still FP8 must not reach dot_scaled.
470
+ unprocessed = make_layer()
471
+ unprocessed.weight_scale = torch.nn.Parameter(
472
+ w_scale[:N, : K // 32].contiguous(), requires_grad=False
473
+ )
474
+ calls = []
475
+ real = rocm_native._mxfp8_dot_scaled_linear
476
+
477
+ def _spy(*a, **k):
478
+ calls.append(1)
479
+ return real(*a, **k)
480
+
481
+ rocm_native._mxfp8_dot_scaled_linear = _spy
482
+ try:
483
+ out = kernel.apply_weights(unprocessed, x)
484
+ finally:
485
+ rocm_native._mxfp8_dot_scaled_linear = real
486
+ assert bool(calls) is aligned
487
+ assert _relerr(out.float(), out_ref.float()) < tol
488
+
489
+
490
# ── EP expert_mask handling for the FlyDSL (AITER_MXFP8) MoE ────────────────
491
# The map/mask choice lives in ``RoutedExperts.expert_map``: it hands AITER
492
# experts (``consumes_expert_mask``) the precomputed 0/1 ``expert_mask`` and
vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
+13
-5
@@ -4,9 +4,9 @@
4
5
Consumes the FP8 E4M3 weights + E8M0 block scales directly (no dequant-to-BF16);
6
activations are MXFP8-quantized per token. Uses the CDNA4 hardware microscaling
7
-matrix cores. Falls back (via the kernel selector) to the BF16
8
-``EmulationMxfp8LinearKernel`` on archs without native MX or for shapes with
9
-``K % 128 != 0``.
7
+matrix cores. ``dot_scaled`` tiles K by 128; a weight whose K is not a multiple
8
+of that is dequantized to BF16 once at load and served by a plain linear
9
+instead, since ``can_implement`` does not filter on K.
10
"""
11
12
import torch
@@ -194,6 +194,9 @@ def _select_cfg(M, N, K):
194
return (128, 128, 256, 8, 2) if K % 256 == 0 else (128, 256, 128, 8, 3)
195
196
197
+_DOT_SCALED_K_ALIGN = 128
198
+
199
+
200
class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
201
"""Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``."""
202
@@ -219,6 +222,8 @@ class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
222
N, K = weight.shape
223
scale_k = K // MXFP8_BLOCK_SIZE
224
weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
225
+ if K % _DOT_SCALED_K_ALIGN != 0:
226
+ weight = dequant_mxfp8_to_bf16(weight.contiguous(), weight_scale)
227
layer.weight = Parameter(weight.contiguous(), requires_grad=False)
228
layer.weight_scale = Parameter(weight_scale, requires_grad=False)
229
@@ -235,10 +240,13 @@ class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
240
)
241
out_shape = (*x.shape[:-1], layer.weight.shape[0])
242
x2d = x.reshape(-1, x.shape[-1])
238
- if x2d.shape[-1] % 128 == 0:
243
+ if layer.weight.element_size() >= 2:
244
+ out = torch.nn.functional.linear(x2d, layer.weight.to(x.dtype))
245
+ elif x2d.shape[-1] % _DOT_SCALED_K_ALIGN == 0:
246
out = _mxfp8_dot_scaled_linear(x2d, layer.weight, layer.weight_scale)
247
else:
241
- # dot_scaled tiling needs K % 128 == 0; dequantize fallback otherwise.
248
+ # process_weights_after_loading dequantizes these weights once, so
249
+ # this runs only if it did not; dot_scaled would be invalid here.
250
w_bf16 = dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale)
251
out = torch.nn.functional.linear(x2d, w_bf16).to(x.dtype)
252
out = out.reshape(out_shape)