fix(moe): MxFp8 weight scales must be 128x4 interleaved after the row shuffle (#4087) - #4093
Open
YangXu1990uiuc wants to merge 1 commit into
Open
YangXu1990uiuc wants to merge 1 commit into
YangXu1990uiuc wants to merge 1 commit into
Conversation
… shuffle (gh #4087) The MxFp8 unified-MoE weight preparation applied only the row permutation to W1/W2 scale factors, but the trtllm-gen kernels consume the legacy shuffle_matrix_sf_a layout: row permutation THEN 128x4 block_scale_interleave. The runner forces use_shuffled_weight=True for MxFp8 and the C++ launcher hard-requires it, so ~99.8% of scale bytes were read from the wrong interleaved address. The corruption stayed sparse and sub-tolerance in existing tests because random weights repeat E8M0 exponents (misplaced reads often land on an equal exponent), the original layout unit test hard-coded the permutation-only layout as expected, and the functional test's atol=0.05 dwarfs the ~0.0025 output scale. Fix: interleave the permuted scales for both W1 and W2, matching the legacy pipeline byte-for-byte. Tests: tests/moe/test_mxfp8_weight_prepare_layout.py — the w2 (non-gated) recipe is compared byte-exact against legacy shuffle_matrix_sf_a using sentinel data where every scale byte is distinct (with random data, misplacements cancel); a gated-path guard asserts the shipped layout is not permutation-only. Validated on B200: 4/4 previously-failing fuzzer configs pass with this change (full 16-config sweep in the PR). AI-assisted (root cause identified via coordinate-dump forensics and an external review pass; A/B validated on live SM100 hardware). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
YangXu1990uiuc
requested review from
IwakuraRein,
aleozlx,
jiahanc,
nv-yunzheq,
samuellees and
yzh119
as code owners
July 22, 2026 13:16
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Contributor
📝 WalkthroughWalkthroughThe MXFP8 weight preparation path now validates scale-layout dimensions and applies ChangesMXFP8 scale layout
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/moe/test_mxfp8_weight_prepare_layout.py`:
- Around line 40-88: Strengthen the regression tests by exercising
prepare_trtllm_fp8_block_weights rather than only composing
block_scale_interleave and permutation helpers directly. In both w2 and gated
w3/w1 cases, compare the API’s two returned scale tensors with the
sentinel-derived expected legacy/interleaved layouts, including shape and
byte-content checks. Add invalid-dimension parametrized cases that assert
prepare_trtllm_fp8_block_weights raises ValueError.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 6b02798c-7b28-4860-b457-6f799922a810
📒 Files selected for processing (2)
flashinfer/fused_moe/prepare.pytests/moe/test_mxfp8_weight_prepare_layout.py
Comment on lines
+40
to
+88
| @pytest.mark.parametrize("rows,cols", [(1024, 768), (256, 512), (2048, 1024)]) | ||
| def test_w2_sf_shuffle_matches_legacy(rows, cols): | ||
| """Non-gated (w2) scale path: permute+interleave == shuffle_matrix_sf_a.""" | ||
| _skip_unless_sm100() | ||
| from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache | ||
| from flashinfer.quantization.fp4_quantization import shuffle_matrix_sf_a | ||
|
|
||
| dev = torch.device("cuda:0") | ||
| sf = _sentinel_sf(rows, cols // 32, dev) | ||
|
|
||
| cache = {} | ||
| permute_sf = get_w2_permute_indices_with_cache(cache, sf, 128, num_elts_per_sf=32) | ||
| got = block_scale_interleave(sf[permute_sf.to(dev)].contiguous()) | ||
|
|
||
| want = shuffle_matrix_sf_a(sf, 128, num_elts_per_sf=32) | ||
|
|
||
| assert got.shape == want.shape, f"{got.shape} vs {want.shape}" | ||
| n_diff = int((got != want).sum()) | ||
| assert n_diff == 0, f"{n_diff}/{want.numel()} scale bytes misplaced" | ||
|
|
||
|
|
||
| @pytest.mark.parametrize("rows,cols", [(2 * 768, 1024), (2 * 256, 512)]) | ||
| def test_w31_sf_is_interleaved(rows, cols): | ||
| """Gated (w3_w1) scale path: the result must be 128x4-interleaved, i.e. | ||
| NOT equal to the permutation-only layout #4026 produced (the exact gh | ||
| #4087 regression), and elementwise-preserving (a permutation of bytes).""" | ||
| _skip_unless_sm100() | ||
| from flashinfer.fused_moe.core import _maybe_get_cached_w3_w1_permute_indices | ||
|
|
||
| dev = torch.device("cuda:0") | ||
| sf = _sentinel_sf(rows, cols // 32, dev) | ||
|
|
||
| cache = {} | ||
| permute_sf = _maybe_get_cached_w3_w1_permute_indices( | ||
| cache, sf, 128, num_elts_per_sf=32, is_gated_act_gemm=True | ||
| ) | ||
| permuted_only = sf[permute_sf.to(dev)].contiguous() | ||
| got = block_scale_interleave(permuted_only) | ||
|
|
||
| # Regression guard: the shipped layout must NOT be the permutation-only | ||
| # layout #4026 produced (the exact gh #4087 regression). | ||
| assert not torch.equal( | ||
| got.flatten()[: permuted_only.numel()], permuted_only.flatten() | ||
| ), "interleave was a no-op -- layout regression (gh #4087)" | ||
| if got.numel() == permuted_only.numel(): | ||
| # Pure relayout (no padding): byte multiset must be preserved. | ||
| assert torch.equal( | ||
| got.flatten().sort().values, permuted_only.flatten().sort().values | ||
| ), "interleave lost/duplicated scale bytes" |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Exercise prepare_trtllm_fp8_block_weights in these regressions.
Both tests construct got directly from the expected primitives, so they still pass if prepare.py skips interleaving, uses the wrong permutation, or reshapes the returned bytes incorrectly. Invoke the preparation API and compare its two returned scale tensors against the sentinel-derived legacy layouts; also add invalid-dimension cases asserting the new ValueError.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/moe/test_mxfp8_weight_prepare_layout.py` around lines 40 - 88,
Strengthen the regression tests by exercising prepare_trtllm_fp8_block_weights
rather than only composing block_scale_interleave and permutation helpers
directly. In both w2 and gated w3/w1 cases, compare the API’s two returned scale
tensors with the sentinel-derived expected legacy/interleaved layouts, including
shape and byte-content checks. Add invalid-dimension parametrized cases that
assert prepare_trtllm_fp8_block_weights raises ValueError.
feih-nv
reviewed
Jul 23, 2026
| raise ValueError( | ||
| "MXFP8 requires hidden_size and intermediate_size divisible by 32." | ||
| ) | ||
| if hidden_size % 128 != 0 or (2 * intermediate_size) % 128 != 0: |
Collaborator
There was a problem hiding this comment.
is (2 * intermediate_size) % 128 != 0 correct? I heard intermediate_size=64 cannot be supported
Contributor
|
Related to #3882? |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
Fixes the MxFp8 half of #4087. The unified-MoE MxFp8 weight preparation applied only the row permutation to W1/W2 scale factors, but the trtllm-gen kernels consume the legacy
shuffle_matrix_sf_alayout: row permutation then 128x4block_scale_interleave. The runner forcesuse_shuffled_weight=Truefor MxFp8 and the C++ launcher hard-requires it, so ~99.8% of scale bytes were read from the wrong interleaved address. Errors stayed sparse and sub-tolerance in existing tests because random weights repeat E8M0 exponents (a misplaced read often lands on an equal exponent and cancels), the original layout unit test hard-coded the permutation-only layout as expected, and the functional test'satol=0.05dwarfs the ~0.0025 output scale. Root-cause credit: external review pass on the #4087 coordinate-dump forensics.Changes
prepare.pyMxFp8 branch: interleave the permuted scales for both W1 and W2, then reshape back to the logical(rows, cols)view the runner validates (the kernel consumes the raw interleaved bytes) — the legacy pipeline's exact idiom.hidden_sizeand2*intermediate_sizedivisible by 128 (the same constraint the legacy path carried implicitly); clearValueErrorinstead of a confusing downstream reshape failure.tests/moe/test_mxfp8_weight_prepare_layout.py: the w2 (non-gated) recipe is compared byte-exact against legacyshuffle_matrix_sf_ausing sentinel data where every scale byte is distinct (with random data, misplacements cancel — the sentinel is what makes the test able to fail); a gated-path guard asserts the shipped layout is not permutation-only.Validation (B200-class SM100, main @ 6104afe)
🔍 Related Issues
Fixes the MxFp8 layout half of #4087. The deepseekfp8 reference calibration and the nvfp4-FL s120 failure remain tracked there.
🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
Tests