跳到主要内容

动态稀疏选择

用一个便宜的打分器为每个 query 挑出最相关的 key——比固定窗口准，比 dense 省

核心要点：

  • 动态稀疏按内容打分选 key，不靠固定位置
  • NSA 三分支：压缩覆盖全局、选择挑重点、滑窗顾局部
  • NSA 选择分支复用压缩分数，零额外打分且端到端可训
  • lightning indexer 低秩 + ReLU 打分，成本仅主注意力约 1.6%
  • top-k 是全局排序，分布式须全局视野

本篇是注意力机制章实现卷的**稀疏化（动态）**篇，与按位置选 key 的 07-滑窗与固定稀疏 相对。本篇是 lightning indexer + top-k selection 的唯一主词条——序列压缩篇 09 的 CSA、interconnect 的上下文并行篇都引用本篇，不重写。NSA 整体（含其压缩分支）也在本篇，09 序列压缩篇只在对比时链回。

名词定义​

名词定义
动态稀疏 (content-based sparse)用打分器按内容为每个 query 选最相关的 key，与位置无关
Lightning indexerDeepSeek 的轻量打分器：低秩 query + ReLU 加权和，头数与精度远小于主注意力
Index scoreindexer 给每个 (query, key) 算的相关性标量分数
top-k selection按 index score 选分最高的 $k$ 个 key 进主注意力
原生稀疏 (natively trainable)训练时就用稀疏（端到端可微），而非推理才套稀疏
块级 gating给每个 token 块算亲和分、选 top-k 块（MoBA）
门控融合 (gating)多分支输出按 learned 标量权重加权合并
全局视野 (global view)top-k 须看到全序列所有 key 的分数才能正确排序

@tbl-attn-dyn-glossary 本篇名词定义

动态稀疏和固定稀疏差在哪？​

动态稀疏花一点打分成本，按内容为每个 query 选 key，因此能抓固定窗口漏掉的远程相关 token。固定稀疏（07）赌"相关 key 在固定位置"，动态稀疏不赌位置——先给候选 key 打分，再选分高的。

代价是多了一个打分器和一次 top-k。本篇的核心问题就是：怎么让打分足够便宜（否则省下的主注意力又被打分吃掉）、怎么让选择端到端可训（否则只能推理时硬套，质量受损）。NSA 和 DSA 给出两条互补答案。

NSA 三分支怎么分工覆盖全局、重点和局部？​

NSA（Native Sparse Attention）对每个 query 并行跑压缩、选择、滑窗三条分支，用 learned gate 加权合并[1]。三分支分工明确：

$$\begin{equation} \mathbf{o}_t^* = \sum_{c \in \{\text{cmp},\,\text{slc},\,\text{win}\}} g_t^c \cdot \mathrm{Attn}(q_t,\, \tilde{K}_t^c,\, \tilde{V}_t^c) \label{eq:attn-dyn-nsa-gate} \end{equation}$$

每条分支有独立 sigmoid gate $g_t^c = \mathrm{sigmoid}(\mathrm{MLP}(q_t))$，三个 gate 不约束加和为 1，比 softmax 灵活。

图 4.3: NSA 对每个 query 并行跑压缩/选择/滑窗三分支，各带 sigmoid gate 加权合并（g_t^c = sigmoid(MLP(q_t))，三 gate 不约束和为 1）；红虚线为核心创新——选择分支复用压缩分支的 softmax 分数 p^cmp 挑重要块，零额外打分

压缩分支：粗粒度全局​

把历史 token 按重叠块用可学习 MLP $\varphi$ 池化成压缩 token，以 $\sim t/d$ 个 key 覆盖全序列:

$$\begin{equation} \tilde{K}_t^{\text{cmp}} = \{\varphi(\mathbf{k}_{id+1:\,id+l}) \mid 1 \le i \le \lfloor(t-l)/d\rfloor\} \label{eq:attn-dyn-nsa-cmp} \end{equation}$$

块长 $l=32$、步长 $d=16$（重叠，防块边界信息碎片化）。$\varphi$ 是 MLP + 块内位置编码，让池化感知顺序而非退化成 bag-of-words。算力从 $O(t^2)$ 降到 $O(t^2/d)$。

选择分支：复用压缩分数挑重点​

NSA 的核心创新——不另训重要性网络，直接复用压缩分支已算的 attention 分数估计块重要性[1]。压缩分支的 softmax 输出 $\mathbf{p}_t^{\text{cmp}}$ 直接（或经累加映射）当作选择块的重要性分：

$$\begin{equation} \mathbf{p}_t^{\text{slc}}[j] = \sum_{m=0}^{l'/d-1} \sum_{n=0}^{l/d-1} \mathbf{p}_t^{\text{cmp}}\!\left[\tfrac{l'}{d} j + m + n\right] \label{eq:attn-dyn-nsa-slc} \end{equation}$$

