第 7 章 KV Cache 优化¶
开场问题：第 4 章的公式已经算出来了，长上下文下 KV Cache 会超过权重本身。
这一章解决：现代模型省 KV 的三条路——砍头数、压维度、跨层共享。它们互相正交，可以叠加使用。
回到行程助手：第 4 章算过，历史对话的 K/V 缓存随轮数线性增长。一张卡能同时陪多少用户聊、每个人能聊多长，基本由它决定。这一章讲现代模型怎样在效果几乎不掉的前提下，把这块缓存缩到原来的几分之一。
这里省的是缓存里的头数和维度，不是让每个 token 少走一批权重。模型依然稠密，前向计算量仍约是 \(2\times\) 全部参数。本章要挡住的失败是：上下文一长，KV 比权重还占显存。
7.1 先算一下：MHA 的开销有多大¶
标准多头注意力里，每个 Q 头都配一个自己的 K 头和 V 头。KV 显存 = 2 × 2 × L × h × d_h × n × B。
以 32 头、36 层、head_dim 128 的模型算：每 token 每层 16 KB，全模型 576 KB/token，32K 上下文单请求就是 19 GB。这个数字在 2022 年之前是可以忍的（上下文都短），2023 年之后彻底不能忍。
从张量形状看，MHA 是
省 KV 的关键不是减少 Q 头——Q 只在当前步短暂存在——而是让 \(K,V\) 的头数 \(n_{\mathrm{kv}}\) 小于 \(h\)。缓存缩放比恰好是 \(n_{\mathrm{kv}}/h\)。
7.2 MQA：省到极致的代价¶
所有 Q 头共用一组 K/V。KV 显存直接除以头数，32 头就是省 32 倍。
此时 \(Q\in\mathbb R^{B\times h\times n\times d_h}\)，而 \(K,V\in\mathbb R^{B\times1\times n\times d_h}\)，计算时把同一组 K/V 广播给所有 Q 头。问题是容量可能下降：所有头共享同一套“索引和内容”，只能靠不同 Q 去形成差异。掉点多少与模型、数据和训练方式有关，不能笼统断言一定明显；在强检索或长上下文任务上风险通常更高。
7.3 GQA：当前的默认¶
折中方案：把 Q 头分成若干组，每组共享一组 K/V。分组数常取 8，KV 开销降到 MHA 的 1/4 左右，质量损失在噪声范围内。
GQA 只规定头数关系 \(n_{\mathrm{kv}}\le h\) 以及组内如何共用 K/V。QK-Norm、RoPE、支路入口的 RMSNorm 与它正交：可以一起装进同一条注意力支路（装配顺序见第 8 章），但不改本章的缓存形状公式，因此下面只写投影之后的头维。
一般地，每组有 \(g=h/n_{\mathrm{kv}}\) 个 Q 头，KV 相对 MHA 缩小 \(g\) 倍。比如 \(h=32,n_{\mathrm{kv}}=8\)，是 4 个 Q 头共享一组 KV，缓存降到 1/4。 实现时的形状链是：
q_proj → [B, n, 32, 128]
k_proj → [B, n, 8, 128]
v_proj → [B, n, 8, 128]
│ 每个 KV 头服务连续 4 个 Q 头
attention kernel 内广播（或 repeat_interleave）
↓
逻辑上的 K/V [B, n, 32, 128]
高效 kernel 不会真的复制 4 份 K/V，而是在读取时按 q_head // 4 映射到 KV 头。手写参考实现可用 repeat_interleave(g, dim=head_axis)，但必须先确认布局是 [B,n,h,d_h] 还是 [B,h,n,d_h]；广播错轴或把分组顺序打乱，形状仍然正确，logits 却无法对齐。
选 8 还有工程原因：若张量并行也是 8 路，每张卡可拥有一组或整数个 KV 头，避免复制；但它不是固定法则，n_kv 应与头数、并行度和质量消融一起定。质量、显存、部署三方较易平衡，所以 GQA 成为 2023 年之后的默认。
7.4 MLA：换个维度压¶
DeepSeek 的思路不是减少头数，而是把 K/V 投影到一个低维隐向量上，缓存里只存这个隐向量，用的时候再升维还原。这是第 0 章那个"低秩分解"的直接应用。
若把所有头拼接后的 KV 记为 \(k_t,v_t\)，MLA 先缓存低维向量
缓存从每 token 的 \(2h d_h\) 个元素变成约 \(d_c\) 个元素（外加一小段位置相关 key）。推理时无需真的完整还原：例如
可把上投影吸收到 query 一侧。这就是“矩阵吸收”省掉显式升维的原因。
RoPE 是位置相关变换，无法同固定的 \(W^{UK}\) 一起预先吸收，因此每头拆成“带 RoPE 的小部分 + 不带 RoPE 的压缩部分”。MLA 可以比 GQA 压得更狠，同时保留更丰富的头级表示；代价是训练、缓存布局和 kernel 都更复杂，这也是它没能像 GQA 一样普及的主要原因。
7.5 跨层共享 KV¶
第三条路：注意到相邻层的注意力模式高度相似，那就让若干层复用同一份 KV（CLA、YOCO 等）。显存按共享层数成比例下降，且与 GQA / MLA 完全正交，可以叠加。
这里要区分两种做法：**跨层注意力（CLA）**让一组相邻层直接复用某层产生的 K/V；YOCO 类结构把网络拆成 self-decoder 与 cross-decoder，让后半段跨层复用前半段生成的全局 KV。前者改动小，后者能把 prefill 与缓存复用做得更彻底，但层之间耦合更强。这是 2024 年之后进入主流视野的方向，在端侧和超长上下文场景里尤其有吸引力。
7.6 对比¶
| 方案 | KV缓存压缩比例（相对MHA） | 训练难度 | 推理Kernel复杂度 | 长文本能力 | 简要备注 |
|---|---|---|---|---|---|
| MHA（标准多头注意力） | 100%（基线，无压缩） | 低 | 低 | 优 | 原始Transformer注意力，KV头与Q头一一对应，显存开销最大；当前很少用于大模型推理 |
| MQA（多查询注意力） | 约\(1/h\)（h为Q头数，h=32则≈3.1%） | 中 | 低 | 中 | 全部Q头共享唯一一组K/V；压缩最强，但长文本、细粒度检索任务容易掉点；早期PaLM、Falcon使用 |
| GQA（分组查询注意力） | \(1/g\)（g为每组Q头数量，g=4则≈25%） | 低 | 中 | 优 | 工业界当前默认方案；Q头分组，组内共享KV；质量损失基本落在噪声区间；LLaMA2、Qwen、GLM主流标配 |
| MLA（多头潜在注意力） | 约7%~10%（DeepSeek-V2配置） | 高 | 高 | 优+ | 不减少头数，将KV压缩为低维隐向量；RoPE需要拆分成两段处理；压缩上限高于GQA，长上下文表现更好；Kernel与训练流程复杂 |
| 跨层共享KV（CLA/YOCO） | \(1/S\)（S为共享层数，S=4则≈25%） | 中偏高 | 中偏高 | 优 | 正交方案，可叠加GQA/MLA；多层复用同一份KV；显存按共享层数下降；代价是prefill阶段层间依赖变强，prefill吞吐受损 |
小结：MHA是性能基线，但KV显存开销过高；MQA压缩最激进，代价是长文本能力下降；GQA在质量、显存、工程落地三者取得最佳平衡，是当下工业界首选。MLA通过低秩隐向量进一步突破KV压缩上限，更适合超长上下文场景，但训练与内核实现成本显著提升。跨层共享KV属于正交优化手段，可以和GQA、MLA叠加使用，在显存预算极度紧张的端侧/超长序列场景使用，需要接受prefill阶段的性能损失。
动手：改一处¶
把注意力从 MHA 换成 GQA（n_kv=8），用第 4 章的公式重算 KV 显存并实测验证，记下质量和显存的变化。
核心问答¶
Q1：GQA 的分组数为什么常取 8？
它是常见经验点，不是通用最优值：不少配置在 8 个 KV 头附近能兼顾质量与缓存，而 8 路张量并行也容易切分。实际要同时满足 n_kv 整除/兼容 TP、KV 预算和质量消融。
Q2：MQA 到底差在哪？ 所有 Q 头共享一份 K/V，压缩最激进，因此头级 K/V 表达容量最小。它不等于多头完全失效——Q 头仍然不同——但在部分长文本和检索任务上更容易掉点；是否明显必须看模型和训练方法。
Q3：MLA 缓存的到底是什么？ 是一个低维隐向量（DeepSeek-V3 里是 512 维量级），不是 K 和 V 本身。用的时候通过学好的升维矩阵还原出各头的 K、V。缓存量与头数无关，这是它能压得比 GQA 狠的原因。
Q4：MLA 和 RoPE 为什么冲突，怎么解决？ 低秩压缩要求能把升维矩阵提前吸收进 Q 的投影里，而 RoPE 是位置相关的旋转，夹在中间无法被吸收。解法是把每个头的维度拆成两段：一小段照常做 RoPE 并直接缓存，其余部分走低秩压缩路径。
Q5：三条路能一起用吗？ 能，它们分别作用在"头数""每头维度""层数"三个不同的轴上，互相正交。实践中是否全上取决于实现复杂度和目标场景，端侧模型往往叠得更狠。