GQA/MQA注意力：KV共享加速推理

可视化演示

展示SCAN电梯调度算法的磁头移动过程：磁头沿一个方向移动到端点后反向，依次处理请求队列中的磁道访问请求。

背景动机

图书馆管理员的烦恼，从每一个研究员都要求“人手一套复印本”开始。在标准多头注意力机制中，模型有多个查询头，每个查询头都独立地对同一个序列进行注意力计算，相当于每个研究员都要为同一批参考书单独复印一套。这套“复印本”在推理时就是 KV 缓存——每个 token 的 Key 和 Value 向量都要被保存下来，供后续 token 查询。序列越长，复印本就越多，书架的负担就越重。

自回归推理是另一个放大器。大模型生成每个新 token 时，都要重新读取之前所有 token 的 Key 和 Value，这就像每写一句话，都要把前面写过的所有参考书重新翻一遍。当序列长度达到几千甚至几万时，KV 缓存的大小就会爆炸。

更深层的瓶颈在于内存带宽。解码阶段每次只生成一个 token，计算量比训练时小得多，此时真正限制速度的不是算力，而是把 KV 缓存从显存搬到计算单元所需要的时间——就像图书馆里最慢的不是复印机本身，而是来回搬书的走廊。KV 缓存越大，搬书时间越长，推理就越慢。这就是为什么缩小 KV 缓存体积成为优化推理性能的第一突破口。

MQA方案

管理员的第一个激进方案是 MQA，把复印本从人手一套缩减为全馆共享唯一一套。MQA 的核心理念是：既然所有查询头查询的是同一个序列，为什么不让它们共享完全相同的 Key 和 Value？于是，所有查询头“排队”查阅同一组 KV。这一改动将 KV 缓存压缩到原来的 1/H（H 为查询头数）。

缓存体积的压缩直接解除了显存焦虑。以前需要 32 套复印本，现在只需要 1 套。在批处理大小为 B 的场景下，MQA 的显存优势更加明显，因为每个 batch 都只需要维护一套 KV。这意味着可以塞进更大的 batch，吞吐量随之提高。

然而，MQA 的激进换来了质量的下降。当所有查询头被迫共享同一组 K 和 V 时，模型表达不同“关系”的能力被削弱了——就像所有研究员都必须在同一本参考书上做笔记，但每个人的笔记重点不同，挤在一起难免互相干扰。实践中，MQA 在部分任务上会出现精度下降，尤其在需要精细区分多种关系的场景里。

MQA 的另一个问题是训练不稳定和实现复杂。虽然推理时缓存大幅缩小，但训练时多个查询头共享 K 和 V 会让梯度传播路径变得复杂，影响收敛速度。因此，MQA 虽然提出来得更早，但在实际大模型中的采用率远不如后来提出的 GQA。

GQA分组

GQA 是管理员深思熟虑后的折中方案：与其让 8 个研究员挤 1 本参考书，或者每人复印 1 本，不如把他们分成小组，每组共享一套复印本。具体来说，GQA 将查询头分成 G 组，每组共享一组独立的 K 和 V 头。当 G = 1 时退化为 MQA，当 G = H 时等价于 MHA，因此 GQA 是两者之间的连续插值。

\begin{aligned}

O &= \text{softmax}\left(\frac{Q G^Q (K G^K)^\top}{\sqrt{d_k}}\right) V G^V \\

\end{aligned}

1. 分组重塑 Q：将查询头按组规整成新的形状，组数为 G。

$$Q \in \mathbb{R}^{B \times G \times (H/G) \times S \times d}$$

2. 广播 K 和 V：让每组共享同一套 K 和 V，通过 expand 操作实现零拷贝广播。

$$K, V \in \mathbb{R}^{B \times G \times S \times d}$$

3. 组内注意力：每个查询头只与所在组的 KV 做缩放点积注意力，把计算限制在小组内部。

分组数 G 是一个超参数，直接决定了“多人共享程度”。G 越大，KV 缓存越大，模型表达能力越强；G 越小，KV 缓存越小，速度越快。Llama 2 70B 选择了 G = 8、H = 64，将 KV 缓存压缩为原来的 1/8，获得了速度与质量的双赢。

