DeepSeek NSA 与 DSA:从原生稀疏注意力到细粒度 token 选择
长上下文推理的核心瓶颈之一是 Attention。上下文越长,模型需要处理的历史 token 越多,Prefill 阶段的注意力计算会随序列长度快速增长;Decoding 阶段虽然每次只生成一个 token,但每一步都要访问不断增长的 KV Cache。DeepSeek 最近两条和稀疏注意力相关的工作,分别从不同角度试图解决这个问题:
- NSA(Native Sparse Attention):一篇独立的稀疏注意力论文,强调“原生可训练”和“硬件友好”。它把注意力拆成压缩、选择、滑动窗口三条路径。
- DSA(DeepSeek Sparse Attention):DeepSeek-V3.2 / V3.2-Exp 中引入的稀疏注意力机制,基于 MLA,用 Lightning Indexer 为每个 query 动态选择 Top-K 历史 KV 条目。
二者名字相近,思想也有交集,但不是同一个机制。NSA 更像一套通用的、从训练阶段就引入的分层稀疏注意力架构;DSA 更像 DeepSeek-V3.2 在 MLA 之上做的生产级长上下文稀疏化方案。
1. 为什么稀疏 Attention 仍然很难?
Full Attention 的好处是简单:每个 query 可以看见所有历史 key/value,表达能力强。但长上下文下代价很高:
- Prefill 阶段需要处理完整序列,标准 attention 的计算量随序列长度接近平方增长。
- Decode 阶段每生成一个 token 都要访问历史 KV Cache,序列越长,单步访问越重。
- 简单的固定稀疏模式,例如只看滑动窗口,虽然便宜,但会丢掉远距离的重要信息。
- 很多稀疏方法理论 FLOPs 下降,但因为随机访存、block 不对齐、训练不稳定,实际 GPU 加速不明显。
因此,一个可用的稀疏注意力机制通常要同时满足三个条件:能保留长距离关键信息、能被稳定训练、能在 GPU 上真的跑快。
2. NSA:Native Sparse Attention
NSA 的出发点是:不要把稀疏注意力当作 Full Attention 训练好之后的推理补丁,而是把它设计成一个 原生可训练(natively trainable) 的注意力结构,并且让稀疏模式天然适配 GPU 的 block 计算。
这张图回答了一个最基本的问题:稀疏之后是不是只是“跑得快但质量差”?左侧把 Full Attention 和 NSA 放在通用评测、长上下文任务、推理任务上比较,NSA 并没有因为少算 attention score 而明显掉点,平均结果甚至更好。右侧看的是 64K 长度下的系统效率,横向比较 Decoding、Forward、Backward 三个阶段。NSA 都有加速,说明它不只是 decode-time trick,而是希望训练和推理都能从稀疏结构中受益。
2.1 三条注意力路径
NSA 的核心结构是三条并行分支:Compression、Selection、Sliding Window。
这张图可以从左到右读:
- Compression 分支:把连续 token block 压缩成更少的 block-level 表示。它提供粗粒度的全局视野,让模型不用逐 token 扫完整上下文,也能捕捉长距离语义。
- Selection 分支:从历史 block 中选择重要区域,并保留更细粒度的信息。它补足 Compression 分支可能丢失的关键细节。
- Sliding Window 分支:始终保留最近一段上下文。局部上下文在语言建模里非常重要,如果没有单独的窗口分支,模型可能会把局部模式“偷懒”塞进其他分支,影响长距离稀疏模式的学习。
- Gated Output:三条分支分别得到 attention 输出后,通过门控机制融合。不同 query 可以动态调整三类信息源的权重。
右侧三块 mask 展示了三种注意力模式。绿色区域是需要计算 attention score 的位置,白色区域可以跳过。Compression 看的是压缩后的全局块,Selection 看的是被选中的重要块,Sliding Window 看的是最近局部窗口。NSA 的稀疏性不是单一规则,而是三种视角的组合。
用一个核心公式就能概括 NSA 的设计。设 $\mathcal{C}=\{\text{cmp},\text{slc},\text{win}\}$,分别表示压缩、选择和滑动窗口三条分支,NSA 的输出是三条分支 attention 结果的加权和:
$$ \mathbf{o}_t^*= \sum_{c\in\mathcal{C}} g_t^c\cdot \operatorname{Attn}(\mathbf{q}_t,\tilde{K}_t^c,\tilde{V}_t^c) $$这里 $\mathbf{q}_t$ 是当前 token 的 query;$\tilde{K}_t^c,\tilde{V}_t^c$ 是第 $c$ 条分支重新组织后的 key/value;$g_t^c\in[0,1]$ 是 gate score,由模型根据当前 token 动态决定。这个公式的直觉很简单:模型不是被迫只用一种稀疏规则,而是可以在“全局摘要”“重要细节”“最近上下文”之间分配注意力预算。
Full Attention 会让每个 query 看所有历史 KV。NSA 则让每条分支各自构造一小部分更有用的 KV。只要三条分支加起来的 KV 数量远小于历史长度,每个 query 要计算的 attention score 就会显著减少。关键是这些 KV 不是随便砍掉的,而是用三种明确角色来组织。
为了避免后面符号混淆,可以先记住这几个约定:
- $t$ 表示当前 query token 的位置,所有历史 token 都在 $:t$ 这个前缀里。
- $c$ 表示三条分支之一,取值是 $\text{cmp}$、$\text{slc}$ 或 $\text{win}$。
- $\tilde{K},\tilde{V}$ 上面的波浪线表示“重新组织后的 KV”,不是原始完整 KV。
- $g_t^c$ 是门控权重,决定当前 token 对第 $c$ 条分支的依赖程度。
- NSA 后面出现的窗口大小 $w$ 只属于 Sliding Window 分支,不要和 DSA 里的 $w_{t,j}^I$ 混淆。
2.2 三条分支分别怎么构造?
Compression 分支把连续 token block 压成一个 block-level key/value。论文里的做法可以理解为:用一个可学习的小网络 $\varphi(\cdot)$ 把长度为 $l$ 的连续 key/value block 映射成一个压缩表示;相邻 block 可以用 stride $d$ 滑动,通常 $d\lt l$,这样 block 之间有重叠,减少边界切断信息的问题。这个分支牺牲 token 级精度,换来长距离上下文的粗粒度全局视野。
对应到图里的 Compress 模块,可以写成:
$$ \tilde{K}_t^{\text{cmp}}= \left\{ \varphi(\mathbf{k}_{id+1:id+l}) \mid 0\leq i\leq \left\lfloor\frac{t-l}{d}\right\rfloor \right\} $$这个公式里的 $i$ 是 block 编号,$l$ 是每个 block 的长度,$d$ 是相邻 block 起点之间的 stride。因此 $\mathbf{k}_{id+1:id+l}$ 表示第 $i$ 个连续 token block。$\varphi$ 是把这个 block 压成一个 compressed key 的小网络。图中 Compression 分支下面那几个浅蓝色小块,就可以理解成这些压缩后的 block 表示。它们数量少,所以能提供便宜的全局视野;但它们已经不是原始 token 级别的信息,所以需要 Selection 分支补细节。
Selection 分支解决 Compression 可能丢细节的问题。它同样先把历史 KV 切成连续 block,但不是把每个 block 都压缩掉,而是给 block 打重要性分数,选出 Top-N 重要 block,再把这些 block 内的细粒度 KV 保留下来参与 attention。换句话说,Compression 负责“把全局读薄”,Selection 负责“把关键位置读细”。
可以把 Selection 抽象成两步:
$$ \mathcal{I}_t= \operatorname{TopN}\left( \operatorname{score}(\mathbf{q}_t,\operatorname{blocks}(\mathbf{k}_{:t})) \right) $$$$ \tilde{K}_t^{\text{slc}}= \operatorname{Concat}\left( \{\mathbf{k}_j\mid j\in \mathcal{I}_t\} \right) $$第一行表示:当前 query $\mathbf{q}_t$ 给历史 block 打分,选出重要 block 的索引集合 $\mathcal{I}_t$。这里的 Top-N 指“选出分数最高的 N 个 block”,N 是 selection 分支的稀疏预算,不是序列长度。第二行表示:把这些被选中 block 里的原始 key 拼起来,作为 selected attention 的输入。图中 Selection 分支里绿色高亮的 block,就是 $\mathcal{I}_t$ 选中的区域。注意它选的是连续 block,不是完全随机的单点 token,这也是它更适合 GPU 的原因。
Sliding Window 分支保留最近 $w$ 个 token,也就是 $\mathbf{k}_{t-w:t},\mathbf{v}_{t-w:t}$。这条分支看起来简单,但很重要。局部模式通常学习得更快,如果不单独给局部上下文一个分支,模型可能会让 Compression 或 Selection 分支“偷学”局部模式,反而削弱它们学习长距离信息的能力。NSA 把三种信息源隔离开,再通过 gate 融合,是为了让局部、重要细节、全局摘要各自承担清晰角色。
它的公式最简单:
$$ \tilde{K}_t^{\text{win}}=\mathbf{k}_{t-w:t},\quad \tilde{V}_t^{\text{win}}=\mathbf{v}_{t-w:t} $$也就是图中 Sliding 分支右侧那段最近窗口。它不负责长距离检索,只负责保证当前 token 永远能看到最近上下文。
2.3 为什么 NSA 要 block-wise?
稀疏注意力不能只看理论复杂度。GPU 喜欢连续、规则、块状的数据访问;如果每个 query 随机挑一堆 token,理论上少算了很多,但实际可能被 gather/scatter 和访存拖垮。NSA 因此把选择做成 blockwise,并为稀疏 pattern 设计专门 kernel。
这张图展示的是 NSA 为什么强调“硬件对齐”。绿色块表示已经搬到 SRAM 上的数据,蓝色块表示仍在 HBM 上的数据。Kernel 外层按 GQA group 组织 query,内层根据稀疏 pattern 取对应 KV block,然后尽量在 SRAM 中完成 attention 计算。这样做的目的不是单纯减少 FLOPs,而是减少高成本的 HBM 访问,并让计算以连续 block 的形式喂给 Tensor Core。
这张延迟图的横轴是上下文长度,纵轴是 kernel latency。随着序列变长,FlashAttention-2 需要处理的完整 attention 区域越来越大;NSA 只处理稀疏 block,因此增长更慢。真正值得注意的是两条曲线的间距随长度扩大而变大:这说明稀疏化在长上下文下收益更明显,也说明 NSA 的 blockwise kernel 把理论上的稀疏收益转成了实际延迟收益。
3. DSA:DeepSeek Sparse Attention
DSA 是 DeepSeek-V3.2 / V3.2-Exp 中的稀疏注意力机制。DeepSeek-V3.2 基本沿用 DeepSeek-V3.1-Terminus 的 MoE + MLA backbone,核心架构变化是加入 DSA。相比 NSA 的三分支结构,DSA 更聚焦:先用一个轻量 indexer 预测哪些历史 token 重要,再只对 Top-K 历史 KV 做主 attention。
3.1 DSA 依赖 MLA 的压缩 KV 表示
理解 DSA 前,需要先理解 DeepSeek 系列的 MLA。MLA 的核心是把 key/value cache 压缩成 latent 表示,从而降低 decode 阶段的 KV Cache 读取压力。
这张图左侧是 MHA mode,右侧是 MQA mode。MHA mode 会从压缩 latent 中恢复出每个 head 各自的 key/value,更适合训练和 prefill,因为这些阶段更偏 compute-bound,完整矩阵计算更容易吃满 GPU。MQA mode 则让多个 query heads 共享同一组 latent KV entry,更适合 decode,因为 decode 更受 KV Cache 读取约束,少读 KV 比少算一点 FLOPs 更重要。
DeepSeek-V3.2 的 DSA 是在 MLA 上实例化的。技术报告中特别提到,出于 kernel 效率考虑,DSA 基于 MLA 的 MQA mode:每个 latent KV entry 可以被多个 query heads 共享。这样主 attention 可以直接在压缩 latent KV 上做稀疏选择,而不是回到完整 per-head KV。
3.2 Lightning Indexer + Top-K Selector
DSA 的核心流程如下:
- 对当前 query,Lightning Indexer 先快速给所有历史 KV entry 打分。
- Top-K Selector 选出得分最高的 k 个历史 KV entry。
- 主 MLA attention 只在这 k 个 entry 上计算。
图中绿色路径就是 DSA 相对 MLA 增加的部分:
- Lightning Indexer 接收当前 query 相关表示和历史 key 相关表示,输出每个历史位置的 index score。
- Top-K Selector 根据这些分数选出最相关的 KV entries。
- Core Attention 不再对所有历史 KV 做 attention,而只对选中的 Top-K 集合做 Multi-Query Attention。
DSA 的核心就是让 Lightning Indexer 先回答一个便宜的问题:“当前 query 最可能需要看哪些历史位置?” 为了理解公式,先沿着图中的绿色路径看一遍数据如何产生:
- 图最下方的
Input Hidden是当前层输入的 hidden state。记当前位置 $t$ 的 hidden state 为 $\mathbf{h}_t\in\mathbb{R}^d$,历史位置 $s$ 当时进入这一层的 hidden state 为 $\mathbf{h}_s\in\mathbb{R}^d$。 - 当前的 $\mathbf{h}_t$ 经过轻量 query-side 投影,产生 $H^I$ 个 indexer query vectors 以及 $H^I$ 个 head weights。抽象地写就是:
- 每个历史 token 在它到达该层时,都会由 $\mathbf{h}_s$ 经过 key-side 投影生成一个 indexer key:
每个位置生成的 $\mathbf{k}_s^I$ 都会存入该层专用的小型 indexer key cache。因此在 decode 到位置 $t$ 时,不需要重新保存或计算所有历史 hidden states,只需要读取位置 $t$ 之前已经缓存的 indexer keys,也就是 $\mathbf{k}_1^I,\mathbf{k}_2^I,\ldots,\mathbf{k}_{t-1}^I$。
这里的 $f_Q^I$ 和 $f_K^I$ 表示 indexer 的轻量投影路径,是为了说明变量来源而写的抽象记号;它们不是论文中新增加的独立公式。实际实现中,query-side 还会复用 MLA 的 low-rank query latent。
Lightning Indexer 不是一个单头打分器,而是一个很小的多头打分网络。DeepSeek-V3.2 的公开配置里,index_n_heads=64,index_head_dim=128:它用 64 个低维 query heads 从不同子空间判断历史 token 是否重要,但所有 heads 对比的是同一份共享 indexer key。这种结构类似 MQA:多个 query 视角共享一个 key 表示,既增加打分表达能力,又避免为历史 token 缓存 64 份 indexer keys。
准备好当前 query 侧的 $\mathbf{q}_{t,j}^I,w_{t,j}^I$ 和历史 key cache 后,技术报告的 index score 写成:
$$ I_{t,s}= \sum_{j=1}^{H^I} w_{t,j}^I\cdot \operatorname{ReLU}\left( \mathbf{q}_{t,j}^I\cdot \mathbf{k}_s^I \right) $$其中 $j$ 遍历 $H^I$ 个 indexer query heads,公开配置中 $H^I=64$;$\mathbf{q}_{t,j}^I,\mathbf{k}_s^I\in\mathbb{R}^{d^I}$,公开配置中 $d^I=128$。上标 $I$ 表示这些量属于 Indexer,不是主 MLA attention 的 query/key。
对每个历史位置 $s$,第 $j$ 个 query head 先计算 $\mathbf{q}_{t,j}^I\cdot\mathbf{k}_s^I$。不同 heads 共享同一个 $\mathbf{k}_s^I$,但因为 query 投影不同,它们仍能从不同角度衡量相关性。ReLU 丢弃负相关结果,$w_{t,j}^I$ 再动态加权各个 head 的正向证据。$w_{t,j}^I$ 也是由当前 $\mathbf{h}_t$ 经 $f_Q^I$ 产生的,不是人工常数,也不是 NSA Sliding Window 的窗口大小 $w$。
所以整条公式的含义是:Lightning Indexer 先用多个低维 heads 分别给历史位置 $s$ 打分,ReLU 只保留正相关证据,再由 $w_{t,j}^I$ 把这些 head 的证据加权合并成最终 index score $I_{t,s}$。这个分数越高,历史位置 $s$ 越可能进入 Top-K。它的设计目标不是精确复现主 attention 的全部计算,而是用一个足够便宜的“预筛选器”找出值得主 attention 精算的候选 KV。
得到所有历史位置的 $I_{t,s}$ 后,DSA 只保留 Top-K 分数对应的 KV entries:
$$ \mathcal{S}_t= \{s\mid I_{t,s}\in \operatorname{TopK}(I_{t,:})\} $$主 attention 则只在这个集合上计算:
$$ \mathbf{u}_t= \operatorname{Attn} \left( \mathbf{h}_t, \{\mathbf{c}_s\mid s\in \mathcal{S}_t\} \right) $$这里 $\mathbf{h}_t$ 是当前 token 的 hidden state,$\mathbf{c}_s$ 是 MLA 中缓存的 latent KV entry。也就是说,Lightning Indexer 仍然扫描历史 token,但它只是一个轻量评分器;真正昂贵的主 MLA attention 只看 $\mathcal{S}_t$ 里的 Top-K KV。
如果序列长度是 $L$,每个 query 只选 $k$ 个历史 token,那么主 attention 的复杂度从 $O(L^2)$ 接近变成 $O(Lk)$。这里的 $k$ 是 Top-K 选择预算,不是 key 向量;在 DeepSeek-V3.2 相关资料中,常见配置是 index_topk=2048。对于 128K 上下文,2048 只占约 1.6%,长上下文下节省非常明显。
需要注意的是,DSA 并不是不看全局上下文。Lightning Indexer 仍然会为当前 query 扫描历史候选 token,只是这个评分网络很轻量,并且可以用 FP8、低 head 数和专门 kernel 实现。真正昂贵的主 attention 被限制在 Top-K 子集上。
3.3 Indexer 怎么训练?
DSA 不是硬编码一个 Top-K 规则,而是训练一个 indexer 去模仿主 attention 的选择偏好。技术报告里分成两个阶段。
第一阶段是 Dense Warm-up。此时仍然保留 dense attention,冻结主模型,只训练 Lightning Indexer。做法是把主 attention 的分数汇总成一个目标分布,再用 KL divergence 让 indexer 的分数分布去接近它。直觉上,这一步是在教 indexer:“如果 full attention 真正看了所有历史 token,它会更关注哪些位置?”
这个阶段的训练目标可以写成:
$$ \mathcal{L}^I= \sum_t \mathbb{D}_{\mathrm{KL}} \left( p_{t,:}\;\middle\|\; \operatorname{Softmax}(I_{t,:}) \right) $$其中 $p_{t,:}$ 是 dense attention 汇总出来的“老师分布”,$\operatorname{Softmax}(I_{t,:})$ 是 Lightning Indexer 对所有历史位置打分后得到的分布。KL divergence 越小,说明 indexer 越能复现 full attention 的偏好。图里的 Lightning Indexer 在这个阶段还不真正决定稀疏 attention,它先学会“打分要像主 attention”。
第二阶段是 Sparse Training。此时启用 Top-K 稀疏选择,并且主模型也一起适配稀疏模式。训练时仍然让 indexer 在被选中的 token 集合上对齐主 attention 的偏好,但主模型本身通过语言建模 loss 学会在稀疏 attention 下工作。
启用 Top-K 后,目标变成只在被选中的集合 $\mathcal{S}_t$ 上对齐:
$$ \mathcal{L}^I= \sum_t \mathbb{D}_{\mathrm{KL}} \left( p_{t,\mathcal{S}_t}\;\middle\|\; \operatorname{Softmax}(I_{t,\mathcal{S}_t}) \right) $$这和图中的 Top-k Selector 对应:进入主 attention 的只剩 $\mathcal{S}_t$ 里的 KV entries,所以 indexer 也必须在这个子集上学会排序。与此同时,主模型用语言建模 loss 适应“只能看 Top-K”的新条件。
这里的重点是:DSA 不是简单地把 full attention 截断成 Top-K,而是先让 indexer 学会模仿 dense attention 的偏好,再让主模型和选择器一起适应稀疏模式。否则 Top-K 选择器可能会在训练早期选错关键 token,导致模型能力退化。
3.4 DSA 对成本的影响
DeepSeek-V3.2 技术报告给出了与 DeepSeek-V3.1-Terminus 的成本对比。下面两张图分别是 Prefill 和 Decode 成本随 token position 增长的变化。
这两张图要看长上下文区域的斜率。蓝线是 DeepSeek-V3.1-Terminus,橙线是 DeepSeek-V3.2。Prefill 图里,蓝线随 token position 明显上升,说明长序列 prefill 仍然越来越贵;橙线增长慢很多,说明 DSA 抑制了 attention 随长度增长的成本。Decode 图里差距更直观:蓝线随着位置增长接近线性变贵,而橙线保持较低斜率,说明每个新 token 不再需要对完整历史做昂贵 attention。
DSA 不是完全消除 KV Cache,也不是让 indexer 免费。它真正改变的是成本结构:便宜的 Lightning Indexer 仍然全局扫描,昂贵的主 attention 只看学习出来的 Top-K 子集。长上下文越长,这个替换越划算。
4. NSA 与 DSA 的关系和差异
NSA 和 DSA 都试图回答同一个问题:长上下文下,模型真的需要对每个 query 看完整历史吗? 它们都认为答案是否定的,但实现方式不同。
| 维度 | NSA | DSA |
|---|---|---|
| 来源 | Native Sparse Attention 论文 | DeepSeek-V3.2 / V3.2-Exp 技术报告 |
| 稀疏粒度 | 更偏 block-wise / 分层稀疏 | 更偏 token-wise / KV entry Top-K |
| 核心结构 | Compression + Selection + Sliding Window 三分支 | Lightning Indexer + Top-K Selector + sparse MLA |
| 是否依赖 MLA | 不依赖,作为通用稀疏 attention 架构 | 是,在 DeepSeek-V3.2 中基于 MLA/MQA mode 实现 |
| 训练视角 | 强调 natively trainable 和 backward kernel | 通过继续训练让 indexer 与主模型适配 |
| 硬件重点 | blockwise sparse kernel,Tensor Core / SRAM 友好 | indexer、top-k、sparse MLA / FlashMLA 路径 |
可以把 NSA 理解为“分层稀疏注意力架构”:粗粒度压缩保留全局视野,细粒度选择保留关键细节,滑动窗口保留局部上下文。它的稀疏 pattern 更结构化,更容易做 blockwise kernel。
DSA 则更像“MLA 上的动态 token 选择器”:每个 query 根据内容动态选择 Top-K 历史 KV。它的好处是选择更细粒度、更内容相关;代价是 top-k 和稀疏 gather 对 kernel、缓存局部性和调度提出更高要求。
5. 什么时候更关注 NSA,什么时候更关注 DSA?
如果从研究和架构设计角度看,NSA 更值得关注。它系统性地讨论了稀疏注意力为什么要原生训练、为什么要 blockwise、为什么要三条分支,以及如何让 forward/backward 都加速。
如果从 DeepSeek 当前模型落地角度看,DSA 更直接。DeepSeek-V3.2 使用的是 DSA:它在已有 MLA backbone 上加 Lightning Indexer 和 Top-K token selection,目标是在长上下文场景下显著降低 prefill 和 decode 成本,同时尽量维持模型能力。
二者的共同启发是:未来的长上下文模型不会只靠扩大 KV Cache 或堆更多显存。更可能的方向是让模型自己学会“当前 token 真正需要看哪些历史信息”,并且把这个选择过程设计得足够硬件友好。
参考资料
- Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
- DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models
- DeepSeek-V3.2 Model Documentation
- DeepSeek Sparse Attention (DSA) — NVIDIA cuDNN