A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 02:30 UTC
[Bugfix] Fix FlashInfer CUTLASS MXFP4-MXFP8 MoE by restoring swizzled scale (#42089)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
May 8, 2026 at 18:59 UTC
8bcd8a260cd19cc534f1773ba241c3a718764452
12 files changed
+30
-27
tests/kernels/moe/test_trtllm_nvfp4_moe.py
+1
-1
@@ -111,7 +111,7 @@ def test_trtllm_fp4_moe_no_graph(
111
# The TRT-LLM FP4 MoE kernel rejects swizzled (padded) activation
112
# scales — its numel-based vec_size check requires numel == M*K/16.
113
# Match what oracle/nvfp4.py does for this backend.
114
- is_nvfp4_scale_swizzled=False,
114
+ is_scale_swizzled=False,
115
)
116
117
score = torch.randn((m, e), device="cuda", dtype=dtype)
tests/kernels/moe/utils.py
+2
-2
@@ -402,7 +402,7 @@ def make_test_quant_config(
402
per_act_token_quant: bool = False,
403
block_shape: list[int] | None = None,
404
make_gate: bool = True,
405
- is_nvfp4_scale_swizzled: bool = True,
405
+ is_scale_swizzled: bool = True,
406
) -> tuple[torch.Tensor, torch.Tensor, FusedMoEQuantConfig]:
407
(_, w1, w1_s, w1_gs), (_, w2, w2_s, w2_gs) = make_test_weights(
408
e,
@@ -443,7 +443,7 @@ def make_test_quant_config(
443
# TODO: make sure this is handled properly
444
g1_alphas=(1 / w1_gs) if w1_gs is not None else None,
445
g2_alphas=(1 / w2_gs) if w2_gs is not None else None,
446
- is_nvfp4_scale_swizzled=is_nvfp4_scale_swizzled,
446
+ is_scale_swizzled=is_scale_swizzled,
447
),
448
)
449
vllm/model_executor/layers/fused_moe/config.py
+10
-6
@@ -245,7 +245,7 @@ class FusedMoEQuantConfig:
245
_a2: FusedMoEQuantDesc
246
_w1: FusedMoEQuantDesc
247
_w2: FusedMoEQuantDesc
248
- is_nvfp4_scale_swizzled: bool = True
248
+ is_scale_swizzled: bool = True
249
250
# MXFP4-specific TRTLLM parameters for SwiGLU activation clamping.
251
# These correspond to gemm1_alpha, gemm1_beta, gemm1_clamp_limit
@@ -502,7 +502,7 @@ class FusedMoEQuantConfig:
502
w1_zp: torch.Tensor | None = None,
503
w2_zp: torch.Tensor | None = None,
504
weight_dtype: torch.dtype | str | None = None,
505
- is_nvfp4_scale_swizzled: bool = True,
505
+ is_scale_swizzled: bool = True,
506
gemm1_alpha: float | None = None,
507
gemm1_beta: float | None = None,
508
gemm1_clamp_limit: float | None = None,
@@ -535,7 +535,9 @@ class FusedMoEQuantConfig:
535
- w2_bias: Optional biases for w1 (GPT OSS Triton).
536
- w1_zp: Optional w1 zero points for int4/int8 quantization.
537
- w2_zp: Optional w2 zero points for int4/int8 quantization.
538
- - is_nvfp4_scale_swizzled: Whether to swizzle the nvfp4 scale swizzling.
538
+ - is_scale_swizzled: Whether the activation scale-factor layout is
539
+ swizzled. Pass through to the underlying quantization kernel for
540
+ dtypes that distinguish layouts (nvfp4, mxfp8). Defaults to True.
541
- gemm1_alpha: Optional MXFP4 TRTLLM SwiGLU alpha parameter.
542
- gemm1_beta: Optional MXFP4 TRTLLM SwiGLU beta parameter.
543
- gemm1_clamp_limit: Optional MXFP4 TRTLLM SwiGLU clamp limit.
@@ -571,7 +573,7 @@ class FusedMoEQuantConfig:
573
_w2=FusedMoEQuantDesc(
574
weight_dtype, w_shape, w2_scale, g2_alphas, w2_zp, w2_bias
575
),
574
- is_nvfp4_scale_swizzled=is_nvfp4_scale_swizzled,
576
+ is_scale_swizzled=is_scale_swizzled,
577
gemm1_alpha=gemm1_alpha,
578
gemm1_beta=gemm1_beta,
579
gemm1_clamp_limit=gemm1_clamp_limit,
@@ -715,6 +717,7 @@ def mxfp4_mxfp8_moe_quant_config(
717
gemm1_beta: float | None = None,
718
gemm1_clamp_limit: float | None = None,
719
mx_alignment: int = 0,
720
+ is_scale_swizzled: bool = True,
721
) -> FusedMoEQuantConfig:
722
"""
723
Construct a quant config for mxfp4 activations and mxfp4 weights.
@@ -728,6 +731,7 @@ def mxfp4_mxfp8_moe_quant_config(
731
gemm1_beta=gemm1_beta,
732
gemm1_clamp_limit=gemm1_clamp_limit,
733
mx_alignment=mx_alignment,
734
+ is_scale_swizzled=is_scale_swizzled,
735
)
736
737
@@ -796,7 +800,7 @@ def nvfp4_moe_quant_config(
800
w2_scale: torch.Tensor,
801
w1_bias: torch.Tensor | None = None,
802
w2_bias: torch.Tensor | None = None,
799
- is_nvfp4_scale_swizzled: bool = True,
803
+ is_scale_swizzled: bool = True,
804
) -> FusedMoEQuantConfig:
805
"""
806
Construct a quant config for mxfp4 activations and nvp4 weights.
@@ -814,7 +818,7 @@ def nvfp4_moe_quant_config(
818
per_act_token_quant=False,
819
per_out_ch_quant=False,
820
block_shape=None,
817
- is_nvfp4_scale_swizzled=is_nvfp4_scale_swizzled,
821
+ is_scale_swizzled=is_scale_swizzled,
822
)
823
824
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+1
-1
@@ -569,7 +569,7 @@ def make_fp8_moe_quant_config(
569
a1_scale=a1_scale,
570
a2_scale=a2_scale,
571
block_shape=block_shape,
572
- is_nvfp4_scale_swizzled=False,
572
+ is_scale_swizzled=False,
573
)
574
575
# All other backends use normal config.
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+4
@@ -1421,6 +1421,7 @@ def make_mxfp4_moe_quant_config(
1421
gemm1_clamp_limit=swiglu_limit,
1422
)
1423
elif mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8:
1424
+ # TRTLLM kernel expects non-swizzled mxfp8 activation scales.
1425
return mxfp4_mxfp8_moe_quant_config(
1426
w1_bias=w1_bias,
1427
w2_bias=w2_bias,
@@ -1430,8 +1431,10 @@ def make_mxfp4_moe_quant_config(
1431
gemm1_beta=gemm1_beta,
1432
gemm1_clamp_limit=swiglu_limit,
1433
mx_alignment=256,
1434
+ is_scale_swizzled=False,
1435
)
1436
elif mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8:
1437
+ # CUTLASS kernel expects swizzled mxfp8 activation scales.
1438
return mxfp4_mxfp8_moe_quant_config(
1439
w1_bias=w1_bias,
1440
w2_bias=w2_bias,
@@ -1440,6 +1443,7 @@ def make_mxfp4_moe_quant_config(
1443
gemm1_alpha=gemm1_alpha,
1444
gemm1_beta=gemm1_beta,
1445
gemm1_clamp_limit=swiglu_limit,
1446
+ is_scale_swizzled=True,
1447
)
1448
elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_FP8:
1449
# W4A8: MXFP4 weights + static FP8 activations
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+1
-1
@@ -461,7 +461,7 @@ def make_nvfp4_moe_quant_config(
461
# NOTE(rob): this is a hack until the MoE kernels
462
# create their own quant configs. TRTLLM kernel
463
# does not accept swizzled input quant scales.
464
- is_nvfp4_scale_swizzled=(
464
+ is_scale_swizzled=(
465
backend
466
not in (
467
NvFp4MoeBackend.FLASHINFER_TRTLLM,
vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py
+1
-1
@@ -241,7 +241,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
241
quant_dtype=quant_config.quant_dtype,
242
per_act_token_quant=False,
243
block_shape=quant_config.block_shape,
244
- is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled,
244
+ is_scale_swizzled=quant_config.is_scale_swizzled,
245
)
246
247
return (
vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py
+2
-5
@@ -106,7 +106,7 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
106
quant_config.quant_dtype,
107
quant_config.per_act_token_quant,
108
quant_config.block_shape,
109
- is_fp4_scale_swizzled=False, # delay swizzle to after comm
109
+ is_scale_swizzled=False, # delay swizzle to after comm
110
mx_alignment=quant_config.mx_alignment,
111
)
112
@@ -129,10 +129,7 @@ class FlashInferNVLinkOneSidedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeMo
129
if a1q_scale is not None:
130
a1q_recv, a1q_scale_recv, topk_ids_recv, topk_weights_recv = recv_payloads
131
# Apply scale interleaving only for CUTLASS (not TRT-LLM)
132
- if (
133
- quant_config.quant_dtype == "nvfp4"
134
- and quant_config.is_nvfp4_scale_swizzled
135
- ):
132
+ if quant_config.quant_dtype == "nvfp4" and quant_config.is_scale_swizzled:
133
a1q_scale_recv = a1q_scale_recv.view(-1, a1q_scale_recv.shape[-1])
134
a1q_scale_recv = a1q_scale_recv.view(torch.uint8)
135
a1q_scale_recv = nvfp4_block_scale_interleave(a1q_scale_recv)
vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_two_sided.py
+2
-2
@@ -173,7 +173,7 @@ def flashinfer_alltoall_dispatch(
173
# which makes the scales tensor different shape than
174
# the hidden states, breaking the A2A kernel. So, we
175
# delay the swizzling until after the A2A.
176
- is_fp4_scale_swizzled=False,
176
+ is_scale_swizzled=False,
177
mx_alignment=quant_config.mx_alignment,
178
)
179
@@ -194,7 +194,7 @@ def flashinfer_alltoall_dispatch(
194
)
195
196
# Swizzle after the A2A if MoE kernel expects swizzled scales.
197
- if quant_config.quant_dtype == "nvfp4" and quant_config.is_nvfp4_scale_swizzled:
197
+ if quant_config.quant_dtype == "nvfp4" and quant_config.is_scale_swizzled:
198
if x_sf.element_size() == 1:
199
x_sf = x_sf.view(torch.uint8)
200
x_sf = nvfp4_block_scale_interleave(x_sf)
vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py
+2
-2
@@ -39,7 +39,7 @@ def _quantize_and_setup_dispatch(
39
quant_dtype=quant_config.quant_dtype,
40
per_act_token_quant=quant_config.per_act_token_quant,
41
block_shape=quant_config.block_shape,
42
- is_fp4_scale_swizzled=False,
42
+ is_scale_swizzled=False,
43
mx_alignment=quant_config.mx_alignment,
44
)
45
@@ -59,7 +59,7 @@ def _unwrap_scale_and_prepare_for_moe(
59
assert scales is not None and len(scales) == 1
60
a1q_scale = scales[0]
61
# Apply swizzling after a2a if the MoE kernel needs it.
62
- if quant_config.quant_dtype == "nvfp4" and quant_config.is_nvfp4_scale_swizzled:
62
+ if quant_config.quant_dtype == "nvfp4" and quant_config.is_scale_swizzled:
63
assert a1q_scale is not None
64
if a1q_scale.element_size() == 1:
65
a1q_scale = a1q_scale.view(torch.uint8)
vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py
+1
-1
@@ -30,7 +30,7 @@ def _quantize_input(
30
quant_dtype=quant_config.quant_dtype,
31
per_act_token_quant=quant_config.per_act_token_quant,
32
block_shape=quant_config.block_shape,
33
- is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled,
33
+ is_scale_swizzled=quant_config.is_scale_swizzled,
34
mx_alignment=quant_config.mx_alignment,
35
)
36
vllm/model_executor/layers/fused_moe/utils.py
+3
-5
@@ -256,7 +256,7 @@ def moe_kernel_quantize_input(
256
quant_dtype: None | torch.dtype | str,
257
per_act_token_quant: bool,
258
block_shape: list[int] | None = None,
259
- is_fp4_scale_swizzled: bool = True,
259
+ is_scale_swizzled: bool = True,
260
ocp_mx_scheme: str | None = None,
261
quantization_emulation: bool = False,
262
mx_alignment: int = 0,
@@ -296,9 +296,7 @@ def moe_kernel_quantize_input(
296
return _int8_quantize(A, A_scale, per_act_token_quant, block_shape)
297
elif quant_dtype == "nvfp4":
298
if not quantization_emulation:
299
- return _nvfp4_quantize(
300
- A, A_scale, is_sf_swizzled_layout=is_fp4_scale_swizzled
301
- )
299
+ return _nvfp4_quantize(A, A_scale, is_sf_swizzled_layout=is_scale_swizzled)
300
else:
301
A = ref_nvfp4_quant_dequant(A, A_scale, block_size=16)
302
return A, None
@@ -322,7 +320,7 @@ def moe_kernel_quantize_input(
320
A_scale,
321
per_act_token_quant,
322
block_shape,
325
- is_sf_swizzled_layout=False,
323
+ is_sf_swizzled_layout=is_scale_swizzled,
324
mx_alignment=mx_alignment,
325
)
326
elif quant_dtype == "mxfp6_e3m2":