[ch-wan]

This PR is part of a stack (oldest at bottom):

• [Refactor] moe: Scope shared-expert configuration to each model #44913
• [Refactor] moe: Split the two meanings of num_fused_shared_experts #44912
• [Fix] moe: Renormalize per-rank shared-slot routing over routed experts only #44904
• -> [Fix] moe: Count routed_scaling_factor once when top-k folds it #44901

Motivation

DeepSeek-style MoE blocks scale the routed experts' output by routed_scaling_factor (rsf; 2.5 for DeepSeek-V3). Some quant methods and runners fold the factor into the top-k weights instead (FusedMoE.should_fuse_routed_scaling_factor_in_topk: ModelOpt NVFP4 on non-Marlin runners, FP8 on CUTLASS or FlashInfer TRT-LLM routed, unquantized TRT-LLM routed, and the MXFP4 methods that declare it), and on ROCm two aiter routes fold it on their own: grouped routing with a correction bias (aiter_biased_grouped_topk, or the unified Triton router) and ungrouped sqrtsoftplus routing (aiter topk_gating / rocm_router_gate, DeepSeek-V4's gate). Grouped routing without a bias, the other ungrouped routes and HashTopK do not. Once the factor is in the top-k weights, no later step may apply it again, and a fused shared expert's weight must not compensate for a multiply that never happens.

Each step after top-k decided this for itself, and six of them got it wrong (static reading at main):

1. Two-batch overlap. DeepseekV2MoE.op_output, the TBO epilogue, skipped the multiply only on aiter, while forward_deepep also skips it when the quant method folds the factor. With --enable-two-batch-overlap, an all-to-all backend and a folding quant method (for example ModelOpt NVFP4 with --moe-a2a-backend deepep --moe-runner-backend flashinfer_cutedsl; nothing in server-arg resolution rejects it), every MoE layer's routed output was scaled by rsf², 6.25 instead of 2.5 for DeepSeek-V3.
2. Per-rank shared-expert slot. remap_topk_for_per_rank_shared_slots gave the fused shared expert a weight of 1/rsf unless on aiter (its comment noted the gap). On the per-rank slot path (DeepEP-class and MegaMoE backends with --enforce-shared-experts-fusion) with a folding quant method nothing scales the output afterwards, so the always-on shared expert counted 0.4x in every MoE layer. DeepSeek-V4.1's vision routing worked around this for MegaMoE in [1/N][DSv41] Enable DP attention with MegaMoE #43228; the other paths did not.
3. Waterfill. WaterfillBalancer always used 1/rsf, wrong both with a folding quant method and on aiter. --enable-waterfill turns shared-experts fusion back on even for runners that disable it (such as flashinfer_cutedsl), so ModelOpt NVFP4 with DeepEP and Waterfill reaches it.
4. Unified grouped router on aiter. With SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK=1 on aiter, biased_grouped_topk_gpu folded the factor only when given fused shared slots, but since MiniMax-M3: allow shared-experts fusion on ROCm gfx942 and newer #36576 select_experts gives it none on aiter. The routed weights then lacked the factor while every later step assumed it, leaving the routed output 1/rsf too small.
5. GLM, dots3 and bailing all-to-all epilogues on aiter. The forward_deepep of glm4_moe, glm4_moe_lite and the dots3 model, and bailing_moe_v3's all-to-all scaling, applied the multiply unless the quant method's fold flag was set, also on aiter where their biased grouped routing already folds it (GLM's gates always carry a bias). With a DeepEP-class backend on ROCm (mori, for glm4_moe) the routed output was scaled twice and a fused shared expert by rsf. For bailing_moe_v3 this became true for the unified router too once item 4 is fixed.
6. Routes that do not fold, on aiter. DeepseekV2MoE treated every aiter route as folded. DeepSeek-V4's hash layers (HashTopK uses the factor only for the 1/rsf shared weight) and grouped routing without a correction bias (DeepSeek-V2's group_limited_greedy) therefore got the factor zero times on aiter, in forward_normal and the all-to-all epilogues alike.

Modifications

topk.py now owns the answer:

• is_routed_scaling_factor_fused_in_topk(topk_config): whether the routed weights a TopK with that config emits already carry the factor: its fold flag, or on aiter one of the two folding routes above. A HashTopK folds only for its own flag.
• fused_shared_expert_weight(routed_scaling_factor=..., fused_in_topk=...): the shared slot's weight derived from that answer (1.0 when folded or when there is no factor, else 1/rsf).

How each site decides now:

