在大型语言模型(LLM)的推理过程中,为了提高生成速度,通常会采用一种名为 KV Cache 的优化技术。KV Cache 可以缓存中间计算结果,避免重复计算,从而显著提升推理效率。本文将深入探讨 LLaMA 模型中 KV Cache 的实现逻辑,包括训练和推理阶段的具体操作。
1. 什么是 KV Cache
1.1 为什么需要 KV Cache
在自回归模型中,模型会逐词生成文本。每生成一个新词,都需要进行一次完整的 Transformer 前向计算,包括计算注意力权重。然而,对于已经生成的词,它们的 Key (K) 和 Value (V) 向量在计算注意力时会被重复使用。
KV Cache 的作用就是缓存已经计算过的 Key 和 Value 向量,当生成新的词时,只需要计算新的 Query 向量,并使用缓存的 K 和 V 向量进行注意力计算,从而避免了对整个序列的重复计算。
1.2 KV Cache 的工作原理
具体来说,KV Cache 会缓存每一层 Transformer 编码器或解码器中的 K 和 V 向量。在推理过程中:
- 首次计算: 模型会计算完整输入序列的 K、V 和