KV Cache：大模型推理加速的核心机制

当 AI 模型生成文本时，它们经常会重复许多相同的计算，这会降低速度。KV Cache 是一种技术，它可以通过记住之前步骤中的重要信息来加快此过程。模型无需从头开始重新计算所有内容，而是重复使用已经计算过的内容，从而使文本生成更快、更高效。

从矩阵运算角度理解 KV Cache

让我们从最基础的注意力机制开始。标准的 self-attention 计算公式大家都很熟悉：

Attention(Q,K,V) = softmax(QK^T/√d_k)V

在实际应用中，随着上下文长度的增加，这个计算会变得非常昂贵。比如当我们有 10,000 个 token 时，QK^T 会产生一个 10,000×10,000 的巨大矩阵。在自回归生成过程中，每次预测新 token 都需要重新计算整个注意力矩阵。但仔细观察会发现，对于已经生成的 token，它们的 K 和 V 向量在每次计算中都是相同的。

以"我爱大模型"的生成过程为例：

• 第1步：输入"我"，预测"爱"
• 第2步：输入"我爱"，预测"大"
• 第3步：输入"我爱大"，预测"模"
• 第4步：输入"我爱大模"，预测"型"

在第4步计算时，"我爱大模"这四个字的 K 和 V 值与前面步骤中计算的完全相同。如果每次都重新计算，就是巨大的浪费。

从矩阵维度来看，每一步的计算实际上是在上一步的基础上增加一行一列。masking 机制确保了只有下三角部分参与计算，这意味着上一步计算的结果可以完全复用。

KV Cache 的核心思想是：缓存之前计算过的 K 和 V 向量，每次只计算新增 token 的部分。

为什么不缓存 Q？ 因为 Q 向量始终是当前步骤新生成的 token，没有复用价值。每次生成新 token 时，Q 是查询的量，即该值是基于每次的新 token 计算的。

具体来说：

• 第1步：计算并缓存"我"的 K 和 V
• 第2步：只计算"爱"的 K 和 V，与缓存的"我"组合使用
• 第3步：只计算"大"的 K 和 V，与缓存的"我爱"组合使用
• 依此类推…

这样，每一步的计算量从 O(n²) 降低到 O(n)，其中 n 是当前序列长度。

让我们更仔细地看看这个过程：

第2步详细分析：
当我们需要预测"大"时，传统方法会重新计算"我"和"爱"之间的所有注意力关系。但实际上：

• "我"的 K 和 V 向量在第1步已经计算过
• 我们只需要计算"爱"与"我"的关系，以及"爱"与自己的关系
• 缓存的"我"的向量直接复用，避免重复计算

第3步时：

• "我"和"爱"的 K、V 向量都已缓存
• 只需计算"大"与之前所有 token 的关系
• 上一步形成的注意力权重矩阵的上三