大模型推理加速：看图学KV Cache

KV Cache是Transformer标配的推理加速功能，transformer官方use_cache这个参数默认是True，但是它只能用于Decoder架构的模型，这是因为Decoder有Causal Mask，在推理的时候前面已经生成的字符不需要与后面的字符产生attention，从而使得前面已经计算的K和V可以缓存起来。

我们先看一下不使用KV Cache的推理过程。假设模型最终生成了“遥遥领先”4个字。

当模型生成第一个“遥”字时，input=”<s>“, “<s>“是起始字符。Attention的计算如下：

为了看上去方便，我们暂时忽略scale项 $\sqrt{d}$， 但是要注意这个scale面试时经常考。

如上图所示，最终Attention的计算公式如下，（softmaxed 表示已经按行进行了softmax）:

$${\color{orange}{Att_1}}(Q, K, V) = \text{softmax}({\color{orange}{Q_1}} K_1^T) \vec{V_1} = \text{softmaxed}({\color{orange}{Q_1}} K_1^T) \vec{V_1} \$$

当模型生成第二个“遥”字时，input=”<s>遥”, Attention的计算如下：

当 $QK^T$ 变为矩阵时，softmax 会针对 行 进行计算。写详细一点如下，softmaxed 表示已经按行进行了softmax。

假设 ${\color{orange}{Att_1}}(Q, K, V)$ 表示 Attention 的第一行， ${\color{orange}{Att_2}}(Q, K, V)$ 表示 Attention 的第二行，则根据上面推导，

其计算公式为：

$$ \begin{aligned} {\color{orange}{Att_1}}(Q, K, V) &= \text{softmaxed}({\color{orange}{Q_1}} K_1^T) \vec{V_1} \ {\color{red}{Att_2}}(Q, K, V) &= \text{softmaxed}({\color{red}{Q_2}} K_1^T) \vec{V_1} + \text{softmaxed}({\color{red}{Q_2}} K_2^T) \vec{V_2 } \end{aligned} \$$

你会发现，由于 $Q_1 K_2^T$ 这个值会mask掉，

  • $Q_1$ 在第二步参与的计算与第一步是一样的，而且第二步生成的 $V_1$ 也仅仅依赖于 $Q_1$ ，与 $Q_2$ 毫无关系。
  • $V_2$ 的计算也仅仅依赖于 $Q_2$ ，与 $Q_1$ 毫无关系。

当模型生成第三个“领”字时，input=”<s>遥遥”Attention的计算如下：

详细的推导参考第二步，其计算公式为：

$$ \begin{aligned} {\color{orange}{Att_1}}(Q, K, V) &= \text{softmaxed}({\color{orange}{Q_1}} K_1^T) \vec{V_1} \ {\color{red}{Att_2}}(Q, K, V) &= \text{softmaxed}({\color{red}{Q_2}} K_1^T) \vec{V_1} + \text{softmaxed}({\color{red}{Q_2}} K_2^T) \vec{V_2 } \ {\color{purple}{Att_3}}(Q, K, V) &= \text{softmaxed}({\color{purple}{Q_3}} K_1^T) \vec{V_1} + \text{softmaxed}({\color{purple}{Q_3}} K_2^T) \vec{V_2 } + \text{softmaxed}({\color{purple}{Q_3}} K_3^T) \vec{V_3 } \end{aligned} \$$

同样的， $Att_k$ 只与 $Q_k$ 有关。

当模型生成第四个“先”字时，input=”<s>遥遥领”Attention的计算如下：

$$ \begin{aligned} {\color{orange}{Att_1}}(Q, K, V) &= \text{softmaxed}({\color{orange}{Q_1}} K_1^T) \vec{V_1} \ {\color{red}{Att_2}}(Q, K, V) &= \text{softmaxed}({\color{red}{Q_2}} K_1^T) \vec{V_1} + \text{softmaxed}({\color{red}{Q_2}} K_2^T) \vec{V_2 } \ {\color{purple}{Att_3}}(Q, K, V) &= \text{softmaxed}({\color{purple}{Q_3}} K_1^T) \vec{V_1} + \text{softmaxed}({\color{purple}{Q_3}} K_2^T) \vec{V_2 } + \text{softmaxed}({\color{purple}{Q_3}} K_3^T) \vec{V_3 } \ {\color{brown}{Att_4}}(Q, K, V) &= \text{softmaxed}({\color{brown}{Q_4}} K_1^T) \vec{V_1} + \text{softmaxed}({\color{brown}{Q_4}} K_2^T) \vec{V_2 } + \text{softmaxed}({\color{brown}{Q_4}} K_3^T) \vec{V_3 } + \text{softmaxed}({\color{brown}{Q_4}} K_4^T) \vec{V_4 } \end{aligned} \$$

