KV Cache

KVCache 是一种在大型语言模型（LLM）推理过程中广泛使用的优化技术。它的核心思想是缓存先前计算过的键（Key）和值（Value），从而在生成新 Token 时避免重复计算，显著提升模型的推理速度。本文将介绍 KVCache 的基本原理、工作方式及其优缺点。

1. 背景：Transformer 与自注意力机制

要理解 KVCache，首先需要回顾 Transformer 模型中的核心组件：自注意力（Self-Attention）机制。在自注意力层中，输入序列的每个 Token 都会生成三个向量：查询（Query, Q）、键（Key, K）和值（Value, V）。这三个向量是通过将 Token 的嵌入表示与三个权重矩阵 W_Q, W_K, W_V 相乘得到的。

对于输入序列 X = (x_1, x_2, \dots, x_t)，在第 i 个位置的 Q, K, V 向量计算如下：
Q_i = x_i W_Q \\ K_i = x_i W_K \\ V_i = x_i W_V

然后，模型通过计算 Q 和 K 的点积来得到注意力分数，这个分数决定了在生成当前位置的输出时，应该对输入序列中的其他 Token 给予多少关注。最后，将这些分数与 V向量加权求和，得到该位置的输出。整个过程可以由以下公式表示：

其中 d_k 是 K 向量的维度。

2. LLM 推理中的计算冗余

大型语言模型通常采用自回归（Autoregressive）的方式进行文本生成。这意味着模型逐个生成 Token，并且每生成一个新 Token，都会将其附加到输入序列中，作为下一次生成步骤的上下文。

假设模型已经生成了 t 个 Token，即 (x_1, x_2, \dots, x_t)。为了生成第 t+1 个 Token，模型需要计算当前 Token x_t 的 Query 向量 Q_t，并让它与序列中所有历史 Token 的 Key 向量 (K_1, K_2, \dots, K_t) 进行注意力计算。

在没有优化的情况下，每一步生成过程都需要重新计算所有先前 Token 的 K 和 V 向量。例如：

• **第 1 步**：输入 x_1，计算 Q_1, K_1, V_1。
• **第 2 步**：输入 (x_1, x_2)，重新计算 (K_1, V_1) 和 (K_2, V_2)。
• **第 t 步**：输入 (x_1, \dots, x_t)，重新计算 (K_1, V_1), \dots, (K_t, V_t)。

可以看到，K_1, V_1, \dots, K_{t-1}, V_{t-1} 在每一步都被重复计算，这造成了巨大的计算浪费，尤其是在序列很长时，推理速度会变得非常缓慢。

3. KV Cache 的工作原理

KV Cache 的核心思想就是为了解决上述的计算冗余问题。它通过缓存（Cache）已经计算过的 K 和 V 向量来避免重复计算。

具体来说，KV Cache 的工作流程如下：

1. 在生成第 1 个 Token 时，计算并保存 K_1 和 V_1 到缓存中。此时缓存为 (K_1, V_1)。
2. 在生成第 2 个 Token 时，模型只需要计算 Q_2, K_2, V_2。然后从缓存中读取 K_1, V_1，并将 K_2, V_2 追加到缓存中。注意力计算在 (Q_2) 和 (K_1, K_2) 之间进行。此时缓存更新为 (K_1, V_1), (K_2, V_2)。
3. 在生成第 t 个 Token 时，模型只需为当前 Token x_t 计算 Q_t, K_t, V_t。然后，将它们与缓存中所有历史 Token 的 (K_1, \dots, K_{t-1}) 和 (V_1, \dots, V_{t-1}) 拼接起来，形成完整的 K 和 V 矩阵。

注意力计算仅需使用当前的 Q_t 和拼接后的 K_{\text{all}}, V_{\text{all}}：

通过这种方式，每一步的计算量从与整个序列长度的平方相关降低到与序列长度的线性相关，极大地提升了推理效率。

4. KV Cache 的优缺点

优点

• **显著提升推理速度**：避免了对历史 Token 的重复计算，使得自回归生成过程更加高效。
• **降低计算复杂度**：在每一步生成中，将注意力计算的复杂度从 O(t^2) 降低到 O(t)，其中 t 是序列长度。

缺点

• **增加内存消耗**：KV Cache 需要存储所有历史 Token 的 K 和 V 向量。对于长序列、大模型（层数多、头数多）或大批量（Batch Size），缓存会占用大量显存。
• **内存管理**：缓存的大小与序列长度成正比。对于非常长的文本，需要考虑缓存的管理策略，例如使用滑动窗口（Sliding Window）来限制缓存的大小，只保留最近的 Token 的 K/V。

5. 总结

KV Cache 是一种简单而高效的优化策略，它通过“以空间换时间”的方式，解决了 LLM 在自回归推理中的性能瓶颈。通过缓存历史的键和值向量，KV Cache 避免了大量的冗余计算，使得长文本的生成成为可能。如今，它已成为几乎所有主流 LLM 推理框架的标准配置。

参考文献

[1] Vaswani, A., et al. (2017). "Attention Is All You Need." Advances in Neural Information Processing Systems 30.
[2] 大模型推理加速：看图学KV Cache：https://zhuanlan.zhihu.com/p/662498827
[3] 最通俗介绍大模型的KVCache优化技术：https://www.bilibili.com/video/BV1U9zBYZEg9/?share_source=copy_web&vd_source=08199c1c3936467e69c683e0ea7700e5