跳转至

第八章 KV Cache与PagedAttention

导读

LLM 自回归生成时，每个新 token 都要 attend 到之前所有 token。如果每次都重算历史 K/V，算力浪费 O(L²)。KV Cache 是把历史 K/V 缓存复用的优化；PagedAttention 是 vLLM 提出的分页式 KV Cache 管理方案，是现代推理引擎的核心。本章解释二者。

8.1 没有 KV Cache 时的问题

考虑生成 "今天天气不错"：

  • 生成 "今" → 输入只有 <bos>
  • 生成 "天" → 输入 <bos> 今
  • 生成 "天" → 输入 <bos> 今天
  • ……

每生成一个新 token，所有历史 token 的 K、V 都要重算一次。复杂度 O(L² · d)。对 L=4096，前向算 4096 次重复计算——巨大浪费。

8.2 KV Cache 原理

KV Cache对比

关键洞察：自回归生成时，历史 token 的 K、V 不变——只需要算新 token 的 Q、K、V，把新 K、V append 到 cache，做一次 attention 即可。

# 伪代码：自回归生成
kv_cache = []   # 初始为空

for step in range(max_new_tokens):
    Q_new, K_new, V_new = model.qkv_projection(new_token)

    kv_cache.append((K_new, V_new))                # 只 append，不重算
    K_all = torch.cat([k for k, v in kv_cache])    # 拼接
    V_all = torch.cat([v for k, v in kv_cache])

    output = attention(Q_new, K_all, V_all)        # 只算新 Q 与全部 K/V
    new_token = sample(output)

每步复杂度从 O(L² · d) 降到 O(L · d)。

8.3 KV Cache 的显存代价

KV Cache 不是免费的——它要存历史所有层的 K、V：

\[\text{KV Cache 大小} = 2 \times L_{\text{layers}} \times L_{\text{seq}} \times N_{\text{kv\_heads}} \times d_{\text{head}} \times \text{dtype\_bytes} \times \text{batch}\]

例：Llama-3 70B（80 层，8 KV heads，d=128，fp16，batch=32，seq=8192）：

\[2 \times 80 \times 8192 \times 8 \times 128 \times 2 \times 32 \approx 86 \text{ GB}\]

仅 KV Cache 就 86GB——比模型权重（70B fp16 = 140GB）的一半还多！长上下文 + 大 batch 时 KV Cache 是显存主瓶颈。

这就是为什么现代模型从 MHA 改成 GQA / MQA——KV heads 减少 4-8 倍，KV Cache 同比例下降。

8.4 减小 KV Cache 的方法

方法 压缩比 代价
GQA / MQA 4-8× 质量略降
KV Cache 量化（FP8 / INT4） 2-4× 质量略降
MLA（Multi-head Latent Attention） ~4× DeepSeek 提出，需要重训
Sliding window 取决窗口 长依赖丢失
Token eviction 视策略 召回策略敏感
Prompt Cache 共享前缀 实现复杂

MLA 是 DeepSeek-V2/V3 提出的方案：把 K/V 投影到低维 latent 空间存储，需要时再恢复——4 倍压缩且质量几乎无损，但需要模型架构层面改动，不能直接套到已有模型。

8.5 PagedAttention：分页管理

传统 KV Cache 实现的问题：

PagedAttention

传统连续分配：每个请求预留 max_len 的连续显存，造成：

  1. 碎片：实际生成短，预留长；
  2. 限 batch：max_len × batch 受显存约束，利用率低；
  3. 无法共享前缀：相同 system prompt 各自存一份。

PagedAttention（vLLM 核心创新）借鉴 OS 虚拟内存分页：

  • KV Cache 切成固定大小 block（如 16 tokens / block）；
  • block table 记录每个 seq → 物理 block id 列表（类似页表）；
  • 按需分配 block，无需预留 max_len；
  • 显存利用率 90%+（vs 传统 30-50%）；
  • 共享前缀：相同 system prompt 共享 block，写时复制（CoW）；
  • 配合 Continuous Batching：长短请求混跑。

8.6 Continuous Batching

传统静态批处理必须等所有请求完成才释放 batch 槽位——短请求 GPU 空转，长请求拖累。

Continuous Batching（vLLM、TGI、TRT-LLM 都支持）：

  • 每 step 重组 batch；
  • 完成的请求立即出，新请求立即入；
  • 配合 PagedAttention，KV Cache 块可重排；
  • GPU 利用率 95%+。

