浅析大模型推理十二篇之KV Cache

0 阅读7分钟

想象一个用户打开一个网站,这个网站有 logo、CSS 文件、JavaScript 文件、字体和图片。
如果没有缓存,用户每次访问新页面时,浏览器都要从服务器重新下载同一个 logo、同一份 CSS、同一份 JavaScript 和同一份字体,这将是极大的浪费。

Logo 没有变,CSS 文件没有变,JavaScript 文件没有变,为什么要重新下载?

浏览器使用缓存来解决这个问题,同样的原理也应用在这里:我们将已计算的 K 和 V 向量存储在缓存中。
在 Prefill 阶段,当模型努力计算第一个 token 时,由于它已经可以访问整个输入 prompt,它会将 prompt token 的 K 和 V 向量缓存起来。

KV cache 就像浏览器缓存,但缓存的不是图片、CSS 和 JavaScript 文件,而是先前 token 的 Key 和 Value 向量。

这避免了重复计算,使逐 token 生成变得快得多。

kvcache_1.png

示例

用同一个 prompt 来理解:

The capital of France is

在第一次前向传播中,模型为所有 5 个 token 计算 K 和 V 向量:

The      → K1, V1
capital  → K2, V2
of       → K3, V3
France   → K4, V4
is       → K5, V5

模型将这些 K 和 V 向量存储在内存中。

这些存储的内存就称为 KV cache

现在模型预测第一个输出 token:

Paris

由于有 KV cache:

The      → 使用缓存的 K1, V1
capital  → 使用缓存的 K2, V2
of       → 使用缓存的 K3, V3
France   → 使用缓存的 K4, V4
is       → 使用缓存的 K5, V5
Paris    → 计算新的 K6, V6

模型不再为所有 6 个 token 重新计算 K 和 V,只需为新 token 计算 K 和 V。

然后模型预测下一个 token:

.

context 变为:

The capital of France is Paris.

只有新 token . 需要新的 K 和 V 计算:

The      → 使用缓存的 K1, V1
capital  → 使用缓存的 K2, V2
of       → 使用缓存的 K3, V3
France   → 使用缓存的 K4, V4
is       → 使用缓存的 K5, V5
Paris    → 使用缓存的 K6, V6
.        → 计算新的 K7, V7

对比之前的例子:没有 KV cache 时,每输出一个 token 都需要重新计算所有 token 的 K 和 V 向量;而有了 KV cache,我们只需为输出 token 计算 K 和 V 向量。

Prefill 与 Decode 的计算特性

在 Prefill 阶段,我们计算 K 和 V 向量并将它们存储到 KV cache 中,因此这个操作是 compute-bound(计算受限)的。

在 Decode 阶段,我们从 KV cache 中读取值的频率远高于为 token 计算 K 和 V 向量的频率。

假设输出为 6 个 token,那么其中 5 个 token 我们从 KV cache 中读取 K 和 V 向量,只为第 6 个 token 计算向量。因此这个操作是 memory-bound(内存受限)的。

为什么叫 KV Cache?—— Q 去哪了

每个刚接触 KV Cache 的人都会问:既然 K、V 能缓存,Q 呢?

类比:Q 是「当前 token 提出的问题」,K/V 是「历史 token 的目录卡片 + 内容摘要」。第二个问题不需要「记得第一个问题」,只需要 同一套卡片和书架

角色含义生命周期
Q (Query)“我该关注谁?”一次性:用完即弃
K (Key)“我是什么索引?”持久:后续每个 token 都要查
V (Value)“关注我你能拿到什么?”持久:后续每个 token 都要读

Decode 第 t 步:

Q_t 需要查询: K_1..K_t      ← 要所有历史 Key
Q_t 需要读取: V_1..V_t      ← 要所有历史 Value
Q_t 不需要:   Q_1..Q_{t-1}  ← 过去的 Query 对新 token 无意义

因果掩码决定了 Q 的一次性:Q_{t+1} 查询范围是 [1, t+1],永远不会再用 Q_t 去查别人,所以:缓存一种不会被再次访问的东西,纯属浪费显存。

量级直觉(LLaMA-2 70B,GQA,seq=4096):

  • K + V ≈ 1.25 GB
  • 若硬存 Q(64 个 Q head)≈ +5 GB,且这 5 GB 永远不会被读

Prefill / Decode:KV Cache 如何降低计算量

1. 计算量下降

Prefill:输入整段 prompt,并行算所有 token 的 K/V,写入 cache,产出第一个输出 token。
Decode:每步只输入 上一步的新 token,算它的 Q/K/V;把新 K/V append 进 cache;用「历史 K/V + 当前 K/V」算 attention,再吐下一个 token。

没有 KV Cache:单步 attention 约 O(N²);生成 L_gen 个 token 的累计代价近似 O(L_gen³)
有 KV Cache:Decode 单步降至 O(N);总代价近似 O(P² + P·L_gen + L_gen²)P 为 prompt 长度, 来自 Prefill)。

简化伪代码:

