第 1 周 · Day 4：KV Cache 与 MQA / GQA

时长 8 小时：§1 KV Cache（3h，含 roofline）· §2 MQA / GQA（2h）· §3 实现（2h）· §4 推理系统一瞥 · §A–§C（1h）

对应正课 Lecture 2 的 MHA / MQA / GQA 部分。学习计划里的原话是”不懂 KV Cache 就不懂 GQA 为什么存在”，所以今天先花一半时间把 KV Cache 讲透。

配套代码 w1day4_checks.py（纯 numpy）：因果性 ⇒ 过去的 K/V 不变、增量解码与全量重算等价（MHA/GQA/MQA）、缓存前后的计算量、四个模型的缓存显存与 70B 估算题、算术强度与 roofline、MHA→GQA 均值池化、绝对位置编码让缓存失效、拼接 vs 预分配

§B 参考实现（code/ 目录）：w1d4_kvcache_gqa.py（在 Day 3 的 mini-GPT 上加缓存与 GQA，含正确性自检和速度测量）。

今天的目标：能证明 KV Cache 为什么成立；能在 30 秒内估出任意模型的缓存显存；能用”算术强度”解释为什么解码慢、为什么 GQA 有效；能写出一个与全量前向逐位一致的缓存生成。

§0 今天的主线

0.1 从训练转向推理

前三天关心的都是训练：一次前向算完 $T$ 个位置，因果掩码保证不泄露。今天换到推理，Day 1 §5.4 说过，推理必须一个 token 一个 token 地来。这带来一条清晰的因果链：

逐 token 生成 ──► 朴素做法每步重算整个前缀，O(T²) 的浪费       ← §1.1
     │
     ▼
KV Cache：过去的 K、V 不会变，存起来                           ← §1.2–1.3
     │  代价：显存随 层数 × 头数 × 长度 × batch 线性增长       ← §1.4
     ▼
解码的瓶颈变成"从显存读字节"，KV Cache 是其中越来越大的一块    ← §1.5
     │
     ▼
MQA / GQA：让多个 query 头共享 K、V，缓存砍到 1/8 甚至更少     ← §2

每一步都是被上一步的代价逼出来的。 这和 Week 0 Day 4 的发明史是同一种叙事，只是这次发生在推理侧。

0.2 今天要用到的旧结论

§1 KV Cache

1.1 动机：朴素生成在反复做同样的事

Day 3 的 generate 每一步都把整个上下文前向一遍。生成第 $t$ 个 token 时，前 $t-1$ 个 token 的全部中间结果，上一步刚刚算过一模一样的一遍。

w1day4_checks.py 第 3 组算了一层（$d=4096$）生成 $T$ 个 token 的累计计算量：

不缓存是 $O(T^2d^2+T^3d)$，缓存是 $O(Td^2+T^2d)$。生成越长，浪费越离谱。

1.2 直觉与证明：过去的 K、V 为什么不会变

直觉。 往序列末尾追加一个新 token，旧 token 的表示会变吗？在因果模型里不会：每个位置只能看它自己和它之前的位置，新 token 在所有旧 token 的”未来”，它们根本看不见它。

证明（对层数归纳）。 记 $h_j^{(l)}$ 为第 $j$ 个位置在第 $l$ 层的输出。

• 第 0 层：$h_j^{(0)}=\text{wte}[x_j]+\text{wpe}[j]$，只依赖 $x_j$ 和 $j$。
• 假设第 $l$ 层的 $h_j^{(l)}$ 只依赖 $x_{\le j}$。第 $l+1$ 层：
  • LayerNorm、FFN 是逐位置的：$j$ 位置的输出只用 $j$ 位置的输入。
  • 注意力：$j$ 位置的输出 $=\sum_{i\le j}A_{ji}v_i$，因果掩码保证只用到 $i\le j$ 的 $k_i,v_i$，而它们由 $h_i^{(l)}$（$i\le j$）算出，只依赖 $x_{\le i}\subseteq x_{\le j}$。
• 所以 $h_j^{(l+1)}$ 只依赖 $x_{\le j}$。$\blacksquare$

