Appearance
KV Cache & Prompt Caching
KV Cache 是 Transformer 模型推理的核心优化技术，通过缓存中间计算结果避免重复计算，将自回归生成的复杂度从 O(n²) 降低到 O(n)。Prompt Caching 则进一步在系统层面缓存重复的输入，是生产环境降低成本的关键手段。
核心原理
为什么需要 KV Cache
在 Transformer 的自注意力机制中，每生成一个新 token 时都需要计算该 token 与所有已生成 token 之间的关系：
无 KV Cache: 每步都重新计算所有历史 token 的 Q, K, V
有 KV Cache: 只计算新 token 的 Q，复用缓存的 K, V复杂度对比：
- 无 KV Cache: O(n²) 每步，总共 O(n³)
- 有 KV Cache: O(n) 每步，总共 O(n²)
KV Cache 的工作机制
在每个 Transformer 层中，为每个已生成的 token 缓存其 Key 和 Value 向量：
python
# 简化示意
for i in range(seq_len):
q = W_q @ x[i] # 新 token 的 Query
k = W_k @ x[i] # 新 token 的 Key
v = W_v @ x[i] # 新 token 的 Value
K_cache[i] = k # 缓存 Key
V_cache[i] = v # 缓存 Value
scores = q @ K_cache[:i+1].T # 与所有缓存的 Key 计算
output = softmax(scores) @ V_cache[:i+1]内存优化策略
KV Cache 的内存挑战
KV Cache 是推理时最大的内存消耗来源之一：
| 参数 | 计算公式 | 示例 (Llama 3 70B, 8K 上下文) |
|---|---|---|
| 隐藏维度 d | 模型配置 | 8,192 |
| 层数 L | 模型配置 | 80 |
| 头数 h | 模型配置 | 64 |
| KV Cache 大小 | 2 × L × h × d × seq_len × bytes | ~8 GB |
对一个 70B 模型运行 8K 上下文，KV Cache 可能占用 8-16 GB 显存。
优化方法
1. GQA / MQA (Grouped-Query / Multi-Query Attention)
减少 KV 头的数量：
- MQA: 所有 Query 头共用一组 K, V
- GQA: 多个 Query 头分组共用 K, V
- 效果: 内存减少 4-8x，质量损失很小
2. KV Cache 量化
将 KV Cache 从 FP16 量化到 INT8 或 INT4：
- KV Cache Quantization: 通常保持质量的同时减少 50-75% 内存
- 实现: vLLM、SGLang 等框架已支持
3. 滑动窗口 (Sliding Window)
只缓存最近的 N 个 token：
- 适用: 长上下文但远距依赖较少的任务
- 例子: Mistral 的 Sliding Window Attention
4. 压缩技术
- H2O (Heavy Hitter Oracle): 只保留重要的 KV 对
- StreamingLLM: 保持窗口内的重要 token
- Compressed Sparse Attention (CSA): DeepSeek 的方案，压缩率达 98%
Prompt Caching
与 KV Cache 的区别
| 特征 | KV Cache | Prompt Caching |
|---|---|---|
| 层级 | 模型内部 | 系统/应用层 |
| 缓存对象 | 中间激活值 | 完整输入的 KV Cache |
| 生命周期 | 单次推理 | 跨请求复用 |
| 实现复杂度 | 模型内建 | 需要系统支持 |
实现方式
客户端缓存
python
# 简化示意
import hashlib
cache = {}
def generate_with_cache(prompt, model):
prompt_hash = hashlib.sha256(prompt.encode()).hexdigest()
if prompt_hash in cache:
# 复用缓存的 KV Cache
kv_cache = cache[prompt_hash]
return model.generate(prompt, kv_cache=kv_cache)
else:
# 首次计算并缓存
result, kv_cache = model.generate(prompt)
cache[prompt_hash] = kv_cache
return result服务端缓存 (Prefix Caching)
- vLLM Prefix Caching: 自动检测和复用相同前缀的 KV Cache
- SGLang RadixAttention: 基于 Radix Tree 的高效前缀匹配
- OpenAI Prompt Caching: API 层面的缓存，降低 50% 成本
应用场景
| 场景 | 节省比例 | 说明 |
|---|---|---|
| 多轮对话 | 30-50% | 系统指令重复使用 |
| 批量处理 | 60-80% | 相同的上下文模板 |
| RAG 系统 | 40-60% | 文档片段重叠 |
| Agent 工作流 | 50-70% | 工具描述重复 |
工具与框架
| 框架 | KV Cache 优化 | Prompt Caching |
|---|---|---|
| vLLM | PagedAttention + Prefix Caching | 支持 |
| SGLang | RadixAttention | 支持 |
| TensorRT-LLM | In-flight Batching + KV Cache Mgmt | 部分支持 |
| llama.cpp | KV Cache 量化 | 支持 |
| TGI | 标准 KV Cache | 支持 |
Synthesis
KV Cache 是 Transformer 推理中"不可或缺"的优化，它将自回归生成从不可行变为高效。随着上下文窗口增长，KV Cache 的内存消耗已成为主要瓶颈，GQA、量化、压缩等技术正在持续推进。Prompt Caching 则是在系统层面进一步提升效率的关键，尤其在企业级应用中能显著降低成本。
Related Pages
- Model Inference & Deployment — 模型推理与部署
- Speculative Decoding — 推理加速技术
- Model Quantization — 模型量化
- Mixture of Experts — MoE 架构中的 KV Cache 优化
Sources
"Efficient Memory Management for Large Language Model Serving with PagedAttention" (SOSP 2023)
"FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness" (NeurIPS 2022)
vLLM 官方文档
SGLang 技术文档
Large Language Model (LLM) — 大语言模型的核心定义、技术原理与发展历程