作者：昇腾实战派
DeepSeek知识地图：https://blog.csdn.net/weixin_45216014/article/details/156450562?spm=1011.2415.3001.5331

1. 传统的多头注意力机制

首先回顾一下标准的多头注意力机制（Multi-Head Attention，MHA），在标准的多头注意力机制下，首先需要根据隐状态$h_t$分别计算计算$q_t$、$k_t$、$v_t$（注意这里都是列向量）：

$$q_t=W^Q{h}_t$$

$$k_t=W^K{h}_t$$

$$v_t=W^V{h}_t$$

其中，$h_t\in {\mathbb{R}}^d$是attention层第t个token对应的输入，$d$是embedding的维度；$q_t,k_t,v_t\in {\mathbb{R}}^{d_h{n}_h}$，$d_h$和$n_h$分别表示每个头的维度（head dim）和头数（head num）；$W^Q,W^K,W^V\in R^{d_h{n}_h\times d}$。
然后$q_t,k_t,v_t$会被切分成$n_h$个头来进行多头注意力的计算：

$$[q_{t,1};q_{t,2};...;q_{t,n_h}]=q_t,$$

$$[k_{t,1};k_{t,2};...;k_{t,n_h}]=k_t,$$

$$[v_{t,1};v_{t,2};...;v_{t,n_h}]=v_t,$$

$$o_{t,i}=\sum _{j=1}^{t}Softma{x}_j(\frac{q_{t,i}^T{k}_{j,i}}{\sqrt{d_h}})v_{j,i},$$

$$u_t=W^O[o_{t,1};o_{t,2};...;o_{t,n_h}],$$

其中，$q_{t,i},k_{t,i},v_{t,i}\in {\mathbb{R}}^{d_h}$表示第i个attention头的query、key和value，$W^O\in {\mathbb{R}}^{d\times d_h{n}_h}$表示输出的映射矩阵。

在推理的时候由于存在KVCache，所有的key和value都需要被缓存，所以传统的多头注意力机制对于每一个token都需要缓存$2n_h{d}_{h}l$个元素，$l$表示模型的层数。在模型的部署阶段KV cache带来的显存压力是一个限制序列长度和batchsize的巨大瓶颈。

2. 推理时的KV cache

这一节简单介绍KV cache的原理。在推理阶段，最朴素的方法是每次推理把输出的最后一个token拼接到prompt的末尾作为下次推理的输入。然后用拼接之后的句子再次调用前向推理，我们可以循环调用多次推理直至结束。图2.1给出了这个过程的示意。

图 2.1 但是，不难发现，在这个流程中存在着大量的重复冗余计算。现在我们从公式出发，看一下这个过程。

• 第一步：
  对于起始输入的prompt X = { x 0 } X=\{x_0\} X={x0​}，分别经过$W^Q$、$W^K$和$W^V$的矩阵运算，获得 Q = { W Q x 0 } = { q 0 } Q=\{W^Qx_0\}=\{q_0\} Q={WQx0​}={q0​}， K = { W K x 0 } = { k 0 } K=\{W^Kx_0\}=\{k_0\} K={WKx0​}={k0​}， V = { W V x 0 } = { v 0 } V=\{W^Vx_0\}=\{v_0\} V={WVx0​}={v0​}，接着进行attention计算，$o_0=softmax(\frac{q_0^T{k}_0}{\sqrt{D}})v_0$，然后计算输出，$u_t=W_O{o}_0$，最后对$u_t$进行采样我们就获得了$x_1$。
