KV缓存：LLM推理加速的核心技术

每次使用 ChatGPT 或 Claude 时，你一定会注意到第一个 token 的出现时间明显更长，随后的 token 几乎瞬间流出。

在后台，这其实是一个经过深思熟虑的工程决策，称为 KV 缓存（KV caching），目的是让 LLM 推理更快。

在进入技术细节之前，我们先对比一下有缓存和无缓存的 LLM 推理过程：

现在我们从基本原理出发，理解它的工作机制。

第一部分：LLM 如何生成 token

Transformer 处理所有输入 token，并为每个 token 生成一个隐藏状态。这些隐藏状态被投影到词汇空间，产生 logits（词汇表中每个单词对应一个分数）。

GIF

但只有最后一个 token 的 logits 才是关键的。从中采样得到下一个 token，将其追加到输入中，重复该过程。

GIF

关键在于：为了生成下一个 token，你只需要最近一个 token 的隐藏状态。其他所有隐藏状态都是中间副产品。

第二部分：Attention 实际上计算什么

在每个 Transformer 层内，每个 token 都会得到三个向量：查询（Q）、键（K）和值（V）。Attention 将查询与键相乘得到分数，然后用这些分数对值进行加权。

GIF

现在只关注最后一个 token。

GIF

QK^T 的最后一行使用了：

• 最后一个 token 的查询向量
• 序列中的所有键向量

该行最终的 attention 输出使用了：

• 相同的查询向量
• 所有键和值向量

因此，要计算我们唯一需要的隐藏状态，每个 attention 层都需要最新 token 的 Q，以及所有 token 的 K 和 V。

第三部分：其中涉及的重叠计算

生成 token 50 需要 token 1 到 50 的 K 和 V 向量。生成 token 51 需要 token 1 到 51 的 K 和 V 向量。

token 1 到 49 的 K 和 V 向量已经计算过。它们没有变化。同样的输入，同样的输出。然而，模型每一步都从头重新计算它们。

每步都会产生 O(n) 的冗余计算，整个生成过程浪费了 O(n²) 的计算量。

第四部分：解决方案

不必在每一步重新计算所有 K 和 V 向量，而是存储它们。对于每个新 token：

• 只计算最新 token 的 Q、K 和 V。
• 将新的 K 和 V 追加到缓存中。
• 从缓存中检索所有先前的 K 和 V 向量。
• 使用新的 Q 与完整的缓存 K 和 V 运行 attention。

GIF

这就是 KV 缓存。每层每一步存储一个新的 K 和一个新的 V。其他所有内容来自内存。

Attention 计算仍然随序列长度扩展（你要关注所有键和值）。但生成 K 和 V 的昂贵投影只对每个 token 执行一次，而不是每步执行一次。

第五部分：首个 token 的延迟（Time-to-First-Token）

现在你能理解为什么第一个 token 慢了。

当你发送提示时，模型会在单次前向传播中处理整个输入，计算并缓存每个 token 的 K 和 V 向量。这是预填充阶段，也是请求中计算最密集的部分。

缓存预热后，每个后续 token 只需要单次前向传播，且只处理一个 token。

初始延迟被称为首个 token 延迟（TTFT，Time-to-First-Token）。更长的提示意味着更长的预填充时间，也就意味着更长的等待。优化 TTFT（分块预填充、推测解码、提示缓存）本身是一个深入的话题，但根本规律始终不变：构建缓存成本高，读取缓存成本低。

第六部分：权衡

KV 缓存用计算换取内存。每一层为每个 token 存储 K 和 V 向量。对于 Qwen 2.5 72B（80层，32K上下文，隐藏维8192），单个请求的 KV 缓存可能占用数 GB 的 GPU 内存。在数百个并发请求下，它往往超过模型权重本身。

这就是分组查询注意力（GQA）和多查询注意力（MQA）存在的原因：在查询头之间共享键/值头，减少内存开销，且质量损失极小。

这也是为什么将上下文长度翻倍很困难。窗口翻倍，每个请求的 KV 缓存翻倍，并发用户数减少。

还有另一种称为分页注意力的方法可以解决这个问题，我最近在这里讨论过：

tl;dr（摘要）

KV 缓存消除了自回归生成过程中的冗余计算。之前的 token 总是产生相同的 K 和 V 向量，因此计算一次并存储即可。每个新 token 只需要自己的 Q、K 和 V，然后 attention 针对完整缓存运行。

实际应用中可获得5倍加速。代价是 GPU 内存，这成为大规模部署时的瓶颈。每个 LLM 服务框架（vLLM、TGI、TensorRT-LLM）都基于这个思想。

GIF

以上就是全部内容！

• 原文链接： x.com/_avichawla/status/...
• 登链社区 AI 助手，为大家转译优秀英文文章，如有翻译不通的地方，还请包涵～