北大袁境阳：稀疏注意力机制让模型 10 倍加速——Attention

NSA（Native Sparse Attention，即原生稀疏注意力机制）是自然语言处理（NLP）和深度学习模型优化‌方向上，多次被研究者们讨论和引用的概念，也是 Attention 领域内的关键研究之一。

本次对谈，我们邀请到了和 DeepSeek 梁文锋团队合作，共同撰写 ACL 2025 最佳论文Native Sparse Attention: Hardware-Aligned and Natively Sparse Attention 的第一作者，北京大学计算机学院的袁境阳博士。

当下一代模型需要 128k 或 1M 上下文时，到底哪种注意力机制能生存下来？

如果我们想让模型真正 “记住全部历史”，能持续推理几十分钟，Attention 需要被怎样重新调整？

是否存在一种注意力结构，能够在预训练阶段就以稀疏方式工作，同时在推理阶段获得 10× 以上加速且不掉精度？

带着对上述问题的好奇，我们就论文中的创新和思考，与境阳博士展开深度对话。以下是论文创新性研究及深度对谈的部分整理，阅读时间需要约 20 分钟。

Enjoy

超长上下文能力，正在成为新一代大模型能力的分水岭。

市场要求模型必须能读懂整段代码、长篇文档，甚至在数万 Token 的链式推理中保持连贯。但随着序列长度不断攀升，标准全注意力的计算与访存消耗迅速成为系统瓶颈。比如，在 64k 上下文场景下，Softmax 注意力往往会吞噬 70%–80% 的解码延迟，使得训练成本与推理成本双双被限制在一个不甚理想的上限。

尽管近年来涌现出诸多稀疏注意力的尝试，但它们大多只在推理阶段生效，与现代 GPU 架构不完全对齐，也缺乏端到端训练能力，使理论收益难以真正落地。

本篇工作的核心 NSA（Native Sparse Attention），尝试从根本处重写这一结构性矛盾。

它不是在全注意力模型上事后裁剪，而是在架构内部重新组织信息流：遥远的 Token 被聚合成若干连续块，以 “块级摘要”（Block）的方式捕获全局轮廓；这些块再经过轻量的评分，从中挑出少数关键块，保留重要细节；而靠近当前 Token 的局部上下文，仍以滑动窗口的方式被完整建模。

这三条注意力路径并行计算，最后通过可学习的门控机制融汇为单一输出。与此同时，这套结构以 “连续块” 为基本操作单元，专为 GPU 访存的特性与 Tensor Core（NVIDIA 研发的专用处理核心）调度，设计了新的 Triton（一种开源的编程语言和编译器）内核，使稀疏注意力的训练、Prefill 与解码都能真正获得加速。

在一个由 GQA 与 MoE 组成的 27B 模型上，NSA 自预训练阶段以稀疏形式参与计算：在 64k 上下文下，训练前向最高可达全注意力的约 9 倍速度，反向可达 6 倍。在解码阶段，KV 访存量下降至原来的约十分之一，对应接近理论上限的 11.6× 加速。更重要的是，性能并未因稀疏性而折损。在 MMLU、GSM8K、DROP 等通用基准上，整体不逊于全注意力；在 LongBench、64k Needle-in-a-Haystack 等长上下文任务上取得更明显的优势；在 AIME 24 数学推理实验中，即使在 8k 或 16k 的生成窗口内，基于 NSA 的模型仍展现出比全注意力更可靠的链式推理精度。

NSA 优于多种现有的稀疏注意力方法

这些结果共同指向一个耐人寻味的结论：对于长上下文模型而言，性能与效率并不一定是此消彼长的选择。

应该尝试让模型学会 “把算力花在应该花的地方”。

NSA 的第一层创新在于，它将稀疏性从 “优化手段” 提升为 “结构原则”。

以往的稀疏注意力，往往依赖启发式规则或后处理剪枝，模型本体依旧由全注意力主导。而在 NSA 中，注意力被拆解为三条并行路径：压缩路径描摹全局轮廓，选择路径负责保留关键细节，滑动窗口路径则确保模型不会失去对局部上下文的敏感性。

这种分工让模型可以在不同尺度上理解同一段上下文，使稀疏性成为结构学习过程的自然结果，而非外加限定。更重要的是，三条路径之间通过可学习的门控机制彼此协调，使模型能够在预训练中自我塑形、逐步形成最优的稀疏行为模式。

第二层创新来自它与现代 GPU 的深度耦合。

