DeepSeek Sparse Attention（DSA）快速洞察（DeepSeek-V3.2）

DeepSeek-V3.2 技术报告中提到的 DeepSeek Sparse Attention (DSA) 是一种旨在提升上下文处理效率的稀疏注意力机制。

一、核心思想

传统 Transformer 模型的自注意力机制具有 $O(L^2)$ 的计算复杂度（其中 $L$ 是序列长度），这在处理长文本（如 128K tokens）时计算成本非常高。DSA的核心思想是：对于序列中的每一个 query token，并不需要关注所有之前的上文 token，而只需关注其中最关键的一小部分（k个）。通过这种方式，将核心注意力计算复杂度从 $O(L^2)$ 降低到了 $O(Lk)$。

二、Multi-Head Latent Attention

1. Multi-Head Attention

在标准多头注意力（MHA）中，通过投影矩阵 $W^Q,W^K,W^V\in {\mathbb{R}}^{d_h{n}_h\times d}$ 将输入的第 $t$ 个 token ${\mathbf{h}}_t\in {\mathbb{R}}^d$ 转化为 query，key 和 value（${\mathbf{q}}_t=W^Q{\mathbf{h}}_t,{\mathbf{k}}_t=W^K{\mathbf{h}}_t,{\mathbf{v}}_t=W^V{\mathbf{h}}_t,{\mathbf{q}}_t,{\mathbf{k}}_t,{\mathbf{v}}_t\in {\mathbb{R}}^{d_h{n}_h}$）。其中 $d$ 是输入 token 的嵌入维度，维度 $d_h{n}_h$ 表示在多头注意力机制中，${\mathbf{q}}_t,{\mathbf{k}}_t,{\mathbf{v}}_t$ 如何被切分为 $n_h$ 个头，每个头的维度为 $d_h$。
$$\begin{array}{r}[{\mathbf{q}}_{t,1};{\mathbf{q}}_{t,2};\cdots ;{\mathbf{q}}_{t,n_h}]={\mathbf{q}}_t,\\ [{\mathbf{k}}_{t,1};{\mathbf{k}}_{t,2};\cdots ;{\mathbf{k}}_{t,n_h}]={\mathbf{k}}_t,\\ [{\mathbf{v}}_{t,1};{\mathbf{v}}_{t,2};\cdots ;{\mathbf{v}}_{t,n_h}]={\mathbf{v}}_t,\\ {\mathbf{o}}_{t,i}=\sum _{j=1}^t{\text{Softmax}}_j(\frac{{\mathbf{q}}_{t,i}^T{\mathbf{k}}_{j,i}}{\sqrt{d_h}}){\mathbf{v}}_{j,i},\\ \mathbf{u}=W^O[{\mathbf{o}}_{t,1};{\mathbf{o}}_{t,2};\cdots ;{\mathbf{o}}_{t,n_h}]\end{array}$$
其中，${\mathbf{q}}_{t,i},{\mathbf{k}}_{t,i},{\mathbf{v}}_{t,i}\in {\mathbb{R}}^{d_h}$ 表示第 $i$ 个头的 QKV 值，$W^O\in {\mathbb{R}}^{d\times d_h{n}_h}$ 是输出投影矩阵。在推理过程中，每一个 token 需要 KV cache 的大小为 $2n_h{d}_{h}l$，$l$ 注意力层数。

KV cache 是 Transformer 模型中 MHA 采用的一种推理加速技术，通过存储中间键值对避免重复计算。注意图中展示的多个头（query）都来源于同一个 token， 每一个头都对应有自己的 $W_i^Q,W_i^K,W_i^V$。MHA 会带来高昂的内存开销，成为系统瓶颈，减少 KV 的常见思路是将不同的查询对应相同的 KV。但这类方法在性能上始终与标准多头注意力存在差距。后来在 DeepSeek-V2 模型中提出的 MLA 实现了突破，该创新方案在显著减少 KV 缓存需求的同时，反而获得了更优越的性能表现。

2. Low-Rank Key-Value Joint Compression

