KV cache 以及 Attention 各种变种

KV cache

自回归的时候，只要 attention_map 中的最后一行用于计算 [B, 1, d_model]

优化 decoder 部分的 Q @ K @ V 部分（cache 的是 decoder 的casual attention 部分）

由于auto regressive的特性，每次的输出都会重新拼到尾部作为输入的一部分

所以之前的计算是重复的，可以缓存起来

例如下图中的白色部分就是重复的，每次只需要拿最新生成的 token 的 Q 进行计算即可

于是可以很容易回答一题常见的面试题：为什么没有 Q cache？

• 因为每次的计算就是那最新的 Q 来计算

KV cache 实际上就是用空间换时间，所以后续很多工作都是优化 KV cache 过大的问题

对于 MHA，由于每个 head 都有一对 KV，所以需要缓存的数据量很大

Multi-Query Attention

一个多头中只有一份KV

• MHA 中，输入经过 $W_Q,W_K,W_V$ 后，将 d_model 维的embedding split为 num_head * head_dim，构成多头的输入

• 在MQA中，输入经过 $W_Q,W_K,W_V$ 后，只将 Q 切分(与原来一样)，K和V是通过线性变换降维到head_dim，所以这 n_heads 个头的 KV是同一个

但是这样子做明显会减少模型的表现能力

Group-Query Attention

GQA 的想法是取 MHA 和 MQA 的折中

• 既然MQA对效果有点影响，MHA缓存又存不下
• 只使用一套共享的 KV 效果，那就改为多套 KV

具体做法为：

• 将 Q 分为多个 group，每个 group 内贡献一组 KV
• 这样子就不是所有 Q 都是同一套了

MLA

Multi-Head Latent Attention (MLA) 的核心思想是通过低秩联合压缩技术，减少 K 和 V 矩阵的存储和计算开销。