推论：第 $j$ 个位置在每一层的 $k_j=W_Kh_j$、$v_j=W_Vh_j$ 一旦算出，就永远不变。 追加多少新 token 都不影响它们。

w1day4_checks.py 第 1 组：前缀 12 个 token vs 13 个 token，两层里前 12 个位置的 K、V 最大差恰好是 0（不是”很小”，是逐位相同）。

为什么只缓存 K、V，不缓存 Q？ 生成下一个 token 只需要最后一个位置的输出。它的 query 是新算的；它要去看的是所有位置的 key 和 value。旧位置的 query 只在算旧位置自己的输出时用过，以后再也用不上。

为什么 BERT 没有 KV Cache？ 双向注意力下，新 token 会改变所有旧 token 的表示，归纳的第二步不成立。这也是 decoder-only 在推理上的结构性优势之一（Day 2 §4.3）。

1.3 算法：prefill 与 decode

prefill：把整个 prompt（长 P）一次前向，和训练时一样并行，
         顺便把每层每个位置的 K、V 写进缓存。输出最后一个位置的 logits。
decode ：循环：
           采样出新 token x_t
           只把 x_t（1 个 token）送进模型，位置编号是 t
           每层：算 x_t 的 q、k、v → k、v 追加进缓存 → q 对缓存里全部 t+1 个 key 做注意力
           输出 x_t 位置的 logits

两个阶段的性格完全不同：prefill 一次处理 $P$ 个 token，矩阵乘又大又满，是计算密集的；decode 每步只有 1 个 token，矩阵乘退化成矩阵-向量乘，是访存密集的（§1.5）。现代推理系统会把这两个阶段分开调度（§4）。

正确性。 缓存省的是重复计算，不是近似。第 2 组：prefill 8 个 token 再逐个解码 12 个，与一次前向 20 个 token 的 logits 最大差 $3.6\times10^{-15}$（MHA）、$6.2\times10^{-15}$（GQA）、$5.8\times10^{-15}$（MQA），全是浮点舍入。写缓存实现时，一定要用这个”与全量前向逐位比对”的测试。

1.4 显存：一个 30 秒的公式

$\boxed{\ \text{KV Cache 字节数}=2\times L\times n_{kv}\times d_{\text{head}}\times(\text{batch}\times\text{序列长})\times\text{每数字节数}\ }$

2 是 K 和 V 各一份；$n_{kv}$ 是 K/V 头数（MHA 时等于 $h$，此时 $n_{kv}d_{\text{head}}=d$）。

w1day4_checks.py 第 4 组（bf16）：

Llama 2 7B 每个 token 要半兆字节。 而 70B 模型每 token 的缓存反而比 7B 小，因为它用了 GQA。

学习计划的估算题：“70B、GQA 8 组、batch 32、序列 8192，KV Cache 大致占多少显存？”

$2\times80\times8\times128\times(32\times8192)\times2\ \text{字节}=80\ \text{GiB}$

同样配置换成 MHA（64 个 KV 头）是 640 GiB，MQA 是 10 GiB。对照：70B 的 bf16 权重本身约 130 GiB。长上下文 + 大 batch 时，KV Cache 可以和模型本身一样大甚至更大。

1.5 为什么解码是访存瓶颈：roofline

动机。 缓存把计算量降了三个数量级，可解码仍然慢。为什么？

算术强度。 一个操作的算术强度 = FLOPs ÷ 需要从显存读写的字节数。GPU 有两个上限：峰值算力（FLOP/s）和显存带宽（字节/s）。两者之比叫拐点：算术强度高于拐点，瓶颈是算力；低于拐点，瓶颈是带宽，GPU 大部分时间在等数据。

A100 80GB：bf16 稠密峰值 312 TFLOP/s，HBM 带宽 2.04 TB/s，拐点约 153 FLOP/字节。

解码一步做了什么？ 为了处理 batch 里的每个新 token，要把全部权重从显存读一遍（$2N$ 字节），做 $2N$ 次运算。batch 为 $B$ 时，权重读一次、用 $B$ 次；但每条序列的 KV Cache 只被它自己的 query 用，读多少用多少。

