【Transformer 与注意力机制】49｜KV Cache：推理为什么是 O(n) 不是 O(n²)

训练一个 Decoder-only Transformer 时，给它一整段 token，模型可以并行计算每个位置的 logits。虽然 causal mask 禁止第 \(i\) 个位置看未来，但所有位置的计算可以放在一个大矩阵里完成。推理时情况完全不同：模型必须先生成第一个 token，才能把它放回上下文，再生成第二个 token。后一个 token 依赖前一个 token，自回归生成天然是串行的。

这就带来一个问题：如果每生成一个新 token，都把整个前缀重新跑一遍 Transformer，代价会非常夸张。生成第 1000 个 token 时，你会重新计算前 999 个 token 的所有中间表示；生成第 1001 个 token 时，又重新算一遍前 1000 个 token。大量计算被浪费在已经不变的历史上。

KV Cache 解决的正是这个重复计算。它利用一个事实：在自回归 decode 中，历史 token 的 Key 和 Value 一旦算出，后面不会变。新 token 需要拿自己的 Query 去看历史 K/V，但历史 K/V 本身可以缓存起来。

本篇能让你学会三件事：

1. 为什么训练和推理阶段的 attention 计算形态完全不同；
2. K/V 为什么能缓存，Q 为什么不能用同样方式缓存；
3. KV Cache 如何把重复计算换成显存占用，并成为长上下文推理的核心瓶颈。

一、自回归生成每一步到底在重复什么

假设模型已经看到前缀：

The cat sat on the

现在要预测下一个 token。朴素做法是把这五个 token 全部送进模型，得到最后一个位置的 logits，采样出 mat。下一步上下文变成：

The cat sat on the mat

如果继续朴素做法，就把六个 token 全部重新送进模型。前五个 token 的 embedding、每层 attention 里的 K/V、每层 FFN 输出，都会被重复计算。可是从 causal mask 的角度看，前五个 token 的表示并不会因为后面生成了 mat 而改变。历史不能看未来，这是自回归模型的基本约束。

所以真正需要新增计算的，是新 token 在每一层产生的 Query、Key、Value，以及它对历史 K/V 的 attention。历史 token 的 K/V 可以直接复用。

这就是 KV Cache 的核心：每生成一个 token，就把它在每层产生的 K/V 追加到 cache；下一步 decode 时，新 token 的 Q 会 attend 到所有 cached K/V。

二、K/V 为什么可以缓存，Q 为什么不能同样缓存

Self-attention 中，每一层都会从隐藏状态投影出 Query、Key、Value：

\[ Q=XW_Q,\quad K=XW_K,\quad V=XW_V \]

对历史 token 来说，隐藏状态在 decode 过程中已经确定。由于 causal mask 禁止历史位置看未来，新生成 token 不会反过来改变历史位置的表示。因此历史 token 在每一层的 \(K\) 和 \(V\) 也不会改变。

新 token 不一样。每一步生成时，新 token 是当前最后一个位置，它需要产生自己的 Query 去查询所有历史 Key，并用对应 Value 加权求和。这个 Query 和当前 token 的隐藏状态有关，而当前 token 是刚生成的，不能提前知道。

所以 cache 的对象是历史 K/V，不是未来 Q。可以把它想成检索系统：历史 token 已经被编码成可检索的 key 和可读取的 value；当前 token 产生一个 query，去这些历史槽位里取信息。

每一层都需要自己的 KV Cache。第 10 层的 K/V 不是第 1 层的 K/V，因为它们来自不同层的隐藏状态。多层 Transformer 的 cache 规模也因此随层数线性增长。

三、prefill 和 decode：推理其实分两个阶段

LLM 推理通常分成两个阶段：prefill 和 decode。

prefill 阶段处理用户输入的 prompt。假设 prompt 有 \(n\) 个 token，模型需要一次性计算这 \(n\) 个 token 的表示，并为每层生成对应 K/V cache。这个阶段很像训练时的 forward：可以用大矩阵并行计算，attention 仍然涉及 prompt 内部的三角关系。

decode 阶段每次只生成一个或少量新 token。对每一步，新 token 的 Query 会和已有 cache 中所有 Key 做 attention，然后产生下一个 logits。生成出的 token 再追加进 cache，进入下一步。

