Repository navigation
[Attn] Rename CP backends and relocate model attention modules - #5445
Merged
Merged
Conversation
meta-cla
Bot
added
the
CLA Signed
acisseJZhong
added this pull request to stack #5446
October 16, 2026 07:49
acisseJZhong
force-pushed
the
rename_attention_backends
branch
from
October 16, 2026 08:07
a88ad32 to
cc9f290
Compare
acisseJZhong
marked this pull request as ready for review
October 16, 2026 08:07
acisseJZhong
requested review from
IvanKobzarev,
SherlockNoMad,
aditvenk,
fegin,
felipemello1,
pzhan9,
sanketpurandare,
tianyu-l,
wconstab and
wwwjn
as code owners
October 16, 2026 08:07
tianyu-l
approved these changes
Oct 16, 2026
acisseJZhong
force-pushed
the
rename_attention_backends
branch
from
October 16, 2026 22:21
cc9f290 to
dd8f415
Compare
acisseJZhong
added a commit
that referenced
this pull request
Oct 19, 2026
…imi K3 and dpskV3 (#4760) Stack order: #5445 -> #4760. Naming and file moves are in #5445; this PR focuses on MLA CP optimization. **The original optimization idea was inspired by this [MLA Ulysses CP proposal](https://gist.github.com/fegin/bd9bd9e9147ae1b11375dfb68815f78a).** Shape suffixes: `T` tokens, `H` heads, `L` KV latent, `N` non-positional key, `R` shared key, `V` value; `K=N+R`, `P=N+V`, `F=K+P`. Suffixes describe logical dimensions; collectives change their sharding. - **K/V all-gather CP** - Before: all-gather materialized K and V separately. K repeats `k_shared_TR` across attention heads. - After: all-gather normalized compressed KV latents and head-shared `k_shared_TR`, then apply `wkv_b` and materialize K/V locally. This trades repeated projection work for lower communication volume. ```python # Before: materialize K before communication. kv_THP = reshape_per_head(wkv_b(kv_c_normed_TL)) k_nope_THN, v_THV = unpack(kv_THP) k_THK = concat(k_nope_THN, expand_per_head(k_shared_TR)) k_THK = all_gather(k_THK) v_THV = all_gather(v_THV) # After: communicate compressed latents before projecting K/V. kv_c_normed_TL = all_gather(kv_c_normed_TL) k_shared_TR = all_gather(k_shared_TR) kv_THP = reshape_per_head(wkv_b(kv_c_normed_TL)) k_nope_THN, v_THV = unpack(kv_THP) k_THK = concat(k_nope_THN, expand_per_head(k_shared_TR)) ``` - **Ulysses CP** - Before: use three all-to-alls for materialized Q, K, and V. K repeats `k_shared_TR` across attention heads. - Current: all-gather headless `k_shared_TR`, then use one all-to-all for `q_THK` + packed `kv_THP`. ```python # Before: materialize K, then move Q, K, and V from token to head sharding. kv_THP = reshape_per_head(wkv_b(kv_c_normed_TL)) k_nope_THN, v_THV = unpack(kv_THP) k_THK = concat(k_nope_THN, expand_per_head(k_shared_TR)) q_THK = all_to_all_token_to_head(q_THK) k_THK = all_to_all_token_to_head(k_THK) v_THV = all_to_all_token_to_head(v_THV) out_THV = attention(q_THK, k_THK, v_THV) out_THV = all_to_all_head_to_token(out_THV) # Current: combine q_THK and kv_THP in one all-to-all. kv_THP = reshape_per_head(wkv_b(kv_c_normed_TL)) k_shared_TR = all_gather(k_shared_TR) packed_q_kv_THF = all_to_all_token_to_head( concat(q_THK, kv_THP) ) q_THK, kv_THP = split(packed_q_kv_THF) k_nope_THN, v_THV = unpack(kv_THP) k_THK = concat( k_nope_THN, expand_per_head(k_shared_TR), ) out_THV = attention(q_THK, k_THK, v_THV) out_THV = all_to_all_head_to_token(out_THV) ``` - **Shared MLA backend and checkpoints** - DeepSeek V3, Kimi K2.7, and Kimi K3 share the compact MLA backend. `MLAAttention` owns `wkv_b` and delegates to a Q/K/V attention kernel; common all-gather communication lives in `distributed/context_parallel.py`. - HF adapters map directly to `attention.inner_attention.wkv_b.weight`, replacing state-dict renaming hooks. HF keys remain compatible. Native checkpoints use the nested key; old `attention.wkv_b.weight` keys are no longer automatically remapped. - **Sharding and backend compatibility** - The KV projection explicitly asserts its CP layout: replicated after latent all-gather, token-sharded before Ulysses all-to-all. The CP transform copies the nested projection config when selecting the backend; the local head reshape preserves that layout. - `fused_mla` rejects CP backends at model construction because its fused forward bypasses the compact MLA communication wrapper. Fused MLA with CP is not supported; remove the override to use the configured CP backend. - GraphTrainer precompile delegates metadata construction to the Decoder; batch-invariant settings traverse Flex configs directly, including those nested inside MLA. - The TorchTitan vLLM wrapper currently rejects MLA. Kimi K3 MLA inference parity and RL integration tests are disabled until support is restored. ## Tests After rebase/split at [`4d5064c40c`](4d5064c): 197 CPU tests and the four-GPU MLA sharding test passed. Non-Pyrefly lint passed; Pyrefly has the same 53 diagnostics as main. GraphTrainer validation remains blocked by the installed PyTorch missing `_is_cacheable_effect`. ### Numerical comparison [`3e6c5c2b1a`](3e6c5c2): DeepSeek V3 and text-only Kimi K3 debug models, 2 H100s, CP=2, BF16/Flex, seed 42, deterministic, 256 tokens, 100 steps. Identical initial weights and a fixed synthetic batch; AdamW LR `1e-3`, gradient clipping `1.0`. Relative differences between the materialized-K/V reference and compact MLA on the same revision. Loss is global mean cross-entropy; grad norm is FP32 L2 before clipping. | Model | CP backend | Loss Δ, step 10 | Grad norm Δ, step 10 | Loss Δ, step 100 | Grad norm Δ, step 100 | | --- | --- | ---: | ---: | ---: | ---: | | DeepSeek V3 | K/V all-gather | 0.067% | 0.371% | 0.541% | 0.116% | | DeepSeek V3 | Ulysses | 0.103% | 0.218% | 0.056% | 0.128% | | Kimi K3 | K/V all-gather | 2.321% | 11.125% | 0.462% | 0.628% | | Kimi K3 | Ulysses | 0.262% | 2.466% | 0.363% | 0.322% | All runs completed 100 steps with finite loss and gradients. First-step logits and loss are **bitwise identical** in all four comparisons, but gradients are not: first-step relative L2 gradient differences range from **0.1475% to 0.2716%**. Changed BF16 gradient accumulation leads to different subsequent optimizer updates. This fixed-batch comparison does not establish convergence on a training dataset. ### Performance Retested at [`3e6c5c2b1a`](3e6c5c2) on MAST: DeepSeek-V3 16B, **32 H100 GPUs**, CP=8/FSDP=4/EP=8, BF16 mixed precision, Flex attention, C4 test data, CUDA graphs disabled, 100 steps. PyTorch `2.15.0a0+git8646abd`. The baseline materializes K/V before communication on the same revision and runtime. TPS/GPU is the **rank-0 mean +/- step SD over steps 21-99**, from logged throughput. One run per case. #### K/V all-gather CP 32,768 tokens per DP rank; HeadTail load balancing. | Path | Input collectives | TPS/GPU (mean +/- SD) | vs. baseline | | --- | --- | ---: | ---: | | [Before: materialized K/V](https://www.internalfb.com/msl/studio/runs/mast/tt-pr4760-3e6c5c2-all-gather-before-32-jessicazhong-r2) | 2 all-gathers | 3817.5 +/- 1021.8 | - | | [Current: compressed KV latent + shared key](https://www.internalfb.com/msl/studio/runs/mast/tt-pr4760-3e6c5c2-all-gather-current-32-jessicazhong-r2) | 2 all-gathers | 3970.8 +/- 1009.7 | +4.02% | #### Ulysses CP 8,192 tokens per DP rank; no CP load balancing. | Path | Input collectives | TPS/GPU (mean +/- SD) | vs. baseline | | --- | --- | ---: | ---: | | [Before: separate Q, K, V](https://www.internalfb.com/msl/studio/runs/mast/tt-pr4760-3e6c5c2-ulysses-before-32-jessicazhong-r2) | 3 all-to-alls | 1000.6 +/- 289.5 | - | | [Current: combined Q + packed KV, shared key](https://www.internalfb.com/msl/studio/runs/mast/tt-pr4760-3e6c5c2-ulysses-current-32-jessicazhong-r2) | 1 all-to-all + 1 all-gather | 1098.7 +/- 303.9 | +9.80% | [Archived notes.](https://www.internalfb.com/phabricator/paste/view/P2549350949 )
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.
Stack: this PR -> #4760.
Separate naming and file organization from the MLA CP optimization in #4760: rename CP/vLLM backends consistently with
InnerAttention, move DeepSeek/Kimi MLA classes into model-specificattention.pyfiles, and update callers, including Muse Spark. Computation, parameter names, and initialization order are unchanged.Validation: 163 CPU tests passed. DeepSeek V3 and text-only Kimi K3 debug models match main bitwise over 10 BF16/Flex training steps (seed 42, deterministic): initial/final weights, outputs, losses, gradients, and grad norms. Non-Pyrefly lint hooks passed; Pyrefly reports the same 53 diagnostics as main. GraphTrainer validation is blocked by the installed PyTorch missing
_is_cacheable_effect.