DeepSeek NSA：Sparse Attention 不是新概念，但这次真的做对了
DeepSeek NSA：Sparse Attention 不是新概念，但这次真的做对了
原文：Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
1. 前言
Sparse Attention 这个概念老得很——BigBird、Longformer、Reformer，2020 年那波论文基本把”哪些 token 可以不算”这件事在算法层面研究透了。
但真正好用的 sparse attention 工程实现，现在才开始出现。
这不是在贬前人，是实事求是：以前的 sparse attention 在 GPU 上的实现大多很尴尬——跳过了一些 token 的计算，但 memory access pattern 还是按 dense 来，HBM 带宽没节省，SRAM 利用率也差，结果理论 FLOPs 降了不少，实际 wall-clock time 省不了多少。
DeepSeek 这篇 NSA（Native Sparse Attention）做的事情是：把 sparse attention 的算法设计和 GPU kernel 实现一起考虑，让 sparsity 真正落地为实际加速，同时支持从头预训练而不是事后近似。
64k 长度序列上，对比 FlashAttention-2：decode 快 11.6×，forward 快 9×，backward 快 6×。而且模型质量不降反升——在通用 benchmark 平均分上超过 Full Attention baseline。
2. 老 Sparse Attention 为什么没做好
先说清楚问题在哪，这样才能理解 NSA 解决了什么。
Attention 的计算复杂度是 $O(N^2)$，序列长度翻倍，计算量翻四倍。Sparse attention 的思路是：不是所有 token 对之间的 attention 都有意义，跳过那些”不重要的”token，把 $O(N^2)$ 压下来。
算法层面没问题。但落到 GPU 上：
GPU 喜欢连续、规整的内存访问。 你告诉它”第 17、43、128、512 号 token 的 KV 我要，其他不要”，GPU 就得在 HBM 里东一块西一块地 load，内存访问变成 random access，带宽利用率极差。哪怕你 FLOPs 少了 90%，memory bound 没解决，速度未必快多少。
另一个问题是可训练性。 很多 sparse attention 是训练完一个 full attention 模型，推理时再做近似（比如 H2O、StreamingLLM）。这种方式不可能做到最优——模型本身没有针对 sparse 模式优化过，近似误差是固定存在的。
NSA 两个问题一起解决。
3. NSA 的三路并行架构
NSA 把 attention 拆成三个并行分支，每个分支处理不同粒度的信息，最后 gated 加权合并：
3.1 Compressed Attention（粗粒度全局感知）
这一路把整个序列的 KV 按块压缩——每个块的 token 用一个 MLP 压成一个代表向量，然后对这些压缩后的 KV 做 attention。
作用：保留全局上下文的粗粒度感知。模型能”看到”整个历史序列的大致走向，不会因为 sparsity 完全丢失远距离信息。
计算复杂度：从 $O(N^2)$ 降到 $O(N \cdot N/b)$，$b$ 是块大小。
3.2 Selected Attention（精细 top-n 块选择）
这一路做精细选择：先用 compressed attention 的分数估算哪些块”最重要”，然后 top-n 挑出来，对这些选定块的原始 KV 做 full attention。
作用：精确捕捉真正重要的远距离依赖。比如你在写代码，当前 token 可能需要精确地 attend 到几百行之前定义的变量——compressed 那路能粗略找到它在哪，selected 这路负责精确计算。
关键设计：选择的粒度是块而非单个 token，这是 hardware-aware 的核心。块级选择意味着 memory access 是连续的，GPU 可以高效地 load 整块 KV，而不是 random access 单个 token。
3.3 Sliding Window Attention（局部精确）
这一路对当前 token 最近的 $w$ 个 token 做完整 attention。
作用：局部连贯性。语言本质上是局部有依赖的，相邻 token 的关联最强。这路保证了不会因为 selected attention 的块选择遗漏掉紧邻的上下文。
实现上这也是最简单的一路，标准 causal attention with a fixed window，FlashAttention 直接支持。
4. Hardware-Aware Kernel 设计
三路架构是 NSA 能真正快起来的算法基础，但真正的加速来自 kernel 实现。
如下图，NSA 的 Triton kernel 设计：
两个关键设计：
1. GQA 分组共享 KV load（Grid Loop）
NSA 基于 GQA（Grouped Query Attention）架构，一组 query head 共享同一套 KV。Kernel 的 Grid Loop 按 GQA group 来遍历 query，这样同一组内的 query head 对同一块 KV 的 load 只发生一次，然后 SRAM 里的 KV 被这一组所有 query head 复用。
结果：HBM → SRAM 的 load 次数大幅减少，bandwidth 利用率提升。
2. 块级 KV 选择（Inner Loop）
Selected attention 的稀疏索引是按块存储的连续 index，Inner Loop 按这个 index 顺序 load 对应的 KV 块进 SRAM，然后做计算。因为选的是整块，load 都是顺序 coalesced access，不存在 random access 的问题。
这和”每次 load 一个 token 的 KV”有本质区别。
Table 4 的分析显示，64k 序列下 NSA 的 memory access volume 仅是 FlashAttention-2 的约 1/9，对应了 forward 9× 的实测加速。
5. 端到端可训练：不是事后近似
NSA 支持从头预训练，这一点值得单独说一下。
三路分支的 gated 合并参数、selected attention 的块选择逻辑，在训练时全部参与梯度更新。模型在训练过程中就学会了”哪些块重要、如何分配三路的权重”，而不是在推理时用启发式规则去近似一个 full attention 模型。
实验用的是 27B 总参数（3B active，MoE 架构），同等 token 预训练下，NSA 的 loss 曲线和 Full Attention 几乎重合，甚至略低。
6. 实验结果
6.1 通用 Benchmark
如下表，NSA 在大多数通用 benchmark 上超过了同等规模的 Full Attention baseline：
平均分 0.456 vs Full Attention 的 0.443。这是在模型整体计算量大幅下降的情况下做到的。核心原因：稀疏注意力减少了无意义 token 对的干扰，某种程度上相当于对 attention pattern 做了一次”去噪”。
6.2 长上下文：Needle-in-a-Haystack
64k 上下文的 Needle-in-a-Haystack 测试，NSA 全绿：
所有位置、所有深度 accuracy = 1.0。这说明 Compressed + Selected 的双路设计确实覆盖了长距离依赖，没有因为 sparsity 丢失关键信息。
6.3 实际加速
如下图，对比 Triton-based FlashAttention-2：
- 8k 序列：forward 2.1×，backward 1.1×（短序列时 sparsity 收益还不明显）
- 32k 序列：forward 6.3×，backward 3.4×
- 64k 序列：forward 9.0×，backward 6.0×
序列越长，收益越大——这正是 sparse attention 应有的表现，$O(N^2)$ vs $O(N \cdot k)$ 的差距在长序列上才显著放大。
7. 和以往 Sparse Attention 的本质区别
| 维度 | 老方案（H2O/StreamingLLM 等） | NSA |
|---|---|---|
| 训练方式 | 训练 Full Attention，推理时近似 | 原生稀疏预训练 |
| sparsity 粒度 | token 级（random access） | block 级（coalesced access） |
| 实际加速 | 理论 FLOPs 降，wall-clock 未必 | 显著真实加速（6-9×） |
| 质量 | 通常有精度损失 | 达到甚至超过 Full Attention |
最本质的区别就一句话：以前是把 dense 模型变稀疏，NSA 是把稀疏作为训练目标本身。
8. 个人 Take
这篇工作让我想起 FlashAttention 刚出来时的感觉——大家都知道 attention 很慢，也都在想各种算法层面的优化，但 FlashAttention 的核心洞察是”IO-awareness 比 FLOPs 更重要”，通过 tiling 和 online softmax 在不改变计算结果的前提下大幅减少 HBM 访问，从而真正快起来。
NSA 做的事情类似：sparse attention 的算法价值大家认可，但落地的 gap 在于 memory access pattern。NSA 用块级选择把 random access 变成 sequential access，再配合 GQA 的 KV 共享，让 GPU 能高效执行稀疏计算。
算法 + 硬件要一起考虑，这条路 DeepSeek 走得很稳。
当然，现在 NSA 的实验主要是 27B MoE 模型，在更大规模（百亿 dense 模型）和更长序列（128k+）上的表现还需要观察。另外，块选择的准确性依赖 compressed attention 的估算质量，在某些极端 attention pattern 的任务上是否会出问题，也是个值得追的问题。
但作为”把 sparse attention 真正做到生产可用”的工作，NSA 应该是 2025 年最值得关注的 attention 效率论文之一。
如果这篇文章涉及的 LLM 推理效率优化你想系统深入，可以看看我之前出版的《动手学 AutoML：从 NAS 到大语言模型优化实战》，书里有专章讲 LLM 推理效率和参数高效微调，和本文讨论的 attention 优化工程背景有直接关联。
