Source-linked AI summary
Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
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
TL;DR
Long-context modeling is increasingly important, but vanilla attention creates severe computational and latency costs. NSA addresses this gap with hierarchical, hardware-aligned, natively trainable sparse attention, matching or exceeding Full Attention performance while accelerating training and inference.
Problem
Efficient sparse attention must convert theoretical computation reductions into hardware-aligned speedups across inference stages while supporting trainable operators for lower-cost end-to-end training.
Method
NSA combines hierarchical token compression, blockwise token selection, sliding-window attention, specialized hardware-aligned kernels, and efficient backward operators in a trainable sparse architecture.
Results
NSA matches or exceeds Full Attention across general benchmarks, long-context evaluations, and reasoning tasks while delivering substantial speedups across decoding, forward, and backward stages.
Takeaways & Limitations
NSA demonstrates that sparse attention can support efficient long-context training and inference without sacrificing the reported modeling performance of Full Attention.
Takeaways & Limitations
Phase-restricted sparse methods may leave either prefilling or decoding at computational costs comparable to Full Attention, limiting acceleration for phase-dominated workloads.
Abstract
from arXiv · showhide
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.
1. Introduction
Long-context modeling is important but vanilla attention becomes a major latency bottleneck at 64k contexts. NSA combines hierarchical sparse attention, hardware-aligned kernels, and trainable operators to improve efficiency while maintaining or exceeding Full Attention performance.
- NSA surpasses Full Attention on average across general benchmarks, long-context tasks, and reasoning evaluation despite using sparse attention.
- NSA delivers substantial speedups over Full Attention for 64k-length decoding, forward propagation, and backward propagation.
- Attention computation accounts for 70–80% of total latency when decoding 64k-length contexts.
- NSA processes keys and values through compressed coarse-grained tokens, selectively retained fine-grained tokens, and sliding windows for local context.
- NSA balances arithmetic intensity through Tensor Core- and memory-access-aware blockwise sparse attention, with efficient backward operators enabling end-to-end training.
2. Rethinking Sparse Attention Methods
Existing sparse attention methods often reduce theoretical computation without consistently reducing latency or supporting efficient end-to-end training. NSA is motivated by these inference and training limitations and redesigns sparsity around hardware compatibility and trainability.
- Many sparse attention methods apply sparsity mainly during inference while retaining pretrained Full Attention backbones, limiting native exploitation of sparse attention.
- Phase-restricted methods leave prefilling or decoding near Full Attention cost, limiting acceleration across workloads dominated by either phase.
- Per-head KV-cache selection conflicts with GQA because memory access follows the union of selections within each query-head group.
- Inference-only sparsity can deviate from pretrained optimization trajectories, while discrete selection operations prevent gradients from flowing through token selection.
- Token-granular selection causes non-contiguous KV-cache access, preventing efficient adaptation of FlashAttention and lowering hardware utilization during training.
- NSA addresses inference-efficiency and training-viability limitations through a natively sparse framework combining computational redesign with operator implementation.
3. Methodology
NSA remaps full key-value histories into a compact, query-dependent representation using compression, block selection, and a sliding window. Its design combines hierarchical sparse attention with hardware-aware execution that balances computation and memory access.
- 3.4. Kernel Design: NSA’s kernel design targets different hardware bottlenecks by reducing computation during training and prefilling and memory access during decoding.The implementation uses group-wise KV sharing to eliminate redundant transfers and balances workloads across GPU streaming multiprocessors, improving arithmetic intensity.
- 3.2. Overall Framework: NSA replaces full key-value histories with compact, dynamically constructed representations for each query, maintaining N_t ≪ t.The remapped representations combine compression, selection, and sliding-window strategies through learned gates.
- 3.3.1. Token Compression: Block compression aggregates sequential keys or values into representations that capture block-level information while reducing attention computation.A learnable MLP with intra-block positional encoding maps each block to one compressed key, with analogous compressed values.
- 3.3.2. Token Selection: Token selection ranks attention-relevant blocks and retains tokens from the top-n sparse blocks, using blockwise accesses for hardware-efficient computation.Importance scores can be induced from intermediate compression-token attention scores, and shared scores support consistent selection across grouped query heads.
- 3.3.3. Sliding Window: A dedicated sliding-window branch preserves recent local context while separating local and long-range information sources into independently gated attention branches.Independent keys and values for compression, selection, and window branches help prevent shortcut learning and gradient interference with minimal overhead.
8. Conclusion
NSA is a hardware-aligned sparse attention architecture for efficient long-context modeling. It combines hierarchical token compression with blockwise token selection to accelerate training and inference while maintaining Full Attention performance.
- NSA integrates hierarchical token compression with blockwise token selection in a trainable architecture.
- NSA accelerates both training and inference while maintaining Full Attention performance.
- NSA matches Full Attention on general benchmarks, exceeds it in long-context evaluations, and enhances instruction-based reasoning.
- The architecture achieves measurable computational-latency reductions and significant speedup.
A. Example of AIME Results
The passages present two logarithmic-equation problems and outline algebraic transformations intended to solve them. One worked solution concludes with m+n=131, while another derives expressions for x and y in terms of z.
- A. Example of AIME Results: One problem asks for m+n after expressing a reduced fraction with relatively prime positive integers m and n.The stated conclusion is m=115, n=16, and m+n=131.
- A. Example of AIME Results: A separate substitution sequence reports z=2^-5/12, y=2^-3/8, and x=2^-7/24.The passages then compute x^4, y^3, and z^2 before multiplying them and taking a base-2 logarithm.
- A. Example of AIME Results: The passages describe converting logarithmic equations to exponential form, substituting resulting expressions, and solving the resulting system.They also mention checking an equation and using Cramer’s rule or matrix inversion in one solution path.
- A. Example of AIME Results: The worked algebra derives x=10z and y=5/2 z from the logarithmic system.These expressions are presented after applying logarithm rules and manipulating the two equations.
- A. Example of AIME Results: Another problem gives x,y>1 satisfying log_x(yx)=10 and log_y(x^4y)=10, then asks for xy.The solution begins by converting the logarithmic equations into exponential form.