KV cache 缓存什么？为什么能加速？
占多少显存？
先弄清楚：每个概念是什么
- Autoregressive decode
- 一次只生成 1 个 token，拼回输入再生成下一个。
- KV cache
- causal mask 下，前面 token 的 K、V 不会被后面的 token 改变，所以存起来直接用。Q 只用当前这一个，以前的 q 早就用完了，所以不存。
- Prefill
- 整个 prompt 一次并行 forward，把所有 K、V 填进 cache，并给出第一个输出 token。矩阵乘很大，所以 compute-bound。
- Decode
- 每步只处理 1 个新 token，却要把全部 weights 和整个 cache 从 GPU memory 读一遍，所以 memory-bound。
30 秒回答
With a causal mask, the keys and values of earlier tokens never change, so each layer caches them. Each decoding step computes q, k and v only for the new token, appends k and v, and attends one query over all cached keys: O(n·d) per layer instead of re-running the whole prefix. Prefill fills the cache in parallel and is compute-bound; decode handles one token per step and is memory-bound. The price is memory: 2 × layers × KV heads × head dim × bytes per token, times length and batch, which can exceed the weights. GQA, paging and quantization reduce it.
关键公式与手算
2 = K 和 V；\(L\) 层，\(n_{kv}\) 个 KV head，context 长度 \(T\)，batch \(B\)。
LLaMA 2 7B（32 层，\(d = 4096\)，MHA，fp16）：
- 每 token：2 × 32 × 4096 × 2 bytes = 512 KiB
- \(T = 4096\)：一条序列 2 GiB；\(B = 32\)：64 GiB，约是 weights（6.7B × 2 bytes ≈ 13.4 GB）的 5 倍
- 换成 GQA 的 8 个 KV head（LLaMA 3 8B）：每 token 128 KiB，小 4 倍，见 MQA、GQA、MLA
省了多少计算：decode 第 \(t\) 步每层，有 cache 是 \(O(td + d^2)\)，没有 cache 要重算整个前缀，是 \(O(t^2d + td^2)\)。推导见 Transformer 复杂度。
手撕 KV cache（NumPy，单层单 head）
class KVCache:
def __init__(self, max_len, d):
self.K = np.zeros((max_len, d)) # 预分配，不每步拷贝
self.V = np.zeros((max_len, d))
self.t = 0 # 已缓存几个 token
def step(self, x, Wq, Wk, Wv):
# x: (d,) 新 token；只为它算 q、k、v
q, k, v = x @ Wq, x @ Wk, x @ Wv
self.K[self.t], self.V[self.t] = k, v
self.t += 1
K, V = self.K[:self.t], self.V[:self.t]
s = K @ q / np.sqrt(q.size) # 1 个 query 对 t 个 key
p = np.exp(s - s.max())
return (p / p.sum()) @ V逐个 token 调 step 的输出和一次算完整的 causal attention 相同。多层就每层一个 cache。
必答点
- 缓存每层前面 token 的 K、V，不存 Q：causal 下它们不会变
- 每步只为新 token 算 q、k、v，1 个 query 对 \(t\) 个 key，每层 \(O(td)\)
- prefill 并行、compute-bound；decode 逐 token、memory-bound
- GPU memory \(= 2\cdot L\cdot n_{kv}\cdot d_{head}\cdot \text{bytes}\cdot T\cdot B\)，long context 的主要瓶颈
- 缩小：MQA / GQA / MLA、PagedAttention、KV cache quantization
追问
用了 RoPE，cache 里存的是转之前还是转之后的 k？
转之后的。k 的位置在写入时就定了，以后不用再转。
PagedAttention 解决什么问题？
以前的系统按最大长度给每个请求预留一整块连续 GPU memory，只有 20.4%–38.2% 真正存了 K、V。vLLM 按固定大小的 block 分页，浪费接近 0，throughput 提升 2–4 倍。
KV cache 最大能到多少？
Pope et al.：500B+ 的 MHA 模型，batch 512、context 2048 时 KV cache 共 3 TB，是 weights 的 3 倍，每步读 cache 比读 weights 还久。
还有什么办法省？
KV cache quantization，DeepSeek-V2 平均每个元素 6 bits。多个请求共用同一段 prompt 前缀的 cache。只保留最近一段窗口的 K、V（sliding window）。
来源Pope et al. 2022 §1（p.1）、§2–2.2（p.3）· Shazeer 2019 §2.4–2.4.1（pp.3–4）· Kwon et al. 2023, vLLM Abstract、§1（p.2）· Touvron et al. 2023a, LLaMA Table 2（p.3）· Touvron et al. 2023b, Llama 2 §2.2、Table 1（pp.5–6）· DeepSeek-V2 §2.1.3（p.8）、§3.2.3（p.16）