再做 blockwise top-$n$ 选出重要块（$l'=64$、$n=16$，约 1024 token）。

对比:为什么复用分数而非独立打分器?独立重要性网络（如 k-means 聚类）引入离散操作，梯度回传需 surrogate loss、训练不稳。复用 softmax 分数完全可微，无额外 loss，端到端反传[1]。

滑窗分支：防学习捷径​

滑窗分支（$w=512$）不是为了覆盖局部（选择分支已会选局部块），而是防止前两支被迫兼顾局部而分工模糊。有了专职滑窗，压缩分支专注全局、选择分支专注跨段跳跃，三者梯度不互相干扰[1]。

原生可训练 + 硬件对齐​

NSA 训练时就用稀疏，而非推理才套——这是"native"的含义。三分支全可微：压缩 MLP 梯度正常流、选择的 top-$n$ 是离散但 NSA 不对它求梯度（选择只决定"往哪做 attention",attention softmax 本身可微，选中块梯度照常回传）、gate 可微[1]。

blockwise 设计对 GPU 友好：选中的块是连续段,gather 到 SRAM 是 contiguous read 而非逐 token scatter；配 group-centric loading（GQA 下一次加载同组所有 query head）。

性能（64K 上下文 vs FlashAttention-2）:有效 KV 从 65536 降到约 5632,decode 11.6×、prefill 9.0×、训练 6.0×；质量不降反升——LongBench +0.032、AIME@16K +0.054，长上下文与数学推理超 full attention[1]。

lightning indexer 怎么便宜地给全序列打分？​

lightning indexer 是 DeepSeek 的轻量打分器，用 ReLU 加权和给每个 (query, key) 算分，成本远小于主注意力[2]。它是 DSA（V3.2）首创、CSA（V4）沿用的选择机制，本节是其唯一主词条。

"lightning"来自三处优化：

  • 低秩 query（V4 CSA 的优化）:V4 的 indexer 把 query 先降到低维 latent 再上投影出 indexer head，而非直接从 $d$ 维展开（V3.2 DSA 则直接投影）[3]:
$$\begin{equation} \mathbf{c}_t^Q = \mathbf{h}_t W^{DQ} \in \mathbb{R}^{d_c}, \qquad [\mathbf{q}_{t,1}^I; \ldots; \mathbf{q}_{t,n_h^I}^I] = \mathbf{c}_t^Q W^{IUQ} \label{eq:attn-dyn-indexer-lowrank} \end{equation}$$
  • 轻头 + 低精度:indexer 用少量头、小头维、FP8 精度，且所有 indexer head 共享一组 key（MQA 结构）;V4 CSA 取 $n_h^I=64$ 头、头维 $c^I=128$[3]。
  • ReLU 加权打分（非 softmax）:
$$\begin{equation} I_{t,s} = \sum_{h=1}^{n_h^I} w_{t,h}^I \cdot \mathrm{ReLU}\!\left(\mathbf{q}_{t,h}^I \cdot K_s^I\right) \label{eq:attn-dyn-index-score} \end{equation}$$

$w_{t,h}^I$ 是从 $\mathbf{h}_t$ 投影的可学习 per-head 标量权重。

选出 top-$k$（$k=2048$）后，主注意力只在这 $k$ 个 key 上做标准 softmax attention，复杂度 $O(n \cdot k)$。

indexer 怎么训:两阶段 continued pre-training[2]。阶段 1 dense warm-up（冻结主模型，只用主注意力分数当目标、KL 散度训 indexer 学会识别重要 token）；阶段 2 sparse training（indexer 与主模型同步训，但梯度 detach 分离——indexer 只由 indexer loss 驱动、主模型只由 LM loss 驱动，防互相干扰）。

成本与新瓶颈:indexer 仍要对全部 $n$ 个 key 算分，是 $O(n^2)$ 全量扫描，但常数极小（头数、头维、精度都远低于主注意力）。长序列下这个 $O(n^2)$ 反成新热点——TRT-LLM 为此用 radix-select 替代 torch.topk,top-k 阶段快 7.41×[4]。

为什么打分用 ReLU 不用 softmax？​

softmax 把分数归一化为和为 1 的分布，抹平了 key 之间的绝对量级差异——而 top-k 选择恰恰需要绝对量级[2]。

  • softmax 的问题:即使一批 key 的绝对相关度都很低,softmax 仍把它们归一成接近均匀的分布，无法区分"都不相关"和"都相关"。top-k 要的是"哪些 key 真正相关"的绝对信号。
  • ReLU 的优势:负值置 0、正值保原始大小，高分 key 在聚合 $I_{t,s}$ 中贡献更大,top-k 据此做有意义的全局排序;ReLU 输出非负,FP8/FP4 量化无符号位浪费；逐元素操作无需全序列归约，对 $O(n^2)$ 扫描更友好。