和之前类似，不再赘述。

看上面图和公式，我们可以得出结论：

  1. 当前计算方式存在大量冗余计算。
  2. $Att_k$ 只与 $Q_k$ 有关。
  3. 推理第 $x_k$ 个字符的时候只需要输入字符 $x_{k-1}$即可。

我们每一步其实之需要根据 $Q_k$ 计算 $Att_k$ 就可以，之前已经计算的Attention完全不需要重新计算。但是 $K$ 和 $V$ 是全程参与计算的，所以这里我们需要把每一步的 $K,V$ 缓存起来。所以说叫KV Cache好像有点不太对，因为KV本来就需要全程计算，可能叫增量KV计算会更好理解。

下面4张图展示了使用KV Cache和不使用的对比。

下面是gpt里面KV Cache的实现。其实明白了原理后代码实现简单的不得了,就是concat操作而已。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
if layer_past is not None:
        past_key, past_value = layer_past
        key = torch.cat((past_key, key), dim=-2)
        value = torch.cat((past_value, value), dim=-2)
    
    if use_cache is True:
        present = (key, value)
    else:
        present = None
    
    if self.reorder_and_upcast_attn:
        attn_output, attn_weights = self._upcast_and_reordered_attn(query, key, value, attention_mask, head_mask)
    else:
        attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)

最后需要注意当sequence特别长的时候，KV Cache其实还是个Memory刺客。

比如batch_size=32, head=32, layer=32, dim_size=4096, seq_length=2048, float32类型，则需要占用的显存为（感谢网友指正） 2 * 32 * 4096 * 2048 * 32 * 4 / 1024/1024/1024 /1024 = 64G。

缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA

前几天，幻方发布的DeepSeek-V2引起了大家的热烈讨论。首先，最让人哗然的是1块钱100万token的价格，普遍比现有的各种竞品API便宜了两个数量级，以至于有人调侃“这个价格哪怕它输出乱码，我也会认为这个乱码是一种艺术”；其次，从模型的技术报告看，如此便宜的价格背后的关键技术之一是它新提出的MLA（Multi-head Latent Attention），这是对GQA的改进，据说能比GQA更省更好，也引起了读者的广泛关注。

MHA

MHA（Multi-Head Attention），也就是多头注意力，是开山之作《Attention is all you need》所提出的一种Attention形式，可以说它是当前主流LLM的基础工作。在数学上，多头注意力MHA等价于多个独立的单头注意力的拼接，假设输入的（行）向量序列为$\boldsymbol{x}_1,\boldsymbol{x}_2,\cdots,\boldsymbol{x}_l$，其中$\boldsymbol{x}_i\in\mathbb{R}^d$，那么MHA可以形式地记为

$$\begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}t^{(h)}\right] \[10pt] \boldsymbol{o}t^{(s)} = Attention\left(\boldsymbol{q}t^{(s)}, \boldsymbol{k}{\leq t}^{(s)} ,\boldsymbol{v}{\leq t}^{(s)}\right)\triangleq\frac{\sum{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)\boldsymbol{v}i^{(s)}}{\sum{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)} \[15pt] \boldsymbol{q}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d\times d_k}\ \boldsymbol{k}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_k^{(s)}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_k^{(s)}\in\mathbb{R}^{d\times d_k} \ \boldsymbol{v}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_v},\quad \boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d\times d_v} \end{gathered} \tag1$$

简单起见，这里省略了Attention矩阵的缩放因子。实践上，常见的设置是$d_k = d_v = d / h$，对于LLAMA2-7b有$d=4096, h=32, d_k = d_v = 128$，LLAMA2-70b则是$d=8192,h=64, d_k = d_v = 128$

由于这里只考虑了主流的自回归LLM所用的Causal Attention，因此在token by token递归生成时，新预测出来的第$t+1$ token，并不会影响到已经算好的$\boldsymbol{k}{\leq t}^{(s)} ,\boldsymbol{v}{\leq t}^{(s)}$，因此这部分结果我们可以缓存下来供后续生成调用，避免不必要的重复计算，这就是所谓的KV Cache。

而后面的MQA、GQA、MLA，都是围绕“如何减少KV Cache同时尽可能地保证效果”这个主题发展而来的产物。

瓶颈

一个自然的问题是：为什么降低KV Cache的大小如此重要？