Site	Before	Now
DeepseekV2MoE.forward_deepep	fold flag or aiter	routed_scaling_factor_fused_in_topk, set once in __init__ from the router it built (the helper for TopK, the flag for HashTopK)
DeepseekV2MoE.forward_normal, forward_normal_dual_stream	skip the multiply on aiter	skip it on aiter only when that attribute says the route folded
DeepseekV2MoE.op_output (TBO)	aiter only	the same attribute, through the epilogue it now shares with forward_deepep (_scale_routed_and_add_shared)
MegaMoE, DeepGEMM and FlyDSL (_run_mega_routed)	fold flag; fold flag or aiter	the same attribute
DeepSeek-V4 TopK fold flag	True on aiter, else fold flag	unchanged; comment corrected (aiter's topk_gating folds anyway, the flag covers the other routes)
remap_topk_for_per_rank_shared_slots	1.0 on aiter, else 1/rsf	fused_shared_expert_weight from the config's answer; HashTopK now passes its fold flag in that config
aiter fused append on the per-rank path	literal 1.0	fused_shared_expert_weight (still 1.0 there)
WaterfillBalancer	always 1/rsf	takes shared_weight, which prepare_moe_topk computes with fused_shared_expert_weight from the TopK's answer or the HashTopK's flag
unified grouped router (biased_grouped_topk_gpu)	folds on aiter only when given shared slots	always folds on aiter, like aiter_biased_grouped_topk
DeepSeek-V4.1 vision routing on MegaMoE	its own override of the shared weight to 1.0	removed; the remap now writes 1.0
forward_deepep of glm4_moe, glm4_moe_lite, dots3; bailing_moe_v3's all-to-all scaling	fold flag	the helper on their TopK config (same on CUDA)

Existing normalized CUDA routes without a folding quant method retain their scaling behavior. The per-token-bias route now retains the configured fold even without normalization and emits a routed-sum-independent shared coefficient. On aiter, routes that do not fold now get the factor once instead of zero times in DeepseekV2MoE. The global AITER append retains its existing default. DeepSeek's nonfolding global route supplies an inverse-rsf weight, multiplied by its existing EP replication factor, because that model restores rsf after the fused MoE output. One edge changes: with no rsf (None or 0) and no fold, the per-rank remap now writes 1.0 instead of leaving whatever the router put in the shared column.

The dynamic per-token-bias router now uses the same folding policy as the static biased AITER route. Its torch implementation applies an explicit fold even without normalization and writes the shared coefficient independently of the routed-score sum. The eager/fused remap equivalence test exercises folding and nonfolding configurations explicitly instead of inferring folding from the platform.

Accuracy Tests

Unit tests:

• test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py (CPU), extended:
  • Waterfill's shared weight, built through prepare_moe_topk, and the per-rank slot's are 1.0 exactly where the route folded: post-MoE multiply, folded by the flag, aiter biased grouped routing, and aiter grouped routing without a bias (1/rsf);
  • DeepseekV2MoE.__init__ (built with a stub gate and experts) derives routed_scaling_factor_fused_in_topk from the experts' fold flag and the route: biased or unbiased grouped routing, DeepSeek-V4's gate, and hash layers;
  • the TBO epilogue (DeepseekV2MoE.op_output) scales the routed output once, with and without a shared output.
• test/registered/kernels/ops/moe/test_jit_grouped_topk.py (GPU): on aiter the unified router's routed weights carry rsf when given no shared slots; elsewhere they do not unless asked.
• test/registered/moe/test_hash_topk.py: the per-rank shared slot of a HashTopK keeps 1/rsf on aiter too.
• Stubs that model a DeepseekV2MoE gained the new attribute: test_mega_moe_deepgemm_api.py, test_mega_moe_decode_capacity.py, test/registered/moe/test_hash_topk.py.

Runs on one 8x H200 machine (Python 3.12, torch 2.14.1), each file in its own process, CPU-registered files with the GPUs hidden and CUDA-registered files on one GPU:

• Before the fix (86e67f7d93 with only the first commit's new tests applied): test_fused_shared_expert_scaling.py ran 9 tests with 4 failures and 1 error, exactly the new cases: the folded per-rank slot (0.4 instead of 1.0), Waterfill when folded and on aiter (0.4 instead of 1.0), and both folded TBO epilogue cases. test_jit_grouped_topk.py: the aiter case failed and the other passed (1 failed, 7 passed, 4 skipped single-group cases that CUDA does not route).
• After the first commit (on 86e67f7d93), 91 unit test files for these areas (unit/layer_boundary/, unit/layers/moe/, unit/models/test_deepseek_*.py, unit/models/test_shared_experts_fusion_gates.py, unit/eplb/, unit/batch_overlap/, unit/lora/test_lora_moe_inplace_unit.py, moe/test_topk_*.py) all exited 0: 626 unittest and 186 pytest cases passed, 13 skipped by platform gates.
• First commit's head (e19a5728f8): the targeted files (the tests above, test_fused_append_remap_per_rank_shared_slots.py on GPU, test_stage_producers_leave_sums.py, test_moe_output_reduction.py, the GLM and dots tests) all passed.
• Second commit (route-aware answer), its new tests against the first commit's code: test_fused_shared_expert_scaling.py failed exactly the route cases (aiter grouped routing without a bias in the construction test and in the shared-weight table, and an aiter DeepSeek-V4 hash layer; 3 failures), and test_hash_topk.py's aiter case failed (shared weight 1.0 instead of 0.4).
• Second commit's head (bd4de12cfa): the same 13 targeted files all passed, glm4_moe, glm4_moe_lite, dots3 and bailing_moe_v3 import, and the same 91 area files as above all exited 0 (626 unittest and 186 pytest cases passed). The tested tree is identical to the pushed commit's.

Owed, not run:

• A GSM8K and long-generation accuracy A/B on a folding configuration under TBO, for example DeepSeek-V3/R1 ModelOpt NVFP4 with --moe-a2a-backend deepep --moe-runner-backend flashinfer_cutedsl --enable-two-batch-overlap.
• The same A/B with fused shared experts on the per-rank path: --moe-a2a-backend megamoe --enforce-shared-experts-fusion, or DeepEP with --enable-waterfill.
• ROCm aiter runs: DeepSeek-V3 with SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK=1, DeepSeek-V4 (hash layers), and a GLM-4.5-class Glm4MoeForCausalLM with --moe-a2a-backend mori. No unit test reaches the GLM, dots3 and bailing all-to-all epilogues.

Additional regression validation at 6c1b70e909 (one H200; torch 2.14.1+cu130, FlashInfer 0.7.2, Triton 3.8.0, sgl-kernel 0.4.10):

• All 12 scaling unit tests pass with GPUs hidden. Running the new cases against the previous production code produces six numerical failures: dynamic bias with/without normalization, unnormalized dynamic shared coefficients in both folding modes, and DeepSeek's unbiased global shared route at EP1/EP4.
• All six registered CUDA append/remap tests pass. With the AITER policy flag forced while using real CUDA kernels, the previous equivalence test fails in seven cases; the repaired test passes in both folding modes. This checks policy compatibility, not AMD execution.
• Applicable pre-commit hooks pass.

Speed Tests and Profiling

No kernel changes. The decisions are made once at construction or read from a stored boolean, as before.

Related

• Supersedes [Fix] Skip redundant routed_scaling_factor in TBO DeepseekV2MoE.op_output #35152 and [Fix] shared expert under-weighted when rsf is pre-folded in topk weights #35163, which were closed for inactivity; thanks to their author.
• Follows [AMD] fix(moe): correct fused shared-expert scaling on aiter/DeepEP path (mori all-to-all) #28237, which fixed the aiter case of the per-rank shared slot.
• [Perf] moe: apply the routed scaling factor inside the fused gate for FP8 MegaMoE #44878 (open) makes FP8 with MegaMoE fold the factor in top-k; with fused shared experts on that path it relies on the per-rank slot fix here.
• Left out, same family but a different decision (whether the runner applied the factor, not top-k), each needing a ROCm run: DeepseekV2MoE.forward_normal multiplies the factor on HIP without aiter although the Triton runner's final sum already applies it there ([Fix] [XPU] glm4_moe_lite: Skip model-side routed scaling on XPU (fix GLM-4.7-Flash accuracy) #43891, open, addresses the XPU case of this pattern in glm4_moe_lite); and the TP paths (forward_normal) of dots3 and bailing_moe_v3 never multiply, so on aiter a route that does not fold (dots3 without a correction bias, bailing with ungrouped or softmax routing) gets no factor.

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.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

CI fixture follow-up at cdf052beb7

The ROCm 10 MI35x job failed at faaf4bb4e6 because test_deepseek_v4_amd_tp4.py::_moe_stub bypasses DeepseekV2MoE.__init__ and omitted routed_scaling_factor_fused_in_topk. The stub now supplies that construction state. The other five test files in that partition passed. This fixes the test fixture; production executable behavior is unchanged.

On one H200, the original fixture reproduces the same AttributeError with the AITER epilogue flags forced; after repair, all four normal/dual-stream and shared-TP1 combinations pass through the actual forward methods. This checks the Python fixture/epilogue contract, not native ROCm collectives or graph replay. The complete AMD TP4 file still needs CI validation.

TopKConfig documents its total ordinary output width and the global AITER absolute shared coefficient versus the standard-path multiplier. Eager/fused subtest labels include the folding mode. Duplicate ordinary model-context scopes and a derived routed-count property remain optional follow-ups. Applicable pre-commit hooks pass. No native AMD rerun or model accuracy result is claimed for this follow-up.

Landing cleanup (2026-10-19)

Current head: 8f95be3e63, rebased onto current main.

Trimmed repetitive comments and docstrings while retaining scaling and layout contracts. The cleanup preserves the executable AST; configured pre-commit checks passed. Earlier test results above retain their recorded revisions. No new accelerator or full-model accuracy run is claimed for this cleanup/rebase.

The final PR in this stack moves the extended routing matrix and compile-context check to opt-in Extra coverage while keeping the focused regressions on the normal path.

---

CI States

Latest PR Test (Base): ❌ Run #39344262524
Latest PR Test (Extra): ❌ Run #39344262390
Latest PR Test (AMD ROCm 10): ❌ Run #39344262320

[xutizhou]

Adding an A/B for the Waterfill item (#3 in the description) that the PR lists as owed: "DeepEP with --enable-waterfill".

Setup. 8x H20-3e, DeepSeek-V3 (FP8 checkpoint), --tp 8 --ep-size 8 --moe-a2a-backend deepep --deepep-mode auto --moe-runner-backend deep_gemm --enable-waterfill, image lmsysorg/sglang:nightly-dev-cu13-20261019-3c560089 with main @ 617c8b6 mounted. Hopper cannot run the ModelOpt NVFP4 configuration, so I exercised the folded branch with a test-only switch that makes _fuses_routed_scaling_factor_in_topk() return True. On FP8 + deep_gemm + DeepEP that is mathematically equivalent to the real folding quant methods: the router multiplies the routed weights by routed_scaling_factor (checked: row sums go 1.0 -> 2.5) and forward_deepep skips the epilogue multiply, so the only remaining question is the shared slot's weight.

Two metrics: GSM8K (sglang.test.few_shot_gsm8k, 1319 questions, 5-shot) and a teacher-forced probe (mean NLL per token over the answer region of 64 GSM8K examples, 13,252 tokens; greedy continuations of 20 prompts).

Case	code	rsf folded in top-k	Waterfill	shared slot weight	GSM8K acc	NLL / token
reference	main	no	off (shared expert computed separately)	–	0.938	0.7056
default FP8 path	fixed	no	on	0.4 (= 1/rsf)	0.933 / 0.938 (two runs)	0.7037
folded, unfixed	main	yes	on	0.4	0.936 / 0.935 (two runs)	0.8579
folded, fixed	fixed	yes	on	1.0	0.936	0.7016

"fixed" = WaterfillBalancer(shared_weight=fused_shared_expert_weight(routed_scaling_factor=..., fused_in_topk=is_routed_scaling_factor_fused_in_topk(topk_config))) in prepare_moe_topk, i.e. the same change as this PR's waterfill.py / moe_ep_setup.py hunks (I carried the two helpers over with the same names and semantics).

Observations:

• The 0.4x shared expert is invisible to GSM8K accuracy: all four configurations land in 0.933–0.938, run-to-run noise. Please don't use GSM8K alone to sign this one off.
• The teacher-forced probe sees it clearly: +21.6% NLL/token (0.706 -> 0.858), and every one of the 64 examples gets worse (+31.5 +/- 11.7 nats per example). Greedy outputs stay fluent, which is why accuracy does not move.
• With the shared weight derived from the fold, the folded path is back at the reference (0.7016 vs 0.7056; per-example differences are +/- 2.6 nats around zero, half better, half worse, i.e. bf16/FP8 kernel noise). The default non-folding path is unchanged by the fix (shared weight stays 1/rsf), as expected.

I'll also add a GPU regression test for Waterfill (test/registered/moe/test_waterfill.py: expansion layout, rank selection, shared weight in both scaling modes, and a DeepseekV2MoE.forward_deepep end-to-end check with a linear stand-in for the experts) so that later changes around top-k / epilogue scaling trip a test instead of a silent 0.4x.

[xutizhou]

Follow-up: the Waterfill regression test mentioned above is #45088 (test only, depends on this PR; on current main exactly the folded-mode cases fail, with this PR's Waterfill change it passes on H20).