从零手搓大模型（深挖细节GQA ）现代大模型为什么不用普通 MHA

GQA，全称Grouped-Query Attention，是现代 LLM 非常常见的 attention 结构。Llama、Qwen、Gemma 等模型都大量使用它。

一句话理解：

Query head 仍然很多，但 Key/Value head 变少，让多个 query head 共享同一组 K/V。

这样可以明显减少 KV cache 显存占用，同时保持接近 MHA 的建模能力。

1. 从 MHA 说起

第 2 章的 Multi-Head Attention 里，每个 head 都有自己的：

Q, K, V

如果有 32 个 attention heads，那么通常也有 32 组 K 和 32 组 V。

这叫 MHA：

Multi-Head Attention

问题是：推理时 KV cache 要保存所有层、所有 token、所有 head 的 K/V。

长上下文下，KV cache 会非常大。

2. GQA 的核心想法

GQA 做了一件很聪明的事：

保留多个 query heads 减少 key/value heads 让多个 query heads 共享同一组 K/V

例如：

num_heads = 6 num_kv_groups = 3

那么：

head 1, 2 共享 KV group 1 head 3, 4 共享 KV group 2 head 5, 6 共享 KV group 3

如果：

num_heads = num_kv_groups

就退化成普通 MHA。

如果：

num_kv_groups = 1

就是 MQA，也就是 Multi-Query Attention，所有 query heads 共享一组 K/V。

3. 为什么 GQA 省显存

KV cache 大小主要和n_kv_heads有关。

公式可以理解成：

KV cache bytes ≈ batch_size × sequence_length × head_dim × num_layers × 2 # K 和 V × bytes_per_elem × n_kv_heads

普通 MHA：

n_kv_heads = n_heads

GQA：

n_kv_heads = n_kv_groups

所以只要n_kv_groups比n_heads少，KV cache 就会明显变小。

例子：

n_heads = 32 n_kv_groups = 8

KV cache 大约少 4 倍。

4. 代码入口：GroupedQueryAttention

在gpt_with_kv_gqa.py里，核心类是：

classGroupedQueryAttention(nn.Module):

初始化里有两个关键断言：

assertd_out%num_heads==0assertnum_heads%num_kv_groups==0

第一个保证总维度可以平分给多个 query head。

第二个保证 query heads 可以平分给多个 KV groups。

5. MHA 和 GQA 的投影层差异

普通 MHA 的 Q/K/V 通常都是：

W_query:d_in->d_out W_key:d_in->d_out W_value:d_in->d_out

GQA 里：

self.W_query=nn.Linear(d_in,d_out,...)self.W_key=nn.Linear(d_in,num_kv_groups*head_dim,...)self.W_value=nn.Linear(d_in,num_kv_groups*head_dim,...)

也就是说：

Q 还是完整的 num_heads K/V 只生成 num_kv_groups 份

这就是 GQA 的关键。

6. shape 变化

假设：

batch = b num_tokens = T num_heads = H num_kv_groups = G head_dim = D

query：

queries.shape = (b, H, T, D)

key/value 在展开前：

keys.shape = (b, G, T, D) values.shape = (b, G, T, D)

因为G < H，所以 K/V 比 MHA 少。

7. repeat_interleave 的作用

attention 计算时，query heads 有 H 个，但 key/value groups 只有 G 个。

所以代码会把 K/V 扩展到 H 个 head：

keys=keys_base.repeat_interleave(self.group_size,dim=1)values=values_base.repeat_interleave(self.group_size,dim=1)

其中：

self.group_size=num_heads//num_kv_groups

例如：

num_heads = 6 num_kv_groups = 3 group_size = 2

K/V 会从：

[K1, K2, K3]

变成：

[K1, K1, K2, K2, K3, K3]

这样每个 query head 都有对应的 K/V 可以做 attention。

8. GQA 的 attention 计算仍然一样

扩展后，attention 公式没有变：

attn_scores=queries @ keys.transpose(2,3)attn_weights=torch.softmax(attn_scores/keys.shape[-1]**0.5,dim=-1)context_vec=(attn_weights @ values).transpose(1,2)

最终输出仍然是：

(batch, num_tokens, d_out)

所以 GQA 对后面的 TransformerBlock 来说是可替换的。

9. GQA 和 KV cache 的关系

GQA 在训练时也减少参数和计算，但它在推理阶段最明显的收益来自 KV cache。

因为 cache 保存的是 K/V。

普通 MHA：

每个 attention head 都要存 K/V

GQA：

多个 query heads 共享少量 K/V groups

因此长上下文推理时，GQA 可以显著减少显存。

10. memory_estimator_gqa.py 怎么理解

这个脚本用来估算 MHA 和 GQA 的 KV cache 显存差异。

README 示例中：

uv run memory_estimator_gqa.py\--emb_dim4096--n_heads32--n_layers32\--context_length32768--n_kv_groups4\--batch_size1--dtypebf16

如果n_heads=32，n_kv_groups=4，则 K/V head 数减少 8 倍。

在长上下文下，这种差异非常大。

11. GQA 的代价

GQA 不是无脑越少越好。

如果num_kv_groups太少，很多 query heads 共享同一组 K/V，模型表达能力可能下降。

所以真实模型通常取中间值：

MHA: n_kv_groups = n_heads GQA: 1 < n_kv_groups < n_heads MQA: n_kv_groups = 1

GQA 是在质量和效率之间折中。

12. 和 Llama/Qwen/Gemma 的关系

这个 bonus 虽然是在 GPT 架构上改 GQA，但真实现代模型里，GQA 已经很常见。

你后面看：

• Llama 3
• Qwen3
• Gemma
• 推理框架的 KV cache 显存优化

都会看到类似设计。

所以这个 bonus 是从“教学版 GPT”过渡到“现代 LLM 架构”的重要桥梁。

13. 建议

推荐这样看：

1. 理解 GQA 省的是 K/V。
2. 对比gpt_with_kv_mha和gpt_with_kv_gqa。
3. 重点看GroupedQueryAttention。
4. 盯住三个参数：

num_heads num_kv_groups group_size

5. 跑或阅读memory_estimator_gqa，感受上下文长度变大后的显存差异。

一句话总结：

GQA 保留多个 query heads，但减少 key/value heads，用较小质量代价换取明显的 KV cache 显存收益。