[sglang-miles] Cherry-pick #40448, #40648, #40978, #40979, #40980 for MiMo-V2.6-Flash colocated RL by guapisolo · Pull Request #40977 · sgl-project/sglang

guapisolo · 2026-09-23T22:25:28Z

Summary

Cherry-pick five upstream PRs so sglang-miles serves MiMo-V2.6-Flash as a colocated RL engine.

Motivation

Miles trains MiMo-V2.6-Flash-RL and rolls out on sglang-miles (radixark/miles#3648 and its stack). Without these picks, sglang-miles cannot load MXFP4 experts (#40448), rejects the split layout (#40978), keeps SWA KV resident after release (#40979), crashes routing-replay graph capture (#40980), and can hang after a graph-pool release/resume (#40648).

Usage

python scripts/run_mimo_v2_6_flash.py --mode rl --model-name mimo26-p4-bf16 [--sglang-precision mxfp4_w4a16_linear]

Miles PRs select this head with ci-sglang-pr: #40977.

Design Notes

Mental model: Miles starts a colocated engine on a BF16 split or MXFP4 checkpoint; the config accepts the split layout, SWA pools pause with the memory saver, TopK captures routed experts per layer, NCCL skips pausable graph buffers, and each sync repacks Marlin experts in place.
Base: rebuilt on sglang-miles e5bba1f9c4; every pick applied without conflicts.
[Feature] Xiaomi MiMo-V2.6/MiMo-V2.6-Pro day0 support #40448: b106483a76.
[Fix] MiMo-V2: accept checkpoints with the split attention layout #40978: bc4479cdaf.
[Fix] Allocate SWA KV pools inside the memory-saver region #40979: 13426c664b, cd3bbc937c, 38a8d24598.
[Fix] MoE: require TopK layer_id to ensure routed expert captures #40980: e322a5f040, 253ca412f0; its Nemotron-H commit is already here as [sglang-miles] Pass Nemotron-H layer ID to routing replay #38718.
fix(nccl): disable graph buffer registration for pausable graph pools #40648: 85c10c4758 through 0fe95448fb, including the two commits pushed upstream on 09-27.
[Fix] MiMo-V2: attend sliding_window_size keys in SWA layers #40981 dropped: it is closed, so SWA layers keep sliding_window_size=config.sliding_window_size.
Marlin reload (sglang-miles only): 43b7fed54a repacks MXFP4 Marlin experts in place on each sync via the [sglang-miles] Kimi K3 colocated RL: LoRA fixes and in-place MXFP4 Marlin reload #37704 helpers.
Fused qkv (sglang-miles only): the same commit regroups a BF16 fused qkv_proj below attention TP 4.

Verification

Versus the previous head b7b31c7c10, the pick files differ only by the [Fix] MiMo-V2: attend sliding_window_size keys in SWA layers #40981 removal, the two new fix(nccl): disable graph buffer registration for pausable graph pools #40648 commits, and sglang-miles's own kimi_k3.py change.
pre-commit run --files on the 32 changed files: all hooks pass; all 48 TopK( calls pass layer_id.
H200 host, CPU only: the 9 touched modules import; test_server_args.py, test_mimo_v2_fused_qkv_loader.py: passed (239 with test_topk.py); test_topk.py has 8 failures from missing sgl_kernel CPU ops, the same on plain sglang-miles.
Earlier revisions: BF16 fused qkv TP2/TP1 log-prob diff 5.55/7.46 before the regroup, 0.020/0.017 after; Gate B 0.024–0.027 (BF16 engine).
Not run on this head: GPU tests (the devbox GPUs are busy).

Review Focus

topk.py: out-of-tree TopK callers without layer_id now fail at init.
mxfp4_marlin_moe.py: the reload assumes the loader writes only the unpadded slice.

🤖 Generated with Claude Code

CI States

Latest PR Test (Base): ❌ Run #36522699365
Latest PR Test (Extra): ❌ Run #36522699124
Latest PR Test (AMD ROCm 10): ❌ Run #36522699294

MiMoV2Config fills attention_projection_layout="split" when a checkpoint omits the field, so every split q/k/v checkpoint was rejected. Split projections load through the stacked-params path; treat "split" like an unset layout. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> (cherry picked from commit cf67edd)

SWAKVPool hard-coded enable_memory_saver=False for its full and SWA sub-pools and the MHA sliding-window configurator never passed the flag, so for SWA models the KV cache was allocated outside torch_memory_saver's region and release_memory_occupation(kv_cache) was a no-op. Measured on an H200 TP2 engine of a hybrid SWA model: 85.3 GiB after load, 32.9 GiB after release(all), 85.3 GiB after resume(weights), i.e. only the weights were pausable and a colocated Miles trainer had no memory left for its own init. Thread the flag through SWAKVPool and pass the server setting from the configurator, like every other KV pool class does. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> (cherry picked from commit 100f590)

…he memory-saver region The hybrid MLA SWA builder passes explicit sub-pool kwargs to SWAKVPool with enable_memory_saver hard-coded to False, so neither MLA/DSA sub-pool was allocated inside the torch_memory_saver region and release_memory_occupation(kv_cache) left their KV cache resident. Pass the server setting there, as the MHA SWA builder now does. IndexKeyCache allocated the DSA index-K buffers outside memory_saver_adapter.region(), so every DSA pool, hybrid or not, kept its index-K resident after a release even with the flag on. Allocate them inside the KV_CACHE region, like the MLA buffers of the same pool. Measured by pausing the kv_cache tag on pools built through the configurator builders on one H200 (48 layers, 12 full / 36 SWA): hybrid MLA freed 0 -> 6.0 GiB of 5.9 GiB; hybrid DSA freed 0 -> 6.4 GiB of 6.3 GiB, 0.39 GiB of which is index-K. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> (cherry picked from commit 34ba83f)

The comment restated why the change was made and claimed an unconditional invariant that the parameter's False default does not enforce; the PR body carries that rationale. No behavior change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> (cherry picked from commit c7db69c)

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 commit 3c529b2)

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> (cherry picked from commit 56c23fb)

…ausable With --enable-memory-saver and SGLANG_MEMORY_SAVER_CUDA_GRAPH the decode graphs are captured into a torch_memory_saver region. That region is a cuMem (VMM) allocation, so NCCL's default graph registration (NCCL_GRAPH_REGISTER=1) binds the NVLS multicast and network paths of every captured collective to the pool's physical pages for the lifetime of the graph. release_memory_occupation(cuda_graph) unmaps and releases those pages and resume maps new ones at the same addresses, so after a release/resume cycle each replay reads the new pages through the addresses and the released pages through the registered paths. Ranks receive different all-reduce results and the TP group deadlocks once their batches diverge: a two-node TP16/EP16 MoE engine produced 0/512 completions with all 16 schedulers stuck in copy_done.synchronize(). Multi-node TP groups disable the custom all-reduce, which already skips its own IPC registration under this region (tms_cudagraph), so every all-reduce goes through NCCL and its registration was the uncovered path. Reuse the existing guard for the TP LM-head all-to-all and set NCCL_GRAPH_REGISTER=0 in the launcher process whenever the graph pool is pausable; an explicit value still wins with a warning that names the reason. Paired bare-server runs with the same seven-step release/resume lifecycle: unpatched 0/512 with a 16-rank watchdog and 2048 NCCL registrations of the 150 MiB graph-pool span; patched 512/512 with zero registrations and no throughput change. An RL colocated-training run (four TP16/EP16 engines, weight updates and graph release/resume every step) completed three rollout/train cycles. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> (cherry picked from commit 76fc17c)

…s decode graphs (#40649) Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> (cherry picked from commit b59d898)

Restore the upstream TP LM-head paragraph of _disable_nccl_graph_buffer_registration unchanged and state the pausable graph pool and DP-attention hazards as the live constraints only; the reproduction history stays in the PR description. The pausable-pool predicate covers every CUDA graph captured into the torch_memory_saver region, not only decode graphs. Comment-only change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> (cherry picked from commit 683927b)

handle_cuda_graph_config runs before handle_data_parallelism and always leaves cuda_graph_config set, so _dp_attention_replays_decode_graphs no longer needs the defensive getattr and None check; a pipeline reorder now fails loudly instead of silently keeping NCCL graph registration on. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> (cherry picked from commit 5aa4433)

Fold the two paragraphs this PR added to the _disable_nccl_graph_buffer_registration docstring into one sentence per hazard; the call-site reasons and the PR description carry the rest. The upstream TP LM-head paragraph is unchanged. Comment-only change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Restore TestTpLmHeadAllToAllNcclGraphRegister._resolve to its upstream form. test_without_all_to_all_env_is_untouched now keeps the DP-attention decode-graph rule out of its check by stubbing _dp_attention_replays_decode_graphs instead of forcing eager decode. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…sed qkv below TP4 Miles syncs weights into a colocated MiMo-V2 engine every step. Two load paths broke that for MXFP4 rollout engines: - Mxfp4MarlinMoEMethod.process_weights_after_loading replaces the expert Parameters with repacked Marlin tensors and renames the scales, so the first update_weights_from_tensor wrote checkpoint-shaped tensors into the Marlin layout and failed in FusedMoE._load_w2. restore_weights_before_loading now views the Marlin storage in the checkpoint layout (weight_scale_inv aliases, zeroed padding) and post-processing repacks into the same storage with reuse_parameter_storage (#37704 helpers), keeping the addresses that CUDA graphs and the memory saver hold. - load_mimo_v2_qkv_proj_weight concatenated a rank's kv-head shards in checkpoint order ([q_i|k_i|v_i]...) below the checkpoint's fused TP, while QKVParallelLinear computes with [q heads|k heads|v heads]; only FP8 codes were regrouped later by _resolve_deferred_qkv_scale_inv. A BF16 fused qkv at attention TP 1/2 therefore served wrong attention. Non-FP8 weights now regroup at load from the layer config; block-FP8 codes keep checkpoint order for the resolver; the loader raises when the shard sizes cannot be derived (MTP names), instead of loading silently. Verified on H200 with the Miles 4-layer partial: after a full weight reload the Marlin engine's experts compare byte-exact and greedy outputs match; the BF16-linear MXFP4 checkpoint at attention TP2/TP1 vs TP4 went from 5.55/7.46 to 0.020/0.017 mean input log-prob diff, the FP8 checkpoint stays at 0.047. New CPU test test/registered/unit/models/test_mimo_v2_fused_qkv_loader.py. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

… 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>

The MiMo launcher imported the private _exclusive_path_lock from base_backend. Make it public as exclusive_path_lock, re-export it from command_utils, and use U.exclusive_path_lock in the launcher and its regression test. Docs section 3.2 still said the MiMo-V2 SGLang fixes were not in the image. sgl-project/sglang#40977 merged into sglang-miles, whose 8035002 carries them, and radixark/miles:dev installs sglang-miles, so the section now names that commit and the image. Move the TP8 engine note from the recipe default (4, used by the 4-layer partial) to the full-model recipe that sets 8. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

github-actions Bot added the memory-pool label Sep 23, 2026

This was referenced Sep 23, 2026

[Fix] MiMo-V2: accept checkpoints with the split attention layout #40978

Merged

[Fix] Allocate SWA KV pools inside the memory-saver region #40979

Merged

github-actions Bot mentioned this pull request Sep 23, 2026

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

Merged

This was referenced Sep 23, 2026

[Fix] MiMo-V2: attend sliding_window_size keys in SWA layers #40981

Closed

feat(models): add MiMo-V2.6-Flash-RL support radixark/miles#3648

Merged

guapisolo force-pushed the jiajun/mimo26-rl-sglang-40448 branch from 7e858b4 to b7b31c7 Compare September 26, 2026 03:18

guapisolo changed the title ~~[sglang-miles] Cherry-pick #40448 (MiMo-V2 MXFP4 MoE) and MiMo-V2.6-Flash colocated RL fixes~~ [sglang-miles] Cherry-pick #40448, #40648, #40978, #40979, #40980, #40981 for MiMo-V2.6-Flash colocated RL Sep 26, 2026

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

guapisolo mentioned this pull request Sep 28, 2026

feat(mimo): serve MiMo-V2.6 rollouts from MXFP4 engines radixark/miles#3744

Draft

ollybbmonster and others added 14 commits September 28, 2026 21:32

[Feature] support bf16 MoE router and mxfp4 MoE for MiMo V2 (#40448)

b106483

fix(nccl): disable graph buffer registration when DP attention replay…

7afe725

…s decode graphs (#40649) Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> (cherry picked from commit b59d898)

guapisolo changed the title ~~[sglang-miles] Cherry-pick #40448, #40648, #40978, #40979, #40980, #40981 for MiMo-V2.6-Flash colocated RL~~ [sglang-miles] Cherry-pick #40448, #40648, #40978, #40979, #40980 for MiMo-V2.6-Flash colocated RL Sep 29, 2026

guapisolo force-pushed the jiajun/mimo26-rl-sglang-40448 branch from b7b31c7 to 43b7fed Compare September 29, 2026 04:40

guapisolo marked this pull request as ready for review September 29, 2026 21:11

guapisolo requested review from Ying1123, hnyls2002 and merrymercy as code owners September 29, 2026 21:11

guapisolo requested review from Alisehen, AniZpZ, BBuf, Edwardf0t1, FlamingoPg, Fridge003, HaiShaw, OrangeRedeng, Qiaolin-Yu, alphabetc1, b8zhong, ch-wan, hanming-lu, hebiao064, huangtingwei9988, hzh0425, ispobock, mmangkad, xiezhq-hermann and yizhang2077 as code owners September 29, 2026 21:11

guapisolo merged commit 57e7d0e into sglang-miles Sep 29, 2026
78 of 87 checks passed

guapisolo deleted the jiajun/mimo26-rl-sglang-40448 branch September 29, 2026 21:11

guapisolo mentioned this pull request Oct 1, 2026

[sglang-miles] Sync #40648: drop the DP-attention NCCL graph registration trigger #41972

Merged

Navigation Menu

Repository navigation

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

[sglang-miles] Cherry-pick #40448, #40648, #40978, #40979, #40980 for MiMo-V2.6-Flash colocated RL#40977
guapisolo merged 14 commits into
sglang-milesfrom
jiajun/mimo26-rl-sglang-40448

guapisolo commented Sep 23, 2026 •

edited by github-actions Bot

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 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

Motivation

Usage

Design Notes

Verification

Review Focus

CI States

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

guapisolo commented Sep 23, 2026 •

edited by github-actions Bot

Loading