iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Bugfix][Quantization] Fix MXFP4 conversion for FlashInfer CUTLASS (#51038)

Browse Source 
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
This commit is contained in:
Gabriel Wu
2026-08-05 23:19:22 -07:00
committed by GitHub
co-authored by Yongye Zhu
parent c0202c5603
commit 2e35c529b8
2 changed files with 137 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/moe/test_ocp_mx_moe.py
+77
View File
@@ -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])
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+60 -2
View File
@@ -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."
)