w1day4_checks.py 第 5 组（7B、MHA、A100）：

（最后一行已超过单卡 80GB，只为说明量级。）

读这张表：

1. batch=1 时算术强度约 1，比拐点低两个数量级。一个 7B 模型在 A100 上单条生成，理论上限约 140 token/s，这个上限由带宽决定，与算力无关。
2. 加大 batch 能摊薄权重读取：batch 8 时吞吐上限升到 892。这是推理服务要做批处理的根本原因。
3. 但 KV Cache 摊不薄。 batch 64、序列 4096 时，KV 读取（128 GiB）是权重（13 GiB）的十倍，算术强度又掉回 6.8。上下文越长，瓶颈越集中在 KV Cache 上。

意义。 推理优化的主线因此是”少读字节”：GQA/MQA 砍 KV 头（§2）、KV Cache 量化、权重量化（Lecture 4）。这也解释了一个常被误解的点：GQA 省的是访存和显存，不是计算（§2.5）。

1.6 代价与边界

显存：§1.4 的公式，随长度和 batch 线性增长。

碎片化：每条序列的最终长度事先不知道。如果按最大长度预分配，大部分空间浪费；如果动态增长，显存碎片化。vLLM 的 PagedAttention（Kwon et al. 2023）借用操作系统的分页思想，把缓存切成固定大小的块按需分配，显著提高了可容纳的并发数。

绝对位置编码 + 滑动窗口 = 缓存失效。 这是本周代码会实际碰到的问题。Day 3 的 mini-GPT 用可学习位置嵌入，上下文超过 block_size 后要”裁剪”，也就是窗口向后滑。但一滑，每个 token 的位置编号都变了，而位置信息已经烤进了 K 里。w1day4_checks.py 第 7 组：同一个 token 在旧窗口（位置 1）和新窗口（位置 0）里的 K 最大差 9.53。缓存只能整个重建。 参考实现里，超过 block_size 后每步都重新 prefill 一次，结果与朴素生成逐 token 相同，但加速也没了。

RoPE 让这个问题消失：它的打分只依赖相对位置，位置编号可以一直往上数，旧的 K 不用重算（Day 5 §2）。Mistral 7B 的滑动窗口 + 滚动缓存正是建立在这个性质上。这是 RoPE 在工程上最被低估的好处之一。

§2 MQA 与 GQA

2.1 动机

§1.4–1.5 把问题说得很清楚了：KV Cache 的大小和读取量都正比于 KV 头数 $n_{kv}$。多头注意力（Day 1 §4）里每个头都有自己的 K、V，$n_{kv}=h$。

能不能让多个 query 头共享同一套 K、V？ query 头还是 $h$ 个，每个头仍然有自己的”关注模式”，只是它们去同一套键和值里查表。

2.2 MQA：所有头共享一套

Multi-Query Attention（Shazeer 2019，论文标题就叫 One Write-Head is All You Need）：$h$ 个 query 头，只有 1 个 K/V 头。

$\text{head}_i=\mathrm{Attention}(XW_Q^{(i)},\ XW_K,\ XW_V),\qquad i=1..h$

收益：KV Cache 缩小到 $1/h$（70B 估算题从 640 GiB 到 10 GiB）。代价：质量有下降，且后续研究报告过训练不稳定。

2.3 GQA：分组共享，在两个极端之间插值

Grouped-Query Attention（Ainslie et al. 2023）：把 $h$ 个 query 头分成 $G$ 组，每组共享一个 K/V 头，$n_{kv}=G$。

MHA  (G = h = 8)：   q0 q1 q2 q3 q4 q5 q6 q7
                     k0 k1 k2 k3 k4 k5 k6 k7        每个 query 头一套 K/V

GQA  (G = 2)：       q0 q1 q2 q3 │ q4 q5 q6 q7
                        k0        │    k1           每 4 个 query 头共享一套

MQA  (G = 1)：       q0 q1 q2 q3 q4 q5 q6 q7
                              k0                    全部共享一套