关键约束：必须配合 PagedAttention 才高效——传统连续 KV Cache 无法在 batch 变化时高效重排。

8.7 Prefill vs Decode 阶段

LLM 推理两个阶段：

阶段 计算特性 优化重点
Prefill 处理 prompt，O(L²) 算力 算力密集，提升 MFU
Decode 逐 token 生成，O(L) 算力但单 token 显存带宽密集

Disaggregated Serving：把 prefill 与 decode 分到不同 GPU 池——prefill 用算力强的卡，decode 用带宽优的卡。提高整体吞吐量，是当前服务系统优化方向之一（如 DeepSeek、Moonshot 的部署架构）。

8.8 Prefix Cache / Prompt Cache

实际业务中常见**固定 system prompt + 变化 user query** 模式：

[system prompt: 2000 tokens 固定]
[user query: 50 tokens 变化]

每请求都重算 2000 token 的 prefill 是巨大浪费。

Prefix Cache：

  • 把常见 system prompt 的 KV Cache 缓存（LRU）；
  • 新请求若前缀命中，直接复用 cache，只 prefill 后半段；
  • vLLM 的 enable_prefix_caching=True 一键开启；
  • 业务收益：相同 system prompt 的请求延迟降 5-20×。

8.9 Speculative Decoding

另一个相关优化——投机解码：

  1. 用小模型（draft model）一次生成 K 个候选 token；
  2. 大模型（target model）一次性验证这 K 个；
  3. 接受前缀中所有概率与目标对齐的 token，剩下重新采。

收益：大模型单次前向能验证多 token，吞吐量 1.5-3× 提升，质量无损（数学等价）。

变体：

  • N-gram 投机：用 n-gram 表替代小模型，免额外模型；
  • Self-speculative：大模型自己分早期层预测，后期层验证；
  • Medusa：在原模型加多个头并行预测多 token。

8.10 工程视角的 KV Cache 调优

项 调优
模型选 GQA 而非 MHA KV Cache 减 4-8×
启用 prefix cache system prompt 复用
KV 量化为 FP8 2× 压缩
选择支持 PagedAttention 的引擎 vLLM / SGLang / TRT-LLM
调 max_num_seqs / max_num_batched_tokens 平衡 batch 与延迟
长上下文用 MLA 或 sliding window 减 KV 增长

工程实战要点

  • KV Cache 是显存大头：长上下文 + 大 batch 时甚至超过模型本身；
  • GQA 是降低 KV 的最简单杠杆：选模型时优先；
  • PagedAttention + Continuous Batching 是 vLLM 性能秘密：相比 HF transformers 默认实现 2-4× 吞吐；
  • Prefix Cache 业务收益巨大：固定 system prompt 场景必开；
  • Speculative Decoding 是免费午餐：质量无损，吞吐 +；
  • Prefill / Decode 分离是前沿方向：大型部署正在采用；
  • max_num_seqs 调优：太大导致 KV OOM，太小吞吐低。

小结

  • KV Cache 缓存历史 K/V，避免重复计算，是 LLM 推理的标准优化；
  • KV Cache 显存 = 2 × L × seq × KV_heads × d × dtype × batch；
  • GQA / MQA 减小 KV heads，是最简单的 KV 优化杠杆；
  • KV 量化（FP8/INT4）2-4× 压缩；
  • MLA 是 DeepSeek 的架构级 KV 压缩；
  • PagedAttention 借鉴 OS 虚拟内存分页，显存利用率 30% → 90%+；
  • Continuous Batching 配合 PagedAttention 让长短请求混跑；
  • Prefix Cache 复用 system prompt 的 KV；
  • Speculative Decoding 质量无损提升吞吐；
  • Prefill/Decode 分离是前沿方向。

练习题

  1. 没有 KV Cache 时，生成 L 个 token 的总复杂度是多少？有 KV Cache 后呢？
  2. 估算 Llama-3 8B（32 层，8 KV heads，d=128，fp16，batch=8，seq=4096）的 KV Cache 大小。
  3. 为什么 GQA 能减小 KV Cache？减小几倍？
  4. PagedAttention 借鉴了 OS 的什么机制？解决了什么问题？
  5. Continuous Batching 相对 Static Batching 的核心改进是什么？
  6. Prefix Cache 在什么业务场景下收益最大？
  7. Speculative Decoding 为什么能"质量无损"地加速？
  8. Prefill 与 Decode 阶段的计算特性有何不同？为什么服务系统可能要分离它们？

下一章：第九章 推理引擎