ACL 2025 最佳论文解读 | LLM注定是稀疏架构本文结合算法创新与硬件友好的优化设计，实现了高效的长上下文建模。

前言本文结合算法创新与硬件友好的优化设计，实现了高效的长上下文建模。NSA采用动态分层稀疏策略，将粗粒度的token压缩与细粒度的token选择相结合，从而同时保留全局上下文感知与局部精确性。

欢迎关注公众号CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。

本文转载自AI TIME论道

仅用于学术分享，若侵权请联系删除

论文题目： Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention

内容简介

长上下文建模是下一代语言模型的关键能力，但标准注意力机制的高计算开销带来了显著的计算挑战。稀疏注意力为在保持模型能力的同时提升效率提供了一个有前景的方向。本文提出NSA（Natively trainable Sparse Attention，原生可训练稀疏注意力），该方法结合了算法创新与硬件友好的优化设计，实现了高效的长上下文建模。NSA采用动态分层稀疏策略，将粗粒度的token压缩与细粒度的token选择相结合，从而同时保留全局上下文感知与局部精确性。本文在稀疏注意力设计上提出了两项关键创新：（1）通过算术强度均衡的算法设计实现了显著加速，并结合现代硬件进行了实现优化；（2）支持端到端训练，在降低预训练计算开销的同时不损失模型性能。实验结果（如图1所示）表明，基于NSA预训练的模型在通用基准测试、长上下文任务以及指令推理任务上均能保持甚至超越全注意力模型的表现。同时，在64k长度序列下，NSA在解码、前向传播与反向传播阶段均实现了显著的加速，有效验证了其在整个模型生命周期中的高效性。

论文地址：arxiv.org/html/2502.1…

论文创新点：

通过精心设计的算法平衡了计算密度，并针对现代硬件设计了专用算子，显著提升了运行速度，
实现了端到端的训练模式，在确保模型性能的前提下大幅降低了预训练的计算量。
同时，NSA是一个专为硬件优化的系统，打破了性能与成本之间的权衡取舍，推动高效LLM发展到新的阶段。

正文

Introduction

本文主要介绍NSA工作的创新点及其实验结果。NSA的研究动机来源于大模型在处理长文本时面临的挑战。首先，二次复杂度的Softmax Attention在训练与推理过程中带来极大的计算开销，并成为主要的效率瓶颈。为了解决全量Attention的高计算成本，学术界和工业界提出了多种稀疏化方法。然而，在NSA之前的诸多方法仍存在明显问题。

一方面，虽然稀疏化策略在理论上能够降低Attention的计算复杂度，但这种优势往往停留在理论层面，难以真正实现训练的加速。造成这一现象的原因主要有两点。其一，许多算法设计没有充分考虑硬件特性。例如，它们采用了过于细粒度的稀疏策略，但缺乏相应的高效内核实现，导致在实际运行中无法发挥优势。其二，部分架构设计与传统的多头注意力（MHA）紧密绑定，无法迁移至更加高效的结构，如GQA等。这不仅使得其推理延迟较高，也削弱了利用新型架构优势的可能性。

另一方面，在稀疏方法的研究早期，后训练的重要性尚未被广泛认识，许多工作仅是在训练完成的模型上附加稀疏化模块。这种方式虽然能够在一定程度上减少计算量，但并未解决如何在训练过程中实现原生的稀疏化支持，从而限制了稀疏方法的实际应用潜力。基于这一问题，NSA尝试提出一种能够原生支持训练的稀疏Attention机制。

设计可训练的稀疏Attention有两个主要出发点。其一，目标不仅在于加速推理阶段，同时也希望在训练阶段实现算力与时间的节约，从而缓解大规模模型训练的成本压力。其二，在训练过程中引入稀疏机制有助于保持甚至提升模型性能。相较之下，如果在训练完成后再对模型结构进行修改，即便删除的部分路径信息量较少，也可能破坏模型已学习到的结构，从而导致性能下降。因此，通过在训练阶段直接引入稀疏化，NSA能够更好地兼顾效率与性能。

The NSA Solution

