fix(moe): support Llama4 NVFP4 router input weights on SM120 by janbernloehr · Pull Request #35504 · sgl-project/sglang

janbernloehr · 2026-08-19T09:56:05Z

Motivation

Llama 4 constructs its MoE layers with apply_router_weight_on_input=True. For ModelOpt NVFP4 on SM120/SM121, SGLang selects the FlashInfer CUTLASS MoE runner, which previously rejected this configuration during server warmup.

For top-1 routing, the required semantics can be implemented exactly by multiplying each token's hidden states by its router weight before activation quantization and passing unit final scales to CUTLASS. This follows the existing pattern in the AITER MoE runner.

The fix remains intentionally restricted to top-1 routing. With top-k greater than one, each selected expert would require a differently scaled copy of the input.

Modifications

Add top-1 router-weight pre-scaling to the common FlashInfer CUTLASS MoE execution path.
Replace the kernel's output-side router scaling with float32 unit token_final_scales when pre-scaling is active.
Normalize all token_final_scales passed through this CUTLASS path to float32, as required by the binding.
Remove the ModelOpt NVFP4 and FlashInfer CUTLASS assertions that rejected apply_router_weight_on_input=True.
Preserve explicit rejection for:
- top-k greater than one;
- activations already quantized before dispatch, including the FlashInfer CUTLASS FP4 all-gather path.
Add focused CPU unit tests covering:
- top-1 activation pre-scaling;
- float32 unit final scales;
- preservation of expert IDs;
- top-k greater than one rejection;
- pre-quantized dispatch rejection.

The normal apply_router_weight_on_input=False path is unchanged apart from enforcing the CUTLASS binding's float32 scale requirement.

Accuracy Tests

Added:

test/registered/unit/layers/moe/test_flashinfer_cutlass_router_weight.py

The tests compare the transformed hidden states and final scales against the expected top-1 formulation and verify that unsupported configurations fail explicitly.

Local validation:

python -m py_compile \
  python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py \
  python/sglang/srt/layers/quantization/modelopt_quant.py \
  test/registered/unit/layers/moe/test_flashinfer_cutlass_router_weight.py

Result: passed.

All commit-time pre-commit checks passed, including Python AST validation, isort, Ruff, Black, codespell, and registered-test validation.

The focused pytest test could not be executed in the local development environment because PyTorch is not installed. No real-checkpoint accuracy test was run for this checkout.

The equivalent patch and its SM120/SM121 validation evidence are documented in #34192. That validation used dummy weights for end-to-end liveness, so it is not claimed as a model-accuracy result.

Speed Tests and Profiling

No speed test or profiling was run for this checkout.

The equivalent patch validation documented in #34192 completed both SM120 scenarios without errors and did not show a decode throughput regression. Those results used dummy weights and a separate validated build, so they are included only as prior evidence, not as performance results for this commit.

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations. No documentation update is needed because this fixes existing runner semantics without changing a user-facing API.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

Review and Merge Process

Ping Merge Oncalls to start the process. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ✅ Run #35553178662
Latest PR Test (Extra): ❌ Run #35553178439
Latest PR Test (AMD ROCm 10): ❌ Run #35553178589

nvpohanh

[by Codex] One inline review finding.

nvpohanh · 2026-08-31T06:58:34Z

+
+    # Activations are already quantized and all-gathered before the runner on
+    # this path, so the multiply is no longer expressible here.
+    if getattr(dispatch_output, "hidden_states_scale", None) is not None:


[by Codex] Severity: style | Confidence: High

Please avoid getattr here. Both StandardDispatchOutput and FlashinferDispatchOutput define hidden_states_scale; annotate this helper with their explicit union (or a small protocol) and read dispatch_output.hidden_states_scale directly so the supported structure is type-checked.

nvpohanh · 2026-09-01T06:49:41Z

/tag-and-rerun-ci

nvpohanh · 2026-09-04T07:23:05Z

@janbernloehr could you fix the conflicts?

b8zhong · 2026-09-04T09:42:05Z

Just delete it

nvpohanh · 2026-09-07T10:44:20Z

/rerun-failed-ci

nvpohanh · 2026-09-08T08:43:39Z

/rerun-failed-ci

nvpohanh · 2026-09-09T03:09:08Z

/rerun-failed-ci

nvpohanh · 2026-09-09T06:38:07Z

/rerun-failed-ci

nvpohanh · 2026-09-14T08:45:40Z

/rerun-failed-ci

nvpohanh · 2026-09-14T12:29:28Z

/rerun-failed-ci

nvpohanh · 2026-09-16T06:29:43Z

/rerun-failed-ci

nvpohanh · 2026-09-22T02:42:01Z

