KV Cache 与 PagedAttention

KV Cache = 自回归解码时把已生成 token 的 Key/Value 向量缓存起来、避免每步重复计算历史 attention 的推理显存结构。每步解码只需计算新 token 的 Q,与缓存的历史 K/V 做 attention,新 K/V 追加进缓存。缓存显存公式(每元素 FP16 占 2 字节): $\text{KV Bytes} = 2 \cdot 2 \cdot n_{layers} \cdot n_{kv\_heads} \cdot d_{head} \cdot s \cdot b$ —— 第 1 个 2 = K 与 V 各一份,第 2 个 2 = FP16 每元素 2 字节,$s$ = 序列长度,$b$ = batch。数值例 (LLaMA-3 70B: 80 层、GQA 8 个 KV heads、$d_{head}=128$,32K 上下文、8 并发): $2 \times 2 \times 80 \times 8 \times 128 \times 32768 \times 8 \approx 85.9\text{ GB} \approx 86\text{ GB}$;若换回同规模 MHA(64 个 KV heads)则约 $687\text{ GB}$ —— 正好差 $64/8 = 8$ 倍,这就是 GQA 压缩 KV 的直观收益。PagedAttention (vLLM): KV 不再按连续长数组预分配,而是切成固定大小的 block(默认每块 16 个 token,即 $16 \times n_{kv\_heads} \times d_{head}$ 个元素),用页表 (block table) 记录逻辑位置 → 物理块映射,像操作系统分页一样按需分配: 消除内部碎片(静态预分配超出实际生成长度)与外部碎片(不同长度序列交错、先来后走导致空闲空间无法拼接),并把「预留显存」降到几乎为 0,显存利用率从连续分配的约 60-80% 提升到接近 100%。

所有 LLM 推理服务标配(vLLM/SGLang/TensorRT-LLM 内核),长上下文、高并发批量解码是显存主瓶颈;面试高频“手算 KV cache 显存”“PagedAttention 为什么省显存”“GQA/MLA 如何压缩”。

attention 每步与全部历史计算复杂度为 $O(s)$,不缓存则总计算量 $O(s^2)$ 且重复;KV 缓存把解码复杂度降到每步 $O(1)$ 增量计算,但显存随 $s \times b$ 线性增长(长序列/大并发下可达数十 GB,占单卡显存 30-60%)。PagedAttention 用分页消除碎片、逼近 100% 显存利用率,使 70B 模型 16K 上下文能批量服务;GQA/MLA 进一步把 $n_{kv\_heads}$ 压缩 8 倍/更低,直接线性降低 KV 显存。

手算 KV cache 显存: 写出公式 $\text{KV Bytes} = 2\cdot2\cdot n_{layers}\cdot n_{kv\_heads}\cdot d_{head}\cdot s\cdot b$ 并解释两个 2 的含义;给定 80 层、8 个 KV heads、$d_{head}=128$、16K 上下文、batch 16,算每序列/整批多少 GB。

PagedAttention 为什么能消除显存碎片: 16 token/块的页表机制;内部碎片与外部碎片分别怎么产生、分页后如何消失;显存利用率从约多少提升到多少。

GQA/MLA 如何压缩 KV 显存: MHA → GQA 共享 KV heads 后公式里哪一项变小;LLaMA-3 70B (64→8) 对应 8 倍降幅;MLA 更进一步(低秩投影压缩)的原理。

Prefill 与 Decode 阶段 KV 缓存的写入与复用差异: prefill 一次性写入全部 s 个 KV、decode 每步追加一个;为什么长 prompt 的显存峰值出现在 prefill 末尾。

KV 缓存与量化/压缩的精度权衡: 为什么 KV 量化 (如 INT8) 比权重量化更敏感;与 FlashAttention、StreamingLLM 滑动窗口如何组合控制显存。