[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

With fused shared experts in per-rank slots (--enforce-shared-experts-fusion with a DeepEP-class or megamoe a2a backend), every MoE layer renormalized its routed weights over one routed expert too many on CUDA, so the routed experts it kept summed to about 0.85-0.97 instead of 1.0.

How it happened:

1. On the per-rank slot path select_experts gives the router no shared slot (num_fused_shared_experts=0), because _post_process_topk_ids appends the shared column afterwards. On aiter it also gives the router the routed count, top_k - num_fused_shared_experts. Off aiter every router call kept top_k (num_routed_topk if _use_aiter else top_k), which counts the shared slot.
2. So the router selected top_k routed experts, one more than the model's num_experts_per_tok, and renormalized their weights to sum to 1.
3. remap_topk_for_per_rank_shared_slots then overwrote the last column with the home-rank shared expert.
4. The routed experts kept sum to sum(top K) / sum(top K+1), per token, in every layer. The CUDA routers return their top-k sorted by score, so the kept experts are the right ones; only their weights shrink.

#36515 (non-grouped routers) and #40943 (grouped routers) gave the router 0 shared slots on this path. That fixed aiter, where the router already gets the routed count, and left CUDA in the state above. Before them, CUDA gave the router top_k with the shared slot, which selects the routed count and renormalizes over it.

Kept routed weight sum per token, one H200, 64 tokens of random logits, EP8, rank 3, one shared expert:

Router path	Before	After
grouped sigmoid + bias, 256 experts, top-8 (DeepSeek-V3 shape; FlashInfer router after the fix, Triton fallback before)	0.873-0.913	1.0000
grouped sigmoid + bias, 384 experts in 8 groups, top-8	0.873-0.914	1.0000
grouped, one group, sigmoid + bias, 384 experts, top-8 (Kimi-K2 shape)	0.873-0.919	1.0000
grouped softmax, no bias, 64 experts, top-6	0.886-0.967	1.0000
ungrouped sqrtsoftplus + bias, 256 experts, top-6 (DeepSeek-V4)	0.850-0.885	1.0000
ungrouped sigmoid + bias, 256 experts, top-8	0.877-0.912	1.0000
ungrouped softmax, 128 experts, top-8	0.909-0.949	1.0000
HashTopK (DeepSeek-V4 hash layers), torch and fused kernel	1.0000	1.0000
DeepSeek-V4.1 vision routing on MegaMoE	1.0000	1.0000

In every row the routed experts equal routing without fusion (64/64 rows) before and after; the shared column is the home-rank slot with weight 1/routed_scaling_factor.

Affected: CUDA, and other platforms without aiter, with shared-experts fusion on a per-rank backend (deepep, mooncake, mori, pplx, megamoe); for the DeepSeek family that takes --enforce-shared-experts-fusion. The original per-rank normalization bug does not affect aiter, Waterfill, the global fused slot (no a2a backend), runs without fusion, HashTopK and the DeepSeek-V4.1 vision routing, whose routers already select the routed experts alone.

Two related problems on the same path are fixed here too:

• FlashInfer's DeepSeek router ignores renormalize. fused_topk_deepseek always normalizes the selected weights, and its ninth parameter is launch_with_pdl, which sglang's wrapper called renormalize. With routed-only sizing the per-rank top-8 path now fits that router (top-9 did not), so a renormalize=False configuration would have been normalized. Configurations without fusion or with the global slot already were, within that router's shapes. Grouped sigmoid routing with a bias, 256 experts in 8 groups, top-8, renormalize=False, zero logits (each raw weight is sigmoid(0) = 0.5): before this PR 0.125 per expert without fusion and 0.5 with a per-rank slot; with this PR's first version 0.125 for both; now 0.5 for both.
• Padded rows on the per-rank path. The append shifted a padded row's -1 to -1 - num_fused_shared_experts and gave it a live home-rank shared slot, so padded rows did not leave top-k as -1. HashTopK and the vision routing already mask after their append.

The renormalize=False guard also exposes the Triton global shared-slot path. That path previously wrote sum(selected routed scores) / rsf before normalization, so disabling normalization made the shared output depend on the routed-score sum (4x for eight zero-logit sigmoid selections). It now writes a shared weight of 1 when rsf is folded, otherwise 1/rsf, after processing the routed weights. The dynamic-bias torch route receives the corresponding repair from #44901.

Modifications

• select_experts decides once, for every router, what the router is asked for: where _post_process_topk_ids appends the shared columns (aiter, per-rank slots), the routed count with no shared slot; otherwise top_k with the shared slots, as before. The values are unchanged on aiter, on the global slot and without fusion.
• remap_topk_for_per_rank_shared_slots takes the routed columns and appends the shared ones, the contract fused_append_remap_shared_experts_deepep already has on aiter. Routers that still write shared placeholders (the per-token-bias router, the DeepSeek-V4.1 vision routing, HashTopK) have them dropped before the append. On CUDA the EPLB remap splits off shared columns only when the router wrote them, and padded rows are masked again after the append.
• biased_grouped_topk_gpu takes the FlashInfer branch only when renormalizing; otherwise the Triton router, which honors renormalize, serves the call. The wrapper's last parameter is renamed launch_with_pdl, and the global shared columns that branch appends take fused_shared_expert_weight(...).
• RoutedExpertsCapturer.capture pads routed-only ids (the device width minus the fused shared experts) to its device width and raises on any other width; the host side already keeps only num_experts_per_tok columns. The per-rank path now captures routed-only ids, and without the pad --enable-return-routed-experts would fail with a shape mismatch.
• The fused aiter kernel's docstrings name its eager twin.

Tests

On one H200 (driver 595.71.05, torch 2.14.1+cu130, flashinfer 0.7.2, triton 3.8.0), for this branch (03d13ebf8b), its base #44901 (bd4de12cfa) and this PR's previous head merged with that base:

• test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py::test_per_rank_slot_routes_like_no_fusion (CPU): select_experts with a per-rank slot keeps the routed experts and weights of routing without fusion, summing to 1. Fails on the base, passes here.
• test/registered/moe/test_fused_append_remap_per_rank_shared_slots.py (CUDA):
  • TestSelectExpertsPerRankSharedSlots::test_routes_like_no_fusion: the same property through five CUDA routers (the DeepSeek-V3 grouped router, the Triton grouped fallback, torch grouped routing, the DeepSeek-V4 sqrtsoftplus router, softmax). Fails on the base for all five, passes here.
  • test_unnormalized_weights_stay_raw: with renormalize=False the routed weights equal the raw sigmoid scores of the reference selection, with and without a per-rank slot. Fails on the previous head (both cases) and on the base (without fusion), passes here.
• test/registered/moe/test_topk_padded_region.py::TestSelectExpertsPerRankSharedSlotPadMask (CUDA): padded rows leave the per-rank path fully -1, for a router masked after routing and one that masks itself. Fails on the base and on the previous head, passes here.
• test/registered/unit/eplb/test_waterfill_eplb.py::test_eplb_remaps_every_column_of_routed_only_per_rank_ids (CPU): EPLB remaps every routed column of routed-only per-rank ids, and the recorder sees them. Fails on the base, passes here.
• test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py::TestCaptureRoutedIdsOnly (CPU): routed-only ids land in the routed columns, with and without the DeepEP gather; any other width raises. Fails on the base (shape RuntimeError) and, for the rejection case, on the previous head; passes here.
• The remap cases in test_fused_shared_expert_scaling.py and the kernel/eager equivalence case pass routed-only columns.
• The MoE, DeepSeek, EPLB and state-capturer unit files (test/registered/unit/layers/moe/, unit/models/test_deepseek_*.py, unit/models/test_shared_experts_fusion_gates.py, unit/eplb/, unit/state_capturer/, moe/test_hash_topk.py) plus the GPU routing tests (moe/test_topk_padded_region.py, moe/test_topk_renormalize_degenerate.py, kernels/ops/moe/test_jit_grouped_topk.py, kernels/ops/moe/test_moe_fused_gate.py, kernels/ops/moe/test_round_robin_ep_balance.py, moe/test_fused_append_remap_per_rank_shared_slots.py): 55 files, each in its own process, CPU-registered ones with the GPUs hidden. All pass on the base and on this branch.
• The tables above: select_experts, HashTopK and vision_topk with the real kernels, on the base, the previous head and this branch.

Additional validation at 4676743e28 (one H200, same library versions):

• The integrated scaling file passes all 13 CPU tests; the append/remap and runtime-routing file passes all nine CUDA tests.
• The new global renormalize=False regression fails against the previous production code for both folding modes (4x shared coefficient) and passes after the kernel repair.
• The fused-gate reference suite passes all 39 cases; the degenerate routing file passes all seven cases. The reference writes the shared coefficient independently after routed normalization/scaling.
• Applicable pre-commit hooks pass.

Owed

• A GPU accuracy eval on an affected configuration, before and after: for example DeepSeek-V3 with --moe-a2a-backend deepep --enforce-shared-experts-fusion, or DeepSeek-V4 on MegaMoE (SM100). Not run.
• ROCm: not run. The aiter paths give the router the same arguments as before (code reading).

CI fixture follow-up at beb15ea999

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.

The kernel comparison now covers normalization and explicit folding independently; the complete standalone router test file passes all 52 cases on one H200. Its stale normalization explanation and CUDA JIT parity claim are corrected. The CUDA counterpart retains its existing unnormalized shared-weight behavior; caller search found tests and benchmarks, with no serving caller.

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: 39bdf12a49, rebased onto main after #44901 merged.

Trimmed repetitive comments and docstrings. Configured cleanup checks passed before landing. Rebase resolution keeps main's removal of obsolete CUDA-JIT comparison prose and preserves routed-only append behavior, including empty output widths. Earlier test results above retain their original revisions; no new accelerator or full-model run is claimed for this cleanup/rebase.

The broader routing matrix moves to opt-in Extra coverage in #44913; the focused CPU and append/remap regressions remain on the normal path.

---

CI States

Latest PR Test (Base): ❌ Run #39344592460
Latest PR Test (Extra): 🚫 Run #39344592427
Latest PR Test (AMD ROCm 10): ❌ Run #39344592456