大模型学习-基础篇（七）KV cache 简单了解

系列目录-大模型学习篇

大模型学习-基础篇
大模型学习-实践篇

问题场景（Problems）

针对attention的改进，我们前面提到了Attention的key和value可以保存在cache中，对于后面的token计算，我们只需要计算当前token的注意力值。这样可以大大减少模型的计算量。随着在KV cache的改进方向上前进，目前已经有了很多版本，例如MQA，GQA，MLA。本文参考了苏神的知乎文章，不涉及证明过程。

其实，这里需要说明的是，我们所讲的k和v是指输入x左乘Wk和Wv以后算得值。生成过程中，每一轮新增加的这个token会使得序列长度+1，但是你前面的token是没变的，所以前面的每个token的k和v都是算好的。

为什么优化KV Cache这么重要

在大模型推理过程中，降低KV Cache的大小至关重要，因为它直接影响显存的使用效率。

模型的推理通常在GPU上进行，而显存空间有限。显存的分配主要用于模型参数、前向计算激活值（固定部分）以及KV Cache（动态增长部分）。随着输入长度增加，KV Cache的占用量也增加，最终可能超出单张GPU甚至整台设备的显存。

因此，优化KV Cache可以实现更多内容的推理，同时减少跨设备通信带来的性能瓶颈。更少设备上的推理意味着更快的推理速度和更高的吞吐量，从而有效降低推理成本。

MQA（Multi-head Query Attention）

思路：直接让所有Attention Head共享同一个K、V
使用MQA的模型包括PaLM、StarCoder、Gemini等。MQA直接将KV Cache减少到了原来的$1/h$，从节省了大量的显存。

从效果方面看，大部分任务的损失都比较有限，通过进一步训练也许可以弥补。此外，MQA由于共享了K、V，将会导致Attention的参数量减少了将近一半，这部分参数量可以增加到FFN/GLU模块上，能起到弥补作用。

GQA（Grouped Query Attention）

思路：将所有Head分为n个组，每组共享同一对K、V，GQA通过调整Attention的组数提供了从MHA到MQA的自然过渡，当组数为1时就是MQA，而组数等于head数时就是MHA。当组数介于两者之间时，GQA能够压缩KV Cache，同时保留更多的自由度，以达到较好的性能与压缩率平衡。

使用GQA的模型：Meta的LLAMA2-70B和LLAMA3、ChatGLM2和ChatGLM3。LLAMA2/3-70B中的GQA通常设置为组数=8，这种配置有助于提升推理效率。由于70B体量的模型无法在单张卡上运行，需要多张卡协同计算。通过GQA的分组机制，就可以让每张卡分别负责一组K、V对应的Attention Head，从而尽可能减少卡间通信，实现更高的推理效率。

MLA（Multi-head Latent Attention）

对于这部分的证明过程，笔者看得还不是特别明白。只能简单介绍一下MLA的方法思路。

思路：首先将向量分为两份，一份作为k，一份作为v。然后每份再均分成g份，这两个g份各自赋值h/g次，总共得到h个k和h个v，作为h个自注意力头的k和v，从而增加了性能表现。在训练阶段，维持不变。推理阶段，通过等价变换吸收$q_i$和$k_i$中的投影矩阵，最后通过$c_i$来表示$q_i$和$k_i$，并且是所有头共享的，从而减少了KV cache。

注意：
1.在MLA的矩阵的恒等变换中，存在一个步骤，只有理论上（在无限精度下）才成立，所以如果从fp32换到fp16降低精度后会造成模型效果的损失。
2.由于吸收投影矩阵的变换中，如果直接在MLA中使用RoPE，则需要在合并的投影矩阵中间插入一个与$t-i$（位置差）有关的项。所以RoPE用于MLA后，不能减少KV cache。

deepseek的团队和苏神交流后，最终采用了一种混合的方法——每个Attention Head的Q、K新增个维度用来添加RoPE，其中K新增的维度每个Head共享。在MLA的最终版本里，为了减少训练期间参数量和相应的梯度，deepseek团队还将Q的输入也改为了低秩投影形式。

总结

关于KV cache的改进，也许能决定大模型api应用价格是否白菜化。目前deepseek声称能做到1元/1B Tokens，笔者之前做过一些数据集的自动标注工作，大概2万条数据集，句子平均50个tokens，感觉数据量加起来也刚好到1B的Tokens。如果以1元/1B Tokens的价格，标注100万条数据也才50元，期待未来吧。