原生稀疏注意力(DeepSeek NSA)

原生稀疏注意力(DeepSeek NSA)

本节摘要:64K token 时注意力吃掉 decode 延迟的 70~80%。每个开源实验室都有修它的计划。DeepSeek 的 NSA(ACL 2025 最佳论文)是站住的那个:三条并行注意力分支——压缩粗粒度 token、选择保留细粒度 token、滑动窗口捕局部上下文——通过学习的门组合。它硬件对齐(核友好)、原生可训(预训练时就用,不是推理时贴上),64K 解码比 FlashAttention 还快,质量持平或更优。本节从零搭三分支,展示稀疏为何端到端可微——你会算 NSA vs 全注意力在 64K 的计算节省,理解「原生可训」为何是关键区分。

学习目标

阅读完本节,你应当能够:

1. 陈述 NSA 的三个注意力分支各自捕获什么。
2. 解释为什么 NSA「原生可训」,而之前的稀疏注意力只能推理时用。
3. 算 NSA vs 全注意力在 64K 上下文的计算节省(压缩块大小与选择 top-k 的函数)。
4. 在纯 Python 短合成序列上实现三分支组合,验证门控行为。

一、问题与直觉

全注意力在序列长 N 时代价 O(N²) 时间、每层 O(N) KV 缓存。64K token 时计算与内存带宽数字灾难性。NSA 论文的理论估算:64K 时注意力占总 decode 延迟 70~80%。下游一切——TTFT、token/秒、每百万 token 成本——都被注意力成本主导。

稀疏注意力是显然答案。之前的尝试分两桶:固定模式稀疏(滑动窗口、步幅、块局部)丢信息、长程召回任务失败;推理时稀疏(KV 缓存剪枝、H2O、StreamingLLM)作用在密集注意力上预训练的模型,只恢复潜在加速的一小部分,因为模型从没被要求通过稀疏模式路由信息。

原生稀疏注意力(Yuan 等,DeepSeek+PKU+UW,ACL 2025 最佳论文,arXiv:2502.11089)两者都做:预训练时模型学的稀疏模式,实现为核对齐的算法,推理时真正交付计算节省。两年后,NSA 或其直系后代是每个前沿长上下文模型的默认注意力。

三条并行分支

对每个查询,NSA 跑三次注意力,针对 KV 缓存的三种不同视图:

1. 压缩分支:token 按块大小 l(典型 32 或 64)分组,每块通过小学习 MLP 压缩成单个摘要 token,查询对这些压缩 token 注意,获得整个序列的粗粒度视图,O(N/l) 个键每查询。

2. 选择分支:用压缩分支的注意力分数识别与当前查询最相关的 top-k 块,读取这些块未压缩的细粒度 token,查询对所有这些注意。压缩分支注意力是选择的路由信号。

3. 滑动窗口分支:查询注意最近 W 个 token(典型 512)获局部上下文,捕获其他两支可能漏的结构密集短程模式(语法、局部共指)。

三分支输出通过每位置学习的门组合:out = g_cmp·out_cmp + g_sel·out_sel + g_win·out_win。门权 g_cmp、g_sel、g_win 来自查询上的小 MLP,不必和为 1,可独立加权各分支。

为什么「原生可训」

之前的稀疏注意力要么固定模式(模型无法适应),要么推理时贴上(模型预训练时用密集注意力,稀疏是部署时的近似)。NSA 在预训练时就开启稀疏——模型从头学通过三分支路由信息。关键:选择分支的 top-k 块选择是可微的(通过压缩分支注意力的软门),梯度能流回教模型哪些块该选。这让 NSA 不是近似,而是模型原生用的注意力。

计算节省

全注意力每查询 O(N) 键。NSA:压缩分支 O(N/l)、选择分支 O(k·l)(k 个块每块 l token)、滑动窗口 O(W)。64K 上下文、l=64、k=16、W=512:全注意力 65536 键,NSA 约 1024(压缩)+1024(选择)+512(窗口)≈2560 键,约 25 倍少。实测 NSA 在 64K decode 比全注意力快得多,质量持平或更优。

硬件对齐

稀疏模式选得让实际 GPU 核达理论加速——块结构契合 GPU 内存访问模式,压缩与选择用融合核避免零散内存访问。这是 NSA 相比之前稀疏方法的另一关键区分:它不只是数学上稀疏,是核友好的稀疏。

二、从零实现

Step 1:压缩分支

def compressed_branch(Q, K, V, block_size=64):
    # K、V 按 block_size 分块,每块 MLP 压成摘要
    n_blocks = K.shape[0] // block_size
    K_cmp = compress_blocks(K, n_blocks)              # (n_blocks, d),粗视图
    V_cmp = compress_blocks(V, n_blocks)
    return softmax(Q @ K_cmp.T / sqrt(d)) @ V_cmp
​