这个区分非常重要。很多人说“KV Cache 让推理从 \(O(n^2)\) 变成 \(O(n)\)”，更准确的说法是：在逐 token decode 阶段，每新增一个 token，不再重新计算完整前缀，而是对已有历史做一次增量 attention。 对生成长度 \(T\) 来说，每步仍然要看越来越长的历史，所以总成本仍然随上下文增长；只是避免了“每步重跑全部前缀”的更糟糕做法。

prefill 更容易吃满 GPU，因为它有较大的矩阵计算；decode 更容易受内存读写和串行依赖限制，因为每步 batch 可能小、计算粒度细，还要不断读取 KV Cache。

四、KV Cache 的内存成本

KV Cache 用计算换显存。它避免重复计算历史 K/V，但必须把每层、每个 token、每个 head 的 K/V 保存下来。

粗略公式可以写成：

\[ \mathrm{CacheBytes} = 2 \times L \times B \times S \times H_{kv} \times D_h \times \mathrm{bytes} \]

其中：

• \(2\) 表示 Key 和 Value 两份；
• \(L\) 是层数；
• \(B\) 是 batch size；
• \(S\) 是缓存序列长度；
• \(H_{kv}\) 是 KV head 数；
• \(D_h\) 是每个 head 的维度；
• \(\mathrm{bytes}\) 是每个元素占用字节数，比如 FP16/BF16 通常是 2。

这个公式说明了为什么长上下文推理很快会被 cache 卡住。序列长度翻倍，cache 显存翻倍；batch size 翻倍，cache 也翻倍；层数越深、hidden size 越大，cache 越贵。

参数量是固定成本，KV Cache 是随请求动态增长的成本。一个服务同时处理很多长上下文请求时，即使模型参数已经量化，cache 也可能成为主要显存压力。

五、为什么 KV Cache 让训练和推理像两种程序

训练时，目标是高吞吐处理大量 token。通常会把一个 batch 的完整序列送进模型，用并行矩阵运算一次算完所有位置的 logits。虽然 causal mask 存在，但 GPU 看到的是大块矩阵计算。

推理时，尤其是 decode 阶段，目标是低延迟地产生下一个 token。每一步都要读 cache、算当前 token、写新 cache，然后等待采样结果进入下一步。这种循环结构让推理更像在线系统，而不是纯训练算子。

KV Cache 也改变了 batching。prefill 请求长度不同，decode 生成速度不同，有的用户已经生成完，有的用户还在继续。服务系统必须动态把不同请求拼在一起，避免 GPU 空转。这就是 continuous batching、PagedAttention 等技术出现的背景。

所以不要把 KV Cache 看成一个简单数组。它是推理系统的核心状态：决定能同时服务多少请求，决定长上下文能开多大，决定 batch 如何调度，也决定显存碎片如何管理。

六、长上下文为什么会把 KV Cache 推到台前

长上下文模型带来的第一感觉是“窗口变大了”。但从推理系统看，窗口变大意味着每个请求的 KV Cache 更长。一个 4K 上下文的请求和一个 128K 上下文的请求，对 cache 的压力完全不是一个量级。

更麻烦的是，多轮对话会不断累积上下文。用户每问一轮，历史消息都要保留在 prompt 中，prefill 或 cache 复用策略必须处理越来越长的前缀。Agent 场景更复杂：工具调用、网页内容、代码文件、日志和中间推理都可能塞进上下文。

KV Cache 让 decode 不必重算历史，但它没有让历史免费。历史被换成了显存中的 K/V。长上下文真正的代价，经常不是“模型能不能看见这么多 token”，而是“服务系统能不能为这么多并发请求保存这些 token 的 cache”。

这也是为什么长上下文优化会同时出现很多方向：更高效的 attention kernel、GQA/MQA 减少 KV head、cache 量化降低 bytes、滑动窗口限制历史、检索增强把部分信息放到外部系统、以及分页式 cache 管理减少碎片。

七、PagedAttention 与 cache 管理

PagedAttention 的直觉来自操作系统虚拟内存。传统 cache 管理如果为每个请求分配连续大块显存，就容易产生碎片：有些请求短，有些请求长，有些提前结束，有些继续生成。显存中会出现很多不好复用的小洞。

