KV Cache#

1. 问题的来源#

Decoder-only 模型（如 GPT）的生成过程是自回归的——每次只生成一个 token，然后把它加入上下文，再生成下一个。

生成第 $t$ 个 token 时，需要对前面所有 $t-1$ 个 token 做注意力。**如果不做优化：**每生成一个新 token，都要对整个序列重新计算所有 token 的 Q、K、V。前面的 token 的 K 和 V 实际上已经算过了，但又重新算了一遍——纯粹浪费。

生成长度为 $L$ 的序列，总的 K/V 计算量是 $1 + 2 + 3 + \dots + L = O(L^2)$，其中大量是重复计算。

2. KV Cache 的做法#

核心思想：把每一层已经算过的 K 和 V 缓存起来，后续 step 直接复用。

1. 生成第 1 个 token：计算 $Q_1, K_1, V_1$，缓存 $K_1, V_1$
2. 生成第 2 个 token：只计算新 token 的 $Q_2, K_2, V_2$，把 $K_2, V_2$ 拼接到缓存中，用 $Q_2$ 和 $[K_1, K_2]$ 做注意力
3. 生成第 $t$ 个 token：只计算 $Q_t, K_t, V_t$，缓存中已经有 $[K_1, \dots, K_{t-1}]$

$$K_{\text{cache}} = [K_1; K_2; \dots; K_{t}], \quad V_{\text{cache}} = [V_1; V_2; \dots; V_{t}]$$

$$\text{attn}_t = \text{softmax}\left(\frac{q_t \cdot K_{\text{cache}}^T}{\sqrt{d_k}}\right) V_{\text{cache}}$$

每步只需计算 1 个新 token 的 Q/K/V，K/V 从缓存读取

K/V 的计算量从 $O(L^2)$ 降到 $O(L)$——每步只算 1 个 token 的 K/V。

3. 代价：显存#

KV Cache 是经典的空间换时间。每一层、每个头都要缓存 K 和 V 矩阵，而且随着生成长度增长，缓存越来越大。

以 LLaMA-7B 为例（32 层、32 头、$d_k = 128$），生成 2048 个 token 时：

$$\text{KV Cache} = 2 \times 32 \times 32 \times 128 \times 2048 \times 2 \text{ bytes (FP16)} \approx 1 \text{ GB}$$

2 = K和V两个矩阵，32 = 层数，32 = 头数，128 = $d_k$，2048 = 序列长度，2 = FP16 每个数 2 字节

对于更大的模型（如 70B）或更长的序列（如 128K），KV Cache 的显存开销会达到几十 GB，成为推理的主要瓶颈。

4. 解决方案的演进#

为了减小 KV Cache，社区发展出了一系列方法：

4.1 MQA → GQA → MLA#

MQA（Multi-Query Attention）

**做法：**所有 Q 头共享 1 组 KV。

**效果：**KV Cache 缩小到 $1/32$，但太极端，性能下降明显。

**名字含义：**Q 还是多头（Multi-Query），KV 缩成了唯一一组——从 Q 的视角命名。

GQA（Grouped-Query Attention）

**做法：**32 个 Q 头分成 8 组，每组共享 1 组 KV。

**效果：**KV Cache 缩小到 $1/4$，性能损失很小。

**定位：**MHA 和 MQA 的折中，目前主流方案。

MLA（Multi-head Latent Attention）

**做法：**不减头数，而是把 KV 低秩压缩到极低维潜在空间。

**效果：**压缩比 ~50×，性能反而超过 MHA。

**思路转变：**从"减少 KV 头数"转向"压缩 KV 维度"。

4.2 MLA 详解（以 DeepSeek-V2 为例）#

第零步：回顾标准 MHA 的计算与缓存

先区分 整个多头投影的宽度 与 单个 head 的点积维度。MHA 为每个 head 独立计算 $q_{t,i}k_{s,i}^{\top}$，不能把所有 head 拼成一个向量再做一次点积，否则会把各头混为一个注意力分布。

MLA 用共享低维 latent $c_s=h_sW_{DKV}$ 表示内容 KV。以行向量约定，单个 head 有

$$k^C_{s,i}=c_sW_{UK,i},\qquad v_{s,i}=c_sW_{UV,i}.$$