接下来关注NSA在性能上的表现。论文中提供的实验结果显示，NSA在多个方面均取得了优异的效果。在三个类别的基准测试中，包括通用任务基准、长文本处理以及推理任务基准，NSA的性能均达到了与全量Attention相当的水平，表现出良好的适应性和鲁棒性。在效率方面，NSA在推理阶段的前向与反向解码过程中均实现了显著的速度提升，这一优势与其稀疏度设计相契合，并能够与全量Attention在速度上形成对照。为了实现上述性能与效率的平衡，NSA在训练过程中还特别设计并实现了高效算子，以充分发挥稀疏Attention的加速潜力。

Core Methodology

NSA的算法设计框架首先基于对输入序列的分块处理，即采用Block-wise 的方式将序列划分为若干部分。随后，引入了三分支的注意力机制，分别对应于Compression Attention、Selection Attention与Sliding Window Attention，它们在不同层面上对序列信息进行建模。

在第一个分支中，Compression Attention会将每个由相邻token组成的block进行压缩操作。例如，四个block的token会被映射为四个压缩后的compression token。在这一过程中，仅对K和V进行压缩，而Q保持不变，所有分支共享原始粒度的Q。这意味着若压缩率为α，则Compression Attention的计算复杂度仅为全量Attention的1/α。与其他方法不同，NSA的压缩方式在保持性能方面更具优势，因为压缩K、V相较于压缩Q对模型性能的影响更小。Compression Attention的主要作用有两个：一是以较低的计算量为模型提供全局概览，从而增强全局感知能力；二是为Selection Attention提供索引。具体而言，Compression Attention的注意力得分可用于选择最重要的若干block，并作为后续Selection Attention的输入。

在第二个分支中，Selection Attention会根据Compression Attention的得分选择top-k的重要block，并将这些block的原始token拼接起来，形成新的注意力输入。其余不重要的block则不会进入该分支的计算。最终，Selection Attention使用原始的Q对这些重要block的K和V进行建模，从而形成稀疏注意力的核心部分。

第三个分支为Sliding Window Attention，用于提供局部依赖的建模能力。该分支确保模型能够捕捉近距离的上下文信息，从而避免Compression Attention和Selection Attention在训练早期过度依赖局部路径而导致学习短路。实验结果表明，若去除 Sliding Window Attention，模型的表达能力会有所下降，因此该分支对于整体性能具有关键作用。最终，三个分支的输出通过门控加权的方式进行融合，形成NSA的最终输出，其注意力mask的设计如论文图示所示。

此外，在实现过程中还需特别注意因果性约束。对于自回归任务中的next token prediction，模型必须严格遵循因果结构。在Compression Attention中，位于某个block内的Q token并不能访问该block中包含未来信息的K和V，以避免信息泄露。只有当Q位于block的最后一个位置时，才允许其访问对应的压缩信息。在Selection Attention的mask设计中，算法默认保留最近的一个block，并固定选取最远的一个attention sink。同时，在实现最近block的因果mask时，也需要严格遵循因果性限制。相比之下，Sliding Window Attention的实现则较为直接，其主要作用是稳定模型对局部信息的学习。最终，三个分支共同构成了NSA的核心设计架构。

Key Innovation

Natively Trainable Design

NSA在设计中重点解决了稀疏注意力难以实现可训练性的问题。传统的稀疏注意力方法之所以不可训练，主要原因在于selection过程本身不可导，因此无法通过梯度反向传播进行优化。为了解决这一问题，NSA借鉴了Mixture-of-Experts（MoE）的思路，将梯度的传导路径与索引的选择路径区分开来。具体而言，负责选择block索引的过程由Compression Attention分支完成，该过程对应图中的黑色箭头，但这一选择操作无法直接传递梯度。

为了使模型能够对索引的选择进行学习，NSA为Compression Attention单独设计了一个额外的输出路径。通过该路径，Compression Attention的特征能够独立地产生可学习的输出，从而间接引导Selection Attention的索引选择过程。换言之，虽然索引的选择操作不可导，但通过为Compression Attention增设可训练的输出，NSA保证了索引选择能够由神经网络进行优化。正是这一机制，使得Selection Attention的索引能够被训练所支配，从而实现了稀疏注意力的原生可训练性。这一设计思路构成了NSA方法的核心创新之一。

