现代 AI 的“推理加速之谜”：从 KV Cache 到 PagedAttention 的演进

在 LLM 推理优化领域，人们经常讨论 PagedAttention 或 Speculative Decoding，但这些技术解决的底层问题其实只有一个：如何高效地管理和利用 KV Cache（Key-Value Cache）。如果你想理解为什么大模型推理如此吃显存，以及为什么上下文越长速度越慢，KV Cache 是唯一的答案。

1. 为什么需要 KV Cache？

LLM 生成文本是自回归（Autoregressive）的：每生成一个新 token，都需要将之前所有生成的 token 作为输入重新喂给模型。

在 Transformer 的注意力机制中，每个 token 都会产生三个向量：Query (Q)、Key (K) 和 Value (V)。
- Query：当前 token “想要寻找什么”。
- Key：历史 token “能提供什么”。
- Value：历史 token “实际包含的内容”。

计算过程是：$\text{Attention}(Q, K, V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V$。

关键点在于：对于已经生成过的历史 token，它们的 $K$ 和 $V$ 向量在后续步骤中是完全不变的。如果每次生成新 token 都要重新计算一遍所有历史 token 的 $K$ 和 $V$，计算量将随序列长度呈平方级增长 $\mathcal{O}(n^2)$。

为了避免重复计算，我们将每一步产生的 $K$ 和 $V$ 存储在显存中，这就是 KV Cache。这样，每一步只需计算当前新 token 的 $Q, K, V$，然后直接从缓存中读取之前的 $K, V$ 进行矩阵乘法。复杂度因此降为 $\mathcal{O}(n)$。

2. KV Cache 的代价：显存黑洞

虽然 KV Cache 提升了速度，但它带来了巨大的显存压力。

计算公式

一个模型的 KV Cache 大小取决于：$\text{batch_size} \times \text{seq_len} \times \text{num_layers} \times \text{num_heads} \times \text{head_dim} \times \text{precision} \times 2$ (Key 和 Value 各一份)。

实战量化

以 Llama-3-8B 为例（FP16）：
- 层数: 32, 头数: 32, 每个头维度: 128.
- 单个 token 在单层产生的 KV 大小 = $2 \times 32 \times 128 \times 2\text{ bytes} = 16\text{ KB}$。
- 全模型单 token KV 大小 = $32\text{ layers} \times 16\text{ KB} = 512\text{ KB}$。

如果 batch_size=32且 seq_len=4096：
$32 \times 4096 \times 512\text{ KB} \approx 67\text{ GB}$。

这意味着即使模型权重本身只占约 15GB，为了支持高并发的长文本推理，你可能需要一张 A100 (80GB) 或更多显卡，仅仅是为了存放这些“记忆碎片”。

3. 从“暴力存储”到“智能管理”的演进

面对显存压力，工业界演进出了三种主流方案：

A. GQA (Grouped-Query Attention) —— 从结构上减负

传统的 MHA 每个 Query 头对应一个 KV 头。GQA 将 Query 分组，每组共享一对 KV 头（如 Llama-3 使用）。这在保持性能的同时显著降低了 KV Cache 的物理尺寸。

B. PagedAttention —— 从内存管理上优化

传统的 KV Cache 要求连续内存空间，导致严重的内存碎片化（Internal Fragmentation）。vLLM 推出的 PagedAttention 将 KV Cache 分页存储（类似操作系统的虚拟内存），允许非连续存储且动态按需分配，将显存利用率提升至接近 100%。

C. 量化 (Quantization) —— 从精度上压缩

将 KV Cache 从 FP16 量化到 INT8 或 FP8。这可以直接将显存占用减半或更多，而对模型生成质量的影响在可接受范围内。

总结

KV Cache 是 LLM 推理的“空间换时间”典范。它解决了自回归生成的冗余计算问题，但也成为了限制吞吐量和上下文长度的最大瓶颈。从 GQA 到 PagedAttention 再到量化，AI 系统工程的演进本质上就是在与这块巨大的“记忆碎片”做斗争。