fix(moe): align the CUTLASS NVFP4 fc1 scale-row check with the kernel layout

[elwhyjay]

The CUTLASS NVFP4 fused MoE binding rejects intermediate_size values that are multiples of 64 but not of 128, for example 1856 (Nemotron-3.5-Lightning-30B-A3B), with fc1 weight block size must be (num_experts_on_rank, inter_size, hidden_size // 4 // block_scale_vector_size).

As far as I can tell, the fc1 row term of that check aligns inter_size with MinKDimAlignmentNVFP4 (64), while both producers of that tensor pad the rows to MinNDimAlignmentNVFP4 (128): getOffsetWeightSF in the kernel and prepare_cutlass_nvfp4_weights / CutlassNvfp4Runner.pack_inputs on the Python side (round_up(gemm1_rows, 128)). The fc2 check in the same function already uses the row alignment for hidden_size. When I is a multiple of 128 the two constants agree, which is probably why the common shapes never hit this.

This PR aligns the full fc1 row count (inter_size, or inter_size * 2 when gated) with MinNDimAlignmentNVFP4 in both branches, and adds a test at intermediate sizes 96 and 192 for ReLU2 and SwiGLU so the padded-row shapes are exercised. The gated form aligns the doubled row count rather than doubling the aligned half, matching the kernel offset and the Python padding (for I=1856 gated the tensors carry 3712 rows, not 3840).

I might be missing a reason the stricter constant was intended. If NVFP4 is meant to require 128-multiples, the cleaner fix would be a Python-side rejection in CutlassNvfp4Runner, as the MXFP paths already do, and I am happy to switch to that instead.

🔍 Related Issues

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

On an RTX PRO 6000 Blackwell (SM120), pytest tests/moe/test_unified_moe_cutlass.py -k "nvfp4 or mxfp8" gives 64 passed and 11 skipped (other architectures), including the new intermediate_size 96/192 x ReLU2/SwiGLU cases. The Nemotron shape (H=2688, I=1856, E=128, top_k=6, ReLU2, NVFP4 W4A4) also passes the flashinfer_benchmark.py --refcheck comparison at 1, 512 and 4096 tokens, with CUTLASS latencies of 59 / 566 / 957 us, in line with the numbers reported in #4646.

Summary by CodeRabbit

• Bug Fixes

  • Corrected NVFP4 shape validation for gated and non-gated mixture-of-experts models.
  • Improved handling of intermediate sizes and FC1 scale padding.
  • Ensured supported configurations execute successfully with finite, numerically accurate outputs.

• Tests

  • Added coverage for gated and non-gated activations using intermediate sizes of 96 and 192.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Advanced

Run ID: 489e5636-4435-4c41-aa4e-a8a22a429837

📥 Commits

Reviewing files that changed from the base of the PR and between 94b9234 and e770d12.

📒 Files selected for processing (2)

• csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu
• tests/moe/test_unified_moe_cutlass.py

Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review.

---

📝 Walkthrough

Walkthrough

The change corrects NVFP4 FC1 intermediate-dimension alignment for gated and non-gated activations. Tests now cover intermediate sizes 96 and 192, including scale padding, execution, finite outputs, and numerical agreement.

Changes

NVFP4 alignment validation

Layer / File(s)	Summary
Update FC1 dimension alignment csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu	Gated and non-gated NVFP4 validation now uses MinNDimAlignmentNVFP4 for the intermediate dimension.
Add parameterized NVFP4 coverage tests/moe/test_unified_moe_cutlass.py	The test helper accepts configurable intermediate sizes. Tests cover ReLU2 and SwiGLU cases with scale-row padding, execution, finite outputs, and quantized-reference comparisons.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to e770d

The alignment update is covered for padded gated and non-gated NVFP4 shapes, with no actionable merge-blocking risk identified.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description check	✅ Passed	The description clearly explains the NVFP4 validation issue, the alignment fix, affected branches, tests, benchmark results, and reviewer consideration. The Related Issues section is empty, but no iss…
Title check	✅ Passed	The title clearly and concisely identifies the main change: aligning the CUTLASS NVFP4 fc1 scale-row check with the kernel layout.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}