Hardware-Aligned System for Real Speedups

在算法设计完成之后，我们进一步来看NSA在算子实现上的思路。整体而言，NSA在实现过程中始终紧贴硬件特性展开设计，以确保在现代GPU上获得最大化的加速效果。首先，之所以选择Block-based的稀疏模式，是因为GPU在连续内存访问上有着更优的优化能力。如果K和V的选择能够保持块状连续性，那么在训练与推理阶段，算子在硬件层面会具有更高的效率和更好的适配性。

进一步来看，NSA在Q的处理方式上也进行了专门的设计。由于NSA的稀疏结构是Q-aware的，每个Q所选择的KV集合不同，这与Flash Attention的设计形成差异。如果依旧采用Flash Attention中以block Q作为计算单元的方式，那么在计算时就需要将一个block中所有Q的KV并集加载到内存中，造成了存储与计算上的大量浪费，也使得算子在复杂度与效率上难以保证。因此，NSA放弃了block Q的处理方式，而是选择以单个Q作为循环单元，使其独立attend对应的KV，从而保证了所有被选择的KV都能参与计算，同时避免了HBM访问的浪费。

然而，如果仅以单个Q作为计算单元，那么与block K的交互会退化为矩阵与向量的乘法，难以充分利用Tensor Core的优势。为此，NSA通过在Q中引入group维度，使多个Q能够共享同一个K集合。这样，group维度与query维度可以共同构成矩阵，从而与K矩阵进行高效的矩阵乘法运算，充分发挥Tensor Core的计算能力。这一设计也使得NSA可以高效融入GQA这类高效推理Attention结构中。

在完成Q的分组后，NSA进一步将其分片。具体而言，Q按照token粒度进行切分，得到的每个循环阶段的Q为一个的矩阵，其中表示group的大小。相比之下，K和V的分片方式更为直接，与Flash Attention保持一致，即沿着序列方向进行划分。每个group中的K按照block size切分为矩阵单元，与对应的Q在内循环中逐步加载，计算完成后输出结果。整体算子的forward计算遵循经典流程：从SRAM加载K，与Q进行计算，再将结果输出到HBM中。

最后，NSA的稀疏结构在并行负载均衡上也展现出优势。其Top-k的选择模式使得不同Q在内循环中加载的KV数量保持均衡。除了因因果结构导致前几个Q选择较少的情况外，其余Q大多会横向选择相同数量的K块，例如三个block。这样的设计使得GPU的每个SM在执行时负载接近一致，从而保证了整体计算过程的并行效率与调度的平衡性。

Performance Evaluation

在完成算子设计之后，我们进一步在大规模模型上验证了NSA的训练效果。实验基于一个27B参数规模的模型展开，该模型采用了DeepSeek MOE的架构，因此其激活参数规模约为3B，相对较小。在实验设计中，我们确保了Full Attention与NSA模型在参数量上的公平对比。NSA在实现时，尽管引入了三分支的Attention结构，但其仍然复用了同一套Query，同时每个分支仅独立维护一对Key和Value矩阵。由于NSA基于GQA架构，因此额外的K和V矩阵规模极小，使得整体参数量与Full Attention保持在相近水平。

在计算复杂度上，NSA相较于Full Attention有着显著优势，尤其是在长序列任务中，这种差异愈加明显。随着文本长度的增加，NSA能够在保持有效性的同时，显著减少计算量，而Full Attention在计算开销上则呈现线性增长。换言之，从计算角度来看，Full Attention在实验中事实上拥有天然的优势。然而，即便在这一优势条件下，NSA依然展现出了更优的训练表现。

具体而言，在loss曲线的对比中，NSA在整体训练过程中始终表现得比Full Attention更优，且收敛过程平稳，不存在因稀疏结构引入的不稳定性问题。这说明NSA不仅在计算效率上具备优势，同时在模型收敛性与优化效果上也保持了良好的表现。

