LLM中的KV缓存，清晰详解#

每次使用ChatGPT或Claude时，你一定都注意到了：第一个词元的出现明显更慢，而后续的词元几乎瞬间就能流式输出。

这背后是一个刻意的工程决策，称为KV缓存，其目的是让LLM推理更快。

在深入技术细节之前，我们先来对比一下使用和不使用KV缓存的LLM推理过程：

现在，让我们从基本原理出发，理解它是如何工作的。

第一部分：LLM如何生成词元#

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

但只有最后一个词元的逻辑值才重要。你从中采样，得到下一个词元，将其附加到输入中，然后重复这个过程。

这是关键洞察：要生成下一个词元，你只需要最近一个词元的隐藏状态。所有其他隐藏状态都是中间副产品。

在每个Transformer层内部，每个词元获得三个向量：查询向量（Q）、键向量（K）和值向量（V）。注意力机制将查询向量与键向量相乘得到分数，然后用这些分数对值向量进行加权。

现在只关注最后一个词元。

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

该行最终的注意力输出使用了：

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

生成第50个词元需要第1到50个词元的K和V向量。生成第51个词元需要第1到51个词元的K和V向量。

第1到49个词元的K和V向量已经被计算过了。它们没有改变。相同的输入，相同的输出。然而，模型在每一步都从头开始重新计算它们。

这是每一步O(n)的冗余计算。在整个生成过程中，就是O(n²)的浪费计算。

与其在每一步重新计算所有的K和V向量，不如将它们存储起来。对于每个新词元：

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

注意力计算仍然随序列长度扩展（你需要关注所有的键和值）。但生成K和V的昂贵投影操作，每个词元只发生一次，而不是每一步都发生。

现在你就能明白为什么第一个词元慢了。

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

一旦缓存预热完成，每个后续词元只需要一个词元的一次前向传播。

这个初始延迟被称为首词元生成时间（TTFT）。提示词越长意味着预填充时间越长，等待时间也就越长。优化TTFT（分块预填充、推测解码、提示词缓存）本身就是一个深入的话题，但其动态原理总是相同的：构建缓存很昂贵，从缓存读取很廉价。

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

这就是为什么存在分组查询注意力（GQA）和多查询注意力（MQA）：在查询头之间共享键/值头，减少内存占用，同时质量损失最小。

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

还有一个想法叫分页注意力，它解决了这个问题，我最近在这里讨论过：

> 引用的推文 > https://t.co/zu8HupqrL5 > https://x.com/i/web/status/2031624056072712547

KV缓存消除了自回归生成过程中的冗余计算。之前的词元总是产生相同的K和V向量，所以你计算一次并存储它们。每个新词元只需要自己的Q、K和V。然后，注意力机制针对完整的缓存运行。

在实践中带来5倍的速度提升。代价是GPU内存，这在规模上成为主要约束。每个LLM服务栈（vLLM、TGI、TensorRT-LLM）都建立在这个思想之上。

讲解完毕！

如果你喜欢这个教程：

找到我 → @_avichawla

我每天分享关于数据科学、机器学习、LLM和RAG的教程和见解。