[dev] [DeepSeek-v4] Part 1: Hybrid Attention with CSA and HCA - #4458
Conversation
|
/claude review |
|
/ok to test 503445a |
|
/claude strict-review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| from megatron.core.transformer.transformer_config import TransformerConfig | ||
| from megatron.core.utils import get_pg_size, nvtx_range_pop, nvtx_range_push | ||
|
|
||
|
|
||
| # --------------------------------------------------------------------------- | ||
| # Helper functions for index computation |
There was a problem hiding this comment.
[SUGGESTION Naming] _get_window_topk_idxs_cached and _get_compress_topk_idxs_cached use lru_cache keyed on device_str: str. Since lru_cache uses the string representation of the device, this is correct. However, the cache never gets invalidated. For long training runs where seqlen changes (e.g., curriculum learning, variable-length batching), the cache will grow up to maxsize=8 entries per function. This is fine for the current maxsize but worth noting if the maxsize is increased.
Also, since torch.Tensor objects held in the cache will pin GPU memory, consider using weakref or clearing the cache between epochs if memory becomes a concern.
There was a problem hiding this comment.
Leave it for now. Will consider it when we add THD support.
There was a problem hiding this comment.
Review Summary
CRITICAL: 0 | IMPORTANT: 4 | SUGGESTION: 5
Important Findings
-
[IMPORTANT Correctness] Grouped output projection broken with TP > 1 (
deepseek_v4_hybrid_attention.py:188-198):linear_o_group_projis allocated with the fullo_groups * o_lora_rankrows but viewed as(o_local_groups, o_lora_rank, -1). When TP > 1, the inferred last dimension becomestp * group_proj_in_sizeinstead ofgroup_proj_in_size, causing the einsum to contract over a mismatched dimension. The weight must be either allocated per-TP-shard or sliced before the einsum. -
[IMPORTANT Correctness] KV layernorm shape mismatch with TP > 1 (
deepseek_v4_hybrid_attention.py:489-507):linear_kv_up_projis aColumnParallelLinearwithgather_output=False, outputtingv_head_dim / tpper rank. Butkv_layernormis created withhidden_size=v_head_dim(full). The learnable parameters and normalization stats will mismatch when TP > 1. -
[IMPORTANT Correctness]
clip_qkassumes MLA multi-head weight layout (deepseek_v4_hybrid_attention.py:863-887):_clip_kv_proj_weightreshapes assuming(num_heads, qk_head_dim + v_head_dim, ...), but DSv4's single-headlinear_kv_up_projhas shape(v_head_dim, hidden_size). This will crash or produce wrong results ifqk_clipis ever enabled. -
[IMPORTANT Compatibility] Missing CLI arguments (
arguments.py): Only--csa-compress-ratiosis added.--csa-window-size,--csa-compress-rotary-base, and--csa-dense-mode(mentioned in the PR description) are missing from the argument parser.
Notable Design Changes
-
Triton RoPE kernel output layout change: The fused MLA RoPE kernels (
_mla_rope_fwd_inplace_kernel,_mla_rope_fwd_kv_split_kernel) now preserve the interleaved layout on write (previously they de-interleaved). The unfused path inrope_utils.pywas updated to match. This is backward-compatible for attention computation (QK^T dot product is permutation-invariant for matching Q/K layouts), but changes the intermediate tensor layout visible to downstream code. Backward-compat aliases are provided. -
Inverse RoPE: The
inverse=Trueflag correctly negates sin in both forward and backward kernels. The Jacobian analysis confirms the backward gradient computation is correct for the inverse case.
Overall Assessment
Risk level: Medium. The core algorithm (CSA + HCA hybrid attention, inverse RoPE, compressed sparse attention with learned indexing) is well-structured and the numerical correctness for TP=1 looks solid. The Triton kernel changes are consistent across forward/backward, Q/KV paths. The DSA indexer loss extensions (custom causal mask, row-validity masking) are correctly implemented.
The main risks are the TP > 1 bugs in the grouped output projection and KV layernorm — these would produce silently wrong results or crash at TP > 1. Since the PR targets dev and notes it's under development, these may be known limitations, but they should be tracked.
The test coverage is good: unit tests for helper functions, Compressor, CSAIndexer, CompressedSparseAttention (window-only, compressed, dense mode), and the full DSv4HybridSelfAttention forward/backward paths.
|
/ok to test 6f5b92f |
|
Hello @hxbai I'm also very interested in implementing sparse attention, and I'd like to know if CP/SP support will be added in the future? |
|
/ok to test 8470a3b |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25173867126 |
### PR Category <!-- One of [ Train | Inference | Compress | Serve | RL | Core | Hardware | CICD | Tools | Others ] --> [Train] Most of codes are copied from Megatron-LM Dev branch. The dev branch is different with main branch or release version. Megatron LM PR: DeepSeek-V4: NVIDIA#4458 NVIDIA#4481 NVIDIA#4518 mHC: NVIDIA#2943 ### PR Types <!-- One of [ User Experience | New Features | Bug Fixes | Improvements | Performance | Breaking Change| Deprecations | Test Case | Docs | Others ] --> [New features] ### PR Description <!-- Describe what you’ve done --> Add DeepSeek V4 model into FlagScale and Megatron-FL Supported: 1. CSA and HCA 2. Hash Router 3. mHC 4. Engram(optional) Unsupported: 1. Sqrtsoftpuls router score function. ✅ 2. mHC recompute. ✅ 3. Overlap_grad_reduce and overlap_param_gather when Zero 1. ✅ 4. Any infra optimizations. ### NOTE: This is only a draft pr, please reivew to give more suggestions. such as: 1. File structure. - All modules are moved into Megatron-FL ### Next plan: 1. Distributed training. ✅ 3. Muon optimizer with Zero 1 adaptation. 🚧 4. Low precision is out of scope of this pr, limited by resource. 5. Maybe context parallel for sparse attention. 6. Welcome to give more suggestions. --------- Co-authored-by: Hongxiao Bai <hongxiaob@nvidia.com> Co-authored-by: Yuzhong Wang <yuzhongw@nvidia.com>
### PR Category <!-- One of [ Train | Inference | Compress | Serve | RL | Core | Hardware | CICD | Tools | Others ] --> [Train] Most of codes are copied from Megatron-LM Dev branch. The dev branch is different with main branch or release version. Megatron LM PR: DeepSeek-V4: NVIDIA/Megatron-LM#4458 NVIDIA/Megatron-LM#4481 NVIDIA/Megatron-LM#4518 mHC: NVIDIA/Megatron-LM#2943 ### PR Types <!-- One of [ User Experience | New Features | Bug Fixes | Improvements | Performance | Breaking Change| Deprecations | Test Case | Docs | Others ] --> [New features] ### PR Description <!-- Describe what you’ve done --> Add DeepSeek V4 model into FlagScale and Megatron-FL Supported: 1. CSA and HCA 2. Hash Router 3. mHC 4. Engram(optional) Unsupported: 1. Sqrtsoftpuls router score function. ✅ 2. mHC recompute. ✅ 3. Overlap_grad_reduce and overlap_param_gather when Zero 1. ✅ 4. Any infra optimizations. ### NOTE: This is only a draft pr, please reivew to give more suggestions. such as: 1. File structure. - **All modules are moved to Megatron-FL. Only model_builder is left in Flagscale.** - Delete Engram related CI or not? ### Next plan: 1. Distributed training. ✅ 3. Muon optimizer with Zero 1 adaptation. 😢 4. Low precision is out of scope of this pr, limited by resource. 5. Maybe context parallel for sparse attention. 6. Welcome to give more suggestions. --------- Co-authored-by: zhaoyingli <86812880+zhaoyinglia@users.noreply.github.com>
Add the unfused compressed sparse attention algorithm, inert configuration fields, and focused algorithm coverage without enabling a model variant. Reconstructed from NVIDIA#4458 and its corrections as part of the frozen feature set in NVIDIA#5795. Signed-off-by: Deyu Fu <Deyu.Foo@gmail.com>
Enable SBHD-only DSv4 Hybrid attention orchestration, model-variant validation, backend-explicit specs, and MTP layer-ratio plumbing. Reconstructed from NVIDIA#4458, NVIDIA#4518, NVIDIA#5018, and NVIDIA#5526 as part of the frozen feature set in NVIDIA#5795. Signed-off-by: Deyu Fu <Deyu.Foo@gmail.com>
5 participants
What does this PR do ?
We will create several PRs to functionally support DeepSeek-v4 training. This is the first one.
Add DeepSeek-v4 Hybrid Attention with CSA and HCA.
We reuse some arguments from DSA to avoid duplication.
DeepSeek-V4 training has several phases, and here are the setting examples (use DSv4-Flash size):
--experimental-attention-variant dsv4_hybrid--csa-window-size 128--csa-compress-ratios ([0,0,4]+[128,4]*20+[0])--csa-compress-rotary-base: 40000--dsa-indexer-n-heads 64--dsa-indexer-head-dim 128--dsa-indexer-topk 512--csa-dense-mode--experimental-attention-variant dsv4_hybrid--csa-window-size 128--csa-compress-ratios ([0,0,4]+[128,4]*20+[0])--csa-compress-rotary-base: 40000--dsa-indexer-n-heads 64--dsa-indexer-head-dim 128--dsa-indexer-topk 512--experimental-attention-variant dsv4_hybrid--csa-window-size 128--csa-compress-ratios ([0,0,4]+[128,4]*20+[0])--csa-compress-rotary-base: 40000--dsa-indexer-n-heads 64--dsa-indexer-head-dim 128--dsa-indexer-topk 512--dsa-indexer-use-sparse-lossIssue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.
For MRs into `dev` branch
The proposed review process for `dev` branch is under active discussion.MRs are mergable after one approval by either
eharper@nvidia.comorzijiey@nvidia.com.