• 第二步：
  我们把$x_1$拼接在prompt最后面，prompt变成了 X = { x 0 , x 1 } X=\{x_0 , x_1\} X={x0​,x1​}，分别经过$W^Q$、$W^K$和$W^V$的矩阵运算，获得 Q = { W Q x 0 , W Q x 1 } = { q 0 , q 1 } Q=\{W^Qx_0, W^Qx_1\}=\{q_0, q_1\} Q={WQx0​,WQx1​}={q0​,q1​}， K = { W K x 0 , W K x 1 } = { k 0 , k 1 } K=\{W^Kx_0, W^Kx_1\}=\{k_0, k_1\} K={WKx0​,WKx1​}={k0​,k1​}， V = { W V x 0 , W V x 1 } = { v 0 , v 1 } V=\{W^Vx_0, W^Vx_1\}=\{v_0, v_1\} V={WVx0​,WVx1​}={v0​,v1​}，接着进行attention计算，$o_0$的计算和前一步一模一样，$o_1=softmax(\left[\begin{array}{c}\frac{q_1^T{k}_0}{\sqrt{D}};\frac{q_1^T{k}_1}{\sqrt{D}}\end{array}\right])\left[\begin{array}{c}{v}_0\\ v_1\end{array}\right]$。

图2.2给出了两步的简单示意图，忽略了一些计算细节。

图2.2

到这里我们可以发现在第二步计算$o_1$的时候与$o_0$完全没有关系，而且只依赖于$k_0$、$k_1$、$q_0$、$q_1$，而$k_0$和$q_0$的结果在第一步已经计算过，没有必要重复计算。因此，KV cache的思想就是把每一步计算得到的$k_t$和$v_t$缓存下来，供下一步使用，避免重复计算。有了KV cache之后，我们的计算步骤就变成了图2.3这样，其中虚线表示缓存的KV cache。

图2.3

最后每个token的$k_t$和$v_t$的维度都是$n_h{d}_h$，模型总共有$l$层，因此每个token都需要缓存$2n_h{d}_{h}l$个元素。这会造成缓存的压力，所以引出了MLA这个方法。

3. 低秩的key-value联合压缩

3.1 QKV的低秩压缩

Multi-head Latent Attention（MLA）的核心是对Key和Value进行低秩联合压缩以减少KV cache：

$$c_t^{KV}=W^{DKV}{h}_t,$$

$$k_t^C=W^{UK}{c}_t^{KV},$$

$$v_t^C=W^{UV}{c}_t^{KV},$$

其中$c_t^{KV}\in {\mathbb{R}}^{d_c}$表示key和value压缩之后的隐向量；$d_c(\ll d_h{n}_h)$表示KV压缩之后的维度；$W^{DKV}\in {\mathbb{R}}^{d_c\times d}$表示下采样压缩矩阵，即将$h_t$的维度从$d$压缩到$d_c$；$W^{UK},W^{UV}\in {\mathbb{R}}^{d_h{n}_h\times d_c}$是key和value的上采样矩阵，它们将$c_t^{KV}$的维度从$d_c$恢复到$d_h{n}_h$，即恢复成多头注意力中相同的key和value。恢复之后的key和value可以继续进行attention计算。这么做的目的可以使MLA在推理的时候只需要缓存$c_t^{KV}$，因此它的KV cache只有$d_{c}l$个元素，大大少于MHA。

此外，在推理时，由于$W^{U}K$可以吸收到$W^Q$中，$W^{U}V$可以吸收到$W^O$中，我们在计算attention时甚至不需要单独计算出key和value。

除了压缩key和value，为了减少训练时的激活量，MLA对query也做了低秩压缩，即便它无法减少KV cache：

$$c_t^Q=W^{DQ}{h}_t,$$

$$q_t^C=W^{UQ}{c}_t^Q,$$

其中$c_t^Q\in {\mathbb{R}}^{d_c^{\prime }}$是query压缩后的向量；$d_c^{\prime }(\ll d_h{n}_h)$表示query压缩之后的维度；$W^{DQ}\in {\mathbb{R}}^{d_c^{\prime }\times d}$和$W^{UQ}\in {\mathbb{R}}^{d_h{n}_h\times d_c^{\prime }}$表示下采样和上采样矩阵，和KV是类似的。

3.2 无位置编码的Attention计算推导

在引入了低秩压缩之后，attention的计算也随之改变，接下来我们先在不考虑位置编码的条件下对新的attention计算进行推导。

