NSA 原生稀疏注意力（一）：压缩-选择-滑窗三分支架构

背景与问题：长上下文推理的注意力困境#

下一代大语言模型的核心能力之一是把上下文窗口推到几万、几十万甚至上百万 token：仓库级代码补全要读完整代码库，多轮 agent 任务要维持跨数千 token 的对话状态，以 OpenAI o 系列、DeepSeek-R1 为代表的推理模型会在回答前生成几万 token 的思考链（chain-of-thought）。上下文一长，标准注意力（full attention）的二次复杂度就变成了无法回避的瓶颈。

回顾一下问题出在哪。对序列中第 $t$ 个 query token $\mathbf{q}_t$ ，注意力要对它之前的所有 key $\mathbf{k}_{:t}$ 计算相似度、对 value $\mathbf{v}_{:t}$ 做加权求和。序列长度为 $t$ 时，一个 query 就要处理 $t$ 个键值对，整个序列就是 $O(t^2)$ 。当 $t$ 到达 64k 时，仅注意力计算一项就占解码总延迟的 70%–80%（论文给出的估计值）。也就是说，到了长上下文场景，注意力本身取代了矩阵乘法，成为最大的性能账单项。

这里需要引入一个贯穿全篇的概念：算术强度（arithmetic intensity），即计算量（FLOPs）与访存量（bytes）的比值。每块 GPU 都有一个由峰值算力与内存带宽之比决定的临界算术强度：高于它，任务受算力限制（compute-bound），低于它，任务受带宽限制（memory-bound）。Transformer 的预填充（prefill）和训练阶段对整段序列做矩阵乘，算术强度高，是 compute-bound 的；而自回归解码每步只生成一个 token，却要把整个 KV 缓存（KV cache）从 HBM 读一遍，算术强度极低，是 memory-bound 的。结论很直接：训练和预填充阶段要省的是计算，解码阶段要省的是内存访问——两个阶段的优化目标并不相同。

图片来源：NSA 论文 Figure 1。左图：九项通用基准、长上下文任务与推理任务的平均分对比，稀疏的 NSA 在平均分上反超全注意力基线；右图：64k 序列下解码（decoding）、前向（forward）、反向（backward）三个阶段的加速比。

既然 query 往往只与少数 key 高度相关，一个自然的想法是稀疏注意力（sparse attention）：只计算真正重要的查询-键对，把理论计算量降下来。大量观察支持这个方向——softmax 注意力分数本身就有很强的稀疏性，例如 MagicPIG 论文中引用的分析显示，top 20% 的注意力权重只覆盖了约 70% 的注意力分数总和；预训练模型注意力图的可视化也表明，高分区域倾向于聚成连续块状（blockwise clustering），相邻 key 的重要性往往相近。

不过，稀疏化这条路在 2024 年前后已经有过大量尝试：KV 驱逐（KV cache eviction）、块级 KV 选择（blockwise KV selection）、采样/聚类/哈希式选择等等。这个站里此前拆解过的 QEvict（三档可恢复驱逐）、BLASST（复用在线 softmax 统计量的动态块稀疏）、MAC-Attention（查询间注意力复用）都属于这一类。它们与 NSA 的分水岭只有一个：这些方法都是在已经预训练好的全注意力模型上、于推理阶段事后施加稀疏性——模型本身的注意力结构没有被改变，稀疏只是部署时的一个加速技巧。2025 年 2 月，DeepSeek 与北京大学联合发布了一篇论文，把问题的定义整个换掉了——稀疏不应该是一个事后加的补丁，而应该是模型从预训练第一天起就「原生」具备的性质。这就是本文的主角 NSA（Native Sparse Attention，原生稀疏注意力），该论文获得 ACL 2025 最佳论文奖（论文 2025 年 2 月发表于 arXiv，同年 7 月在维也纳举行的 ACL 2025 上被评为四篇最佳论文之一，作者包括北大袁境阳（第一作者）、DeepSeek 的曾望鼎、梁文锋与北大的张铭教授等）。