GQA 的巧妙之处在于它几乎不牺牲质量。研究发现，组内共享 KV 带来的表达力损失远小于全部共享，因为多数注意力模式在组内就能被充分表达。这个结果像极了图书馆的“小组讨论区”——几个人的小组共享参考书，互相补充笔记，效率和质量都有保障。

KV缓存

KV 缓存是 GQA 的直接受益者。在推理时，每生成一个新 token，模型都需要把新 token 的 Key 和 Value 追加到缓存中，并在下一次生成时读取全部缓存。这个“写入 + 读取”的过程是内存带宽的主要消费者。KV 缓存的大小直接决定了单条请求的显存占用，进而决定了显卡能同时容纳多少条请求（batch size）。

应用 GQA 后，KV 头数从 H 降到 G，缓存体积相应变为：

压缩效果有多显著？以 Llama 2 70B 为例，该模型有 80 层、每层 8 个 KV 头、每个头 128 维。一条 4096 token 的序列，其 KV 缓存大约是 2 × 80 × 8 × 128 × 4096 ≈ 670 MB。如果不采用 GQA 而是用标准 64 个 KV 头，这个数字会膨胀到约 5.4 GB——一张 A100 80G 的显卡瞬间被吃掉了一小半。

KV 缓存的缩小还带来一个连锁收益：推理引擎可以在缓存中容纳更长的上下文。长上下文场景（如多轮对话、文档分析）是 KV 缓存压力的重灾区，GQA 让这些场景从“显存爆掉”变为“勉强可跑”甚至“流畅运行”。这也是为什么当前主流大模型几乎全部采用 GQA 而非 MHA。

性能收益

共享 KV 的最终目标不是省显存本身，而是把省下来的资源转化为吞吐量。推理吞吐量 = 每秒生成的 token 数，它受到两个因素制约：单次请求的延迟和并发请求的批处理规模。GQA 在这两个维度上都有贡献。

从延迟维度看，KV 缓存变小意味着每个生成步骤需要从显存搬运的数据量减少，内存带宽瓶颈被显著缓解。自回归解码本质上是“内存受限”任务，计算量只是附属，搬运 KV 缓存的时间常常占了大头。缓存缩到 1/8，搬运时间也近似缩到 1/8，单次生成延迟大幅下降。

这个 2～3 倍的吞吐量提升，是实测数据支撑的结论。Google 的 PaLM、Meta 的 Llama 2、Mistral 7B 等模型在采用 GQA 后，都报告了相似的收益。对于服务提供商而言，吞吐量翻倍意味着同样的显卡可以服务两倍的用户，成本直接减半。

GQA 的均衡策略还体现了一个重要理念：性能优化往往不是“快就好”，而是在多个目标之间寻找最佳平衡点。MQA 最快但质量受损，MHA 质量最好但速度受限，GQA 选择了中间地带，用略高一点的 KV 缓存换来了接近无损的质量，这才是工业界最需要的方案。

工程应用

GQA 在工程界的普及速度令人瞩目。Meta 在 Llama 2 系列中首次大规模应用 GQA，其 70B 版本采用 8 个 KV 头，而 7B 版本则用 32 个 KV 头（即接近 MHA）。这种“大模型更加需要压缩”的设计思路，恰好印证了 KV 缓存压力随模型规模增大而急剧上升的规律。

主流推理引擎对 GQA 的支持已经相当成熟。Hugging Face Transformers、vLLM、TensorRT-LLM 等都内置了 GQA 的实现，并且在 KV 缓存管理上做了精细优化：内存池化、连续显存分配、Cache Manager 动态调度。对于应用开发者而言，无需手动实现 GQA，只需在加载模型时指定正确的配置，引擎会自动构建共享 KV 的注意力计算图。

在长上下文部署场景中，GQA 的地位不可替代。例如部署一个支持 32K 上下文的对话模型，如果使用 MHA，KV 缓存可能占据超过 20GB 显存，导致 batch size 被压到个位数；而使用 GQA 后，缓存缩减使 batch size 可以扩大数倍，服务吞吐量实现质的飞跃。

除了 Llama 2，Mistral 7B 和 Mixtral 8x7B 也采用 GQA，其 7B 模型使用 8 个 KV 头，配合滑动窗口注意力（Sliding Window Attention），在长文本任务上取得了极佳的效果。这代表了当前大模型架构的一个明确趋势：在注意力机制上做“减法”，用结构优化换取效率。

