跳转至

KV Cache Fundamentals

导言

KV Cache 是自回归大模型推理中最重要的运行时状态之一。它把各层历史 token 已经算出的 Key 和 Value 留在缓存里,使后续步骤只处理新 token,而不必反复计算整个前缀。

这不是免费的加速:模型用持续增长的显存驻留换取更少的重复计算。理解这一交换,才能解释为什么生成会加快、长上下文会吃显存、并发容量可能下降,以及为什么修改提示词中间的 token 会使其后的缓存失效。

先记住结论

KV Cache 可以压缩成一句话:每一层都保存历史 token 的 K/V;生成下一 token 时只计算当前 Q/K/V,用当前 Q 查询全部历史 K/V,并把新 K/V 追加到缓存。

它有效的前提是因果注意力:较早位置看不到未来位置,因此未来 token 的出现不会反过来改变已经计算好的历史 K/V。历史 Query 已经完成了当时的查询任务,未来步骤只需要新的 Query,所以通常不存在 Q Cache

它的直接收益与代价分别是:

  • 收益:减少重复计算。 单步 decode 不再重跑整个历史前缀。
  • 代价:增加持久显存。 每个活跃请求、每一层、每个历史位置都要保留 K/V。
  • 约束:缓存必须与前缀完全一致。 提示词、位置、模型权重或相关配置改变后,不能把不一致的旧状态当成新前缀继续使用。

小黑把历史 Key 和 Value 放进缓存转盘

自绘认知示意图:小黑不再把整卷历史反复送回机器,而是保留成对的 K/V 抽屉,每次只加入新 token 的一对卡片;省下重复计算的同时,抽屉也会持续占用空间。

从重复阅读到增量生成

假设模型已经处理了 今天天气很好,准备继续生成。没有缓存时,下一步会再次让每一层处理整个前缀;生成一个新 token 后,又要处理更长的完整前缀。相同历史位置的 K/V 会被一次次重新投影。

有了 KV Cache,第一次处理 今天天气很好 时,各层把这些位置的 K/V 保存下来。下次只让新 token 经过模型,并用它的 Query 与历史 Key 计算匹配程度,再按权重聚合历史 Value。

可以把三者暂时理解成:

  • Query:当前在找什么。 它属于当前查询动作。
  • Key:什么查询可能关注我。 它用于计算匹配分数。
  • Value:关注我以后取走什么。 它被注意力权重加权汇总。

类比不能代替机制

K/V 不是自然语言索引卡或数据库记录,而是每个 Transformer 层产生的高维张量。注意力也不是找到唯一记录,而是对可见历史位置计算权重并聚合 Value。

Q、K、V 到底是什么

在第 l 层,设输入隐藏状态为 \(X^l\in\mathbb{R}^{B\times S_{new}\times D}\)B 是 batch,S_new 是本次新处理的 token 数,D 是隐藏维度。三个投影为:

\[ Q^l=X^lW_Q^l,\qquad K^l=X^lW_K^l,\qquad V^l=X^lW_V^l \]

reshape 并拆成注意力头后,当前张量 shape 是:

Q_new: [B, Hq,  S_new, d]
K_new: [B, Hkv, S_new, d]
V_new: [B, Hkv, S_new, d]

其中 Hq 是 Query Head 数,Hkv 是 KV Head 数,d 是 Head Dim。为了先看清主机制,下面的具体例子采用 Hq = Hkv;共享 KV Head 的架构只改变 head 映射和缓存大小,不改变“新 Query 读取历史 K/V”的核心逻辑。

历史缓存和新张量沿 token 轴拼接:

\[ K_{\le t}^l=\operatorname{concat}(K_{<t}^l,K_t^l,\text{token axis}) \]
\[ V_{\le t}^l=\operatorname{concat}(V_{<t}^l,V_t^l,\text{token axis}) \]

当前输出的语义计算为:

\[ P_t^l=\operatorname{softmax}\left(\frac{Q_t^l(K_{\le t}^l)^\top}{\sqrt d}+M_t\right), \qquad O_t^l=P_t^lV_{\le t}^l \]

M_t 是因果或 padding mask,P_t^l 沿历史 token 轴归一化。矩阵收缩发生在 d 维:Q_t^l [B,Hq,S_new,d] 与转置后的 Key 相乘,得到 [B,Hq,S_new,T];再与 Value 相乘,回到 [B,Hq,S_new,d]

核心对象如下:

对象 Shape 生产者 → 消费者 生命周期
当前层输入 X^l [B,S_new,D] 上一子层 → Q/K/V 投影 临时 activation
当前 Q_t^l [B,Hq,S_new,d] Q 投影与位置编码 → QK MatMul 当前 forward
当前 K_t^l,V_t^l [B,Hkv,S_new,d] K/V 投影与位置处理 → cache update 写入缓存前为临时 activation
历史 K_{<t}^l,V_{<t}^l [B,Hkv,T_past,d] 过去步骤 → 当前与未来注意力 跨 decode step 持久 cache
概率 P_t^l [B,Hq,S_new,T] QK、scale、mask、Softmax → PV MatMul 临时 activation;融合 Kernel 可不完整物化
层输出 Y_t^l [B,S_new,D] PV、Head Merge、输出投影 → 后续子层 临时 activation

