Attention

摘要 扩展有效上下文长度对于推动大语言模型 (LLM) 向通用人工智能 (AGI) 发展至关重要。然而，传统注意力机制中固有的计算复杂度呈二次方增长，带来了令人望而却步的开销。现有的方法要么强加了强偏置结构，例如特定于任务的 sink 或窗口注意力，这些方法通常是任务特定的，要么彻底地将注意力机制修改为线性近似，但这些方法在复杂推理任务中的性能仍未得到充分探索。 在这项工作中，我们提出了一种遵循“更少结构”原则的解决方案，该方案允许模型自主决定关注的位置，而无需引入预定义的偏见。我们引入了块注意力混合（MoBA），这是一种创新的方法，它将专家混合（MoE）的原则应用于注意力机制。 这种新颖的架构在长文本任务中表现出卓越的性能，同时提供了一个关键优势：能够在完整注意力和稀疏注意力之间无缝切换，从而提高效率，且不会有损性能。MoBA 已经部署用于支持 Kimi 的长文本请求，并展示了在大语言模型（LLM）高效注意力计算方面的显著进展。 1 引言 现有的方法通常利用预定义的结构约束来利用这种稀疏性，例如基于汇聚节点的方法 (G. Xiao et al. 2023) 或滑动窗口注意力机制 (Beltagy et al. 2020)。虽然这些方法可能有效，但它们往往是高度特定于任务的，可能会阻碍模型的整体泛化能力。 替代方法，例如稀疏注意力机制，以 Quest (Tang et al. 2024)、Minference (H. Jiang et al. 2024) 和 Re-trievalAttention (Di Liu et al. 2024) 为例，在推理时选择 Token 的子集。 尽管这些方法可以减少长序列的计算量，但它们并没有实质性地减轻长上下文模型的高昂训练成本，因此难以有效地将大语言模型扩展到数百万 Token 的上下文。 最近，另一种有希望的替代方案以线性注意力模型的形式出现，例如 Mamba (Dao and Gu 2024)、RKWV (Peng, Alcalá, et al. 2023)、Peng, Goldstein, et al. 2024 和 ReNet (Sun et al. 2023)。 这些方法用线性近似代替传统的基于 softmax 的注意力，从而减少了长序列处理的计算开销。 然而，由于线性和传统注意力之间存在显着差异，因此调整现有的 Transformer 模型通常会产生高昂的转换成本 (Mercat et al. 2024, J. Wang et al. 2024, Bick et al. 2025, M. Zhang et al. 2024) 或需要从头开始训练全新的模型 (A. Li et al. 2025)。 更重要的是，它们在复杂推理任务中的有效性证据仍然有限。 ...

