[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 <zyy1102000@gmail.com>
Gabriel Wu committed
Aug 6, 2026 at 14:19 UTC
2e35c529b8decf545b2b037569e5b31640e093dd
2 files changed
+137
-2
tests/kernels/moe/test_ocp_mx_moe.py
+77
@@ -1014,6 +1014,83 @@ def test_flashinfer_cutlass_mxfp4_mxfp8_fused_moe(
1014
check_accuracy(ref, out, atol=0, rtol=0.3, percent=0.8)
1015
1016
1017
+@pytest.mark.parametrize(
1018
+ ("backend_name", "interleaves_weights"),
1019
+ [
1020
+ ("FLASHINFER_CUTLASS_MXFP4_MXFP8", False),
1021
+ ("FLASHINFER_CUTLASS_MXFP4_BF16", True),
1022
+ ],
1023
+)
1024
+@pytest.mark.skipif(not has_flashinfer(), reason="flashinfer is required")
1025
+def test_convert_standard_mxfp4_weights_for_flashinfer_cutlass(
1026
+ monkeypatch, backend_name, interleaves_weights
1027
+):
1028
+ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
1029
+ Mxfp4MoeBackend,
1030
+ convert_weight_to_mxfp4_moe_kernel_format,
1031
+ )
1032
+
1033
+ w13 = torch.arange(16, dtype=torch.uint8).reshape(1, 4, 4)
1034
+ w2 = torch.arange(8, dtype=torch.uint8).reshape(1, 4, 2)
1035
+ w13_scale = torch.arange(8, dtype=torch.uint8).reshape(1, 4, 2)
1036
+ w2_scale = torch.arange(4, dtype=torch.uint8).reshape(1, 4, 1)
1037
+ w13_bias = torch.arange(4, dtype=torch.float32).reshape(1, 4)
1038
+ w2_bias = torch.arange(4, dtype=torch.float32).reshape(1, 4)
1039
+ interleaved_weights = []
1040
+ interleaved_scales = []
1041
+
1042
+ def record_weight_interleave(weight, quant_type):
1043
+ assert quant_type == "fp4"
1044
+ interleaved_weights.append(weight.clone())
1045
+ return weight
1046
+
1047
+ def record_scale_interleave(scale):
1048
+ interleaved_scales.append(scale.clone())
1049
+ return scale
1050
+
1051
+ monkeypatch.setattr("flashinfer.block_scale_interleave", record_scale_interleave)
1052
+ monkeypatch.setattr(
1053
+ "flashinfer.fused_moe.interleave_moe_weights_for_sm90_mixed_gemm",
1054
+ record_weight_interleave,
1055
+ )
1056
+ monkeypatch.setattr(
1057
+ "flashinfer.fused_moe.interleave_moe_scales_for_sm90_mixed_gemm",
1058
+ record_scale_interleave,
1059
+ )
1060
+ converted = convert_weight_to_mxfp4_moe_kernel_format(
1061
+ getattr(Mxfp4MoeBackend, backend_name),
1062
+ types.SimpleNamespace(),
1063
+ w13,
1064
+ w2,
1065
+ w13_scale,
1066
+ w2_scale,
1067
+ w13_bias,
1068
+ w2_bias,
1069
+ )
1070
+
1071
+ expected_w13 = torch.cat([w13[:, 2:], w13[:, :2]], dim=1)
1072
+ expected_w13_scale = torch.cat([w13_scale[:, 2:], w13_scale[:, :2]], dim=1)
1073
+ torch.testing.assert_close(converted[0], expected_w13)
1074
+ torch.testing.assert_close(converted[1], w2)
1075
+ torch.testing.assert_close(converted[2], expected_w13_scale)
1076
+ torch.testing.assert_close(converted[3], w2_scale)
1077
+ expected_w13_bias = torch.cat([w13_bias[:, 2:], w13_bias[:, :2]], dim=1).to(
1078
+ torch.bfloat16
1079
+ )
1080
+ torch.testing.assert_close(converted[4], expected_w13_bias)
1081
+ torch.testing.assert_close(converted[5], w2_bias.to(torch.bfloat16))
1082
+
1083
+ if interleaves_weights:
1084
+ assert len(interleaved_weights) == 2
1085
+ torch.testing.assert_close(interleaved_weights[0], expected_w13)
1086
+ torch.testing.assert_close(interleaved_weights[1], w2)
1087
+ else:
1088
+ assert not interleaved_weights
1089
+ assert len(interleaved_scales) == 2
1090
+ torch.testing.assert_close(interleaved_scales[0], expected_w13_scale)
1091
+ torch.testing.assert_close(interleaved_scales[1], w2_scale)
1092
+
1093
+
1094
@pytest.mark.parametrize("topk", [1, 4])
1095
@pytest.mark.parametrize("num_experts", [32])
1096
@pytest.mark.parametrize("num_tokens", [1, 128])
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+60
-2
@@ -27,6 +27,9 @@ from vllm.model_executor.layers.fused_moe.config import (
27
mxfp4_w4a16_moe_quant_config,
28
ocp_mx_moe_quant_config,
29
)
30
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
31
+ swap_w13_to_w31,
32
+)
33
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import _swizzle_mxfp4
34
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
35
OCP_MX_BLOCK_SIZE,
@@ -1278,7 +1281,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1281
]:
1282
"""Convert loaded weights into backend-specific kernel format.
1283
1281
- Supports DeepGEMM, TRTLLM MXFP8, Triton and Marlin backends.
1284
+ Supports DeepGEMM, FlashInfer, TRTLLM MXFP8, Triton and Marlin backends.
1285
"""
1286
from vllm.platforms.rocm import on_gfx1250
1287
@@ -1569,6 +1572,60 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1572
w13_bias,
1573
w2_bias,
1574
)
1575
+ elif mxfp4_backend in (
1576
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
1577
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
1578
+ ):
1579
+ # Standard checkpoints store fused gate/up tensors as [w1; w3], while
1580
+ # FlashInfer CUTLASS consumes [w3; w1]. Keep weights, scales, and bias
1581
+ # in the same order before applying the backend-specific interleave.
1582
+ w13_weight = swap_w13_to_w31(w13_weight.data)
1583
+ w13_weight_scale = swap_w13_to_w31(w13_weight_scale.data)
1584
+ if w13_bias is not None:
1585
+ b1, b3 = torch.chunk(w13_bias.data, 2, dim=-1)
1586
+ w13_bias = torch.cat([b3, b1], dim=-1).to(torch.bfloat16)
1587
+ if w2_bias is not None:
1588
+ w2_bias = w2_bias.data.to(torch.bfloat16)
1589
+
1590
+ if mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8:
1591
+ from flashinfer import block_scale_interleave
1592
+
1593
+ w13_scale_shape = w13_weight_scale.shape
1594
+ w13_weight_scale = block_scale_interleave(
1595
+ w13_weight_scale.view(torch.uint8)
1596
+ ).reshape(w13_scale_shape)
1597
+
1598
+ w2_scale_shape = w2_weight_scale.shape
1599
+ w2_weight_scale = block_scale_interleave(
1600
+ w2_weight_scale.data.view(torch.uint8)
1601
+ ).reshape(w2_scale_shape)
1602
+ else:
1603
+ from flashinfer.fused_moe import (
1604
+ interleave_moe_scales_for_sm90_mixed_gemm,
1605
+ interleave_moe_weights_for_sm90_mixed_gemm,
1606
+ )
1607
+
1608
+ w13_weight = interleave_moe_weights_for_sm90_mixed_gemm(
1609
+ w13_weight.contiguous(), "fp4"
1610
+ )
1611
+ w2_weight = interleave_moe_weights_for_sm90_mixed_gemm(
1612
+ w2_weight.data.contiguous(), "fp4"
1613
+ )
1614
+ w13_weight_scale = interleave_moe_scales_for_sm90_mixed_gemm(
1615
+ w13_weight_scale.to(torch.uint8)
1616
+ )
1617
+ w2_weight_scale = interleave_moe_scales_for_sm90_mixed_gemm(
1618
+ w2_weight_scale.data.to(torch.uint8)
1619
+ )
1620
+
1621
+ return (
1622
+ w13_weight,
1623
+ w2_weight,
1624
+ w13_weight_scale,
1625
+ w2_weight_scale,
1626
+ w13_bias,
1627
+ w2_bias,
1628
+ )
1629
elif mxfp4_backend in (
1630
Mxfp4MoeBackend.XPU,
1631
Mxfp4MoeBackend.EMULATION,
@@ -1586,7 +1643,8 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1643
else:
1644
raise ValueError(
1645
f"Unsupported mxfp4_backend for Mxfp4MoEMethod: {mxfp4_backend}. "
1589
- f"Expected TRTLLM, Triton, AITER, XPU, or emulation backend."
1646
+ "Expected TRTLLM, FlashInfer CUTLASS, Triton, AITER, XPU, or "
1647
+ "emulation backend."
1648
)
1649
1650