Feature Request: Adding NSA and other sparse attention mechanisms

[csking101]

Is your feature request related to a problem? Please describe.
Sparse attention mechanisms like Native Sparse Attention (NSA) (from DeepSeek's "Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention") address the quadratic attention problem by reducing attention complexity while preserving model quality, but Megatron-LM currently lacks support for these methods.

Describe the solution you'd like
I will extend the existing DSA implementation, which is under megatron/core/transformer/experimental_attention_variant/dsa.py. I'll add a design where the token compression, top-k token selection and sliding window features can be configured per layer.

Describe alternatives you've considered
There are external libraries doing the same, however we want to have it in-house.

Additional context

• The NSA paper: arxiv.org/abs/2502.11089
• DeepSeek V3.2 reference implementation: github.com/deepseek-ai/DeepSeek-V3.2-Exp
• The existing experimental DSA implementation in this repo (megatron/core/transformer/experimental_attention_variant/dsa.py) already implements key components: top-k indexing, KL-divergence indexer loss, and Hadamard rotation. This provides a strong foundation to build upon.

[csking101]

Hey @sbhavani , can I work on this issue, I have created the same, after consulting the roadmap for Megatron Core.

[Phlip79]

@csking101 DSA is already fully implemented and integrated. Do you have a specific use-case for NSA, that isn't satisfied by DSA?

[csking101]

@Phlip79 , actually DSA does a subset of what NSA aims to do currently. The current implementatoin provides a learned indexer, that scores all the KV toeksn and takes the top-k relevant ones for each query, trained by KL divergence loss against true attention scores. This is required by NSA.

However, NSA would extend the same by adding two things:

• A token compression section, which compresses blocks of KV tokens into summary representations.
• A sliding window section, which guarantees attention to the nearby tokens

All 3 branches are required for optimal performance, the learned per-query gate weighs these contributions and proceeds with the calculations. Basically, DSA isn't complete.

[sbhavani]

@csking101 thank you for volunteering to implement NSA! We have a kernel in cudnn-frontend that we are evaluating: https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/native_sparse_attention. We are right now running convergence tests on an 8B model and will share a public branch soon.

In the meantime, please feel free to share any thoughts you have on this kernel or your ideas for implementation. We also have many more open issues if you'd like to work on something else!

[csking101]

@csking101 thank you for volunteering to implement NSA! We have a kernel in cudnn-frontend that we are evaluating: https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/native_sparse_attention. We are right now running convergence tests on an 8B model and will share a public branch soon.

In the meantime, please feel free to share any thoughts you have on this kernel or your ideas for implementation. We also have many more open issues if you'd like to work on something else!

Hey @sbhavani , thanks for the reply, I was unaware about this since I am new to the ecosystem😅. I will check it out, thanks! I will look at other open issues as well.

Edit: I had a look at the implementation in the cudnn library. I had a couple of questions:

• In the NSA paper, the output is a weighted sum of the three attention branches. The weights for the sum, are given by a light-weight MLP, that is trained on a per-query basis. Is it not included in this library because it doesn't call for a custom CUDA kernel?
• As a follow up to the previous question, since you haven't implemented gating, are you planning to directly implement the same, when you write the central API for NSA?
• On the other hand, all these kernels are implemented for a forward pass, are you planning on adding kernels for a backward pass as well? The gating feature, in this instance would require training I believe.
• Is my understanding right, that the cudnn code is used as an accelerator for certain compute heavy parts of the operation during inference time?

Edit 2: Also @sbhavani , if you have any issues that haven't been assigned, please point me to the same, I am looking to learn more and contribute!