众所周知，一般情况下LLM的推理都是在GPU上进行，单张GPU的显存是有限的，一部分我们要用来存放模型的参数和前向计算的激活值，这部分依赖于模型的体量，选定模型后它就是个常数；另外一部分我们要用来存放模型的KV Cache，这部分不仅依赖于模型的体量，还依赖于模型的输入长度，也就是在推理过程中是动态增长的，当Context长度足够长时，它的大小就会占主导地位，可能超出一张卡甚至一台机（8张卡）的总显存量。

在GPU上部署模型的原则是：能一张卡部署的，就不要跨多张卡；能一台机部署的，就不要跨多台机。这是因为“卡内通信带宽 > 卡间通信带宽 > 机间通信带宽”，由于“木桶效应”，模型部署时跨的设备越多，受设备间通信带宽的的“拖累”就越大，事实上即便是单卡H100内SRAM与HBM的带宽已经达到了3TB/s，但对于Short Context来说这个速度依然还是推理的瓶颈，更不用说更慢的卡间、机间通信了。

所以，减少KV Cache的目的就是要实现在更少的设备上推理更长的Context，或者在相同的Context长度下让推理的batch size更大，从而实现更快的推理速度或者更大的吞吐总量。当然，最终目的都是为了实现更低的推理成本。

要想更详细地了解这个问题，读者可以进一步阅读《FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness》、《A guide to LLM inference and performance》、《LLM inference speed of light》等文章。

MQA

MQA，即“Multi-Query Attention”，是减少KV Cache的一次非常朴素的尝试，首次提出自《Fast Transformer Decoding: One Write-Head is All You Need》，这已经是2019年的论文了，这也意味着早在LLM火热之前，减少KV Cache就已经是研究人员非常关注的一个课题了。

MQA的思路很简单，直接让所有Attention Head共享同一个K、V，用公式来说，就是取消MHA所有的$\boldsymbol{k},\boldsymbol{v}$ 上标${}^{(s)}$：

$$ \begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}t^{(h)}\right] \[10pt] \boldsymbol{o}t^{(s)} = Attention\left(\boldsymbol{q}t^{(s)}, \boldsymbol{k}{\leq t}^{\color{#ccc}{\smash{\bcancel{(s)}}}} ,\boldsymbol{v}{\leq t}^{\color{#ccc}{\smash{\bcancel{(s)}}}}\right)\triangleq\frac{\sum{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}}{}^{\top}\right)\boldsymbol{v}i^{\color{#ccc}{\smash{\bcancel{(s)}}}}}{\sum{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}}{}^{\top}\right)} \[15pt] \boldsymbol{q}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d\times d_k}\ \boldsymbol{k}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}} = \boldsymbol{x}_i\boldsymbol{W}_k^{\color{#ccc}{\smash{\bcancel{(s)}}}}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_k^{\color{#ccc}{\smash{\bcancel{(s)}}}}\in\mathbb{R}^{d\times d_k} \ \boldsymbol{v}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}} = \boldsymbol{x}_i\boldsymbol{W}_v^{\color{#ccc}{\smash{\bcancel{(s)}}}}\in\mathbb{R}^{d_v},\quad \boldsymbol{W}_v^{\color{#ccc}{\smash{\bcancel{(s)}}}}\in\mathbb{R}^{d\times d_v} \end{gathered} \tag2$$

使用MQA的模型包括PaLM、StarCoder、Gemini等。很明显，MQA直接将KV Cache减少到了原来的$1/h$，这是非常可观的，单从节省显存角度看已经是天花板了。

效果方面，目前看来大部分任务的损失都比较有限，且MQA的支持者相信这部分损失可以通过进一步训练来弥补回。此外，注意到MQA由于共享了K、V，将会导致Attention的参数量减少了将近一半，而为了模型总参数量的不变，通常会相应地增大FFN/GLU的规模，这也能弥补一部分效果损失。

GQA

然而，也有人担心MQA对KV Cache的压缩太严重，以至于会影响模型的学习效率以及最终效果。为此，一个MHA与MQA之间的过渡版本GQA（Grouped-Query Attention）应运而生，出自论文《GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints》，是去年的工作。

事后看来，GQA的思想也很朴素，它就是将所有Head分为$g$ 组（$g$ 以整除$h$），每组共享同一对K、V，用数学公式表示为

$$\begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}t^{(h)}\right] \[10pt] \boldsymbol{o}t^{(s)} = Attention\left(\boldsymbol{q}t^{(s)}, \boldsymbol{k}{\leq t}^{\color{red}{(\lceil sg/h\rceil)}} ,\boldsymbol{v}{\leq t}^{\color{red}{(\lceil sg/h\rceil)}}\right)\triangleq\frac{\sum{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{\color{red}{(\lceil sg/h\rceil)}}{}^{\top}\right)\boldsymbol{v}i^{\color{red}{(\lceil sg/h\rceil)}}}{\sum{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{\color{red}{(\lceil sg/h\rceil)}}{}^{\top}\right)} \[15pt] \boldsymbol{q}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d\times d_k}\ \boldsymbol{k}_i^{\color{red}{(\lceil sg/h\rceil)}} = \boldsymbol{x}_i\boldsymbol{W}_k^{\color{red}{(\lceil sg/h\rceil)}}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_k^{\color{red}{(\lceil sg/h\rceil)}}\in\mathbb{R}^{d\times d_k} \ \boldsymbol{v}_i^{\color{red}{(\lceil sg/h\rceil)}} = \boldsymbol{x}_i\boldsymbol{W}_v^{\color{red}{(\lceil sg/h\rceil)}}\in\mathbb{R}^{d_v},\quad \boldsymbol{W}_v^{\color{red}{(\lceil sg/h\rceil)}}\in\mathbb{R}^{d\times d_v} \end{gathered} \$$

