diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py index 182672597a..ec8f1fd4ea 100644 --- a/tests/kernels/test_minimax_m3_amd_ops.py +++ b/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 diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py index 9714f92dff..29739afae3 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/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)