Research

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

← All research
Preprint Third-party arXiv·February 16, 2025

This is third-party work — not authored by or affiliated with DarcStar Technologies.

By Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, Y. X. Wei, Lean Wang, Zhiping Xiao, Yuqing Wang, Chong Ruan, Ming Zhang, Wenfeng Liang, Wangding Zeng

Research

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.

DarcStar commentary

Our perspective on this work and how it informs our research — not part of the original paper.

NSA answers the objection we raise about most sparse attention: that it's bolted on after the fact. The common failure mode is to train a dense model and then sparsify it at inference — and the quality quietly leaks out the seam. NSA is trained sparse from the first step, end to end, so the model learns to work with its own selection rather than around a pattern imposed on it later. That “native, not post-hoc” order of operations is exactly the one we believe in.

The other thing it gets right is that it treats the hardware as a first-class constraint, not an afterthought. The selection is arithmetic-intensity-balanced and block-structured so it actually maps onto the GPU — because a sparsity pattern that saves FLOPs but starves the tensor cores is not a speedup, it's a slower correct answer. And the design is hierarchical: coarse-grained compression to keep global context in view, fine-grained selection for local precision, so it isn't betting everything on a single routing decision.

Where we'd still lean on it is the same place as any learned-selection scheme: the gate decides where the model looks, but nothing yet guarantees every token stays reachable in a bounded number of hops. Trained-in selection makes the usual case strong; it doesn't by itself make the worst case safe. Pair it with a small always-on backbone that guarantees reach and you've closed the gap between “learns to attend well” and “cannot structurally strand what matters.”

As it stands, NSA is the most complete answer on this shelf to how you build content-adaptive sparsity — trained-in and hardware-true — and that's most of the battle.