LLM 中的 KV 缓存

KV Caching in LLMs

KV 缓存（KV caching）是加速 LLM 推理的一种常用技术。为了让你有个直观的感受，先看看下图中推理速度的差异：

• 使用 KV 缓存 → 9 秒
• 不使用 KV 缓存 → 40 秒（约慢 4.5 倍，而且随着生成的 token 越来越多，这个差距还会继续拉大）。接下来，我们通过图示直观地理解一下 KV 缓存是如何工作的。

要理解 KV 缓存，我们首先得知道 LLM 是如何输出 token 的。

如上图所示：

• Transformer 会为所有 token 生成隐藏状态（hidden state）。
• 这些隐藏状态会被投影到词表（vocab）空间。
• 最后一个 token 的 logits 会被用来生成下一个 token。
• 后续 token 重复这一过程。

因此，要生成一个新 token，我们只需要最近这个 token 的隐藏状态，其他 token 的隐藏状态都用不上。

接下来，我们看看在 Transformer 层内部，最后一个隐藏状态是如何通过注意力机制计算出来的。

在注意力计算中，我们首先要做 query 矩阵和 key 矩阵的乘积，其中最后一行涉及的是最后一个 token 的 query 向量，以及全部的 key 向量：

推理过程中，其他 query 向量都用不上。同样地，最终注意力结果的最后一行，涉及的也是最后一个 query 向量，以及全部的 key 和 value 向量。看下面这张图会更直观：

上面这个洞见说明，网络中的每一次注意力计算，要生成一个新 token，实际上只需要：

• 最后一个 token 的 query 向量。
• 全部的 key 和 value 向量。

但这里还有一个关键的洞见：在生成新 token 的过程中，此前所有 token 对应的 KV 向量都不会发生变化。

因此，我们只需要为上一步生成的那个 token 生成新的 KV 向量，其余的 KV 向量都可以直接从缓存中取出，从而节省计算和时间。这就是 KV 缓存！换句话说：与其反复重复计算所有上下文 token 的 KV 向量，不如把它们缓存起来。

生成一个 token 的流程如下：

• 为上一步生成的 token 计算 QKV 向量。
• 从缓存中取出其余所有的 KV 向量。
• 计算注意力。
• 把新生成的 KV 值存入缓存。

可以看出，这样做能在推理时节省大量时间。事实上，这也是为什么 ChatGPT 生成第一个 token 会比后续 token 慢一些——在那短暂的停顿期间，系统正在计算提示词的 KV 缓存。话虽如此，KV 缓存也会占用大量内存。以 Llama3-70B 为例，它的参数是：

• 总层数 = 80
• 隐藏层维度（hidden size）= 8k
• 最大输出长度 = 4k

由此可得：

• 每个 token 在 KV 缓存中占用约 2.5 MB。
• 4k 个 token 将占用 10.5 GB。

用户越多 → 占用的内存也越多。