第 09 章　GQA、MLA 与跨层 KV 共享

生成一句话时，模型每次只增加一个 token，却要反复读取历史的 key 和 value。第 08 章用缓存避免了重复投影，但历史越长，缓存本身就越大。本章接着问：这些历史信息，能不能少存一些？

GQA、MLA 和跨层共享分别从三个方向入手：GQA 减少 KV 头，MLA 减少每个 token 的存储维度，跨层共享减少独立保存 KV 的层数。这样做也改变了模型能独立表示哪些信息。因此，它们都不是可以随手打开的推理开关；对任意已训练的 MHA 模型直接换用这些结构，不能期待输出保持不变。

本章源码基准为固定提交中的 GQA、MLA 和 KV sharing。以下推导和小实验为本书原创。阅读结果时要分清来源：上游的随机模型生成和 GPU 时间不属于本书实测；正文保留的实验误差也不能当作速度或语言质量的证据。

09.1 先把缓存账算在正确的对象上

设批量为 $B$ ，历史长度为 $T$ ，层数为 $L$ ，query 头数为 $H$ ，每头维度为 $d$ ，每个元素占 $s$ 字节。若 key 和 value 维度相同，普通 MHA 缓存占 $2 B T L H d s$ 字节。这个式子只计算保存下来的 key 和 value，不包括参数、临时注意力分数、框架内存池和其他激活。因此，缓存少四倍，并不代表进程显存也少四倍。

例如 $B = 1, T = 1024, L = 8, H = 8, d = 64, s = 2$ ，得到 16,777,216 字节，即 16 MiB。GB 使用十进制，GiB 使用二进制。读估算器结果时，除了代入的参数，也要核对输出单位。还有一处实现差异：GQA 等估算脚本采用向上取整头宽的方式，容许不能整除的输入；实际模型通常要求整除。因此，估算器能打印结果，并不保证同一组参数能用来实例化模型。

09.2 GQA：多个问题，可以共用一套历史记录

令 KV 头数为 $G$ ，要求 $H$ 能被 $G$ 整除，每组含 $r = H / G$ 个 query 头。输入 $X$ 的形状为 $[B, T, D]$ 。三个投影分别输出 $Q : [B, T, H d]$ 、 $K, V : [B, T, G d]$ ；拆头以后，query 第 $h$ 个头读取 KV 第 $g (h) = ⌊ h / r ⌋$ 个头：

A_{h} = softmax (\frac{Q_{h} K_{g (h)}^{⊤}}{\sqrt{d}} + M) V_{g (h)} .

$M_{t, j}$ 在 $j > t$ 时为负无穷，否则为零。同一组 query 头虽然读取相同的 KV，每个头却仍有自己的投影，算出的相关性分数可以不同。共享的是历史表示，不是注意力权重。缓存缩为 $2 B T L G d s$ ，与 MHA 之比为 $G / H$ 。当 $G = H$ 时退化为 MHA； $G = 1$ 时是 MQA。GQA 原论文研究了从多头模型继续训练的路径，论文中的质量结论依赖训练方案，不能由本章的等价性测试推出。

源码 GroupedQueryAttention.__init__ 用 num_kv_groups 表示 KV 头数，建立输出宽度为 num_kv_groups * head_dim 的 W_key 和 W_value。forward 先缓存未展开的 KV，再用 repeat_interleave(group_size, dim=1) 对齐 query 头。若两个 KV 头记作 A、B，四个 query 头应读取 A、A、B、B。若误用 repeat，排列就成了 A、B、A、B。张量形状完全一样，却把 query 头配给了不同的 KV 头；仅检查形状，发现不了这个错误。

一个容易混淆的细节来自同目录的估算器：其 CLI --n_kv_groups 用作 每个 KV 头分摊的 query 头数，以 n_heads // n_kv_groups 推出 KV 头数。同名参数在两处指向不同的量，不能把估算器的命令行参数原封不动搬入模型字典。记笔记或改配置时，先明确写出 $H, G, r$ ，再对应到具体函数参数，就不容易把组数和每组大小弄反。

缓存推理时， $Q$ 长度可能只有 1，而 $K$ 长度为全部历史。ptr_current_pos 决定 query 的绝对位置；此时若对一个 $1 \times T$ 矩阵套用从零开始的上三角模板，就会把当前 query 误当成序列开头，只让它看见第一个 key。reset_cache 同时清空 K、V 和位置计数；新请求如果只重置其中一个，会产生跨请求污染或位置错位。