模型只在连续 Token Block 上进行选择与加载，这种块状结构与 GPU 的访存方式天然兼容。论文针对 GQA/MQA 的特性，设计了新的 Triton 内核：同一 GQA 组内的多个注意力头共享稀疏后的 KV 索引，所有需要参与计算的 KV Block 被一次性、连续地加载到 SRAM（静态随机存取存储器）中。这种 “按组加载、按块计算” 的方式最大化了 Tensor Core 的利用率，避免了以往稀疏方案中 “大量跳跃访存” 导致的硬件浪费。

第三层创新则关乎 “如何学习稀疏”。

它复用压缩路径在中间生成的注意力分布来近似推断每个块的重要性，并通过连续化的 Top-n 选择机制使这一过程可微，从而能够在预训练中通过反向传播直接学习稀疏模式。这意味着模型可以在庞大的训练语料中自动探索 “应该保留哪些块、应该忽略哪些块”，无需人工设计启发式规则，也避免了离散选择带来的训练不稳定。

以下是对谈原文的整理：

绿洲：境阳博士好，非常感谢您接受今天的访谈。我们先从一个更个人的起点开始：能否请您介绍一下您的研究背景，以及为什么最终会走到 “高效注意力” 这一方向？

袁博士：好的，我也很高兴能和大家一起讨论我们的工作。

我目前在北京大学攻读博士，所在的院系是计算机学院，主要从事大规模模型的架构研究。我自己的话，更聚焦的方向是高效注意力（Efficient Attention）。我相信您的团队之所以找到我，也是因为我围绕高效注意力，或者说稀疏注意力展开的研究工作。

整个研究脉络和初心，其实源自一个很现实的趋势。随着模型变得越来越智能，它们需要处理的上下文长度也越来越长。这个意味着不仅是能够接收长文本输入，更是要在长序列中保持连续推理能力。而长上下文处理中的效率瓶颈，一直是限制模型进一步发展的核心障碍。

绿洲：我们是不是可以理解为，“高效注意力” 并不是一个孤立的技术点，而是长上下文时代里绕不过去的系统性问题？

袁博士：没错。我们在研究中反复观察到：模型之所以难以扩展到很长的上下文，很大一部分原因来自注意力的固有成本。简单来说就是只要注意力机制的复杂度与序列长度平方关系增长，计算量就不可避免地爆炸。

而在稠密注意力（Dense Attention）中，每个Token 都要与所有其他 Token 做关联计算。当然了，我们可以看到这些 Pair-wise （全对偶）连接，但它们真正起作用的比例，其实是相当低的。

换句话说，大部分注意力计算本身并不 “值得” 做。

这也是我们重新思考稀疏注意力（Sparse Attention）的起点。

与此同时，还有另一条重要的观察路径来自语言学和模型行为本身。我们发现，在一个大型语言模型内部，词与词之间的关联呈现出非常明显的 “稀疏性”，也就是说并不是每个词都与其他所有词有强关联。所以当模型在处理一个 Token 时，它真正依赖的上下文其实往往是结构性的、局部性的或者块状的，而不是一种完全无差别的密集集合。这种稀疏结构不仅是语言本身的特征，也是模型内部逐渐学习到的规律。

绿洲：听起来像是在说，语言的统计结构决定了 Attention 不可能是均匀密集的？

袁博士：可以这么说。

其实从 Softmax 的性质来看，我们也能得到相似的推论。Softmax 本质上是一个竞争性的激活函数，也就是说，这种函数在归一化之后，权重天然呈现 “少数突出、多数衰减” 的分布。这意味着 Attention 权重本来就会自发变稀疏，只不过传统的密集注意力把所有 Pair 都算了一遍，没有把这种稀疏性利用起来。

正是基于这些观察，我们才会去探索：既然注意力的有效部分天然稀疏、模型内部结构天然稀疏、Softmax 的行为也天然稀疏，那为什么不让注意力在架构层面就成为一种“原生稀疏”的机制？

这就是我们后来做原生稀疏注意力机制的动机。

绿洲：在我们刚才的讨论中，您其实已经触及到论文的核心动机：我们到底想解决一个什么样的问题？

那如果把工作拆开来看，NSA 的三分支：压缩、选择、滑动窗口，它们分别承担了怎样的角色？为什么需要把注意力拆成这三条路径？

袁博士：如果顺着我们前面的讨论继续展开，一个直接的问题就是：稀疏注意力最关键的并不是 “如何变稀疏”，而是 “如何选出应当被保留下来的那一小部分 Token”。

