115 lines
14 KiB
Markdown
115 lines
14 KiB
Markdown
# NSA：原生可训练、硬件对齐的稀疏注意力
|
||
|
||
> **原文标题**：Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
|
||
> **arXiv**：[2502.11089](https://arxiv.org/abs/2502.11089) ｜ **发布**：16 Feb 2025（最后修订 27 Feb 2025, v2）
|
||
> **作者/机构**：Jingyang Yuan、Huazuo Gao、Damai Dai、Junyu Luo、Liang Zhao、Zhengyan Zhang、Zhenda Xie、Y. X. Wei、Lean Wang、Zhiping Xiao、Yuqing Wang、Chong Ruan、Ming Zhang、Wenfeng Liang、Wangding Zeng；DeepSeek-AI（深度求索）联合北京大学（Peking University, PKU-Anker LLM Lab）与华盛顿大学（University of Washington）
|
||
> **原文 PDF**：[03-NSA-原生稀疏注意力-2502.11089.pdf](./03-NSA-原生稀疏注意力-2502.11089.pdf)
|
||
|
||
## 一、摘要（全文翻译）
|
||
|
||
长上下文建模对下一代语言模型至关重要，然而标准注意力机制的高昂计算开销带来了严峻的计算挑战。稀疏注意力为在保持模型能力的同时提升效率提供了颇具前景的方向。我们提出 NSA——一种原生可训练的稀疏注意力（Natively trainable Sparse Attention）机制，它把算法创新与硬件对齐（hardware-aligned）的优化结合起来，以实现高效的长上下文建模。NSA 采用动态分层稀疏策略，将粗粒度的 token 压缩与细粒度的 token 选择相结合，从而同时保持全局上下文感知能力与局部精度。我们的方法以两项关键创新推进了稀疏注意力的设计：（1）借助算术强度（arithmetic intensity）平衡的算法设计，并配合面向现代硬件的实现优化，我们取得了显著的加速；（2）我们实现了端到端训练，在不牺牲模型性能的前提下减少了预训练算力开销。如图 1 所示，实验表明：用 NSA 预训练的模型在通用基准、长上下文任务与指令式推理上均保持或超越了全注意力（Full Attention）模型。同时，在 64k 长度序列上，NSA 在解码、前向传播与反向传播各阶段相比全注意力都取得了显著加速，验证了其效率贯穿整个模型生命周期。
|
||
|
||
## 二、核心贡献
|
||
|
||
1. **提出原生可训练的分层稀疏注意力框架 NSA**：把每个查询对应的 key-value 重映射为三条并行分支——粗粒度 token 压缩（compression）、细粒度 token 选择（selection）、局部滑动窗口（sliding window），再用可学习门控加权融合，在 $N_t \ll t$ 的高稀疏度下同时保住全局感知与局部精度。
|
||
2. **以算术强度平衡为准则做硬件对齐设计**：区分「训练/预填充是算力受限、解码是访存受限」两种优化目标；实现上按 GQA 分组加载查询、组内共享稀疏 KV 块，用 Triton 写专用核，使计算削减真正转化为延迟下降，覆盖预填充、解码与训练三个阶段。
|
||
3. **稀疏模式在训练中原生学习，而非事后套用**：压缩与选择算子可反向传播，无需聚类、哈希等离散操作；27B 总参数 / 3B 激活参数模型的预训练损失全程低于全注意力基线。
|
||
4. **性能不降反升**：通用基准 9 项中 7 项超过全注意力（平均 0.456 vs 0.443），LongBench 平均分 0.469 为全部方法最高（高出全注意力 +0.032、高出 Exact-Top +0.046）。
|
||
5. **全生命周期加速与长上下文能力增强**：64k 序列下解码最高 11.6×、前向 9.0×、反向 6.0× 加速，且序列越长优势越大；64k 大海捞针（Needle-in-a-Haystack）检索全位置满分，AIME 24 上 NSA-R 相比全注意力基线在 8k/16k 生成预算下分别高出 +0.075 / +0.054。
|
||
|
||
## 三、方法与技术要点
|
||
|
||
### 3.1 为何要「原生」稀疏：对现有稀疏注意力的再思考（论文第 2 节）
|
||
|
||
**高效推理的幻觉（2.1 节）。** 一是**阶段受限**：H2O 只在解码阶段稀疏，预填充仍需昂贵的注意力图计算与索引构建；MInference 只做预填充稀疏。二者至少有一个阶段的计算量与全注意力相当，在书本摘要（预填充主导）或长思维链（解码主导）等负载下都无法全程加速。二是**与现代架构不兼容**：Quest 让每个头独立选择 KV 子集，在多头注意力（MHA）下计算与访存都稀疏；但在分组查询注意力（GQA）等共享 KV 架构下，同组各头实际需要的访存量是各自选择的**并集**——算得少了，但读得没少。
|
||
|
||
**可训练稀疏的神话（2.2 节）。** 事后（post-hoc）套用稀疏迫使模型偏离预训练轨迹：已有工作表明 top 20% 的注意力仅覆盖总注意力分数的 70%，预训练形成的「检索头」在剪枝下非常脆弱。训练侧有两类硬伤：ClusterKV 的 k-means 聚类、MagicPIG 的 SimHash 选择等**离散操作**切断计算图，梯度无法流经 token 选择；HashAttention 的 token 粒度选择虽理论可训，却要从 KV cache 加载大量彼此离散的 token，非连续访存令 FlashAttention 的分块计算失效，被迫退回低硬件利用率。
|
||
|
||
### 3.2 总体框架：三路重映射与门控融合（3.1、3.2 节）
|
||
|
||
标准注意力为 $\mathbf{o}_t=\operatorname{Attn}(\mathbf{q}_t,\mathbf{k}_{:t},\mathbf{v}_{:t})$，其中 $\alpha_{t,i}=e^{\mathbf{q}_t^{\top}\mathbf{k}_i/\sqrt{d_k}}$。NSA 的核心思路是把 $\mathbf{k}_{:t},\mathbf{v}_{:t}$ 换成更紧凑、信息密度更高的表示：
|
||
|
||
$$\tilde{K}_t=f_K(\mathbf{q}_t,\mathbf{k}_{:t},\mathbf{v}_{:t}),\qquad \tilde{V}_t=f_V(\mathbf{q}_t,\mathbf{k}_{:t},\mathbf{v}_{:t})$$
|
||
|
||
$$\mathbf{o}^{*}_t=\operatorname{Attn}(\mathbf{q}_t,\tilde{K}_t,\tilde{V}_t)=\sum_{c\in\mathcal{C}}g_t^{c}\cdot\operatorname{Attn}(\mathbf{q}_t,\tilde{K}_t^{c},\tilde{V}_t^{c})$$
|
||
|
||
映射集合仅取三种：$\mathcal{C}=\{\text{cmp},\text{slc},\text{win}\}$；门控 $g_t^{c}\in[0,1]$ 由输入特征经 MLP 加 sigmoid 得到。重映射后总长度 $N_t=\sum_{c\in\mathcal{C}}\operatorname{size}[\tilde{K}^{c}_t]$，靠 $N_t\ll t$ 维持高稀疏比。
|
||
|
||
### 3.3 三条分支的算法设计（3.3 节）
|
||
|
||
**Token 压缩（3.3.1）。** 把连续的 key/value 分块聚合为块级表示：$\tilde{K}^{\text{cmp}}_t=\{\varphi(\mathbf{k}_{id+1:id+l})\mid 0\leqslant i\leqslant\lfloor\frac{t-l}{d}\rfloor\}$，其中 $l$ 为块长、$d$ 为块间滑动步长，$\varphi$ 是带块内位置编码的可学习 MLP；通常取 $d<l$ 以缓解信息碎片化。该分支代价低，负责粗粒度全局扫描。
|
||
|
||
**Token 选择（3.3.2）。** 只靠压缩表示会丢失细粒度信息，故再挑出最重要的若干连续块。选块而非选 token 有两个理由：GPU 对连续块访问的吞吐远高于随机索引读，块计算也能充分利用 Tensor Core（FlashAttention 的基本原则）；且注意力分数具有空间连续性，相邻 key 的重要性往往相近（已有工作与本文可视化均证实）。重要性分数**几乎不额外花钱**——直接复用压缩分支产生的中间注意力分数：
|
||
|
||
$$\mathbf{p}_t^{\text{cmp}}=\operatorname{Softmax}\left(\mathbf{q}_t^{T}\tilde{K}_t^{\text{cmp}}\right)$$
|
||
|
||
当压缩块与选择块切分一致（$l'=l=d$）时，选择块重要性分数 $\mathbf{p}_t^{\text{slc}}$ 就等于 $\mathbf{p}_t^{\text{cmp}}$；切分不一致时按空间关系求和聚合：
|
||
|
||
$$\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]$$
|
||
|
||
对 GQA/MQA 这类跨头共享 KV 的模型，组内必须做**一致的块选择**，否则解码时访存量退化为并集，故先把组内各头分数相加：${\mathbf{p}_t^{\text{slc}}}'=\sum_{h=1}^{H}\mathbf{p}_t^{\text{slc},(h)}$。随后取 top-$n$ 块：
|
||
|
||
$$\mathcal{I}_t=\{i\mid\operatorname{rank}({\mathbf{p}_t^{\text{slc}}}'[i])\leqslant n\},\qquad \tilde{K}^{\text{slc}}_t=\operatorname{Cat}\left[\{\mathbf{k}_{il'+1:(i+1)l'}\mid i\in\mathcal{I}_t\}\right]$$
|
||
|
||
由于分数来自可微的注意力计算（而非聚类或哈希），选择路径的梯度可自然回流。
|
||
|
||
**滑动窗口（3.3.3）。** 局部模式学得最快，容易「短路」学习过程，使模型忽视压缩与选择分支。因此专门划出滑动窗口分支 $\tilde{K}^{\text{win}}_t=\mathbf{k}_{t-w:t},\ \tilde{V}^{\text{win}}_t=\mathbf{v}_{t-w:t}$ 显式承担局部上下文，让另两条分支专心学习各自特征。为进一步阻断捷径学习，三条分支各自持有**独立的 key/value 投影**，仅在门控处汇总——以极小开销换取「局部识别」与「长程识别」之间的梯度隔离。
|
||
|
||
### 3.4 硬件对齐的核设计（3.4 节）
|
||
|
||
压缩与滑窗分支可直接复用 FlashAttention-2 核，难点在稀疏选择分支。若沿用 FlashAttention 把时间上连续的查询块载入 SRAM 的策略，块内不同查询可能指向互不相交的 KV 块，访存低效。NSA 的关键改动是**更换查询分组方式**：对每个查询位置，把同一 GQA 组内**所有查询头**（共享同一组稀疏 KV 块）一起载入 SRAM。具体三点：
|
||
|
||
1. **组为中心的加载**：内层循环加载该组在位置 $t$ 的全部头查询 $Q\in\mathbb{R}^{[h,d_k]}$ 及其共享的稀疏块索引 $\mathcal{I}_t$；
|
||
2. **共享 KV 抓取**：按 $\mathcal{I}_t$ 顺序把连续 KV 块 $K\in\mathbb{R}^{[B_k,d_k]},V\in\mathbb{R}^{[B_k,d_v]}$ 载入 SRAM，$B_k$ 为满足 $B_k\mid l'$ 的核块大小；
|
||
3. **网格上的外层循环**：内层长度（正比于所选块数 $n$）对不同查询块几乎相同，故把查询/输出循环交给 Triton 的 grid 调度器，简化并优化核实现。
|
||
|
||
组内共享消除冗余 KV 搬运，并把负载均衡到各流式多处理器（SM），逼近最优算术强度。
|
||
|
||
## 四、主要实验结果
|
||
|
||
**预训练配置**：GQA + MoE 骨干，27B 总参数 / 3B 激活参数，30 层，隐层 2560；GQA 分 4 组、共 64 个头，$d_q=d_k=192$、$d_v=128$；MoE 采用 DeepSeekMoE，72 个路由专家 + 2 个共享专家，top-k 取 6，首层 MoE 换成 SwiGLU 形式的 MLP 以保证稳定。NSA 超参：压缩块长 $l=32$、滑动步长 $d=16$、选择块长 $l'=64$、选择块数 $n=16$（其中固定激活首个块与 2 个局部块）、窗口 $w=512$。全注意力与稀疏模型均在 8k 长度文本上预训练 270B token（引言处写作 260B token），随后在 32k 文本上用 YaRN 做继续训练与监督微调，两者都训到完全收敛。
|
||
|
||
| 基准 / 任务 | 本模型（NSA） | 对比对象 | 说明 |
|
||
|---|---|---|---|
|
||
| 通用基准平均（9 项） | 0.456 | 全注意力 0.443 | 9 项中 7 项领先 |
|
||
| MMLU / MMLU-PRO / CMMLU | 0.565 / 0.286 / 0.587 | 0.567 / 0.279 / 0.576 | 知识类基本持平或略优 |
|
||
| BBH / GSM8K / MATH / DROP | 0.521 / 0.520 / 0.264 / 0.545 | 0.497 / 0.486 / 0.263 / 0.503 | 推理类提升明显（DROP +0.042、GSM8K +0.034） |
|
||
| MBPP / HumanEval | 0.466 / 0.348 | 0.482 / 0.335 | 代码类一降一升 |
|
||
| LongBench 平均 | 0.469 | 全注意力 0.437；Exact-Top 0.423；Quest 0.392；InfLLM 0.383；H2O 0.303 | 全部方法最高 |
|
||
| HPQ / 2Wiki（多跳 QA） | 0.437 / 0.356 | 0.350 / 0.305 | 分别 +0.087 / +0.051 |
|
||
| LCC（代码理解） | 0.232 | 0.163 | +0.069 |
|
||
| PassR-en（段落检索） | 0.905 | 0.830 | +0.075 |
|
||
| 64k 大海捞针检索 | 全位置满分 | — | 分层稀疏同时覆盖全局扫描与局部精确定位 |
|
||
| AIME 24（8k / 16k 生成预算） | NSA-R 0.121 / 0.146 | Full Attention-R 0.046 / 0.092 | 8k 下 +0.075，16k 下 +0.054 |
|
||
| 训练加速（64k，Triton 同后端） | 前向 9.0× / 反向 6.0× | FlashAttention-2 | 序列越长优势越大 |
|
||
| 解码访存量（8k→64k，等价 token 数） | 2048 / 2560 / 3584 / 5632 | 全注意力 8192 / 16384 / 32768 / 65536 | 预期加速 4× / 6.4× / 9.1× / 11.6× |
|
||
|
||
通用基准上 NSA 在 9 项中 7 项超过全注意力，推理类提升最突出，说明稀疏预训练迫使模型聚焦最关键信息、过滤无关通路的噪声；仅 MBPP（−0.016）与 MMLU（−0.002）落后。长上下文侧，在与基线同等稀疏度（每查询激活 2560 token，含前 128 个 token 与 512 个局部 token，对应 NSA 在 32k 序列下的平均激活量）下取得最高 LongBench 平均分，多跳问答与代码理解提升最大，表明原生稀疏使各组件同步适配。AIME 24 上 NSA-R 在两个预算下均优于全注意力基线，16k 预算下仍有提升（0.121→0.146），表明扩展推理深度未出现灾难性遗忘。
|
||
|
||
## 五、结论、局限与影响
|
||
|
||
- **结论**：NSA 把分层 token 压缩与块级 token 选择整合进可训练架构，在保持全注意力性能的同时实现训练与推理双向加速：通用基准与基线相当，长上下文与推理能力更强，并伴随可测量的延迟下降。
|
||
- **局限性**：论文未设专门的局限性章节。可辨识的自述局限有两点：通用基准上 NSA「可能无法在较短序列上充分发挥效率优势」；注意力分数块状聚类的确切性质「仍需进一步研究」。此外，长上下文对比中作者排除了 LongBench 上所有模型得分都过低的若干子集（其无法提供有意义的比较）。
|
||
- **影响**：该工作把「稀疏注意力能否既快又可训」从推理期工程推进到预训练架构层面；其算术强度协同设计范式与压缩+选择+滑窗的三分支结构，成为后续 DeepSeek 系模型（如 DeepSeek-V3.2 中的 DSA 稀疏注意力）的直接前身；长上下文能力不再必须付出全注意力的二次方代价。
|
||
|
||
## 六、关键术语对照表
|
||
|
||
| English | 中文 |
|
||
|---|---|
|
||
| Native Sparse Attention (NSA) | 原生稀疏注意力 |
|
||
| Full Attention | 全注意力 |
|
||
| Arithmetic Intensity | 算术强度 |
|
||
| Token Compression | token 压缩（块级语义聚合） |
|
||
| Token Selection | token 选择 |
|
||
| Sliding Window | 滑动窗口 |
|
||
| Blockwise Selection | 块级选择 |
|
||
| Importance Score | 重要性分数 |
|
||
| Grouped-Query Attention (GQA) | 分组查询注意力 |
|
||
| Multi-Query Attention (MQA) | 多查询注意力 |
|
||
| Multi-Head Attention (MHA) | 多头注意力 |
|
||
| KV-cache Eviction | KV 缓存驱逐 |
|
||
| FlashAttention | FlashAttention（分块精确注意力实现） |
|
||
| Tensor Core | 张量核心 |
|
||
| Gating Mechanism | 门控机制 |
|
||
| Needle-in-a-Haystack | 大海捞针（长上下文检索测试） |
|