在性能评测方面，我们首先在一系列综合性benchmark上对NSA与Full Attention进行了对比。结果显示，在九个基准测试中，NSA在其中七项任务上取得了优势，整体平均分数也优于Full Attention。对于一个稀疏化设计的注意力机制而言，这种在性能层面的全面领先是极为难得的。进一步地，我们也在若干具有代表性的长文本benchmark上进行了深入对比。在图中的测试中，即便在64K序列长度下，NSA依然能够稳定捕捉关键信息，展现出较强的长程建模能力。在LongBench的评测结果中，NSA相较于现有的多种稀疏注意力方法（包括Exact-TopK等）均表现出明显优势，且在平均表现上同样优于Full Attention。整体来看，NSA在十余个benchmark的综合评测结果中，都呈现出稳定且一致的优越性。最后，我们还在reasoning相关的任务上对NSA与Full Attention进行了更具挑战性的对比。通过对模型进行RFT并在AIME benchmark上进行测试，我们观察到NSA在推理能力上甚至展现出了超越Full Attention的性能。这一结果进一步表明，NSA不仅在效率与长文本任务中具有优势，同时在复杂推理场景下也能够提供额外的收益，显示出良好的通用性与潜力。

Efficiency Analysis

在效率方面，我们将NSA的自研稀疏Attention kernel与同样基于相同Backbone实现的Flash Attention kernel进行了对比。实验结果表明，在forward阶段，NSA随着序列长度的增加展现出显著的速度优势。在64K序列长度下，forward端的速度提升可达到9倍；在backward阶段，同样在64K的情况下，NSA相较于Full Attention也实现了6倍的加速。值得注意的是，这种性能提升会随着序列长度的增长而进一步扩大，呈现出“越长越优”的趋势。

这一现象的原因可以通过NSA内部三分支Attention的计算时间分布得到解释。随着序列变长，compression attention的计算占比逐渐增加，而sparse attention的时间开销则呈现出线性增长的趋势。这与其计算复杂度分析相吻合：在固定top-k的前提下，sparse attention的复杂度是O(N)，而compression attention由于采用固定倍率压缩，其复杂度则为O(N²/α)，其中α为压缩率。因此，在较大压缩比（如32:1或64:1）的设定下，NSA的渐进速度提升会逐渐达接近理论值，即压缩分支的的压缩比。

在推理阶段，NSA同样展现出优越的效率表现。解码过程的主要瓶颈在于访存，而NSA由于缓存需求显著低于Full Attention，使得其解码速度能够随着序列长度进一步提升。为此，NSA在设计时采用了GQA架构，并在组内共享Key与Value，从而有效降低了访存开销。如果采用MHA或各个head独立选择token，将导致缓存读取量成倍增加，在最坏情况下需要读取的token数量为head数×top-k block×block size，这会严重拖慢解码速度。因此，NSA的设计不仅在理论上符合复杂度优势，同时在训练和推理过程中也充分考虑了硬件实现与访存瓶颈，确保了整体的高效性。

Visualization

为了进一步理解NSA稀疏选择的有效性，我们对attention score进行了可视化分析。从可视化的热力图中可以观察到，在tensor score的分布中存在明显的block-wise聚类现象，即模型在不同位置对若干相邻block的注意力呈现集中化的模式。这一现象间接验证了NSA在block粒度上进行压缩与选择的合理性，也从直观上说明了稀疏选择机制能够捕捉到序列内部具有代表性的结构模式。

Conclusion

综合来看，NSA的核心贡献体现在两个方面。其一，在设计理念上，NSA的整个算法框架充分考虑了训练、推理以及硬件实现的需求，提出了一种在实践中高度可行且硬件友好的三分支Attention架构。其二，在性能表现上，NSA通过Natively trainable的稀疏机制，不仅在训练过程中保持了稳定的收敛特性，还在多个评测基准上取得了优于甚至超越Full Attention的效果。这表明，稀疏Attention的潜力不仅仅在于计算与存储效率的提升，更能够在可训练性保障的前提下实现对模型性能的进一步增强。我们也因此相信，未来的稀疏结构研究将越来越多地转向对可训练性的探索与实现，使其能够在效率与性能之间取得更好的平衡。