为什么已有稀疏注意力方法不够：两个「幻觉」#

论文的第二节标题起得很尖锐：一个是「高效推理的幻觉」（The Illusion of Efficient Inference），一个是「可训练稀疏性的神话」（The Myth of Trainable Sparsity）。NSA 的设计动机正是建立在对这两点的批判之上。

幻觉一：理论稀疏没有变成实际加速。 很多方法只在部分阶段省计算。H2O 这类 KV 驱逐方法只在自回归解码阶段省内存，预填充阶段仍然要完整计算注意力图、构建驱逐索引，成本一点没省；MInference 则只针对预填充阶段做稀疏化，解码阶段退化为全注意力。也就是说，无论工作负载以哪个阶段为主（总结长文档是预填充为主，长思考链推理是解码为主），总有一个阶段还在付出全注意力的代价。更麻烦的是与先进注意力架构的兼容问题：现代模型普遍使用 GQA（Grouped-Query Attention，分组查询注意力）或 MQA（Multi-Query Attention，多查询注意力）共享 KV 头来压低解码期的内存带宽，但很多稀疏方法（论文以 Quest 为例）由每个注意力头独立选择自己的 KV 子集——在 MHA 下这没问题，在 GQA 下，同一个组内所有 query 头共享一份 KV 缓存，实际需要加载的 KV 块变成组内所有头选择结果的并集。计算是稀疏了，内存访问却依然很满。可见「理论计算减少」和「延迟下降」之间隔着硬件与架构的一整层现实。

幻觉二：稀疏性训练不起来。 更根本的问题在于，现有方法几乎都只在推理期生效，无法参与训练。原因有两类：其一，离散操作不可导。ClusterKV 依赖 k-means 聚类，MagicPIG 依赖 SimHash 哈希，聚类中心或哈希桶的选择都是计算图上的硬断点，梯度无法流经「选哪个簇/哪个桶」这一步，模型也就永远学不会最优的稀疏模式。其二，即使理论上可导，token 粒度的选择在工程上也低效。以 HashAttention 为例，它按单个 token 粒度选择，选中结果的访存是非连续的、散布在 KV 缓存各处，而 FlashAttention 之所以快，恰恰建立在连续分块访问和块状计算之上——token 粒度选择让快速注意力技巧全部失效，只能退回低硬件利用率的实现。

此外还有第三层代价：对预训练好的模型事后修剪，等于强迫它偏离已收敛的优化轨迹。模型当初是被训练成「每对 query-key 都要算」的，它的检索头（retrieval head）等信息结构依赖完整注意力，突然剪掉一部分，能力必然受损。

把这两点合起来，就能得到 NSA 的两条设计准则：

硬件对齐（hardware-aligned）：块状（blockwise）稀疏模式 + 针对 Tensor Core 和 GQA 共享 KV 结构优化的 kernel，让理论稀疏真正变成延迟下降；
原生可训练（natively trainable）：稀疏模式由端到端梯度训练学出来，而不是事后施加。

用一张表把 NSA 与推理期稀疏方法的差异收拢起来，后面几节的内容都可以对照着这张表读：

维度	H2O / Quest / MInference 等推理期方法	NSA
稀疏施加时机	预训练完成后、推理时	预训练一开始即原生生效
能否端到端训练	否（离散操作截断梯度，或只在推理期使用）	是（top-n 退化为访存操作，梯度走压缩分支）
覆盖的推理阶段	通常只有一个阶段（解码或预填充）	训练、预填充、解码全生命周期
稀疏粒度与访存	token 级随机访存，或逐 head 块选择（与 GQA 组内共享冲突）	块级连续访存，选择按 GQA 组统一
稀疏模式来源	启发式规则（重尾统计、min-max 估计）	由损失函数驱动的梯度学习

NSA 总体框架：把 KV 重映射成信息更致密的表示#

先回顾标准注意力。对输入序列，注意力函数的定义是：

\mathbf{o}_{t}=\operatorname{Attn}\left(\mathbf{q}_{t},\mathbf{k}_{:t},\mathbf{v}_{:t}\right)

