本文深入解析了 LLM 推理中 KV Cache 的工作原理。文章首先阐述了自回归模型的特性，指出在推理过程中，历史 token 的 Key 和 Value 矩阵会被重复计算。通过对比无缓存和有缓存的推理步骤，揭示了只需缓存 K 和 V 矩阵的原因：Q 矩阵仅需当前 token 的最新行，无需跨步复用。文章还分析了 Prefill 和 Decode 阶段的计算与访存特征，探讨了显存占用及 PagedAttention 等优化策略，确认 KV Cache 是一种无损的性能优化手段。
KV Cache 在 LLM 推理系统中已成为必不可少的优化技术。其核心思想是通过缓存中间结果来避免重复计算，但在实际应用中，关于缓存的具体内容（为何只有 K 和 V）往往令人困惑。
要彻底理解 KV Cache，需要厘清以下三个问题：
如果我们从矩阵乘法的角度观察 Attention 运算的过程，这些问题的答案将变得清晰。
LLM 的推理过程被称为自回归（Autoregressive），即模型上一步的输出会被当作下一步的输入。例如，用户输入的 prompt 是 "recite the first law"，模型产生的第一个 token 是 "A"，然后输出的 "A" 会添加到用户的 prompt 后面，再次输入到模型中，模型输出 "robot" 这个 token。如此反复，直到模型输出结束词。
我们把运行一次模型称为一个 step，每一个 step 生成一个新的 token。这种机制决定了随着生成的进行，输入序列的长度不断增加。
下面，我们从最朴素的推理方法出发，逐步推导出 KV Cache 的优化方法。假设我们的模型有两个 attention 层。
我们先不考虑任何 Cache，观察推理情况。以'你好'作为 prompt 输入到模型中。
模型的 embedding 层会将 token 转化为对应的 embedding，得到一个 2x4 的矩阵 embedding1，其中第一行是'你'的 embedding 表示，第二行是'好'的 embedding 表示。
之后，我们进入到第一个 attention 层。embedding1 会与 attention 层中的三个权重矩阵相乘，分别得到 Q、K 和 V 矩阵。
在 attention layer 中，根据计算公式对 Q、K 和 V 矩阵进行计算。首先，Q 和 K 矩阵相乘，得到一个 attention 矩阵，其中包括了每两个词之间的注意力值。需要注意的是，'你'这个 token 对'好'这个 token 的注意力值被 mask 了，即每一个 token 只计算与之前的 token 的注意力值，而不计算之后的 token 的注意力值。
然后，attention 矩阵与 V 矩阵相乘，再把相乘结果输入到 FFN 层中，就得到了新的 embedding2。第二个 attention 层也是如此类推，并输出 embedding3。
最后，embedding3 的最后一行，也就是'好'这个 token 对应的 embedding 被输入到一个分类器中，预测下一个 token。得到的预测结果是'啊'。到此，我们就生成了第一个 token，即完成了 step0。
在预测了 token'啊'之后，我们将'你好啊'再次输入到模型中。
除了 token 数量从 2 变成了 3 之外，需要执行的运算与 Step 0 完全相同。输入的 embedding 经过两层 attention 层之后，得到 embedding3。embedding3 的最后一行被输入到分类器中，预测得到 token'!'。
一个细节是，Step 1 和 Step 0 的红色部分是相同的。比如说，Step 1 的 K1 矩阵前两行与 Step 0 的 K1 矩阵的前两行是相同的。这是因为 K1 矩阵由 embedding1 矩阵和模型的权重矩阵相乘得到，而两个 Step 的权重矩阵是相同的，唯一不同的地方在于 Step 1 的 embedding1 比 Step 0 多了最后一行，所以两个 Step 的 QKV 矩阵的前两行是相同的。
观察图 1 和图 3 可以发现，模型之后的 Classifier 只会使用 embedding3 矩阵的最后一行作为输入，并输出预测结果。
由此我们可以发现，实际上，每一个 embedding 矩阵只有最新输入的 token 的行是有效的。比如，在图 3 的 embedding3 中，只有 step 0 产生的新 token'啊'对应的第三行是有用的，前两行都是无效的。因此，我们可以删除这些无用的行，以节约计算和存储开销。
因为 embedding2 和 embedding3 矩阵是通过 attention 和 V1 矩阵相乘得到的，所以我们也需要删除 attention 矩阵的重复部分。而 attention 矩阵是由 Q 矩阵和 K 矩阵相乘得到，所以我们还需要删除 Q 矩阵的重复部分。
通过删除无效计算，我们的模型推理过程可以简化为仅计算当前 token 相关的部分。
经过仔细观察我们可以发现，在优化后的图 6 中，Step 1 和 Step 2 的红色部分是完全相同的，只有白色的部分是不相同的。比如说，Step 1 的 Q1 矩阵、K1 矩阵和 V1 矩阵和 Step 0 的 Q1 矩阵、K1 矩阵和 V1 矩阵的前两行是完全相同的。
那么我们其实完全可以在 Step 0 的时候把这些矩阵 Cache 下来，然后在 Step 1 的时候 load 进来，从而节省大量的计算。
在 Step 0 中，我们把红色部分的数据 Cache 下来，然后在 Step 1 的 load 进行计算。在图 7 的 Step 1 中，红色部分的数据都是历史缓存，只有白色部分是新产生的数据。我们将缓存数据和新产生的数据拼接起来，即可得到一个新的矩阵。当然，在 Step 1 中，我们也需要缓存新产生的数据，以供后面的 Step 使用。
在 KV Cache 中，我们把第一个 Step，即 Step 0，称为是 Prefill 阶段，因为这个阶段是在'填充'KV Cache。而之后的所有 Step 称为是 Decode 阶段。
通过观察我们可以发现：
大模型推理系统基本上要解决下面的问题：
在实际工程中，KV Cache 的显存占用是一个关键指标。对于一个 Batch Size 为 B，最大序列长度为 S，隐藏层维度为 H 的模型，每个 token 的 KV Cache 大小约为 2 * S * B * H * sizeof(float)。如果使用 FP16，每个 token 大约占用 2 * S * B * H / 1024 KB。
当序列长度很长时，KV Cache 会迅速消耗显存，导致 OOM（Out Of Memory）。为了解决这个问题，现代推理框架引入了 PagedAttention 等机制，将 KV Cache 分块管理，允许非连续的物理内存分配，从而减少碎片化并提高利用率。
此外，量化技术（如 INT8, INT4）也可以用于压缩 KV Cache 的存储，进一步降低显存需求，但需要在精度损失和推理速度之间取得平衡。
基于上述分析，我们可以回答开始提到的三个问题：
缓存的是什么？ 缓存的是 K 和 V 矩阵。
为什么是 KV Cache，而不是 QKV Cache？ 如图 5 所示，我们只需要 embedding 的最后一行，所以我们也只需要 attention 的最后一行，因此只需要 Q 矩阵的最后一行。由于 Q 矩阵的每一行对应当前的输入 token，它不需要跨步复用，因此不缓存 Q 矩阵的原因是没必要缓存，我们只需要 Q 矩阵的最后一行即可。
KV Cache 对模型的输出有任何的影响吗？ 因为输入到 Classifier 的 embedding 是没有改变的，只是计算路径不同，所以模型的输出也不会有任何变化。这是一种无损优化。
KV Cache 是提升 LLM 推理效率的核心技术之一。理解其背后的原理，不仅有助于优化现有的推理服务，也为未来更高效的架构设计奠定了基础。随着硬件的发展和算法的演进，KV Cache 的管理方式也在不断进化，从简单的线性缓存到分页管理，再到量化压缩，每一步都在推动大模型应用的落地。
微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
生成新的随机RSA私钥和公钥pem证书。 在线工具，RSA密钥对生成器在线工具，online
基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。 在线工具，Mermaid 预览与可视化编辑在线工具，online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具，Base64 字符串编码/解码在线工具，online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具，Base64 文件转换器在线工具，online
将 Markdown（GFM）转为 HTML 片段，浏览器内 marked 解析；与 HTML转Markdown 互为补充。 在线工具，Markdown转HTML在线工具，online