这里的$\lceil\cdot\rceil$ 上取整符号。GQA提供了MHA到MQA的自然过渡，当$g=h$ 就是MHA，$g=1$ 就是MQA，当$1 < g < h$ ，它只将KV Cache压缩到$g/h$，压缩率不如MQA，但同时也提供了更大的自由度，效果上更有保证。GQA最知名的使用者，大概是Meta开源的LLAMA2-70B，以及LLAMA3全系列，此外使用GQA的模型还有TigerBot、DeepSeek-V1、StarCoder2、Yi、ChatGLM2、ChatGLM3等，相比使用MQA的模型更多（ChatGLM虽然在它的介绍中说自己是MQA，但实际是$g=2$ GQA）。

在llama2/3-70B中，GQA的$g=8$，其他用了GQA的同体量模型基本上也保持了这个设置，这并非偶然，而是同样出于推理效率的考虑。我们知道，70B这个体量的模型，如果不进行极端的量化，那么不可能部署到单卡（A100/H100 80G）上。单卡不行，那么就能单机了，一般情况下一台机可以装8张卡，刚才我们说了，Attention的每个Head实际上是独立运算然后拼接起来的，当$g=8$ ，正好可以每张卡负责计算一组K、V对应的Attention Head，这样可以在尽可能保证K、V多样性的同时最大程度上减少卡间通信。

MLA

有了MHA、MQA、GQA的铺垫，我们理解MLA（Multi-head Latent Attention）就相对容易一些了。DeepSeek-V2的技术报告里是从低秩投影的角度引入MLA的，以至于有部分读者提出“为什么LoRA提出这么久了，直到MLA才提出对KV Cache低秩分解的做法”之类的疑问。

然而，笔者认为低秩投影这个角度并不贴近本质，因为要说低秩投影的话，事实上只要我们将GQA的所有K、V叠在一起，就会发现GQA也相当于在做低秩投影：

$$\underbrace{\left[\boldsymbol{k}_i^{(1)},\cdots,\boldsymbol{k}_i^{(g)},\boldsymbol{v}_i^{(1)},\cdots,\boldsymbol{v}i^{(g)}\right]}{\boldsymbol{c}_i\in\mathbb{R}^{g(d_k+d_v)}} = \boldsymbol{x}_i \underbrace{\left[\boldsymbol{W}_k^{(1)},\cdots,\boldsymbol{W}_k^{(g)},\boldsymbol{W}_v^{(1)},\cdots,\boldsymbol{W}v^{(g)}\right]}{\boldsymbol{W}_c\in\mathbb{R}^{d\times g(d_k+d_v)}} \$$

这里我们将所有$\boldsymbol{k}_i^{(s)},\boldsymbol{v}_i^{(s)}$ 在一起记为$\boldsymbol{c}_i$，相应的投影矩阵也拼在一起记为$\boldsymbol{W}_c$，注意到一般都有$d_c = g(d_k+d_v) < d$，所以$\boldsymbol{x}_i$ $\boldsymbol{c}_i$ 变换就是一个低秩投影。所以，MLA的本质改进不是低秩投影，而是低秩投影之后的工作。

Part 1

GQA在投影之后做了什么呢？首先它将向量对半分为两份分别作为K、V，然后每一份又均分为$g$ ，每一份复制$h/g$ ，以此来“凑”够$h$ Attention Head所需要的K、V。我们知道分割、复制都是简单的线性变换，所以MLA的第一个想法是将这些简单的线性变换换成一般的线性变换，以增强模型的能力：

