KV-Cache技术小结（MHA,GQA,MQA,MLA)

type

status

date

slug

summary

1 背景

KV-cache技术是目前LLM，VLLM等自回归模型常用的避免冗余计算的手段。但引入该技术需要额外的存储成本。原生的kv-cache所需的存储成本与生成的token长度成正比，是目前长文本生成的主要瓶颈之一。目前针对如何降低KV-cache的存储成本激起大量研究者广泛关注。GQA (Group Query Attention)，MQA (Multi Query Attention)，MLA (Multi-Head Latent Attention)是目前常用的方法。本文将从经典的casual attention出发，阐述kv-cache的必要性，及目前常见优化kv-cache的手段。

TL,DR

不同attention方法 KV cache的存储单元数量

ㅤ	KV-cache	存储单元数量
Casual Attention
`MHA`
`GQA`
`MQA`
`MLA`

2 经典Casual-Attention的KV-Cache工作机制

假定当前层attention的输入为为sequence的长度。通过3个线性层得到query ，key，value。

随后通过标准的casual attention机制，得到输出。

在训练阶段，为了通过teaching forcing技巧进行并行化训练，引入一个causal mask ，来保证位置的token只看的到的token。这个阶段没有kv-cache。

在生成阶段。token是按序生成，在模型内部体现在的每一行是依次输出的。

对于第一个token的生成只依赖,第二个token的生成只依赖，依次类推。

对每一个，attention的计算如下

如果考虑位置编码，上式写改写为,表示位置编码函数

从上面的计算流程不难看出，的生成只依赖当前位置的query,依赖前面所有位置的key和value。为了得到最naive的做法是：生成位置的token时，将作为Attention的输入，以此保证能够被正确计算。Naive的做法没有kv-cache。

但从上面的计算流程我们不难看出，需要的中已经在的计算中被计算。因此可以能把算好的保存起来，在位置只需计算，再与前面的进行拼接就可以得到。这样大大减少了冗余的计算量。这就是kv-cache的核心motivation。（公式中被“框起来”的部分是可以cache的。）

用kv-cache的生成思路，

生成第个token时，此时attention层输入，输出。是缓存的kv-cache 生成第个token时，此时attention层输入，输出 … 生成第个token时，此时attention层输入，输出 …

kv-cache能够显著降低attention的计算量，但随着生成token的增多，kv-cache所需的存储成本呈线性增加，导致GPU的显存成为生成长度的瓶颈。

3 Multi-Head Attention（`MHA`） KV-Cache工作机制

paper: https://arxiv.org/pdf/1706.03762

MHA是上面的一个推广。假定MHA的输入为为sequence的长度。假定有个head，每个head投影的维度为

通过线性层的矩阵计算，得到不同head下的

表示head query矩阵在序列位置为处的向量，其他符号记法类似。

在生成阶段每一个head经过attention计算后的位置的输出如下（依序生成）,

循环Attention head ,可以计算所有head 时刻的输出最后将不同head的输出进行拼接，得到最终的输出

对于MHA而言，的计算缓存的kv-cache为

4 Group Query Attention（`GQA`） KV-Cache工作机制

paper: https://arxiv.org/pdf/2305.13245

GQA的attention计算机制与MHA一致。有所区别的是，GQA为了降低KV-Cache的存储，将attention的head分为了组，同一组共享kv-cache。

是向上取整符号。若,那么共享这个group的key，value。

同样，在生成阶段依序生成。每一个head经过attention计算后的位置的输出如下,

Loop Attention head ,可以计算所有head 时刻的输出最后将不同head的输出进行拼接，得到最终的输出

对于GQA而言，的计算缓存的kv-cache为 ,相比标准的MHA，KV-cache降低了倍

5 Multi Query Attention（`MQA`） KV-Cache工作机制

paper: https://arxiv.org/pdf/1911.02150

MQA是GQA的一个特例。当时，即所有head的query共享同一组key, value，此时的GQA成为MQA。

对于MQA而言，的计算缓存的kv-cache为 ,相比标准的MHA，KV-cache降低了倍。

6 Multi-Head Latent Attention (MLA)的工作机制

MLA是deepseek提出的一项针对kv-cache的优化。

paper: https://arxiv.org/abs/2405.04434

（一）先抛开位置编码

假定MLA的输入为为sequence的长度。假定有个head，每个head投影的维度为。初步来看未引入位置编码的MLA像是引入了一个低秩分解矩阵(类似LoRA的做法)的MHA

head 的Q,K,V计算过程如下

矩阵维度变换说明

ㅤ	dimension
（下标kv表示key-value的compress latent编码投影矩阵，`/A`类比LORA的A矩阵)
（`/B` 类比LORA的B矩阵， ``)表示attention第个head


（上标q表示 query的compress latent code）
（表示 key-value的compress latent code）

在生成阶段,每一个head经过attention计算后的位置的输出如下（依序生成)

式（9）和MHA的generate阶段的形式相同，当然可以通过缓存实现kv-cache。

但MLA提供了一个新的方法（式10），只需要缓存即可，相比原始方法的kv-cache的存储单元数量从降低到。但这个方法需要引入两个矩阵乘法的计算量。因为不大，引入的计算量是可以接受的。（还有一种更为巧妙的方法能规避计算量增加的问题，在（二）中介绍）

（二）引入位置编码的MLA

这个形式也是deepseek论文中提出的形式。有了上面的基础，再理解就很简单了。同样假定MLA的输入为为sequence的长度。假定有个head，每个head投影的维度为

head 的Q,K,V计算过程如下

矩阵维度变换说明

ㅤ	dimension
(下标kv表示key-value的compress latent编码投影矩阵，`/A`类比LORA的A矩阵)
(`/B` 类比LORA的B矩阵， ``)表示attention第个head




（上标q表示 query的compress latent code）
（表示 key-value的compress latent code）
（上标C表示compression的首字母“C”）	(不含位置编码的query)
（上标R是RoPE的`R`）	(包含位置编码的query)