GQA分组注意力：Transformer推理优化

可视化演示

展示MHA/GQA/MQA三种注意力头结构的KV共享机制，以及KV Cache与推理效率的对比。

MHA内存瓶颈

先回到我们的演唱会现场。多头注意力（MHA）是这个场馆的“全明星速记员制”：每位明星（query 头）都配备一名专属速记员（对应的 key/value 头）。这种“一人一记”的豪华配置确实让记录质量无懈可击，可问题也随之而来——当明星越来越多，速记员团队的规模就变得极其庞大，场馆的后台（显存）根本堆不下这么多人手。

这就是 MHA 在推理阶段遭遇的尴尬。Transformer 解码时，需要把历史上所有 token 的键值对（Key 和 Value）缓存下来——这就是 KV Cache。在 MHA 中，每个查询头都有一份独立的键值缓存，假设模型有 32 个注意力头，每层就要存 32 份键值对。乘以 32 层、乘以序列长度、乘以批大小，显存瞬间就被塞满了。

更致命的是内存带宽瓶颈。解码阶段每一步只生成一个 token，却要把缓存中的所有键值从头到尾读取一遍。哪怕计算本身只需要几毫秒，读数据的时间却要几十毫秒——就像演唱会观众只有一位，却要把整个速记员团队叫醒来翻阅所有档案。这种“算的少、读的多”的模式，让内存带宽成了真正的短板。MHA 对显存和带宽的挥霍，正是我们要解决的第一道难题。

回到演唱会：MHA 就像给每位明星配一个速记员，显存就是后台的休息室。明星越多，休息室越拥挤，最终走廊上都站满了人，进出都困难。

MQA激进简化

既然“一人一记”这么贵，为什么不试试“全场只用一个速记员”呢？这就是多查询注意力（MQA）的思路——所有查询头共享同一组键值头。相当于演唱会雇了一位金牌速记员，同时为台上所有的明星服务。

这个方案的效果立竿见影。KV Cache 的占用直接从“头数份”缩减到“1 份”。如果原本有 32 个查询头，那么 KV 缓存的内存直接降为原来的 1/32！内存带宽压力也同步大减，解码速度飙升。理论上，这是效率的极致。

然而，代价同样明显。想象这位速记员要同时记录 32 位明星的每句话、每个动作，他能记住谁说了什么？注意力头本应从不同角度（语法、语义、指代）捕捉信息，现在所有头共用一份键值对，信息多样性大打折扣。实验表明，MQA 与 MHA 相比，在质量上确实有所下滑，尤其是在长文本任务上。

更麻烦的是，MQA 的模型训练收敛速度偏慢，且很难通过事后微调从 MHA 直接转换。它像一个“一步到位”的激进改革：省了成本，但伤了元气。MQA 把“省”做到了极致，却也把“好”牺牲了太多，这不是我们想要的终点。

MQA 就像把速记员团队裁到只剩一人，省钱是省到了极致，可这位速记员记不过来。我们需要一个既省钱又不降质的方案。

GQA折中智慧

回到演唱会后台，主办方想通了一个道理：我们不需要给每位明星配速记员，也不需要全场合用一个，而是可以把明星分成 4 组，每组配一名速记员。这就是分组查询注意力（GQA）的核心思想。

GQA 的做法非常优雅：假设模型有 32 个查询头，我们设定分组数为 8，那么每组包含 4 个查询头，共享一组键值对。这样一来，KV Cache 的占用就是 MHA 的 8/32 = 1/4，也就是减少了 75% 的显存。而由于 4 个查询头共享一组 KV，它们捕捉信息的多样性虽然比 MHA 略逊，但远比 MQA 丰富。

为什么 GQA 能保持较好的质量？因为注意力头的多样性需求主要集中在 Query 侧——查询角度可以丰富多彩，而 Key/Value 侧记录的事实信息没必要重复太多。比如一段话里的“苹果”，既可以指水果，也可以指公司，但“苹果”这个名词本身对应的键值信息是确定的，不同头只是从不同角度查询它而已。因此，共享 KV 并不会造成严重的信息丢失。

GQA 还具备一个极佳的工程特性：分组数是一个可调节的超参数。设为 1 时退化为 MQA，设为与头数相同时退化为 MHA。它像一个可伸缩的旋钮，让我们在质量和效率之间自由滑动。Llama 2 之所以选择 GQA，正是看中了这种灵活性。

GQA 的分组共享，就像是给演唱会配了 8 位速记员，每位负责 4 位明星。成本降了，记录质量还基本保留——两全其美。

KV Cache机制

前面反复提到 KV Cache，它是理解 GQA 价值的基石。你可以把 KV Cache 想象成演唱会的“舞台提词器”：当明星（模型）需要回忆之前说过什么、观众反应如何时，不用重新翻看整场录影，而是直接看提词器上的摘要。

在 Transformer 解码过程中，每生成一个新 token，都需要与之前所有 token 计算注意力。如果不缓存，就得把之前的 token 全部重新计算一遍键值——这相当于每次说一句话都要重新读完整个剧本，计算量是灾难性的。KV Cache 的聪明之处在于：把历史 token 的 Key 和 Value 矩阵保存下来，新 token 只需计算自己的 Query 与缓存的 Key/Value 做注意力运算。