KV Cache 的完整机制

三层定义

层次 本文采用的定义
概念 复用因果自注意力中不会被未来 token 改写的历史 K/V。
通用机制 每层保存历史 K/V;本轮计算新 Q/K/V,更新缓存,再让新 Q 读取完整可见历史。
框架实现 固定版本的 Hugging Face Llama 在 RoPE 后调用 past_key_values.updateDynamicLayer 沿倒数第二维追加 K/V。3 2

固定源码中的关键调用可以缩写为以下五行;这里只去掉了 shape 处理和分派代码,投影、RoPE 与 cache update 的先后关系保持不变:

query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)

完整的教学伪代码如下。它既适用于 S_new = prompt length、空 cache 的 prefill,也适用于 S_new = 1、已有 cache 的 decode;没有省略 KV Cache 正确性所依赖的投影、位置编码、拼接、mask、归一化、聚合和输出步骤。

def kv_attention(hidden_new, layer_cache, params, positions, causal_mask):
    batch, new_tokens, model_dim = hidden_new.shape

    q = linear(hidden_new, params.WQ)
    k = linear(hidden_new, params.WK)
    v = linear(hidden_new, params.WV)

    q = reshape_transpose(q, batch, params.Hq, new_tokens, params.head_dim)
    k = reshape_transpose(k, batch, params.Hkv, new_tokens, params.head_dim)
    v = reshape_transpose(v, batch, params.Hkv, new_tokens, params.head_dim)

    q, k = apply_rope(q, k, positions)

    keys_all = concat(layer_cache.keys, k, axis="token")
    values_all = concat(layer_cache.values, v, axis="token")
    layer_cache.keys = keys_all
    layer_cache.values = values_all

    scores = matmul(q, transpose_last_two_axes(keys_all))
    scores = scores * params.attention_scale
    scores = scores + causal_mask
    probabilities = softmax(scores, axis="history_token")
    heads = matmul(probabilities, values_all)

    merged = transpose_reshape(heads, batch, new_tokens, model_dim)
    output = linear(merged, params.WO)
    return output, layer_cache

KV Cache 的物理原理、因果逻辑、执行流程、生命周期与张量流

自绘五联技术图,依据 Hugging Face Transformers `c7cf04b1` 的 `cache_utils.py:113-146` 与 `modeling_llama.py:179-277`。A 看重复前缀如何变成持久缓存,B 看因果成立条件与代价,C 看 prefill/decode 顺序,D 看 cache 跨 token 存活直到请求结束,E 看实际张量、算子、shape 与写回路径。

这张图需要按五个视角阅读:

  1. 物理前后对比: 没有缓存时,历史 token 每步重新进入各层;有缓存时,只把新 token 送入计算,并从缓存读取历史 K/V。
  2. 因果逻辑: causal mask 使历史表示不依赖未来,因此旧 K/V 可以复用;直接效果是减少重算,新代价是缓存驻留与一致性管理。
  3. 执行流程: prefill 先建立初始缓存;decode 每步执行投影、RoPE、追加、Attention、采样,再进入下一步。
  4. 组件与生命周期: 临时 Q 和注意力概率在本轮结束后释放,K/V 从 prefill 开始跨多步存活,请求结束或驱逐时才释放。
  5. 张量数据流: 双圈圆点代表持久 cache,普通圆点代表 activation,方块代表具体算子;任何 shape 变化都由 reshape、transpose 或 concat 明确承担。

适用与归属

KV Cache 适用于因果自回归推理。缓存语义由模型 Attention 层定义,生命周期与内存布局通常由生成框架或推理服务管理,具体 Kernel 决定分数、Softmax 和聚合张量是否物化。

接入一个新模型或 Attention 变体时,至少要适配:K/V 的真实 shape、位置编码施加位置、每层 cache 类型、mask 与 cache position、共享 Head 映射、最大上下文和失效策略。不能假设所有模型都缓存标准的完整 K/V 张量。

证据边界

KV Cache 没有被本文当作一篇独立论文提出的方法。上面的机制由因果注意力语义和固定版本官方实现共同支撑;源码能证明该版本如何更新和消费缓存,不能证明所有框架采用相同物理布局,也不能给出跨硬件通用加速比例。

Prefill 与 Decode

同一个 kv_attention 在两个阶段接收不同的 S_new

阶段 本次输入 Cache 状态 主要工作 输出
Prefill 完整 prompt,S_new = S_prompt 通常为空 并行处理 prompt,建立每层初始 K/V 首次 logits 与初始 cache
Decode 通常一个新 token,S_new = 1 已包含 prompt 和历史输出 只计算新 Q/K/V,读取历史并追加缓存 下一 token logits 与增长后的 cache