$$\begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}t^{(h)}\right] \[10pt] \boldsymbol{o}t^{(s)} = Attention\left(\boldsymbol{q}t^{(s)}, \boldsymbol{k}{\leq t}^{(s)} ,\boldsymbol{v}{\leq t}^{(s)}\right)\triangleq\frac{\sum{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)\boldsymbol{v}i^{(s)}}{\sum{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)} \[15pt] \boldsymbol{q}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d\times d_k}\ \boldsymbol{k}_i^{(s)} = \boldsymbol{c}_i\boldsymbol{W}_k^{(s)}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_k^{(s)}\in\mathbb{R}^{d_c\times d_k} \ \boldsymbol{v}_i^{(s)} = \boldsymbol{c}_i\boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_v},\quad \boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_c\times d_v} \[10pt] \boldsymbol{c}_i = \boldsymbol{x}_i \boldsymbol{W}_c\in\mathbb{R}^{d_c},\quad \boldsymbol{W}_c\in\mathbb{R}^{d\times d_c} \end{gathered} \$$

然而，理论上这样是能增加模型能力，但别忘了GQA的主要目的是减少KV Cache，出于节省计算和通信成本的考虑，我们一般会缓存的是投影后的$\boldsymbol{k}_i, \boldsymbol{v}_i$ 不是投影前的$\boldsymbol{c}_i$ $\boldsymbol{x}_i$，而MLA的这个做法，通过不同的投影矩阵再次让所有的K、V Head都变得各不相同，那么KV Cache的大小就恢复成跟MHA一样大了，违背了GQA的初衷。

对此，MLA发现，我们可以结合Dot-Attention的具体形式，通过一个简单但不失巧妙的恒等变换来规避这个问题。首先，在训练阶段还是照常进行，此时优化空间不大；然后，在推理阶段，我们利用

$$\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top} = \left(\boldsymbol{x}_t\boldsymbol{W}_q^{(s)}\right) \left(\boldsymbol{c}_i\boldsymbol{W}_k^{(s)}\right){}^{\top} = \boldsymbol{x}_t\left(\boldsymbol{W}_q^{(s)}\boldsymbol{W}_k^{(s)}{}^{\top}\right)\boldsymbol{c}_i^{\top} \$$

这意味着推理阶段，我们可以将$\boldsymbol{W}_q^{(s)}\boldsymbol{W}_k^{(s)}{}^{\top}$ 并起来作为Q的投影矩阵，那么$\boldsymbol{c}_i$ 取代了原本的$\boldsymbol{k}_i$，同理，在$\boldsymbol{o}_t$ 面我们还有一个投影矩阵，于是$\boldsymbol{v}_i^{(s)} = \boldsymbol{c}_i\boldsymbol{W}_v^{(s)}$ $\boldsymbol{W}_v^{(s)}$ 可以吸收到后面的投影矩阵中去，于是等效地$\boldsymbol{v}_i$ 可以用$\boldsymbol{c}_i$ 替，也就是说此时KV Cache只需要存下所有的$\boldsymbol{c}_i$ 行，而不至于存下所有的$\boldsymbol{k}_i^{(s)}$、$\boldsymbol{v}_i^{(s)}$。注意到$\boldsymbol{c}_i$ ${}^{(s)}$ 关，也就是说是所有头共享的，即MLA在推理阶段它可以恒等变换为一个MQA。

再次强调，本文的主题是一直都是减少KV Cache，那到目前为止，MLA做到了什么呢？答案是通过不同的投影矩阵来增强了GQA的能力，并且推理时可以保持同样大小的KV Cache。那么反过来，如果我们只需要跟GQA相近的能力，那么是不是就可以再次减少KV Cache了？换言之，$d_c$ 必要取$g(d_k+d_v)$，而是取更小的值（DeepSeek-V2取了512），从而进一步压缩KV Cache，这就是MLA的核心思想。

（注：这里有一个细节，就是$\boldsymbol{W}_q^{(s)}\boldsymbol{W}_k^{(s)}{}^{\top}$ 并成一个矩阵的恒等变换，理论上只有在无限精度下才成立，实际上如果我们使用单精度尤其是BF16的话，经过变换后的精度损失往往还是挺明显的，经过多层累积后可能放大到比较可观的程度，这里可能要根据实际误差看要不要做一些后处理。）

Part 2

