Fix MXFP8 MoE weight sizing for non-gated models - #36097
Open
elvischenv wants to merge 1 commit into
Open
Conversation
create_fp8_moe_weight_ hardcoded is_concat=True, so w13 was always sized as 2 * intermediate (gate + up). Non-gated MoEs fuse only up into w13, leaving the upper half as uninitialised torch.empty that no weight loader writes. NemotronH (relu2, checkpoint has up_proj/down_proj and no gate_proj) is the only MoE model in sglang with is_gated=False. It scored 0.066 gsm8k under --quantization mxfp8 versus 0.943 in bf16; with this fix it scores 0.945. Derive is_concat from the layer's gating instead. Add a CPU unit test covering both the sizing primitive and the call site. It fails on the unfixed tree and passes with the fix.
elvischenv
force-pushed
the
nemotron-mxfp8-acc-fix
branch
from
August 23, 2026 15:43
f7ecbbd to
f57f24b
Compare
elvischenv
marked this pull request as ready for review
August 23, 2026 15:43
elvischenv
requested review from
Alisehen,
AniZpZ,
BBuf,
Edwardf0t1,
FlamingoPg,
HaiShaw,
OrangeRedeng,
b8zhong,
ch-wan and
mmangkad
as code owners
August 23, 2026 15:43
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
1 participant
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.
Problem
create_fp8_moe_weight_hardcodedis_concat=True, sow13was always sized as2 * intermediate(gate + up). Non-gated MoEs fuse only up intow13, leaving theupper half as uninitialised
torch.emptythat no weight loader writes.NemotronH (
relu2, checkpoint hasup_proj/down_projand nogate_proj) is theonly MoE model in sglang with
is_gated=False.Fix
Derive
is_concatfrom the layer's gating.Result
Nemotron-3-Ultra-550B, GB300, TP8,
--quantization mxfp8, gsm8k 5-shot:Isolated by leaving the MoE in bf16 while keeping dense mxfp8, which scored 0.9439
and pointed at the MoE path.
Test
test_fp8_moe_weight_gating.py(CPU, ~1.5s) covers the sizing primitive and thecall site. Fails on the unfixed tree, passes with the fix.
CI States
Latest PR Test (Base): ❌ Run #32649456510
Latest PR Test (Extra): ❌ Run #32649456361
Latest PR Test (AMD ROCm 7.2): ❌ Run #32649456507