写完 Multi-Head Attention 和 Causal Attention 之后，我一直觉得 Attention 这条线已经差不多了——公式拆过了，routing space 的 factorization 也想明白了，Causal Mask 的因果约束也理清了。
直到最近看 Llama 2 的模型配置，发现一行让我停下来的参数：
1num_attention_heads: 64
2num_key_value_heads: 8Query Heads 64 个，KV Heads 只有 8 个？这两个数不一样？
我之前写 MHA 的时候，脑子里的默认假设一直是：每个 Query Head 对应一个独立的 KV Head，\(H_Q = H_{KV}\)。这个假设突然不成立了。
顺着这条线查下去，发现这个设计叫 Grouped-Query Attention（GQA），而它的极端形式——所有 Query Heads 共享同一个 KV Head——叫 Multi-Query Attention（MQA）。GQA 介于 MHA 和 MQA 之间。
让我花了时间想的问题是：如果多个 Query Heads 共享同一份 K/V，不同 Head 的 Attention Distribution 还能不同吗？不同的话，为什么不全用 MQA？
这篇把这条线理一遍。
MHA 的核心结构：简短回顾#
这部分不重新推导，QKV 那篇和 MHA 那篇已经写得比较完整了。这里只摘出后面需要用到的几个要点。
Multi-Head Attention 的设计动机：单个 Attention Head 只有一套 Q/K matching space，只能形成一个 routing distribution。但一个 token 和上下文之间可能存在多种不同性质的关联——局部搭配、远距离依赖、句法结构。把这些全压进同一个 matching space 要求不同的 relational needs 互相妥协。MHA 拆出多个独立的 Q/K/V projection，让每个 Head 在自己的子空间里独立 route、独立 read。
一个上次写 MHA 时提过、这里需要再强调的细节：每个 Head 的投影是从完整的 \(d_{\text{model}}\) 维输入学习一个 \(d_{\text{head}}\) 维投影，不是对原始 hidden state 做机械切片。 经典配置 \(d_{\text{model}} = 512\)，\(h = 8\)，\(d_{\text{head}} = 64\)。8 个 Head 的 score computation 总量 \(8 \times L^2 \times 64 = L^2 \times 512\)，和一个 512 维 single-head 同数量级。MHA 用 representation factorization 换 relational diversity，不是粗暴叠计算量。
到这里为止，MHA 的一个隐含假设是：Query Heads 的数量和 KV Heads 的数量相等。 每个 Head \(i\) 的计算形如：
$$ O_i = \operatorname{Attention}(Q_i, K_i, V_i) $$8 个 Query Heads，8 套独立的 K/V projection，8 份独立的 K/V 数据。
训练的时候这不太引人注目。但推理的时候，这 8 份 K/V 数据的存储和读取开始变成一个具体的资源问题。
自回归推理的 KV Cache：为什么历史 K/V 需要保存#
Causal Attention 那篇的核心结论是：Causal Mask 保证位置 \(i\) 只能看到 \(j \le i\) 的 token，历史 token 的表示不受未来 token 影响。
这个性质在推理阶段有一个直接的工程含义：如果历史 token 的 K 和 V 不会因为新 token 的加入而改变，那它们算过一次之后就可以缓存起来，不用每生成一个新 token 都重新算整个 sequence 的 K/V。
这就是 KV Cache。
具体来说，自回归推理分成两个阶段：
Prefill：用户输入 Prompt（比如 4 个 token），模型做一次完整的 forward pass。每一层 Transformer 计算全部 token 的 Q、K、V，完成 Attention，并把这一层的 K 和 V 保存下来。Prefill 结束时，最后一个位置的输出预测出第一个生成 token。
Decode：逐 token 生成。每一步只有一个新 token 输入，模型只需要计算这个新 token 的 Q、K、V。新 token 的 Q 需要和这一层所有历史位置的 K 做 matching，得到 routing weights，然后从所有历史位置的 V 做 weighted read。计算完成后，新 token 的 K 和 V 追加到这一层的 Cache 中。
有几个容易忽略的细节：
- 每一层 Transformer 都有自己独立的 KV Cache。 不是整个模型共享一份。一个 32 层的模型就有 32 份独立的 KV Cache。
- 新 token 的 Q 仍然需要重新计算。 缓存的只是历史的 K 和 V，不是 Q。Q 描述的是"当前 token 需要什么"，每个新 token 的需求不同，无法复用。
- Prefill 最后一个位置预测出的第一个生成 token，它自己的 K/V 通常要等到下一次 Decode forward 时才产生。 不要把"刚被预测出来但还没进入下一轮 forward 的 token"误认为已经在 KV Cache 里了。
- 为什么缓存有效？ 因为 Causal Attention。历史 token 在位置 \(j\) 的 \(K_j\) 和 \(V_j\)，是由 \(j\) 位置的输入和该层的权重决定的，不依赖任何 \(j\) 之后的 token。新 token 的加入不会改变历史位置的 K 和 V。如果用的是 RoPE 这类位置编码，K 中已经嵌入了该 token 的绝对位置信息，缓存后的 K 和固定位置对应，不需要重新编码。
Decode 阶段的问题就变得明确了：每生成一个 token，模型需要在每一层读取该层所有历史位置的 K 和 V。随着生成长度增长，需要读取的历史 KV 数据量持续增加。
传统 MHA 下，每一层的 KV Cache 包含所有 KV Heads 的数据。KV Heads 越多，每一层的 Cache 越大，每一步 Decode 需要读取的数据越多。
这就引出了一个具体的问题：这些 KV Cache 到底有多大？
算一笔 KV Cache 的显存账#
之前写过的显存心智模型讨论过推理时的显存组成，但没有把 KV Cache 单独拎出来算。这里补上。
KV Cache 的总大小由一个公式决定：
$$ M_{KV} = 2 \times N_L \times B \times S \times H_{KV} \times D_H \times b $$各项含义：
- \(N_L\)：Transformer 层数
- \(B\)：Batch Size
- \(S\)：缓存的序列长度（Context Length）
- \(H_{KV}\)：KV Head 数量
- \(D_H\)：每个 Head 的维度
- \(b\)：每个元素的字节数（BF16 = 2 bytes）
- \(2\)：Key 和 Value 各一份
用一个具体的配置算：
1Layers = 32
2Q Heads = 32
3KV Heads = 32 (MHA)
4Head Dim = 128
5Seq Length = 8192
6Batch Size = 1
7Precision = BF16 (2 bytes)4 GiB。而这仅仅是 KV Cache 本身，不包含模型权重（这个配置下大约 7-8B 参数，BF16 约 14-16 GiB）、Activation workspace、Runtime 开销等。
换句话说，一个 7B 左右的模型在 8K context 下，KV Cache 就要占掉接近模型权重 1/4 的显存。如果 context 拉到 32K，KV Cache 直接翻四倍到 16 GiB。如果在线服务需要同时处理多个请求（\(B \gt 1\)），KV Cache 还要乘 batch size。
这个公式里有两个维度在增长时特别让人不安：\(S\)（context 越长 cache 越大）和 \(H_{KV}\)（KV Head 越多每层 cache 越大）。\(S\) 是用户需求驱动的，很难压缩。那 \(H_{KV}\) 呢？
如果让所有 Query Heads 共享同一份 K/V？#
这就是 Multi-Query Attention（MQA）的出发点，由 Noam Shazeer 在 2019 年的论文 Fast Transformer Decoding: One Write-Head is All You Need 中提出。
MQA 的改动很简单：保留所有 Query Heads，但只用一个 KV Head。所有 Query Heads 共享同一份 K 和 V：
$$ O_i = \operatorname{Attention}(Q_i, K, V) $$和 MHA 对比：
$$ \text{MHA: } O_i = \operatorname{Attention}(Q_i, K_i, V_i) \quad \text{（每个 Head 独立的 K/V）} $$$$ \text{MQA: } O_i = \operatorname{Attention}(Q_i, K, V) \quad \text{（全部共享一份 K/V）} $$我第一次看到这个设计时的直觉反应是：K/V 都共享了，不同 Head 的 Attention 结果岂不是一样的？
但仔细想了想，发现不对。Attention Distribution 的差异来自 Q 和 K 的 matching，不只是来自 K。不同 Head 有不同的 \(Q_i\)（因为各自有独立的 \(W_Q^{(i)}\)），即使面对同一个 K，它们产生的 routing weights 也不同：
$$ \operatorname{softmax}(Q_0 K^T) \neq \operatorname{softmax}(Q_1 K^T) $$通常成立，因为 \(Q_0 \neq Q_1\)。
用一个直观的例子：假设 sequence 有 4 个位置，K 是共享的。Head 0 的 Query 更关注语法结构，它的 Q 和位置 1、2 的 K 的 compatibility 更高；Head 1 的 Query 更关注语义相关性，它的 Q 和位置 3、4 的 K 的 compatibility 更高。两个 Head 面对同一组 Key，生成了完全不同的 attention weight 分布，从 Value 中读取了不同的加权组合。
MQA 共享的是 Key/Value 的表示，不是 Attention Distribution。 多套 Attention Distribution 仍然存在，因为 Q 仍然是多套独立的。
然后 MQA 在资源上的收益就很直观了：
- KV Cache 缩小到原来的 \(1/H_Q\)。 原来 32 个 KV Heads 变成 1 个，KV Cache 从 4 GiB 降到 128 MiB。
- K/V Projection 的参数量和计算量减少。 原来需要 32 套 \(W_K\) 和 \(W_V\)，现在只需要 1 套。
- Decode 阶段读取历史 KV 的数据量减少。 每一步 Decode 只需要从 HBM 读取一份 K/V，而不是 32 份。
有一个需要区分的地方：逻辑上 32 个 Query Heads 都在"使用"同一份 K/V，但物理上 GPU 是否只从 HBM 读取一次，取决于具体的 Kernel 实现。如果 32 个 Head 的 Attention 计算被分配到不同的 CTA（Cooperative Thread Array），每个 CTA 可能各自从 HBM 读取一次 K/V，也可能通过 L2 Cache 或 Shared Memory 实现数据复用。之前写 GPU 存储层级的时候讨论过 HBM 和片上存储的关系——逻辑上的共享和物理上的数据搬运是两个层面的事情。
不过，即使不考虑 Kernel 实现的细节，KV Cache 的总量确实从 32 份降到了 1 份，这个存储层面的缩减是确定的。
为什么不是所有模型都用 MQA#
MQA 把 KV Cache 压到了极致。一个自然的问题：既然效果这么好，为什么不让所有 LLM 都用 MQA？
我后来想明白的一点是：MQA 的 KV 共享不是没有代价的。
MHA 中，不同 Head 各自有独立的 \(W_K^{(i)}\) 和 \(W_V^{(i)}\)，可以学到不同的 K/V 表示。Head 0 的 Key 空间可以强调语法信息，Head 1 的 Key 空间可以强调语义信息。这些不同的表示给模型提供了更多的 relational structure 自由度。
MQA 把所有这些压成了一个 Key 空间和一个 Value 空间。所有 Query Heads 的 matching 都发生在同一个 K 空间里，所有 read 都从同一个 V 空间取数据。这等于把 K/V 侧的表示多样性完全去掉了。
GQA 论文（Ainslie et al., 2023, GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints）报告了一个经验性的观察：在部分任务上，MQA 相比 MHA 确实存在质量损失。具体的影响大小依赖模型规模、训练数据和下游任务，不能一概而论——但方向上，把 KV 表示自由度压缩到极致确实有模型能力的 trade-off。
这就引出了 Grouped-Query Attention（GQA）的设计空间。
GQA 不是一个全新的 architecture，而是在 MHA 和 MQA 之间的连续谱上选择一个中间点。做法很直接：把 Query Heads 分成若干组，每组共享一个 KV Head。
比如 32 个 Query Heads，8 个 KV Heads：每 4 个 Query Heads 共享同一个 KV Head。
GQA 减少的是 KV Heads 的数量，不是 Query Heads 的数量。 32 个独立的 Attention Distribution 仍然存在（因为 32 个 Query Heads 各自有独立的 \(W_Q\)），只是它们面对的 K/V 表示从 32 套减到了 8 套。
用一个统一的公式把三种架构联系起来：
$$ \text{GroupSize} = \frac{H_Q}{H_{KV}} $$- MHA：\(\text{GroupSize} = 1\)（每个 Q Head 有自己的 KV Head）
- GQA：\(1 \lt \text{GroupSize} \lt H_Q\)（每组 Q Heads 共享一个 KV Head）
- MQA：\(\text{GroupSize} = H_Q\)（所有 Q Heads 共享唯一的 KV Head）
GQA 不是数学上证明的最优点。它是一个设计参数——\(H_{KV}\) 可以根据模型规模和工作负载调整。Llama 2 70B 选了 \(H_Q = 64, H_{KV} = 8\)（GroupSize = 8），这是在质量和推理效率之间做出的一个具体的工程选择。
三种架构的统一对比#
沿用前面的配置（32 Layers，32 Q Heads，\(D_H = 128\)，\(S = 8192\)，\(B = 1\)，BF16），把三种架构的 KV Cache 放在一起：
| 架构 | Q Heads | KV Heads | GroupSize | KV Cache | 相比 MHA |
|---|---|---|---|---|---|
| MHA | 32 | 32 | 1 | 4 GiB | 基准 |
| GQA-8 | 32 | 8 | 4 | 1 GiB | 减少 75% |
| MQA | 32 | 1 | 32 | 128 MiB | 减少 96.9% |
这张表最直观的信息是：KV Cache 的大小和 \(H_{KV}\) 成正比。 GQA-8 把 KV Heads 从 32 减到 8，KV Cache 从 4 GiB 降到 1 GiB；MQA 降到 1 个 KV Head，Cache 直接到 128 MiB。
下面这张图从视觉上呈现三者的规模差异：
需要强调的是：KV Cache 的缩减比例不能直接等同于推理延迟或吞吐量的改善比例。 KV Cache 变小意味着存储占用减少、Decode 阶段的 KV 读取量减少，但 Attention 的核心计算（\(QK^T\)、softmax、Attention \(\times\) V）并没有按同样的比例减少——Query Heads 数量没变，score computation 的总量取决于 \(H_Q\) 而不是 \(H_{KV}\)。实际的端到端性能还受 Kernel 实现、Batch Size、Context Length、FFN 等其他算子的影响。
三种架构的 Head 映射关系如下图所示。注意 Query Heads 的数量始终保持 8 个（这里用的是简化的 8 Head 示例），只有 KV Heads 的数量在变化：
几个值得记住的 Insight#
这篇文章的核心就一条线：MHA 最初把 Attention Distribution 的数量和 KV 表示的数量绑定在了一起；MQA 和 GQA 把它们解耦了。
更具体地拆成三点：
Attention 分布的数量与 KV 表示的数量可以独立设定。 MHA 的设计中，每个 Query Head 都配了一套独立的 K/V，这不是数学上的必然要求——多个 Q Heads 共享 K/V 之后仍然能产生不同的 routing weights，因为 Attention Distribution 由 Q 和 K 的 matching 共同决定，不是仅由 K 决定。MQA 和 GQA 利用了这个事实。
模型架构的设计需要考虑推理生命周期，不只是单次 forward 的表达能力。 传统 MHA 从表示能力的角度是自洽的——更多独立的 K/V Heads 意味着更多的 relational structure。但自回归推理把这个 one-shot 的计算变成了一个持续的过程：Decode 阶段每生成一个 token，模型就要在每一层读取所有历史位置的全部 KV Heads。KV Cache 不是一次性开销，而是一个随 context 增长、随并发请求增多而持续膨胀的资源。这使得 KV Head 数量成为一个值得单独优化的设计参数。
减少存储成本不等于按比例减少计算或延迟。 这一点在 Arithmetic Intensity 那篇里已经讨论过：Decode 阶段典型的 memory bound 特性意味着减少 KV 读取量确实能改善性能瓶颈；但 Prefill 阶段更接近 compute bound，KV Cache 的大小对 Prefill 延迟的影响就小得多。推理性能是一个多因素的系统问题，不是某一项资源缩减的线性映射。
后续还有几个相关但本文不展开的方向：
- FlashAttention 通过 tiling 和 online softmax 减少 Attention 计算中的 HBM IO，和 GQA/MQA 减少 KV 存储量是两个不同层面的优化，可以叠加。
- PagedAttention（vLLM）解决的是 KV Cache 的动态内存管理问题：多个请求的 KV Cache 长度不同、会动态增长，如何避免碎片化？
- Prefill 和 Decode 的 GPU 性能瓶颈为什么不同？前者更接近 compute bound，后者更接近 memory bound，两者对 KV Head 数量的敏感度也不一样。
参考资料#
- Vaswani et al., Attention Is All You Need, 2017 — https://arxiv.org/abs/1706.03762
- Shazeer, Fast Transformer Decoding: One Write-Head is All You Need, 2019 — https://arxiv.org/abs/1911.02150
- Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, 2023 — https://arxiv.org/abs/2305.13245
- Touvron et al., Llama 2: Open Foundation and Fine-Tuned Chat Models, 2023 — https://arxiv.org/abs/2307.09288
- 从动态读取到 QKV：Self-Attention 的公式是怎么长出来的
- 一组 QKV 就够了吗：Multi-Head Attention 到底在拆什么
- 为什么 Self-Attention 不能偷看未来？从 Causal Mask 到自回归训练
- FlashAttention 的核心思想：为什么 Attention Matrix 根本不需要存下来？