Prefill 仍需要处理 prompt 内部的因果注意力,注意力计算量会随 prompt 长度快速增长;KV Cache 主要消除的是后续 decode 对已经处理过的前缀进行重复投影和重复层计算

每层缓存不能互换。第 l 层的 K/V 来自第 l 层输入隐藏状态;即使 token ID 相同,不同层的表示、投影权重和位置处理也不同。

旧 Query 不缓存,是因为位置 j 的 Query 只用于生成位置 j 的注意力输出。到位置 t 时,发起新查询的是 Q_t,但被它反复访问的是历史 K_{\le t}V_{\le t}

速度与显存交换

设当前前缀长度为 T。如果每一步都重新执行完整前缀,单步 dense attention 需要处理约 T × T 的位置关系;缓存后,单个新 Query 只与 T 个历史 Key 比较,单步 attention 的序列项从约 \(O(T^2)\) 变为 \(O(T)\)1

这不等于生成成本变成常数。随着历史增长,新 Query 仍要读取更多 K/V;若连续生成 N 个 token,缓存路径的增量 attention 总量仍包含 \(1+2+\cdots+N\),约为 \(O(N^2)\)

标准逐层 K/V 的理论字节数近似为:

\[ M_{KV}=2\times L\times B\times H_{kv}\times T\times d\times b \]

其中第一个 2 代表 K 和 V,L 是层数,B 是同时缓存的序列数,Hkv 是 KV Head 数,T 是缓存长度,d 是 Head Dim,b 是每个元素的字节数。

例如 L=32B=1Hkv=8T=4096d=128、BF16 的 b=2

M_KV = 2 × 32 × 1 × 8 × 4096 × 128 × 2 bytes
     = 536,870,912 bytes
     = 512 MiB

这是张量元素账,不是服务峰值显存。真实峰值还包括模型权重、临时 activation、Attention workspace、内存分配器保留块、对齐与分页尾部浪费、采样状态和运行时余量:

\[ M_{peak}=M_{model}+M_{KV}+M_{activation}+M_{workspace}+M_{allocator}+M_{runtime} \]

因此“启用 KV Cache 后单请求更快,但并发容量下降”是合理现象:每个活跃请求都保留更多持久状态,显存限制下可同时驻留的请求可能减少。不过这不是普遍的吞吐结论;减少重算也可能让服务在相同时间内完成更多请求,必须在同模型、同长度分布、同精度、同硬件和同调度策略下测量。

缓存何时失效

设位置 j 在第 l 层的隐藏状态满足:

\[ h_j^l=f_l(x_{\le j}) \]

如果把位置 i 的 token 从 x_i 改成 x_i',所有 j\ge i 的条件前缀都可能变化。由于后续 K/V 由这些隐藏状态投影得到:

\[ k_j^l=h_j^lW_K^l,\qquad v_j^l=h_j^lW_V^l \]

所以位置 i 及其后的旧缓存通常必须重算。位置 i 之前的前缀没有变化,仍可复用。这比“逻辑链条被破坏”更精确:被破坏的是缓存张量与新条件前缀之间的数值一致性。

如果修改不是替换而是插入或删除,后续 position id 也会移动;采用 RoPE 等位置编码时,即使后续文本相同,位置变化也会进一步改变 Key。

可以安全复用的通常是完全相同的前缀,并且还要求模型权重、适配器、位置编码规则、精度路径及影响 Attention 的配置一致。语义相似不等于缓存可复用。

常见误解

  • KV Cache 不是原文副本。 它是每一层 Attention 产生的数值状态,不是 token 文本或向量数据库。
  • KV Cache 不是模型长期记忆。 它服务当前请求或可验证的公共前缀,请求结束后通常可以释放。
  • KV Cache 不让长上下文免费。 显存随缓存长度增长,单步 Attention 仍需读取更长历史。
  • 训练中的 K/V activation 不等于推理 cache。 训练通常一次并行处理完整序列,为反向传播保存或重算 activation,而不是跨生成步骤持续 append。官方文档也明确警告缓存应只用于推理。1

实践判断

观察一个 KV Cache 方案时,先回答四个问题:

  1. 缓存什么: 是每层标准 K/V、压缩状态,还是其他 Attention 状态?
  2. 缓存在哪里: GPU、CPU,还是跨设备分片?
  3. 活多久: 一个 decode step、整个请求、共享前缀,还是会被滑窗或驱逐?
  4. 省了什么又付出什么: 减少了哪些重算,新增多少设备内存、带宽、搬运或一致性成本?

如果需要继续比较训练 K/V activation、普通 KV Cache、递归状态和序列并行,可阅读 Attention Cache and Sequence Parallelism

参考资料


  1. Hugging Face Transformers, How caching works, fixed revision c7cf04b1, accessed 2026-08-24. 

  2. Hugging Face Transformers, DynamicLayer.update, fixed revision c7cf04b1

  3. Hugging Face Transformers, LlamaAttention.forward and eager attention, fixed revision c7cf04b1

评论