在使用大语言模型 (LLM) 时，尤其是在模型规模越来越大，生成序列越来越长的情况下，推理速度和效率至关重要。KV 缓存 (KV Caching) 是一种关键的优化技术，能够显著提高生成速度，并降低内存占用。本文将深入剖析 KV 缓存 的原理、工作机制以及适用场景，帮助读者更好地理解和运用这项技术。
什么是 KV 缓存？从 Transformer 的 Q、K、V 说起
要理解 KV 缓存，首先需要回顾 Transformer 模型中 Attention 机制的基本概念。Transformer 中的每个 Token 都会被转换成三个向量：
- Query (Q)：代表当前 Token 正在“查询”的信息。
- Key (K)：代表每个 Token 可以“提供”的信息。
- Value (V)：代表每个 Token 的实际“内容”。
这些向量通过线性投影 (矩阵乘法) 从 Token 嵌入中生成。模型使用这些向量计算自注意力 (Self-Attention)，从而确定每个 Token 应该关注序列中其他 Token 的程度。
用公式表示：Attention(Q, K, V) = softmax(Q × K.T / sqrt(d)) × V，其中 d 是 K 向量的维度。这意味着对于每个 Query，模型将其与所有 Key 进行比较，将结果转化为权重，然后使用这些权重来混合 Value。
无 KV 缓存的 Transformer 工作方式：效率瓶颈
传统的 Transformer (例如 BERT 或训练期间的 GPT) 会一次性处理整个序列。对于每个 Token，都会从头开始计算其 Q、K 和 V。在训练阶段，由于可以访问完整的输入序列，这种方式没有问题。
然而，在推理阶段，尤其是文本生成任务中，我们通常无法提前获得完整的序列。模型需要逐个生成 Token。如果没有 KV 缓存，生成第 100 个 Token 将需要重新计算前 99 个 Token 的 Q、K 和 V 向量，并重复进行 Attention 计算。这显然是极其低效的。想象一下，你正在用一个大语言模型创作一篇长篇小说，每生成一个新词，都要让模型重新思考之前已经写好的所有内容，这无疑会大大降低创作效率。
KV 缓存：避免重复计算，加速推理
KV 缓存 解决了上述问题。它通过在生成输出时存储 Key (K) 和 Value (V) 向量，避免了重复计算。当需要生成下一个 Token 时，KV 缓存 的工作流程如下：
- 仅计算新 Token 的 Query (Q) 向量。
- 重用之前 Token 存储的 K 和 V 向量 (KV 缓存)。
- 使用新的 Q 向量和缓存的 K/V 向量计算 Attention。
这种方法显著减少了冗余计算，从而大幅提升了生成速度，尤其是在生成长序列时效果更为明显。例如，在生成一篇几千字的报告时，KV 缓存 可以让模型的响应速度提升数倍。
为什么只缓存 K 和 V，而不是 Q？
Query (Q) 向量依赖于当前的 Token，而当前 Token 在每一步都会发生变化，因此必须每次都重新计算。另一方面，先前生成的 Token 的 K 和 V 向量不会改变，可以安全地重用。这就是为什么只缓存 K 和 V 向量的原因。此外，模型权重 (例如 W_Q、W_K 或 W_V) 无需缓存，因为它们是固定参数，已经存储在内存中。
KV 缓存的具体工作流程
以下是带有 KV 缓存 的典型生成循环：
- 第一个 Token：计算 Q、K、V → 将 K 和 V 存储在 KV 缓存 中。
- 第二个 Token：计算 Q → 重用缓存的 K/V → 计算 Attention。
- 第三个 Token：计算 Q → 重用不断增长的 K/V → 生成下一个 Token。
- 重复上述过程。
实际上，KV 缓存 会逐个 Token 增长，并为每个 Transformer 层和 Attention Head 单独维护。像 Hugging Face Transformers 这样的框架会自动处理这些细节，简化了 KV 缓存 的使用。例如，在使用 Transformers 库构建聊天机器人时，可以轻松启用 KV 缓存 来提高对话流畅度。
KV 缓存的性能提升：从 O(n²) 到 O(n)
KV 缓存 带来的性能提升是巨大的。如果没有缓存，生成长输出意味着重复计算所有先前 Token 的 Attention。有了 KV 缓存，每个新 Token 的计算量几乎保持不变，而与输出的长度无关。
在性能方面：
- 无 KV 缓存：时间复杂度为 O(n²)
- 有 KV 缓存：生成的时间复杂度为 O(n)
这里的 n 代表生成的 Token 数量。可以看到，KV 缓存 将时间复杂度从平方级别降低到线性级别，这对于生成长文本的场景来说至关重要。KV 缓存 是 ChatGPT 等模型能够快速响应的关键因素之一，即使在处理数千个 Token 时也是如此。
KV 缓存会影响准确性吗？
一般来说，KV 缓存 在功能上与每次都重新计算 K 和 V 向量是相同的。它只是避免冗余计算的一种快捷方式。因此，KV 缓存 本身通常不会影响模型的准确性。
然而，在一些极端情况下 (例如量化模型或近似 Attention 方法)，可能会出现微小的差异，但这些差异通常可以忽略不计。例如，在使用低精度量化来压缩模型时，KV 缓存 可能会略微放大量化误差，但这通常不会对最终结果产生显著影响。
不适用 KV 缓存的场景
KV 缓存 是一种强大的优化技术，但并非总是适用。以下是一些不应该使用 KV 缓存 的情况：
-
训练或微调：在训练期间，模型需要处理完整的序列，并且计算可以在所有 Token 上并行化。KV 缓存 会引入不必要的复杂性，而没有任何好处。
-
双向模型 (例如 BERT)：这些模型会同时查看过去和未来的 Token。KV 缓存 仅适用于从左到右 (自回归) 的生成。
-
动态上下文任务：在 RAG (检索增强生成) 或输入上下文在生成过程中发生变化的任务中，缓存过时的 K/V 向量可能会导致不正确的输出。例如，在一个需要根据实时信息生成摘要的系统中，如果缓存了过时的信息，可能会导致摘要内容不准确。
-
低内存环境：KV 缓存 会消耗大量的 GPU 内存，尤其是在处理长上下文和大型模型时。例如，在资源受限的边缘设备上部署大语言模型时，可能需要权衡 KV 缓存 带来的性能提升和内存占用之间的关系。
适用 KV 缓存的场景
以下是应该使用 KV 缓存 的场景：
- 自回归推理：GPT 风格的模型逐个生成 Token。
- 聊天机器人、代码完成、摘要等：任何需要在生成过程中考虑延迟和速度的应用。
- 长文本生成：需要生成较长篇幅的文章、报告、故事等。
简单来说，任何注重生成速度和效率的自回归生成任务都适合使用 KV 缓存。
优化 KV 缓存：进一步提升性能
除了直接使用 KV 缓存 外，还可以通过一些策略进一步优化其性能：
-
量化 (Quantization)：使用更低精度的数据类型 (例如 int8 或 float16) 来存储 K 和 V 向量，从而减少内存占用和计算量。例如，可以将 K 和 V 向量从 float32 量化到 int8，从而将内存占用减少 4 倍。
-
稀疏化 (Sparsity)：在 K 和 V 向量中引入稀疏性，即大部分元素为零。这样可以利用稀疏矩阵运算来加速计算。例如，可以使用剪枝 (Pruning) 技术来移除 K 和 V 向量中不重要的元素，从而实现稀疏化。
-
PagedAttention: 是一种用于解决大模型 KV Cache 存储和管理问题的新技术。它借鉴了操作系统中的分页机制，将 KV Cache 在逻辑上分成固定大小的块（Page），并在非连续的物理内存上存储这些块。PagedAttention 的优势在于它可以有效避免 KV Cache 的连续内存碎片问题，从而提高了内存的利用率。与传统的连续内存分配方法相比，PagedAttention 可以显著减少内存浪费，并支持更长的序列长度。
总结
KV 缓存 是一种简单但功能强大的技术，可以显著提高大语言模型在推理过程中的速度。通过重用模型已经“知道”的信息，它可以避免浪费性的重复计算，并帮助高效地将生成扩展到数千个 Token。KV 缓存 几乎已经成为了大型语言模型应用的标配技术。
但是，与任何优化一样，它应该在适合的地方使用，并避免在不适合的地方使用。如果您正在构建或微调大语言模型，了解 KV 缓存 的工作原理 (以及何时不使用它) 可以帮助您在性能和内存使用方面做出更明智的选择。随着大语言模型技术的不断发展，相信 KV 缓存 以及相关的优化技术将会在未来发挥更加重要的作用。