理解大模型推理中的 KV Cache

1. 从因果解码器（causal decoder）说起

KV Cache 的可行性源于 decoder-only 模型（如 GPT、LLaMA 等）中所使用的因果掩码（Causal Mask）机制。

在 causal decoder 架构下，token 序列中每个位置的 attention 计算仅依赖于其前文信息。
这是通过一个下三角掩码矩阵（主对角线右上角元素全为负无穷，softmax 计算后变为全 0）实现的，确保当前 token 无法看到之后的 token，即"只能看左边"。

正因为这种严格的因果性，先前所有 token 的 Key 和 Value 向量一旦计算完成，便可安全缓存。
在后续生成新 token 的步骤中，可直接复用这些缓存结果，而无需重新计算历史 token 的 K、V，通过空间换时间来提升推理生成效率。

Decoder 自回归生成 token 过程

Decoder 自回归生成 token 过程 [1]

2. 推理生成的两个阶段：Prefill 与 Decode

在深入探讨 KV Cache 和 PD 分离之前，我们首先需要理解大语言模型（LLM）推理过程的两个主要阶段。
即整个推理流程可划分为以下两阶段：

  • Prefill 阶段（预填充阶段）：该阶段属于计算密集型操作。模型并行处理输入序列（prompt）中的所有 token，生成相应的 KV Cache，并输出第一个 token；
  • Decode 阶段（解码阶段）：该阶段属于内存访问密集型操作。模型基于已生成的 KV Cache，以自回归方式依次生成后续的每个 token，每次前向传播仅产生一个新 token。

生成式 LLM 推理的两个阶段

生成式 LLM 推理的两个阶段：prefill 阶段是计算密集型的，而 decode 阶段是访存密集型的。[2]

3. 有 KV Cache 和无 KV Cache 的差异

一个标准的 Transformer Decoder Block 主要由自注意力（Self-Attention）和前馈网络（FFN，Feed Forward Network）两个核心模块构成，中间通过残差连接与层归一化（Add & Norm）进行衔接。
为简化分析，我们暂不讨论 Add & Norm 这些计算量小的操作。
FFN 本质上是全连接层（也可扩展为混合专家模型 MoE 形式），这里也不考虑。
本文重点放在注意力机制的计算过程上。

decoder-only 模型结构

decoder-only 模型结构 [5]

正如前文所述，大语言模型（LLM）的推理过程是指用户输入一个提示（Prompt），例如序列 [t₀, t₁, t₂]，模型基于该输入自动生成后续内容，完成应答或续写任务。

LLM 推理生成新 token，同时缓存新 token 的 KV

LLM 推理生成新 token，同时缓存新 token 的 KV [6]

3.1. Prefill 阶段的计算过程

为简化描述，此处不考虑多头机制（Multi-Head）。
假设输入 Prompt 包含 3 个 Token，则 Decoder Block 的输入为一个第一维 size 为 3 的张量。
该输入分别与权重矩阵 W_q、W_k、W_v 相乘，得到每个 Token 对应的 Query、Key 和 Value 向量。
随后进行以下计算：

  1. 将 Q 与 K 相乘，得到注意力分数矩阵；
  2. 对该分数矩阵进行掩码 Softmax 处理，确保因果性；
  3. 将结果与 V 相乘，汇总为注意力输出。

最终所得即为当前层的自注意力输出。图示如下（作图参考[7]）：

prefill 阶段 attention 计算流程

prefill 阶段 attention 计算流程

Prefill 阶段最终输出一个 token：在最后一个 decoder block 中，仅使用 QKV 输出的最后一行进行后续计算生成该 token。

Prefill 阶段仅使用最后一行 QKV 输出生成 token

是否使用 KV Cache 对 token 生成效率影响显著。先简单看看不使用 KV Cache 的生成过程。

3.2. 不使用 KV Cache：重复执行 Prefill

不使用 KV Cache 时，每生成一个新 token，就将其拼接到当前序列末尾，将整个新序列再次输入模型，执行完整的 Prefill 过程，以计算下一个 token。
该过程不断重复，实现自回归生成。
这种方式存在大量冗余计算：每次生成时，都需要为所有历史 token 重新计算 QKV 向量。
事实上，在 Decoder-only 模型中，每个 token 的 K 和 V 值一旦计算完成，在后续生成中不会改变。
因此，可将这些 KV 值缓存起来，后续 Decode 时直接复用，避免重复计算，从而显著减少计算量。

3.3. 使用 KV Cache：高效 Decode

使用 KV Cache 后，每次生成新 token 时无需再输入完整的历史序列，只需将前一个生成的 token 作为输入即可生成下一个 token，计算过程如下图所示（使用 t3 生成 t4，t0-t2 的 KV 已缓存）：

decode 阶段生成新 token 的 attention 计算

decode 阶段生成新 token 的 attention 计算

由上图可看出，在生成每个新 token 时，仅需使用当前 token 的 Query 与前序（包括自己）所有 token 的 Key 和 Value 进行计算，即可得到注意力输出。
前序 token 的 Query 并不会参与当前 attention 计算，因此无需缓存。

对比总结：

  • 没有 KV Cache：在解码第 t 个 token 时，需要将第 t 个 token 与之前所有 t-1 个 token 重新计算一次注意力。
    这意味着 attention 计算量会随着生成长度 n 呈平方级（O(n²)）增长，同时 FFN 计算量也会线性增长。
  • 有 KV Cache：在解码第 t 个 token 时，之前所有词元的 K、V 已缓存好。
    模型只需要计算当前新词元的 QKV，去与 Cache 中所有 K 做注意力计算，再与 V 加权求和。
    计算量仅随生成长度 n 呈线性（O(n)）增长，同时 FFN 计算量稳定。

4. PD 分离与 KV Cache 传输

鉴于 Prefill 与 Decode 阶段对计算和访存带宽等资源的需求存在显著差异，目前许多 LLM 推理服务采用 PD 分离的部署方式，即将两个阶段分别部署于不同的服务器。
这种架构引入了 KV Cache 的传输需求：Prefill 阶段产生的 KV Cache 需传输至 Decode 服务器以供使用。
有关 KV Cache 传输的性能问题，后续再进一步学习研究。

5. 参考资料