Repository navigation
[Fix] MiMo-V2: accept checkpoints with the split attention layout - #40978
Merged
Merged
Conversation
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>
ziang-and
pushed a commit
to zianglih/sglang
that referenced
this pull request
Oct 2, 2026
… 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 `layout in (None, "split")` condition and the ValueError text already name the accepted layouts; the PR body records why "split" needs no fused-qkv TP check. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
alexnails
approved these changes
Oct 6, 2026
guapisolo
force-pushed
the
jiajun/mimo-v2-accept-split-layout
branch
from
October 7, 2026 09:47
31360f9 to
15bbf29
Compare
Collaborator
Author
|
/rerun-test test/registered/unit/configs/test_model_config.py test/registered/e2e/models/test_mimo_v2_flash.py |
Contributor
|
🚀 🚀 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
MiMo-V2 checkpoints with separate q/k/v projections now load.
Symptom & Reproduction
q_proj/k_proj/v_projstops at model-config validation withValueError: MiMoV2 hf_config has unsupported attention_projection_layout='split'; expected 'fused_qkv' or unset.attention_projection_layoutinconfig.json, then runpython -m sglang.launch_server --model-path <split checkpoint> --trust-remote-code --tp 4; expected: the server starts; observed: the error above.Root Cause
MiMoV2Configfillsattention_projection_layout="split"whenconfig.jsonomits it.get_mimo_v2_fused_qkv_expected_tp_size(model_config.py:122) returns early only forNone.fused_qkvthen raisesValueError, so split checkpoints never load.Fix
Treat
"split"like an unset layout. Split projections already load through the stacked-params path ofMiMoV2ForCausalLM.load_weights(("qkv_proj", "q_proj", "q")), and the fused-qkv TP constraint does not apply to them.fused_qkvcheckpoints and unknown values behave as before.Verification
mainatc389cbaf99without conflicts; no new test results are claimed for the rebased head.pre-commitonmodel_config.pypassed on the previous head; it has not been rerun on this rebased head.Historical sglang-miles report (#40977): the same patch loaded and served BF16 split MiMo-V2.6-Flash-RL conversions (4-layer partial and full model) through Miles colocated RL. This does not validate the rebased
mainhead.test/registered/unit/configs/test_model_config.py(CPU) andtest/registered/e2e/models/test_mimo_v2_flash.py(8-GPU H200): pending targeted CI runs after publication. Existing tests do not specifically cover the converted MiMo-V2.6 split checkpoint.Review Focus
get_mimo_v2_fused_qkv_expected_tp_size: returningNoneforsplitmust not skip a TP check that split checkpoints need.Generated with Claude Code
CI States
Latest PR Test (Base): ❌ Run #37603000643
Latest PR Test (Extra): ❌ Run #37603000386
Latest PR Test (AMD ROCm 10): ❌ Run #37603000634