LLM推理入门指南②：深入解析KV缓存
内容提要
本文讨论了在大语言模型推理中优化KV缓存的方法，以降低注意力计算的成本。通过缓存键和值张量，KV缓存减少了重新计算过去标记的需求，使注意力计算能够随着总序列长度线性扩展。文章介绍了KV缓存的实现和优化策略，并讨论了内存消耗和挑战。最后，提出了减少KV缓存大小的方法，例如减少注意力头的数量和对KV缓存进行量化。
延伸解读
KV缓存：以内存换计算的权衡
KV缓存通过存储键和值张量，避免在生成每个新词元时重新计算过去词元的键和值，从而将注意力计算复杂度从二次方降至线性。然而，这种优化以内存占用为代价：缓存大小随批处理大小和总序列长度线性增长，可能超过模型权重本身。例如，Llama-2-7B在FP16下每词元约0.5MB，28k词元缓存即与模型权重相当。因此，KV缓存是计算与内存之间的折衷，并非一劳永逸。
缓存淘汰策略的挑战与风险
当GPU内存不足时，系统可能采用缓存淘汰策略丢弃旧对话的KV缓存。若用户恢复对话，则需重新计算整个历史缓存，产生与总对话长度成二次方的计算成本。因此，设计高效的淘汰策略需权衡两个高成本选项：消耗稀缺的GPU内存和带宽，或付出二次方级计算量。这要求在实际部署中根据场景（如对话空闲时间）动态调整策略，以平衡资源与延迟。
减少KV缓存大小的主要方法
文章总结了多种减少KV缓存内存占用的方法：减少注意力头数量（如MQA、GQA）、对KV缓存进行量化（如INT8、FP8）、使用局部注意力或滑动窗口（如Mistral-7B）、以及缓存压缩策略（如H2O、Scissorhands、FastGen）。此外，高效内存管理（如PagedAttention）和跨请求缓存共享（如RadixAttention）也能显著提升内存利用率。这些方法可单独或组合使用，以缓解GPU内存压力。
KV缓存对推理性能的实际影响
KV缓存虽降低了计算量，但增加了内存数据传输。在自回归解码阶段，推理受内存带宽限制，大量时间花费在加载权重和不断增长的KV缓存上，而非实际计算，导致GPU利用率低下。因此，减少KV缓存大小不仅能节省内存，还能减少数据传输，从而降低延迟并提高吞吐量。这也是量化、压缩和高效内存管理成为推理优化关键方向的原因。
Q&A
KV缓存在大语言模型推理中有什么作用？
KV缓存通过缓存键和值张量，减少了对过去标记的重新计算需求，使注意力计算能够线性扩展，从而降低计算成本。
KV缓存的内存消耗会有多大？
KV缓存的内存消耗可能非常大，尤其是在处理长序列时，每个词元的内存消耗约为1MB，且随着总序列长度线性增长。
如何优化KV缓存以降低内存消耗？
可以通过减少注意力头的数量和对KV缓存进行量化来优化KV缓存，从而降低内存消耗。
掩码技术在KV缓存中起什么作用？
掩码技术使得每个生成步骤只需计算最后生成的词元的键和值，从而减少冗余计算，提高效率。
KV缓存的实现如何影响注意力计算的复杂度？
KV缓存的实现将注意力计算的复杂度从二次方降低到线性，显著提高了计算效率。
在使用KV缓存时，如何管理GPU内存？
可以通过使用多个GPU或将负载转移到CPU内存来缓解GPU内存不足的问题，同时采用高效的内存管理策略。