换句话说，稀疏注意力的核心，不在稀疏本身，而在于索引（Index），也就是说如何找到最关键、最相关的 Token。

这一步几乎带有一种 “Oracle 性质”，因为在运用稠密注意力进行计算时，我们会发现模型自身实际上已经暴露了哪些 Token 应该被重点关注。如果你已经把密集注意力的分数都算出来了，那其实也就失去了稀疏的意义了。所以在 NSA 研究中，我们把 “如何建立 Index” 这一问题放在了整个方法的中心。

这也就是为什么压缩分支会被放在最前面。

它的任务不是替代注意力，而是负责提供一个较轻量但可靠的评分来源：在 “压缩注意力” 里，我们对原始序列进行块级聚合，通过一次更廉价的 Full Attention，得到这些 Block 的粗粒度重要程度分数（Importance Score）排序。这些分数不会十分精细，但足够作为后续选择分支的 “初步判断”。

绿洲：也就是说，压缩并不是为了直接得到输出，而是为了提供一个可学习的、轻量级的 “预筛选”？

袁博士：对，如果我们只依赖选择分支，从一开始就要让模型决定每个 Token 的重要性，那它的学习难度会非常高。而压缩的输出，它能帮助模型把搜索空间从 Token-level 减少到 Block-level，使整个选择过程更稳定、更可控。

在进入真正的稀疏分支之前，我们需要做的第二件事，就是把压缩后的重要程度分数作为稀疏注意力的索引信号。这意味着我们不仅仅是保留最重要的 Token（或 Token Block），而是把这些位置作为稀疏注意力的主要计算位置，将有限的算力集中在那里。

绿洲：为什么不能让压缩分支同时承担 “选择” 和 “输出” 两件事？为什么一定要再单独做一个选择分支（Selection Branch）？

袁博士：这是一个非常关键的点。原因在于压缩本身不是一个可逆过程，它是一种 “丢信息” 的操作。你只能用它来判断哪些 Block 重要，但不能依赖它来恢复细节。因此，它无法直接作为最终的注意力的输出，因为那样模型将无法学习 Token-level 的精细关联。

所以我们必须让选择分支成为一个独立的 Attention 分支，为真正的稀疏区域提供细粒度的注意力计算。它的存在确保了模型在缩减计算量的同时，不会失去对关键信息的敏感性。

绿洲：那滑动窗口分支呢？为什么需要一个专门处理局部上下文的分支？

袁博士：这是我们在训练早期观察到的另一个结构性现象：模型在学习长上下文之前，往往先依赖区域内的密集信息，类似语言中的短距离依赖结构。

如果在这个阶段，让学习网络只能接触 “压缩” 与 “选择” 后的稀疏信号，它会过早陷入一种 “过度依赖长距离注意力、忽略局部语义” 的状态。这样的训练过程是不稳定的。

因此我们引入第三个分支，也就是 “滑动窗口注意力”分支（Sliding Window Attention）。它的角色非常明确：在训练初期，为模型提供一个稳定可靠的 “位置注意力”（Local Attention）通道，让模型先建立自然语言里最常见、也最稳固的短程关系。这个分支等于为模型 “兜底”，让整个系统不会在训练早期因为信号稀疏而失去方向。

绿洲：所以，这就是最终的 NSA 三分支架构了：压缩（Compression）负责预筛选、选择（Selection）负责关键细节、滑动窗口（Sliding Window）负责稳定训练。

袁博士：是的。我们在论文里也强调 NSA 的 “trinominal structure”，因为这三种注意力在训练和推理过程中承担着完全不同的职责，简单来说可以理解为：

压缩（Compression）负责 “在哪里看”；

选择（Selection）负责 “看什么”；

滑动窗口（Sliding Window）负责 “确保模型能够顺利学会如何看”。

如果缺少某一条路径，整个原生稀疏注意力的训练就无法成立，这也是为什么以前的稀疏注意力方法难以端到端训练。它们缺少一个稳定的学习过程。

所以，在这个意义上，NSA 的真正创新不是 “把注意力变稀疏”，而是 “让模型能够学习如何以稀疏的方式运作”。

绿洲：我们前面谈了 NSA 在算法层面的设计。

在论文里，你们也强调了对 “闪电注意力机制” 内核（Flash Attention Kernel）局限性的重新审视，能否在这个话题上展开聊一下？

