Repository navigation
[Model] LongCat-Flash: scale MLA norms while loading and drop the post-load sweep - #60080
Merged
aoshen02 merged 3 commits intoOct 6, 2026
Merged
Conversation
aoshen02
requested review from
AndreasKaratzas,
DarkLight1337,
Harry-Chen,
khluu and
ywang96
as code owners
October 5, 2026 17:03
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
…t-load sweep FlashModel.load_weights ended with a sweep over every attention module that block-dequantized kv_b_proj into self_attn.w_kc / w_vc, which nothing reads (MLAAttention.process_weights_after_loading derives W_UK_T / W_UV from kv_b_proj), and scaled q_a_layernorm / kv_a_layernorm in place behind _mla_*_lora_scaled flags. LongCatFlashMTP.load_weights had the same sweep. The flags make the scaling depend on how the weights arrive: - AutoWeightsLoader calls FlashModel.load_weights once per contiguous run of model.* names. LongCat-Flash-Chat-FP8 streams 16 model.* tensors, lm_head, then the rest, so the first call scales the not-yet-loaded norms and sets the flags, and the second call loads the real norms unscaled (GSM8K 0.000). - reload_weights loads the norms again and the flags skip the scaling (LongCat-Flash-Lite: GSM8K 0.841 -> 0.036 after a reload). - With FP8 attention (LongCat-Flash-Lite-FP8) the dead block_dequant asserts on the shards a layerwise reload passes in, so reload_weights fails. Scale the norms in their weight loaders instead (composed_weight_loader, as plamo3 does), so every load writes checkpoint * factor, and remove the sweep from both files. load_weights now only loads. FlashModel.load_weights also skipped MTP weights with ".mtp." in name, but the names arrive as mtp.layers.* (AutoWeightsLoader strips "model."), so a checkpoint with MTP layers failed with KeyError. Skip the mtp. prefix once, at the top of the loop. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02
force-pushed
the
wu/longcat-scale-norms-on-load
branch
from
October 6, 2026 00:19
f973ba7 to
acb7cf8
Compare
ZJY0516
approved these changes
Oct 6, 2026
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #93009 for commit |
Contributor
CI selector (shadow): 6 test steps (26 jobs) instead of 59 (91 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (6)
Would skip (today's rules run them) (56)
Would add (today's rules do not run them) (3)
AMD mirrors: would skip (55)
AMD mirrors: would add (0)none 2 changed files · base |
Contributor
Author
|
/ci run |
aoshen02
enabled auto-merge (squash)
October 6, 2026 07:04
|
✅ Triggered Buildkite CI #93046 for commit |
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
3 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.
Overview
LongCat-Flash's
load_weightsends with a post-processing sweep. It block-dequantizeskv_b_projintow_kc/w_vc, which nothing reads, and scales the MLA LoRA norms in place behind once-only flags. This PR removes the sweep and scales the norms in their weight loaders. It also fixes the MTP-weight skip in the target loader, without which LongCat-Flash-Chat cannot load. It is a prerequisite for #57733:load_weightsnow only loads.Claims
mainits MLA LoRA norms end up unscaled because the post-load sweep runs before they are loaded (details below).".mtp." in name, butAutoWeightsLoaderstrips themodel.prefix, so the names aremtp.layers.…and loading fails withKeyError: 'mtp.layers.0.hnorm.m.weight'.reload_weightskeeps the MLA LoRA norms scaled. Onmainthe_mla_*_lora_scaledflags survive the reload, so the reloaded norms stay unscaled and accuracy collapses.reload_weightsno longer crashes for FP8 checkpoints with quantized attention (LongCat-Flash-Lite-FP8): the removedblock_dequantofkv_b_projasserts on the reloaded shard.w_kc/w_vcare written and never read;MLAAttention.process_weights_after_loadingderivesW_UK_T/W_UVfromkv_b_projitself. The removed block also leftwunbound for FP8 withweight_block_size=None; no public LongCat checkpoint reaches that branch.Validation
GB200, real weights,
vllm serve,--enforce-eager(see below). GSM8K: 1319 questions, 5-shot, greedy. "Control" is the image's LongCat loader (same sweep and flags asmain) plus only the MTP-prefix fix, so that Chat can load at all.reload_weights→ GSM8Kblock_dequant:assert n_tiles == x_s.shape[0])main's MTP): GSM8K, draft acceptance, after reloadpre-commitpasses.Setup: the vLLM image
nightly-dev-arm64-cu130-ac9126e58aa7plus this change, overlaid file by file with SHA-256 gates. On this imagetorch.compilefails for LongCat on both arms (a Dynamo data-dependent assert in the zero-expert path,moe_runner._maybe_add_zero_expert_output), hence--enforce-eager; the change is load-time only.Details
Why Chat-FP8 scores 0 on
main.AutoWeightsLoadercallsFlashModel.load_weightsonce per contiguous run ofmodel.*names. LongCat-Flash-Chat-FP8 streams 16model.*tensors, thenlm_head, then the remaining 73046, so the sweep runs at the end of the first call: it scales the norms before they are loaded and sets the flags. The second call loads the real norms, and the flags skip the scaling. A debug build shows layer 0'sq_a_layernormsumming to 1892.8 onmainand 3785.6 with this PR (the expected factor √(6144/1536) = 2), andkv_a_layernorm6.74 vs 23.34 (√(6144/512) ≈ 3.46).Norm scaling.
scale_mla_lora_norms_on_loadsets theweight_loaderofq_a_layernorm.weight/kv_a_layernorm.weighttocomposed_weight_loader(default_weight_loader, lambda x: x * factor)when the config enables the scaling (the same pattern as plamo3). Every load, including a reload, writes checkpoint × factor, so no flag is needed. It is applied to both attention modules of each decoder layer and to the MTP block.MTP skip.
FlashModel.load_weightsnow skips names starting withmtp.once, at the top of the loop, instead of three".mtp." in namechecks that never matched.LongCatFlashMTPloads those weights.Not a duplicate. No open PR changes these loaders. #53811 proposes removing the native LongCat implementation in favour of the Transformers backend; it is not merged and has conflicts. This PR fixes the native implementation while it exists.
AI assistance. Written with Claude Code. The submitter reviewed every changed line and ran the tests and GPU checks above.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
🤖 Generated with Claude Code