其中 $\mathbf{q}_{t}$ 是第 $t$ 个位置的 query， $\mathbf{k}_{:t}$ 、 $\mathbf{v}_{:t}$ 是此前所有位置的 key 和 value，注意力函数展开为：

\operatorname{Attn}\left(\mathbf{q}_{t},\mathbf{k}_{:t},\mathbf{v}_{:t}\right)=\sum_{i=1}^{t}\frac{\alpha_{t,i}\mathbf{v}_{i}}{\sum_{j=1}^{t}\alpha_{t,j}},\quad \alpha_{t,i}=e^{\frac{\mathbf{q}_{t}^{\top}\mathbf{k}_{i}}{\sqrt{d_{k}}}}

$\alpha_{t,i}$ 是 query 与第 $i$ 个 key 之间的注意力权重（softmax 归一化前的指数形式）， $d_{k}$ 是 key 的特征维度。这里的求和遍历全部 $t$ 个历史位置——正是 $O(t^2)$ 的来源。

NSA 的基本想法是：与其让 query 面对全部 $t$ 个原始键值对，不如先把 KV 重映射成一组「更紧凑、信息更致密」的表示 $\tilde{K}_{t}$ 、 $\tilde{V}_{t}$ ，再让 query 只对这组表示做注意力：

\tilde{K}_{t}=f_{K}(\mathbf{q}_{t},\mathbf{k}_{:t},\mathbf{v}_{:t}),\quad \tilde{V}_{t}=f_{V}(\mathbf{q}_{t},\mathbf{k}_{:t},\mathbf{v}_{:t})

\mathbf{o}_{t}^{*}=\operatorname{Attn}\left(\mathbf{q}_{t},\tilde{K}_{t},\tilde{V}_{t}\right)

映射函数 $f_{K}$ 、 $f_{V}$ 是 query 相关的、动态构造的：不同的 query 会得到不同的重映射结果。NSA 一共设计了三种映射策略，分别记为 $\mathcal{C}=\{\text{cmp},\text{slc},\text{win}\}$ ——压缩（compression）、选择（selection）、滑窗（sliding window），每条策略的输出都做一次注意力，最后用可学习的门控（gating）加权融合：

\mathbf{o}_{t}^{*}=\sum_{c\in\mathcal{C}}g_{t}^{c}\cdot\operatorname{Attn}\left(\mathbf{q}_{t},\tilde{K}_{t}^{c},\tilde{V}_{t}^{c}\right)

门控分数 $g_{t}^{c}\in[0,1]$ 由输入特征经过一个小型 MLP 加 sigmoid 激活得到，每个位置 $t$ 独立计算。注意这里三个门控彼此独立、不要求加和为 1——模型的意图是让每条分支各自「按需放大或缩小」，而不是在一份固定预算里分蛋糕。

设 $N_{t}$ 为所有分支重映射后键/值的总数：

N_{t}=\sum_{c\in\mathcal{C}}\operatorname{size}[\tilde{K}_{t}^{c}]

NSA 要维持高稀疏率，即保证 $N_{t}\ll t$ 。这三条分支各司其职：压缩分支用极小代价把全局「读一遍」，选择分支把真正关键的位置「读细」，滑窗分支兜住最近的局部上下文。整体架构如图 2 所示。

图片来源：NSA 论文 Figure 2。左半部分：给定一个 query（图中紫色方块），历史 key/value 被组织成三条并行路径——压缩注意力（compressed attention，粗粒度）、选择注意力（selected attention，细粒度关键块）、滑窗注意力（sliding window，最近局部），输出经门控融合；右半部分：三条分支各自产生的注意力模式可视化，可见压缩分支覆盖全局、选择分支集中到高相关块、滑窗分支聚焦尾部。

下面逐个分支展开。

分支一：压缩注意力——花小钱把全局扫一遍#

