大模型推理加速：看图学KV Cache (opens new window)

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）:

$$At{t}_1(Q,K,V)=\text{softmax}(Q_1{K}_1^T)\overrightarrow{V_1}=\text{softmaxed}(Q_1{K}_1^T)\overrightarrow{V_1}$$

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

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

假设 $At{t}_1(Q,K,V)$ 表示 Attention 的第一行， $At{t}_2(Q,K,V)$ 表示 Attention 的第二行，则根据上面推导，

其计算公式为：

$$\begin{array}{rl}At{t}_1(Q,K,V)& =\text{softmaxed}(Q_1{K}_1^T)\overrightarrow{V_1}\\ At{t}_2(Q,K,V)& =\text{softmaxed}(Q_2{K}_1^T)\overrightarrow{V_1}+\text{softmaxed}(Q_2{K}_2^T)\overrightarrow{V_2}\end{array}$$

你会发现，由于 $Q_1{K}_2^T$ 这个值会mask掉，

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

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

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

$$\begin{array}{rl}At{t}_1(Q,K,V)& =\text{softmaxed}(Q_1{K}_1^T)\overrightarrow{V_1}\\ At{t}_2(Q,K,V)& =\text{softmaxed}(Q_2{K}_1^T)\overrightarrow{V_1}+\text{softmaxed}(Q_2{K}_2^T)\overrightarrow{V_2}\\ At{t}_3(Q,K,V)& =\text{softmaxed}(Q_3{K}_1^T)\overrightarrow{V_1}+\text{softmaxed}(Q_3{K}_2^T)\overrightarrow{V_2}+\text{softmaxed}(Q_3{K}_3^T)\overrightarrow{V_3}\end{array}$$

同样的， $At{t}_k$ 只与 $Q_k$ 有关。

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

$$\begin{array}{rl}At{t}_1(Q,K,V)& =\text{softmaxed}(Q_1{K}_1^T)\overrightarrow{V_1}\\ At{t}_2(Q,K,V)& =\text{softmaxed}(Q_2{K}_1^T)\overrightarrow{V_1}+\text{softmaxed}(Q_2{K}_2^T)\overrightarrow{V_2}\\ At{t}_3(Q,K,V)& =\text{softmaxed}(Q_3{K}_1^T)\overrightarrow{V_1}+\text{softmaxed}(Q_3{K}_2^T)\overrightarrow{V_2}+\text{softmaxed}(Q_3{K}_3^T)\overrightarrow{V_3}\\ At{t}_4(Q,K,V)& =\text{softmaxed}(Q_4{K}_1^T)\overrightarrow{V_1}+\text{softmaxed}(Q_4{K}_2^T)\overrightarrow{V_2}+\text{softmaxed}(Q_4{K}_3^T)\overrightarrow{V_3}+\text{softmaxed}(Q_4{K}_4^T)\overrightarrow{V_4}\end{array}$$

和之前类似，不再赘述。

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

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

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

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

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

https:// (opens new window)

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 (opens new window)

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

MHA

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

$$\begin{array}{}\text{(1)}& \begin{array}{c}{\mathit{o}}_t=[{\mathit{o}}_t^{(1)},{\mathit{o}}_t^{(2)},\cdots ,{\mathit{o}}_t^{(h)}]\\ {\mathit{o}}_t^{(s)}=Attention({\mathit{q}}_t^{(s)},{\mathit{k}}_{\le t}^{(s)},{\mathit{v}}_{\le t}^{(s)})\triangleq \frac{\sum _{i\le t}\exp \left({\mathit{q}}_t^{(s)}{\mathit{k}}_i^{(s)}{}^{\mathrm{\top }}\right){\mathit{v}}_i^{(s)}}{\sum _{i\le t}\exp \left({\mathit{q}}_t^{(s)}{\mathit{k}}_i^{(s)}{}^{\mathrm{\top }}\right)}\\ {\mathit{q}}_i^{(s)}={\mathit{x}}_i{\mathit{W}}_q^{(s)}\in {\mathbb{R}}^{d_k},{\mathit{W}}_q^{(s)}\in {\mathbb{R}}^{d\times d_k}\\ {\mathit{k}}_i^{(s)}={\mathit{x}}_i{\mathit{W}}_k^{(s)}\in {\mathbb{R}}^{d_k},{\mathit{W}}_k^{(s)}\in {\mathbb{R}}^{d\times d_k}\\ {\mathit{v}}_i^{(s)}={\mathit{x}}_i{\mathit{W}}_v^{(s)}\in {\mathbb{R}}^{d_v},{\mathit{W}}_v^{(s)}\in {\mathbb{R}}^{d\times d_v}\end{array}\end{array}$$

