Appearance
KV Cache 缓存 Transformer 推理中已计算的 Key 和 Value 向量,避免重复计算,是自回归生成的核心优化。
为什么需要 kv-cache
Transformer 自回归生成时,每生成一个新 token 都需要对之前所有 token 做注意力计算。KV Cache 将已计算的 K/V 向量缓存起来,使得每步只需计算当前 token 的 Q/K/V,然后与缓存的 K/V 做注意力。
核心原理
Prefill 阶段:处理整个 prompt
token[0..N] → 计算 K[0..N], V[0..N] → 存入 cache
Decode 阶段:逐个生成
token[N+1] → 计算 Q, K, V → K,V 追加到 cache
→ Attention(Q, K_cache, V_cache) → 输出 logits内存布局(每层):
K_cache: [max_positions, kv_hidden_dim]
V_cache: [max_positions, kv_hidden_dim]关键优化:
- GQA — Grouped Query Attention 减少 KV head 数
- Prefix Caching — 复用已处理的 prompt 前缀
- Rolling Cache — 淘汰旧位置以支持无限生成
在源码中的实现
src/llama-memory.cpp— KV Cache 管理逻辑src/llama-kv-cache.cpp— KV Cache 核心实现src/llama-kv-cells.h— Cache 单元定义src/llama-context.cpp— decode 时的 cache 更新
相关概念
- batch-decode — 批量解码利用 KV Cache
- backend — 不同后端的 cache 内存管理
- ggml — cache 使用 GGML 张量存储