一切似乎都很完美，看上去一个又好又省的理想设计就要出炉了。不过别急，当我们再深入思考一下就会发现，到目前为止的MLA有一个难以绕开的缺陷——不兼容RoPE（旋转位置编码）。

刚才我们说了，MLA之所以能保持跟GQA一样大小的KV Cache，其关键一步是“将$\boldsymbol{W}_q^{(s)}\boldsymbol{W}_k^{(s)}{}^{\top}$ 并成一个（跟位置无关的）矩阵作为Q的投影矩阵”，但如果加了RoPE的话，这一步就无法实现了。这是因为RoPE是一个跟位置相关的、$d_k\times d_k$ 分块对角矩阵$\boldsymbol{\mathcal{R}}_m$，满足$\boldsymbol{\mathcal{R}}_m\boldsymbol{\mathcal{R}}n^{\top}=\boldsymbol{\mathcal{R}}{m-n}$，MLA加入RoPE之后会让$\boldsymbol{W}_q^{(s)}\boldsymbol{W}k^{(s)}{}^{\top}$ 间多插入了一项$\boldsymbol{\mathcal{R}}{t-i}$：

$$\boldsymbol{q}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_q^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\quad,\quad\boldsymbol{k}_i^{(s)} = \boldsymbol{c}_i\boldsymbol{W}_k^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i} \ \boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top} = \left(\boldsymbol{x}_t\boldsymbol{W}_q^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_t}\right) \left(\boldsymbol{c}_i\boldsymbol{W}_k^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\right){}^{\top} = \boldsymbol{x}_t\left(\boldsymbol{W}q^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}{t-i}}\boldsymbol{W}_k^{(s)}{}^{\top}\right)\boldsymbol{c}_i^{\top} \$$

这里的$\boldsymbol{W}q^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}{t-i}}\boldsymbol{W}_k^{(s)}{}^{\top}$ 无法合并为一个固定的投影矩阵了（跟位置差$t-i$ 关），从而MLA的想法无法结合RoPE实现。

前段时间，笔者也很荣幸跟DeepSeek团队讨论过这个问题，但这个问题可以说非常本质，所以当时笔者实际上也没能提出什么有效的建议。最简单的方式是放弃RoPE，换用其他基于Attention Bias的位置编码，如ALIBI，但DeepSeek的实验显示它明显不如RoPE（注意，MLA不是不能加RoPE，而是加了RoPE之后无法用恒等变换技巧来减少KV Cache），笔者也提议过换Sandwich，它不像ALIBI单调衰减到负无穷，估计效果会好些，但感觉是治标不治本。还有一个折中的办法是将$\boldsymbol{q}_i$ 输入也改为$\boldsymbol{c}_i$，然后RoPE加在$\boldsymbol{c}_i$ 后，即

$$\boldsymbol{q}_i^{(s)} = \boldsymbol{c}_i\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\boldsymbol{W}_q^{(s)},\quad\boldsymbol{k}_i^{(s)} = \boldsymbol{c}_i\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\boldsymbol{W}_k^{(s)} \$$

这样$\boldsymbol{\mathcal{R}}_i$ 可以吸收到$\boldsymbol{c}_i$ 去，但这样就没有$\boldsymbol{\mathcal{R}}_m\boldsymbol{\mathcal{R}}n^{\top}=\boldsymbol{\mathcal{R}}{m-n}$ 运算了，此时的RoPE不再是通过绝对位置实现相对位置，而单纯是在Q、K上加绝对位置，让模型自己想办法提炼相对位置信息。

最后发布的MLA，采取了一种混合的方法——每个Attention Head的Q、K新增$d_r$ 维度用来添加RoPE，其中K新增的维度每个Head共享：