$G=h$ 就是 MHA，$G=1$ 就是 MQA。GQA 是一个旋钮，在质量（接近 MHA）和缓存大小（接近 MQA）之间找平衡点。

实现上：K/V 投影只输出 $G\cdot d_{\text{head}}$ 维；做注意力之前，把每个 K/V 头复制 $h/G$ 份，让第 $g$ 个 K/V 头服务第 $g\cdot\frac hG$ 到 $(g+1)\cdot\frac hG-1$ 个 query 头（repeat_interleave）。复制只发生在计算时；缓存里存的是复制前的 $G$ 个头，这才是省显存的地方。

w1day4_checks.py 第 6 组：$G=h$ 时 GQA 与 MHA 的输出最大差 0，确实是同一个东西。

2.4 从 MHA 转换到 GQA：均值池化 + 继续训练

问题：已经有一个训好的 MHA 模型，想换成 GQA，必须从头训练吗？

GQA 论文的办法：把同一组里的 $h/G$ 个 K（和 V）投影矩阵取平均作为新的共享 K/V 头，然后用原预训练计算量的约 5% 继续训练（uptraining）。论文比较过几种初始化，均值池化效果最好。

为什么平均能工作？ 第 6 组给了一个对照：

随机权重的头彼此无关，平均之后信息大量丢失。训练好的模型里，很多头是冗余的（Day 1 §4.6 引用的剪枝研究），同组头相近时平均的损失小得多，剩下的差距靠少量继续训练补回来。

2.5 GQA 省了什么，没省什么

GQA 是一个推理优化，它针对的正是 §1.5 找到的瓶颈。

2.6 为什么现代模型都选 GQA

Llama 2 只在 34B 和 70B 上用了 GQA（7B、13B 仍是 MHA）；Llama 3 从 8B 起全部用 GQA，8 个 KV 头；Mistral 7B 也是 8 个 KV 头。8 个 KV 头几乎成了事实标准：相对 MHA，缓存缩小 4–8 倍，质量损失很小。MQA 则因为质量和稳定性问题较少被大模型采用。

2.7 更进一步（知道即可）

• MLA（DeepSeek-V2）：不减少头数，而是把 K、V 压缩成一个低维的潜向量再缓存，推理时再解压。缓存比 GQA 更小。
• 滑动窗口注意力（Mistral）：每层只看最近 $W$ 个 token，缓存大小封顶为 $W$。多层叠加后感受野仍可以覆盖更远（第 $l$ 层能间接看到 $l\cdot W$ 远）。
• 跨层共享 KV：相邻层共用一份缓存。

它们都在回答同一个问题：怎样让每个 token 占用更少的缓存字节。

§3 实现：给 mini-GPT 加缓存与 GQA

参考代码 code/w1d4_kvcache_gqa.py 复用 Day 3 的 GPT 和 Block，只换注意力类。几个值得单独讲的细节：

3.1 缓存怎么存：预分配，不要拼接

class KVCache:
    def __init__(self, B, n_kv, max_len, dk):
        self.k = torch.zeros(B, n_kv, max_len, dk)      # 一次分配到最大长度
        self.v = torch.zeros(B, n_kv, max_len, dk)
        self.pos = 0
    def update(self, k_new, v_new):
        T = k_new.size(2)
        self.k[:, :, self.pos:self.pos + T] = k_new
        self.v[:, :, self.pos:self.pos + T] = v_new
        self.pos += T
        return self.k[:, :, :self.pos], self.v[:, :, :self.pos]   # 视图，不复制

最直观的写法是每步 torch.cat([cache, k_new])。问题是每次拼接都要把已有的全部缓存复制到新分配的内存里。第 8 组：生成 4096 个 token，拼接累计搬运 839 万个位置，预分配只写 4096 个，差 2048 倍，还伴随反复的显存分配和碎片。

3.2 位置：增量解码时位置编号不是 0

新 token 的位置是它在整个序列里的下标 pos，不是 0。Day 3 的 GPT.embed(idx, start) 专门留了 start 参数：

