NSA：Native Sparse Attention
——硬件对齐且原生可训练的稀疏注意力
(Table 4)
(Figure 6)
(Table 1)
(Figure 5)
摘要
NSA 是 DeepSeek-AI 与北京大学、华盛顿大学于 2025 年 2 月发布的稀疏注意力架构。论文首先诊断既有稀疏注意力两条路线的系统性缺陷：推理阶段后验稀疏化存在"效率幻觉"（阶段受限、与 MQA/GQA 不兼容），而可训练稀疏方案普遍遭遇性能退化与反传低效。NSA 的回答是"原生稀疏"：将压缩、选择、滑窗三分支并行纳入前向计算图，用可学习门控聚合，配合专为 GQA/MQA 设计的 Triton 内核实现硬件对齐。实验上，270B token 预训练的 27B/3B-ACTIVE 模型在通用基准 9 项中 7 项超过 Full Attention（平均 0.456 vs 0.443），LongBench 平均 0.469 领先，64k 检索满分，数学 CoT 推理大幅超越全注意力基线；同时取得 64k 下最高前向 9.0×、解码 11.6× 加速。该工作确立了"稀疏注意力必须从预训练原生习得"这一设计原则。（推断：其三分支门控范式已成为后续长上下文架构的重要参照系。）
一、问题定位
1.1 计算瓶颈的现实约束
Transformer 解码的核心矛盾随上下文长度急剧放大：softmax 注意力在 64k 长度解码时占总延迟约 70–80%（p2 §1）。这意味着无论 MoE 如何压缩激活参数量、无论 KV 缓存管理多么精巧，只要注意力仍是 O(t) 全序列扫描，长上下文推理的经济性就存在硬天花板。稀疏注意力因此成为必然方向——问题不在"要不要稀疏"，而在"如何稀疏才能既快又不掉点"。
1.2 第一类缺陷：高效推理的幻觉（§2.1）
- 阶段受限稀疏：Quest、H2O 等后验稀疏化方法只在推理期做 KV 选择，预训练模型从未在稀疏注意力分布下学习过，导致分布外漂移与性能退化。
- 理论缺陷的实证：top 20% 注意力仅能覆盖约 70% 的注意力分数总分（p4 §2.2，引 Chen et al. 2024b/MagicPIG）——按单 query 视角挑"重要 token"无法逼近完整注意力质量。
- 与 MQA/GQA 不兼容：现代大模型普遍采用 GQA/MQA 共享 KV 以压内存，但逐 query head 的独立 KV 选择会在共享缓存下产生访问 union 膨胀，抵消理论收益。
1.3 第二类缺陷：可训练稀疏的迷思（§2.2）
此前尝试把稀疏模式放进训练的做法普遍失败，原因包括性能退化、训练效率需求未满足、注意力模式中存在不可训练组件（如离散 top-k 操作不可导）、以及稀疏操作带来的低效反向传播。换言之，社区不是没想过原生稀疏，而是缺少一条能同时保住"质量不掉、训练稳、硬件跑满"三条底线的工程路径。
1.4 问题重述
NSA 把问题重新表述为：设计一种从预训练第一天就原生运行的稀疏注意力，使其（a）稀疏模式可端到端学习；（b）与 GQA/MQA 架构兼容；（c）算术强度对齐 GPU 内存层次以兑现真实加速。这是一个算法-系统协同设计问题，而非单纯的注意力近似问题。
二、概念内核
2.1 原生稀疏（Native Sparsity）
论文核心概念主张是 §2.3 的 "Native Sparsity as an Imperative"，内涵三层：其一，稀疏性不是推理期的压缩技巧，而是模型能力的一部分，必须参与预训练的梯度流；其二，只有原生训练才能让模型学会"往哪里看"，这解释了为什么 NSA 在需要全局路由的多跳任务上收益最大（LongBench 多跳 QA HPQ +0.087）；其三，loss 曲线持续低于 Full Attention（Figure 4）表明原生稀疏不只是"少算"，还可能是一种隐式正则。（推断：最后一点与 MoE 文献中"稀疏带来泛化红利"的现象形成跨领域呼应，但本文未给出因果解释，应视为经验观察。）
2.2 三分支互补假设
NSA 的机制设计建立在"注意力需求天然分层"的假设上：粗粒度全局语义扫描（压缩分支）、细粒度精确检索（选择分支）、局部连续上下文（滑窗分支）覆盖不同粒度的信息需求。三分支使用相互独立的 K/V 投影并分别隔离计算，经学习门控聚合：
独立投影的设计动机是防止局部模式"走捷径"主导学习、避免梯度互相干扰（§3.3 明确陈述）。
2.3 算术强度作为一等公民
与传统"先设计算法再优化内核"不同，NSA 把 arithmetic intensity（算术强度，§3.1）作为设计输入：GPU 上随机索引读取吞吐远低于连续块访问，Tensor Core 需要 FlashAttention 式分块才能打满。因此稀疏单元被定义为"块"而非"token"——块级选择既是注意力分数空间连续性的统计事实（相邻 key 重要性相近），也是硬件友好的工程必然。算法与内核在同一次设计循环中互相约束，这是标题中 "Hardware-Aligned" 的真正含义。
2.4 门控聚合
g∈[0,1] 由输入特征经 MLP+sigmoid 产生，保持重映射后的键值总数 Nt ≪ t 以维持高稀疏率。门控让三分支的贡献成为逐 query 动态决策而非静态超参，这是"可训练"在机制层的落点。（推断：门控亦可解读为软性注意力分配器，其行为与 §6.2 可视化分析互为印证。）
2.5 概念坐标系：NSA 在谱系中的位置
把已有方法放进"决策时机 × 稀疏粒度"二维坐标系：按时机分——推理期后验（H2O、Quest、Exact-Top）vs 训练期原生（NSA）；按粒度分——逐 token（随机访问、硬件不友好）vs 块级（连续访问、Tensor Core 友好）。NSA 占据"原生 × 块级"象限，并以三分支覆盖粒度光谱——压缩分支是极粗粒度端，滑窗是不选择的细粒度局部端，选择分支居中。（推断：这一定位解释了为何对比实验同时包含两类基线——H2O 代表 token 级后验，InfLLM 代表块级检索增强。）
三、方法拆解
3.1 总体框架（§3.2）
对每个 query qt，将原始 (k,v) 重映射为三种紧凑表示后分别做注意力，再加权聚合。框架约束是在维持高稀疏率的同时让每个分支都可导、可并行、可融合进单个内核。
3.2 Token Compression（压缩分支）
以块长 l=32、步长 d=16 聚合连续 key 块，φ 为带块内位置编码的可学习 MLP，将块内 keys 映射为单个压缩 key。步长小于块长（d<l）有意为之：重叠窗口缓解信息碎片化，保证块边界处语义不被截断。该分支提供粗粒度全局语义扫描。
3.3 Token Selection（选择分支）
- 块级选择：以 l′=64 块为单位做 top-n 而非逐 token 选择。依据有二：GPU 对连续块访问吞吐远高于随机索引读取；注意力分数空间连续性意味着块内重要性相近，块级近似几乎无损。
- 重要性分数零开销复用：直接复用压缩分支中间注意力分数 p̃cmp 诱导选择块重要性（式 8–9），块划分不一致处按空间关系聚合。"一鱼两吃"消除了独立打分器的额外计算与显存开销。
- GQA/MQA 兼容：组内所有 query head 的重要性分数求和共享（式 10），强制组内一致块选择，最小化解码期 KV 加载——直接回应 §2.1 的不兼容缺陷。
- Top-n 参数：n=16，含固定激活的 1 个初始块 + 2 个局部块，保证最基本的全局锚点与局部衔接。
3.4 Sliding Window（滑窗分支）
窗口 w=512 显式承接局部上下文。语言建模中局部依赖占主导是公认事实，单独设分支避免全局分支被迫承担局部职责。
3.5 内核设计（§3.4）
基于 Triton，面向 GQA/MQA 共享 KV 架构。压缩与滑窗分支可直接复用 FlashAttention-2 内核；稀疏选择分支定制内核有三特征：
- Group-Centric Data Loading：内循环加载 GQA 组内全部 query head 进 SRAM，共享同一稀疏 KV 块索引；
- Shared KV Fetching：按索引顺序加载连续 KV 块，Bk 整除 l′；
- Outer Loop on Grid：query/output 循环放入 Triton grid 调度器，内循环长度≈选中块数 n，各 query 块几乎一致，均衡 SM 负载。
三项设计共同消除冗余 KV 传输，达成近最优算术强度。
3.6 关键超参汇总
| 参数 | 值 | 含义 |
|---|---|---|
| l_cmp | 32 | 压缩块长 |
| d | 16 | 压缩步长（d<l 重叠） |
| l′ | 64 | 选择块长 |
| n | 16 | top-n 选择块数 |
| w | 512 | 滑窗宽度 |
| 激活预算 | 2560 | 各基线统一激活 token 数（=NSA 处理 32k 序列平均激活数，另含 leading 128 + local 512，p12 §4.3） |
3.7 训练动力学视角：为什么独立投影重要
若三分支共享投影，滑窗分支承载的高频局部梯度会持续流过共享参数，把全局分支的键空间拉向局部模式——模型退化为"昂贵的滑窗"。独立隔离后各分支键值子空间各自特化，门控在聚合层学习分工。（推断：这与多任务学习中"共享底层易被主任务劫持、参数隔离+门控融合更稳"的经验一致；Figure 4 的 loss 优势可能部分源于此设计避免了训练早期的模式竞争。）此外，压缩分支的中间分数被选择分支复用，两分支虽投影独立但信息通路耦合，形成"先粗扫、后精选"的级联语义。
四、证据链
4.1 训练设置与 loss 证据
骨干为 27B 总参/3B 激活的 GQA+MoE（DeepSeekMoE：72 routed + 2 shared experts，top-6）；30 层，hidden 2560，GQA 4 组 64 heads，d_q=d_k=192，d_v=128；首层 MoE 换 SwiGLU MLP 保稳。以 270B tokens、8k 长度预训练至完全收敛，再用 YaRN 于 32k 续训。证据链第一环：NSA 的 loss 曲线稳定下降且持续低于 Full Attention（Figure 4）——原生稀疏不仅可行，还有质量增益。
4.2 通用基准（Table 1）
平均 0.456 vs Full Attention 0.443，9 项中 7 项领先：
| 基准 | FullAttn | NSA | Δ |
|---|---|---|---|
| MMLU | 0.567 | 0.565 | -0.002 |
| MMLU-PRO | 0.279 | 0.286 | +0.007 |
| CMMLU | 0.576 | 0.587 | +0.011 |
| BBH | 0.497 | 0.521 | +0.024 |
| GSM8K | 0.486 | 0.520 | +0.034 |
| MATH | 0.263 | 0.264 | +0.001 |
| DROP | 0.503 | 0.545 | +0.042 |
| MBPP | 0.482 | 0.466 | -0.016 |
| HumanEval | 0.335 | 0.348 | +0.013 |
（推断：知识记忆型任务持平、多步推理型任务提升的模式，与"稀疏迫使模型学会选择性路由"的解释一致。）
4.3 长上下文（Table 2 / Figure 5）
64k Needle-in-a-Haystack 全位置 100% 检索准确率。LongBench 排名：NSA 0.469 > FullAttn 0.437 > Exact-Top 0.423 > Quest 0.392 > InfLLM 0.383 > H2O 0.303。NSA 超 Full Attention +0.032、超 Exact-Top +0.046；子项 HPQ +0.087（0.437 vs 0.350）、2Wiki +0.051、LCC +0.069（0.341 vs 0.294）、PassR-en +0.075（0.232 vs 0.163）。作者归因于原生端到端稀疏预训练 + 分层机制兼顾全局/局部。值得注意：Exact-Top 用真实注意力分数做 oracle 式选择仍输给 NSA 0.046——构成"原生训练 > 后验选择"命题最有力的对照证据。
4.4 复杂推理（Table 3）
因小模型 RL 效果有限，改用 DeepSeek-R1 蒸馏，SFT 10B tokens 的 32k 数学推理轨迹。AIME24（temperature 0.7、top-p 0.95、每题 16 采样取平均）：生成上限 8k 时 NSA-R 0.121 vs FullAttention-R 0.046（+0.075），16k 时 0.146 vs 0.092（+0.054）。长思维链场景下注意力路由质量被急剧放大——差距超过通用基准一个数量级。
4.5 效率（§5 / Figure 6 / Table 4）
8-GPU A100、与同后端 Triton 版 FlashAttention-2 对比：前向最高 9.0×、反向 6.0×（64k；随长度单调增强，8k 时仅约 1.1×/2.1×）。解码为 memory-bound，加速比近似线性于内存访问缩减：
| 上下文长度 | NSA 等效加载 token | 期望加速 |
|---|---|---|
| 8k | 2048 | 4× |
| 16k | 2560 | 6.4× |
| 32k | 3584 | 9.1× |
| 64k | 5632 | 11.6× |
（对照：Full Attention 需全长加载；Figure 1 右图同标 11.6×。）
4.6 反面证据与讨论（§6）
作者披露曾尝试 ClusterKV 类聚类方案并放弃，三重瓶颈：动态聚类开销、MoE 专家并行负载不均放大失衡、强制周期性重聚类与分块顺序训练协议冲突。这段失败记录本身是重要证据：原生稀疏并非任意可训练方案都成立，硬件-系统协同约束会淘汰大量理论上优雅的设计。§6.2 可视化证实注意力分数空间连续性是块选择的分布基础。
五、批判性分析
5.1 规模外推的未验证性
全部训练实验止步于 27B/3B-ACTIVE、270B tokens——相对前沿模型是小尺度。更大模型的注意力头数更多、GQA 组更宽、任务混合更复杂，三分支门控是否仍能学到同样干净的分工未有证据。（推断：以 DeepSeek 的工程文化，V3/R1 世代若已内部验证类似机制，论文刻意不提的可能性不能排除，但公开证据层面此为空白。）
5.2 CoT 实验的替代解释
AIME24 的巨大差距来自蒸馏 SFT 路线而非从头 RL，且每题 16 采样取平均。教师模型 R1 的注意力结构是稠密的；学生 NSA 模仿超长思维链反而更好，可能解释包括：稀疏正则抑制对无关上下文的过拟合、或全注意力基线在小模型容量下更易被长序列淹没。论文未做消融区分这些假设，0.121 vs 0.046 的因果归属不完全清晰。
5.3 评测口径的选择性
LongBench 因低分剔除了部分子集（p12）；基线稀疏度按"激活 token 数对齐"而非 FLOPs 或 wall-clock 严格对齐——激活数对齐对 NSA 略有利，其压缩分支重映射 KV 更小但计算模式不同，等 token 不等于等计算。此外通用基准多为短文本，作者自认 NSA 在短序列上无法充分发挥效率优势（p11）——质量结论主要来自短文本、效率结论主要来自长文本，两者的适用域并未完全重合。
5.4 超参的手工性与敏感性黑箱
l_cmp/d/l′/n/w = 32/16/64/16/512 全部人工设定，论文未提供自适应机制或系统敏感性消融。这些超参间存在耦合（n 个 l′-块的总预算与 w 的比例决定局部/全局资源分配），换到不同模型宽度或数据分布下是否需要重新调参无从判断。（推断：n=16 × l′=64 ≈ 1024 的选择预算加 w=512 落在"千级全局+半千级局部"的经验区间，这种整数美感暗示超参经过较多手工迭代而非原理推导。）
5.5 硬件边界的局限
不支持 MHA 解码（memory-intensive，p8 §3.4）；内核对比限于 A100 + Triton 后端，未覆盖 H100/H800 及其他加速器。A100 上 memory-bound 的加速逻辑在新代际硬件（更大 SRAM、更高带宽比）下是否保持同样曲线存疑。（推断：Group-Centric 加载策略对 SRAM 容量的依赖可能在更大 KV 头维度下触及上限。）
5.6 与后验方法的公平性问题
Exact-Top 使用真实注意力分数仍输给 NSA，是支持原生训练的最强论据；但 Quest/InfLLM/H2O 等基线均为推理期方法，它们的预训练在全注意力下完成——比较对象没有获得"同等原生训练机会"。严格的对照应包含其他原生可训练稀疏方案的同类实现，论文对此类基线的缺失未作充分说明。
六、启示与影响
6.1 设计原则的确立
NSA 最持久的贡献是把"原生稀疏"从尝试性做法升格为必要性原则：后验稀疏化的理论缺口（top 20% 仅覆盖 70% 分数质量）与 Exact-Top 对照实验共同封死了"推理期补救"路线的主导地位。后续长上下文工作都需回答"你的稀疏模式是否经过原生训练"这个问题。
6.2 算法-系统协同设计范式的样板
论文展示了把 arithmetic intensity 当设计输入、让算法结构服从内存层次的完整方法论：块级稀疏单元、组中心数据加载、grid 外循环调度，每一步都是统计性质（分数空间连续性）与硬件性质（连续访问、Tensor Core 分块）的双重约束求解。（推断：Triton 而非 CUDA C++ 的选型降低了社区复现门槛，与其开源生态策略一致。）
6.3 对推理经济的直接影响
11.6× 的 64k 解码期望加速若在产品级模型兑现，直接改变长上下文 API 的成本结构与最大盈利上下文窗口设定。（推断：DeepSeek 后续产品的长上下文定价策略与本文技术储备之间存在合理关联，但属商业推断。）解码 memory-bound 的本质意味着此类收益随带宽瓶颈长期有效。
6.4 开放问题
- 超参自适应化：块长/窗口随层深或数据动态调整；
- 与 MoE 负载均衡的联合设计（ClusterKV 失败教训表明两者耦合）；
- RL 阶段而非仅蒸馏 SFT 下的表现；
- MHA 及新硬件的支持矩阵；
- 多模态/检索增强等 KV 异构场景中三分支门控的行为。
6.5 一句话定位
NSA 用一份"算法+内核+训练配方"的全栈证明回答了一个悬而未决的问题：稀疏注意力可以在不牺牲甚至小幅提升质量的前提下兑现接近线性的理论收益——前提是把稀疏性当作模型的原生能力来训练，而不是当作推理期的近似技巧来补救。