KV Cache 通过缓存历史 K、V 矩阵，将 LLM 自回归推理复杂度从 $O(n^2)$ 降至 $O(n)$，是推理优化的核心技术。
核心原理
工作流程
| 阶段 | 操作 |
|---|
| Prefill | 计算所有 Prompt Token 的 K、V 并缓存 |
| Decode | 仅计算新 Token 的 K、V，追加到缓存 |
复杂度对比
| 场景 | 无 KV Cache | 有 KV Cache |
|---|
| 计算复杂度 | $O(n^2 \cdot d_k)$ | $O(n \cdot d_k)$ |
| 显存占用 | $O(n^2)$ | $O(n)$ |
显存计算公式
$$\text{Memory} = 2 \times 2 \times \text{num_heads} \times \text{seq_len} \times \text{head_dim} \times \text{batch_size}$$
优化技术
| 技术 | 效果 |
|---|
| 量化（INT4/INT8） | 显存降低 2-4 倍 |
| 滑动窗口 | 支持超长序列 |
| PagedAttention | 显存利用率提升 2-3 倍 |
| 分层缓存 | 支持 100k+ 序列 |
Hugging Face 使用
1
2
3
4
5
6
| from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-2-7b-hf",
use_cache=True # 启用 KV Cache
)
|
核心数学性质
- 等价性：KV Cache 与标准注意力计算结果完全一致
- 线性复杂度：解码阶段复杂度与序列长度线性相关
- 增量可扩展：新 Token 的 K、V 可直接追加
发展演进
| 阶段 | 核心问题 | 代表技术 |
|---|
| 基础版（2017-2020） | 重复计算 | 简单缓存 |
| 优化版（2021-2023） | 显存瓶颈 | 量化、滑动窗口、PagedAttention |
| 系统级（2024-） | 存储/调度限制 | LMCache、Mooncake、存算分离 |
张芷铭的个人博客
Comments