进阶← 返回知识库
KV Cache（键值缓存）
Transformer 推理优化：缓存已计算的 Key 和 Value 矩阵，生成新 token 时只算增量并拼接，避免重复计算历史 token。将每步注意力计算复杂度从 O(n²) 降到 O(n)。
核心公式
📖 符号说明
| 符号 | 含义 | 说明 | 取值范围 |
|---|---|---|---|
| 历史 Key 缓存 | 前 t-1 步已计算的 Key 矩阵 | (t-1, d_k) | |
| 新 token 的 Key | 当前 token 经 W_K 投影的结果 | (1, d_k) | |
| 完整 Key 矩阵 | 拼接后的全序列 Key | (t, d_k) |
前置知识
KV Cache（键值缓存）
一句话理解
KV Cache 就是给 Attention 装了一个"备忘录"——算过的 Key 和 Value 不重算，下次直接查表。
为什么需要 KV Cache
自回归生成时，每生成一个新 token 都需要算一次 Self-Attention。而 Self-Attention 需要所有位置的 Key 和 Value 矩阵。没有缓存时：
- 生成第 1 个 token：计算 1 个位置的 K/V
- 生成第 2 个 token：重新计算 2 个位置的 K/V
- 生成第 n 个 token：重新计算 n 个位置的 K/V
前面的 token 没有变化，K/V 也不会变——全部重算是纯粹的浪费。
核心思想
缓存历史 token 的 K 和 V，新 token 只需算自己的 Knew 和 Vnew，然后拼接到缓存：
K1:t=[KcacheKnew],V1:t=[VcacheVnew]新 token 的 Query Qnew 与完整的 K1:t 做点积，得到对所有历史 token 的注意力权重。
复杂度对比
| 无缓存 | 有 KV Cache | |
|---|---|---|
| 第 t 步计算量 | O(t2d)（全量 Q×K） | O(td)（1×t 点积） |
| 生成 n 个 token 累计 | O(n3d) | O(n2d) |
| 额外内存 | 无 | O(n⋅d) 存 K/V |
时间换空间：用 O(nd) 的内存，省下 O(n) 倍的计算。
代码实现
python
class CachedAttention:
def __init__(self, W_Q, W_K, W_V):
self.W_Q, self.W_K, self.W_V = W_Q, W_K, W_V
self.cache_K = None
self.cache_V = None
def __call__(self, x):
Q = x @ self.W_Q # 只算新 token 的 Query
K_new = x @ self.W_K
V_new = x @ self.W_V
# 拼接缓存
if self.cache_K is not None:
K = np.vstack([self.cache_K, K_new])
V = np.vstack([self.cache_V, V_new])
else:
K, V = K_new, V_new
self.cache_K, self.cache_V = K, V # 更新缓存
scores = Q @ K.T / np.sqrt(K.shape[-1])
weights = softmax(scores)
return weights @ V
def reset(self):
self.cache_K = self.cache_V = None
为什么不需要缓存 Q？
Query 只用于当前 token 和历史 Key 做匹配——用完即弃，下一步的 Query 是新 token 的，和上一步没有关系。而 Key 和 Value 会被所有后续 token 反复使用，所以只有 K/V 值得缓存。
KV Cache 的代价
KV Cache 用内存换速度——对于大模型来说内存开销不可忽视。以 GPT-3（175B）为例，生成 2048 token 时 KV Cache 约占 3GB 显存。第 6 篇会介绍量化等技术来缓解这个问题。