首先由$h_t$获得query的压缩向量$c_t^Q$

$$c_t^Q=W^{DQ}{h}_t,$$

然后由$c_t^Q$上采样获得真正的query $q_t^C$，里面包含了所有的query头

$$[q_{t,1}^C;q_{t,2}^C;...;q_{t,n_h}^C]=q_t^C=W^{UQ}{c}_t^Q,$$

以同样的方式获得key和value，需要注意的是，k和v共用一个压缩向量$c_t^{KV}$，但使用不同的上采样矩阵

$$c_t^{KV}=W^{DKV}{h}_t,$$

$$[k_{t,1}^C;k_{t,2}^C;...;k_{t,n_h}^C]=k_t^C=W^{UK}{c}_t^{KV},$$

$$[v_{t,1}^C;v_{t,2}^C;...;v_{t,n_h}^C]=v_t^C=W^{UV}{c}_t^{KV},$$

接下来，计算q和k相关性，即attention分数，我们以$q_i^C$和$k_j^C$为例：

$$(q_i^C{)}^T{k}_j^C=(W^{UQ}{c}_i^Q{)}^T{W}^{UK}{c}_j^{KV}=(c_i^Q{)}^T(W^{UQ}{)}^T{W}^{UK}{c}_j^{KV}$$

到这一步我们可以发现，如果推理阶段，计算$q_i^C$和$k_j^C$的相关性时甚至不需要提前计算出并存储$k_j^C$，只需要有$c_j^{KV}$就可以了。因为推理时$(W^{UQ}{)}^T$和$W^{UK}$是已知的，我们只需要按照上面的式子从左往右计算，就可以避免计算出key本身，减少激活显存。在DeepSeek-V2原文中提到了矩阵的“吸收”，即$W^{UK}$不再是先用来和$c_j^{KV}$计算出key，而是和$(W^{UQ}{)}^T$先计算，它被$(W^{UQ}{)}^T$“吸收”了，此外还有一种方式是直接把两个矩阵合并为一个，因为推理的时候权重是已知的。这就是为什么MLA可以只缓存$c_t^{KV}$的原因。

最后我们来推导输出$u_t$：

$$\begin{gathered} u_t=W^O{o}_t, \\ =W^O\sum _{j=1}^{t}Softma{x}_j(\frac{q_t^T{k}_j}{\sqrt{d_h}})W^{UV}{c}_j^{KV}, \\ =W^O(a_{t1}{W}^{UV}{c}_1^{KV}+a_{t2}{W}^{UV}{c}_2^{KV}+,...,+a_{tt}{W}^{UV}{c}_t^{KV}), \\ =W^O{W}^{UV}(a_{t1}{c}_1^{KV}+a_{t2}{c}_2^{KV}+,...,+a_{tt}{c}_t^{KV}) \\ =W^O{W}^{UV}\sum _{j=1}^{t}Softma{x}_j(\frac{q_t^T{k}_j}{\sqrt{d_h}})c_j^{KV}, \\ =W^{OUV}\sum _{j=1}^{t}Softma{x}_j(\frac{q_t^T{k}_j}{\sqrt{d_h}})c_j^{KV} \end{gathered}$$

这里$a_{tj}$表示$q_t$和$k_j$的注意力分数，方便起见这里我们保持$q_t^T{k}_j$为原样不做展开，注意$o_t$中包含了多个头。我们可以发现在计算输出的时候同样不需要计算出并存储$v_j$，只需要$c_j^{KV}$就够了，因为$W^{UV}$也可以被$W^O$吸收。

至此，基本解释了为什么推理时可以只存储$c_j^{KV}$。

4. 解耦位置编码

目前的推导还没有引入位置编码，因此我们先尝试加入位置编码看看。位置编码同样可以表示成矩阵乘法，我们用$R_t{k}_t^C$表示对$k_t^C$应用位置编码，然后我们来尝试计算query和key的相关性。