class KVCache:
    def __init__(self):
        self.cache = {"key"None"value"None}

    def update(self, key, value):
        if self.cache["key"is None:
            self.cache["key"] = key
            self.cache["value"] = value
        else:
            # 在 seq 维拼接;若布局是 [B, heads, seq, dim],改 dim=2
            self.cache["key"] = torch.cat([self.cache["key"], key], dim=1)
            self.cache["value"] = torch.cat([self.cache["value"], value], dim=1)

2. 显存占用下降

Memory_KV ≈ 2 × b_kv × L × B × S × H × (N_kv / N_attn)
  • 2:K 和 V 各一份
  • b_kv:字节数(FP16=2)
  • L:层数;B:并发请求数;S:平均序列长度
  • H:hidden size;N_kv / N_attn:KV head 数 / Q head 数(GQA 时 <1)

关键S × B 是乘积关系。长上下文 × 高并发会把 KV Cache 推到 超过模型权重

Qwen2.5-7B(FP16)为例,每 token 约 56 KB

batchseq_lenKV Cache对比权重 (~14 GB)
12,0480.11 GB<1%
832,76814.3 GB≈ 权重
3232,76857 GB≈ 4× 权重

3. PagedAttention:如何管好 KV 显存碎片

问题:为每个请求按 max_model_len 连续预分配 KV →

  • 预留浪费:分了 32K,实际只生成了 3K
  • 外部碎片:短请求释放的洞拼不出下一块大连续区

传统方案有效 KV 利用率常只有 20–40%

方案(对应操作系统虚拟内存):

OSPagedAttention
物理页框KV block(固定 block_size 个 token)
页表Block Table(logical → physical)
MMU 硬件翻译attention kernel 软件查表 gather
按需调页block 按需分配、用完回收

结果:浪费压到 <4%(主要是每个请求最后一个 block 未填满),同等显存吞吐常提升 2–4×

顺带解锁

  • Prefix Caching:多个请求的 block table 可指向同一物理 block(相同前缀共享)
  • Offloading:按 block 粒度换入换出,比搬动整条连续大 tensor 灵活
  • 量化:block 内部做 FP8/INT8,kernel 侧反量化

Prompt Caching:跨请求的「前缀 KV 复用」

前面讲的 KV Cache 解决的是 同一次生成里 不重算历史,真实产品里还有一层更「省钱」的缓存:多次 API 调用共享相同前缀

1. Prompt Caching 是什么

长且稳定 的前缀(system prompt、工具定义、知识库、前几轮对话)算出的 KV(或等价中间状态)留在服务端;下一次请求若前缀 逐字节一致,就直接复用,只对 变化的后缀 做 Prefill。

2. OpenAI 和 Anthropic 提供的 Prompt Caching 方式

OpenAI

  • ≥1024 token 的公共前缀自动参与缓存,按 128 token 粒度递增命中
  • 命中部分显著降 TTFT 与输入费用(官方示例:TTFT 最多降 ~80%,输入成本最多降 ~90%)
  • 关键工程纪律:静态内容放前面,时间戳等易变内容放 metadata;工具/示例顺序也要稳定

Claude

  • cache_control 显式标记缓存断点(前缀字节级匹配)
  • 顺序固定:tools → system → messages;前面一改,后面全失效
  • 成本:写缓存约为 1.25× 输入价,读缓存约为 0.1×;官方长 prompt 场景成本可降 ~90%、TTFT 降 ~85%

自托管(vLLM)

  • 在 PagedAttention 之上做 Automatic Prefix Caching:相同前缀的 block 直接复用,天然省 Prefill。

3. KV Cache ≠ Prompt Caching(对比收束)

KV CachePrompt Caching / Prefix Cache
作用域单次请求内:第 1→N 个输出 token跨请求:多次调用共享相同前缀
存什么各层 Attention 的 K、V前缀对应的 KV / 中间状态(按前缀哈希复用)
解决什么Decode 不重算历史,O(N²) → O(N)重复前缀的 Prefill 只付一次,TTFT/成本下降
用户是否可控框架自动(如 use_cache=TrueOpenAI 自动+prompt_cache_key;Claude 需 cache_control
典型瓶颈显存容量 + HBM 带宽前缀稳定性、路由到同一缓存、TTL
关系是基础设施是 KV Cache 思想在「服务层」的延伸

4. 总结

KV Cache 让「这一次生成」别重算;Prompt Caching 让「下一次请求」别重复 Prefill 相同前缀。两者叠加,才构成现代 LLM 服务「快且便宜」的完整缓存栈。


参考

  1. pub.towardsai.net/llm-inferen…
  2. github.com/ForceInject…
  3. github.com/ForceInject…
  4. github.com/ForceInject…
  5. github.com/ForceInject…
  6. huggingface.co/blog/not-la…
  7. huggingface.co/blog/kv-cac…
  8. huggingface.co/blog/kv-cac…
  9. openai.com/index/api-p…
  10. developers.openai.com/api/docs/gu…
  11. developers.openai.com/cookbook/ex…
  12. claude.com/blog/prompt…
  13. claude.com/blog/lesson…
  14. vllm.ai/blog/2023-0…
  15. arxiv.org/abs/2309.06…