加速大模型推理：深入探究MQA、GQA、MLA（DeepSeek）、KV缓存技术

回顾：多头注意力机制
为什么LLM推理是串行的
KV缓存的挑战
2019年——多查询注意力机制（Multi Query Attention）
2023年5月——分组查询注意力机制（Grouped Query Attention）
2023年——Falcon模型中的修改后的MQA
2024年5月——DeepSeek V2——多头潜在注意力机制（Multi Head Latent Attention）
FastGen：自适应KV缓存（ICLR 2024）：

正如DeepSeek - V2论文所提到的，KV缓存似乎确实是推理时间的瓶颈。到目前为止，为处理这个问题而提出的方法，如多查询注意力（MQA，2019[5]）和GQA（Group Query Attention，2023年5月[6]）都是以性能降低为代价的。但在DeepSeek V2论文中提出的多头潜在注意力（MLA）[4]的情况下，情况并非如此。它实际上比原始Transformers论文中的普通多头注意力（MHA）具有更好的性能（Vaswani等人，2017 [2]）。

Llama 2 & 3.1使用GQA，而Falcon模型使用修改后的MQA。最后，另一篇关于KV缓存的有趣论文是ICLR 2024论文：“模型告诉你要丢弃什么：LLM的自适应KV缓存压缩”[8]。

让我们逐一介绍这些论文中的一些要点。在回顾了推理时间的挑战之后，我还将尝试解释KV缓存的问题。

回顾——多头注意力机制（2017，Vaswani）

为了回顾自我注意的工作原理，我们将输入嵌入与W（q）、W（k）、W（v）矩阵（我们在训练过程中学习的矩阵）相乘，分别得到查询向量Q、键向量K和值向量。

self - attention层的最终输出是：

在Transformer中，我们使用h个不同的W（k）、W（q）、W（v）集合来获得h个不同的自注意力头集。引用Vaswani等人（2017）的原始Transformer论文 [2]：“多头注意力使模型能够共同关注来自不同位置的不同表示子空间的信息。”

所以在h个头和模型维度d的情况下，键、查询维度是d/h，以保持整体计算成本相似。

分析MHA的挑战

一个词的答案是：推理速度慢。让我们在下面详细探讨一下。

所以，需要注意的一个关键点是：Transformer的训练总是并行进行的，即使对于那些具有解码器架构（是自回归）的模型也是如此。如果您不确定如何使用已知的真实数据来完成训练，我建议您查看“教师强迫”的概念，或者说Bishop的《深度学习》一书中的第12.3.1节（“解码器Transformer”）。

为什么LLM推理是串行的

在推理的情况下，与上述训练场景相反，输出Token总是按顺序解码。为了解释这样做的后果，让我引用2019年MQA论文[5] ：在（并行）训练过程的情况下（同时使用一些掩码）：“首先，我们从序列中的n个不同位置生成查询。这些查询都与相同的键和值交互。”

但是，由于解码的自回归性质，情况并非如此。当前第i个位置的自注意力层的token受（i - 1）个位置的输出（即它产生的键值对）和其他先前输出token的影响。但是由于第（i - 1）个标记刚刚输出，我们需要在每个时间戳将新的相应K - V对添加到内存中。因此，无法进行并行计算。因此推理速度很慢。下一节将对此进行更多介绍。

（注意：我们不会每次都加载新的W（k），W（v）矩阵，它们对于给定的头部和层组合是固定的。）

KV缓存的挑战

通过将所有键值对都保存在内存中，我们避免了重新计算，但我们付出的代价是额外的内存消耗。Adaptive KV缓存论文 [8] 中很好地解释了由此带来的挑战。

所以，我们希望将所有先前Token的键值对保留在内存中，以避免重新计算。但是这个KV缓存会随着模型大小和生成长度而显著扩展。但是如果它超过了GPU容量怎么办？我们别无选择，只能将其“卸载”到CPU。示例包括Microsoft针对多个GPU的“深度速度推理”（2022年6月）和针对单个GPU的“FlexGen”（2023年3月）。

“虽然这些方法有助于减轻使用KV缓存对稀缺GPU内存的压力，但由于许多设备上GPU和CPU之间的PCIe带宽有限，将KV缓存卸载到CPU/NVMe仍然会增加生成推理性能的重要开销。”——自适应KV缓存论文 [7]

因此，正如他们在MQA论文的“性能分析”部分所展示的那样，在推理的情况下，内存带宽成为现代计算硬件的主要瓶颈。

多查询注意力机制（2019）的建议

如果我们有一个用于键值的头，同时保持查询部分不变，那么推理解码会变得更快。因此，如果我们在MHA中有（h d k）、（h d v）、（h d q），现在新的维度是（d k）、（d v）和（h d q），其中h是头的数量。

因此，内存需求减少了h倍。实验结果表明，MQA的性能仅比基线（MHA）稍差，远优于削减h、d或k的维度。（请查看下表3）

2023年5月——分组查询注意力机制

这是MHA和MQA之间的权衡。MQA的所有头都有1个键值对，MHA所有头都有h对——而GQA有g对，其中g << h。它比MHA快得多，性能也比MQA好得多。

