Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention

Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao +10 more
2/16/2025
cs.CLcs.AIcs.LG

No official code found. Best implementation: fal-ai-community/NativeSparseAttention (63 stars, last commit 2025-02-19). 1 more below.

Abstract

Long-context modeling is crucial for next-generation language models, yet the high computational cost of standard attention mechanisms poses significant computational challenges. Sparse attention offers a promising direction for improving efficiency while maintaining model capabilities. We present NSA, a Natively trainable Sparse Attention mechanism that integrates algorithmic innovations with hardware-aligned optimizations to achieve efficient long-context modeling. NSA employs a dynamic hierarchical sparse strategy, combining coarse-grained token compression with fine-grained token selection to preserve both global context awareness and local precision. Our approach advances sparse attention design with two key innovations: (1) We achieve substantial speedups through arithmetic intensity-balanced algorithm design, with implementation optimizations for modern hardware. (2) We enable end-to-end training, reducing pretraining computation without sacrificing model performance. As shown in Figure 1, experiments show the model pretrained with NSA maintains or exceeds Full Attention models across general benchmarks, long-context tasks, and instruction-based reasoning. Meanwhile, NSA achieves substantial speedups over Full Attention on 64k-length sequences across decoding, forward propagation, and backward propagation, validating its efficiency throughout the model lifecycle.

View on arXivView PDF

Code Implementations(2)

research impl of Native Sparse Attention (2502.11089)

634PythonFeb 19, 20251 years ago
clinno0616/nsaCommunity54%
Apache-2.0

Implementation of the Proposed Method for Deep Seek Native Sparse Attention (NSA, arXiv 2502.11089)

20PythonFeb 23, 20251 years agoApache-2.0
deepseeknative-sparse-attentionnsa

Cite this paper

@article{yuan2025native,
  title  = {Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention},
  author = {Jingyang Yuan and Huazuo Gao and Damai Dai and Junyu Luo and Liang Zhao and Zhengyan Zhang and Zhenda Xie and Y. X. Wei and Lean Wang and Zhiping Xiao and Yuqing Wang and Chong Ruan and Ming Zhang and Wenfeng Liang and Wangding Zeng},
  year   = {2025},
  eprint = {2502.11089},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2502.11089v2}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2502.11089.svg)](https://codeofpaper.com/paper/2502.11089)

Discussion