摘要 长上下文建模对于下一代语言模型至关重要，然而标准注意力机制的高计算成本带来了显著的计算挑战。稀疏注意力为提高效率同时保持模型能力提供了有前景的方向。我们提出了NSA（可原生训练的稀疏注意力机制），通过算法创新与硬件特性对齐的优化相结合，实现了高效的长上下文建模。NSA采用动态分层稀疏策略，结合粗粒度Token压缩与细粒度Token选择，在保持全局上下文感知能力的同时确保局部精度。该方法通过两大关键创新推进稀疏注意力设计：（1）通过算术强度平衡的算法设计实现显著加速，并针对现代硬件进行实现优化；（2）支持端到端训练，在保持模型性能的前提下减少预训练计算量。如图1所示，实验表明采用NSA预训练的模型在通用基准测试、长上下文任务和基于指令的推理中均保持或超越全注意力模型。与此同时，NSA在处理64k长度序列时，在解码、前向传播和后向传播阶段相较全注意力机制均实现了显著加速，有效验证了该方案在整个模型生命周期中的效率优势。 介绍 一种实现高效长文本建模的自然方法是利用 softmax 注意力的固有稀疏性。通过选择性地计算关键的查询-键对，可以在保持性能的同时显著降低计算开销。最近的进展通过各种策略展示了这种潜力，包括：KV-cache 驱逐方法、分块 KV-cache 选择方法，以及基于采样、聚类或哈希的选择方法。尽管这些策略很有前景，但现有的稀疏注意力方法在实际部署中往往表现不足。许多方法未能实现与其理论收益相当的加速；此外，大多数方法主要关注推理阶段，缺乏有效的训练时支持来充分利用注意力的稀疏模式，从而无法充分挖掘注意力的稀疏模式。为了解决这些局限性，有效部署稀疏注意力必须应对两个关键挑战：（1）硬件对齐的推理加速：将理论计算量的减少转化为实际的速度提升，需要在预填充和解码阶段进行硬件友好的算法设计，以缓解内存访问和硬件调度瓶颈；（2）训练感知的算法设计：通过可训练的算子实现端到端计算，以降低训练成本，同时保持模型性能。这些要求对于实际应用至关重要，以实现快速的长文本推理或训练。综合考虑这两个方面，现有方法仍然存在明显的差距。 为了实现更有效和高效的稀疏注意力机制，我们提出了 NSA，一种原生可训练的稀疏注意力架构，它集成了分层 Token 建模。如图2 所示，NSA 通过将键 (key) 和值 (value) 组织成时间块，并通过三个注意力路径处理它们来减少每个查询的计算量：压缩的粗粒度 Token、选择性保留的细粒度 Token 和用于局部上下文信息的滑动窗口。然后，我们实现专门的内核以最大限度地提高其实际效率。NSA 引入了两个与上述关键要求相对应的核心创新：（1）硬件对齐系统：优化块状稀疏注意力，以充分利用 Tensor Core 并优化内存访问，确保平衡的算术强度。（2)训练感知设计：通过高效的算法和反向算子实现稳定的端到端训练。这种优化使 NSA 能够支持高效部署和端到端训练。我们通过对真实世界语言语料库的综合实验来评估 NSA。在使用 260B Token 的 27B 参数 Transformer 主干上进行预训练后，我们评估了 NSA 在通用语言评估、长上下文评估和思维链推理评估中的性能。我们进一步比较了 A100 GPU 上内核速度与优化的 Triton 实现。实验结果表明，NSA 实现了与完整注意力基线相当或更优越的性能，同时优于现有的稀疏注意力方法。此外，与完整注意力相比，NSA 在解码、前向和后向阶段都提供了显著的加速，并且加速比随着序列长度的增加而增加。这些结果验证了我们的分层稀疏注意力设计有效地平衡了模型能力和计算效率。 重新思考稀疏注意力方法 大多数方法主要在推理过程中应用稀疏性，同时保留预训练的完整注意力骨干网络，这可能会引入架构偏差，从而限制它们充分利用稀疏注意力的优势。在介绍我们原生的稀疏架构之前，我们通过两个关键视角系统地分析这些局限性。 高效推理的错觉 尽管在注意力计算中实现了稀疏性，但许多方法未能实现相应的推理延迟降低，这主要是由于两个挑战：阶段性稀疏。诸如 H2O 之类的方法在自回归解码期间应用稀疏性，同时在预填充期间需要计算密集型预处理（例如，注意力图计算、索引构建）。相比之下，像 MInference 这样的方法仅专注于预填充稀疏性。这些方法未能实现跨所有推理阶段的加速，因为至少一个阶段的计算成本与完整注意力机制相当。这种阶段专业化降低了这些方法在以预填充为主的工作负载（如书籍摘要和代码补全）或以解码为主的工作负载（如长链式思维推理）中的加速能力。 与高级注意力架构的不兼容。 一些稀疏注意力方法无法很好地适配现代高效解码架构，例如多查询注意力（MQA）和分组查询注意力（GQA）。这些架构通过在多个查询头之间共享 KV，显著降低了解码过程中的内存访问瓶颈。 例如，像 Quest 这样的方法中，每个注意力头独立地选择其 KV-cache 子集。 虽然这种方法在多头注意力（MHA）模型中表现出一致的计算稀疏性和内存访问稀疏性，但在基于 GQA 等架构的模型中，情况则有所不同。在 GQA 架构中，KV-cache 的内存访问量取决于同一 GQA 组内所有查询头选择的并集。 这种架构特性意味着，虽然这些方法可以减少计算操作，但所需的 KV-cache 内存访问量仍然相对较高。 这种限制迫使我们面临一个关键选择：虽然一些稀疏注意力方法减少了计算量，但它们分散的内存访问模式与高级架构中高效的内存访问设计相冲突。 这些限制的出现是因为许多现有的稀疏注意力方法侧重于 KV-cache 减少或理论计算减少，但难以在高级框架或后端中实现显著的延迟降低。 这促使我们开发算法，将先进的架构和硬件高效的实现相结合，从而充分利用稀疏性来提高模型效率。 ...