然而，KV Cache 也是有代价的——它占显存。我们以 Llama 2-7B 为例：该模型有 32 层，每层有 32 个注意力头，每头维度 128。当序列长度为 2048 时，如果采用 MHA，KV Cache 的占用约为 32 × 32 × 128 × 2048 × 2 字节 ≈ 512MB。这还只是一个请求，如果批处理 8 个请求，就是 4GB！显存瞬间被缓存吞没。

而 KV Cache 的减半机制，正是通过减少 KV 头数实现的。在 MHA 中，32 个查询头对应 32 组 KV；在 GQA 中，这 32 个查询头只对应 8 组 KV，缓存体积直接缩到原来的 1/4。缓存小了，内存带宽压力也同步下降，解码速度自然提升。

这个机制就像演唱会后台不断增加的道具箱。MV Cache 就是为了避免每次演出都重搬道具，但道具箱本身也占地方。GQA 的妙处在于：减少道具箱的数量，但不影响演出效果。

分组比例设计

既然 GQA 的分组数 $G$ 可以自由调节，那到底选多少合适？是做事的另一大关键。如果 $G$ 选太小（比如 1），就成了 MQA，质量受损；如果 $G$ 选太大（比如 32），就等于 MHA，显存压力又回来了。

实践中，最常见的配置是 $G=8$。Llama 2 全系列模型都采用“32 个查询头 + 8 个 KV 头”的配置。为什么是 8？这背后有硬件适配的考量：现代 GPU 的并行计算单元通常以 8 或 16 为粒度，选择 8 组 KV 头可以让每组查询头数量（32/8=4）在硬件上实现高效的内存访问模式。

不同分组比例对性能的影响也呈现鲜明的权衡曲线。把 $G$ 从 32 降到 8，KV Cache 内存下降 75%，而质量损失几乎可以忽略（约 0.1% 的困惑度增加）。继续降到 4，内存再降一半，但质量损失开始显性化。到 2 或 1 时，质量下降已经能让用户感知了。这就是所谓的“边际收益递减”——前期的压缩几乎无痛，后面的压缩则开始伤害效果。

此外，分组比例的选择还要考虑任务性质。对于短文本的对话场景（序列长度几百 token），KV Cache 的绝对占用不大，$G$ 大一点没关系；但对于长文本生成（如论文摘要、代码生成），序列长度动辄几万 token，KV Cache 每压缩一点都能省出大量显存，此时选择较小的 $G$ 更划算。

分组比例的设计就像演唱会速记员人数的取舍：速记员太多，后台拥挤（显存高）；太少，记录质量堪忧（效果差）。8 人组是一个经过实践检验的“黄金人数”。

解码阶段提速

解码阶段，才是 GQA 真正发光的舞台。理解这一点，要先搞懂 Transformer 生成文本的方式：它是逐 token 生成的，就像作家一个字一个字地写文章。每写一个字，都要回顾之前已经写下的所有内容。

在每一步解码中，模型都要执行这样的操作：将当前 token 的 Query 与 KV Cache 中所有的 Key 计算注意力分数，再与对应的 Value 加权求和。关键在于——KV Cache 中存储的数据量通常远大于计算量。一个 7B 模型，每一步解码的浮点运算量只有几百万次，但要读取的 KV 数据却有几十 MB。GPU 的算力再强，也追不上内存搬运的速度，这就是“内存带宽瓶颈”。

GQA 怎么解决这个问题？减少 KV 头数 = 减少需要搬运的数据量。假设原本每步要读取 32 组 KV（MHA），现在只需读取 8 组（GQA），数据搬运量直接减少 75%。每一步省下 75% 的带宽，累积到几千步的解码过程，生成速度的提升是肉眼可见的。

这就是为什么 GQA 能显著提升推理吞吐量的原因。它不仅省显存（能塞下更大的批大小），还省带宽（每步解码更快）。两者叠加，推理吞吐量自然水涨船高。这也是 Llama 2 选择 GQA 的底层逻辑——它不是为了让模型训练得更快，而是为了让部署后的推理成本更低、响应更快。

解码的提速原理，相当于让速记员只把“小组摘要”递给明星，而不是把整场演唱会的所有记录都抱过来。小组摘要够用，还更省时间。

Llama工程实践

理论说得再好，最终要看工程实践。Meta 的 Llama 2 系列模型，是 GQA 大规模落地的最佳注脚。Llama 2-7B 和 13B 模型采用 MHA，而到了 34B 和 70B 的大规模版本，则切换到了 GQA。

为什么大模型更需要 GQA？因为 KV Cache 的开销与模型规模成正比——层数越多、头数越多、维度越大，KV Cache 的膨胀就越恐怖。70B 模型如果用 MHA，KV Cache 能撑爆几乎所有的显存，根本没法做长文本推理。Llama 2-70B 采用 GQA 后，KV Cache 显存直接除以 4，让 70B 模型可以在有限的 GPU 上流畅运行。

