3 min read#attention#kv-cache#inference#sglang

KV Cache

Apr 7, 2026

键值缓存(KV Cache) 是 Transformer 自回归推理中的核心优化：把历史 token 已计算的 Key 和 Value 向量缓存起来，后续生成新 token 时直接复用，避免对整段前缀重复计算。

为什么需要 KV Cache

在标准 Transformer 的 self-attention 中，生成第 $t$ 个 token 需要计算它对前面所有 token 的注意力：

$o_t = \text{Attention}(q_t,\; K_{1:t},\; V_{1:t}) = \text{softmax}\!\left(\frac{q_t K_{1:t}^\top}{\sqrt{d_k}}\right) V_{1:t}$

如果没有缓存，每生成一个新 token 都要重新计算所有历史 token 的 K 和 V——即把整段前缀重新做一遍前向传播。KV Cache 的做法是：历史 token 的 $K$ 和 $V$ 只算一次就存下来，后续步骤直接从缓存中读取。

逐步示例：生成 “The cat is cute”

以一个单头、单层的简化 Transformer 为例。设 prompt 为 "The cat"（2 个 token），模型依次生成 "is" 和 "cute"。

记 $W_Q, W_K, W_V$ 为注意力层的投影矩阵， $x_i$ 为第 $i$ 个 token 的隐状态。每个 token 的 Q/K/V 向量为：

$q_i = W_Q x_i, \quad k_i = W_K x_i, \quad v_i = W_V x_i$

阶段一：Prefill “The cat”

将 prompt 的两个 token 一次性送入模型，并行计算所有 Q/K/V：

$\underbrace{q_1, q_2}_{\text{用于 prefill 阶段的注意力计算}}, \quad \underbrace{k_1, k_2, v_1, v_2}_{\text{写入 KV Cache}}$

Prefill 阶段的注意力矩阵（因果掩码下）：

$A_{\text{prefill}} = \text{softmax}\!\begin{pmatrix} q_1 k_1^\top & -\infty \\ q_2 k_1^\top & q_2 k_2^\top \end{pmatrix} \cdot \frac{1}{\sqrt{d_k}}$

Prefill 结束后，KV Cache 的状态：

$\boxed{\text{Cache} = \{(k_1, v_1),\; (k_2, v_2)\}}$

阶段二：Decode “is”（第 3 个 token）

模型只处理新 token $x_3$ （"is"），计算它的 Q/K/V：

$\underbrace{q_3}_{\text{新计算}}, \quad \underbrace{k_3, v_3}_{\text{新计算，追加到 Cache}}$

注意力计算时， $k_1, k_2, v_1, v_2$ 从缓存读取，不重新计算：

$o_3 = \text{softmax}\!\left(\frac{q_3 \cdot [\overbrace{k_1, k_2}^{\text{从 Cache 读取}}, \overbrace{k_3}^{\text{新计算}}]^\top}{\sqrt{d_k}}\right) \cdot \begin{bmatrix} \color{gray}{v_1} \\ \color{gray}{v_2} \\ v_3 \end{bmatrix}$

Decode 结束后，将 $k_3, v_3$ 追加到缓存：

$\boxed{\text{Cache} = \{(k_1, v_1),\; (k_2, v_2),\; (k_3, v_3)\}}$

阶段三：Decode “cute”（第 4 个 token）

同样，只处理新 token $x_4$ （"cute"）：

$\underbrace{q_4}_{\text{新计算}}, \quad \underbrace{k_4, v_4}_{\text{新计算，追加到 Cache}}$

$o_4 = \text{softmax}\!\left(\frac{q_4 \cdot [\overbrace{k_1, k_2, k_3}^{\text{从 Cache 读取}}, \overbrace{k_4}^{\text{新计算}}]^\top}{\sqrt{d_k}}\right) \cdot \begin{bmatrix} \color{gray}{v_1} \\ \color{gray}{v_2} \\ \color{gray}{v_3} \\ v_4 \end{bmatrix}$

