Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
Journal:
arXiv
Published Date:
Feb 16, 2025
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.