$$(q_i^C{)}^T{R}_j{k}_j^C=(W^{UQ}{c}_i^Q{)}^T{R}_j{W}^{UK}{c}_j^{KV}=(c_i^Q{)}^T(W^{UQ}{)}^T{R}_j{W}^{UK}{c}_j^{KV}$$

从公式中可以看出如果我们对$k_j^C$应用位置编码，则公式中$W^{UK}$将与位置敏感的RoPE矩阵耦合。由于当前生成token相关的RoPE矩阵将位于$W^{UQ}$和$W^{UK}$之间，而RoPE矩阵是位置敏感的，即不同的序列位置t拥有不同的RoPE矩阵，又因为矩阵乘法不服从交换律，因此，$W^{UK}$在推理时无法再被吸收到$W^{UQ}$中，就是说$W^{UK}$、RoPE矩阵和$W^{UQ}$无法合并成一个矩阵。这意味着导致在推理过程中，我们必须对所有前缀token重新计算key，这将极大地影响推理效率。

为此，MLA解耦了位置编码，用额外的query头$q_{t,i}^R\in {\mathbb{R}}^{d_h^R}$和共享的key头$k_t^R\in {\mathbb{R}}^{d_h^R}$来携带位置编码（其实可以理解成向量中的一部分维度携带位置编码），这里$d_h^R$表示query和key中携带位置编码的每个头的维度。

解耦了位置编码之后，MLA的计算就变成了下面的过程：

首先通过$W^{QR}{c}_t^Q$生成出query中携带位置编码的部分，并计算位置编码

$$[q_{t,1}^R;q_{t,2}^R;...;q_{t,n_h}^R]=q_t^R=RoPE(W^{QR}{c}_t^Q),$$

同样计算出key携带位置编码的部分

$$k_t^R=RoPE(W^{KR}{h}_t),$$

与不带位置编码的部分进行拼接，获得真正的query和key

$$q_{t,i}=[q_{t,i}^C;q_{t,i}^R],$$

$$k_{t,i}=[k_{t,i}^C;k_t^R],$$

进行attention计算和输出的计算

$$o_{t,i}=\sum _{j=1}^{t}Softma{x}_j(\frac{q_{t,i}^T{k}_{j,i}}{\sqrt{d_h+d_h^R}})v_{j,i}^C,$$

$$u_t=W^O[o_{t,1};o_{t,2};...;o_{t,n_h}]$$

这里$W^{QR}\in {\mathbb{R}}^{d_h^R{n}_h\times d_c^{\prime }}$，$W^{KR}\in {\mathbb{R}}^{d_h^R\times d}$，这两个矩阵分别用来生成query和key中携带位置编码的部分。$W^{QR}$将$c_t^Q$的维度从$d_c^{\prime }$映射到$d_h^R{n}_h$，而$W^{KR}$直接从$h_t$生成key携带旋转位置编码的部分。RoPE(.)表示使用位置编码，即乘上矩阵$R$；[. ; .]表示拼接。**所以所谓的解耦其实就单独生成一块带位置编码的部分，然后拼接到不带位置编码的部分上。**可以发现query是每个头$q_{t,i}$都会有不同的携带位置编码的向量$q_{t,i}^R$，而key的所有头$k_{t,i}$共用同一个向量 $k_t^R$。在推理时，携带位置编码的这部分$k_t^R$也需要被缓存。

这里我们再解释一下为什么解耦了位置编码，就不再影响矩阵的吸收了。有了位置编码以后$q_{t,i}^T{k}_{j,i}$的计算变成了$[q_{t,i}^C;q_{t,i}^R]$和$[k_{t,i}^C;k_t^R]$的相乘，其结果等价于$q_{t,i}^C{k}_{t,i}^C+q_{t,i}^R{k}_t^R$，而$q_{t,i}^C{k}_{t,i}^C$就是前文分析的不带位置编码的相关性计算。

到此我们就可以引出完整的计算过程了，图中蓝色的部分是推理时需要被缓存的

最后，结合论文中的图可以更直观的理解其过程，图中标出了公式中各个矩阵乘法的位置。