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.
核心贡献 · Key contributions
提出 NSA，一种原生可训练的稀疏注意力机制，采用动态分层令牌压缩与选择。 Proposes NSA, a natively trainable sparse attention with dynamic hierarchical token compression and selection.
通过块状内存访问和以组为中心的内核设计实现硬件对齐的加速。 Achieves hardware-aligned speedups via blockwise memory access and group-centric kernel design.
支持稀疏注意力的端到端训练，在不损失性能的情况下减少预训练计算量。 Enables end-to-end training of sparse attention, reducing pretraining computation without performance loss.
在长上下文任务和思维链推理基准上优于全注意力。 Outperforms full attention on long-context tasks and chain-of-thought reasoning benchmarks.
在 64k 长度序列上实现高达 11.6 倍解码加速和 9.0 倍前向加速。 Demonstrates up to 11.6x decoding speedup and 9.0x forward speedup on 64k-length sequences.
集成滑动窗口、压缩和选择分支，并通过学习门控实现稳定训练。 Integrates sliding window, compression, and selection branches with learned gating for stable training.
局限 · Limitations
评估仅限于 27B 参数的 MoE 模型，未验证向更大模型的扩展性。 Evaluation limited to 27B-parameter MoE model; scalability to larger models not verified.
需要自定义 Triton 内核，未探索与其他硬件后端的兼容性。 Requires custom Triton kernels; compatibility with other hardware backends not explored.
稀疏模式依赖于块状选择，可能无法捕获细粒度的令牌级稀疏性。 Sparse pattern relies on blockwise selection; may not capture fine-grained token-level sparsity.
未测试超长序列（>64k）的性能，潜在退化未知。 Performance on extremely long sequences (>64k) not tested; potential degradation unknown.
门控和压缩 MLP 的训练开销可能抵消小模型的效率提升。 Training overhead from gating and compression MLPs may offset efficiency gains in small models.