09.3 MLA：先存一个低维向量，需要时再展开

GQA 让多个头共用 KV，MLA 则进一步改变缓存中保存的表示。教学版 MLA 令潜在宽度为 $c$ ，用三个矩阵建立分解：

C = X W_{D}, K = C W_{U K}, V = C W_{U V},

其中 $W_{D} : [D, c]$ ，两个上投影矩阵为 $[c, H d]$ ，缓存 $C : [B, T, c]$ 。缓存大小变为 $B T L c s$ 。例如 $H d = 512, c = 128$ ，每 token 的 KV 存储从 1024 个数缩到 128 个数。但投影 $W_{D} W_{U K}$ 的秩至多为 $c$ ；低维表示因此限制了模型能表达的投影；它并不保证任意旧 key 都能无损压缩。

源码 MultiHeadLatentAttention.forward 的顺序很重要：先对新输入产生 query 与 latent_new，把 latent 追加到 cache_c_kv，再对 全部历史 latent 执行 W_UK、W_UV。这段实现减少了跨解码步保存的数据，但计算时仍会形成完整的 KV 临时张量，而且每个解码步都要重新上投影历史。于是，cache_c_kv.numel() 只能描述潜在缓存的元素数，不能代表整个注意力的峰值内存；字节数减少了多少，也不能直接换算成速度提升。

对一个头，暂忽略偏置与位置旋转，有

Q (C U_{K})^{⊤} = (Q U_{K}^{⊤}) C^{⊤}, P (C U_{V}) = (P C) U_{V} .

先把 query 投到潜在空间，再与 $C$ 点乘，可不显式重建 K；value 的上投影也能后移。这两步利用的是矩阵乘法结合律，没有引入近似。难点在于位置处理：如果对展开的 key 施加依赖位置的旋转，旋转与上投影一般不能随意交换，便不能直接照搬这套变换。DeepSeek-V2 的完整 MLA 使用解耦的旋转位置部分；本目录保留 GPT 的绝对位置 embedding，直接从输入生成 query，没有完整复现论文的 query 低秩支路和解耦 RoPE。完整模型的缓存公式还应计入位置相关 key，而不能只剩 $c$ 。

本书的 latent_and_sharing() 同时检查 $X (W_{D} W_{U K}) = (X W_{D}) W_{U K}$ 和上述投影吸收。原稿记录的最大误差均为 $1.78 \times 10^{- 15}$ 。这项检查的范围很窄：在给定低秩参数后，比较两种计算顺序是否等价。即使检查通过，也不能据此认定低秩模型能替代任意密集模型。

09.4 跨层共享：消费者仍然提问，生产者负责写入

设只有 $P$ 层产生 KV，其余层引用生产者缓存。共享以后，每层仍用自己的当前隐藏状态形成 query，因此不同层可以向同一份历史提出不同的查询。缓存大小为 $2 B T P G d s$ 。这种方法减少层方向的独立历史表示；理论背景可见 Cross-Layer Attention。

教学脚本 MultiHeadAttentionWithKVSharing.forward 接收 shared_kv。它为 None 时执行 K/V 投影和追加，并返回 new_shared_kv；非空时只读传入张量，返回的新共享值为 None。GPTModel.forward 顺序遍历 ModuleList：前 $P$ 层独立生产，后面的层复用最近的生产者。读到这里，需要把具体共享规则和“跨层共享”这个总称区分开。这里没有让任意相邻层两两共享，也没有采用 Gemma 4 按局部/全局类型寻找生产者的规则；后者见第 21 章。

生产者交出的 KV 已经包含当前 token，消费者若再追加一遍，每步历史长度就会增长两次。共享还要求消费者只读：就地旋转或修改 key，会同时改变其他层正在读取的对象。本章小例检验三位消费者的 data_ptr 相同，再改变 query，确认共享缓存仍能产生不同输出。

教学实现保留消费者的 W_key/W_value 模块。它们在共享分支中不参与计算，但仍可能占参数内存。所以，检查实现时要分别问清：有没有计算该投影，有没有缓存该层 KV，以及有没有保留该参数。前两项省下来了，不代表参数内存也随之消失。上游测试分别验证所有层生产时与 MHA 一致、只有生产层留下缓存、估算器用生产层数计费；这比只检查输出形状更有针对性。

09.5 从注意力接回 GPT：哪些步骤没有改变

