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

tests/kernels/test_minimax_m3_amd_ops.py

@@ -424,6 +424,69 @@ def test_mxfp8_linear_emulation_bf16_at_load(
     assert _relerr(out.float(), out_ref.float()) < 2e-2
 
 
+# --------------------------------------------------------------------------- #
+# Native MXFP8 linear: the K % 128 != 0 weight is dequantized once at load
+# --------------------------------------------------------------------------- #
+@requires_gfx950
+# (5120, 576) and (5120, 288): DSv4.1 shared-expert down_proj at TP4 and TP8.
+@pytest.mark.parametrize("shape", [(5120, 576), (5120, 288), (256, 2048)])
+@torch.inference_mode()
+def test_mxfp8_rocm_native_unaligned_k_dequantizes_at_load(shape):
+    """``dot_scaled`` tiles K by 128, so an unaligned weight must not reach it.
+
+    ``process_weights_after_loading`` converts those to BF16 once; a weight that
+    never went through it has to fall back per step rather than take the
+    invalid path, which is what keying the dispatch on the weight dtype alone
+    would do.
+    """
+    from vllm.model_executor.kernels.linear.mxfp8 import rocm_native
+    from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
+        Mxfp8LinearLayerConfig,
+    )
+
+    N, K = shape
+    aligned = K % 128 == 0
+    torch.manual_seed(0)
+    w_bf16 = torch.randn(N, K, device=DEVICE, dtype=torch.bfloat16) * 0.1
+    w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False)
+    x = torch.randn(7, K, device=DEVICE, dtype=torch.bfloat16) * 0.5
+    out_ref = torch.nn.functional.linear(x, dequant_mxfp8_to_bf16(w_fp8, w_scale))
+    # dot_scaled re-quantizes the activation, the dequant paths do not.
+    tol = 5e-2 if aligned else 2e-3
+
+    def make_layer():
+        layer = torch.nn.Module()
+        layer.weight = torch.nn.Parameter(w_fp8.clone(), requires_grad=False)
+        layer.weight_scale = torch.nn.Parameter(w_scale.clone(), requires_grad=False)
+        return layer
+
+    kernel = rocm_native.RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig())
+    layer = make_layer()
+    kernel.process_weights_after_loading(layer)
+    assert (layer.weight.element_size() >= 2) is not aligned
+    assert _relerr(kernel.apply_weights(layer, x).float(), out_ref.float()) < tol
+
+    # An unaligned weight that is still FP8 must not reach dot_scaled.
+    unprocessed = make_layer()
+    unprocessed.weight_scale = torch.nn.Parameter(
+        w_scale[:N, : K // 32].contiguous(), requires_grad=False
+    )
+    calls = []
+    real = rocm_native._mxfp8_dot_scaled_linear
+
+    def _spy(*a, **k):
+        calls.append(1)
+        return real(*a, **k)
+
+    rocm_native._mxfp8_dot_scaled_linear = _spy
+    try:
+        out = kernel.apply_weights(unprocessed, x)
+    finally:
+        rocm_native._mxfp8_dot_scaled_linear = real
+    assert bool(calls) is aligned
+    assert _relerr(out.float(), out_ref.float()) < tol
+
+
 # ── EP expert_mask handling for the FlyDSL (AITER_MXFP8) MoE ────────────────
 # The map/mask choice lives in ``RoutedExperts.expert_map``: it hands AITER
 # experts (``consumes_expert_mask``) the precomputed 0/1 ``expert_mask`` and

vllm/model_executor/kernels/linear/mxfp8/rocm_native.py

@@ -4,9 +4,9 @@
 
 Consumes the FP8 E4M3 weights + E8M0 block scales directly (no dequant-to-BF16);
 activations are MXFP8-quantized per token. Uses the CDNA4 hardware microscaling
-matrix cores. Falls back (via the kernel selector) to the BF16
-``EmulationMxfp8LinearKernel`` on archs without native MX or for shapes with
-``K % 128 != 0``.
+matrix cores. ``dot_scaled`` tiles K by 128; a weight whose K is not a multiple
+of that is dequantized to BF16 once at load and served by a plain linear
+instead, since ``can_implement`` does not filter on K.
 """
 
 import torch
@@ -194,6 +194,9 @@ def _select_cfg(M, N, K):
     return (128, 128, 256, 8, 2) if K % 256 == 0 else (128, 256, 128, 8, 3)
 
 
+_DOT_SCALED_K_ALIGN = 128
+
+
 class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
     """Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``."""
 
@@ -219,6 +222,8 @@ class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
         N, K = weight.shape
         scale_k = K // MXFP8_BLOCK_SIZE
         weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
+        if K % _DOT_SCALED_K_ALIGN != 0:
+            weight = dequant_mxfp8_to_bf16(weight.contiguous(), weight_scale)
         layer.weight = Parameter(weight.contiguous(), requires_grad=False)
         layer.weight_scale = Parameter(weight_scale, requires_grad=False)
 
@@ -235,10 +240,13 @@ class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
             )
         out_shape = (*x.shape[:-1], layer.weight.shape[0])
         x2d = x.reshape(-1, x.shape[-1])
-        if x2d.shape[-1] % 128 == 0:
+        if layer.weight.element_size() >= 2:
+            out = torch.nn.functional.linear(x2d, layer.weight.to(x.dtype))
+        elif x2d.shape[-1] % _DOT_SCALED_K_ALIGN == 0:
             out = _mxfp8_dot_scaled_linear(x2d, layer.weight, layer.weight_scale)
         else:
-            # dot_scaled tiling needs K % 128 == 0; dequantize fallback otherwise.
+            # process_weights_after_loading dequantizes these weights once, so
+            # this runs only if it did not; dot_scaled would be invalid here.
             w_bf16 = dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale)
             out = torch.nn.functional.linear(x2d, w_bf16).to(x.dtype)
         out = out.reshape(out_shape)