$$\begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}t^{(h)}\right] \[10pt] \boldsymbol{o}t^{(s)} = Attention\left(\boldsymbol{q}t^{(s)}, \boldsymbol{k}{\leq t}^{(s)} ,\boldsymbol{v}{\leq t}^{(s)}\right)\triangleq\frac{\sum{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)\boldsymbol{v}i^{(s)}}{\sum{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)} \[15pt] \boldsymbol{q}i^{(s)} = \left[\boldsymbol{x}i\boldsymbol{W}{qc}^{(s)}, \boldsymbol{x}i\boldsymbol{W}{qr}^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}i}\right]\in\mathbb{R}^{d_k + d_r},\quad \boldsymbol{W}{qc}^{(s)}\in\mathbb{R}^{d\times d_k},\boldsymbol{W}{qr}^{(s)}\in\mathbb{R}^{d\times d_r}\ \boldsymbol{k}i^{(s)} = \left[\boldsymbol{c}i\boldsymbol{W}{kc}^{(s)}, \boldsymbol{x}i\boldsymbol{W}{kr}^{\color{#ccc}{\smash{\bcancel{(s)}}}}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}i}\right]\in\mathbb{R}^{d_k+d_r},\quad \boldsymbol{W}{kc}^{(s)}\in\mathbb{R}^{d_c\times d_k}, \boldsymbol{W}{kr}^{\color{#ccc}{\smash{\bcancel{(s)}}}}\in\mathbb{R}^{d\times d_r} \ \boldsymbol{v}_i^{(s)} = \boldsymbol{c}_i\boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_v},\quad \boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_c\times d_v} \[10pt] \boldsymbol{c}_i = \boldsymbol{x}_i \boldsymbol{W}_c\in\mathbb{R}^{d_c},\quad \boldsymbol{W}_c\in\mathbb{R}^{d\times d_c} \end{gathered} \$$

这样一来，没有RoPE的维度就可以重复“Part 1”的操作，在推理时KV Cache只需要存$\boldsymbol{c}_i$，新增的带RoPE的维度就可以用来补充位置信息，并且由于所有Head共享，所以也就只有在K Cache这里增加了$d_r$ 维度，原论文取了$d_r = d_k / 2 = 64$，相比原本的$d_c=512$，增加的幅度不大。

Part 3

最后有一个细节，就是MLA的最终版本，还将Q的输入也改为了低秩投影形式，这与减少KV Cache无关，主要是为了减少训练期间参数量和相应的梯度（原论文说的是激活值，个人表示不大理解）所占的显存：

$$\begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}t^{(h)}\right] \[10pt] \boldsymbol{o}t^{(s)} = Attention\left(\boldsymbol{q}t^{(s)}, \boldsymbol{k}{\leq t}^{(s)} ,\boldsymbol{v}{\leq t}^{(s)}\right)\triangleq\frac{\sum{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)\boldsymbol{v}i^{(s)}}{\sum{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)} \[15pt] \boldsymbol{q}i^{(s)} = \left[\boldsymbol{c}i’\boldsymbol{W}{qc}^{(s)}, \boldsymbol{c}i’\boldsymbol{W}{qr}^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}i}\right]\in\mathbb{R}^{d_k + d_r},\quad \boldsymbol{W}{qc}^{(s)}\in\mathbb{R}^{d_c’\times d_k},\boldsymbol{W}{qr}^{(s)}\in\mathbb{R}^{d_c’\times d_r}\ \boldsymbol{k}i^{(s)} = \left[\boldsymbol{c}i\boldsymbol{W}{kc}^{(s)}, \boldsymbol{x}i\boldsymbol{W}{kr}^{\color{#ccc}{\smash{\bcancel{(s)}}}}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}i}\right]\in\mathbb{R}^{d_k+d_r},\quad \boldsymbol{W}{kc}^{(s)}\in\mathbb{R}^{d_c\times d_k}, \boldsymbol{W}{kr}^{\color{#ccc}{\smash{\bcancel{(s)}}}}\in\mathbb{R}^{d\times d_r} \ \boldsymbol{v}_i^{(s)} = \boldsymbol{c}_i\boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_v},\quad \boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_c\times d_v} \[10pt] \boldsymbol{c}_i’ = \boldsymbol{x}_i \boldsymbol{W}_c’\in\mathbb{R}^{d_c’},\quad \boldsymbol{W}_c’\in\mathbb{R}^{d\times d_c’} \ \boldsymbol{c}_i = \boldsymbol{x}_i \boldsymbol{W}_c\in\mathbb{R}^{d_c},\quad \boldsymbol{W}_c\in\mathbb{R}^{d\times d_c} \ \end{gathered} \$$

注意$\boldsymbol{k}_i^{(s)}$ 的第二项，带RoPE的部分，其输入还是$\boldsymbol{x}_i$ 不是$\boldsymbol{c}_i$，这里保持了原论文的设置，不是笔误，$d_c’$ 论文的取值是1536，跟$d_c=512$ 同。同时，我们把带RoPE的MHA放在下面，方便大家对比：

$$\begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}t^{(h)}\right] \[10pt] \boldsymbol{o}t^{(s)} = Attention\left(\boldsymbol{q}t^{(s)}, \boldsymbol{k}{\leq t}^{(s)} ,\boldsymbol{v}{\leq t}^{(s)}\right)\triangleq\frac{\sum{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)\boldsymbol{v}i^{(s)}}{\sum{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)} \[15pt] \boldsymbol{q}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_q^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d\times d_k}\ \boldsymbol{k}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_k^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}_i}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_k^{(s)}\in\mathbb{R}^{d\times d_k} \ \boldsymbol{v}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_v},\quad \boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d\times d_v} \end{gathered} \$$

