Repository navigation
[Model][DeepSeek-V4.1] Permute mega-attention wq_b / wo_a while loading - #60064
Merged
aoshen02 merged 4 commits intoOct 6, 2026
Merged
Conversation
FlashMLA mega attention (the SM100 default DSv4.1 attention) permutes wq_b rows and wo_a columns, plus their MXFP8 scales, into the kernel's chunk-interleaved layouts. That ran as a post-load step (finalize_loaded_weights, via finalize_mega_attn_weights) guarded by a _fused_layouts_ready flag, and only produced the right result when it ran before the per-layer quant hooks, which is what calling it from inside load_weights happens to do. Whenever it runs after the quant hooks, the tensors are already in the kernel layout (CuteDSL transposes wq_b to [K, N]; the wo_a BMM hook reshapes to 3D and packs scales), so permute_wq_b_ derives the wrong head_dim and q_fused_permutation fails with "IndexError: index 1280 is out of bounds for dimension 0 with size 1280". This already happens for DSv4.1 dummy load on main, where load_weights is never called and the model-level process_weights_after_loading runs after the quant hooks, and it would hit real loads too once model-level finalization moves to the loader hook (vllm-project#57733). The flag also meant a second load_weights pass (e.g. a weight update) left the new wq_b / wo_a unpermuted. Make the permutation part of loading the checkpoint tensor instead: wrap the weight loaders of wq_b.weight / wq_b.weight_scale / wo_a.weight / wo_a.weight_scale with composed_weight_loader so each loaded shard is gathered into the fused layout right after it is written. Each of these tensors is loaded by a single weight_loader call (no stacked mapping), so the full local shard is present when the permutation runs. The quant hooks then see already-permuted tensors, exactly as main's real load does, so kernel-visible tensors are unchanged. finalize_loaded_weights, _fused_layouts_ready, finalize_mega_attn_weights and DSpark's _finalize_attn are removed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02
requested review from
AndreasKaratzas,
WoosukKwon,
mgoin,
tlrmchlsmth,
yewentao256 and
zyongye
as code owners
October 5, 2026 14:07
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
mergify
Bot
added
deepseek
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai>
zyongye
approved these changes
Oct 6, 2026
Contributor
Author
|
/ci run |
aoshen02
enabled auto-merge (squash)
October 6, 2026 04:11
|
✅ Triggered Buildkite CI #93026 for commit |
Contributor
CI selector (shadow): 21 test steps (56 jobs) instead of 33 (46 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 (21)
Would skip (today's rules run them) (24)
Would add (today's rules do not run them) (12)
AMD mirrors: would skip (29)
AMD mirrors: would add (1)
5 changed files · base |
6 tasks done
This was referenced Oct 9, 2026
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
DeepSeek-V4.1's FlashMLA mega attention (the SM100 default) permutes
wq_brows andwo_acolumns into the kernel's layout as a post-load step. That step is only correct when it runs before the linear quant hooks. This PR does the permutation in the weight loaders instead, so it no longer depends on post-load order. It is a prerequisite for #57733.Claims
IndexError: index 1280 ...; TP4: a size-mismatchRuntimeError).reload_weights(layerwise reload) for mega attention. Onmainthe reloadedwq_b/wo_aare written in checkpoint order and_fused_layouts_readyskips the permutation, so accuracy collapses. Layerwise reload replays the loaders recorded at model init, so with this PR the reloaded shards are permuted again.mainon all 4 ranks.finalize_loaded_weights,_fused_layouts_ready,finalize_mega_attn_weightsand DSpark's_finalize_attn(+62 / −94).Validation
DeepSeek-V4.1-Flash, real weights,
vllm serve, GB200, mega attention (FLASHMLA_MEGA_ATTN_DSV41, confirmed in the logs). GSM8K: 1319 questions, 5-shot, greedy. MegaMoE unless noted.mainreload_weights, GSM8K--load-format dummyIndexError: index 1280 is out of bounds for dimension 0 with size 1280--load-format dummyRuntimeError: The size of tensor a (327680) must match the size of tensor b (1280)With MegaMoE,
reload_weightscollapses on both arms for a separate reason: the model-level post-load step (MegaMoE and mHC state) is not rerun after the layerwise reload finishes. #57733 reruns it; with #57733 and this PR, the same TP4 MegaMoE reload measures 0.901 → 0.898, while #57733 without this PR measures 0.900 → 0.002. The DP4 reload is not measurable on any arm: layerwise reload re-materializes the CPU-offloaded Engram table on the GPU and runs out of memory (pre-existing).tests/kernels/test_dsv41_mega_attn_layouts.py: 2 passed. It loadswq_b/wo_aand their MXFP8 scales through the permuting loader and checks the permuted GEMMs against a reference that permutes the activation instead.pre-commitpasses.Setup: the vLLM image
nightly-dev-arm64-cu130-ac9126e58aa7plus this commit, overlaid file by file with SHA-256 gates.Details
Root cause.
permute_wq_b_deriveshead_dimfromweight.shape[0]. After the CuteDSL linear hook,wq_bis transposed to[K, N], sohead_dimis wrong andq_fused_permutationindexes out of bounds. Thewo_aBMM hook likewise reshapeswo_ato 3D and packs its scales. A real load calls the permutation from insideload_weights, before the hooks, so it works there. A dummy load never callsload_weights. The model-levelprocess_weights_after_loadingthen runs after the hooks and crashes.Fix.
permute_on_load(param, perm, dim)wraps the parameter'sweight_loaderwithcomposed_weight_loader, so each loaded shard is gathered into the fused layout right after it is written. The wrapping happens inDeepseekV4MegaAttnAttention.__init__, forwq_b.weight,wq_b.weight_scale,wo_a.weightandwo_a.weight_scale. Each of these tensors is loaded by a singleweight_loadercall (no stacked mapping), so the whole local shard is present when the permutation runs. The quant hooks then see already-permuted tensors, exactly asmain's real load does.The permutations are computed from the module's own
n_local_heads/head_diminstead of tensor shapes.n_local_heads // n_local_groups == WV_GROUP_SIZEis already enforced in__init__.Not a duplicate. I searched open PRs for
mega attention wq_b,DeepSeek-V4.1 dummy load,fused_layout permute,finalize_loaded_weightsandV4.1 mega attn. #57977 touchesflash_mla_mega_attn.pyfor MXFP8 activations but does not change the weight permutation.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