MLA 的核心思想是将投影矩阵分解为两个低秩矩阵：$W=W^U{W}^{DKV}$，其中 $W^{DKV}\in {\mathbb{R}}^{d_c\times d}$ 是键和值的下投影矩阵（down-projection matrix），$W^U\in {\mathbb{R}}^{d_h{n}_h\times d_c}$ 是上投影矩阵（up-projection matrix），且 $d_c\ll d_h{n}_h$。下投影矩阵将键和值压缩进一个隐向量 ${\mathbf{c}}_t^{KV}=W^{DKV}{\mathbf{h}}_t,{\mathbf{c}}_t^{KV}\in {\mathbb{R}}^{d_c}$。因为对每一个 token 而言只需要储存 ${\mathbf{c}}_t^{KV}$ 而不用储存 ${\mathbf{k}}_t$ 和 ${\mathbf{v}}_t$，所以需要的内存开销从 $2n_h{d}_{h}l$ 减少到 $d_{c}l$。而键和值通过隐向量 ${\mathbf{c}}_t^{KV}$ 计算得到：
$$\begin{gathered} {\mathbf{k}}_t^C=W^{UK}{\mathbf{c}}_t^{KV}, \\ {\mathbf{v}}_t^C=W^{UV}{\mathbf{c}}_t^{KV}, \end{gathered}$$
其中，$W^{UK},W^{UV}\in {\mathbb{R}}^{d_h{n}_h\times d_c}$ 分别表示键和值的上投影矩阵。关键之处在于，在推理过程中 $W^{UK}$ 被吸收进 $W^Q$，而 $W^{UV}$ 被吸收进 $W^O$，因此我们无需显示地计算 ${\mathbf{k}}_t^C,{\mathbf{v}}_t^C$。这里吸收的意思是：在推理前，预先将多个投影矩阵相乘合并，从而避免推理过程中显示生成中间变量，但这并不意味着训练时候只训练一个矩阵，这些投影矩阵在训练时依旧是分开训练的。
$$\begin{array}{rl}{\mathbf{q}}_i^T{\mathbf{k}}_j& =(W^Q{\mathbf{h}}_i{)}^T{W}^{UK}{\mathbf{c}}_j^{kv}\\ & ={\mathbf{h}}_i^T((W^Q{)}^T{W}^{UK}){\mathbf{c}}_j^{kv}\\ & ={\mathbf{h}}_i^{T}A{\mathbf{c}}_j^{kv}\end{array}$$

此外，在训练过程中还采用了查询的低秩压缩技术，以降低激活内存的使用。
$$\begin{gathered} {\mathbf{c}}_t^Q=W^{DQ}{\mathbf{h}}_t, \\ {\mathbf{q}}_t^C=W^{UQ}{\mathbf{c}}_t^Q. \end{gathered}$$

3. 解耦 RoPE

DeepSeek-V2 采用旋转位置编码（RoPE）：
$$\begin{array}{rl}{\mathbf{q}}_i^T{\mathbf{k}}_j& =(W^Q{\mathbf{h}}_i{)}^T{\text{RoPE}}_{\Theta ,j-i}(W^K{\mathbf{h}}_j)\\ & ={\mathbf{h}}_i^T(W^Q{)}^T{\text{RoPE}}_{\Theta ,j-i}(W^{UK}{W}^{DKV}{\mathbf{h}}_j)\end{array}$$
其中，${\text{RoPE}}_{\Theta ,j-i}(\cdot )$ 表示应用 $\text{RoPE}$ 的运算操作，$\Theta$ 是预定义的参数，而 $i,j$ 分别表示第 $i$ 和第 $j$ 个位置。它作用在 $W^{UK}{W}^{DKV}{\mathbf{h}}_j$ 这个整体上，且依赖于 $i,j$，因此，$W^{UK}$ 不会被吸收进 $W^Q$ 中，这会导致推理过程中的计算成本显著增加。

为了解决这个问题，DeepSeek-V2 提出将 RoPE 解耦为一组独立的 query 和 key：多头查询 ${\mathbf{q}}_{t,i}^R\in {\mathbb{R}}^{d_h^R}$ 和所有头共享的键 ${\mathbf{k}}_t^R\in {\mathbb{R}}^{d_h^R}$，其中 $d_h^R$ 表示解耦后查询与键的每头维度。这种解耦策略本质上会计算两组独立的注意力权重，随后将其相加。完整的多头潜在注意力（MLA）计算流程如下：
$$\begin{array}{rl}[{\mathbf{q}}_{t,1}^R;{\mathbf{q}}_{t,2}^R;\cdots ;{\mathbf{q}}_{t,n_h}^R]={\mathbf{q}}_t^R& =\text{RoPE}(W^{QR}{\mathbf{c}}_t^Q),\\ {\mathbf{k}}_t^R& =\text{RoPE}(W^{KR}{\mathbf{h}}_t),\\ {\mathbf{q}}_{t,i}& =[{\mathbf{q}}_{t,i}^C;{\mathbf{q}}_{t,i}^R],\\ {\mathbf{k}}_{t,i}& =[{\mathbf{k}}_{t,i}^C;{\mathbf{k}}_t^R],\\ {\mathbf{o}}_{t,i}& =\sum _{j=1}^t{\text{Softmax}}_j(\frac{{\mathbf{q}}_{t,i}^T{\mathbf{k}}_{j,i}}{\sqrt{d_h+d_h^R}}){\mathbf{v}}_{j,i}^C,\\ {\mathbf{u}}_t& =W^O[{\mathbf{o}}_{t,1};{\mathbf{o}}_{t,2};\cdots ;{\mathbf{o}}_{t,n_h}],\end{array}$$
其中 $W^{QR}\in {\mathbb{R}}^{d_h^R{n}_h\times d_c^{\prime }},W^{KR}\in {\mathbb{R}}^{d_h^R\times d}$ 分别表示用于生成解耦查询和键的矩阵，$[\cdot ;\cdot ]$ 表示矩阵拼接操作。在推理过程中，维度为 $d_h^R$ 的解耦键 ${\mathbf{k}}_t^R$ 也会被缓存。因此，每个 token 总共需要缓存大小 $(d_c+d_h^R)l$ 的空间。对于 DeepSeek-V2 模型，$d_c=4d_h,d_h^R=\frac{d_h}{2}$，因此每个 token 所需的 KV 缓存大小为 $\frac{9}{2}{d}_{h}l$。

