From 2e35c529b8decf545b2b037569e5b31640e093dd Mon Sep 17 00:00:00 2001 From: Gabriel Wu <13583761+lucifer1004@users.noreply.github.com> Date: Thu, 6 Aug 2026 14:19:22 +0800 Subject: [PATCH] [Bugfix][Quantization] Fix MXFP4 conversion for FlashInfer CUTLASS (#51038) Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Co-authored-by: Yongye Zhu --- tests/kernels/moe/test_ocp_mx_moe.py | 77 +++++++++++++++++++ .../layers/fused_moe/oracle/mxfp4.py | 62 ++++++++++++++- 2 files changed, 137 insertions(+), 2 deletions(-) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 9f4d3370d5..18c70093a1 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -1014,6 +1014,83 @@ def test_flashinfer_cutlass_mxfp4_mxfp8_fused_moe( check_accuracy(ref, out, atol=0, rtol=0.3, percent=0.8) +@pytest.mark.parametrize( + ("backend_name", "interleaves_weights"), + [ + ("FLASHINFER_CUTLASS_MXFP4_MXFP8", False), + ("FLASHINFER_CUTLASS_MXFP4_BF16", True), + ], +) +@pytest.mark.skipif(not has_flashinfer(), reason="flashinfer is required") +def test_convert_standard_mxfp4_weights_for_flashinfer_cutlass( + monkeypatch, backend_name, interleaves_weights +): + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + Mxfp4MoeBackend, + convert_weight_to_mxfp4_moe_kernel_format, + ) + + w13 = torch.arange(16, dtype=torch.uint8).reshape(1, 4, 4) + w2 = torch.arange(8, dtype=torch.uint8).reshape(1, 4, 2) + w13_scale = torch.arange(8, dtype=torch.uint8).reshape(1, 4, 2) + w2_scale = torch.arange(4, dtype=torch.uint8).reshape(1, 4, 1) + w13_bias = torch.arange(4, dtype=torch.float32).reshape(1, 4) + w2_bias = torch.arange(4, dtype=torch.float32).reshape(1, 4) + interleaved_weights = [] + interleaved_scales = [] + + def record_weight_interleave(weight, quant_type): + assert quant_type == "fp4" + interleaved_weights.append(weight.clone()) + return weight + + def record_scale_interleave(scale): + interleaved_scales.append(scale.clone()) + return scale + + monkeypatch.setattr("flashinfer.block_scale_interleave", record_scale_interleave) + monkeypatch.setattr( + "flashinfer.fused_moe.interleave_moe_weights_for_sm90_mixed_gemm", + record_weight_interleave, + ) + monkeypatch.setattr( + "flashinfer.fused_moe.interleave_moe_scales_for_sm90_mixed_gemm", + record_scale_interleave, + ) + converted = convert_weight_to_mxfp4_moe_kernel_format( + getattr(Mxfp4MoeBackend, backend_name), + types.SimpleNamespace(), + w13, + w2, + w13_scale, + w2_scale, + w13_bias, + w2_bias, + ) + + expected_w13 = torch.cat([w13[:, 2:], w13[:, :2]], dim=1) + expected_w13_scale = torch.cat([w13_scale[:, 2:], w13_scale[:, :2]], dim=1) + torch.testing.assert_close(converted[0], expected_w13) + torch.testing.assert_close(converted[1], w2) + torch.testing.assert_close(converted[2], expected_w13_scale) + torch.testing.assert_close(converted[3], w2_scale) + expected_w13_bias = torch.cat([w13_bias[:, 2:], w13_bias[:, :2]], dim=1).to( + torch.bfloat16 + ) + torch.testing.assert_close(converted[4], expected_w13_bias) + torch.testing.assert_close(converted[5], w2_bias.to(torch.bfloat16)) + + if interleaves_weights: + assert len(interleaved_weights) == 2 + torch.testing.assert_close(interleaved_weights[0], expected_w13) + torch.testing.assert_close(interleaved_weights[1], w2) + else: + assert not interleaved_weights + assert len(interleaved_scales) == 2 + torch.testing.assert_close(interleaved_scales[0], expected_w13_scale) + torch.testing.assert_close(interleaved_scales[1], w2_scale) + + @pytest.mark.parametrize("topk", [1, 4]) @pytest.mark.parametrize("num_experts", [32]) @pytest.mark.parametrize("num_tokens", [1, 128]) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index e8371e4ab4..0d8dc570bd 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -27,6 +27,9 @@ from vllm.model_executor.layers.fused_moe.config import ( mxfp4_w4a16_moe_quant_config, ocp_mx_moe_quant_config, ) +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + swap_w13_to_w31, +) from vllm.model_executor.layers.quantization.utils.mxfp4_utils import _swizzle_mxfp4 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( OCP_MX_BLOCK_SIZE, @@ -1278,7 +1281,7 @@ def convert_weight_to_mxfp4_moe_kernel_format( ]: """Convert loaded weights into backend-specific kernel format. - Supports DeepGEMM, TRTLLM MXFP8, Triton and Marlin backends. + Supports DeepGEMM, FlashInfer, TRTLLM MXFP8, Triton and Marlin backends. """ from vllm.platforms.rocm import on_gfx1250 @@ -1569,6 +1572,60 @@ def convert_weight_to_mxfp4_moe_kernel_format( w13_bias, w2_bias, ) + elif mxfp4_backend in ( + Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, + Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, + ): + # Standard checkpoints store fused gate/up tensors as [w1; w3], while + # FlashInfer CUTLASS consumes [w3; w1]. Keep weights, scales, and bias + # in the same order before applying the backend-specific interleave. + w13_weight = swap_w13_to_w31(w13_weight.data) + w13_weight_scale = swap_w13_to_w31(w13_weight_scale.data) + if w13_bias is not None: + b1, b3 = torch.chunk(w13_bias.data, 2, dim=-1) + w13_bias = torch.cat([b3, b1], dim=-1).to(torch.bfloat16) + if w2_bias is not None: + w2_bias = w2_bias.data.to(torch.bfloat16) + + if mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8: + from flashinfer import block_scale_interleave + + w13_scale_shape = w13_weight_scale.shape + w13_weight_scale = block_scale_interleave( + w13_weight_scale.view(torch.uint8) + ).reshape(w13_scale_shape) + + w2_scale_shape = w2_weight_scale.shape + w2_weight_scale = block_scale_interleave( + w2_weight_scale.data.view(torch.uint8) + ).reshape(w2_scale_shape) + else: + from flashinfer.fused_moe import ( + interleave_moe_scales_for_sm90_mixed_gemm, + interleave_moe_weights_for_sm90_mixed_gemm, + ) + + w13_weight = interleave_moe_weights_for_sm90_mixed_gemm( + w13_weight.contiguous(), "fp4" + ) + w2_weight = interleave_moe_weights_for_sm90_mixed_gemm( + w2_weight.data.contiguous(), "fp4" + ) + w13_weight_scale = interleave_moe_scales_for_sm90_mixed_gemm( + w13_weight_scale.to(torch.uint8) + ) + w2_weight_scale = interleave_moe_scales_for_sm90_mixed_gemm( + w2_weight_scale.data.to(torch.uint8) + ) + + return ( + w13_weight, + w2_weight, + w13_weight_scale, + w2_weight_scale, + w13_bias, + w2_bias, + ) elif mxfp4_backend in ( Mxfp4MoeBackend.XPU, Mxfp4MoeBackend.EMULATION, @@ -1586,7 +1643,8 @@ def convert_weight_to_mxfp4_moe_kernel_format( else: raise ValueError( f"Unsupported mxfp4_backend for Mxfp4MoEMethod: {mxfp4_backend}. " - f"Expected TRTLLM, Triton, AITER, XPU, or emulation backend." + "Expected TRTLLM, FlashInfer CUTLASS, Triton, AITER, XPU, or " + "emulation backend." )