三个目录中的 LayerNorm 仍沿特征轴计算均值与总体方差，GELU 使用 tanh 近似，FeedForward 保留扩展再投回的两层结构。TransformerBlock.forward 依次执行归一化、注意力、残差，再归一化、前馈、残差。这些步骤共同构成 GPT 主干；更换 KV 机制时，可以沿着它们核对，避免把注意力之外的改动混进比较。

GPTModel 的 token/position embedding 相加，最后归一化并经输出头得到 $[B, T, V]$ logits。缓存版用 ModuleList 传递 use_cache，记录模型级绝对位置，并通过 reset_kv_cache 重置全部注意力层。generate_text_simple_cached 先处理整个 prompt，再仅处理新增 token；到上下文边界时要重建裁剪后的上下文，不能让绝对位置 embedding 越界。CLI 的 main、parse_args 负责组织运行，计时和设备峰值内存报告帮助观察执行成本。这些输出都不能用来判断模型的语言质量。目录里的 MHA 对照文件保持相同主干，差别集中在投影、缓存与掩码。

memory_estimator_* 的字节函数、参数校验、单位换算和 main 把 09.1—09.4 的式子变成命令行输出；plot_memory_estimates_* 在不同 $T$ 上重复计算这些式子并画线。曲线因此反映公式给出的解析估算，不能当作运行模型采集的测量。本书的图按公式重新绘制，不复用上游 webp。

09.6 可复现实验、练习与答案

完整可运行程序是 examples/architectures/architecture_checks.py，运行 python examples/architectures/architecture_checks.py。原稿注明的运行环境为 Python 3.12.14、PyTorch 2.6.0+cpu，seed 1949、单线程、float64，并将结果记在日志 qa/architectures/algebra_checks.json 中：GQA 分组与展开误差为 0，错误 repeat 被检出；MLA 两项最大误差为 $1.78 \times 10^{- 15}$ ；共享相同查询的误差为 0。该程序不导入上游模型，也不下载权重。当前提供的文件包含检查程序，但缺少这份历史日志，因此上述精确误差只能作为原稿记录保留，不能视为已核验的运行结果。复现时应保存自己的完整输出，再对照每项检查的含义。

补充练习。 ① $H = 12, G = 3$ 时，query 头 7 应读取哪个 KV 头？② 为什么 MLA 的低维缓存与标准 MHA 输出不应直接比较？③ 所有层都设为生产者时，跨层共享模型应满足什么回归测试？④ 为什么显式复制 KV 可能抵消部分 GQA 工作内存收益？

提示与答案。 ① 零基编号， $r = 4$ ， $⌊ 7 / 4 ⌋ = 1$ 。② 等价性需要两边用相同的因子化投影，任意 MHA 参数通常不满足秩约束。③ 同权重、关闭 dropout、相同位置和 mask 时，完整前向及缓存生成均应与 MHA 一致。④ 缓存可以保留 $G$ 个头，但 repeat_interleave 的临时张量恢复了 $H$ 个头；专用内核可避免该展开，教学实现未承诺做到。

若要从代数例子继续检查真实模型类，可以使用 source_attention_checks.py。它导入已审查的 GQA、MLA 和 KV 共享脚本，但不运行 CLI。原稿报告，GQA 与 MLA 的完整前向和分块计算误差均为 0；KV 共享三层模型的误差为 $1.49 \times 10^{- 7}$ ，且消费者没有独立缓存。原稿对应的日志是 qa/architectures/source_attention_checks.json，当前文件中同样缺少这份历史记录。源码可以确认检查对象和断言，却不足以核实当时的精确输出。

参考资料与边界

固定源码：前述三个目录，逐函数映射见配套覆盖台账。
Ainslie 等，GQA（2023）：https://arxiv.org/abs/2305.13245。
DeepSeek-AI，DeepSeek-V2（2024）：https://arxiv.org/abs/2405.04434。
Brandon 等，Cross-Layer Attention（2024）：https://arxiv.org/abs/2405.12981。
未运行：大模型生成、GPU 内存/吞吐比较、从 MHA checkpoint 继续训练。

第 09 章 GQA、MLA 与跨层 KV 共享 ​

09.1 先把缓存账算在正确的对象上 ​

09.2 GQA：多个问题，可以共用一套历史记录 ​

09.3 MLA：先存一个低维向量，需要时再展开 ​

09.4 跨层共享：消费者仍然提问，生产者负责写入 ​

09.5 从注意力接回 GPT：哪些步骤没有改变 ​

09.6 可复现实验、练习与答案 ​

参考资料与边界 ​