袁博士：这是 NSA 中经常被忽略但非常核心的一点。为了理解它，我们要先回到闪电注意力机制本身是如何被设计的。

闪电注意力机制将 Q、K、V 按序列方向切分成若干 Block，再通过块级的连续访存保证内核（Kernel）的高吞吐。这种结构对于密集注意力来说非常理想，因为每个 Q Head 都会访问完整的 K-V 序列，数据连续性强，硬件利用率很高。

但当你尝试让注意力变稀疏时，这种优势一下子变成了限制。

因为稀疏注意力最大的特点，是每个 q 所需要访问的 k-v 块并不是连续的，而是根据 Index 动态选择的。当闪电注意力机制假设 “所有 Head 在相同的 Block 边界读取数据” 时，它就无法满足稀疏注意力机制的这种不规则、动态选择的访问模式。

尤其是在 NSA 这种需要 “Queue-aware” 的稀疏结构中，问题变得更明显。

如果每次 Block 内有多个 Q 头都需要访问不同的 K-V 区域，而这些区域又不是连续的，那么系统就不得不反复从显存中加载大量不连续的 KV 块。随后所有的稀疏优势都会被 I/O 访问放大几十倍甚至上百倍的代价抵消掉。

这也解释了为什么传统的闪电注意力机制内核无法直接支持 NSA：它的假设前提是密集、顺序、规则；而我们需要的是稀疏、动态、不连续。两者在模式上有根本冲突。

绿洲：所以 NSA 的核心是 “想让稀疏能在硬件上跑得起来，并且真的加速”？

袁博士：对可以这么说。稀疏注意力的核心挑战不是 “减少计算量”，而是 “减少访存”。如果你把计算量减少了，但为了找到那些散落的稀疏块而导致 I/O 翻倍，那它永远不会比密集注意力更快。

所以 NSA 必须从硬件访问模式起重新设计。我们做的事情很简单却非常关键：

我们让所有属于同一个 GQA Group 的 Head，在一个解码步骤（Decoding Step）中共享同一组 KV Block。这样一来，你只需要把这些 Block 加载一次到 SRAM，然后所有 Head 都能基于它们做选择。

这背后的直觉就是：用算法统一选择，用硬件统一加载。

对于稀疏注意力来说，这是一种从底层改变游戏规则的尝试。

绿洲：我们回到论文里的另一个重点：你们反复提到 NSA 的 “训练感知（Training-aware）” 设计。可以说一说你们对 “训练” 这件事的重视背后的思考吗？

袁博士：这个问题在过去确实被严重低估了。

过去几年的工作大多把注意力的稀疏性应用在推理（Inference）阶段，而一旦模型回到训练状态，尤其是预训练阶段，稀疏结构就几乎没有办法成立。

这是因为训练与推理在注意力机制的行为上有本质不同。训练阶段的注意力需要学习，学习本身需要 “梯度信号”。如果你从一开始就把注意力剪得很稀疏，它根本无法知道哪些 Token 是重要的：它缺乏一个可以探索的空间。换句话说，它还没学会稀疏，就已经被迫稀疏了。

因此 NSA 在 “如何让模型能够学会稀疏” 这个问题上做了非常多的设计。

首先，我们让压缩分支以一种稳定的方式提供梯度引导。它不是直接替代注意力，而是提供一个关于空间结构的粗粒度监督，让选择分支能够逐步学会 Token-level 的稀疏行为。

其次，我们加入滑动窗口作为一个始终可靠的局部通道。它保证了模型在训练早期不会因为稀疏导致学习断裂。没有这个通道，模型甚至无法学会自然语言本身的短程关联结构。

这个设计的关键是：稀疏不是一个状态，而是一种可以被学习的结构。

如果你希望模型以稀疏的方式运作，那你必须给它一个稳定的学习路径，让它逐步过渡。未来大模型的发展趋势也在不断印证这一点。以前大家把算力几乎全部投入预训练，但随着 Agent、长序列推理、强化学习等能力越来越重要，训练后期的代价正在急剧上升。

这意味着未来很大的优化空间，会来自 “模型在后训练阶段如何高效地继续学习”。如果架构本身不能被稀疏化、不能端到端训练，它就会阻碍未来的能力提升。

NSA 想要解决的，正是这一代模型必须面对的问题：如何让注意力学习技能 “密集”，又能在成熟后以的稀疏的方式高效运作，并且始终保持能力不掉。

