NVIDIA / NVIDIA/TransformerEngine
[Feature Request] Add Native Sparse Attention (NSA) support
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 3.5k
- Forks
- 831
- Avg merge
- 3d 11h
- Merged PRs (30d)
- 65
Description
Is your feature request related to a problem? Please describe.
Training and inference with long sequences (32K+) using dense attention is prohibitively expensive. TE currently offers sliding window attention as the only sparse alternative, but this uses a fixed local pattern that loses long-range dependencies.
Describe the solution you'd like
Add support for Native Sparse Attention (NSA) from the paper Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention. NSA dynamically selects important tokens via learned compression and top-k block selection, maintaining long-range dependencies while being hardware-efficient.
Describe alternatives you've considered
Using the standalone Triton implementation (native-sparse-attention-triton) - Works but doesn't integrate with TE's FP8 quantization, CP, or fused kernels. Requires maintaining separate code paths.
Additional context
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reading the linked Native Sparse Attention paper and the standalone native-sparse-attention-triton implementation. The issue names no TransformerEngine files or tests; done would require an agreed integration of NSA with FP8 quantization, context parallelism, and fused kernels.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100