大模型训练过程中KVCache与MLA

基础内容

在Transformer模型中，每个token有qkv三个属性，分别通过神经网络变换得到。¹
根据Transformer中注意力公式，每个token的q需要和之前所有的k计算注意力，然后经过Softmax函数后乘以之前所有token的V，得到最终的结果。
也就是说，面对新的token,q是向量，KV是矩阵；q和k的点乘可以看做当前token和之前所有token之间的注意力

-
每个token作为K矩阵中一列和V矩阵中一行
为K计算新列，为V增加新行²
注意的是，只有自注意力层进行KVcache的交互，其他部分不需要(position embedding、Laynorm、Feedforward 不需要)

显存大小计算

以Qwen-72B为例³，模型80层，每层64个头，每个头的维度是128.
那么，一个token对应的qv的数量为2*64*80=10240;如果是半精度存储，每个参数为两个字节；那么为10240*2*128/1024/1024=2.5MB
还需要考虑B和长度大小L；有兴趣可以看³

Example
过程中注意
由于生成第一个token前没有KVcache，所以延迟较高，需要为提示中每一个token计算KV矩阵

可以观察到，GPT不同长度的token收费标准不同；主要原因在于处理/生成文本越多，占用的GPU内存多，收费就比较高

改进优化

MQA

kv共享
kv数量从2*层数*头的数量变为2*层数

GQA

kv数量从2*层数*头的数量变为2*层数*分组数

MLA

根据⁴deepseek V3 报告。

参考链接

---

1. https://www.youtube.com/watch?v=80bIUggRJf4 ↩︎

2. https://github.com/wdndev/llm_interview_note/blob/main/06.%E6%8E%A8%E7%90%86/llm%E6%8E%A8%E7%90%86%E4%BC%98%E5%8C%96%E6%8A%80%E6%9C%AF/llm%E6%8E%A8%E7%90%86%E4%BC%98%E5%8C%96%E6%8A%80%E6%9C%AF.md ↩︎

3. https://zhuanlan.zhihu.com/p/16730036197 ↩︎ ↩︎

4. https://arxiv.org/html/2412.19437v1 ↩︎