3、KV Cache 核心原理：解码阶段的重复计算问题、空间换时间思想、缓存命中与复用

3.1 解码阶段的重复计算问题

我们先来看一个最直观的问题。Transformer 做推理时，分为预填充（Prefill）和解码（Decode）两个阶段。预填充阶段一次性把整个 prompt 扔进去，算出所有 token 的中间表示。但到了解码阶段，每次只生成一个 token，然后把这个 token 拼到序列末尾，再重新跑一遍整个序列。

你想想看，这有多浪费？

假设序列长度是 1000，生成第 1001 个 token 时，模型需要重新计算前 1000 个 token 的注意力分数。而前 1000 个 token 的 Key 和 Value，其实在上一步已经算过了。每次生成新 token，前面那些 token 的 K、V 矩阵完全没变，但模型却要重新算一遍。

我刚开始做推理优化时，第一次看到这个现象，心里直呼「这也太笨了吧」。实际上，这就是解码阶段的核心痛点——重复计算。每次自回归生成，都要把历史 token 的注意力全部重算一遍。序列越长，浪费越严重。

核心问题： 解码阶段每步生成一个 token，都需要重新计算所有历史 token 的 Key 和 Value 矩阵，导致计算量随序列长度线性增长，且大量计算是冗余的。

举个例子。假设模型有 12 层，每层 hidden size 是 4096。生成第 100 个 token 时，需要计算 100 个 token 的注意力。生成第 101 个 token 时，又要算 101 个 token 的注意力。前 100 个 token 的 K、V 明明没变，却要重复算 100 次。这 100 次计算，说白了就是白费功夫。

3.2 空间换时间思想

怎么解决？答案很简单——把算过的结果存下来。

这就是 KV Cache 的核心思想：空间换时间。用额外的显存空间，存储历史 token 的 Key 和 Value 矩阵，避免重复计算。

具体做法是这样的：

  • 在预填充阶段，正常计算所有 token 的 K、V，并缓存起来
  • 进入解码阶段后，每生成一个新 token，只计算当前 token 的 K、V
  • 从缓存中取出历史 token 的 K、V，与当前 token 的 Q 做注意力计算
  • 把当前 token 的 K、V 追加到缓存中，供后续步骤使用

我习惯把这个过程想象成「搭积木」。每生成一个 token，就往缓存里多放一块积木。下次生成时，直接从缓存里拿积木用，不用重新搭一遍。

我的经验： 在实际项目中，KV Cache 的显存占用往往成为瓶颈。比如一个 7B 模型，序列长度 2048，batch size 1，KV Cache 大约需要 2GB 显存。如果 batch size 增大到 32，显存占用直接飙到 64GB。所以 KV Cache 不是免费的午餐，它用显存换速度。

从计算角度看，KV Cache 把解码阶段的时间复杂度从 O(n²) 降到了 O(n)。为什么？因为每次只计算当前 token 的注意力，历史 token 的注意力直接从缓存中读取。说白了，就是把重复计算变成了内存访问。

这里有个关键点：KV Cache 只缓存 Key 和 Value，不缓存 Query。因为 Query 每次都在变（当前 token 不同），而 Key 和 Value 是历史 token 的，不会变。

3.3 缓存命中与复用

KV Cache 的命中率，理论上永远是 100%。因为解码阶段是确定性的，每次生成的 token 序列是固定的，缓存中的 K、V 一定被用到。

但实际工程中，情况要复杂得多。

我曾经踩过的坑： 有一次做多轮对话推理，用户连续问了 10 个问题。我天真地以为每次对话都重新构建 KV Cache 就行。结果发现，前 9 轮对话的历史信息全部丢失，模型每次都要重新计算。后来我改成「增量缓存」策略——每轮对话只追加新 token 的 K、V，历史缓存保留。这才把推理速度提上来。

缓存复用有几个典型场景：

  1. 连续生成场景：同一个 prompt 生成多个候选回答时，可以共享预填充阶段的 KV Cache
  2. 多轮对话场景：历史对话的 KV Cache 可以跨轮次复用，只需追加新 token
  3. 前缀共享场景：多个请求共享相同前缀（如系统提示词），可以复用前缀部分的 KV Cache

我建议在实现 KV Cache 时，注意以下几点：

  • 缓存结构要支持动态扩展，因为序列长度会逐渐增加
  • 考虑使用分页缓存（PagedAttention）来减少显存碎片
  • 对于长序列，可以设置缓存淘汰策略（如 LRU）来限制显存占用

下面我用一个简单的示意图，展示 KV Cache 的工作流程：

KV Cache 工作流程示意图 输入序列 预填充阶段 计算所有 token 的 K、V KV Cache 存储历史 K、V 解码阶段 只计算当前 token 的 K、V 读取历史 K、V 追加当前 K、V 生成当前 token 循环直到生成结束

从上图可以看出，KV Cache 的核心就是「存」和「取」。预填充阶段把 K、V 存起来，解码阶段从缓存中取历史 K、V，同时把当前 token 的 K、V 存进去。如此循环，直到生成结束。

嗯，这里要注意一个细节：缓存的是每层 Transformer 的 K、V，不是最后一层的输出。因为注意力计算是在每一层独立进行的，所以每层都需要自己的 KV Cache。

我见过不少新手把 KV Cache 和模型输出混为一谈，结果缓存了错误的数据。记住：KV Cache 是中间结果，不是最终输出。

最后，我总结一下 KV Cache 的核心要点：

维度 说明
解决的问题 解码阶段重复计算历史 token 的 K、V 矩阵
核心思想 空间换时间，用显存存储中间结果
缓存内容 每层 Transformer 的 Key 和 Value 矩阵
缓存时机 预填充阶段全部缓存，解码阶段增量追加
复用场景 连续生成、多轮对话、前缀共享
工程挑战 显存占用大，需要分页管理或淘汰策略
我的建议： 在实际部署时，先估算 KV Cache 的显存需求。公式很简单：显存 = 层数 × 每层 hidden size × 序列长度 × 2（K 和 V） × 每个元素的字节数。比如 32 层、4096 hidden size、2048 序列长度、FP16，就是 32 × 4096 × 2048 × 2 × 2 = 1GB 左右。这个估算能帮你提前规划显存分配。

KV Cache 是 Transformer 推理优化的基石。理解了它，你就能明白为什么推理框架都在拼命优化显存管理，为什么 PagedAttention 能成为主流方案。说白了，所有优化都是在「存」和「取」这两个动作上做文章。

专注资料整理
← 上一章 📖 返回目录 下一章 →