2023年——Falcon模型中的修改后的MQA

简单说明一下时间线，Falcon 40B于2023年5月发布，180B于2023年9月发布。

他们将MQA调整为他们所谓的“多组”——这个想法与GQA非常相似——唯一的区别是这里的组数与“张量并行秩”或GPU的数量相同。

他们的理由是，每个GPU都需要有一个单独的键值副本；他们需要彼此共享梯度，或者由一个GPU计算梯度并与其他所有人通信。他们建议，为什么不为每个GPU设置一个单独的键值对呢？

Falcon论文中提到的一个重要点是，对于模型训练本身，FlashAttention是大幅加快模型训练速度的关键。但这可能是另一天的话题！

2024年5月——DeepSeek V2——多头潜在注意力机制

从上面回顾一下MHA，给定一个输入向量x ，我们将其与W（k）、W（v）、W（q）相乘，分别得到key、value和queries。它们每个都有一个d（h）的维度，并且总共有n（h）个——n（h）是头的数量。

在这里，我们将输入向量投影到一个较低维度的向量中，重新加载它，然后将其投影到实际的、更高维度的键 - 查询向量中。换句话说，给定输入向量x，我们将其与W（DKV）相乘，下投影矩阵产生c，其维度d（c）远低于d（h）*n（h）。

我们可以通过将c（t）与两个不同的上投影矩阵相乘来获得key和value——因此，在推理时，我们只需要加载c（t）。所以这里的KV缓存只有c（t） * l，其中l是层数。

FastGen：自适应KV缓存（发布于2023年10月，ICLR 2024年）

他们的方法实现了“美洲驼1 - 65B的内存减少约40%，美洲驼1 - 30B的内存减少约30%，美洲驼1 - 13B和美洲驼1 - 7B的内存减少约20%”[8]。由于Llama2 - chat版本使用GQA，因此他们使用了带有MHA版本的原始Llama - 1（2023）。

这个想法是这样的：并非所有的注意力头都是一样的！以下是5种类型：

有些只有简短的本地上下文。
有些人只分别关注特定的标记/标点符号。
有些具有逐列稀疏的注意力图——因此我们可以安全地丢弃最不经常参与的标记。
有些人广泛关注所有代币。

那么为什么不丢弃特定头不感兴趣的代币呢？这种自适应压缩的质量损失可以忽略不计（能够恢复95% + 的注意力分数）。

基于上述结构分析，以下是他们遵循的5种不同的KV缓存压缩策略：

仅保留特殊令牌：在KV缓存中仅保留特殊令牌，例如指令令牌、句子起始令牌。（C特殊政策）
标点符号令牌：在KV缓存中仅保留“.”、“：”、“？”等。
位置性：一旦上下文令牌和当前令牌超过决定的阈值，它就会从KV缓存中被驱逐。
频率（Heavy Hitter）：累积注意力总和作为频率，最频繁的标记保存在内存中。
混合策略：上述所有4个策略的不同组合。他们发现在这些策略中，特殊令牌和标点符号策略对于恢复注意力图尤其重要，尽管它们的数量较少。

现在，不同Attention结构在不同层的分布：

他们的观察：在初始层和最终层，有更多的注意力分配给完整的KV缓存。
在中间层，注意力图集中在特殊令牌上（“对于这些注意力头，特殊令牌上的累积注意力分数高于0.95”[8]）。

因此，他们建议在对每个注意力头应用单独的压缩策略之前，我们应该首先检测这种潜在的注意力图结构。

参考文献

《深度学习——基础与概念》- Christopher M. Bishop和Hugh Bishop。第一版，（2023）
Vaswani, A. (2017). Attention is all you need. Advances in Neural Information Processing Systems.
DeepSeek - V3技术报告，2024年12月 — https://arxiv.org/abs/2412.19437
MULTI HEAD LATENT ATTENTION, Liu, A., Feng, B., Wang, B., Wang, B., Liu, B., Zhao, C., … & Xu, Z. (2024). Deepseek - v2: A strong, economical, and efficient mixture - of - experts language model. arXiv preprint arXiv:2405.04434.
MULTI QUERY ATTENTION : Shazeer, N. (2019). Fast transformer decoding: One write - head is all you need. arXiv preprint arXiv:1911.02150.
Ainslie, J., Lee - Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., & Sanghai, S. (2023). GQA: Training generalized multi - query transformer models from multi - head checkpoints. arXiv preprint arXiv:2305.13245.
Almazrouei, E., Alobeidli, H., Alshamsi, A., Cappelli, A., Cojocaru, R., Debbah, M., … & Penedo, G. (2023). The falcon series of open language models. arXiv preprint arXiv:2311.16867.
Ge, S., Zhang, Y., Liu, L., Zhang, M., Han, J., & Gao, J. (2023). Model tells you what to discard: Adaptive kv cache compression for llms. arXiv preprint arXiv:2310.01801.

本文由mdnice多平台发布

加速大模型推理：深入探究MQA、GQA、MLA（DeepSeek）、KV缓存技术