arXiv。 arXiv:2502.11089（Submitted 16 Feb 2025，last revised 27 Feb 2025，v2）
摘要
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.
摘要指出长上下文下标准注意力成本高，稀疏注意力方向有前景。NSA 将算法与硬件对齐优化结合，采用粗粒度压缩与细粒度选择的分层稀疏，并强调两项创新：按算术强度平衡的加速，以及可端到端训练。实验报告在通用、长上下文与指令推理上持平或超过 Full Attention，并在 64k 序列的解码、前向与反向中给出加速。
总结
该工作从一项实践差距出发：理论稀疏往往不能降低墙钟时间，且许多方法仅在推理阶段施加稀疏。NSA 规定三条通路——压缩块、选择块与滑动窗口——并从预训练引入该层次。按块选择是硬件上的取舍：连续加载可以合并访存；共享 KV 的 Head 共享同一选择。在总参数 27B、激活 3B、260B Token 的主干上，平均分被报告为持平或超过完整注意力。A100 上的 Kernel 测量报告 64k 处前向最高约 9.0 倍、反向约 6.0 倍，并按 KV 流量估计解码约 11.6 倍。这些数字不是普遍的端到端产品加速。
详细解读
理论稀疏为何不等于时延
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.
完整因果注意力使每个 Query 与全部历史 Key 比较。训练与 Prefill 大致按 N² 增长。长解码反复读取大型 KV Cache。减少数学上的配对数，并不等于降低时延。摘要已陈述分层稀疏策略；本节其余部分说明该层次必须进入训练，而不能仅在推理阶段施加。
两类瓶颈与三条通路
算术强度——每搬运一字节所完成的 FLOPs——决定 Kernel 受算力限制还是受带宽限制。
This leads to different optimization goals — reducing computation cost during training and prefilling, while reducing memory access during decoding.
仅稀疏解码或仅稀疏 Prefill，将使另一阶段仍按完整注意力计费。若每个 Head 独立选择零散 Token，在 GQA 或 MQA 下 KV 流量可能成为这些选择的并集，内存流量不必随 FLOP 下降。NSA 因此选择连续块，在训练中引入它们，并由适合 Tensor Core 的 Kernel 执行。
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.
压缩注意力经可学习函数将连续 Key/Value 块映射为块表示，以较低分辨率覆盖全部上下文。选择注意力复用粗粒度分数对原始块排序，保留 Top 块并对其 Token 做注意力。滑动窗口保留最近 Token，以避免局部模式独占学习信号。可学习门控融合三路输出。论文保持重映射规模 N_t ≪ t。
质量与速度
预训练采用总参数 27B、激活 3B 的 MoE 主干与 260B Token。通用、长上下文与思维链评测的平均分被报告为持平或超过完整注意力。
NSA achieves perfect retrieval accuracy across all positions in 64k-context needle-in-a-haystack (Kamradt 2023) test.
Needle 检索具有参考价值，仍比跨文档合成大量弱信号更为简单。
在 A100 上，Triton 实现与 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.
论文并按内存受限模型下的 KV 访问量，估计 64k 解码最高约 11.6 倍。该数字不是普遍的端到端产品加速。采样、MoE 路由、通信、批大小与较短上下文都会降低注意力在总耗时中的比例。
证据边界
FlashAttention 讨论的是如何以更少搬运计算全部配对。NSA 讨论的是应保留何种有结构的子集，以及硬件能否执行。压缩与窗口分支可采用 FlashAttention 类 Kernel；选择分支需要块稀疏 Kernel。短上下文可能消除优势。已用稠密注意力训练的检查点，不能通过更换单一模块获得 NSA。在这一论述中，高效注意力设计需要同时规定四项内容：信息层次、学习路径、内存访问模式与 GPU Kernel。缺少这四项的稀疏度陈述是不完整的。