arXiv. arXiv:2502.11089 (Submitted 16 Feb 2025, last revised 27 Feb 2025, v2)
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.
Summary
The paper starts from a practical gap: theoretical sparsity often fails to reduce wall-clock time, and many methods apply sparsity only at inference. NSA specifies a three-branch hierarchy—compressed blocks, selected blocks, and a sliding window—and trains that hierarchy from pretraining. Blockwise selection is a hardware choice: contiguous loads coalesce, and heads that share a KV cache share the same selection. On a 27B-total / 3B-active backbone trained with 260B tokens, average scores are reported as matching or exceeding full attention. Kernel measurements on A100 report up to 9.0 × forward and 6.0 × backward speedup at 64k, with an estimated 11.6 × decoding speedup from KV traffic. Those figures are not a universal end-to-end product speedup.
Detailed reading
Why theoretical sparsity is not latency
Despite these promising strategies, existing sparse attention methods often fall short in practical deployments. Many approaches fail to achieve speedups comparable to their theoretical gains; moreover, most methods lack effective training-time support to fully exploit the sparsity patterns of attention.
Full causal attention compares each query with all earlier keys. Training and prefilling scale roughly as N². Long decoding repeatedly reads a large KV cache. Reducing the number of mathematical pairs is not equivalent to reducing latency. The abstract already states the hierarchical strategy; the remainder of this section specifies why that hierarchy must be trained, not applied only at inference.
Two bottlenecks and three branches
Arithmetic intensity—FLOPs per byte moved—determines whether a kernel is compute-bound or memory-bound.
This leads to different optimization goals — reducing computation cost during training and prefilling, while reducing memory access during decoding.
A method that sparsifies only decoding, or only prefilling, leaves the other phase at full-attention cost. If each head independently selects scattered tokens, GQA or MQA may still read the union of those selections, so memory traffic need not fall with FLOP count. NSA therefore selects contiguous blocks, includes them during training, and executes them with kernels suited to Tensor Cores.
The framework processes input sequences through three parallel attention branches: For a given query, preceding keys and values are processed into compressed attention for coarse-grained patterns, selected attention for important token blocks, and sliding attention for local context.
Compressed attention maps consecutive key/value blocks through a learnable function and provides low-resolution coverage of the full context. Selected attention reuses the coarse scores to rank original blocks, retains the top blocks, and attends to their tokens. Sliding-window attention retains recent tokens so that local patterns do not monopolize learning. Learned gates mix the three outputs. The paper keeps the remapped size N_t ≪ t.
Quality and speed
Pretraining uses a 27B-total / 3B-active MoE backbone and 260B tokens. Average scores are reported as matching or exceeding full attention on general, long-context, and chain-of-thought evaluations.
NSA achieves perfect retrieval accuracy across all positions in 64k-context needle-in-a-haystack (Kamradt 2023) test.
Needle retrieval is informative and still simpler than synthesizing many weak signals across documents.
On A100 GPUs, Triton implementations are compared with FlashAttention-2.
As shown in Figure 6, our NSA achieves progressively greater speedups as context length increases, up to 9.0 × forward and 6.0 × backward speedup at 64k context-length.
The paper further estimates up to 11.6 × decoding speedup at 64k from KV traffic under a memory-bound model. That figure is not a universal end-to-end product speedup. Sampling, MoE routing, communication, batch size, and shorter contexts all reduce attention’s share of total time.
Evidence boundary
FlashAttention asks how to compute all pairs with less movement. NSA asks which structured subset should exist and whether hardware can execute it. Compression and window branches may use FlashAttention-style kernels; the selected branch requires a block-sparse kernel. Short contexts can erase the advantage. A dense pretrained checkpoint does not obtain NSA by exchanging a single module. An efficient attention design, on this account, must specify four quantities together: the information hierarchy, the learning path, the memory-access pattern, and the GPU kernel. Sparsity stated without those four is incomplete.