每个头都存一份 K/V 太贵,怎么办?
本课一问:每个头都存一份 K/V 太贵,怎么办?
这一课先不碰:不存 K、V 本身的做法(第 10 课);服务端的缓存管理(C13);质量评测怎么做才可信(C14)。
从已有出发
上一课的账:每个 token 的缓存 = 全注意力层数 × 2 × KV 头数 × 每头维数 × 字节数。层数和每头维数由模型能力决定,上下文长度由用户决定;剩下能动的,就是「KV 头数」。第 7 课的多头注意力里,每个 Q 头都有自己的一份 K、V——锚点 Qwen3.8 有 24 个 Q 头,可它的 config 里 KV 头只有 4 个。中间差的这些,就是本课的故事。
先自己设计:每个头各存一份(MHA)
按第 7 课原样实现:24 个头,每个头的 K、V 都单独投影、单独缓存。这叫多头注意力(MHA),GPT-2、GPT-3 都这么做。
- ① 朴素方案
- ② 它在哪失败
- ③ 改进方案
问题不只是放不放得下。生成时每一步都要把整份缓存从显存里读一遍,只为算一个 token(上一课的 Dig):缓存越大,每一步越慢,一张卡能同时服务的请求越少。
MQA(Shazeer 2019):所有 Q 头共用 1 份 K、V。缓存直接除以头数:每 token 16 KiB,256K 上下文 4 GiB。原论文的结论是逐 token 解码所需的访存大幅下降,质量只有小幅下降 证据等级 · 实验验证 由对比实验得出的经验结论,只在出处写明的模型、数据和规模下成立。换了条件可能不再成立,引用时要带上实验设置。 本条原文:incur only minor quality degradation from the baseline;当时的翻译与语言模型实验 出处Shazeer 2019 摘要 ↗需联网 ;但后来的工作指出 MQA 可能带来质量下降和训练不稳定 证据等级 · 实验验证 由对比实验得出的经验结论,只在出处写明的模型、数据和规模下成立。换了条件可能不再成立,引用时要带上实验设置。 本条原文:MQA can lead to quality degradation and training instability 出处Ainslie et al. 2023 §1 ↗ 。
GQA(Ainslie 等 2023):折中——把 Q 头分成 g 组,每组共用 1 份 K、V。锚点 Qwen3.8 就是 g = 4:每 6 个 Q 头一组,每 token 64 KiB,256K 上下文 16 GiB。
试三件事:① 在分组图里看清「组数 = KV 头数」:选 GQA 时 24 个 Q 头被涂成 4 种颜色,每种颜色连向同一个 KV 头;② 把上下文拖到原生长度 262,144,对比三种方案的总量,再拖到约 1M(需要开 YaRN)看会翻几倍;③ 用「自定义 g」把 KV 头数从 1 调到 24,看缓存怎么随 g 线性变化——想一想,g 选多少才算划算?答案不在公式里,在实验里(下面)。
改进:分组共用 K、V
GQA 的做法:K、V 的投影只输出 g 个头;计算注意力时,第 h 个 Q 头使用它所在组的那一份 K、V。实现上就是把 K、V 沿头的维度复制成 24 份再照常计算——我们项目里的 repeat_interleave 就是这一步;缓存里存的是复制之前的 g 份。
| 维度 | 含义 | 真实规模 | 小规模 |
|---|---|---|---|
| H | Q 头数 num_attention_heads | 24 | 4 |
| G | KV 头数 = 组数 num_key_value_heads | 4 | 2 |
| T | 序列长度(贯穿例句) | 5 | 5 |
| dh | 每头维数 head_dim | 256 | 3 |
它省了什么、没省什么。缓存与每步读取的数据量除以 6(锚点);K、V 的投影参数也跟着变小——k_proj、v_proj 各是 5,120 × 1,024,而 q_proj 的查询部分是 5,120 × 6,144。没省的是打分和加权求和的计算量:每个 Q 头照样要和 T 个 k 打分 证据等级 · 数学可证明 在写明的前提下可以严格推导出来,不依赖实验。只要前提成立,换数据、换规模仍然成立。 本条字节与参数都是计数;GQA 等价于「组内复制 K/V 权重的 MHA」,项目测试 test_gqa_equals_mha_with_shared_kv_weights 断言了这一点 出处待补 。
质量呢?这要靠实验。GQA 论文把 T5-XXL 从 MHA 改成 MQA 和 GQA-8(8 组),只用原预训练 5% 的算力继续训练,七个任务的平均分:MHA 47.2、GQA-8 47.1、MQA 46.6;表中的推理时间:1.51 秒、0.28 秒、0.24 秒。GQA-8 的质量接近 MHA,速度接近 MQA 证据等级 · 实验验证 由对比实验得出的经验结论,只在出处写明的模型、数据和规模下成立。换了条件可能不再成立,引用时要带上实验设置。 本条T5 编码器-解码器模型、摘要/翻译/问答任务;不是今天的纯解码器大模型 出处Ainslie et al. 2023 表 1 ↗ 。这是 GQA 在 2023–2025 年成为开源大模型默认做法的理由;各家选多少组,没有公式,是按规模和预算实验出来的。到 2026 年,前沿又分出了几条路:不存 K、V 本身(第 10 课 MLA),或者干脆换成不随长度增长的状态、只读一部分缓存(C11)。
我们项目里的对照实验:同样的数据、步数和种子,只改 n_kv_head,看验证 loss 与每 token 缓存字节怎么变(数据由项目导出,在上面计算器的最后一栏显示)。按 2026-09-26 导出的 5 个种子均值:KV 头 8 / 4 / 2 / 1 分别是 4.6247 ± 0.0104、4.6289 ± 0.0153、4.6332 ± 0.0203、4.6300 ± 0.0128,和 kv = 8 的差距只有 +0.004 到 +0.009,全都小于 3σ 阈值 0.020,也小于一个标准误。单种子时看起来「KV 头越少 loss 越低」,补到 5 个种子后这个顺序消失了——kv = 8 的均值反而最低。能说的只有:缓存从每 token 8,192 字节降到 4,096 / 2,048 / 1,024 字节,在这个规模上没测出质量变化。读这张表时先看组件给的「重训波动」:同一个配置重训一次,val loss 也会差一点;几组之间的差距如果和这个波动差不多大,就只能说「没看到质量变化」,不能说「KV 头越少越好」 证据等级 · 实验验证 由对比实验得出的经验结论,只在出处写明的模型、数据和规模下成立。换了条件可能不再成立,引用时要带上实验设置。 本条5 个种子、400 步、3.7M 参数的小模型实验(c3-attention.json › seed_summary);结论只在这个规模上成立 出处待补 。
⚠ 常见误解
「GQA 的『组数』和 KV 头数是两回事,比如 8 组 = 4 个 KV 头。」旧版课程就这样写错过。一组共用一个 KV 头,所以组数就等于 KV 头数:Qwen3.8 是 4 组 = 4 个 KV 头,每组 6 个 Q 头;GQA 论文里的 GQA-8 就是 8 个 KV 头。
「GQA 让注意力算得更快,因为算得更少。」打分的乘法次数一个没少。它快,是因为生成时每一步要从显存读的 K、V 少了——解码阶段的瓶颈在读数据,不在计算(C13 用屋顶线细讲)。
注意力的演进,你在这里
① 训练时:一个 token 怎么从别的 token 那里取信息
- 解决了能读任意长的序列,参数与长度无关;手里只带一个固定大小的状态 h。卡在ht 要等 ht−1 算完:T = 4,096 时每层 4,096 步不能并行,第 1 个 token 的信息要经过 4,095 次更新才到最后;翻译时整句还得挤进一个定长向量——Bahdanau 2014 说的瓶颈。展开岔路 B1
- 解决了解码器每一步回头给源句的每个位置打分、加权取回,不再只靠那一个定长向量:score = vT·tanh(W1q + W2k)。卡在每一对都要单独过 tanh,交不给矩阵乘单元;朴素实现要先落一个 T×T×da 的中间张量:T = 4,096、da = 256 时单头单层 8.0 GiB(BF16),点积只留 T×T 个分数,32 MiB。它还长在 RNN 翻译模型里,串行没解决。展开第 2 课
- 解决了直接比 embedding 时「给了」最像的是它自己(余弦 1.000),小明、苹果、小王排在最后;投影成 Q、K 再点积,它自己排到最后、前三名是苹果、小王、小明,全部配对一次矩阵乘 QKT;取回用单独的 V——两个人名的 K 余弦 0.995,取回 K 分不清是谁。卡在不缩放:锚点每头 256 维,点积标准差 √256 = 16,同一行分差十几二十是常态(e20 ≈ 4.9 亿),softmax 几乎独占,梯度传不回 WQ、WK。
- 解决了softmax 把分数变成正的、和为 1 的权重,只看分差(按占比分,对话样本的 10、3、−4 会给「汽车」−44.4%);除以 √dk = 16,点积方差回到 1。卡在训练时整句一次喂进去:不加掩码,「给了」一行有 32.5% 的注意力直接落在它要预测的答案「小王」上,loss 低得虚假。展开第 3 课
- 解决了旋转后 q·k 只依赖距离 j − i;V 不参与打分,不转(转了会把绝对位置混进取回的内容)。锚点每头 256 维里只转 64 维。卡在旋转矩阵随位置变,夹在展开矩阵和缓存之间就没法事先并进 Q:MLA 因此把 RoPE 单独拆出 64 维另存(DeepSeek-V2 config,第 10 课)。
- 解决了标准 softmax 逼每行分完 100%:「苹果」一行 3 个分数只差零点几(0.04、0.36、0.00),照样分成 30% / 41% / 29%;分母加可学习的 sink logit,或在输出上乘门,这个头就能「这次不取」——锚点开着输出门(attn_output_gate = true)。卡在出口的效果只有实验证据:Gated Attention 原文说门控「缓解」sink,不是消除;DeepSeek-V4 没给 sink logit 的单独消融。展开岔路 B4
② 生成时:K、V 怎么越存越少
- 解决了缓存直接除以头数:锚点形状下每 token 16 KiB,262,144 上下文 4 GiB。卡在质量会掉、训练可能不稳:GQA 论文表 1(T5-XXL,七个任务平均)MHA 47.2 → MQA 46.6。展开第 9 课
- 解决了24 个 Q 头分 4 组、每组 6 个共用 1 个 KV 头(组数 = KV 头数):每 token 64 KiB、262,144 上下文 16 GiB;T5-XXL 上 GQA-8 平均 47.1、推理 0.28 s,质量接近 MHA(47.2、1.51 s),速度接近 MQA(0.24 s)。卡在再减 KV 头就回到 MQA:DeepSeek LLM 67B 已是 GQA-8,每 token 仍要 389,120 B(每头一份的 MHA 是 3,112,960 B),1M token 单条请求 380 GiB。展开第 9 课
- 解决了不存 K、V,每层每 token 只存 512 维潜向量 + 64 维 RoPE = 576 个数,用时再展开:DeepSeek-V2(60 层)每 token 69,120 B、V3(61 层)70,272 B;相对 67B 按 BF16 降 82.2%,官方「−93.3%」还含约 6 bit 的 KV 量化。卡在仍随长度线性增长,打分仍是 T²:GLM-5.3 同是 576 个数、78 层,每 token 89,856 B,原生 1,048,576 上下文时单条请求只算潜向量就要 87.75 GiB。展开第 10 课
- 解决了换掉「按长度存」本身:锚点 64 层里 48 层不存 KV,只存固定大小的状态(FP32 共 144 MiB),与上下文长度无关(这些层是什么、怎么更新,C11 第 4 课揭晓);DeepSeek-V4.1 叠加压缩、稀疏与低精度,官方报告全局 KV 890 B/token。卡在固定状态装不下任意多对键值,锚点仍留 16 层全注意力;稀疏注意力每步只读挑出来的一部分,挑本身也有代价——各自省了什么、没省什么,C11 讲。
证据三问
- 它解决什么问题?在几乎不掉质量的前提下,把 KV 缓存和每步的读取量降到 MHA 的几分之一。
- 数学上为什么可能有效?字节账是计数 证据等级 · 数学可证明 在写明的前提下可以严格推导出来,不依赖实验。只要前提成立,换数据、换规模仍然成立。 本条组件现算 出处待补 ;「几个 Q 头共用一份 K、V 不怎么伤表达」是理论动机——不同的头仍然有不同的 Q,可以对同一组 K 提出不同的问题 证据等级 · 有理论动机 有数学或机制上的理由说明它可能有效,但没有证明它一定更好。最后是否采用,仍然要看实验。 本条没有证明给出最优组数 出处待补 。
- 实验上为什么选了它?GQA-8 在 T5-XXL 上质量接近 MHA、速度接近 MQA 证据等级 · 实验验证 由对比实验得出的经验结论,只在出处写明的模型、数据和规模下成立。换了条件可能不再成立,引用时要带上实验设置。 本条见上 出处Ainslie et al. 2023 表 1 ↗ ;之后各家大模型沿用,组数各不相同(Qwen3.8 是 4,Qwen3-32B 是 8)。
一句话
GQA 让几个 Q 头共用一份 K、V:组数 = KV 头数;省下的是缓存和读取,不是计算。
下一个缺口
沿着这条路继续减 KV 头,尽头就是 MQA,质量会掉。有没有另一个方向:K、V 本身能不能不存,只存一个更小的东西,用的时候再还原?→ 第 10 课:KV 还能再压吗?
你可能会问
这一课还没有收到真实提问;你的提问会记到 05-问答记录/C3-L8.md。
🔍 已经训练好的 MHA 模型,能改成 GQA 吗?
能。GQA 论文的办法叫 uptraining:把每组里各头的 K、V 投影矩阵取平均(论文比较过,平均比随便挑一个头、或重新随机初始化都好),得到 g 个 KV 头,再用原预训练约 5% 的算力继续训一段,让模型适应新结构。上面引用的 T5-XXL 数字就是这样来的。
🔍 Qwen3.8 为什么选 4 组,不是 8 组或 2 组?
官方没有公开这个选择的消融。能确定的是:同一代的 Qwen3.8-Flash-Next 只用了 2 个 KV 头(调研档案 R2 §1.3.2),而上一代 Qwen3-32B 是 8 个——组数是能力、上下文长度和服务成本之间的权衡,由各家的实验决定。项目里的 kv_heads_ablation 就是让你在小模型上亲手看这笔权衡。