x = self.wte(idx) + self.wpe(torch.arange(start, start + T, device=idx.device))

忘了这一点，每个新 token 都以为自己在位置 0，不报错，生成质量却会明显变差。与全量前向逐位比对的测试能立刻抓到它。

3.3 掩码：三种情况

带缓存时，query 的长度 $T$ 和 key 的长度 $S$ 可能不等：

第二行就是 Day 1 §A.1 边界 3 的坑：SDPA 的 is_causal=True 用左上角对齐的下三角，$T=1$ 时只允许看第一个 key。参考代码的自检第 4 项专门复现了这个反例。

第三行出现在”prompt 太长，分几块喂进去”或推测解码（§4）的场景。第 $i$ 个新 query 的绝对位置是 $S-T+i$，它能看的 key 是 $0..S-T+i$，所以允许矩阵是 torch.ones(T, S).tril(diagonal=S - T)。PyTorch 2.2 起也提供了 torch.nn.attention.bias.causal_lower_right 来表达它。

3.4 GQA 的三种写法

# 写法一：repeat_interleave，最直观，但会真的复制出 h 个头的 K/V
k = k.repeat_interleave(h // n_kv, dim=1)

# 写法二：expand + reshape，expand 本身不复制（reshape 时可能复制）
k = k[:, :, None].expand(B, n_kv, h // n_kv, S, dk).reshape(B, h, S, dk)

# 写法三：PyTorch ≥ 2.5，让 SDPA 自己处理
y = F.scaled_dot_product_attention(q, k, v, enable_gqa=True)

三种写法的分组方式必须一致：第 $j$ 个 query 头用第 $\lfloor j/(h/G)\rfloor$ 个 K/V 头。如果你用 repeat 而不是 repeat_interleave，分组会变成”交错”的（第 $j$ 个 query 头用第 $j\bmod G$ 个），不报错，但和别人的权重对不上。参考代码的自检第 8 项检查了 enable_gqa 与 repeat_interleave 一致。

3.5 必须通过的两个测试

1. logits 逐位一致：prefill + 逐 token 解码 vs 一次全量前向（MHA / GQA / MQA 各测一次，再测一次分块 prefill）。
2. 生成逐 token 一致：贪心解码下，缓存版与 Day 3 的朴素 generate 输出完全相同，包括超过 block_size 之后。

参考代码的自检覆盖了这两项。

§4 推理系统一瞥

今天的内容是整个推理优化领域的地基。之后几讲会陆续遇到这些名字，先建立一个地图：

共同的主题：解码是访存密集的，所有这些技术都在想办法让每读一个字节做更多的计算，或者干脆少读字节。

§A 缓存与 GQA：API 用法与边界

A.1 HuggingFace 的缓存

out = model(input_ids, use_cache=True)            # 生成模型默认就是 True
past = out.past_key_values                        # 新版本是 Cache 对象（如 DynamicCache）
out2 = model(next_ids, past_key_values=past, use_cache=True)

• model.generate() 默认使用缓存。训练时要关掉（use_cache=False，或 model.config.use_cache = False），特别是开了梯度检查点时，否则会报警告或浪费显存。
• 旧版本的 past_key_values 是嵌套 tuple：每层一个 (k, v)，形状 (B, n_kv, T, d_head)。新版本（约 4.36 起）引入了 Cache 类，DynamicCache 就是按需增长的实现。写代码时不要假设它是 tuple。
• 自己传 past_key_values 做增量前向时，还要正确传 position_ids 或 cache_position，道理同 §3.2。

A.2 配置里的头数

config.num_attention_heads      # query 头数 h
config.num_key_value_heads      # K/V 头数 G；等于 h 是 MHA，等于 1 是 MQA

看一个模型用不用 GQA、缓存多大，直接看这两个字段，再代入 §1.4 的公式。

A.3 SDPA 的 enable_gqa

PyTorch 2.5 起，F.scaled_dot_product_attention(q, k, v, enable_gqa=True) 接受 K/V 头数少于 Q 头数的输入，内部按 repeat_interleave 的方式分组，避免显式复制。老版本传这个参数会报 TypeError。

A.4 生成循环的其他细节

• torch.inference_mode() 比 no_grad() 更激进（张量不能再参与 autograd），推理时两者都可以，前者略快。
• 每步都 .item() 或 .cpu() 会强制同步（Week 0 Day 5 §A.5），GPU 上测速时要注意；测时间之前要 torch.cuda.synchronize()，否则测到的只是提交 kernel 的时间。
• 批量生成时各序列长度不同，要左补齐 + padding 掩码（Day 1 §3.4 的 NaN 陷阱在这里最容易出现）。

§B 实操

1. 跑通 w1day4_checks.py，8 组结论都看懂。第 5 组（roofline）值得多花时间：换成你自己的 GPU 参数重算一遍。
2. 纸笔（1h）：(a) 写出 §1.2 的归纳证明；(b) 不看表，算出 Llama 3 8B、batch 16、上下文 8192 的 KV Cache；(c) 算出你自己的 GPU 上 7B 模型 batch=1 解码的吞吐上限。
3. 给 mini-GPT 加 KV Cache（2h）：预分配缓存、start 位置、三种掩码。必须通过 §3.5 的两个测试再往下做。
4. 加 GQA（1h）：n_kv_head 可配；实现 MHA → GQA 的均值池化转换，把 Day 3 训好的模型转成 $G=2$、$G=1$，看验证 loss 掉了多少；再用少量步数继续训练，看能恢复多少。
5. 测速（0.5h）：朴素生成 vs 缓存生成，生成长度取 64、128、255，画出每 token 耗时随长度的变化。

参考实现：code/w1d4_kvcache_gqa.py（--selftest 覆盖第 3、4 题的正确性，--bench 做第 5 题）。

§C 自测

动机与直觉层

1. 朴素生成浪费在哪里？累计代价是多少阶？
2. 为什么过去 token 的 K、V 不会变？为什么不缓存 Q？BERT 为什么没有 KV Cache？
3. prefill 和 decode 的性格有什么不同？
4. 解码为什么是访存瓶颈？加大 batch 能缓解什么、缓解不了什么？
5. MQA 和 GQA 的动机是什么？GQA 为什么说是一个”旋钮”？
6. 为什么绝对位置编码会让滑动窗口下的缓存失效？RoPE 为什么能解决？

推导与估算层

7. 写出”KV Cache 不变”的归纳证明。
8. 写出 KV Cache 显存公式，算出 70B、GQA-8、batch 32、序列 8192 的结果；换成 MHA、MQA 呢？
9. 定义算术强度与拐点。算出 A100 的拐点，以及 7B 模型 batch=1 解码的吞吐上限。
10. GQA 省了哪些东西？没省哪些？为什么说它是推理优化？
11. MHA → GQA 的均值池化为什么在训练好的模型上比在随机权重上有效得多？

工程与 API 层

12. 为什么缓存要预分配而不是每步 torch.cat？
13. 增量解码时位置编号应该是多少？忘了会怎样？报错吗？
14. 单 token 解码、分块 prefill 各该用什么掩码？SDPA 的 is_causal=True 在哪种情况下错？
15. repeat 和 repeat_interleave 做 GQA，结果有什么区别？
16. HuggingFace 训练时为什么要关 use_cache？

出关标准

• 能写出 KV Cache 成立的归纳证明，并说清为什么只缓存 K、V
• 能在 30 秒内用公式估出任意模型的 KV Cache 显存，包括 70B 那道题
• 能用算术强度和 roofline 解释解码为什么慢、batch 能摊薄什么
• 能画出 MHA / GQA / MQA 的头分组图，说清 GQA 省了什么、没省什么
• 能说出 MHA → GQA 的转换方法和它为什么可行
• 能解释绝对位置编码为什么让滑动窗口下的缓存失效
• 写出带 KV Cache 和 GQA 的 mini-GPT，通过 §3.5 的两个一致性测试
• 测出缓存带来的加速，并能用 §1.1 的复杂度解释

疑问清单