[2/2] Wiring cuDNN fused DSA kernels support with THD, CP and IndexShare (GLM5.2)

[HollowMan6]

What does this PR do ?

This PR adds the wiring code for cuDNN fused DSA operations and wires it into the backend-neutral hooks introduced by PR 1. It contains the cuDNN-specific wiring code only

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/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"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
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, the Final Review label 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 Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Adds optional fused cuDNN/FlashMLA kernels for DSA attention, including config-level dependency validation and a native-vs-fused parity test to validate correctness.

Changes:

• Add fused DSA kernel wrappers (FlashMLA forward + cuDNN DSA indexer/attention backward + KL-loss paths).
• Add apply_dsa_kernel_fusion config flag with runtime dependency/SM checks.
• Add GPU parity tests comparing native PyTorch DSA vs the module implementation across backends and loss variants.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

File	Description
tests/unit_tests/transformer/experimental_attention_variant/test_dsa_cudnn_native_parity.py	Adds parity test and native reference implementation for DSA.
megatron/core/transformer/transformer_config.py	Adds apply_dsa_kernel_fusion flag and validates fused-kernel dependencies in config init.
megatron/core/transformer/experimental_attention_variant/dsa_cudnn_kernels.py	Introduces fused DSA kernels wrapper and custom autograd for fused indexer+sparse-attn+loss.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 8327e85bfa

ℹ️ 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".

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

[HollowMan6]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: d3592e090f

ℹ️ 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".

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: db277983ba

ℹ️ 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".

[claude]

LGTM

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 35efdad661

ℹ️ 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".

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 318fc1caf3

ℹ️ 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".

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Bravo.

Reviewed commit: 8138d8716b

ℹ️ 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".

[claude]

LGTM

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Another round soon, please!

Reviewed commit: 9832712fea

ℹ️ 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".

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28962390611

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28965756617

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28978107273

[xiaoxi-wangfj]

@HollowMan6
When selecting the cuDNN backend for DSA, in the CP > 1 case, the indexer still falls back to the naive implementation path:
_indexer_topk_from_score_chunks -> _compute_indexer_scores_chunk_with_global_rows
This makes performance very slow.

Is there a cuDNN-accelerated operator implementation available for the DSA indexer when CP > 1?

[HollowMan6]

@xiaoxi-wangfj The _indexer_topk_from_score_chunks -> _compute_indexer_scores_chunk_with_global_rows fall back should not always happen, common packed-THD CP path still uses cuDNN indexer scoring (loss off / sparse indexer loss on). _compute_indexer_scores_chunk_with_global_rows() path is only for broader support, and it's used only when cuDNN cannot represent the required score masking directly, notably:

• Dense indexer loss (loss_coeff>0, sparse_loss=False) with packed/varlen rows because full dense scores are requested.
• Generic varlen/non-packed layouts outside the optimized packed-CP cases.
• Generic score chunking where later chunks need absolute global query-row positions.

Would you let me know more about your use case?

[xiaoxi-wangfj]

@xiaoxi-wangfj The _indexer_topk_from_score_chunks -> _compute_indexer_scores_chunk_with_global_rows fall back should not always happen, common packed-THD CP path still uses cuDNN indexer scoring (loss off / sparse indexer loss on). _compute_indexer_scores_chunk_with_global_rows() path is only for broader support, and it's used only when cuDNN cannot represent the required score masking directly, notably:

• Dense indexer loss (loss_coeff>0, sparse_loss=False) with packed/varlen rows because full dense scores are requested.
• Generic varlen/non-packed layouts outside the optimized packed-CP cases.
• Generic score chunking where later chunks need absolute global query-row positions.

Would you let me know more about your use case?

@HollowMan6
I checked my run in more detail. My case seems to be a non-packed long-context causal CP setup.

The relevant configuration is:

--context-parallel-size 2
--dsa-indexer-n-heads 32
--dsa-indexer-head-dim 128
--dsa-indexer-topk 1024
--dsa-indexer-loss-coeff 0.001
--dsa-indexer-use-sparse-loss
--dsa-kernel-backend cudnn
--dsa-indexer-topk-freq 4
--dsa-indexer-skip-topk-offset 1

In DSAttention.forward, the inputs are:

attn_mask_type=AttnMaskType.causal
attention_mask=None
position_ids=None
packed_seq_params=None

For SEQ_LEN=8192 and CP=2, one rank gets bounds like:

varlen_starts = [0, 0, 0, ...]
varlen_ends = [2049, 2050, ..., 6144]
key_positions = None

Inside _indexer_topk_bshd, the relevant flags are:

return_scores = False
return_topk_scores = True
use_local_indexer_varlen = False
single_packed_thd_sequence = False

So this is not the dense-loss case, and it is not the packed-THD CP path. It falls into the generic non-packed varlen CP path:

dsa_kernels.run_fused_qk_topk_with_loss
-> FusedQKTopKWithSparseLossFunc
-> _indexer_topk_bshd
-> _indexer_topk_from_score_chunks
-> _compute_indexer_scores_chunk_with_global_rows

I also noticed that cuDNN frontend 1.26.0-rc adds q_causal_offsets to the DSA indexer forward API. This seems to match my case, Is q_causal_offsets intended to support this non-packed CP causal case and avoid the _compute_indexer_scores_chunk_with_global_rows fallback? If so, is there a plan to wire this API into the Megatron-LM DSA path?

[xiaoxi-wangfj]

@xiaoxi-wangfj The _indexer_topk_from_score_chunks -> _compute_indexer_scores_chunk_with_global_rows fall back should not always happen, common packed-THD CP path still uses cuDNN indexer scoring (loss off / sparse indexer loss on). _compute_indexer_scores_chunk_with_global_rows() path is only for broader support, and it's used only when cuDNN cannot represent the required score masking directly, notably:

• Dense indexer loss (loss_coeff>0, sparse_loss=False) with packed/varlen rows because full dense scores are requested.
• Generic varlen/non-packed layouts outside the optimized packed-CP cases.
• Generic score chunking where later chunks need absolute global query-row positions.

Would you let me know more about your use case?

@HollowMan6 I checked my run in more detail. My case seems to be a non-packed long-context causal CP setup.

The relevant configuration is:

--context-parallel-size 2 --dsa-indexer-n-heads 32 --dsa-indexer-head-dim 128 --dsa-indexer-topk 1024 --dsa-indexer-loss-coeff 0.001 --dsa-indexer-use-sparse-loss --dsa-kernel-backend cudnn --dsa-indexer-topk-freq 4 --dsa-indexer-skip-topk-offset 1

In DSAttention.forward, the inputs are:

attn_mask_type=AttnMaskType.causal attention_mask=None position_ids=None packed_seq_params=None

For SEQ_LEN=8192 and CP=2, one rank gets bounds like:

varlen_starts = [0, 0, 0, ...] varlen_ends = [2049, 2050, ..., 6144] key_positions = None

Inside _indexer_topk_bshd, the relevant flags are:

return_scores = False return_topk_scores = True use_local_indexer_varlen = False single_packed_thd_sequence = False

So this is not the dense-loss case, and it is not the packed-THD CP path. It falls into the generic non-packed varlen CP path:

dsa_kernels.run_fused_qk_topk_with_loss -> FusedQKTopKWithSparseLossFunc -> _indexer_topk_bshd -> _indexer_topk_from_score_chunks -> _compute_indexer_scores_chunk_with_global_rows

I also noticed that cuDNN frontend 1.26.0-rc adds q_causal_offsets to the DSA indexer forward API. This seems to match my case, Is q_causal_offsets intended to support this non-packed CP causal case and avoid the _compute_indexer_scores_chunk_with_global_rows fallback? If so, is there a plan to wire this API into the Megatron-LM DSA path?

@HollowMan6
Hi HollowMan, Looking forward to your reply.

[HollowMan6]

Thank you for your info, if you are not using packing, then it will fall back to torch implementation with BSHD and CP>1. I didn't really target BSHD involving CP in this PR, as this PR mainly targets THD-packing (as suggested in the PR title). Also, since typically we can have highly variance length between different sequence, then THD should be more memory friendly and compute efficient, since you are already trying to use CP, I would recommend THD as well to maximize the memory saving here.

For the q_causal_offsets API, I haven't looked into it but I don't currently have plan to add CP support for BSHD.