[Bugfix][NVFP4 MoE] Pad gated intermediate to 64 for FlashInfer TRT-LLM shuffle (M%128) (#46880)
Signed-off-by: Mike G <180722391+mikekg@users.noreply.github.com>
Mike G committed
Jul 15, 2026 at 14:37 UTC
2bd8957627bfb5668c46f2bc359bef47d371270c
1 file changed
+7
-1
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+7
-1
@@ -353,7 +353,13 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
353
layer.moe_config.hidden_dim = padded_hidden
354
355
# Align weights for FI NVFP4 MoE kernels.
356
- min_alignment = 16 if is_gated else 128
356
+ # FlashInfer's TRT-LLM block-scale shuffle asserts the gate/up row dim
357
+ # (= up_mult * padded_intermediate, up_mult=2 when gated) is a multiple of
358
+ # 128. So gated needs padded_intermediate % 64 (2*64=128); the old value 16
359
+ # left 2*intermediate a multiple of only 32, so an NVFP4 MoE whose rank-local
360
+ # intermediate is not 128-aligned at TP>1 (e.g. Gemma-4-26B-A4B at tp4) hit
361
+ # `assert M % 128 == 0`. Padded rows are zero -> outputs unchanged.
362
+ min_alignment = 64 if is_gated else 128
363
w13, w13_scale, w2, w2_scale, padded_intermediate = (
364
align_fp4_moe_weights_for_fi(
365
w13, w13_scale, w2, w2_scale, is_act_and_mul, min_alignment