All NV pipelines have passed. @b8zhong @ch-wan could we merge this? Thanks!

…input Upstream sgl-project#35504 made _run_flashinfer_cutlass read runner_config.apply_router_weight_on_input via _prescale_router_weight_on_input. The TestSmallmRuntimeFallback SimpleNamespace stub (from the SM120 small-row NVFP4 kernel pick) predates that field, so all four fallback cases raised AttributeError after the upstream merge. Add it with MoeRunnerConfig's default (False). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

janbernloehr requested review from Alisehen, AniZpZ, BBuf, Edwardf0t1, FlamingoPg, Fridge003, HaiShaw, OrangeRedeng, Ying1123, b8zhong, ch-wan, ispobock, merrymercy and mmangkad as code owners August 19, 2026 09:56

github-actions Bot added the quant LLM Quantization label Aug 19, 2026

janbernloehr force-pushed the jbernloehr/support-llama4-nvfp4-rout branch from 11dc473 to 47b6b56 Compare August 20, 2026 07:22

nvpohanh added the bug Something isn't working label Aug 31, 2026

nvpohanh reviewed Aug 31, 2026

View reviewed changes

github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Sep 1, 2026

janbernloehr force-pushed the jbernloehr/support-llama4-nvfp4-rout branch from bb02d04 to cecbbd4 Compare September 4, 2026 09:31

b8zhong reviewed Sep 4, 2026

View reviewed changes

Comment thread test/registered/unit/layers/moe/test_flashinfer_cutlass_router_weight.py Outdated

b8zhong Sep 4, 2026

Copy link
Copy Markdown

Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just delete it

b8zhong reviewed Sep 4, 2026

View reviewed changes

Comment thread python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py

b8zhong approved these changes Sep 4, 2026

View reviewed changes

fix(moe): support Llama4 NVFP4 router input weights

24f643b

janbernloehr force-pushed the jbernloehr/support-llama4-nvfp4-rout branch from cecbbd4 to 24f643b Compare September 4, 2026 12:23

b8zhong enabled auto-merge (squash) September 4, 2026 12:31

Merge branch 'main' into jbernloehr/support-llama4-nvfp4-rout

80d8337

nvpohanh mentioned this pull request Sep 8, 2026

[Bug] test_expert_pack_mxfp4.py intermittently hangs while loading JIT extension #38408

Closed

nvpohanh added 2 commits September 9, 2026 14:58

Merge branch 'main' into jbernloehr/support-llama4-nvfp4-rout

f8cfcdf

Merge branch 'main' into jbernloehr/support-llama4-nvfp4-rout

55f575e

nvpohanh approved these changes Sep 14, 2026

View reviewed changes

nvpohanh added 4 commits September 15, 2026 08:45

Merge branch 'main' into jbernloehr/support-llama4-nvfp4-rout

d1a6bf4

Merge branch 'main' into jbernloehr/support-llama4-nvfp4-rout

103b3f9

Merge branch 'main' into jbernloehr/support-llama4-nvfp4-rout

a86c9c9

Merge branch 'main' into jbernloehr/support-llama4-nvfp4-rout

8448046

nvpohanh added 4 commits September 17, 2026 11:15

Merge branch 'main' into jbernloehr/support-llama4-nvfp4-rout

f0907a6

Merge branch 'main' into jbernloehr/support-llama4-nvfp4-rout

8687a30

Merge branch 'main' into jbernloehr/support-llama4-nvfp4-rout

96c7449

Merge branch 'main' into jbernloehr/support-llama4-nvfp4-rout

bb84a20

mmangkad approved these changes Sep 22, 2026

View reviewed changes

kpham-sgl disabled auto-merge September 22, 2026 03:18

kpham-sgl merged commit 56fee88 into sgl-project:main Sep 22, 2026
161 of 181 checks passed

Conversation

janbernloehr commented Aug 19, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

Review and Merge Process

CI States

Uh oh!

nvpohanh left a comment

Choose a reason for hiding this comment

Uh oh!

nvpohanh Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

janbernloehr Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

nvpohanh commented Sep 1, 2026

Uh oh!

nvpohanh commented Sep 4, 2026

Uh oh!

b8zhong Sep 4, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

nvpohanh commented Sep 7, 2026

Uh oh!

nvpohanh commented Sep 8, 2026

Uh oh!

nvpohanh commented Sep 9, 2026

Uh oh!

nvpohanh commented Sep 9, 2026

Uh oh!

nvpohanh commented Sep 14, 2026

Uh oh!

nvpohanh commented Sep 14, 2026

Uh oh!

nvpohanh commented Sep 16, 2026

Uh oh!

nvpohanh commented Sep 22, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

janbernloehr commented Aug 19, 2026 •

edited by github-actions Bot

Loading