PagedAttention 把 KV Cache 切成固定大小的 block，用类似分页的方式管理。请求看到的是连续的逻辑序列，底层物理 cache block 可以分散存放。这样可以更灵活地复用显存，减少碎片，并支持动态 batching。

vLLM 的 PagedAttention 工作让这个问题变得非常有名。它提醒大家：LLM 推理优化不只是 kernel 快不快，还包括内存管理、调度和请求生命周期。KV Cache 是状态，状态管理不好，吞吐和并发都会受影响。

这和 FlashAttention 的视角不同但互补。FlashAttention 关注 attention 计算内部如何减少 I/O；PagedAttention 关注推理服务中 KV Cache 如何高效分配和复用。

八、KV Cache 的变体与压缩方向

Multi-Query Attention（MQA）和 Grouped-Query Attention（GQA）是减少 KV Cache 的重要结构。标准 multi-head attention 每个 query head 都有自己的 K/V head；MQA 让多个 query head 共享同一组 K/V；GQA 则在二者之间折中，让一组 query heads 共享一个 K/V head。

从 cache 公式看，减少 \(H_{kv}\) 会线性降低 cache 大小。这就是为什么 GQA 在很多现代大模型中很受欢迎：它在质量和推理效率之间取得折中。

cache quantization 则从 bytes 入手。把 KV Cache 从 FP16/BF16 压到更低精度，可以降低显存和带宽压力，但会引入精度损失，需要小心评估。和权重量化一样，它不是免费午餐。

sliding window、cache eviction 和 attention sink 等方向试图限制模型实际关注的历史范围。它们适合某些长序列场景，但会改变模型可访问的信息。什么时候可以丢历史，丢哪些历史，丢了会不会伤害任务，是需要结合模型和应用判断的问题。

九、关键概念回顾

• KV Cache：自回归推理中缓存每层历史 token 的 Key 和 Value，避免重复计算前缀。
• prefill：处理 prompt 的阶段，生成初始 cache，计算形态更接近训练 forward。
• decode：逐步生成新 token 的阶段，每步读取 cache 并追加新 K/V。
• MQA / GQA：减少 KV head 数，从结构上降低 cache 大小和读写成本。
• PagedAttention：以分页方式管理 KV Cache，减少碎片并支持高效服务调度。
• cache quantization：用更低精度保存 KV Cache，降低显存和带宽压力。

十、常见误解

10.1 “KV Cache 让 attention 不再依赖上下文长度”

不对。decode 每一步仍然要让当前 token attend 到历史 K/V。KV Cache 避免的是反复重算历史 K/V，不是让历史长度免费。

10.2 “Q 也应该缓存起来”

历史位置的 Q 对未来 decode 通常没有复用价值。新 token 需要的是自己的 Q 去查询历史 K/V；历史 K/V 才是可复用状态。

10.3 “训练也可以直接用 KV Cache 省掉 O(n²)”

训练要并行计算所有位置的 loss，和逐 token decode 的数据流不同。KV Cache 主要服务自回归推理中的增量生成，不是训练 attention 的通用替代。

10.4 “cache 越大越好”

更大的 cache 能保留更多上下文，但会占用显存、降低并发、增加带宽压力。长上下文系统必须在能力、成本和延迟之间权衡。

十一、下一步

KV Cache 解释了为什么推理工程不能只看模型参数量。对 Decoder-only 模型来说，历史上下文本身会变成巨大的运行时状态。再往后看，Transformer 的根本局限就更清楚了：attention 的二次关系、decode 的串行性、KV Cache 的线性增长，共同决定了长上下文和低延迟生成的边界。

十二、参考文献

1. Vaswani, A. et al. “Attention Is All You Need.” NeurIPS 2017. Transformer self-attention 与自回归 decoder 的来源。
2. Shazeer, N. “Fast Transformer Decoding: One Write-Head is All You Need.” arXiv:1911.02150, 2019. Multi-Query Attention 的代表工作。
3. Ainslie, J. et al. “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.” EMNLP 2023. Grouped-Query Attention 的代表工作。
4. Kwon, W. et al. “Efficient Memory Management for Large Language Model Serving with PagedAttention.” SOSP 2023. vLLM / PagedAttention 论文。
5. Dao, T. et al. “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” NeurIPS 2022. prefill attention kernel 优化的重要背景。

← 上一篇：48｜从 logits 到文本　|　下一篇：50｜Speculative Decoding →