这是一种从头训练的低秩架构约束，不是把任意已有 MHA 权重无损压缩。解耦的 RoPE key $k^R_s$ 另行缓存；因此每 token、每层的缓存维度为 $d_c+d_r$。例如 DeepSeek-V2 的 $512+64=576$；相比 $2n_hd_h=32768$ 的示例 MHA 配置约小 56.9 倍，实际字节数还取决于 dtype 等因素。

Key 吸收：

$$q^C_{t,i}(k^C_{s,i})^\top=(q^C_{t,i}W_{UK,i}^{\top})c_s^\top.$$

把 query 映射到 latent 维度后，可直接对缓存的 $c_s$ 打分。该等式按 head 成立。若 query 路径含低秩投影和归一化，不能把整个路径武断合并成一个从 $h_t$ 出发的线性矩阵。

**Value 吸收：**先在 latent 空间加权，再恢复输出：

$$\sum_s a_{t,s,i}v_{s,i}=\left(\sum_s a_{t,s,i}c_s\right)W_{UV,i}.$$

所以无需为每个历史 token 恢复完整 V；$W_{UV,i}$ 还可与该 head 对应的输出投影组合。Value 路径里没有 $W_QW_{UV}^{\top}$ 这一操作。

以下是单个 head 的 NumPy 形状示意，省略 batch、mask、query 的具体构造和多头求和。c、k_rope 已包含当前 token 的缓存：

import numpy as np

# q_content: [d_content]; W_UK: [d_c, d_content]
# c: [L, d_c]; q_rope: [d_r]; k_rope: [L, d_r]
q_absorbed = q_content @ W_UK.T
scores = (c @ q_absorbed + k_rope @ q_rope) / np.sqrt(d_content + d_r)
weights = np.exp(scores - scores.max())
weights /= weights.sum()
context = weights @ c
# W_UV: [d_c, d_v]; W_O_head: [d_v, d_model]
output_head = context @ W_UV @ W_O_head

MLA 的低缓存收益和算子性能还取决于 prefill/decode 阶段及具体实现。进一步推导见 MLA 专题。

💡 MLA vs GQA 的本质区别

GQA 是"共享"——多个 Q 头用同一组 KV，减少 KV 的"组数"，但共享同一组的 Q 头看到的 KV 完全相同，表达能力受限。

MLA 是"压缩"——不减少头数，而是把所有头的 KV 信息联合压缩到一个低维空间。恢复时 $W_{UK}$ 为每个头生成不同的 K，保留了多头的表达能力。再加上吸收技巧消除了恢复开销。这解释了 MLA 的缓存效率；精度是否优于 GQA/MHA 取决于具体模型和训练实验，矩阵吸收本身不提供精度保证。

4.3 第二个方向：减少缓存的序列长度#

不是所有历史 token 的 KV 都同样重要。有些方法会选择性地丢弃不重要的 KV：

• 滑动窗口注意力（Mistral 用的）：只缓存最近 $w$ 个 token 的 KV，更早的直接丢掉。假设 $w = 4096$，不管序列多长，KV Cache 的大小永远是 4096。代价是模型看不到很远的历史信息。
• Token 淘汰策略：根据注意力权重判断哪些历史 token "不重要"（权重一直很低），把它们的 KV 从缓存中踢掉，只保留重要的。

4.4 第三个方向：压缩每个 KV 的精度#

• KV 量化：不减少数量，而是把每个 KV 向量从 FP16（2 字节）压缩成 INT8（1 字节）甚至 INT4（0.5 字节）。精度略有损失，但缓存大小减半甚至更多。

**优化 KV Cache 总结起来就三条路：**头数上减少（GQA/MQA）、长度上减少（滑窗/淘汰）、精度上压缩（量化）。目标都是一个：用更少的显存存储历史信息，同时尽量不影响模型效果。

🔥 面试标准回答

KV Cache 的作用是避免自回归生成时重复计算已有 token 的 K/V，用空间换时间。代价是显存开销随序列长度线性增长。GQA 通过让多个 Q 头共享 KV 来减小缓存大小，是现在的主流方案（LLaMA 2/3、Mistral 都在用）。

延伸阅读#

• 推理优化与部署速查：部署时怎样区分延迟、吞吐和显存问题？