[Bugfix][NVFP4 MoE] Pad gated intermediate to 64 for FlashInfer TRT-LLM shuffle (M%128)

[mikekg]

Closes #46879.

Summary

FlashInfer's TRT-LLM NVFP4 MoE weight prep shuffles block-scale rows via
get_shuffle_matrix_sf_a_row_indices, which asserts the gate/up row dim is a
multiple of 128 (epilogue_tile_m). align_fp4_moe_weights_for_fi pads the
intermediate to min_alignment, and the gate/up dim is
up_multpadded_intermediate (up_mult=2 when gated). The caller passes
min_alignment = 16 if is_gated else 128; 16 is the NVFP4 scale-block size
(num_elts_per_sf=16), which aligns the quant scale blocks but is weaker than
the shuffle's 128-row tile. For gated, 2 * round_up(intermediate, 16) is only
guaranteed a multiple of 32, so a gated NVFP4 MoE whose rank-local
intermediate isn't 64-aligned at TP>1 (e.g. nvidia/Gemma-4-26B-A4B-NVFP4 at
-tp 4: 2(704/4)=352, 352%128=96) fails engine init with assert M % 128 == 0.

Why Padding to M=64 (and not the tile size requirement of M%128)

64 is not itself a shuffle tile. Gated activations compute SwiGLU(x) =
SiLU(x·W_gate) ⊗ (x·W_up), and vLLM fuses the two projections into a single
W13 = [W_gate | W_up] — one GEMM whose row dim is up_mult*padded_intermediate
with up_mult=2. The shuffle's 128-row epilogue_tile_m applies to that fused
dim, so the per-shard padded_intermediate only needs to be a multiple of 128
// up_mult = 64 for 2 * padded_intermediate to land on the tile — i.e.
min_alignment = max(16, 128 // up_mult), keeping the NVFP4 scale block (16) as
the floor. The non-gated path (up_mult=1) keeps the full 128.

Fix

min_alignment = 64 for gated, so 2 * padded_intermediate is a multiple of 128.
Padded rows are zero in both weights and scales, so outputs are unchanged.
Non-gated path (min_alignment=128) unchanged; Marlin path (#45295) unaffected.

Validation

• Failure (verified): vllm serve nvidia/Gemma-4-26B-A4B-NVFP4 -tp 4
  --quantization modelopt --trust-remote-code on B200 fails at
  process_weights_after_loading with assert M % 128 == 0 (FLASHINFER_TRTLLM
  backend confirmed selected in the log).
• Fix correctness (by construction): with min_alignment=64,
  padded_intermediate = round_up(176,64) = 192, so the gate/up row dim 2192 =
  384 = 3128 satisfies the assert.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[lucianommartins]

thanks for that, @mikekg.

confirming this diagnosis independently on current main via code analysis — both call sites still pass min_alignment = 16 if is_gated else 128, and the arithmetic holds for nvidia/Gemma-4-26B-A4B-NVFP4 at -tp 4:

• intermediate per partition = 704/4 = 176; gated up_mult = 2
• current: 2 * round_up(176, 16) = 352, 352 % 128 = 96 → assert M % 128 == 0 fails
• with this PR (min_alignment = 64): 2 * round_up(176, 64) = 384, 384 % 128 = 0 ✓

the requirement is up_mult * padded_intermediate ≡ 0 (mod 128), i.e. padded_intermediate ≡ 0 (mod 64) for gated; 64 is also a multiple of the NVFP4 scale block (num_elts_per_sf = 16), so 64 if is_gated else 128 is the minimal correct value. (A self-documenting alternative: min_alignment = max(16, 128 // up_mult).)

one thing worth checking while here: the same min_alignment = 16 if is_gated else 128 pattern also lives in prepare_fp8_moe_layer_for_fi → align_moe_weights_for_fi (flashinfer_utils.py, the non-block FP8 path). If the FP8 TRT-LLM shuffle carries the same 128-row epilogue-tile constraint on the gated dim, a gated FP8 MoE whose rank-local intermediate isn't 64-aligned (e.g. the same 26B-A4B at -tp 4) would hit the identical assert, and this PR wouldn't cover it. Do you have a B200 to check the FP8 path too, or should that be a follow-up?

@ywang96 @Isotr0py - could you please do a review here?

[mikekg]

Fail is an unrelated torch.compile() fusion regression on trunk identified in
tests/compile/fusions_e2e/test_tp2_ar_rms.py -k inductor_partition

[mikekg]

@lucianommartins :

If the FP8 TRT-LLM shuffle carries the same 128-row epilogue-tile constraint on the gated dim, a gated FP8 MoE whose rank-local intermediate isn't 64-aligned (e.g. the same 26B-A4B at -tp 4) would hit the identical assert, and this PR wouldn't cover it. Do you have a B200 to check the FP8 path too, or should that be a follow-up?

I'd suggest putting it with a follow-up that has its own repro. I don't have a model handy to repro for FP8 activations, and I think it's a hard sell to patch a complex code base without hard data. If you can rig up a repro fail and demonstrate that it's fixed, maybe you can submit an appropriate patch?