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 与 PagedAttention」专属刷题排雷

做单选排雷题、推导选项机制，答错自动收录进专属错题本。