现代 LLM 的核心架构设计其四:GQA
多头注意力 MHA#我们在前面展开过标准 Transformer 使用多头注意力机制 个注意力头各自拥有独立的 Q、K、V 投影矩阵而其中每个头独立学习不同的注意力模式。最终的输出是 个头的拼接这本身是为了增加表达能力的合理设置但 KV Cache 出现后KV Cache 需要为每个头单独存储一份 K 和 V。这一结构设计带来了较大的内存压力。2. 多查询注意力 MQA#19 年Shazeer就是 SwiGLU 那位在 Fast Transformer Decoding: One Write-Head is All You Need 提出了一个激进方案即多查询注意力Multi-Query AttentionMQA。在 MQA 中 个 Query Head 共享同一组 K 和 V只有一个 K 头和一个 V 头意思是无论有多少个 Query Head它们查的都是同一份 K 和 V。这样KV Cache 的大小瞬间降到 MHA 的 。对于 64 头的模型直接省了 98.4% 的 KV Cache 内存。但代价也很明显不同 Query Head 已经被证明会关注不同模式把它们绑定到同一份 K、V 上必然损失表达能力。实验结果也印证了这一点MQA 的训练更不稳定在质量敏感的任务上效果有明显下降。3. 分组查询注意力 GQA#目前的主流方案来自 23 年的论文 GQA: Training Generalized Multi-Query Transformer for Multi-Head Attention 它其实更像是前两个方案的折中把 个 Query Head 分成 组每组共享一个 K 头和一个 V 头。 是一个可调参数。这其实是把质量与效率的权衡变成了一个连续可调的超参数你想省多少显存就设置多少组。举个例子假设 那么其对应关系即如下Query Head使用的 KV HeadQ₀KV₀Q₁KV₀Q₂KV₁Q₃KV₁Q₄KV₂Q₅KV₂Q₆KV₃Q₇KV₃于是注意力实际上是这样的然后所有组的输出再拼接到一起看得出来GQA 的改动非常小它只改变了 K 和 V 的投影矩阵列数简单对比如下MHAGQAGQA 的 K、V 列数从 缩小到 Q 保持不动。这意味参数量节省了 同时 KV Cache 也相应缩小。而从实现角度看现代框架通常不会真的复制 K、V。而是在进入注意力计算前先针对头索引 构造一个映射然后计算时直接索引这样来实现只共享内存不会真的复制数据。4. 大模型中的实际配置#GQA 在提出之后迅速成为主流方案。如今绝大多数开源大模型都已经放弃传统 MHA转而采用 GQA 来控制 KV Cache 的规模。一些代表性开源模型如下模型KV HeadQuery HeadQ:KVLLaMA 2 70B8648:1LLaMA 3 8B8324:1LLaMA 3 70B8648:1LLaMA 3 405B812816:1Mistral 7B8324:1Mixtral 8x7B8324:1Qwen 2.5 72B8648:1Gemma 2 9B8162:1值得一提的是8 个 KV Head 几乎成为行业默认值。这是因为对于常见的 个 Query Head 而言8 个 KV Head 已经能够保留足够丰富的注意力模式同时又能让 KV Cache 缩小到原来的 左右这是大量实践下的优解。