二、DSA 的两大核心组件

1. 快速索引器（Lightning Indexer）
2. 细粒度 token 选择机制（Fine-grained Token Selection Mechanism）

其工作原理可以概括为：对于每一个查询 token（${\mathbf{h}}_t\in {\mathbb{R}}^d$），先用快速索引器迅速地为每一个它的前继 token（${\mathbf{h}}_s\in {\mathbb{R}}^d$）计算一个索引分数（$I_{t,s}$），然后根据这个分数，通过细粒度 token 选择机制选出最重要的 $k$ 个 token，最后只在这 $k$ 个 token 上执行标准的注意力计算。

下面我们详细拆解这两个组件。

1. 快速索引器（Lightning Indexer）

这个组件的作用是快速、粗略地评估一个查询 token ${\mathbf{h}}_t$ 与它之前的每一个 token ${\mathbf{h}}_s$ 之间的相关性。计算公式如下：
$$I_{t,s}=\sum _{j=1}^{H^I}{w}_{t,j}^I\cdot \text{ReLU}({\mathbf{q}}_{t,j}^I\cdot {\mathbf{k}}_s^I)$$
其中：

• $j=1,2\cdots ,H^I$：$j$ 表示索引器头的编号，$H^I$ 表示索引器的头数，上标 $I$ 用于区分属于索引器的参数和向量和属于主注意力模型（MLA）的参数和向量。

• ${\mathbf{q}}_t^I\in {\mathbb{R}}^{d^I{H}^I}$ 和 ${\mathbf{k}}_s^I\in {\mathbb{R}}^{d^I}$：索引器专用的查询向量和键向量，分别由 ${\mathbf{h}}_t$ 和 ${\mathbf{h}}_s$ 通过某种变换得到的，它们的维度 $d^I$ 不需要与主注意力模型中的查询/键向量维度相同。

$$\begin{array}{rl}{\mathbf{c}}_t^Q& =W^{DQ}{\mathbf{h}}_t,\\ {\mathbf{q}}_t^I& =W^{UQI}{\mathbf{c}}_t^Q=[{\mathbf{q}}_{t,1}^I;{\mathbf{q}}_{t,2}^I;\cdots ;{\mathbf{q}}_{t,H^I}^I]\\ {\mathbf{k}}_s^I& =W^{KI}{\mathbf{h}}_s\end{array}$$

• $w_t^I\in {\mathbb{R}}^{H^I}$：权重向量，同样由 ${\mathbf{h}}_t$ 衍生而来，用于对不同索引器头的输出进行加权求和。

$$w_t^I=W^{wI}{\mathbf{h}}_t$$

出于对计算效率的考量，选用 $\text{ReLU}$ 作为激活函数。鉴于快速索引器具有较少的头数且支持 FP8 精度实现，其运算效率表现尤为出色。

2. 细粒度 Token 选择机制（Fine-grained Token Selection Mechanism）

这个组件利用索引器的输出，执行真正的“稀疏”选择。

工作流程：

• 对于 ${\mathbf{h}}_t$，快速索引器为其计算出了其与所有前继 token ${\mathbf{h}}_s$ 的索引分数集合 {It,s}\{I_{t,s}\}{It,s​}。
• 从这个集合中，只选出分数最高的 $k$ 个 token。即：{hs∣It,s∈Top-k(It,:)}\{\mathbf{h}_s|I_{t,s}\in \text{Top-k}(I_{t,:})\}{hs​∣It,s​∈Top-k(It,:​)}。
• 从 Key-Value 缓存中，只取出这 $k$ 个被选中的 token 对应的 Key-value 条目，记为 {csKV}\{c_{s}^{KV}\}{csKV​}。
• 最终的注意力输出 ut=Attn(ht,{csKV∣It,s∈Top-k(It,:)})u_t=\text{Attn}(\mathbf{h}_t,\{c_s^{KV}|I_{t,s}\in \text{Top-k}(I_t,:)\})ut​=Attn(ht​,{csKV​∣It,s​∈Top-k(It​,:)})。