[Fix] MoE: require TopK layer_id to ensure routed expert captures by guapisolo · Pull Request #40980 · sgl-project/sglang

guapisolo · 2026-09-23T22:35:23Z

Summary

TopK now requires layer_id, so no MoE model can omit it and break routed-experts capture.

Symptom & Reproduction

Symptom: with --enable-return-routed-experts, MiMo-V2 fails decode CUDA-graph capture with a shape mismatch, Target sizes: [8, 1, 4, 8]; 20 other TopK calls on main (Nemotron-H, Mixtral, Llama 4, DeepSeek, Kimi, ...) omit layer_id the same way.
Reproduction: python -m sglang.launch_server --model-path <MiMo-V2 checkpoint> --trust-remote-code --tp 4 --enable-return-routed-experts; expected: graphs capture; observed: the error above.

Root Cause

TopK.__init__ defaults layer_id to None, so omitting it is silent.
capture_routed_experts_if_allowed forwards layer_id=None to the capturer.
self.buffer[:batch, layer_id, :] (state_capturer/base.py:50) inserts an axis for None.
The copy target becomes [batch, 1, num_layers, top_k], so the copy fails.

Fix

TopK.__init__ takes layer_id: int with no default, so a model that omits it fails at construction with TypeError instead of at capture time. Past fixes went model by model ([Fix] Support mimo state_capturer routing replay #30363, [sglang-miles] Pass Nemotron-H layer ID to routing replay #38718) and did not stop the next model from repeating the mistake.
All 21 calls now pass the id each model already gives its experts. The Nemotron-H line matches [sglang-miles] Pass Nemotron-H layer ID to routing replay #38718 (author Shi Dong kept).
The DBRX, DeepSeek, Xverse and AFMoE MoE blocks gain a layer_id parameter fed by their decoder layer.
The Triton fused-MoE test and benchmark pass layer_id=0.
Unchanged: calls that already passed layer_id, HashTopK, and select_experts.

Verification

pre-commit run --files <changed files> at this head: passed (isort, ruff, ruff-format, codespell, registered-test checks).
AST scan of the tree at this head: all 47 TopK(...) calls pass layer_id (deepseek_v2.py through topk_kwargs), and each passed name is a parameter or an earlier-assigned attribute of the enclosing __init__.
The MiMo-V2 line on the sglang-miles tree ([sglang-miles] Cherry-pick #40448, #40648, #40978, #40979, #40980 for MiMo-V2.6-Flash colocated RL #40977): MiMo-V2.6-Flash engines with --enable-return-routed-experts captured CUDA graphs. The Nemotron-H line ([sglang-miles] Pass Nemotron-H layer ID to routing replay #38718): NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 started with routing replay on 8 H200 GPUs.
Not run on a GPU at this head.

Review Focus

TopK.__init__: dropping the default makes out-of-tree models that build TopK without layer_id fail at init with TypeError; that is the intended guard.
dbrx.py / deepseek.py / xverse_moe.py / afmoe.py: the new MoE-block layer_id parameter has exactly one caller each.

🤖 Generated with Claude Code

CI States

Latest PR Test (Base): ✅ Run #36065710223
Latest PR Test (Extra): ❌ Run #36065709688
Latest PR Test (AMD ROCm 10): ❌ Run #36065710061

guapisolo · 2026-09-24T07:55:59Z

/rerun-test test/registered/moe/test_triton_fused_moe.py test/registered/moe/test_fused_moe.py test/registered/moe/test_topk_padded_region.py test/registered/moe/test_topk_renormalize_degenerate.py test/registered/rl/test_return_routed_experts.py test/registered/e2e/models/test_mimo_v2_flash.py test/registered/basic_perf/test_moe_throughput.py test/registered/e2e/models/test_kimi_linear_models.py test/registered/dllm/test_dllm_batching_fdfo.py test/registered/e2e/models/test_minimax_m25_basic.py test/registered/unit/models/test_granitemoe_split_expert_loading.py test/registered/unit/models/test_shared_experts_fusion_gates.py test/registered/unit/models/test_nemotron_h_mtp_reduction.py

github-actions · 2026-09-24T07:56:28Z

Results for /rerun-test test/registered/moe/test_triton_fused_moe.py test/registered/moe/test_fused_moe.py test/registered/moe/test_topk_padded_region.py test/registered/moe/test_topk_renormalize_degenerate.py test/registered/rl/test_return_routed_experts.py test/registered/e2e/models/test_mimo_v2_flash.py test/registered/basic_perf/test_moe_throughput.py test/registered/e2e/models/test_kimi_linear_models.py test/registered/dllm/test_dllm_batching_fdfo.py test/registered/e2e/models/test_minimax_m25_basic.py test/registered/unit/models/test_granitemoe_split_expert_loading.py test/registered/unit/models/test_shared_experts_fusion_gates.py test/registered/unit/models/test_nemotron_h_mtp_reduction.py:

🚀 1-gpu-h100 (5 tests): ✅ View workflow run

cd test/ && python3 registered/moe/test_triton_fused_moe.py
cd test/ && python3 registered/moe/test_fused_moe.py
cd test/ && python3 registered/moe/test_topk_padded_region.py
cd test/ && python3 registered/moe/test_topk_renormalize_degenerate.py
cd test/ && python3 registered/dllm/test_dllm_batching_fdfo.py

🚀 4-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/rl/test_return_routed_experts.py

🚀 8-gpu-h200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/e2e/models/test_mimo_v2_flash.py
cd test/ && python3 registered/e2e/models/test_minimax_m25_basic.py

🚀 2-gpu-h100 (2 tests): ❌ View workflow run

cd test/ && python3 registered/basic_perf/test_moe_throughput.py
cd test/ && python3 registered/e2e/models/test_kimi_linear_models.py

🚀 ubuntu-latest (3 tests): ✅ View workflow run

cd test/ && python3 registered/unit/models/test_granitemoe_split_expert_loading.py
cd test/ && python3 registered/unit/models/test_shared_experts_fusion_gates.py
cd test/ && python3 registered/unit/models/test_nemotron_h_mtp_reduction.py

guapisolo · 2026-09-24T09:43:23Z

/rerun-test test/registered/e2e/models/test_kimi_linear_models.py

guapisolo · 2026-09-24T09:43:44Z

/rerun-test test/registered/basic_perf/test_moe_throughput.py

github-actions · 2026-09-24T09:43:51Z

Results for /rerun-test test/registered/e2e/models/test_kimi_linear_models.py:

🚀 2-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/e2e/models/test_kimi_linear_models.py

github-actions · 2026-09-24T09:44:07Z

Results for /rerun-test test/registered/basic_perf/test_moe_throughput.py:

🚀 2-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/basic_perf/test_moe_throughput.py

guapisolo · 2026-09-24T10:26:58Z

/tag-run-ci-label

MiMoV2MoE built its TopK without layer_id, so with --enable-return-routed-experts the routed-experts capture could not index its per-layer buffer and the decode CUDA-graph capture failed with a shape mismatch ("Target sizes: [8, 1, 4, 8]"). Pass the layer id like the other MoE models do. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Cherry-picked from #38718 (merged into sglang-miles).

TopK defaulted layer_id to None, so a model that forgot it still built and only failed once routed-experts capture was on: the capturer indexes buffer[:, None, :], which inserts an axis and breaks CUDA-graph capture with a shape mismatch. The same one-line fix has been made model by model (#30363, #38718, and the MiMo-V2 commit above). Make layer_id a required keyword and pass it in the remaining MoE models. The DBRX, DeepSeek, Xverse and AFMoE MoE blocks now take layer_id from their decoder layer. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

ch-wan

LGTM

… MiMo-V2.6-Flash colocated RL (#40977)

… graph registration (sgl-project#40648, sgl-project#40978, sgl-project#40979, sgl-project#40980, sgl-project#40977, sgl-project#41972) Retain colocated training fixes above the release MiMo implementation, including address-stable Marlin reload and mandatory TopK layer IDs. Disable NCCL graph registration for TP LM-head all-to-all and pausable graph pools. Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

guapisolo force-pushed the jiajun/mimo-v2-topk-layer-id branch from 652b47b to 5cd6b5c Compare September 24, 2026 01:57

guapisolo requested review from BBuf, Edwardf0t1, Fridge003, HaiShaw, Ying1123, ch-wan, ispobock and merrymercy as code owners September 24, 2026 01:57

github-actions Bot added the deepseek label Sep 24, 2026

guapisolo changed the title ~~[Fix] MiMo-V2: pass layer_id to the MoE TopK~~ [Fix] MoE: require TopK layer_id (cherry-pick #38718) Sep 24, 2026

guapisolo changed the title ~~[Fix] MoE: require TopK layer_id (cherry-pick #38718)~~ [MoE] Require layer_id in TopK Sep 24, 2026

guapisolo changed the title ~~[MoE] Require layer_id in TopK~~ [Fix] MoE: require TopK layer_id Sep 24, 2026

guapisolo changed the title ~~[Fix] MoE: require TopK layer_id~~ [Fix] MoE: require TopK layer_id to ensure routed expert captures Sep 24, 2026

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

guapisolo force-pushed the jiajun/mimo-v2-topk-layer-id branch from 5cd6b5c to 3bd28b7 Compare September 24, 2026 19:59

guapisolo and others added 3 commits September 24, 2026 15:08

Pass Nemotron-H layer ID to routing replay

de0fc54

Cherry-picked from #38718 (merged into sglang-miles).

guapisolo force-pushed the jiajun/mimo-v2-topk-layer-id branch from 3bd28b7 to 56c23fb Compare September 24, 2026 22:08

guapisolo removed the run-ci CI: run the baseline test suite on this PR label Sep 24, 2026

ch-wan approved these changes Sep 24, 2026

View reviewed changes

guapisolo mentioned this pull request Sep 26, 2026

[sglang-miles] Cherry-pick #40448, #40648, #40978, #40979, #40980 for MiMo-V2.6-Flash colocated RL #40977

Merged

Yunzez mentioned this pull request Sep 29, 2026

[Bug] --enable-return-routed-experts returns all-zero routings on the triton-kernels / flashinfer top-k paths #41743

Open

5 tasks done

guapisolo added a commit that referenced this pull request Sep 29, 2026

[sglang-miles] Cherry-pick #40448, #40648, #40978, #40979, #40980 for…

57e7d0e

… MiMo-V2.6-Flash colocated RL (#40977)

ch-wan merged commit 3d49538 into main Sep 29, 2026
162 of 180 checks passed

ch-wan deleted the jiajun/mimo-v2-topk-layer-id branch September 29, 2026 23:29

This was referenced Sep 30, 2026

[CI Monitor] Daily Report - 2026-09-30 bingxche/sglang-ci-bot#211

Open

[CI Monitor] Daily Report - 2026-10-01 bingxche/sglang-ci-bot#212

Open

Navigation Menu

Repository navigation

[Fix] MoE: require TopK layer_id to ensure routed expert captures - #40980

[Fix] MoE: require TopK layer_id to ensure routed expert captures#40980
ch-wan merged 3 commits into
mainfrom
jiajun/mimo-v2-topk-layer-id

guapisolo commented Sep 23, 2026 •

edited by github-actions Bot

Loading

Uh oh!

guapisolo commented Sep 24, 2026

Uh oh!

github-actions Bot commented Sep 24, 2026 •

edited

Loading

Uh oh!

guapisolo commented Sep 24, 2026

Uh oh!

guapisolo commented Sep 24, 2026

Uh oh!

github-actions Bot commented Sep 24, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Sep 24, 2026 •

edited

Loading

Uh oh!

guapisolo commented Sep 24, 2026

Uh oh!

ch-wan left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

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

Uh oh!

Summary

Symptom & Reproduction

Root Cause

Fix

Verification

Review Focus

CI States

Uh oh!

guapisolo commented Sep 24, 2026

Uh oh!

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

Uh oh!

Uh oh!

guapisolo commented Sep 24, 2026

Uh oh!

guapisolo commented Sep 24, 2026

Uh oh!

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

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

guapisolo commented Sep 24, 2026

Uh oh!

ch-wan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

guapisolo commented Sep 23, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Sep 24, 2026 •

edited

Loading

github-actions Bot commented Sep 24, 2026 •

edited

Loading

github-actions Bot commented Sep 24, 2026 •

edited

Loading