对于 AI 工程师而言，理解 GQA 不仅是为了读懂模型配置，更是为了在实际部署中做出正确决策：何时选择 GQA（推理成本敏感型场景）、如何设置分组数（在质量预算内尽量压缩）、怎样配合其他优化手段（如 KV 量化、PagedAttention）形成组合拳。这些知识直接决定了推理服务的成本与体验。

AI 应用场景

Python 演示

import torchimport torch.nn as nnimport torch.nn.functional as F  class GroupedQueryAttention(nn.Module):    """分组查询注意力 GQA：分组共享 KV，可退化为 MQA/MHA"""     def __init__(self, d_model=512, h_q=8, h_kv=2, dropout=0.0):        super().__init__()        self.h_q = h_q        self.h_kv = h_kv        self.d = d_model // h_q        self.w_q = nn.Linear(d_model, d_model, bias=False)        self.w_k = nn.Linear(d_model, h_kv * self.d, bias=False)        self.w_v = nn.Linear(d_model, h_kv * self.d, bias=False)        self.w_o = nn.Linear(d_model, d_model, bias=False)        self.drop = nn.Dropout(dropout)     def forward(self, x, past_kv=None):        b, s, _ = x.shape        q = self.w_q(x).view(b, s, self.h_q, self.d).transpose(1, 2)        k = self.w_k(x).view(b, s, self.h_kv, self.d).transpose(1, 2)        v = self.w_v(x).view(b, s, self.h_kv, self.d).transpose(1, 2)        if past_kv is not None:            pk, pv = past_kv            k = torch.cat([pk, k], dim=2)            v = torch.cat([pv, v], dim=2)        # 核心：将 q 的 head 维度按组 reshape，再将 kv 扩展回原头数        group = self.h_q // self.h_kv        q = q.view(b, self.h_kv, group, s, self.d)        k_e = k.unsqueeze(2).expand(b, self.h_kv, group, -1, self.d)        v_e = v.unsqueeze(2).expand(b, self.h_kv, group, -1, self.d)        score = (q @ k_e.transpose(-1, -2)) / (self.d ** 0.5)        attn = F.softmax(score, dim=-1)        attn = self.drop(attn)        out = (attn @ v_e).transpose(1, 2).contiguous()        out = out.view(b, s, self.h_q * self.d)        return self.w_o(out), (k, v)  def main():    torch.manual_seed(0)    model = GroupedQueryAttention(d_model=64, h_q=8, h_kv=2)    x1 = torch.randn(1, 5, 64)    x2 = torch.randn(1, 3, 64)    out1, kv = model(x1)    out2, kv2 = model(x2, kv)    # 验证拼接后长度    assert kv2[0].shape[2] == 8, "KV 缓存应增大到 8"    print("GQA 单步输出:", out1.shape)    print("GQA 增量输出:", out2.shape)    print("KV 缓存形状:", kv2[0].shape)  if __name__ == "__main__":    main()

1. q = self.w_q(x).view(...).transpose(1, 2)将 Q 投影后拆成多个查询头，shape 变为 [batch, 头数, 序列, 维度]。
2. self.w_k = nn.Linear(d_model, h_kv * self.d, bias=False)关键设计：仅投影出 h_kv 组 K（如 2 组），而不是 h_q 组，从源头压缩了 KV 缓存。
3. torch.cat([pk, k], dim=2)在序列维度拼接历史 KV 和新 token 的 KV，这是缓存追加的标准操作。
4. q.view(b, self.h_kv, group, s, self.d)将 h_q 个查询头按组 reshape，例如 8 个头 → 2 组 × 4 个头，为后续广播做准备。
5. k.unsqueeze(2).expand(...)零拷贝广播 KV 到每个组内全部查询头。expand 不复制数据，这正是 GQA 高效的原因。
6. (q @ k_e.transpose(-1, -2)) / (self.d ** 0.5)缩放点积注意力，计算每个查询头与其组内 KV 的匹配度。
7. F.softmax(score, dim=-1)对最后一个维度归一化，得到注意力权重。
8. return self.w_o(out), (k, v)返回输出与 KV 缓存，后者将在下一次生成时复用。

常见误区

面试题