Step 2:选择分支

def selected_branch(Q, K, V, cmp_scores, block_size=64, top_k=16):
    # 用压缩分支分数选 top-k 块,读细粒度 token 注意
    top_blocks = top_k_blocks(cmp_scores, top_k)
    K_sel = gather_blocks(K, top_blocks)              # (top_k*block_size, d)
    V_sel = gather_blocks(V, top_blocks)
    return softmax(Q @ K_sel.T / sqrt(d)) @ V_sel
​

Step 3:滑动窗口分支

def sliding_window_branch(Q, K, V, window=512):
    K_win = K[-window:]; V_win = V[-window:]          # 最近 W token
    return softmax(Q @ K_win.T / sqrt(d)) @ V_win
​

Step 4:门控组合

def nsa_attention(Q, K, V):
    out_cmp, cmp_scores = compressed_branch(Q, K, V)
    out_sel = selected_branch(Q, K, V, cmp_scores)
    out_win = sliding_window_branch(Q, K, V)
    gates = branch_gate_mlp(Q)                         # (3,) 每位置三门权
    return gates[0]*out_cmp + gates[1]*out_sel + gates[2]*out_win
​

Step 5:可微性验证

验证选择分支的 top-k 选择梯度可流:压缩分支分数是软的(softmax),门可微,整个 NSA 注意力端到端可微——这是「原生可训」的数学基础。

Step 6:计算节省

def nsa_keys_vs_full(seq_len, block_size=64, top_k=16, window=512):
    full = seq_len
    nsa = seq_len//block_size + top_k*block_size + window
    return full, nsa, full/nsa
# 64K: full=65536, nsa≈2560, 约 25 倍少
​

三、框架对比

NSA 已是 DeepSeek 长上下文模型的默认注意力,DSA(DeepSeek Sparse Attention)是其在 DeepSeek 血统里的后继。对比 FlashAttention(仍是全注意力,只是核优化):NSA 在 64K decode 比它还快,因为它是真稀疏而非优化密集。对比 H2O/StreamingLLM(推理时 KV 剪枝):它们作用在密集预训练模型上,只恢复部分加速;NSA 预训练时就稀疏。对比固定模式稀疏(滑动窗口、Longformer):它们丢信息,长程召回失败;NSA 的压缩+选择分支保全局视图。

四、可复用产物

本节产出 outputs/skill-nsa-config-tuner.md——一个技能,给定模型大小、目标上下文长、硬件,推荐 NSA 的块大小 l、top-k、滑动窗口 W,算预期 decode 加速与质量影响,给出预训练集成建议。

五、练习

1. (Easy) 扩展计算节省分析:绘 NSA 键数 vs 上下文长(8K 到 256K)曲线,对比全注意力,量化加速比随上下文增长。

2. (Medium) 在针在草堆任务上对比 NSA 与全注意力:64K 草堆里藏针,测两者召回准确率,验证 NSA 质量持平。

3. (Medium) 实现块压缩的不同策略(均值池化、MLP、最大池化),对比压缩分支的质量,验证学习 MLP 最优。

4. (Hard) 验证 NSA 的端到端可微:写自动微分测试,确认选择分支的 top-k 选择梯度可流回压缩分支,这是原生可训的数学基础。

5. (Hard) 训一个小 NSA Transformer 与全注意力 Transformer(同等参数),在 8K 上下文的 WikiText 上比困惑度,再外推到 32K 测泛化,验证 NSA 的原生可训优势。

本节要点回顾

1. 64K 时注意力吃 70~80% decode:全注意力 O(N²),长上下文计算与带宽灾难。
2. 三分支并行:压缩(块摘要全局粗视图)、选择(top-k 块细粒度,路由自压缩)、滑动窗口(局部结构)。
3. 门控组合:每位置 MLP 产三门权,独立加权,不必和为 1。
4. 原生可训是关键区分:预训练时就稀疏,模型从头学路由,非推理时贴上的近似。
5. 可微的 top-k:选择分支经压缩分数的软门,梯度流回教模型选哪些块。
6. 64K 约 25 倍少键:全 65536 vs NSA≈2560(1024 压缩+1024 选择+512 窗口)。
7. 硬件对齐:块结构契合 GPU 内存访问,融合核达理论加速,这是核友好的稀疏。
8. 比 FlashAttention 还快:Flash 是优化密集,NSA 是真稀疏,64K decode 更快。
9. 对比推理时剪枝:H2O/StreamingLLM 作用在密集模型上只部分恢复加速,NSA 原生。
10. 对比固定模式稀疏:滑动窗口/Longformer 丢信息长程召回失败,NSA 压缩+选择保全局。

下一节,多 token 预测(MTP):DeepSeek-V3 让每个隐藏态同时预测多个未来 token,14B 额外参数换来更密训练信号 + 推理时的投机解码草稿。