压缩分支的直觉可以类比为「先读目录再决定细读哪一章」：用一个极低的开销建立对全序列的粗粒度认知。实现方式是把连续 token 的 key 按块聚合，用一个小网络把整块压缩成一个「摘要 key」。

形式化地，压缩后的 key 表示为：

\tilde{K}_{t}^{\text{cmp}}=f_{K}^{\text{cmp}}(\mathbf{k}_{:t})=\left\{\varphi(\mathbf{k}_{id+1:id+l})\;\middle|\;0\leqslant i\leqslant\left\lfloor\frac{t-l}{d}\right\rfloor\right\}

逐项解释： $l$ 是压缩块的长度（block length）， $d$ 是相邻块的滑动步长（sliding stride）， $\varphi$ 是一个带块内位置编码（intra-block position encoding）的可学习 MLP，负责把块内 $l$ 个 key 映射成单个压缩 key。 $\tilde{K}_{t}^{\text{cmp}}$ 的维度是 $d_{k}\times\left\lfloor\frac{t-l}{d}\right\rfloor$ ，value 侧 $\tilde{V}_{t}^{\text{cmp}}$ 完全对称。

两个设计细节值得展开。

为什么 $\varphi$ 要带块内位置编码？ 压缩操作本质上是对块内 token 的信息融合，如果 MLP 把块内 $l$ 个 key 当对称的多重集处理，块内 token 的相对顺序信息就丢失了。位置编码把「这是块内第几个 token」的信息注入每个 key，让压缩后的表示保留顺序语义。这与 RoPE 等位置编码的思路一脉相承，只是作用范围限于块内部。

为什么 $d<l$ ，即相邻压缩块要重叠？ 论文里给出的理由是缓解信息碎片化（information fragmentation）。设想一个块恰好把一段语义连贯的文本从中间切开——例如一个长单词、一个完整短语被分到两个相邻块里，如果块与块之间严格衔接（ $d=l$ ），这个语义单元在两个块里都是残缺的，各自压缩后信息都会被稀释。让相邻块重叠 $l-d$ 个 token（即步长小于块长），语义单元至少会在某一个压缩块里保持完整。论文实验配置是 $l=32$ 、 $d=16$ ，即每 16 个 token 生成一个新压缩块，相邻块重叠一半。

开销账也很清楚：每个 query 在压缩分支只需要对约 $\left\lfloor\frac{t-l}{d}\right\rfloor$ 个压缩 key 做注意力。64k 上下文时约 4096 个键，对比全注意力的 65536 个，降了一个数量级。但代价是分辨率：压缩 key 是块级摘要，丢失了块内的细粒度信息——比如「答案在 1024 号块内，但具体是哪个 token」这类精确位置信息。这正是选择分支存在的理由。

分支二：选择注意力——把关键位置读细#

选择分支的目标是以很小的额外开销，找回压缩分支丢失的细粒度信息：对每个 query，从历史 KV 中挑出最相关的若干个连续块，保留块内全部原始 token 做精确注意力。

为什么按「块」选而不是按 token 选#

这是 NSA 与很多 token 级稀疏方法的根本分歧，论文给了两层理由。

硬件层面，现代 GPU 对连续块的访问吞吐远高于随机索引读取；Tensor Core 的矩阵乘单元也是按分块组织计算的。FlashAttention 之所以能在 HBM 带宽之上再榨出性能，块状（blockwise）访存与块状计算是最基本的原则。按 token 随机挑选意味着每次读取都要跨越大段不相关数据，理论稀疏在硬件上根本兑现不了——前面提到的 HashAttention 就是反面教材。

数据层面，注意力分数的分布本身就倾向于空间连续：已有工作（如 MInference）观察到相邻 key 的重要性相近，NSA 论文自己的可视化（论文 Figure 8，画的是预训练 27B 全注意力模型的注意力图）也显示，高注意力区域呈明显的块状聚集。既然「重要的 key 往往聚在一起」，按块选择就不会漏掉太多信息，却能把稀疏模式组织成硬件友好的形状。

重要性分数：白捡的零开销信号#

