[2/2] Add TileLang fused DSA kernels support with THD and CP & Clean up - #5049
Conversation
c52cb66 to
51f051b
Compare
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.
51f051b to
716a699
Compare
716a699 to
14807da
Compare
![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".
14807da to
953f210
Compare
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".
953f210 to
d4ed0ef
Compare
|
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". |
db0579b to
ffe3a16
Compare
|
/claude strict-review |
|
@codex review |
Signed-off-by: Hollow Man <hollowman@opensuse.org>
|
Codex Review: Didn't find any major issues. Can't wait for the next one! Reviewed commit: ℹ️ 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". |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Strict review passed — no significant issues found. LGTM
Scope: 11 net-new files adding the TileLang fused DSA backend (indexer fwd/bwd, sparse-MLA fwd/bwd, indexer-loss KL kernels) plus Python wrappers under ops/, and 1 new unit-test file (1181 lines). No existing files are modified, so there is no checkpoint-format, default-value, or API-contract backward-compat surface at risk.
Verified correctness surfaces:
- Autograd gradient ordering in
IndexerFunction.backwardcorrectly re-maps(grad_q, grad_w, grad_k)→(grad_q, grad_k, grad_w, None×5)to match the 8 forward inputs. - SparseMLA small-head padding (pad→16, slice back,
_zero_invalid_heads) prevents gradient leakage into padded heads (unit-tested). - Fused KL target/normalization matches the backend-neutral
dsa_indexer_lossreference, including all-invalid/empty rows; TP async all-reduce double-buffering uses correct ping-pong slot ordering. - Bucket-padding in fwd/bwd interfaces is symmetric (pad → run → slice) with
-1sentinel fill. - Every unsupported case (non-bf16 dtype, unsupported head count, dim≠576/v_channels≠512, THD/CP the kernel can't handle) returns
Noneand falls back to the unfused path — extensively guard-tested.
Compliance: No new parallel_state.get_*_group() reads in megatron/core; the TP group flows in via ProcessGroupCollection, per CLAUDE.md guidance.
Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 0.
Risk: Low. The change is additive, opt-in via dsa_kernel_backend="tilelang", with clean unfused fallbacks and strong unit coverage (including CUDA reference-match tests). Vendored kernel bodies carry clear upstream-provenance comments. Note the CUDA reference tests only exercise coverage where TileLang + a GPU are available in CI.
| seen["use_relu"] = use_relu | ||
| return torch.arange(6, dtype=torch.float32).view(2, 3) | ||
|
|
||
| monkeypatch.setattr(indexer, "indexer_fwd_interface", fake_indexer_fwd_interface) |
There was a problem hiding this comment.
Are the TileLang indexer forward/backward kernels tested anywhere? we should add numerical-parity tests against the pytorch for both forward and backward, similar to the coverage for the cudnn kernels.
There was a problem hiding this comment.
now added real TileLang numerical-parity tests
| from megatron.core.transformer.transformer_config import TransformerConfig | ||
|
|
||
|
|
||
| def run_fused_qk_topk( |
There was a problem hiding this comment.
Is this a pure wrapper around tilelang_dsa.run_fused_qk_topk?
There was a problem hiding this comment.
Yes, It forwards the backend inputs and converts TileLang’s indices-only return value to the common DSA backend contract (topk_indices, topk_length), with topk_length=None
| return sanitized_indices, topk_scores | ||
|
|
||
|
|
||
| def _build_packed_cp_indexer_inputs( |
There was a problem hiding this comment.
The multi-sequence packed-CP branch duplicates the segment-length, cumulative-offset, and source_indices construction in _indexer_topk_multi_packed_cp_thd. Could we extract that backend-neutral layout construction into dsa_layout.py and let TileLang and cuDNN consume the same metadata? The kernel invocation and bounds/remapping can remain backend-specific, but the packed THD layout should have one authoritative implementation.
There was a problem hiding this comment.
I extracted the backend-neutral packed-CP metadata construction into dsa_layout.build_packed_cp_indexer_layout. Both cuDNN and TileLang now consume the same segment lengths, cumulative offsets, key starts, and source-index mapping. Kernel invocation, bounds handling, and result remapping remain backend-specific.
| ) | ||
|
|
||
|
|
||
| def _canonicalize_topk_scores_for_tp_reduce( |
There was a problem hiding this comment.
Mostly to the sort/gather/mask logic in _sort_valid_topk_indices_and_scores_by_index. This ordering is backend-neutral and correctness-critical. The only difference seems to be: TileLang uses -1 sentinels. cuDNN uses a compact prefix described by topk_length. Could we move it into a shared top-k/masking helper that accepts an explicit validity mask.?
There was a problem hiding this comment.
I moved the common sort/gather/mask behavior into dsa_masking.sort_topk_by_index, which accepts an explicit validity mask. TileLang supplies indices >= 0, while cuDNN supplies validity derived from its compact-prefix topk_length, preserving each backend’s representation without duplicating the ordering semantics.
Signed-off-by: Hollow Man <hollowman@opensuse.org>
|
Codex Review: Didn't find any major issues. You're on a roll. Reviewed commit: ℹ️ 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". |
Strict Review SummaryReviewed all 14 changed files (~4.4k additions) for the TileLang fused DSA kernels, tracing dtype/gradient flow, TP reduction, packed-THD/CP layout remapping, and the availability-gated fallback contract. Findings by severity:
Most impactful:
What's good:
Overall risk: LOW–MODERATE. The feature is experimental and fully gated behind availability + input guards with a correct unfused fallback, so blast radius on existing training runs is minimal. No correctness or backward-compatibility blockers found. |
|
/claude strict-review |
There was a problem hiding this comment.
Strict review passed — no significant issues found. LGTM
I traced the correctness-critical paths in this PR:
- Hook signature contract: The
dsa_kernelsdispatcher passesconfigas a keyword only torun_fused_qk_topk_with_loss; the TileLang wrapper'sdel configis intentional interface parity with the cuDNN backend (not a dead argument). The no-loss and absorbed-attention hooks line up with the dispatcher's call sites. - Process groups: No new direct
parallel_state.get_*_group()reads inmegatron/core— TP access threads through the injectedProcessGroupCollection(pg_collection.tp), consistent with repo guidance. - CP packed-THD segment remap:
_build_packed_cp_indexer_inputssanitizes in local segmented key space then remaps back to global int32 indices; the extractedbuild_packed_cp_indexer_layoutis behavior-preserving vs. the deleted inline cuDNN logic (thesk % (2*cp_size)guard moved intact). - Scratch-buffer aliasing: The streaming-KL double buffer (two slots,
chunk_id & 1) copies each target chunk into a slot before the sharedkl_attn_chunk_sumbuffer is reused, and consumes the pending chunk from the other slot — no stale aliasing. - SparseMLA head padding:
_is_supported_sparse_mla_head_countcorrectly requires head_kv be a power-of-two ≥16 to match the kernel's unbounded padded-head indexing, with caller-side pad-then-slice; unsupported shapes decline to the unfused path. - TP score alignment: top-k slots are sorted by key index before slot-wise TP all-reduce.
- Fallback safety: the entire path is gated on
HAVE_TILELANGwith graceful decline (return None) to the authoritative unfused reference implementation, and dtype boundaries (bf16 compute → fp32 accumulation) are explicit throughout.
Test coverage is thorough (~1300 lines exercising wrappers, masking, segment remap, head-mask helpers, scratch cache, and env parsing). No confirmed correctness, performance, or backward-compatibility defects.
Findings by severity — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 0. Overall risk: low (additive, behind a feature gate with an unfused fallback).
|
🔄 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 |
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.