可以发现，其实在训练阶段，除了多了一步低秩投影以及只在部分维度加RoPE外，MLA与Q、K的Head Size由$d_k$ 成$d_k + d_r$ MHA基本无异。

推理阶段的MLA则改为

$$\begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}\boldsymbol{W}_v^{(1)}, \boldsymbol{o}_t^{(2)}\boldsymbol{W}_v^{(2)}, \cdots, \boldsymbol{o}_t^{(h)}\boldsymbol{W}v^{(h)}\right] \[10pt] \boldsymbol{o}t^{(s)} = Attention\left(\boldsymbol{q}t^{(s)}, \boldsymbol{k}{\leq t}^{(s)} ,\boldsymbol{c}{\leq t}\right)\triangleq\frac{\sum{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)\boldsymbol{c}i}{\sum{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)} \[15pt] \boldsymbol{q}i^{(s)} = \left[\boldsymbol{c}i’\boldsymbol{W}{qc}^{(s)}\boldsymbol{W}{kc}^{(s)}{}^{\top}, \boldsymbol{c}i’\boldsymbol{W}{qr}^{(s)}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}i}\right]\in\mathbb{R}^{d_c + d_r}\ \boldsymbol{k}i^{(s)} = \left[\boldsymbol{c}i, \boldsymbol{x}i\boldsymbol{W}{kr}^{\color{#ccc}{\smash{\bcancel{(s)}}}}\color{#3ce2f7}{\boldsymbol{\mathcal{R}}i}\right]\in\mathbb{R}^{d_c+d_r}\ \boldsymbol{W}{qc}^{(s)}\in\mathbb{R}^{d_c’\times d_k},\boldsymbol{W}{kc}^{(s)}\in\mathbb{R}^{d_c\times d_k},\boldsymbol{W}{qr}^{(s)}\in\mathbb{R}^{d_c’\times d_r},\boldsymbol{W}{kr}^{\color{#ccc}{\smash{\bcancel{(s)}}}}\in\mathbb{R}^{d\times d_r} \[10pt] \boldsymbol{c}_i’ = \boldsymbol{x}_i \boldsymbol{W}_c’\in\mathbb{R}^{d_c’},\quad \boldsymbol{W}_c’\in\mathbb{R}^{d\times d_c’} \ \boldsymbol{c}_i = \boldsymbol{x}_i \boldsymbol{W}_c\in\mathbb{R}^{d_c},\quad \boldsymbol{W}_c\in\mathbb{R}^{d\times d_c} \ \end{gathered} \$$

此时Q、K的Head Size变成了$d_c + d_r$，V的Head Size 则变成了$d_c$，按照原论文的设置，这是$d_k$、$d_v$ 4倍。所以实际上MLA在推理阶段做的这个转换，虽然能有效减少KV Cache，但其推理的计算量是增加的。

那为什么还能提高推理效率呢？这又回到“瓶颈”一节所讨论的问题了，我们可以将LLM的推理分两部分：第一个Token的生成（Prefill）和后续每个Token的生成（Generation），Prefill阶段涉及到对输入所有Token的并行计算，然后把对应的KV Cache存下来，这部分对于计算、带宽和显存都是瓶颈，MLA虽然增大了计算量，但KV Cache的减少也降低了显存和带宽的压力，大家半斤八两；但是Generation阶段由于每步只计算一个Token，实际上它更多的是带宽瓶颈和显存瓶颈，因此MLA的引入理论上能明显提高Generation的速度。

还有一个细节充分体现了这个特性。一般的LLM架构参数满足$h \times d_k = d$，即num_heads * head_size = hidden_size，但DeepSeek-V2不一样，它$d_k=128,d=5120$，但$h=128$，是一般设置的3倍！这是因为MLA的KV Cache大小跟$h$ 关，增大$h$ 会增加计算量和提升模型能力，但不会增加KV Cache，所以不会带来速度瓶颈。

小结

本文简单概述了多头注意力的演变历程，特别是从MHA向MQA、GQA，最终到MLA的变化理念，最后详细展开了对MLA的介绍。在本文中，MLA被视为GQA的一般化，它用投影矩阵的方式替代了GQA的分割、重复，并引入了一个恒等变换技巧来可以进一步压缩KV Cache，同时采用了一种混合方法来兼容RoPE。总的来说，MLA称得上是一种非常实用的注意力变体。