KV 缓存的工作原理：
它对应的 K 和 V 在不同生成步骤中是完全相同的。
这是由两个条件决定的：
模型权重在推理阶段固定；
Causal Mask 保证当前位置只能看到自身以及之前的位置。
因此，KV Cache 的基本思想非常直接：
K 和 V 只计算一次，然后保存起来。之后的生成步骤直接读取缓存，而不是重新计算。
概念本身并不复杂，但真正重要的是：
KV Cache 到底保存了什么？
Forward Pass 如何发生变化？
为什么只保存 K/V，而不保存 Q？
它节省了多少计算？
又带来了多少显存开销？
What Exactly Gets Stored?
KV Cache 究竟存储什么？
对于模型中的每一个 Attention Layer，我们都保存两个矩阵：
K_cache: positions 0, 1, 2, ..., n-1 的所有 Key vectors
V_cache: positions 0, 1, 2, ..., n-1 的所有 Value vectors
也就是说：
Layer 0:
K_cache
V_cache
Layer 1:
K_cache
V_cache
Layer 2:
K_cache
V_cache
...
Layer L:
K_cache
V_cache
需要特别注意：
我们保存 K 和 V，但不保存 Q。
后面会具体解释原因。
The Shape of the Cache
KV Cache 的形状
对于：
单个 Transformer Layer
单个 Attention Head
有：
K_cache shape: [current_sequence_length, head_dim]
V_cache shape: [current_sequence_length, head_dim]
例如：
sequence_length = 1000
head_dim = 128
那么：
K_cache = [1000, 128]
V_cache = [1000, 128]
这里每一行对应一个 token 的 K 或 V。
对于完整模型，可以抽象成：
Full KV cache shape:
[num_layers, 2, current_sequence_length, num_heads, head_dim]
↑ ↑ ↑ ↑ ↑
│ │ │ │ └─ 每个 head 的维度
│ │ │ └─ Attention heads
│ │ └─ 随生成长度不断增加
│ └─ K 和 V
└─ 每个 Transformer Layer 都有自己的 Cache
也就是：
KVCache∈RL×2×N×H×dheadKVCache \in \mathbb{R}^{L\times2\times N\times H\times d_{head}}
其中：
LL：Transformer Layer 数
NN：当前 Sequence Length
HH：Attention Head 数量
dheadd_{head}：每个 Head 的维度
Concrete Example: LLaMA-7B
一个具体例子：LLaMA-7B
博客以 LLaMA-7B 为例：
Layers = 32
Heads = 32
head_dim = 128
dtype = float16
每个数值 = 2 bytes
对于长度为 seq_len 的序列：
CacheSize=num_layers×2×seq_len×num_heads×head_dim×bytesCacheSize = num\_layers \times 2 \times seq\_len \times num\_heads \times head\_dim \times bytes
代入：
32×2×seq_len×32×128×232\times2\times seq\_len\times32\times128\times2
得到：
524,288×seq_len bytes524,288\times seq\_len\ bytes
也就是大约：
0.5MB×seq_len0.5MB\times seq\_len
因此，对于一个：
2048 tokens
的 Sequence：
0.5MB×2048≈1GB0.5MB\times2048 \approx 1GB
也就是说：
单个请求，仅 KV Cache 就可能占用约 1GB。
这也是为什么长上下文推理中，KV Cache 会成为非常重要的显存消耗来源。
How the Forward Pass Changes
KV Cache 如何改变 Forward Pass？
现在分别看：
不使用 KV Cache；
使用 KV Cache。
区别会非常明显。
Without KV Cache
不使用 KV Cache
假设当前已经有：
[t_0, t_1, ..., t_{n-1}]
总共 nn 个 token。
最朴素的实现类似：
def forward_naive(all_tokens):
hidden_states = embed(all_tokens)
# Shape: [n, hidden_dim]
for layer in transformer_layers:
# 为所有 token 重新计算 Q、K、V
Q = hidden_states @ W_Q
# Shape: [n, head_dim]
K = hidden_states @ W_K
# Shape: [n, head_dim]
V = hidden_states @ W_V
# Shape: [n, head_dim]
# 完整 Attention
attention_scores = Q @ K.T / sqrt(d)
# Shape: [n, n]
attention_scores = apply_causal_mask(
attention_scores
)
attention_weights = softmax(
attention_scores
)
attention_output = attention_weights @ V
# Shape: [n, head_dim]
hidden_states = layer.ffn(
attention_output
)
# 生成下一个 token 实际只需要最后一个位置
return hidden_states[-1] @ W_output
问题在于：
假设上一轮已经计算过：
K₀ K₁ K₂ ... Kₙ₋₂
V₀ V₁ V₂ ... Vₙ₋₂
这一轮生成新 token 后，又重新计算：
K₀ K₁ K₂ ... Kₙ₋₂ Kₙ₋₁
V₀ V₁ V₂ ... Vₙ₋₂ Vₙ₋₁
但是：
K₀ ... Kₙ₋₂
V₀ ... Vₙ₋₂
和上一轮完全一样。
这些计算全部是重复计算。
With KV Cache
使用 KV Cache
使用 Cache 后，Decode 阶段每次只传入一个新 token：
def forward_with_cache(new_token, kv_cache):
"""
new_token:
当前新生成的一个 token
kv_cache:
保存每一层历史 K/V
"""
hidden_states = embed(new_token)
# Shape: [1, hidden_dim]
for layer_idx, layer in enumerate(
transformer_layers
):
# 只计算当前 token 的 Q/K/V
Q_new = hidden_states @ W_Q
# Shape: [1, head_dim]
K_new = hidden_states @ W_K
# Shape: [1, head_dim]
V_new = hidden_states @ W_V
# Shape: [1, head_dim]
# 读取这一层历史 K/V
K_cached = kv_cache[layer_idx]['K']
# Shape: [n-1, head_dim]
V_cached = kv_cache[layer_idx]['V']
# Shape: [n-1, head_dim]
# 历史 Cache + 当前 token
K_full = concat([
K_cached,
K_new
])
# Shape: [n, head_dim]
V_full = concat([
V_cached,
V_new
])
# Shape: [n, head_dim]
# 更新 Cache
kv_cache[layer_idx]['K'] = K_full
kv_cache[layer_idx]['V'] = V_full
# 当前 token 的 Q 查询所有 K
attention_scores = (
Q_new @ K_full.T / sqrt(d)
)
# Shape: [1, n]
attention_weights = softmax(
attention_scores
)
# Shape: [1, n]
attention_output = (
attention_weights @ V_full
)
# Shape: [1, head_dim]
hidden_states = layer.ffn(
attention_output
)
return hidden_states @ W_output
核心变化就是：
没有 KV Cache：
输入 n 个 token
↓
重新计算 n 个 Q
重新计算 n 个 K
重新计算 n 个 V
↓
计算 [n × n] Attention
变成：
使用 KV Cache：
输入 1 个新 token
↓
计算 1 个 Q_new
计算 1 个 K_new
计算 1 个 V_new
↓
读取历史 K/V
↓
计算 [1 × n] Attention
The Key Differences
两种方式的核心区别
项目Without KV CacheWith KV Cache每次处理 token 数nn1新计算的 Knn 个1 个新计算的 Vnn 个1 个新计算的 Qnn 个1 个从显存读取 K/V无n−1n-1 个历史 K/VAttention Matrix[n,n][n,n][1,n][1,n]
所以 KV Cache 本质上做了一件事：
用 Memory Storage + Memory Read，换取避免重复计算。
也就是：
Compute⟶MemoryCompute \longrightarrow Memory
这就是 KV Cache 的基本 trade-off。
Visualizing Cache Growth
KV Cache 是如何增长的？
假设 Prompt：
["The", "cat", "sat"]
STEP 0：Prefill
输入：
["The", "cat", "sat"]
三个 token 并行处理。
计算：
K₀ K₁ K₂
V₀ V₁ V₂
KV Cache：
┌──────────────────────────┐
│ K_cache: [K₀, K₁, K₂] │
│ V_cache: [V₀, V₁, V₂] │
└──────────────────────────┘
然后模型输出：
"on"
STEP 1：Decode "on"
输入只有：
["on"]
当前 token 位置：
position = 3
新计算：
Q₃
K₃
V₃
从 Cache 读取：
K₀ K₁ K₂
V₀ V₁ V₂
Attention：
Q₃
↓
[K₀ K₁ K₂ K₃]
即：
Q3[K0,K1,K2,K3]TQ_3 [K_0,K_1,K_2,K_3]^T
计算结束后更新 Cache：
┌──────────────────────────────┐
│ K_cache: [K₀,K₁,K₂,K₃] │
│ V_cache: [V₀,V₁,V₂,V₃] │
└──────────────────────────────┘
模型生成：
"the"
STEP 2：Decode "the"
输入：
["the"]
只计算：
Q₄
K₄
V₄
读取：
K₀ K₁ K₂ K₃
V₀ V₁ V₂ V₃
然后：
Q₄
↓
[K₀ K₁ K₂ K₃ K₄]
即：
Q4[K0,K1,K2,K3,K4]TQ_4 [K_0,K_1,K_2,K_3,K_4]^T
新的 Cache：
┌─────────────────────────────────┐
│ K_cache: [K₀,K₁,K₂,K₃,K₄] │
│ V_cache: [V₀,V₁,V₂,V₃,V₄] │
└─────────────────────────────────┘
然后输出：
"mat"
所以整个 Decode 的模式始终是：
计算新 token 的 K/V
↓
读取全部历史 K/V
↓
当前 Q 与全部 K 做 Attention
↓
利用 Attention Weight 聚合全部 V
↓
把新 K/V 添加到 Cache
↓
生成下一个 token
每生成一个 token：
KV Cache += 一个 token 的 K/V
因此 Cache 大小随 Sequence Length 线性增长。
Why Don't We Cache Q?
为什么不缓存 Q？
一个自然的问题是：
既然 K 和 V 都缓存了，为什么不把 Q 也缓存？
原因就在于自回归生成时 Attention 的工作方式。
Decode 阶段，我们只计算当前新 token 的 Attention。
例如当前生成位置 nn：
QnQ_n
需要和：
K0,K1,…,KnK_0,K_1,\ldots,K_n
进行计算。
但是历史 token：
token 0
token 1
token 2
...
token n-1
不需要