[Bugfix][Test] Fix test_flashinfer_cutlass_mxfp4_fused_moe on sm90 (stale weight/scale interleave) (#46915)
Signed-off-by: wentian-byte <2990624738@qq.com>
xiaolinchen committed
Jun 28, 2026 at 02:30 UTC
35e3850fa9499b99f0b32ee8e9d5551a290d9c54
1 file changed
+14
-15
tests/kernels/moe/test_ocp_mx_moe.py
+14
-15
@@ -659,19 +659,6 @@ def test_trtllm_gen_mxfp4_fused_moe(
659
check_accuracy(ref_result, tg_result, atol=0, rtol=0.3, percent=0.8)
660
661
662
-def _interleave_scales_lastdim_by4(scales: torch.Tensor) -> torch.Tensor:
663
- """Interleave scales on the last dimension by groups of 4, matching
664
- the transformation in mxfp4.py's BF16 (Hopper) path."""
665
- s = scales.to(torch.uint8)
666
- s_shape = s.shape
667
- assert s_shape[-1] % 4 == 0
668
- s = s.reshape(*s_shape[:-1], s_shape[-1] // 4, 4)
669
- # Move the 4-group dimension before the row dimension
670
- permuted = s.permute(0, 2, 1, 3)
671
- # Merge the row dim with the 4-group dim
672
- return permuted.reshape(s_shape[0], s_shape[-1] // 4, s_shape[1] * 4)
673
-
674
-
662
@pytest.mark.parametrize("topk", [1, 4])
663
@pytest.mark.parametrize("num_experts", [32])
664
@pytest.mark.parametrize("num_tokens", [1, 128])
@@ -771,13 +758,25 @@ def test_flashinfer_cutlass_mxfp4_fused_moe(
758
w1_w, w3_w = torch.chunk(w13_q, 2, dim=1)
759
w13_q_swapped = torch.cat([w3_w, w1_w], dim=1)
760
761
+ # SM90 mixed-input GEMM expects weights/scales in an interleaved layout;
762
+ # without it the FP4->BF16 LUT reads bytes from wrong positions for K>128.
763
+ from flashinfer.fused_moe import (
764
+ interleave_moe_scales_for_sm90_mixed_gemm,
765
+ interleave_moe_weights_for_sm90_mixed_gemm,
766
+ )
767
+
768
+ w13_q_swapped = interleave_moe_weights_for_sm90_mixed_gemm(
769
+ w13_q_swapped, quant_type="fp4"
770
+ )
771
+ w2_q = interleave_moe_weights_for_sm90_mixed_gemm(w2_q, quant_type="fp4")
772
+
773
b1, b3 = torch.chunk(bias13.to(torch.float32), 2, dim=-1)
774
w13_b = torch.cat([b3, b1], dim=-1).to(torch.bfloat16)
775
776
w1_s, w3_s = torch.chunk(w13_scale, 2, dim=1)
777
w13_s = torch.cat([w3_s, w1_s], dim=1)
779
- w13_s_inter = _interleave_scales_lastdim_by4(w13_s)
780
- w2_s_inter = _interleave_scales_lastdim_by4(w2_scale)
778
+ w13_s_inter = interleave_moe_scales_for_sm90_mixed_gemm(w13_s)
779
+ w2_s_inter = interleave_moe_scales_for_sm90_mixed_gemm(w2_scale)
780
781
routing_weights = torch.nn.functional.softmax(
782
router_logits, dim=1, dtype=torch.float32