简单起见，这里省略了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，并不会影响到已经算好的${\mathit{k}}_{\le t}^{(s)},{\mathit{v}}_{\le 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》 (opens new window)、《A guide to LLM inference and performance》 (opens new window)、《LLM inference speed of light》 (opens new window)等文章。

MQA

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

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

$$\begin{array}{}\text{(2)}& \begin{array}{c}{\mathit{o}}_t=[{\mathit{o}}_t^{(1)},{\mathit{o}}_t^{(2)},\cdots ,{\mathit{o}}_t^{(h)}]\\ {\mathit{o}}_t^{(s)}=Attention({\mathit{q}}_t^{(s)},{\mathit{k}}_{\le t}^{\bcancel{(s)}},{\mathit{v}}_{\le t}^{\bcancel{(s)}})\triangleq \frac{\sum _{i\le t}\exp \left({\mathit{q}}_t^{(s)}{\mathit{k}}_i^{\bcancel{(s)}}{}^{\mathrm{\top }}\right){\mathit{v}}_i^{\bcancel{(s)}}}{\sum _{i\le t}\exp \left({\mathit{q}}_t^{(s)}{\mathit{k}}_i^{\bcancel{(s)}}{}^{\mathrm{\top }}\right)}\\ {\mathit{q}}_i^{(s)}={\mathit{x}}_i{\mathit{W}}_q^{(s)}\in {\mathbb{R}}^{d_k},{\mathit{W}}_q^{(s)}\in {\mathbb{R}}^{d\times d_k}\\ {\mathit{k}}_i^{\bcancel{(s)}}={\mathit{x}}_i{\mathit{W}}_k^{\bcancel{(s)}}\in {\mathbb{R}}^{d_k},{\mathit{W}}_k^{\bcancel{(s)}}\in {\mathbb{R}}^{d\times d_k}\\ {\mathit{v}}_i^{\bcancel{(s)}}={\mathit{x}}_i{\mathit{W}}_v^{\bcancel{(s)}}\in {\mathbb{R}}^{d_v},{\mathit{W}}_v^{\bcancel{(s)}}\in {\mathbb{R}}^{d\times d_v}\end{array}\end{array}$$

使用MQA的模型包括PaLM (opens new window)、StarCoder (opens new window)、Gemini (opens new window)等。很明显，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》 (opens new window)，是去年的工作。

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

$$\begin{array}{c}{\mathit{o}}_t=[{\mathit{o}}_t^{(1)},{\mathit{o}}_t^{(2)},\cdots ,{\mathit{o}}_t^{(h)}]\\ {\mathit{o}}_t^{(s)}=Attention({\mathit{q}}_t^{(s)},{\mathit{k}}_{\le t}^{(\lceil sg/h\rceil )},{\mathit{v}}_{\le t}^{(\lceil sg/h\rceil )})\triangleq \frac{\sum _{i\le t}\exp \left({\mathit{q}}_t^{(s)}{\mathit{k}}_i^{(\lceil sg/h\rceil )}{}^{\mathrm{\top }}\right){\mathit{v}}_i^{(\lceil sg/h\rceil )}}{\sum _{i\le t}\exp \left({\mathit{q}}_t^{(s)}{\mathit{k}}_i^{(\lceil sg/h\rceil )}{}^{\mathrm{\top }}\right)}\\ {\mathit{q}}_i^{(s)}={\mathit{x}}_i{\mathit{W}}_q^{(s)}\in {\mathbb{R}}^{d_k},{\mathit{W}}_q^{(s)}\in {\mathbb{R}}^{d\times d_k}\\ {\mathit{k}}_i^{(\lceil sg/h\rceil )}={\mathit{x}}_i{\mathit{W}}_k^{(\lceil sg/h\rceil )}\in {\mathbb{R}}^{d_k},{\mathit{W}}_k^{(\lceil sg/h\rceil )}\in {\mathbb{R}}^{d\times d_k}\\ {\mathit{v}}_i^{(\lceil sg/h\rceil )}={\mathit{x}}_i{\mathit{W}}_v^{(\lceil sg/h\rceil )}\in {\mathbb{R}}^{d_v},{\mathit{W}}_v^{(\lceil sg/h\rceil )}\in {\mathbb{R}}^{d\times d_v}\end{array}$$

这里的$\lceil \cdot \rceil$ 上取整符号。GQA提供了MHA到MQA的自然过渡，当$g=h$ 就是MHA，$g=1$ 就是MQA，当$1<g<h$ ，它只将KV Cache压缩到$g/h$，压缩率不如MQA，但同时也提供了更大的自由度，效果上更有保证。GQA最知名的使用者，大概是Meta开源的LLAMA2-70B (opens new window)，以及LLAMA3 (opens new window)全系列，此外使用GQA的模型还有TigerBot (opens new window)、DeepSeek-V1 (opens new window)、StarCoder2 (opens new window)、Yi (opens new window)、ChatGLM2 (opens new window)、ChatGLM3 (opens new window)等，相比使用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也相当于在做低秩投影：

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

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

Part 1

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

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

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

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

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

这意味着推理阶段，我们可以将${\mathit{W}}_q^{(s)}{\mathit{W}}_k^{(s)}{}^{\mathrm{\top }}$ 并起来作为Q的投影矩阵，那么${\mathit{c}}_i$ 取代了原本的${\mathit{k}}_i$，同理，在${\mathit{o}}_t$ 面我们还有一个投影矩阵，于是${\mathit{v}}_i^{(s)}={\mathit{c}}_i{\mathit{W}}_v^{(s)}$ ${\mathit{W}}_v^{(s)}$ 可以吸收到后面的投影矩阵中去，于是等效地${\mathit{v}}_i$ 可以用${\mathit{c}}_i$ 替，也就是说此时KV Cache只需要存下所有的${\mathit{c}}_i$ 行，而不至于存下所有的${\mathit{k}}_i^{(s)}$、${\mathit{v}}_i^{(s)}$。注意到${\mathit{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的核心思想。

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

Part 2

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

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

$$\begin{gathered} {\mathit{q}}_i^{(s)}={\mathit{x}}_i{\mathit{W}}_q^{(s)}{\mathcal{R}}_i,{\mathit{k}}_i^{(s)}={\mathit{c}}_i{\mathit{W}}_k^{(s)}{\mathcal{R}}_i \\ {\mathit{q}}_t^{(s)}{\mathit{k}}_i^{(s)}{}^{\mathrm{\top }}=\left({\mathit{x}}_t{\mathit{W}}_q^{(s)}{\mathcal{R}}_t\right)\left({\mathit{c}}_i{\mathit{W}}_k^{(s)}{\mathcal{R}}_i\right){}^{\mathrm{\top }}={\mathit{x}}_t\left({\mathit{W}}_q^{(s)}{\mathcal{R}}_{t-i}{\mathit{W}}_k^{(s)}{}^{\mathrm{\top }}\right){\mathit{c}}_i^{\mathrm{\top }} \end{gathered}$$

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

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

$${\mathit{q}}_i^{(s)}={\mathit{c}}_i{\mathcal{R}}_i{\mathit{W}}_q^{(s)},{\mathit{k}}_i^{(s)}={\mathit{c}}_i{\mathcal{R}}_i{\mathit{W}}_k^{(s)}$$

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

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

$$\begin{array}{c}{\mathit{o}}_t=[{\mathit{o}}_t^{(1)},{\mathit{o}}_t^{(2)},\cdots ,{\mathit{o}}_t^{(h)}]\\ {\mathit{o}}_t^{(s)}=Attention({\mathit{q}}_t^{(s)},{\mathit{k}}_{\le t}^{(s)},{\mathit{v}}_{\le t}^{(s)})\triangleq \frac{\sum _{i\le t}\exp \left({\mathit{q}}_t^{(s)}{\mathit{k}}_i^{(s)}{}^{\mathrm{\top }}\right){\mathit{v}}_i^{(s)}}{\sum _{i\le t}\exp \left({\mathit{q}}_t^{(s)}{\mathit{k}}_i^{(s)}{}^{\mathrm{\top }}\right)}\\ {\mathit{q}}_i^{(s)}=[{\mathit{x}}_i{\mathit{W}}_{qc}^{(s)},{\mathit{x}}_i{\mathit{W}}_{qr}^{(s)}{\mathcal{R}}_i]\in {\mathbb{R}}^{d_k+d_r},{\mathit{W}}_{qc}^{(s)}\in {\mathbb{R}}^{d\times d_k},{\mathit{W}}_{qr}^{(s)}\in {\mathbb{R}}^{d\times d_r}\\ {\mathit{k}}_i^{(s)}=[{\mathit{c}}_i{\mathit{W}}_{kc}^{(s)},{\mathit{x}}_i{\mathit{W}}_{kr}^{\bcancel{(s)}}{\mathcal{R}}_i]\in {\mathbb{R}}^{d_k+d_r},{\mathit{W}}_{kc}^{(s)}\in {\mathbb{R}}^{d_c\times d_k},{\mathit{W}}_{kr}^{\bcancel{(s)}}\in {\mathbb{R}}^{d\times d_r}\\ {\mathit{v}}_i^{(s)}={\mathit{c}}_i{\mathit{W}}_v^{(s)}\in {\mathbb{R}}^{d_v},{\mathit{W}}_v^{(s)}\in {\mathbb{R}}^{d_c\times d_v}\\ {\mathit{c}}_i={\mathit{x}}_i{\mathit{W}}_c\in {\mathbb{R}}^{d_c},{\mathit{W}}_c\in {\mathbb{R}}^{d\times d_c}\end{array}$$

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

Part 3

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

$$\begin{array}{c}{\mathit{o}}_t=[{\mathit{o}}_t^{(1)},{\mathit{o}}_t^{(2)},\cdots ,{\mathit{o}}_t^{(h)}]\\ {\mathit{o}}_t^{(s)}=Attention({\mathit{q}}_t^{(s)},{\mathit{k}}_{\le t}^{(s)},{\mathit{v}}_{\le t}^{(s)})\triangleq \frac{\sum _{i\le t}\exp \left({\mathit{q}}_t^{(s)}{\mathit{k}}_i^{(s)}{}^{\mathrm{\top }}\right){\mathit{v}}_i^{(s)}}{\sum _{i\le t}\exp \left({\mathit{q}}_t^{(s)}{\mathit{k}}_i^{(s)}{}^{\mathrm{\top }}\right)}\\ {\mathit{q}}_i^{(s)}=[{\mathit{c}}_i^{\prime }{\mathit{W}}_{qc}^{(s)},{\mathit{c}}_i^{\prime }{\mathit{W}}_{qr}^{(s)}{\mathcal{R}}_i]\in {\mathbb{R}}^{d_k+d_r},{\mathit{W}}_{qc}^{(s)}\in {\mathbb{R}}^{d_c^{\prime }\times d_k},{\mathit{W}}_{qr}^{(s)}\in {\mathbb{R}}^{d_c^{\prime }\times d_r}\\ {\mathit{k}}_i^{(s)}=[{\mathit{c}}_i{\mathit{W}}_{kc}^{(s)},{\mathit{x}}_i{\mathit{W}}_{kr}^{\bcancel{(s)}}{\mathcal{R}}_i]\in {\mathbb{R}}^{d_k+d_r},{\mathit{W}}_{kc}^{(s)}\in {\mathbb{R}}^{d_c\times d_k},{\mathit{W}}_{kr}^{\bcancel{(s)}}\in {\mathbb{R}}^{d\times d_r}\\ {\mathit{v}}_i^{(s)}={\mathit{c}}_i{\mathit{W}}_v^{(s)}\in {\mathbb{R}}^{d_v},{\mathit{W}}_v^{(s)}\in {\mathbb{R}}^{d_c\times d_v}\\ {\mathit{c}}_i^{\prime }={\mathit{x}}_i{\mathit{W}}_c^{\prime }\in {\mathbb{R}}^{d_c^{\prime }},{\mathit{W}}_c^{\prime }\in {\mathbb{R}}^{d\times d_c^{\prime }}\\ {\mathit{c}}_i={\mathit{x}}_i{\mathit{W}}_c\in {\mathbb{R}}^{d_c},{\mathit{W}}_c\in {\mathbb{R}}^{d\times d_c}\end{array}$$

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

$$\begin{array}{c}{\mathit{o}}_t=[{\mathit{o}}_t^{(1)},{\mathit{o}}_t^{(2)},\cdots ,{\mathit{o}}_t^{(h)}]\\ {\mathit{o}}_t^{(s)}=Attention({\mathit{q}}_t^{(s)},{\mathit{k}}_{\le t}^{(s)},{\mathit{v}}_{\le t}^{(s)})\triangleq \frac{\sum _{i\le t}\exp \left({\mathit{q}}_t^{(s)}{\mathit{k}}_i^{(s)}{}^{\mathrm{\top }}\right){\mathit{v}}_i^{(s)}}{\sum _{i\le t}\exp \left({\mathit{q}}_t^{(s)}{\mathit{k}}_i^{(s)}{}^{\mathrm{\top }}\right)}\\ {\mathit{q}}_i^{(s)}={\mathit{x}}_i{\mathit{W}}_q^{(s)}{\mathcal{R}}_i\in {\mathbb{R}}^{d_k},{\mathit{W}}_q^{(s)}\in {\mathbb{R}}^{d\times d_k}\\ {\mathit{k}}_i^{(s)}={\mathit{x}}_i{\mathit{W}}_k^{(s)}{\mathcal{R}}_i\in {\mathbb{R}}^{d_k},{\mathit{W}}_k^{(s)}\in {\mathbb{R}}^{d\times d_k}\\ {\mathit{v}}_i^{(s)}={\mathit{x}}_i{\mathit{W}}_v^{(s)}\in {\mathbb{R}}^{d_v},{\mathit{W}}_v^{(s)}\in {\mathbb{R}}^{d\times d_v}\end{array}$$

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

推理阶段的MLA则改为

$$\begin{array}{c}{\mathit{o}}_t=[{\mathit{o}}_t^{(1)}{\mathit{W}}_v^{(1)},{\mathit{o}}_t^{(2)}{\mathit{W}}_v^{(2)},\cdots ,{\mathit{o}}_t^{(h)}{\mathit{W}}_v^{(h)}]\\ {\mathit{o}}_t^{(s)}=Attention({\mathit{q}}_t^{(s)},{\mathit{k}}_{\le t}^{(s)},{\mathit{c}}_{\le t})\triangleq \frac{\sum _{i\le t}\exp \left({\mathit{q}}_t^{(s)}{\mathit{k}}_i^{(s)}{}^{\mathrm{\top }}\right){\mathit{c}}_i}{\sum _{i\le t}\exp \left({\mathit{q}}_t^{(s)}{\mathit{k}}_i^{(s)}{}^{\mathrm{\top }}\right)}\\ {\mathit{q}}_i^{(s)}=[{\mathit{c}}_i^{\prime }{\mathit{W}}_{qc}^{(s)}{\mathit{W}}_{kc}^{(s)}{}^{\mathrm{\top }},{\mathit{c}}_i^{\prime }{\mathit{W}}_{qr}^{(s)}{\mathcal{R}}_i]\in {\mathbb{R}}^{d_c+d_r}\\ {\mathit{k}}_i^{(s)}=[{\mathit{c}}_i,{\mathit{x}}_i{\mathit{W}}_{kr}^{\bcancel{(s)}}{\mathcal{R}}_i]\in {\mathbb{R}}^{d_c+d_r}\\ {\mathit{W}}_{qc}^{(s)}\in {\mathbb{R}}^{d_c^{\prime }\times d_k},{\mathit{W}}_{kc}^{(s)}\in {\mathbb{R}}^{d_c\times d_k},{\mathit{W}}_{qr}^{(s)}\in {\mathbb{R}}^{d_c^{\prime }\times d_r},{\mathit{W}}_{kr}^{\bcancel{(s)}}\in {\mathbb{R}}^{d\times d_r}\\ {\mathit{c}}_i^{\prime }={\mathit{x}}_i{\mathit{W}}_c^{\prime }\in {\mathbb{R}}^{d_c^{\prime }},{\mathit{W}}_c^{\prime }\in {\mathbb{R}}^{d\times d_c^{\prime }}\\ {\mathit{c}}_i={\mathit{x}}_i{\mathit{W}}_c\in {\mathbb{R}}^{d_c},{\mathit{W}}_c\in {\mathbb{R}}^{d\times d_c}\end{array}$$

此时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称得上是一种非常实用的注意力变体。