按块选择的第一步是给每个块打分。直接的做法是每个块算一个「代表 key」跟 query 点积，但这意味着额外的前向计算。NSA 的巧妙之处在于：这个分数根本不需要额外算——压缩分支的注意力分数直接拿来用。

压缩分支在计算 $\operatorname{Attn}(\mathbf{q}_{t},\tilde{K}_{t}^{\text{cmp}},\tilde{V}_{t}^{\text{cmp}})$ 时，中间产物就是 query 与所有压缩 key 的 softmax 分数：

\mathbf{p}_{t}^{\text{cmp}}=\operatorname{Softmax}\left(\mathbf{q}_{t}^{T}\tilde{K}_{t}^{\text{cmp}}\right)

$\mathbf{p}_{t}^{\text{cmp}}$ 的每个元素正好对应一个压缩块的重要性。把压缩块和选择块对齐（块大小一致时两者直接等价），这个分数就自动变成了选择块的重要性分数——压缩分支本来就要算的注意力，顺手完成了选择分支的「路由」，额外开销为零。这是 NSA「先粗后精」方法论的核心：粗粒度注意力既是输出的一部分，又是细粒度选择的寻址信号。

当压缩块与选择块的分块方案不一致时（论文实验里压缩块 $l=32$ 、选择块 $l'=64$ ，确实不一致），需要按空间关系把压缩分数聚合到选择块上。设 $l\leqslant l'$ 、 $d\mid l$ 、 $d\mid l'$ ，选择块 $j$ 的重要性分数为：

\mathbf{p}_{t}^{\text{slc}}[j]=\sum_{m=0}^{\frac{l'}{d}-1}\sum_{n=0}^{\frac{l}{d}-1}\mathbf{p}_{t}^{\text{cmp}}\left[\frac{l'}{d}j-m-n\right]

直觉上，这组双重求和是在把「所有与选择块 $j$ 有空间重叠的压缩块」的分数累加。由于压缩块按步长 $d$ 滑动、选择块按步长 $l'$ 划分，一个选择块 $j$ 覆盖的位置区间 $[jl'+1,(j+1)l']$ 对应 $\frac{l'}{d}$ 个压缩块起点；又因为压缩块自身长度 $l$ 是步长 $d$ 的整数倍、相邻压缩块重叠， $\frac{l}{d}$ 项内层求和把跨块边界的信息补全。式中的下标体现了两个分块网格之间的空间对齐关系。

GQA 组内聚合：让共享 KV 的兄弟头「步调一致」#

现代模型用 GQA 让一个组内的多个 query 头共享同一份 KV 缓存，解码时 KV 只需加载一次、组内所有头复用。但如果组内各个头各自挑选不同的块，这份共享就名存实亡——需要加载的 KV 变成组内所有头选择结果的并集，带宽又涨回去了。因此 NSA 在组内把所有 query 头的重要性分数直接相加：

{\mathbf{p}_{t}^{\text{slc}}}^{\prime}=\sum_{h=1}^{H}\mathbf{p}_{t}^{\text{slc},(h)}

