KV Cache

TODO 需要更加清晰的图示说明

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+ 序列

更根本的压缩来自注意力结构本身：从 MHA 到 MLA 从头结构削减缓存规模，稀疏注意力 直接减少需要 attend 的 token 数。

算子侧的 FlashAttention 与本篇的缓存管理互补，共同构成推理加速主线。

Hugging Face 使用

from transformers import AutoModelForCausalLM
 
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    use_cache=True  # 启用 KV Cache
)

核心数学性质

1. 等价性：KV Cache 与标准注意力计算结果完全一致
2. 线性复杂度：解码阶段复杂度与序列长度线性相关
3. 增量可扩展：新 Token 的 K、V 可直接追加

发展演进

阶段	核心问题	代表技术
基础版（2017-2020）	重复计算	简单缓存
优化版（2021-2023）	显存瓶颈	量化、滑动窗口、PagedAttention
系统级（2024-）	存储/调度限制	LMCache、Mooncake、存算分离