陷阱:ReLU 本身不产生稀疏（多 head 求和后 $I_{t,s}$ 几乎不为 0）。真正的稀疏来自 top-k selector，不是 ReLU。ReLU 只负责让分数可比。

MoBA 怎么把选择粒度放粗到块？​

MoBA（Mixture of Block Attention）把序列分块，用 block-level gating 选 top-k 个块，在选中块的 KV 上做 attention——选择粒度从单 token 放粗到块[5]。亲和分用块内 key 均值池化与 query 的内积计算：

$$\begin{equation} s_i = \langle q,\; \mathrm{mean\_pool}(K[\mathcal{I}_i]) \rangle, \qquad g_i = \begin{cases} 1 & s_i \in \mathrm{Top\text{-}k}(\{s_j\}) \\ 0 & \text{otherwise} \end{cases} \label{eq:attn-dyn-moba-gate} \end{equation}$$

只在 $g_i=1$ 的块上做 attention，当前块强制参与并施 causal mask。

放粗粒度的权衡：

维度token 级（DSA）块级（MoBA）
gating 开销$O(n)$ 每 query 算 $n$ 个分$O(n/B)$ 每 query 算 $n/B$ 个分
选择粒度单 token（细）整块 $B$ token（粗）
硬件友好top-k token 不连续，gather 散列选中块连续，复用 FlashAttention tile
表达力块内最优 token 可精确选中块内 token 全入选或全丢弃

@tbl-attn-dyn-token-vs-block token 级与块级动态选择的权衡

典型配置：块 $B=512$（8K 上下文，选 3 块）到 $B=4096$（1M 上下文，选 12 块，稀疏度 95.31%）;1M token 推理比 full attention 快 6.5×、10M 快 16×[5]。块连续使 gating 复杂度降到 $O(N^2/B)$、主注意力 $O(N \cdot k \cdot B)$。

为什么 top-k 选择必须有全局视野？​

top-k 是对全序列所有 key 的 index score 做全局排序，不知道第 $k+1$ 位之前无法确定哪 $k$ 个入选——这是动态稀疏的算子属性，也是它在分布式下不能用 ring 环传、必须先聚合的根因：

$$\begin{equation} C_t^{\text{selected}} = \{ s \mid I_{t,s} \in \mathrm{Top\text{-}k}(I_{t,\,0:n}) \} \label{eq:attn-dyn-topk-global} \end{equation}$$

这个属性在算子层面是确定性：top-k 本质是全局 argmax，要求对整条序列每个 key 都算出 $I_{t,s}$ 后才能 row-wise 选,不能提前剪枝、不能按块近似（否则退化成 MoBA 式块级选择）。分布式怎么满足全局视野（AllGather 全量 KV 或全量分数、每卡再本地选）是上下文并行的工程问题，详见 interconnect 章上下文并行篇 — 异构 Attention 下的 CP。本篇只锁定"top-k 需全局视野"这一算子事实,CP 篇引用此处。

Takeaway​

知识点核心结论
动态 vs 固定动态按内容打分选 key，抓固定窗口漏掉的远程相关，代价是打分器 + top-k
NSA 三分支压缩覆盖全局 / 选择挑重点 / 滑窗顾局部，gate 加权（不约束和为 1）
NSA 选择创新复用压缩 softmax 分数估块重要性，零额外打分 + 完全可微
NSA 原生稀疏训练即稀疏，不对 top-n 求梯度但 attention 可微；64K 下 decode 11.6× 且质量超 full
lightning indexer低秩 query + ReLU 加权 + FP8，成本约 dense 1.6%；DSA/CSA 共用
ReLU 非 softmax保留分数绝对量级利于 top-k；稀疏来自 top-k selector 非 ReLU
MoBA块级 gating（mean-pool 亲和分），开销 $O(n/B)$，硬件友好换粒度粗
top-k 全局视野top-k 是全局排序，不能提前剪枝，分布式须先聚合

@tbl-attn-dyn-takeaway 动态稀疏选择核心知识点

参考资料​

  1. Yuan et al., Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention, arXiv:2502.11089, 2025. https://arxiv.org/abs/2502.11089
  2. DeepSeek-AI, DeepSeek-V3.2, arXiv:2512.02556, 2025. https://arxiv.org/abs/2512.02556
  3. DeepSeek-AI, DeepSeek-V4 Technical Report, 2026.
  4. NVIDIA, Optimizing DeepSeek-V3.2 on NVIDIA Blackwell GPUs, TensorRT-LLM Tech Blog 15, 2025. https://nvidia.github.io/TensorRT-LLM/blogs/tech_blog/blog15_Optimizing_DeepSeek_V32_on_NVIDIA_Blackwell_GPUs
  5. Lu et al., MoBA: Mixture of Block Attention for Long-Context LLMs, arXiv:2502.13189, 2025. https://arxiv.org/abs/2502.13189