$h$ 是组内的头索引， $H$ 是每组 query 头的数量。聚合后的分数 $\mathbf{p}_{t}^{\text{slc}'}$ 是组内共享的：组内所有头用同一个块选择结果，加载的 KV 集合是「大家共同认可」的，而不是各自为政的并集。这个设计直接回应了前面批评 Quest 的那条缺陷——稀疏选择与 GQA 的共享 KV 结构从此兼容。顺带说明，这也是 GQA/MQA 在 NSA 里的意义所在：论文的效率分析采用 4 个 GQA 组、每组 16 个头，选择按组共享，KV 加载成本在组内被摊销。

「按 GQA 组共享选择」这一架构决策直接决定了 kernel 的访存组织方式：以组为单位把 query 加载进 SRAM，再按组共享的块索引取 KV，天然消除了重复传输。论文 Figure 3 展示了这个 kernel 的整体结构——外层按 query 位置循环（Grid Loop），内层按稀疏 KV 块循环（Inner Loop），绿色块是驻留在 SRAM 上的数据，蓝色块是需要从 HBM 读取的数据。

图片来源：NSA 论文 Figure 3。kernel 按 GQA 组加载 query（Grid Loop），内层循环按组共享的稀疏块索引抓取 KV（Inner Loop），注意力计算在 SRAM 上完成。kernel 的详细实现（group-centric data loading、共享 KV fetching、外层 grid 调度）留待下一篇拆解。

top-n 选择与细粒度注意力#

拿到聚合后的块重要性分数后，按降序排名选出前 $n$ 个块：

\mathcal{I}_{t}=\{i\mid\operatorname{rank}\left({\mathbf{p}_{t}^{\text{slc}}}^{\prime}[i]\right)\leqslant n\}

$\operatorname{rank}(\cdot)$ 表示降序排名（rank=1 对应最高分）， $\mathcal{I}_{t}$ 是被选中块的索引集合。随后把这些块内的原始 key、value 拼接起来：

\tilde{K}_{t}^{\text{slc}}=\operatorname{Cat}\left[\{\mathbf{k}_{il'+1:(i+1)l'}\mid i\in\mathcal{I}_{t}\}\right]

$\tilde{K}_{t}^{\text{slc}}\in\mathbb{R}^{d_{k}\times nl'}$ ，value 侧对称。也就是说，选择分支最终保留了 $n$ 个块、每块 $l'$ 个原始 token 的细粒度 KV，query 对这 $nl'$ 个 token 做完整精度的注意力。论文配置 $l'=64$ 、 $n=16$ ，其中 1 个初始块和 2 个局部块是固定激活的——初始块相当于注意力汇聚点（类似 StreamingLLM 的 attention sink，负责承接全序列对开头的关注），局部块保证对最近位置的连续性，剩下的 13 个块名额交给分数排名。这样每 query 在细粒度分支处理 $16\times64=1024$ 个 token。

把 64k 上下文下的完整选择流程走一遍，整个过程就具体了。设 query 位于第 $s=65536$ 个 token：压缩分支按块长 32、步长 16 生成约 4096 个压缩块，query 对这 4096 个压缩 key 做注意力，得到 4096 维的 softmax 分数向量 $\mathbf{p}_{t}^{\text{cmp}}$ ——这一步同时是压缩分支的输出计算；把分数按选择块（大小 64）聚合，64k 序列共有 $\frac{65536}{64}=1024$ 个候选选择块；组内 16 个 query 头的分数求和后，对这 1024 个块按分数降序排名，取前 16 个（其中 3 个固定激活，13 个按排名）；每个被选块内的 64 个原始 token 拼起来，得到 1024 个细粒度键参与精确注意力。整个流程中，「打分」没有花一分额外计算——它完全是压缩分支的副产品。

分支三：滑窗注意力——把局部模式隔离出去#

有了压缩和选择，还需要滑窗吗？论文给出了一个反直觉但深刻的理由：局部模式会「劫持」学习。自然语言中，语法结构、局部共指等短距离依赖的模式非常强，如果让 query 在全序列范围学习注意力，模型很容易只学会「看附近」，因为局部模式又强又容易拟合，会压制对全局结构的探索。换句话说，局部注意力像一条捷径（shortcut），模型的容量和梯度都被它吸走，压缩分支和选择分支反而学不好。

对策是让滑窗注意力独立成支：保留最近 $w$ 个 token 的原始 KV（ $\tilde{K}_{t}^{\text{win}}=\mathbf{k}_{t-w:t}$ ， $\tilde{V}_{t}^{\text{win}}=\mathbf{v}_{t-w:t}$ ），专门承接局部模式；压缩和选择分支则不用再跟局部模式竞争，专注各自的学习目标。论文配置 $w=512$ 。

为防止三条分支之间再出现捷径式的相互干扰，三个分支使用独立的 KV 投影（independent keys and values），即同一条输入在三个分支里经过三套不同的线性投影才生成各自的 key/value。这样每个分支的输出特征互不重叠，梯度不会在分支间互相污染，门控融合也才有意义——如果三分支共享 KV，门控调节的只是「看哪部分历史」，而不是「以什么视角看」。论文称这一设计「以极小的额外开销换来稳定的学习」。

这里能看到一条清晰的谱系：滑窗加固定激活的初始块，本质上就是 StreamingLLM「attention sink + local window」的结构，但 NSA 把它降级为三个分支之一，而不是全部。

门控融合：让模型自己决定信谁#

回到融合公式：

\mathbf{o}_{t}^{*}=g_{t}^{\text{cmp}}\cdot\operatorname{Attn}(\mathbf{q}_{t},\tilde{K}_{t}^{\text{cmp}},\tilde{V}_{t}^{\text{cmp}})+g_{t}^{\text{slc}}\cdot\operatorname{Attn}(\mathbf{q}_{t},\tilde{K}_{t}^{\text{slc}},\tilde{V}_{t}^{\text{slc}})+g_{t}^{\text{win}}\cdot\operatorname{Attn}(\mathbf{q}_{t},\tilde{K}_{t}^{\text{win}},\tilde{V}_{t}^{\text{win}})

三个门控 $g_{t}^{c}$ 都由 query 位置的输入特征经 MLP + sigmoid 得到，值域 $[0,1]$ ，每个位置、每个分支独立取值。为什么要学门控而不是用固定权重？因为「该信谁」随位置和任务变化：处理长距离推理时全局扫描更重要，模型可能放大压缩分支；做代码补全时最近的语法上下文更关键，滑窗分支的权重会更大；而某些位置可能三种信息都需要，三个门控同时接近 1。固定配比相当于假设所有位置的信息结构相同，显然不符合真实分布。sigmoid 而非 softmax 归一化（即不要求三者之和为 1）让门控具备独立缩放能力，模型可以自由表达「全都要」或「都不要」。

复杂度分析：一个 query 到底看几个键#

把所有分支的键数加起来，每个 query 实际处理的键总数为：

N_{t}=\left\lfloor\frac{s-l}{d}\right\rfloor+nl'+w

其中 $s$ 是已缓存的序列长度。代入论文配置（ $l=32$ 、 $d=16$ 、 $l'=64$ 、 $n=16$ 、 $w=512$ ），各长度下的键数如下表（论文 Table 4 的原始数据）：

上下文长度	全注意力键数	NSA 键数	期望加速比
8k	8192	2048	4.0×
16k	16384	2560	6.4×
32k	32768	3584	9.1×
64k	65536	5632	11.6×

以 64k 为例拆解：压缩分支约 $\left\lfloor\frac{65536-32}{16}\right\rfloor\approx4094$ 个键，选择分支 $16\times64=1024$ 个键，滑窗分支 $512$ 个键，合计约 5630（论文 Table 4 记为 5632，取整口径略有差异）——约是全注意力的 1/11.6。

这个数字的含金量要从两个阶段分别看。解码阶段是 memory-bound 的，每步延迟与 KV 加载量成正比，键数降 11.6 倍意味着解码访存近似线性地降 11.6 倍——这是「内存访问稀疏」而非「计算稀疏」带来的收益，正是论文批评过的那些只省 FLOPs 不省访存的方案给不了的。预填充和训练阶段是 compute-bound 的，注意力矩阵乘的 FLOPs 随键数线性下降，加速比同样接近理论值。还要注意一个趋势：上下文越长，收益越大。8k 时只有 4×，64k 时已到 11.6×，因为 NSA 的三个分支里，压缩分支随序列长度线性增长，选择与滑窗分支则与长度无关——稀疏率随时间自然提高。

为什么 NSA「原生可训练」：梯度是怎么穿过 top-k 的#

前面说过，稀疏方法训练不起来的根子是离散操作（聚类、哈希、top-k）切断了计算图。NSA 的选择分支里明明也有一个离散的 top-n 块选择，凭什么它能端到端训练？

关键在于看穿 top-k 在计算图中的真实角色。回顾选择分支的数据流：压缩分支的分数 $\mathbf{p}_{t}^{\text{cmp}}$ 是连续、可微的，它对全序列做了粗粒度的注意力；top-n 操作只是在这份连续分数上挑出排名靠前的索引，用来决定从内存里加载哪些 KV 块——它只影响「加载什么」，不参与任何数值计算。在前向计算图里，top-n 是一个 no-op：输出 $\tilde{K}_{t}^{\text{slc}}$ 的数值只来自被选中的块的原始 KV，而「选中了谁」这件事本身不产生梯度流经的数值依赖。反向传播时，梯度沿着压缩分支的 softmax 分数正常回传——而压缩分支覆盖了全序列、且完全可微。

所以 NSA 的「可训练性」不是靠重新发明了可微 top-k，而是靠架构设计让选择分支的离散部分变成了纯内存寻址操作。模型通过梯度在压缩分支上学习「如何评价块的重要性」，选择分支则忠实地执行这个评价。这与 SeerAttention 等用辅助损失（auxiliary loss）监督打分器的方案有本质区别：辅助损失需要额外构造监督信号（论文实验里用每个 key 块内注意力分数的均值池化作为块级标签、用 KL 散度监督预测），既增加算子开销，又可能因为辅助目标与主目标不一致而损害性能。

论文用 3B 模型做了消融（论文 Figure 7）：对比三种选择策略的训练损失曲线——辅助损失式选择（近似 SeerAttention 思路）、启发式无参数选择（Quest 式，用 query 与 key 分块的逐维 min-max 做点积打分，论文还尝试了先跑 1000 步全注意力再切换的冷启动方案）以及 NSA。结果是前两者的损失都明显高于 NSA 与全注意力基线——启发式打分召回率低，辅助损失打分优化轨迹歪。这也解释了为什么「事后可微」不够：打分器必须和主模型一起、以端到端的方式被训练。

「原生」还意味着另一件事：稀疏模式是预训练阶段从数据里学出来的，而不是在训练完成后由启发式规则硬剪出来的。后者的风险，论文引用了 MagicPIG 中的观察——top 20% 的注意力权重只覆盖约 70% 的注意力分数总和，也就是说，按分数剪枝很容易剪掉那些看似不显著、却在某些输入上扮演关键角色的「检索头」。原生训练让模型自己在损失函数的驱动下决定哪些块值得保留，稀疏模式与模型其余部分协同演化。

这套设计带来的是全生命周期收益：训练、预填充、解码三个阶段都享受稀疏性。与之相对，H2O、Quest、MInference 这些推理期方案无法参与训练，也就谈不上把稀疏模式的训练收益带回模型本身。

小结#

NSA 的核心可以压缩成一句话：用「压缩分支读全局、选择分支读关键、滑窗分支读局部」的三分支架构替代全注意力，其中选择分支的块重要性分数完全复用压缩分支的注意力中间结果，从而让离散的 top-n 选择退化为纯内存寻址操作，实现端到端原生训练。它的每一步设计都对应着一个具体问题：块状选择对应 GPU 访存连续性与注意力分数的空间聚集性，GQA 组内分数聚合对应共享 KV 的解码带宽，独立 KV 投影对应分支间的捷径学习，可学习门控对应位置相关的信息结构差异。在论文的配置下，64k 上下文每 query 只需处理 5632 个键（全注意力的 1/11.6），解码访存理论加速 11.6 倍，且加速比随序列长度单调上升。

本文是 NSA 系列的第一篇，聚焦算法与架构层。下一篇将进入工程与实验层：基于 Triton 的稀疏注意力 kernel 设计（group-centric data loading、共享 KV fetching、外层 grid 调度），前向 9×/反向 6× 的实测加速与 A100 上的评测方法，27B 模型的九项通用基准、LongBench 与 64k Needle-in-a-Haystack 实验数据，以及 NSA 思想在 DeepSeek V3.2 的 DSA（DeepSeek Sparse Attention）中的工程化落地与演进。

音乐

音乐