$\boxed{\text{Cache} = \{(k_1, v_1),\; (k_2, v_2),\; (k_3, v_3),\; (k_4, v_4)\}}$

每步的计算量对比

步骤	新计算的 Q/K/V	从 Cache 读取的 K/V	注意力规模
Prefill “The cat”	$q_1, q_2, k_1, k_2, v_1, v_2$	无	$2 \times 2$
Decode “is”	$q_3, k_3, v_3$	$k_1, k_2, v_1, v_2$	$1 \times 3$
Decode “cute”	$q_4, k_4, v_4$	$k_1, k_2, k_3, v_1, v_2, v_3$	$1 \times 4$

没有 KV Cache 时，生成 "cute" 要对全部 4 个 token 重算 K/V，注意力规模为 $4 \times 4$ 。有 KV Cache 后降为 $1 \times 4$ ——只有新 token 需要算 Q/K/V，历史 K/V 全部复用。

KV Cache 没有消除的开销

新 token 仍要对历史所有 K/V 做注意力（上例中 $q_4$ 要和 $k_1, k_2, k_3, k_4$ 都算点积），因此单步 decode 的计算量仍随上下文长度线性增长。KV Cache 省掉的是”重复计算历史 K/V”的冗余，而不是注意力本身的 $O(n)$ 开销。

Prefill 与 Decode 的分界

阶段	输入	计算特征	KV Cache 行为
Prefill	整段 prompt（ $n$ 个 token）	并行计算，计算密集	批量写入 $n$ 组 K/V
Decode	每步 1 个新 token	串行逐步，访存密集	读取全部历史 + 追加 1 组 K/V

Prefill 是计算密集(compute-bound)，Decode 是访存密集(memory-bound)。这也是 serving 系统将二者当作不同工作负载调度的原因。

显存占用

对于一个 $L$ 层、 $h$ 头、头维度 $d_k$ 的模型，缓存 $n$ 个 token 的 KV Cache 占用：

$\text{KV Cache Size} = 2 \times L \times h \times d_k \times n \times \text{sizeof(dtype)}$

其中因子 2 来自 K 和 V 各一份。以 LLaMA-2 70B（80 层，64 头 GQA 8 KV 头， $d_k = 128$ ，FP16）为例，缓存 4096 个 token 的 KV Cache 约占：

$2 \times 80 \times 8 \times 128 \times 4096 \times 2\text{B} \approx 1.3\text{ GB}$

这解释了为什么现代架构用 GQA/MLA 减少 KV 头数——直接按比例压缩 KV Cache 的显存占用。

工程实现

Paged Attention

传统实现为每个请求预分配最大序列长度的连续显存，造成严重的内碎片。分页注意力(Paged Attention) 借鉴操作系统虚拟内存的思想，把 KV 存储拆成固定大小的页(page)，按需分配。

在 sglang 中，这通过两层映射实现：

ReqToTokenPool：请求的第 $i$ 个 token → 逻辑槽位
TokenToKVPool：逻辑槽位 → GPU 上的物理地址

动态分配、按需回收，避免为短请求浪费整段显存。

Prefix Cache / Radix Cache

如果两个请求共享相同前缀（如相同的 system prompt），这段前缀的 KV 只需计算一次。sglang 维护 radix tree 结构来查找最长公共前缀并复用已有 cache，直接影响调度器的调度策略。

延迟回收

请求完成后通常先释放引用关系，而不立即擦除底层 KV 数据。如果后续来了共享前缀的请求，这些 cache 仍可被命中复用。只有显存紧张时才真正执行 eviction。

代价与限制

维度	收益	代价
计算	历史 K/V 不重算	单步 decode 仍随上下文线性增长
显存	大幅提升吞吐	需要额外显存保存所有历史 K/V
调度	可做 prefix-aware batching	需要池化、回收和 eviction 逻辑