Repository navigation
[2/2] Add TileLang fused DSA kernels support with THD and CP & Clean up - #5049
Conversation
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Adds a TileLang-backed fused implementation for the experimental DSA attention variant (indexer + SparseMLA + optional streaming sparse-KL), along with a native PyTorch parity test.
Changes:
- Introduces TileLang utility helpers (availability detection, no-op jit, and runtime guard).
- Adds TileLang fused kernels for the sparse MLA forward/backward and indexer forward/backward, plus DSA hook glue.
- Adds a unit test comparing TileLang-backed DSA against a native baseline for outputs and gradients.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit_tests/transformer/experimental_attention_variant/test_dsa_tilelang_native_parity.py | Adds parity test for TileLang vs native DSA outputs/grads across configs. |
| megatron/core/transformer/experimental_attention_variant/ops/tilelang_utils.py | Adds TileLang availability detection + no-op jit + explicit runtime requirement checks. |
| megatron/core/transformer/experimental_attention_variant/ops/tilelang_sparse_mla_fwd.py | Adds TileLang sparse MLA forward kernel + caching + padding/bucketing wrapper. |
| megatron/core/transformer/experimental_attention_variant/ops/tilelang_sparse_mla_bwd.py | Adds TileLang sparse MLA backward + preprocess/postprocess kernels + caching + padding/bucketing wrapper. |
| megatron/core/transformer/experimental_attention_variant/ops/tilelang_indexer_fwd.py | Adds TileLang fused indexer forward kernel + logits cleaning + caching. |
| megatron/core/transformer/experimental_attention_variant/ops/tilelang_indexer_bwd.py | Adds TileLang fused indexer backward kernel + caching + topk padding wrapper. |
| megatron/core/transformer/experimental_attention_variant/ops/tilelang_dsa.py | Adds TileLang-backed DSA hooks (topk, SparseMLA, streaming sparse-KL) + scratch buffer cache. |
| megatron/core/transformer/experimental_attention_variant/ops/sparse_mla.py | Wraps TileLang sparse MLA fwd/bwd in a torch.autograd.Function. |
| megatron/core/transformer/experimental_attention_variant/ops/indexer.py | Wraps TileLang indexer fwd/bwd in an autograd Function and exposes lighting_indexer. |
| megatron/core/transformer/experimental_attention_variant/dsa_tilelang_kernels.py | Exposes TileLang hook entrypoints for the DSA backend integration. |
💡 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: 14807dae90
ℹ️ 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: 953f210933
ℹ️ 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. Already looking forward to the next diff. ℹ️ 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". |
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ffe3a16357
ℹ️ 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: b4804238de
ℹ️ 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".
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30025819135 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30040566352 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30045776196 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30051554542 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30116068044 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30117538144 |
…up (NVIDIA#5049) Signed-off-by: Hollow Man <hollowman@opensuse.org>
…up (NVIDIA#5049) Signed-off-by: Hollow Man <hollowman@opensuse.org>
…up (NVIDIA#5049) Signed-off-by: Hollow Man <hollowman@opensuse.org> (cherry picked from commit 4f65717) Signed-off-by: Robin Zhang <robinz@nvidia.com>
…up (NVIDIA#5049) Signed-off-by: Hollow Man <hollowman@opensuse.org> Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>
What does this PR do ?
Closes #5703
This PR adds the TileLang implementation of the fused DSA operations.
Main changes
megatron/core/transformer/experimental_attention_variant/ops/.Packed THD and CP behavior
The TileLang path uses the metadata and bounds produced by PR 1. When the current TileLang kernel path cannot safely handle a packed THD or CP case, DSA falls back to the unfused implementation instead of changing semantics. This keeps packed-sequence and context-parallel correctness owned by the common DSA path while allowing supported dense cases to use TileLang.
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.