实测数据显示，GQA 带来的提升非常可观。在相同的硬件环境下，Llama 2-70B（GQA）的推理吞吐量比同规模的 MHA 模型提升了大约 2-3 倍。同时，长文本生成的中间态显存峰值显著下降，这使得 Llama 2 能够支持更长的上下文窗口。在工程部署中，这种优势意味着更低的成本、更高的并发。

Llama 之外，Falcon 系列、Mistral 系列也纷纷采用 GQA 或变体。这说明 GQA 已经从学术研究走向工业标准，成为大模型推理优化的事实标配。对于 AI 从业者来说，GQA 不是可学可不学的知识，而是打开高性能推理大门的一把钥匙。

Llama 2 的实践表明：在真正的工业场景中，GQA 不是纸上谈兵的优化技巧，而是撑起大规模模型高效推理的基石工程。

AI 应用场景

Python 演示

import mathimport torchimport torch.nn as nnimport torch.nn.functional as F class MHA(nn.Module):    """多头注意力：每个查询头配一组独立的KV头"""    def __init__(self, d_model, n_heads):        super().__init__()        self.n_heads = n_heads        self.head_dim = d_model // n_heads        self.wq = nn.Linear(d_model, d_model)        self.wk = nn.Linear(d_model, d_model)        self.wv = nn.Linear(d_model, d_model)        self.wo = nn.Linear(d_model, d_model)     def forward(self, x):        B, L, _ = x.shape        q = self.wq(x).view(B, L, self.n_heads, self.head_dim).transpose(1, 2)        k = self.wk(x).view(B, L, self.n_heads, self.head_dim).transpose(1, 2)        v = self.wv(x).view(B, L, self.n_heads, self.head_dim).transpose(1, 2)        attn = F.scaled_dot_product_attention(q, k, v)        return self.wo(attn.transpose(1, 2).reshape(B, L, -1))  class GQA(nn.Module):    """分组查询注意力：n_heads个查询头共享n_kv_heads组KV头"""    def __init__(self, d_model, n_heads, n_kv_heads):        super().__init__()        self.n_heads = n_heads        self.n_kv_heads = n_kv_heads        self.head_dim = d_model // n_heads        self.wq = nn.Linear(d_model, d_model)        self.wk = nn.Linear(d_model, n_kv_heads * self.head_dim)        self.wv = nn.Linear(d_model, n_kv_heads * self.head_dim)        self.wo = nn.Linear(d_model, d_model)     def forward(self, x):        B, L, _ = x.shape        q = self.wq(x).view(B, L, self.n_heads, self.head_dim).transpose(1, 2)        k = self.wk(x).view(B, L, self.n_kv_heads, self.head_dim).transpose(1, 2)        v = self.wv(x).view(B, L, self.n_kv_heads, self.head_dim).transpose(1, 2)        # 分组扩展：将KV头复制到对应的查询组        reps = self.n_heads // self.n_kv_heads        k = k.repeat_interleave(reps, dim=1)        v = v.repeat_interleave(reps, dim=1)        attn = F.scaled_dot_product_attention(q, k, v)        return self.wo(attn.transpose(1, 2).reshape(B, L, -1))  def count_kv_cache(seq_len, batch_size, d_model, n_heads, n_kv_heads):    """计算KV Cache占用（半精度字节数）"""    head_dim = d_model // n_heads    mha_cache = 2 * seq_len * batch_size * n_heads * head_dim * 2    gqa_cache = 2 * seq_len * batch_size * n_kv_heads * head_dim * 2    return mha_cache, gqa_cache  if __name__ == "__main__":    torch.manual_seed(42)    d_model, seq_len, batch_size = 512, 128, 4    n_heads, n_kv_heads = 32, 8     model_mha = MHA(d_model, n_heads)    model_gqa = GQA(d_model, n_heads, n_kv_heads)     x = torch.randn(batch_size, seq_len, d_model)    out_mha = model_mha(x)    out_gqa = model_gqa(x)     mha_cache, gqa_cache = count_kv_cache(seq_len, batch_size, d_model, n_heads, n_kv_heads)    print(f"MHA 输出形状: {out_mha.shape}")    print(f"MHA KV Cache 占用: {mha_cache / 1024:.2f} KB")    print(f"GQA 输出形状: {out_gqa.shape}")    print(f"GQA KV Cache 占用: {gqa_cache / 1024:.2f} KB")    print(f"显存节省比例: {(1 - gqa_cache / mha_cache) * 100:.1f}%")

1. self.wq = nn.Linear(d_model, d_model)查询头保持全量，每个头都有独立的查询变换
2. self.wk = nn.Linear(d_model, n_kv_heads * head_dim)键投影压缩到 n_kv_heads 组，这是 GQA 的核心
3. self.wv = nn.Linear(d_model, n_kv_heads * head_dim)值投影同样压缩到 n_kv_heads 组
4. k.repeat_interleave(reps, dim=1)将 KV 头复制到对应的查询组，完成分组共享
5. F.scaled_dot_product_attention(q, k, v)PyTorch 内置的高效缩放点积注意力
6. count_kv_cache(...)验证 KV Cache 的显存节省比例，MHA 与 GQA 对比一目了然

常见误区

面试题