深入 KV Cache 的运作过程

KV Cache 的工作主要发生在 Transformer 模型的 Decoder Block 中，特别是其多头自注意力（Multi-Head Self-Attention）层。

整个推理过程通常分为两个阶段：预填充阶段 (Pre-fill) 和 解码阶段 (Decoding)。

这个阶段处理用户的整个输入 Prompt（比如 100 个 token）。

输入与投影： 完整的输入序列 $[t_1, t_2, \dots, t_{100}]$ 进入 Transformer 的每一层。
计算 $Q, K, V$ ： 在自注意力层中，模型使用权重矩阵 $W_Q, W_K, W_V$ 对输入向量进行投影，一次性计算出所有 token 的 $Q$ 、 $K$ 和 $V$ 矩阵：
- $K_{\text{full}} = [K_1, K_2, \dots, K_{100}]$
- $V_{\text{full}} = [V_1, V_2, \dots, V_{100}]$
自注意力计算： 模型计算完整的自注意力，生成 Prompt 的编码表示。
缓存 $K$ 和 $V$ ： 关键步骤。 计算得到的

$K_{\text{full}}$

$V_{\text{full}}$

矩阵被存储到 GPU 显存中的 KV Cache 区域。

特点： 这是一个高度并行化的过程（所有 token 同时计算），速度快，但计算量大（二次复杂度 $O(L^2)$ ）。

这个阶段是模型逐个生成新的输出 token。假设模型现在要生成第 101 个 token 。

输入 Q： 模型的输入是上一步生成的最后一个 token

$t_{100}$

。
- 模型计算
  $t_{101}$
  对应的
  $Q_{\text{new}}$
  向量。
获取 K和 V： 模型计算

$t_{101}$

对应的

$K_{\text{new}}$

和

$V_{\text{new}}$

向量。
- 模型从 KV Cache 中读取上一个阶段存储的所有
  $K_{\text{cache}}$
  和
  $V_{\text{cache}}$
  。
拼接 K 和 V：
- 将新的
  
  $K_{\text{new}}$
  
  向量追加到缓存的
  
  $K_{\text{cache}}$
  
  后面，形成完整的 K’ 矩阵：
  
  $K' = [K_{\text{cache}}, K_{\text{new}}]$
- 同样，将
  
  $V_{\text{new}}$
  
  追加到缓存的
  
  $V_{\text{cache}}$
  
  后面，形成完整的 V’ 矩阵：
  
  $V' = [V_{\text{cache}}, V_{\text{new}}]$
注意力计算：模型使用新的

$Q_{\text{new}}$

与拼接后的 K’ 和 V’ 进行注意力计算：

$\text{Attention}_{\text{new}} = \text{Softmax}\left(\frac{Q_{\text{new}} (K')^T}{\sqrt{d_k}}\right) V'$
生成下一个 Token： 注意力输出经过后续的 Feed-Forward 层和 Softmax 预测，生成下一个 token

$t_{101}$

。
更新缓存： $K_{\text{new}}$ 和 $V_{\text{new}}$ 向量被永久保存并追加到 KV Cache 中，供下一个

$token t_{102}$

使用。

特点： 这是一个串行自回归的过程（一次只能生成一个 token）。最重要的是，**每次计算 K’ 和 V’ 时，不需要重新计算 Prompt 部分的 K/V。**计算量大大降低（线性复杂度 O(L)），但因为是串行的，总耗时依赖于生成长度。

维度	无 KV Cache	使用 KV Cache
计算 $K/V$	每生成一个 token，都需要重新计算所有先前 token 的 K/V。	只计算当前新 token 的 K/V，并从缓存中获取历史 K/V。
时间复杂度	在 L长度的序列上，每次计算 K/V 的复杂度是 O(L^2)。	每次生成一个 token 的计算复杂度是 O(L)。
显存代价	极低（仅存储模型权重）。	高，需要存储所有 Transformer 层、所有注意力头的历史 K/V 向量。
推理速度	极慢（尤其在长序列上）。	显著加快。