KV缓存：加速LLM推理
在使用LLM时，推理过程中的速度和效率变得至关重要，尤其是在模型变得越来越大、序列越来越长的情况下。一个关键的优化技术是KV缓存，它能够实现快速且内存高效的生成。
在这篇博客中，我们将分解KV缓存是什么，它是如何在后台工作的，以及何时有用（或无用）。我们会保持语言简单，跳过炒作，专注于清晰度。
1、Transformer中的Q、K、V是什么？
在我们进入缓存之前，让我们快速回顾一下变换器中注意力机制的基本原理。
每个进入变换器的标记都会被转换成三个向量：
- 查询（Q） —— 这个标记正在寻找什么
- 键（K） —— 这个标记有什么可以提供的
- 值（V） —— 标记的实际内容
这些向量是通过线性投影（矩阵乘法）从标记嵌入中创建的。然后模型使用它们来计算自注意，这告诉我们每个标记应该对序列中的其他所有标记注意多少。
公式
Attention(Q, K, V) = softmax(Q × K.T / sqrt(d)) × V
这意味着：对于每个查询，将其与所有键进行比较，将这些分数转换为权重，然后使用这些权重来混合值。
2、没有缓存的transformer是如何工作的？
在常规的变换器（如训练时的BERT或GPT）中，我们一次处理整个序列。对于每个标记，我们从头开始计算它的Q、K和V。然后模型可以关注序列中的所有标记。
在训练期间，我们有完整的输入，这是可以接受的。但在推理过程中，特别是文本生成时，我们不会提前知道完整的序列。
没有缓存的话，生成第100个标记需要反复重新计算前99个标记的Q、K和V。这非常低效。
3、引入KV缓存
KV缓存解决了这个问题。
而不是每次都重新计算键（K）和值（V）向量，我们在生成输出时存储它们。当需要生成下一个标记时：
我们仅计算新标记的Q。 我们重用之前标记的缓存K和V。 我们使用新的Q和缓存的K/V来计算注意力。
这减少了冗余计算，并带来了巨大的加速，特别是在生成长序列时。
4、为什么只缓存K和V（不包括Q）
查询（Q）依赖于当前标记，这在每一步都会变化。因此，每次都需要重新计算。
另一方面，之前生成的标记的K和V向量不会改变。它们已经被计算过了，我们可以直接重复使用。这就是为什么只缓存K和V的原因。此外，不需要缓存像W_Q、W_K或W_V这样的模型权重，因为它们是固定的参数，并且已经在内存中。
5、KV缓存是如何工作的？
典型的生成循环如下所示，带有KV缓存：
- 第一个标记： 计算Q、K、V → 将K和V存储在缓存中。
- 第二个标记： 计算Q → 重用缓存的K/V → 计算注意力。
- 第三个标记： 计算Q → 重用不断增长的K/V → 生成下一个标记。
- 重复。
实际上，缓存会逐个标记地增长，并且为每个Transformer层和注意力头单独维护。HuggingFace Transformers等框架会自动管理这一点。
6、KV缓存有多大帮助？
非常大。
没有缓存，生成一个长输出意味着反复重新计算所有先前标记的注意力。 有了缓存，每个新标记的计算几乎保持不变，无论输出有多长。
从性能角度来看：
没有缓存：O(n²) 的复杂度 有缓存：O(n) 的生成复杂度
KV缓存是ChatGPT能够快速响应的原因之一——即使处理数千个标记也是如此。
7、KV缓存会影响准确性吗？
不会，KV缓存是功能上等同于每次重新计算K和V。它只是避免冗余计算的一种快捷方式。
不过，在一些边缘情况下（比如量化模型或近似注意力方法），可能会出现微小的差异——但通常可以忽略不计。
8、什么时候不应该使用KV缓存
KV缓存很强大，但并不是在所有情况下都适用。以下是一些不应该使用KV缓存的情况：
训练或微调：
- 在训练期间，您处理的是完整的序列。
- 所有标记的计算是并行化的。
- 缓存会引入不必要的复杂性，没有任何好处。
双向模型（如BERT）：
- 这些模型同时查看过去和未来的标记。
- KV缓存仅适用于从左到右（自回归）生成。
动态上下文任务：
- 在RAG中，或者在生成过程中输入上下文发生变化的任务中，缓存过时的K/V可能导致错误输出。
低内存环境：
- KV缓存可能会消耗大量的GPU内存——尤其是当上下文很长且模型很大时。
9、什么时候应该使用KV缓存
- 自回归推理：GPT风格的模型逐个生成标记
- 聊天机器人、代码补全、摘要等
- 任何在生成过程中需要考虑延迟和速度的地方
10、结束语
KV缓存是一种简单但强大的技巧，使LLM在推理过程中运行得更快。通过重用模型已经知道的内容，它避免了浪费性的重新计算，并有助于高效地扩展生成到数千个标记。
但就像任何优化一样，它应该在适合的地方使用——并在不适合的地方避免使用。
如果您正在构建或微调LLM，理解缓存的工作原理（以及何时不该使用它）可以帮助您在性能和内存使用方面做出更明智的选择。
原文链接：What is KV Caching? Making LLMs Lighting Fast
汇智网翻译整理，转载请标明出处