【LLM】KV Cache

KV Cache（Key-Value Cache）是 Transformer 模型推理时的一种优化技术，用于减少计算量并提高效率。

背景

Transformer 模型在生成文本时，每次预测一个新 token 都需要处理整个输入序列，导致重复计算，尤其是在自回归生成任务中（如GPT）。随着序列长度增加，计算量呈平方级增长，影响推理速度。

KV Cache 的作用

KV Cache 通过缓存中间计算结果来避免重复计算，具体来说：

• Key (K) 和 Value (V) 是Transformer注意力机制中的两个矩阵，用于计算注意力权重和上下文表示。
• 在推理时，模型可以缓存这些K和V的值，避免每次生成新token时重新计算。

工作原理

1. 初始计算：对于输入序列，模型计算并存储每个位置的 K 和 V。
2. 缓存使用：生成新token时，模型只需计算当前token的K和V，并与缓存的K和V结合，计算注意力权重和上下文表示。
3. 更新缓存：将新token的K和V加入缓存，供后续步骤使用。

优势

• 减少计算量：避免重复计算，显著提升推理速度。
• 降低内存访问：减少对显存的频繁读写，提高效率。
• 支持长序列生成：通过缓存机制，模型能更高效地处理长序列。

总结

KV Cache通过缓存中间结果，优化了Transformer模型的推理过程，提升了生成任务的效率，尤其在处理长序列时效果显著。

KV Cache 主要用于 自注意力（Self-Attention） 中，而不是 交叉注意力（Cross-Attention）

---

1. 自注意力（Self-Attention）

自注意力是 Transformer 的核心机制，用于计算序列中每个 token 与其他 token 的关系。在解码器中，自注意力的作用是让当前生成的 token 关注之前生成的所有 token。

• KV Cache 的作用：

  • 在自注意力中，模型需要计算每个 token 的 Key (K) 和 Value (V)。
  • 如果没有 KV Cache，每次生成新 token 时，模型都需要重新计算所有 token 的 K 和 V。
  • 通过 KV Cache，模型可以缓存之前 token 的 K 和 V，避免重复计算。

• 适用场景：

  • 自回归生成任务（如 GPT），模型逐词生成输出序列。
  • 每次生成新 token 时，只需要计算当前 token 的 K 和 V，并复用之前缓存的 K 和 V。

---

2. 交叉注意力（Cross-Attention）

交叉注意力通常用于 编码器-解码器架构 中（如 Transformer 用于机器翻译）。它的作用是让解码器的 token 关注编码器的输出。

• KV Cache 的作用：

  • 在交叉注意力中，Key (K) 和 Value (V) 来自编码器的输出，而 Query (Q) 来自解码器的当前 token。
  • 编码器的输出是固定的（因为输入序列是固定的），因此不需要缓存 K 和 V。
  • 每次生成新 token 时，模型只需要从编码器的输出中提取 K 和 V，而不需要重新计算。

• 适用场景：

  • 编码器-解码器任务（如机器翻译），编码器的输出是固定的，解码器逐词生成输出序列。
  • 由于编码器的输出不变，KV Cache 对交叉注意力的优化作用较小。

---

3. KV Cache 的具体应用

• GPT 等纯解码器模型：

  • 只使用自注意力，没有交叉注意力。
  • KV Cache 主要用于缓存自注意力中的 K 和 V。

• 编码器-解码器模型（如 Transformer 用于翻译）：

  • 解码器同时使用自注意力和交叉注意力。
  • KV Cache 主要用于自注意力部分，而交叉注意力部分通常不需要缓存。

---

总结

• KV Cache 主要用于自注意力，而不是交叉注意力。
• 在自注意力中，KV Cache 通过缓存之前 token 的 K 和 V，避免重复计算，从而提高推理效率。
• 在交叉注意力中，由于编码器的输出是固定的，KV Cache 的优化作用较小。