KV Cache 机制详解：它到底是什么，为什么会产生，存储结构长什么样

好，咱们今天来聊聊 KV Cache。说实话，这个机制是 vLLM 长序列推理优化的基石。你想想看，如果没有 KV Cache，那 Transformer 做自回归生成时，每生成一个 token 都得把前面所有的 token 重新算一遍——这效率低得简直没法看。

我个人习惯把 KV Cache 比作「推理过程中的草稿纸」。你解一道复杂的数学题，中间步骤算出来的结果，总得找个地方记下来吧？不然下一步又要从头算起，那不得累死。KV Cache 干的就是这个活。

KV Cache 是什么？

简单来说，KV Cache 就是 Transformer 模型在自回归解码过程中，缓存下来的 Key 矩阵和 Value 矩阵。

在 Transformer 的每一层自注意力（Self-Attention）里，输入序列会通过三个线性变换，生成 Query、Key、Value 三个矩阵。对于当前生成的 token，它的 Query 要和前面所有 token 的 Key 做点积，算出注意力分数，再和 Value 加权求和。

嗯，这里要注意：前面 token 的 Key 和 Value 一旦算出来，就不会再变了。所以我们可以把它们存起来，下次直接用，不用重新算。

核心定义：KV Cache 是 Transformer 每一层、每个注意力头中，已生成 token 对应的 Key 矩阵和 Value 矩阵的缓存副本。

为什么会产生 KV Cache？

这个问题我当年刚接触时也困惑过。说白了，原因就一个：避免重复计算。

自回归生成的过程是这样的：

  1. 输入 prompt，生成第一个 token
  2. 把第一个 token 拼到 prompt 后面，生成第二个 token
  3. 把前两个 token 拼到 prompt 后面，生成第三个 token
  4. ……

你发现没有？每次生成新 token 时，前面的序列其实没变。如果不做缓存，每次都要把整个序列的 Key 和 Value 重新算一遍。假设序列长度是 L，那复杂度就是 O(L²)，这谁受得了？

我记得有一次在项目中调试一个 32K 长度的推理任务，没开 KV Cache，结果一个 token 要等好几秒。我当时就意识到：没有 KV Cache 的长序列推理，基本就是不可用的。

有了 KV Cache 之后，每次推理只需要：

  • 计算当前 token 的 Key 和 Value
  • 把新的 Key/Value 追加到缓存里
  • 用当前 token 的 Query 和缓存里的所有 Key 做注意力计算

复杂度从 O(L²) 降到了 O(L)，这才是工程上可行的方案。

避坑指南：我曾经犯过一个错误——在 batch 推理时，不同序列的长度不一样，我直接把所有序列的 KV Cache 拼在一起，结果注意力计算全乱了。后来才意识到，每个序列的 KV Cache 必须独立管理，不能混用。

存储结构分析

好，咱们来拆解一下 KV Cache 到底长什么样。我习惯用三维的视角来看它。

维度拆解

KV Cache 的存储结构，本质上是一个四维张量：

维度 含义 示例值
Batch Size (B) 同时推理的序列数量 4
Num Heads (H) 每层的注意力头数 32
Sequence Length (L) 当前已生成的 token 总数 4096
Head Dim (D) 每个注意力头的维度 128

所以，对于一层 Transformer，KV Cache 的大小就是：

KV Cache 大小 = B × H × L × D × 2  (Key 和 Value 各一份)

举个例子：B=4, H=32, L=4096, D=128，那单层 KV Cache 就是：

4 × 32 × 4096 × 128 × 2 = 134,217,728 个元素
每个元素用 FP16（2字节）存储，就是 256 MB

如果模型有 32 层，那总共就是 256 MB × 32 = 8 GB。你看，长序列推理时，KV Cache 占的内存比模型权重还大，这是常态。

注意：上面算的是单 batch 的情况。如果 batch size 是 64，那 KV Cache 直接飙到 512 GB。这就是为什么 vLLM 要用 PagedAttention 来做内存管理——不优化的话，显存根本扛不住。

层间结构

每一层 Transformer 都有自己的 KV Cache。它们之间是独立的，不能共享。我画了一张图，帮你理解这个结构：

KV Cache 存储结构示意图 输入序列: [token_1, token_2, ..., token_L] Layer 1 KV Cache: [B, H, L, D] × 2 Layer 2 KV Cache: [B, H, L, D] × 2 Layer N KV Cache: [B, H, L, D] × 2

你看，每一层都独立维护自己的 KV Cache。层与层之间不共享，因为每层的 Key/Value 是通过不同的权重矩阵算出来的，语义完全不同。

内存布局细节

在实际的 GPU 显存里，KV Cache 是怎么排布的？我以 vLLM 的实现为例，给你拆解一下：

  • 连续存储 vs 分页存储：传统方案是连续存储，但 vLLM 用 PagedAttention，把 KV Cache 切成固定大小的块（page），按需分配。这有点像操作系统的虚拟内存。
  • 内存对齐：每个 page 的大小通常是 16 或 32 个 token 的 KV 数据。这样对齐后，GPU 的访存效率更高。
  • 显存碎片：连续存储容易产生碎片，尤其是序列长度变化时。分页存储能有效缓解这个问题。

关键点：KV Cache 的存储结构直接影响推理的吞吐量和延迟。vLLM 之所以能在大 batch、长序列场景下表现优异，核心就在于它用分页方式管理 KV Cache，避免了显存浪费和碎片化。

实际项目中的经验

我在做长序列推理优化时，遇到过几个和 KV Cache 相关的坑，分享给你：

  • 显存爆炸：有一次 batch size 设得太大，KV Cache 直接撑爆了 80G 的 A100。后来我加了个动态 batch 调度，根据当前序列长度动态调整 batch size。
  • 缓存失效：在 beam search 场景下，每条 beam 的 KV Cache 不能共享。我一开始没注意，结果不同 beam 的注意力计算互相污染，生成结果全乱了。
  • 精度问题：KV Cache 用 FP16 存储时，长序列下误差会累积。我后来在关键层改用 FP32 缓存，虽然多占点显存，但精度保住了。

好了，KV Cache 的机制就讲到这里。你理解了它的本质、产生原因和存储结构，后面再学 PagedAttention、内存复用这些高级技巧，就会轻松很多。

专注资料整理

← 上一章 📖 返回目录 下一章 →