[1/2] Enable DSA CP, THD with absorbed MLA - #3674
Conversation
There was a problem hiding this comment.
Pull request overview
This PR extends the experimental DSAttention (DSA) path to support context parallelism (allgather CP) and packed THD masking, adds an “absorbed MLA” integration path, and introduces TileLang-based fused kernels (indexer + sparse MLA) with fallbacks and expanded unit coverage.
Changes:
- Enable DSAttention CP allgather masking and packed THD (varlen) masking, including sparse-KL streaming for indexer loss.
- Integrate absorbed-MLA tensor rewrite in
MultiLatentAttentionand route absorbed execution through DSAttention (with optional fused SparseMLA). - Add TileLang fused kernels/interfaces for indexer and sparse MLA, plus extensive new unit tests for CP/THD/absorbed parity and fused plumbing.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit_tests/transformer/test_attention_variant_dsa.py | Adds coverage for CP layout helpers, packed THD varlen masking parity, fused indexer loop behavior, streaming sparse-KL, and absorbed parity. |
| megatron/core/transformer/transformer_config.py | Removes the config-time guard that previously disallowed DSA with context parallelism. |
| megatron/core/transformer/multi_latent_attention.py | Adds absorbed MLA tensor rewrite and passes position_ids/up_v_weight into DSA for CP + absorbed execution. |
| megatron/core/transformer/experimental_attention_variant/ops/tilelang_sparse_mla_fwd.py | New TileLang sparse-MLA forward kernel + Python interface. |
| megatron/core/transformer/experimental_attention_variant/ops/tilelang_sparse_mla_bwd.py | New TileLang sparse-MLA backward kernels + Python interface. |
| megatron/core/transformer/experimental_attention_variant/ops/tilelang_indexer_fwd.py | New TileLang fused indexer forward kernel + logits “cleaning” kernel. |
| megatron/core/transformer/experimental_attention_variant/ops/tilelang_indexer_bwd.py | New TileLang fused indexer backward kernel + Python interface. |
| megatron/core/transformer/experimental_attention_variant/ops/sparse_mla.py | Autograd wrapper around TileLang sparse-MLA forward/backward. |
| megatron/core/transformer/experimental_attention_variant/ops/indexer.py | Autograd wrapper around TileLang indexer forward/backward and a helper for extracting top-k scores. |
| megatron/core/transformer/experimental_attention_variant/dsa.py | Core DSA updates: CP position/masking helpers, varlen bounds, fused top-k + streaming sparse-KL, scratch caching, absorbed sparse attention routing, and updated loss/masking plumbing. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e0d5681007
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
We are changing our review process and marking all open, unlabeled PRs as draft. This change will go in effect starting once #3659 is merged. Moving forward, all PRs will be required to start as draft PRs. If you wish to get your PR merged, mark your PR as “Ready for review”. Read more about the new process at submit.md. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 58df45053a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
291e439 to
0c51ebe
Compare
chtruong814
added
needs-follow-up
chtruong814
added
needs-follow-up
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e04b1b44cc
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8a3e5ae2f1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 774df97500
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
Codex Review: Didn't find any major issues. Delightful! ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
Greptile SummaryThis PR adds Context Parallelism (allgather) and packed THD format support to DeepSeek Sparse Attention (DSA), along with absorbed Matrix-Latent Attention (MLA) integration. It introduces new layout helpers (
Confidence Score: 4/5PR is mostly safe but contains a P1 crash bug in the absorbed MLA shape assertion that will always fire for any TP > 1 configuration. The rope_utils.py fixes are correct, the CP layout and masking logic is sound, and the loss-scale integration is clean. The single P1 issue (incorrect TP shape assertion in absorbed_mla.py) would cause an unconditional crash for any TP > 1 run, which must be fixed before merging. Everything else is P2 or lower. megatron/core/transformer/experimental_attention_variant/absorbed_mla.py — shape assertion on line 852 must be corrected before merging. Important Files Changed
Sequence DiagramsequenceDiagram
participant AML as AbsorbedMLASelfAttention
participant DSA as DSAttention
participant LAY as dsa_layout
participant MSK as dsa_masking
participant KRN as dsa_kernels
AML->>AML: project Q, compress KV
AML->>DSA: forward(q, k, v, packed_seq_params, up_v_weight)
DSA->>LAY: "get_cp_positions_from_layout / build_packed_*"
LAY-->>DSA: query_pos, key_pos / kv_reorder_idx
DSA->>DSA: all_gather KV (CP allgather)
DSA->>MSK: build_dsattention_forward_mask
MSK-->>DSA: index_mask, varlen bounds
DSA->>KRN: run_fused_qk_topk_with_loss (→ None, unfused fallback)
DSA->>DSA: _unfused_absorbed_dsa_fn: QK scores → topk → sparse gather
DSA->>DSA: einsum V up-projection (up_v_weight)
DSA-->>AML: core_attn_out [sq, b, np, v_channels]
AML->>AML: _apply_absorbed_v_up_projection → output projection
Reviews (4): Last reviewed commit: "Address review issue" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
|
@greptileai review |
|
Split draft PRs opened to reduce review scope:
I left this original PR branch unchanged. |
Signed-off-by: Hollow Man <hollowman@opensuse.org>
|
Close as superseded by split PRs: #3674 (comment) |
What does this PR do ?
This PR refactors the DSv3.2 DSA attention implementation so the matrix-absorbed MLA structure and DSA control flow are backend-neutral. It does not add a TileLang or cuDNN kernel; it creates the shared foundation that those kernels can plug into later.
Need together with #3026
Background
DSA uses an indexer to choose the top-k key/value tokens for each query token, then runs sparse attention over those selected tokens. DSv3.2 also uses matrix absorption for MLA: the key up-projection is absorbed into the query side, and the value up-projection is applied after sparse attention. That absorbed structure is specific to this experimental attention path, so keeping it inside the already large generic
multi_latent_attention.pymakes the original MLA implementation harder to maintain.This PR separates that DSv3.2-specific absorbed MLA path from the generic MLA file and reorganizes
dsa.pyaround a small set of overridable fused-kernel hooks. The default behavior remains the PyTorch/MCore implementation; backend-specific kernels are intentionally left to follow-up PRs.Main changes
megatron/core/transformer/experimental_attention_variant/absorbed_mla.py.megatron/core/transformer/multi_latent_attention.py.Packed THD and CP support
[sequence, batch]layout. This PR builds masks from the packed sequence metadata so DSA top-k selection and sparse attention do not attend across packed sequence boundaries.Contribution process
flowchart LR A[Pre-checks] --> B[PR Tests] subgraph Code Review/Approval C1[Expert Review] --> C2[Final Review] end B --> C1 C2 --> D[Merge]Pre-checks
Core 0.8)Code review
The following process is enforced via the CODEOWNERS file for changes into
megatron/core. For changes outside ofmegatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.For MRs into `main` branch
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!
(Step 1): Add PR label
Expert Review(Step 2): Collect the expert reviewers reviews
Expert Reviewlabel when your PR is ready for review.Final Review might get declined if these requirements are not fulfilled.
(Step 3): Final Review
Final Reviewlabel(Optional Step 4): Cherry-pick into release branch
If this PR also needs to be merged into
core_r*release branches, after this PR has been merged, selectCherry-pickto open a new PR into the release branch.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.Merging your PR
Any member of core-adlr and
core-nemowill be able to merge your PR.