缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA

注意力是transformer很重要的一个概念，如果没有注意力的改进，大模型不会出现如此好的效率和智能，所以对注意力算法的了解也是至关重要

引用： 苏剑林老师的博客 https://kexue.fm/archives/10091

注意力相关论文

❝

MHA（Multi-Head-Attention）： https://arxiv.org/abs/1706.03762

MQA （Multi-Query Attention）： https://arxiv.org/abs/1911.02150

GQA (Grouped-Query Attention) ： https://arxiv.org/pdf/2305.13245

MLA（Multi-head Latent Attention） : https://arxiv.org/abs/2405.04434

FlashAttention : https://arxiv.org/abs/2205.14135

PageAttention： https://arxiv.org/abs/2205.09729

❝

为什么这篇文章我会说“缓存与效果的极限拉扯”呢？众所周知一般情况下LLM的推理都是在GPU上进行，单张GPU的显存是有限的，一部分我们要用来存放模型的参数和前向计算的激活值，这部分依赖于模型的体量，选定模型后它就是个常数；

另外一部分,为了加速推理，我们要用来存放模型的KV
Cache，这部分不仅依赖于模型的体量，还依赖于模型的输入长度，也就是在推理过程中是动态增长的，当Context长度足够长时，KV
Cache的大小就会占主导地位。

注意力机制的优化就是为了 就是为了减少KV Cache，实现更低的推理成本。

Multi-Head-Attention

这一篇也是注意力的开山鼻祖了。多头的多组qkv,可以关注到数据不同维度的特征，唯一的缺点是推理缓存太大了。

MQA （Multi-Query Attention）

MQA是多头注意的一种变体。

MQA的方法是保持Q的初始头数，但K和V只有一个头，这意味着所有Q个头共享相同的K和V，因此称为Multi-Query.

MQA 让所有的头之间 共享 同一份 Key 和 Value 矩阵，每个头只单独保留了一份 Query 参数，从而大大减少 Key 和 Value
矩阵的参数量。使用矩阵乘法 matmul 来广播，使得每个头都乘以这同一个 tensor，以此来实现参数共享

MQA通过以下方法实现推理加速:

1、KV缓存大小减少了h(头数量)，这意味着需要存储在GPU内存中的张量也减少了。节省的空间可以用来增加批大小，从而提高效率。

2、减少了从内存中读取的数据量，从而减少了计算单元的等待时间，提高了计算利用率。

3、MQA有一个相对较小的KV数量，可以放入缓存(SRAM)中。MHA则需要较大的KV数量，不能完全存储在缓存中，需要从GPU内存(DRAM)读取，这很耗时

GQA (Group Query Attention)

然而，也有人担心MQA对KV Cache的压缩太严重，以至于会影响模型的学习效率以及最终效果。为此，一个MHA与MQA之间的过渡版本GQA

他的思想是将所有Head分为g个组（g可以整除h），每组共享同一对K、V

MLA (Multi-head Latent Attention)

MLA 将 KV vectors 压缩为 latent vector ，从而在推理时仅需保存 latent vector 而无需保存 KV cache
( , 为 head dim, 为 #heads) 数学表达式如下:

但问题是不能很好兼容Rope