绿洲：我们刚才从算法和内核的角度讨论了 NSA 的设计。但在本篇工作里，其实还有一个值得注意的视角是硬件的视角，就是说稀疏注意力的出现不仅是算法层面的突破，也来自现代 GPU 架构本身的发展趋势。能否进一步展开？

袁博士：是的，这其实源自行业观察。我们与很多产业团队一起交流时，会发现一个越来越明显的现象：现代 GPU 的硬件特性正在把行业推向稀疏架构。

以现在的商用 GPU 为例，它们的显存带宽其实远没有显存容量那么充裕。显存读取的延迟往往在几十微秒的量级，而显存容量却越来越大。这意味着什么？

意味着如果你的模型设计既浪费带宽又浪费显存，那你在解码阶段每秒钟能处理的 Token 数就会被限制得非常死。你想象一下：在某些国际评测中，一个模型如果一分钟只能生成几十个 Token，它几乎是不具有实用性的。

所以，在高带宽成为稀缺资源、显存容量却相对富余的今天，“密集且无差别地读写 KV” 这种方式就会显得格外低效。

绿洲：这听起来感觉像在 GPU 结构不改变的情况下，大模型行业最终会被迫走向稀疏路线。

袁博士：这么说也可以。

你可以理解为：当显存容量远大于可用带宽时，你唯一能保证模型 TPS（Token per Second）不掉的方法，就是让注意力访问更稀疏、更具选择性。

要么，你就会陷入一个悖论对吧，这个悖论就是说：

要想不浪费显存，你必须 Dense；要想不浪费带宽，你必须 Sparse。

但是这两者本身是冲突的，而稀疏架构是在这个矛盾里唯一能自洽的路径。

直白点说，就是就算你发明了一个非常强大的注意力模型，不管是基于 Linear 的、Softmax 的、还是 MoE 样式的结构，只要它是稠密的，你一定会用满所有显存带宽，TPS 数值下降到无法部署。

此时行业唯一能走通的路，就是构造一种稀疏的架构，让模型的注意力访问密度与硬件带宽相匹配。

这实际上也是我们在 NSA 里看到的逻辑：不是因为稀疏 “更美”，而是因为它 “更现实”。

事实上，当你开始理解 GPU 架构对模型的影响后，你会发现许多 “看似算法的创新”，其本质都是在规避硬件限制。例如我们刚才讨论的 KV 访问问题，就是一个典型例子：在稀疏结构下，如果你不改变 KV 的访问模式，你永远无法实现真正意义上的加速。

很多人误以为这是 “纯算法问题”，但其实它是一个 “带宽问题” 和 “数据布局问题”。

绿洲：在您最近的一些演讲里，您还提到 “后训练（Post-training）阶段的重要性正在迅速上升”，可以和我们分享一下吗？

袁博士：这是一点越来越明显的产业趋势。过去几年我们对模型的算力分配主要集中在预训练阶段，但随着更复杂的智能体（Agent）、检索增强系统（RAG）、长序列推理、RLHF 和各种强化推理范式的兴起，训练后的阶段正在消耗越来越多的算力。

这意味着，如果架构本身不能支持稀疏化的训练，那么后训练阶段的算力成本就会指数级增长，以至于你根本无法把模型推向更复杂、更长期的任务。

所以我认为，NSA 的意义不只是一个 “注意力算法改进”，而是为未来的大模型提供一种新的基础假设：模型在训练与推理的全周期内都需要稀疏化的路径。

从预训练到指令微调，从强化学习到偏好优化，从长期任务执行到 Agent 内部记忆管理，所有这些场景背后都有一个共同点：它们都要求模型能在长序列下保持效率。

如果注意力的密度不随能力降低，那么整个未来生态都会在带宽上被卡住。

因此我认为未来的一个发展方向，就是更多 “后训练架构” 的创新。NSA 是这一方向的早期尝试，它展示了一个可能性：如果架构本身就能吸收稀疏性，那么整个模型生命周期的计算效率，都能被重塑。

受访者及论文简介

本期受访人员：袁境阳博士

北京大学计算机学院博士生、数据科学与工程研究所 Dlib 组。主要研究方向是高效大语言模型和稀疏注意力机制，提出了原生稀疏注意力模型（NSA）等高效大模型架构。曾获北京市优秀毕业生、北京大学优秀毕业生等荣誉。

论文链接：https://arxiv.org/abs/2502.11089

感谢境阳博士受邀参与对谈。

本文来自微信公众号“緑洲资本 Vitalbridge”，作者：参赞生命力，36氪经授权发布。