GQA 分组查询注意力：用分组 KV 头桥接多头质量与多查询速度

笔记日期： 2026-05-27 笔记作者： Zhongzhu Zhou 论文标题： GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints 作者： Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, Sumit Sanghai arXiv： 2305.13245 状态 / Venue： EMNLP 2023

一句话总结

GQA（分组查询注意力）是介于经典多头注意力（MHA，每个查询头独占一对 K、V 头）和激进的多查询注意力（MQA，所有查询头共享同一对 K、V 头）之间的结构化注意力变体。它将 $H$ 个查询头划分为 $G$ 个组，每组共享一对 K、V 头：$G = H$ 时退化为 MHA，$G = 1$ 时退化为 MQA，中间的 $G$ 值则给出质量与解码速度之间的中间取舍。

论文做了三件事：

1. 提出 GQA 这一架构本身。它现在已经是生产级 LLM 的默认选择（LLaMA 2 70B、LLaMA 3、Mistral、Falcon、Qwen、DeepSeek-V2 base），原因正是它能在保持 MHA 质量的同时拿到几乎 MQA 全部的解码带宽收益。
2. 提出 uptraining（迁移训练） 配方，把已有 MHA 检查点转换为 GQA / MQA 检查点。代价只有原始预训练 5% 的算力，关键技巧是把组内每个头的 KV 投影做均值池化作为初始化。
3. 通过实验确定了一条权衡曲线：在 T5 XXL（$H = 64$）上选 $G = 8$，推理速度比 MHA-XXL 快 5.4×，而平均分仅下降 0.1。

论文最重要的一组数字来自 T5 XXL：MHA = 47.2 / 1.51 s 每样本；MQA = 46.6 / 0.24 s；GQA-8 = 47.1 / 0.28 s。GQA-8 比 MQA 多用 0.04 s，但拿回了 0.5 个平均分，这是一个非常锋利的帕累托点。

如果你今天要为一款主打推理服务的模型选注意力变体，默认就用 $G \in \{4, 8\}$ 的 GQA——除非有特别理由不这样做。下面这篇笔记会把背后的原因一层一层拆开。

前置知识

要理解 GQA，必须先理解 MHA 在解码阶段实际做了什么、为什么 KV 缓存 是自回归生成时占据主导地位的内存对象、以及”内存带宽瓶颈”在量化意义上到底意味着什么。本节铺垫这些基础，已经熟悉的读者可以直接跳到 论文做了什么 一节。

1. Transformer 解码器与自回归生成

只有解码器的 Transformer（GPT、LLaMA、Mistral 背后的架构）以一次一个 token 的方式生成文本。给定长度为 $P$ 的提示词，模型先做一次 prefill（前缀填充）的并行前向，得到 $h_1, \ldots, h_P$。之后进入 decode 循环：

$$y_{t+1} = \mathrm{argmax}_v \; \mathrm{softmax}\big(W_{out} h_t\big)_v$$

直到采样到结束符。每一步 decode 都接收一个 token id，做嵌入、走每一层、投到词表、采样、追加。

这与训练阶段截然不同。训练时 $L$ 个 token 在一次并行前向里一起处理；解码时 $L$ 个 token 要 $L$ 次串行前向，每次序列方向上的有效 batch 都是 1。算术强度（每字节内存读到的 FLOP 数）骤降，工作负载从计算瓶颈切换到内存带宽瓶颈。

flowchart LR
    subgraph Prefill["Prefill 阶段：P 个 token 并行"]
        P1[Token 1] --> H1[h_1]
        P2[Token 2] --> H2[h_2]
        P3[Token ...] --> H3[h_...]
        PP[Token P] --> HP[h_P]
    end
    subgraph Decode["Decode 阶段：一次一个 token"]
        D1[h_P] --> Y1[y_P+1]
        Y1 --> D2[一次前向]
        D2 --> Y2[y_P+2]
        Y2 --> D3[一次前向]
        D3 --> Y3[y_P+3]
    end
    HP -.最后一层隐藏态.-> D1

生产环境里成本主要来自 decode 循环。用户输入 50 个 token，模型回 500 个 token：prefill 被分摊到 500 步里，单步 decode 必须够快。

2. 多头注意力 MHA——标准构建块

Transformer 的注意力模块输入 $X \in \mathbb{R}^{L \times d_{model}}$，输出同形状。$H$ 个头，每个头维度 $d_k = d_{model}/H$，标准 MHA 写作：

$$\mathrm{MHA}(X) = \mathrm{Concat}(\mathrm{head}_1, \ldots, \mathrm{head}_H) W^O$$ $$\mathrm{head}_h = \mathrm{softmax}\!\left(\frac{Q_h K_h^\top}{\sqrt{d_k}}\right) V_h$$ $$Q_h = X W_h^Q, \quad K_h = X W_h^K, \quad V_h = X W_h^V$$

其中 $W_h^Q, W_h^K, W_h^V \in \mathbb{R}^{d_{model} \times d_k}$，输出投影 $W^O \in \mathbb{R}^{H d_k \times d_{model}}$。

三个性质决定了后续的分析：

• 每个头独立。 每个头有自己的 $W^Q, W^K, W^V$。$W^O$ 是头之间唯一的交互点。
• 参数量对称。 Q、K、V 投影每层各占 $H d_k d_{model}$ 个参数。表面上是对称的，但对 decode 而言，砍掉 K、V 投影会丢掉 2/3 的投影参数，更关键的是会丢掉 2/2 的 KV 缓存。
• 内积结构。 $\mathrm{head}_h$ 是对内积做 softmax。不同头可以学到关注不同位置、不同句法关系、不同主题。这种多样性是 MHA 在实验上优于单头的根本原因。

下面是 MHA 在 decode 阶段每一步的完整算法，写得这么细是为了之后和 MQA、GQA 的访存模式作对比。

算法 1：MHA 解码（第 t 步，单层）
输入：
  x_t          ∈ R^{1 × d_model}   当前 token 的隐藏态
  W_h^Q, W_h^K, W_h^V，其中 h = 1..H
  K_cache      ∈ R^{L × H × d_k}    之前所有 token 的 K
  V_cache      ∈ R^{L × H × d_k}    之前所有 token 的 V
  W^O          ∈ R^{(H d_k) × d_model}

 1.  对每个 h = 1..H：
 2.      q_h        = x_t · W_h^Q                         # [1 × d_k]
 3.      k_h_new    = x_t · W_h^K                         # [1 × d_k]
 4.      v_h_new    = x_t · W_h^V                         # [1 × d_k]
 5.      K_cache[t, h, :] = k_h_new                       # 追加
 6.      V_cache[t, h, :] = v_h_new                       # 追加
 7.      scores_h   = q_h · (K_cache[:t+1, h, :])^T / √d_k  # [1 × (t+1)]
 8.      attn_h     = softmax(scores_h)                    # [1 × (t+1)]
 9.      out_h      = attn_h · V_cache[:t+1, h, :]         # [1 × d_k]
10.  返回 Concat(out_1, ..., out_H) · W^O                  # [1 × d_model]

逐行解读：

• 第 2 行 把输入投影成 $H$ 个查询向量。便宜：每头 $d_{model} \cdot d_k$ 个 FLOP。
• 第 3–4 行 计算这个 token 的新 K、V 向量，准备追加到缓存。
• 第 5–6 行 写入 KV 缓存。整体上 K 缓存从 $L \cdot H \cdot d_k$ 增长到 $(L+1) \cdot H \cdot d_k$，V 同理。
• 第 7 行 当前查询和所有 $L+1$ 个缓存 K 做点积。这要读取头 $h$ 的整段 K 缓存，也就是 $(L+1) \cdot d_k$ 个值。
• 第 8 行 一个长度 $L+1$ 的 softmax，很便宜。
• 第 9 行 注意力权重乘缓存的 V，又要读 $(H, L)$ 那一块的 $(L+1) \cdot d_k$ 个值。
• 第 10 行 拼接、过 $W^O$。

关键观察：第 7、9 行加起来在所有 $H$ 个头上要从 HBM 读 $2 H (L+1) d_k$ 个 FP16 值。这就是 每步每层的 KV 缓存流量。

3. KV 缓存：推理内存去哪儿了

KV 缓存是解码器的持久状态。如果没有它，每一步 decode 都得把所有历史 token 的 K、V 从头算一遍，复杂度变成 $O(L^2 d_{model})$ FLOP，很快就跑不动。有了它，每步计算降到 $O(L d_{model})$ FLOP，但每步内存流量也是 $O(L d_{model})$ 字节。

给点具体数字。考虑一个大约 7B LLaMA 量级的模型：

• 层数 $N_L = 32$
• 头数 $H = 32$
• 头维 $d_k = 128$
• 序列长 $L = 2048$
• 精度 FP16（2 字节）
• 批 $B = 1$

KV 缓存总大小：

$$S_{KV} = 2 \cdot N_L \cdot H \cdot d_k \cdot L \cdot B \cdot \text{bytes\_per\_value}$$ $$= 2 \cdot 32 \cdot 32 \cdot 128 \cdot 2048 \cdot 1 \cdot 2 = 1{,}073{,}741{,}824 \text{ 字节} = 1 \text{ GiB}$$

如果换成 70B 量级（$N_L = 80, H = 64, d_k = 128$）、$L = 4096$：$S_{KV} = 2 \cdot 80 \cdot 64 \cdot 128 \cdot 4096 \cdot 2 = 10.7$ GiB——单个请求的 KV 缓存就有十几个 GiB。再乘上 32 个并发用户，KV 缓存就会成为 HBM 里最大的对象，甚至比权重还大。

MHA 的 KV 缓存随 $N_L, H, d_k, L, B$ 线性增长。其中 $H$ 和 $L$ 这两项正是长上下文、高并发服务最难的地方。

flowchart TB
    subgraph HBM["GPU HBM（A100 上 80 GB）"]
        W[权重 7B fp16 约 14 GB]
        KV[KV 缓存 batch=32, L=2048<br/>= 32 × 1 GiB = 32 GiB]
        ACT[激活 较小]
    end
    subgraph SM["流式多处理器（计算单元）"]
        ATTN[第 h 个头的注意力]
    end
    HBM -- 2 TB/s --> SM
    SM -- 写回 --> HBM

A100 的 HBM-to-SM 带宽约 2.0 TB/s，H100 约 3.35 TB/s。计算吞吐则高得多：A100 FP16 张量核 312 TFLOP/s。FP16 矩阵乘的算术强度临界点大约在 156 FLOP/字节，解码注意力远远低于这个临界点。

4. 解码阶段的内存带宽瓶颈

我们把单层一次 decode 的 FLOP 和字节都数出来。用上述 7B 数字（$H = 32, d_k = 128, L = 2048, d_{model} = 4096$）：

每层每步 FLOP：

• Q、K、V 投影：$3 \cdot d_{model} \cdot H d_k = 3 \cdot 4096 \cdot 4096 \approx 50.3$ MFLOP
• 注意力分数：$H \cdot L \cdot d_k = 32 \cdot 2048 \cdot 128 \approx 8.4$ MFLOP
• 注意力输出乘 V：同上 $\approx 8.4$ MFLOP
• $W^O$：$d_{model}^2 = 16.8$ MFLOP
• MLP（典型 4× 放大）：$\approx 134.2$ MFLOP

合计约 218 MFLOP，即 $\sim 2.2 \times 10^8$ FLOP。

每层每步字节：

• 权重（Q、K、V、O 投影加 MLP）：约 400 MB
• KV 缓存流量：$2 \cdot H \cdot L \cdot d_k \cdot 2 = 32$ MB
• 激活：很小，约 16 KB

合计每层约 432 MB，32 层就是 13.8 GB 每 token。

A100（2 TB/s）读 13.8 GB 用时 $13.8 / 2000 \approx 6.9$ ms。32 层算力（$\sim 7 \times 10^9$ FLOP）在 312 TFLOP/s 下只要 $\sim 0.022$ ms——比内存读慢 300 多倍。

结论很冷酷：batch=1 下，GPU 有 99.7% 的时间在等内存。要提速只能减字节。减权重很难——那就是模型大小。减 KV 缓存相对容易，因为 KV 缓存存在结构冗余可以被利用：$H$ 个 KV 头的表示是相关的。

flowchart LR
    subgraph Cycle["一次 decode"]
        L1[读权重 约 400 MB] --> C1[计算 Q,K,V]
        C1 --> L2[读 KV 缓存 约 32 MB]
        L2 --> C2[计算注意力]
        C2 --> L3[读 MLP 权重 约 250 MB]
        L3 --> C3[计算 MLP]
        C3 --> L4[写 KV 约 16 KB]
    end
    style L1 fill:#fbb,stroke:#900
    style L2 fill:#fbb,stroke:#900
    style L3 fill:#fbb,stroke:#900
    style C1 fill:#bfb,stroke:#090
    style C2 fill:#bfb,stroke:#090
    style C3 fill:#bfb,stroke:#090

红色方块是带宽受限，绿色方块是计算受限。batch=1 时解码就是一连串红色方块，绿色方块隐藏在它们后面。增大 batch 可以把权重读分摊到多个请求上（权重是共享的），但 KV 缓存是每个请求独占的，会随 batch 线性增长。所以 KV 缓存是唯一可以从结构上攻击的项。

5. 多查询注意力 MQA——极端精简方案

Noam Shazeer 在 2019 年提出 MQA：每层仍有 $H$ 个查询头，但只有一个键头和一个值头：

$$\mathrm{head}_h^{MQA} = \mathrm{softmax}\!\left(\frac{Q_h K^\top}{\sqrt{d_k}}\right) V$$

其中 $K = X W^K$、$V = X W^V$ 在所有 $h$ 间共享。KV 缓存从 $2 H L d_k$ 缩到 $2 L d_k$，缩小 $H$ 倍。对 $H = 32$ 到 $64$ 的模型来说，这是 1–2 个数量级。

实测效果很显著。Shazeer 原文（以及本文在 T5 XXL 上复现）的结论：MQA 在 decode 上能比 MHA 快 6 倍左右，质量损失个位数百分点。代价是：MQA 太激进了——把 32 个 KV 头压成 1 个，丢掉了每头的差异性；在摘要这种任务上的质量损失是可观测的。

GQA 的核心论点：不必在 $H$ 个 KV 头（MHA，最大质量、最大带宽）和 $1$ 个 KV 头（MQA，最大带宽节省、可见质量损失）之间二选一。可以选任何介于两者之间的整数，并选帕累托前沿上那个最优点——实验上看，就是 $G = 8$。

论文做了什么

论文的技术内容由两部分组成。

架构。 提出 GQA。把 $H$ 个查询头划分成 $G$ 个组（一般取 $G$ 是 $H$ 的因数，每组 $H/G$ 个头）。第 $g$ 组共享一个 K 头 $K_g$ 和一个 V 头 $V_g$。KV 缓存每层从 $2 H L d_k$ 降到 $2 G L d_k$，缩小因子 $H/G$。$G = H$ 退化为 MHA，$G = 1$ 退化为 MQA。

迁移训练（uptraining）配方。 不从头训练 GQA 模型（贵——当年 T5 XXL 预训练要约 10000 TPU-day），而是把已有 MHA 检查点转换过来：

1. 第 $g$ 组的 K、V 投影定义为该组 $H/G$ 个 MHA 头投影的均值：$W_g^K = \frac{1}{H/G} \sum_{h \in g} W_h^K$，V 同理。
2. 在原始数据分布上继续训练 $\alpha = 5\%$ 的原始预训练步数。其他参数（Q 投影、$W^O$、MLP、嵌入）从 MHA 检查点继承，正常训练。

论文证明这个 5% 已经够把质量补回 MHA 水平，即使把 64 个 KV 头压成了 8 个。

两部分都重要，但意义不同。架构是真正被用起来的东西；uptraining 配方是让早期采用者敢冒险尝试的东西。今天的前沿模型（LLaMA 3、Mistral）都是直接从头训 GQA，不再需要 uptraining。但 uptraining 这个论证说服了整个领域：“GQA 是正确的结构先验”——如果对 MHA 做均值池化、再用 5% 训练就能恢复质量，那只能说明 MHA 那些额外的 KV 头本来就携带着大量冗余信息。

flowchart LR
    subgraph MHA["MHA：H=8 个 KV 头"]
        Q1[Q1]-->K1[K1]
        Q2[Q2]-->K2[K2]
        Q3[Q3]-->K3[K3]
        Q4[Q4]-->K4[K4]
        Q5[Q5]-->K5[K5]
        Q6[Q6]-->K6[K6]
        Q7[Q7]-->K7[K7]
        Q8[Q8]-->K8[K8]
    end
    subgraph GQA["GQA-2：G=2 个 KV 头"]
        GQ1[Q1]-->GK1[第 1 组 K]
        GQ2[Q2]-->GK1
        GQ3[Q3]-->GK1
        GQ4[Q4]-->GK1
        GQ5[Q5]-->GK2[第 2 组 K]
        GQ6[Q6]-->GK2
        GQ7[Q7]-->GK2
        GQ8[Q8]-->GK2
    end
    subgraph MQA["MQA：1 个 KV 头"]
        MQ1[Q1]-->MK1[共享 K]
        MQ2[Q2]-->MK1
        MQ3[Q3]-->MK1
        MQ4[Q4]-->MK1
        MQ5[Q5]-->MK1
        MQ6[Q6]-->MK1
        MQ7[Q7]-->MK1
        MQ8[Q8]-->MK1
    end

方法详解：分组查询注意力

GQA 架构

设 $H$ 为查询头数，$G$ 是 $H$ 的因数。第 $h$ 个查询头被划入第 $g(h) = \lceil h / (H/G) \rceil$ 组。第 $\ell$ 层的前向变成：

$$Q_h = X W_h^Q \quad (h = 1, \ldots, H)$$ $$K_g = X W_g^K, \quad V_g = X W_g^V \quad (g = 1, \ldots, G)$$ $$\mathrm{head}_h = \mathrm{softmax}\!\left(\frac{Q_h K_{g(h)}^\top}{\sqrt{d_k}}\right) V_{g(h)}$$ $$\mathrm{GQA}(X) = \mathrm{Concat}(\mathrm{head}_1, \ldots, \mathrm{head}_H) W^O$$

相对 MHA 唯一的改动是 K、V 投影的数量：$G$ 而不是 $H$。Q 投影没变。$W^O$ 没变。MLP 没变。这一点是这个架构的关键优点：它是一次最小化的外科手术。

decode 阶段的算法变成：

算法 2：GQA 解码（第 t 步，单层）
输入：
  x_t ∈ R^{1 × d_model}
  W_h^Q（h = 1..H）
  W_g^K, W_g^V（g = 1..G）
  K_cache ∈ R^{L × G × d_k}
  V_cache ∈ R^{L × G × d_k}
  W^O ∈ R^{(H d_k) × d_model}

 1.  对每个 g = 1..G：
 2.      k_g_new = x_t · W_g^K                            # [1 × d_k]
 3.      v_g_new = x_t · W_g^V                            # [1 × d_k]
 4.      K_cache[t, g, :] = k_g_new
 5.      V_cache[t, g, :] = v_g_new
 6.  对每个 h = 1..H：
 7.      q_h = x_t · W_h^Q                                # [1 × d_k]
 8.      g  = ceil(h / (H/G))                              # 查组号
 9.      scores_h = q_h · (K_cache[:t+1, g, :])^T / √d_k
10.      attn_h   = softmax(scores_h)
11.      out_h    = attn_h · V_cache[:t+1, g, :]
12.  返回 Concat(out_1, ..., out_H) · W^O

逐行解读：

• 第 1–5 行 计算并追加 $G$ 个新的 K、V 向量。这是和算法 1 唯一在 cardinality 上不同的地方（MHA 这里是 $H$ 个）。
• 第 6–11 行 像 MHA 一样算 $H$ 个注意力输出，但现在每组内的查询头共享同一片 KV 切片。同组内多个查询头都索引 K_cache[:, g, :]。在 GPU 上这一点至关重要：写得好的内核可以把一组里的 $H/G$ 个查询融合到一起，让每个缓存的 K 向量只读一次就服务于组内所有查询头，而不是每个查询头读一次。

每层每步的 KV 缓存流量从 $2 H L d_k$ 字节降到 $2 G L d_k$ 字节。$H = 64, G = 8$ 时缩小 8 倍。

形式化定义与 KV 缓存节省量

设 $T_{KV}^{MHA} = 2 H L d_k \cdot B \cdot \text{bytes}$ 为 batch=B 下 MHA 单层单步的 KV 流量。GQA 对应是 $T_{KV}^{GQA} = 2 G L d_k \cdot B \cdot \text{bytes}$。比值 $H/G$。

一个 70B 量级模型 $H = 64$。选 $G = 8$ 是 KV 缓存 大小和带宽 同时缩小 8 倍：

模型	$H$	$G$	$L=4096, B=1$ 下 KV 缓存	每步 KV 流量
7B MHA	32	32	1 GiB	32 MB/层
7B GQA-8	32	8	0.25 GiB	8 MB/层
70B MHA	64	64	10.7 GiB	64 MB/层
70B GQA-8	64	8	1.3 GiB	8 MB/层
70B MQA	64	1	0.17 GiB	1 MB/层

长上下文下绝对值会爆炸：$L = 32768$ 时，70B MHA 的 KV 缓存约 85 GiB——单卡 HBM 装不下。GQA-8 缩到约 10.7 GiB，能稳定放下。带宽收益几乎线性地变成 decode 墙钟时间收益，因为工作负载本来就被带宽卡住。

xychart-beta
    title "KV 缓存大小随序列长度变化（70B, fp16）"
    x-axis [1024, 2048, 4096, 8192, 16384, 32768]
    y-axis "KV 缓存 (GiB)" 0 --> 90
    line [2.7, 5.4, 10.7, 21.5, 42.9, 85.9]
    line [0.34, 0.67, 1.34, 2.68, 5.36, 10.7]
    line [0.04, 0.08, 0.17, 0.34, 0.67, 1.34]

三条线从上到下分别是 MHA、GQA-8、MQA。MHA 与 GQA-8 之间正好 8 倍；GQA-8 与 MQA 之间又是 8 倍。

迁移训练（Uptraining）流程

这是第二个贡献，也是早期采用者敢做切换的工程基础。假设你已经训出一个 MHA 模型——花了 $C$ TPU-day。你想发一个推理更快的版本，但不能把 $C$ TPU-day 的成果扔掉。

迁移训练配方：

算法 3：MHA → GQA 检查点转换 + Uptraining
输入：
  含 W_h^K, W_h^V 等参数的 MHA 检查点（h = 1..H）
  目标组数 G（H 的因数）
  Uptraining 预算 α（如 0.05）
  原始预训练步数 T_orig

 1.  group_size = H / G
 2.  对每个 g = 1..G：
 3.      // 组内 KV 投影均值池化
 4.      W_g^K = (1/group_size) · Σ_{h ∈ group(g)} W_h^K
 5.      W_g^V = (1/group_size) · Σ_{h ∈ group(g)} W_h^V
 6.  把检查点中原来的 W_h^K, W_h^V 替换为 G 个均值池化后的 W_g^K, W_g^V
 7.  保持 W_h^Q, W^O, MLP, 嵌入等参数不变
 8.  按原始 Adafactor 优化器、相同 schedule、相同语料继续训练 T_uptrain = α · T_orig 步

逐行解读：

• 第 1–5 行：均值池化是关键。在线性模型下，$W_g^K = \mathrm{mean}_{h \in g} W_h^K$ 就是把组内 $H/G$ 个向量投到它们在参数空间里的质心。如果把每个头看作”真实组级投影 + 等方差噪声”，均值就是最小方差无偏估计。
• 第 6 行：参数形状从 $(H, d_{model}, d_k)$ 改成 $(G, d_{model}, d_k)$。这些参数的优化器状态需要按新形状重置；论文指出由于 uptraining 很短，这不是问题。
• 第 7 行：注意不动 Q 投影。每个查询头保留自己的 $W_h^Q$——查询的多样性是要保留的，只在 K、V 一侧折叠。
• 第 8 行：$T_{uptrain} = 0.05 \cdot T_{orig}$。T5 XXL 原始约 100 万步，uptrain 约 5 万步，论文报告 XXL 转换约 600 TPUv3 chip-day。沿用原优化器、原学习率 schedule、原数据——这点对稳定权重至关重要。

flowchart TB
    MHACkpt[MHA 检查点] --> Group[把 H 个头分成 G 组]
    Group --> Mean[组内对 W_K, W_V 均值池化]
    Mean --> NewCkpt[新检查点，G 个 KV 头]
    NewCkpt --> Uptrain[Uptrain 原 5% 步数]
    Uptrain --> Final[可上线的 GQA 模型]
    style Mean fill:#bdf,stroke:#039
    style Uptrain fill:#fdb,stroke:#930

为什么用均值池化？

为什么挑均值，而不是别的方案？论文消融了三种：

1. 均值池化（选用）：$W_g^K = \frac{1}{H/G} \sum_{h \in g} W_h^K$
2. 取第一个头：$W_g^K = W_{h_0}^K$，$h_0$ 是组内第一个头
3. 随机初始化：从原始初始化分布重新采样

实验上排序：均值池化 ≫ 取第一个 ≫ 随机。直觉上：

均值池化保留最大量预训练信息。 组内 $H/G$ 个头通常编码相似的特征（尤其相邻头、尤其预训练之后），它们的均值是对每个头的良好近似。下游层是按拼接后的多头输出训练出来的，遇到的扰动对每个头来说都很小，对均值来说还几乎为零。

取第一个头丢掉 $H/G - 1$ 个头的学习成果。 如果选第 1 个头代表 $\{1,2,3,4\}$ 整组，第 2、3、4 头学到的东西就全丢了。即使它们和第 1 头高度冗余，那部分残差信息也是有用的，模型还得花算力重新学。

随机初始化把一切都扔了。 这本质上等于”从一个 KV 层被破坏的部分检查点出发，从头训 GQA”。给足够 uptraining 当然能恢复，但代价随 $G$ 增长得糟糕得多。

更严格一点：假设 MHA 训练时每个头的 K 投影是 $W_h^K = W^*_g + \varepsilon_h$，其中 $W^*_g$ 是”真正的组级投影”，$\varepsilon_h$ 是 $\mathbb{E}[\varepsilon_h] = 0, \mathrm{Var}[\varepsilon_h] = \sigma^2 I$ 的噪声。那么均值估计 $\widehat{W}_g^K = \frac{1}{H/G} \sum_h W_h^K$ 的方差是 $\sigma^2 / (H/G) \cdot I$——是 $W^*_g$ 的最小方差无偏估计。挑一个头的方差是 $\sigma^2 I$，差 $H/G$ 倍。随机初始化和 $W^*_g$ 毫无关系。

这个论证是非严格的（MHA 各头并不真的是某个共享均值的噪声拷贝），但论文实验结果与预测一致：均值池化第一，第一头次之，随机最差，且 $H/G$ 越大差距越大。

flowchart LR
    subgraph Group1["MHA 一组里 4 个头"]
        H1[W1_K]
        H2[W2_K]
        H3[W3_K]
        H4[W4_K]
    end
    H1 -->|求和| Plus((+))
    H2 -->|求和| Plus
    H3 -->|求和| Plus
    H4 -->|求和| Plus
    Plus -->|除以 4| Mean[均值池化后的 W_g_K]
    style Mean fill:#bdf,stroke:#039

分组数 G 的选择

GQA 给了你一个旋钮：$G \in [1, H]$，最好是 $H$ 的因数。怎么选？

KV 缓存与解码带宽 随 $G$ 线性变。$G$ 减半，KV 减半。

质量 随 $G$ 单调上升，但收益递减很快。论文数据显示，T5 XXL（$H = 64$）的质量曲线到 $G = 8$ 已基本平坦。

解码速度 取决于带宽瓶颈程度。小 batch + 长上下文下，decode 严重带宽受限，KV 缓存减少基本能 1:1 翻译成提速。大 batch 下，权重读会主导 KV 读，KV 缩减对吞吐的提升变小（但对显存占用还有意义，进而限制了 batch 上限）。

张量并行（TP）切分 引入一个硬约束：TP 度数 $K$ 时，$G$ 最好是 $K$ 的倍数，这样每个 TP rank 能完整持有整数个 KV 头。MQA（$G = 1$）没法切——只能在 $K$ 个 rank 上复制那个唯一的 KV 头，带宽节省被部分抵消。GQA-8 在 TP=8 下让每个 rank 正好持有 1 个 KV 头——干净。

论文结论：T5 XXL（$H = 64$）取 $G = 8$ 是 “favorable middle ground”。现代 LLM 沿用同一配方：LLaMA 2 70B $G = 8$，LLaMA 3 70B $G = 8$（$H = 64$），Mistral 7B $G = 8$（$H = 32$）。

xychart-beta
    title "质量-速度 Pareto 曲线（T5 XXL 示意）"
    x-axis [0.24, 0.28, 0.40, 0.60, 1.00, 1.51]
    y-axis "平均分" 46 --> 48
    line [46.6, 47.1, 47.2, 47.2, 47.2, 47.2]

最左下点是 MQA，紧靠它右上、质量几乎追平 MHA 的是 GQA-8；最右上是 MHA-XXL。Pareto 拐点恰好就在 GQA-8。

GQA 与张量并行推理

把一个 70B 模型用 TP=8 切到 8 张 GPU 上时，每层的矩阵乘按头维度拆分。MHA 下，64 个头 / 8 卡 = 每卡 8 个头，每卡算自己的 $Q_h, K_h, V_h$。

MQA 下只有 1 个 K 头和 1 个 V 头，但 Q 头还是 64 个。Q 头仍可按 8 切，但 K、V 必须在 8 张卡上复制。每卡都重新算一次 K、V 投影（工作量很小）并各自存一份缓存（每卡都是单头 KV 的全量）。累加起来全设备的 KV 缓存是单头版本的 8 倍——虽然仍然只是 MHA 整体的 1/8。

GQA-8 + TP=8 下，每卡正好持有 1 个 KV 头加 8 个对应的 Q 头。KV 被完全切分、零复制：每卡的 KV 缓存是总量的 $1/8$。结构上很干净，这也是 GQA-8 在生产环境如此受欢迎的原因之一。

flowchart TB
    subgraph TPMHA["MHA, TP=8 64Q 64KV 头"]
        D0M[卡 0：8Q + 8KV]
        D1M[卡 1：8Q + 8KV]
        D7M[... 共 8 卡 ...]
    end
    subgraph TPMQA["MQA, TP=8 64Q 1KV 复制"]
        D0Q[卡 0：8Q + KV 副本 1]
        D1Q[卡 1：8Q + KV 副本 2]
        D7Q[... KV 复制 8 份 ...]
    end
    subgraph TPGQA["GQA-8, TP=8 64Q 8KV 切分"]
        D0G[卡 0：8Q + 1KV]
        D1G[卡 1：8Q + 1KV]
        D7G[... 共 8 卡 ...]
    end
    style TPGQA fill:#dfd,stroke:#090

要点：$G$ 通常应是 TP 度数 $K$ 的倍数（或等于）。$G = K$ 是最干净的情况。$G = 8$ 兼容 TP $\in \{1, 2, 4, 8\}$，覆盖大多数部署。

实验

模型与数据集

论文在 T5（编码-解码架构）上评估，具体是 T5 Large 与 T5 XXL 的 v1.1 配置：

• T5 Large：约 7.7 亿参数，$H = 16$, $d_k = 64$, $N_L = 24$（编码 + 解码）
• T5 XXL：约 110 亿参数，$H = 64$, $d_k = 64$（T5 即使在 XXL 也保持 $d_k = 64$，是 T5 设计的有意之处），$N_L = 24$

转换只应用在解码器自注意力和编码-解码交叉注意力上。编码器自注意力保持 MHA——编码器不是带宽瓶颈（它并行处理输入，是计算瓶颈），改了反而会变慢。

数据集：

• CNN/DailyMail（新闻摘要）
• arXiv（长篇科学论文摘要）
• PubMed（生物医学摘要）
• MediaSum（对话摘要）
• MultiNews（多文档摘要）
• WMT EN-DE（机器翻译）
• TriviaQA（开放域问答）

这套数据集是刻意选的：长输入摘要任务对 KV 缓存压力最大（解码端要交叉注意大量编码端 token，交叉注意的 K、V 序列很长），arXiv / PubMed 输入动辄几千 token。如果 GQA 要在哪里翻车，最该翻在这些长上下文摘要任务上。

推理时间测量方法

单 TPU 上测量每样本推理时间。只测 decode——也就是自回归生成那部分的墙钟时间，不算 prefill。这是正确的指标，因为前面分析过，decode 才是带宽瓶颈阶段。

测量条件：贪心解码（无采样开销）、batch=1（带宽瓶颈最严重的情形，也最贴近单用户低延迟场景）。

主要实验结果

论文给出的 T5 XXL（5% uptraining）正式结果：

模型	$T_{infer}$（秒/样本）	平均	CNN/DM	arXiv	PubMed	MediaSum	MultiNews	WMT EN-DE	TriviaQA
MHA-Large	0.37	46.0	42.9	44.6	46.2	35.5	46.6	27.7	78.2
MHA-XXL	1.51	47.2	43.8	45.6	47.5	36.4	46.9	28.4	81.9
MQA-XXL	0.24	46.6	43.0	45.0	46.9	36.1	46.5	28.5	81.3
GQA-8-XXL	0.28	47.1	43.5	45.4	47.7	36.3	47.2	28.4	81.6

定量观察：

• MHA-XXL → MQA-XXL：速度 6.3 倍（1.51 → 0.24 s），平均分降 0.6。
• MHA-XXL → GQA-8-XXL：速度 5.4 倍（1.51 → 0.28 s），平均分仅降 0.1。
• GQA-8-XXL vs MHA-Large：GQA-8 是 XXL 体量（XXL 级质量），但跑得比 Large 还快（0.28 vs 0.37 s）。这是头条结果：GQA-8 让你在 sub-Large 延迟下拿到 XXL 级质量。
• PubMed 与 MultiNews 上 GQA-8 甚至超过 MHA-XXL（47.7 vs 47.5；47.2 vs 46.9）。差距小，可能在噪声范围内，但说明 GQA 不是任何意义上的”弱化版”架构——更少 KV 头带来的隐式正则化在某些任务上甚至略有帮助。

xychart-beta
    title "T5 XXL：质量 vs decode 时间"
    x-axis [0.20, 0.30, 0.40, 0.60, 1.00, 1.51]
    y-axis "平均分" 45.5 --> 47.5
    line [46.6, 47.1, 47.2, 47.2, 47.2, 47.2]

消融一：检查点转换策略

论文消融转换初始化策略：

1. 均值池化：$W_g^K = \frac{1}{H/G} \sum_h W_h^K$
2. 取第一个头：$W_g^K = W_1^K$
3. 随机初始化

结果：均值池化 > 第一个头 > 随机，差距明显。Uptraining 预算很小（1%）时差距最大；很大（5%+）时随机能追上一些，但永远追不平均值。

这是论文最直接的证据：预训练已经在每组内学到了近似共享的结构，而均值池化是把这种共享结构”抽取出来”放进 GQA 的正确方式。

消融二：迁移训练预算

论文扫描 $\alpha \in \{0, 1\%, 5\%, 10\%\}$。主要定性结论：

• $\alpha = 0$：完全不 uptrain。GQA 已经惊人地接近 MHA——均值池化本身就是不错的初始化。MQA 没 uptrain 就差得多。
• $\alpha = 1\%$：大部分差距已经填平。
• $\alpha = 5\%$：GQA 基本饱和。MQA 也很接近但略差。
• $\alpha = 10\%$：边际改进很小。

论文默认 $\alpha = 5\%$——足够便宜（预训练 100 天，再花 5 天转换），也够大到补回大部分可恢复的质量。

xychart-beta
    title "示意：质量随 uptraining α 变化"
    x-axis [0, 1, 5, 10]
    y-axis "平均分" 44 --> 48
    line [46.2, 46.8, 47.1, 47.15]
    line [44.5, 45.8, 46.6, 46.7]

上方曲线是 GQA-8（恢复快、上限接近 MHA），下方曲线是 MQA（恢复慢、上限低于 MHA）。

消融三：分组数量

T5 XXL 上扫描 $G$ 从 1（MQA）到 64（MHA）：

$G$	速度	质量（相对 MHA）
1（MQA）	最快	-0.6
2	快	-0.3
4	快	-0.15
8（GQA-8）	快	-0.1
16	中	-0.05
32	慢	-0.02
64（MHA）	最慢	0（基准）

质量曲线凹形——大部分差距在 $G = 8$ 处已经收敛，之后是长长的平尾。速度曲线在带宽瓶颈区间内大致线性于 $G$：$G$ 翻倍 KV 流量翻倍。

“额外的 $G$ 买到越来越少的质量、付出越来越多的成本” 的交点——T5 XXL（$H = 64$）下就是 $G = 8$。

一份更细的成本账

继续走主线前我把账算到底。这种算术我个人觉得很有澄清作用，而论文里只是带过。

考虑 70B 量级模型一层的 decode，$L = 8192$、batch=1。

MHA（$H = 64$, $d_k = 128$, $d_{model} = 8192$）：

• Q 投影：$1 \cdot d_{model} \cdot H \cdot d_k \approx 67$ MFLOP
• K 投影：$\approx 67$ MFLOP
• V 投影：$\approx 67$ MFLOP
• 注意力分数 $QK^\top$：$H \cdot L \cdot d_k \approx 67$ MFLOP
• softmax·V：$\approx 67$ MFLOP
• $W^O$：$d_{model}^2 \approx 67$ MFLOP
• MLP（4× 放大）：$\approx 540$ MFLOP

合计算力：约 942 MFLOP $\approx 10^9$ FLOP。

字节：

• Q+K+V 权重：$3 \cdot d_{model} \cdot H \cdot d_k \cdot 2 = 3 \cdot 8192 \cdot 8192 \cdot 2 = 400$ MB
• $W^O$ 权重：$134$ MB
• MLP 权重：$\approx 1.07$ GB
• KV 缓存流量：$2 \cdot H \cdot L \cdot d_k \cdot 2 = 268$ MB
• 激活：可忽略

合计每层每步约 1.87 GB。80 层就是约 150 GB 每 token，H100（3.35 TB/s）下约 45 ms。

GQA-8（同模型，$G = 8$）：

• Q 投影不变：约 67 MFLOP，134 MB 权重读
• K 投影：$1 \cdot d_{model} \cdot G \cdot d_k \approx 8.4$ MFLOP，权重 $d_{model} \cdot G \cdot d_k \cdot 2 = 16.8$ MB
• V 投影：同 K，约 8.4 MFLOP，16.8 MB
• 注意力分数：$H \cdot L \cdot d_k \approx 67$ MFLOP（查询仍 $H = 64$ 个）
• softmax·V：同上
• $W^O$、MLP 权重不变
• KV 缓存流量：$2 \cdot G \cdot L \cdot d_k \cdot 2 = 33.5$ MB

合计每层每步约 1.40 GB。80 层就是约 112 GB 每 token，H100 下约 33 ms。

加速比：$45 / 33 \approx 1.36$ 倍。KV 流量缩了 8 倍，但墙钟只快 1.36 倍。原因：$L = 8192$ 在这个 70B 模型上，KV 只占字节预算的约 14%；把它压到约 2% 最多省下这 14%，再加上 K、V 投影权重读变小的一小部分。

再换 $L = 32768$：

• MHA KV 流量：约 1.07 GB / 层
• 总字节：约 2.67 GB / 层，80 层约 214 GB，H100 约 64 ms
• GQA-8 KV 流量：约 134 MB / 层
• 总字节：约 1.50 GB / 层，80 层约 120 GB，H100 约 36 ms

加速比：$64 / 36 \approx 1.78$ 倍，正如线性预测。

更长，$L = 131072$：

• MHA KV 流量：约 4.29 GB / 层
• 总字节：约 5.59 GB / 层，80 层约 447 GB，H100 约 133 ms 每 token
• GQA-8 KV 流量：约 537 MB / 层
• 总字节：约 1.90 GB / 层，80 层约 152 GB，H100 约 45 ms 每 token

加速比：$133 / 45 \approx 2.95$ 倍。

结论：超长上下文下 KV 完全主导字节预算，GQA 越接近满 $H/G = 8$ 倍加速。短上下文下 KV 本来就是小头，GQA 的提速也就小。GQA 的收益随上下文长度增长——这正是它在生产上下文从 2k 一路涨到 128k 之后变得不可或缺的原因。

xychart-beta
    title "70B 模型 H100 上 decode 每 token 时间"
    x-axis [2048, 8192, 32768, 131072]
    y-axis "ms / token" 0 --> 140
    line [33, 45, 64, 133]
    line [32, 33, 36, 45]

上方是 MHA，下方是 GQA-8。差距随上下文增长而拉大。

数字背后的分析

除了原始数字，下面几点值得专门强调。

6.3× 的 MQA 加速比并非 64×。 MQA 把 KV 缓存缩了 64 倍（$H = 64$），但解码墙钟只快 6.3 倍。因为 KV 缓存流量只是每步内存预算的一部分。权重每步还要读，激活每步还要写。把 KV 流量从字节预算的 60% 压到 5%，理论上最多在内存时间上快 2.4×；剩下的 6.3× 来自 K、V 投影计算消失（虽小但 KV 流量缩小后变成相对显著项）和内核层面效应——寄存器压力降低、缓存命中率提升。

为什么 GQA-8 只比 MQA 略慢？ MQA 是 $G = 1$，GQA-8 的 KV 流量是 MQA 的 8 倍。可墙钟差距只是 0.28 vs 0.24 s，约 17%，不是 700%。因为 $G = 8$ 时 KV 流量已经被压到远小于权重流量——再往下挤只能从一小块剩余内存预算里抠。我们已经稳稳进入”权重而非 KV 是瓶颈”的区间。这条直觉很关键：GQA 的任务是把 KV 项压到权重项之下；一旦做到，再压收益就急剧递减。

为什么 5% uptraining 就够？ 因为均值池化是一个出奇好的初始化。均值后的 KV 投影对输入 $x_t$ 的作用是 $K_g x_t = \frac{1}{H/G} \sum_h K_h x_t$——也就是组内每个头 K 输出的均值。下游注意力层原本基于 $H$ 个头的拼接训练，现在看到每组内被池化的注意力值。均值池化精确保留了组内 K、V 输出的一阶矩。下游层是注意力输出的线性组合（$W^O$ 投影）——一阶意义下这些线性组合没变。丢的是高阶结构（方差、单头特异性），而那部分正是 5% uptraining 在补。

配方在不同尺度下怎么变？ 大模型有两个相反的效应。(a) 大模型 $d_{model}$ 更大，权重读增长比 KV 读增长快（KV 只随 $d_k = d_{model}/H$ 缩放）。$d$ 越大，KV 占带宽的比例越小。(b) 大模型常配长序列，KV 绝对值更大。论文（以及 LLaMA-3 时代）的结论：净效应近似尺度不变——GQA-8 在各种尺度上都是好选择。

关联工作：GQA 之后又发生了什么

GQA 发表在 2023 年中，领域没有停下来。下面这条简短的路线图能把 GQA 在更大的故事里放到位。

LLaMA 2（2023 年 7 月）。 Meta 的开源权重模型在 34B 和 70B 上采用 GQA，$G = 8$。这是它在生产规模上的首次大规模部署，让”GQA = 默认”这个观念在社区扎下来。7B 与 13B 仍是 MHA，事后看是个失之交臂的机会——小规模下 GQA 也是明显赢家。

Mistral 7B（2023 年 10 月）。 用 GQA $G = 8$、$H = 32$，即更小的 $H/G = 4$。Mistral 还组合了滑动窗口注意力（每个 token 只关注前 $W = 4096$ 个 token），把 KV 缓存大小从 $O(L)$ 进一步压到 $O(W)$。两者协同：GQA 减少 每个位置 的 KV 流量，滑动窗口减少 位置数。

LLaMA 3（2024 年 4 月）。 所有规模（8B、70B、405B）都用 GQA。8B 是 $H = 32, G = 8$；70B 是 $H = 64, G = 8$；405B 延续。”$G = 8$ 跨尺度不变” 是社区共同收敛到的甜点。

DeepSeek-V2 / V3 + MLA（2024 年 5 月起）。 MLA（Multi-head Latent Attention）是下一代思路：不再把 $H$ 个 KV 头分组成 $G$ 个，而是把 KV 整体压成一个低秩潜变量 $c_t \in \mathbb{R}^{d_c}$，$d_c \ll H d_k$。decode 时只缓存 $c_t$，每头 K、V 通过小矩阵乘 on-the-fly 重建。这是比 GQA 更激进的结构先验——DeepSeek-V2 实验显示同等 KV 缓存预算下 MLA 优于 GQA。代价是单步计算更多、内存流量更少；在内存受限的 decode 上这是赚的。

投机解码（speculative decoding）。 GQA 与投机解码组合干净——目标模型一次验证 $k$ 个草稿 token 时，瓶颈仍是 KV 带宽，GQA 同样有效。事实上 GQA 还提升了投机解码的”加速天花板”——目标模型每尝试一个 token 都更便宜。

Paged Attention（vLLM）。 vLLM 的 PagedAttention 把 KV 缓存按固定块管理以提高显存利用率。块大小按 KV 头粒度切；GQA 的 $G$ 个 KV 头让块更小，打包效率更高。GQA 间接也帮了内存分配器。

总结：GQA 是一个能与几乎所有别的推理优化叠加的原语。它的耐用性来自 “处在正确的抽象层”——不是注意力效率的终极答案，但其他东西都能在它之上往上盖。

局限性与适用边界

GQA 不是免费午餐。几点限制：

1. 质量差距虽小但非零。 0.1 平均分对任一单数据集都在噪声范围内，但跨多个数据集一致，所以很可能是真实但小的回退。如果你工作在最后 0.1 个质量分比 5× 解码加速更重要的场景，留在 MHA 是合理的。

2. 配方的适用范围有限。 T5 是编码-解码架构，论文只把 GQA 用在解码自注意力和交叉注意力上。Decoder-only 模型（如今主流）下，配方直接迁移到解码自注意力。很长编码输入下的交叉注意力是个未充分研究的角落。

3. Uptraining 需要原始训练栈。 你需要原优化器状态、数据加载器、训练 schedule。如果只有权重而没有基础设施，转换会困难得多。（前沿实验室的解法：直接从头训 GQA。）

4. KV 缓存不是唯一瓶颈。 GQA 把 KV 流量压到权重流量之下后，继续压就是浪费。这意味着 GQA 的收益对小模型 + 长上下文（KV 主导）最大，对大模型 + 短上下文（权重主导）最小。

5. 分组方式是固定的。 论文用连续分组（头 1..H/G 进第 1 组等等），没有探索可学习分组。后续工作也没有强证据显示学习分组比连续分组明显更好——经验上分组方式的选择似乎远不如 $G$ 的选择重要。

6. GQA 假设 K、V 应该一起共享。 每组同时有一个 $K_g$ 和一个 $V_g$。论文没消融”$G_K = 8$ 但 $G_V = 16$“这类不对称设计。后续工作（如 DeepSeek 的 MLA）开始探索非对称 KV 压缩。

批判性分析：不足与可改进之处

前面几节总结了论文的主张和它验证得好的地方。这一节反过来看——哪些证据比头条数字暗示的更弱、作者没有测试什么、以及一个更扎实的版本应该长什么样。

实验设计的不足

整个实证论证只建立在一个架构家族（T5）和一类任务（长编码器输入的条件生成）上。 主表里的每一个数字——47.2 / 47.1 / 46.6 的平均分，1.51 / 0.28 / 0.24 秒的时延——全部来自 T5 Large 和 T5 XXL，这是一个在论文 2023 年发表时相对 decoder-only 的 GPT 系模型已经偏小众的编码-解码架构。GQA 如今几乎只用在 decoder-only 模型（LLaMA、Mistral、Qwen）上，论文里却没有一个 decoder-only 实验。论文把从编码-解码交叉注意力（K/V 序列是编码器输出，长度固定、只计算一次）推广到 decoder-only 自注意力（K/V 序列随每个生成 token 增长、必须增量重计算）当成了一件理所当然的事，但这两种设定对 KV 缓存的压力完全不同——交叉注意力的 K/V 在 decode 过程中从不增长，所以“GQA 省带宽是因为 KV 缓存会不断变大”这个论证，对论文自己实验设置的一半（交叉注意力）甚至不适用。真正对今天读者最重要的自注意力数字，论文完全没有报告，这些数字是多年后由社区（LLaMA 2、Mistral）重新验证出来的，而不是由作者本人。

论文任何地方都没有报告置信区间、方差或随机种子。 表 1 里的每一个数字（47.2/47.1/46.6 的对比）都是每个配置仿佛只跑了一次训练/uptraining 得到的单点估计。论文自己的头条主张——“GQA-8 相对 MHA 只掉 0.1 分”——是在一个把七个不同任务分数平均起来的指标上的1分差异；论文没有对 MHA 或 GQA-8 检查点报告任何 run-to-run 方差，所以单凭论文本身无法判断 0.1 分是否在单次训练的噪声带内。考虑到 uptraining 只跑了预训练 5% 的步数、起点又是一个被扰动过的检查点，run-to-run 方差几乎肯定存在，如果这个方差恰好是 ±0.2 分这个量级，论文的核心卖点（“几乎免费的质量”）就会大大打折。

$G$ 的消融实验（$G \in \{1,2,4,8,16,32,64\}$）只在一个模型规模（T5 XXL）上跑过，而且只看平均分，不看分任务。 很有可能某些单个任务（比如更依赖事实回忆而非摘要的 TriviaQA）对 $G$ 的敏感度曲线与平均曲线很不一样。将七个差异很大的任务（摘要、翻译、问答）平均后再报告消融数字，会掩盖掉任务特异的“悬崖”——后续压缩文献（如 SliceGPT、GQA 后续工作）已充分记录过这种现象。缺少专门针对 $G$-扫描的分任务数据（论文只对头条 MHA/MQA/GQA-8 对比给了分任务数据，对完整 $G$ 扫描没有），工程师根本不知道自己的下游任务是靠近悬崖还是安稳地处在平坦部分。

论文从未控制 MHA、MQA、GQA-8 三种配置的总训练算力。 GQA-8-XXL 和 MQA-XXL 都从同一个 MHA-XXL 检查点出发、接受相同的 uptraining 预算（$\alpha=5\%$），但 MHA-XXL 本身却拿到了 100% 的原始预训练预算，没有留一块给“免费的”继续训练。这意味着表 1 的对比并不是“三个等算力检查点”，而是“一个训 $T$ 步的检查点 vs 两个训 $T+0.05T$ 步的检查点”。多出来的 5% 梯度步不是免费的——比 MHA-XXL 基准拿到的总算力还多。更干净的消融应该是把 MHA-XXL 基准也继续多训同样额外的 5% 步数（不改架构），看这个额外训练本身是否就已经补上了大部分差距。论文没有报告这个对照组，所以挑剔的读者无法把“GQA 架构本身好”和“多 5% 训练对什么都有点帮助”完全分开。

作者淡化或回避的局限

论文完全没提从头训练 GQA 与 uptrain 后 GQA 的对比。 论文所有 GQA 结果都来自 uptraining 一个 MHA 检查点。论文没有报告如果用完整预训练预算从头训一个 GQA 模型会怎样。这很重要，因为 uptraining 隐式地占了 MHA 检查点已经学到的表示的便宜；从头训的 GQA 模型得在没有这个起点的情况下自己学会组内共享 KV 表示。论文自己的核心主张（“GQA 能拿回几乎全部 MHA 的质量”）严格来说只在 uptrain 这个体制下成立，但到 GQA 成为行业默认（LLaMA 3、Mistral）时，几乎所有生产部署都是从头训 GQA。论文的核心实证主张和 GQA 在实际中被使用的方式之间存在微妙的错位，论文没有明确指出这个差距。

张量并行那部分的论证是断言，不是实测。 论文关于 TP 切分的那节内容声称 $G$ 是 TP 度数 $K$ 的倍数就能避免 KV 复制开销，但论文没有报告任何 TP 切分下 MQA/GQA-8/MHA 的实测吞吐量或内存数据。这整个 TP 效率论证（包括本文前面重复的那一段）只是从第一原理推导的可信度论证，不是实验。考虑到生产中选 $G=8$（而不是 $G=4$ 或 $G=16$）的动机很大程度依赖于这种 TP 切分干净性，缺少直接测量是个真实的缺口。

Uptraining 预算扫描（$\alpha \in \{0,1,5,10\%\}$）止于 10%，而这个点上 MQA 的曲线还在（微弱地）上升。 论文选 $\alpha=5\%$ 作为默认值，却没测试过 $\alpha=15\%$ 或 $\alpha=20\%$ 是否能让 MQA 完全追上 MHA——这会动摇论文隐含的“MQA 有不可挽回的质量天花板，而 GQA 避开了它”这个主张。也许 MQA 只是需要更多 uptraining 而不是更多 KV 头；论文的实验设计（扫描停在 10%）无法区分这两种假设。

论文没有讨论批处理、高吞吐量服务下会发生什么——而这现在已经是主流部署体制。所有时延数字都是单个 TPU 上 batch=1 的 decode 延迟。正如本笔记自己的“成本账”一节所示，KV 缓存在总内存预算里的占比（进而 GQA 的提速比例）强烈依赖 batch 大小、序列长度、硬件——这些论文一个都没扫。batch=1 是对 GQA 最有利的情况（其他开销如权重读取被摊薄得越少，KV 占字节预算的比例就越大）；在大 batch 下权重读取被多个请求摊薄，GQA 的相对收益可能完全不同，论文没有提供这一点的任何数据。

可以改进的地方

1. 报告 run 间方差。 用至少 3 个不同随机种子重跑 T5 XXL 的 MHA/MQA/GQA-8 训练与 uptraining，报告平均分 ± 标准差。没有这个，论文的核心定量主张（“只掉 0.1 分”）无法与噪声区分。
2. 加一个算力对齐的 MHA 对照组。 把 MHA-XXL 基准也多训同样额外 5% 步数（不改架构），看这是否已经能补上部分 GQA/MQA 看似领先的质量优势。这是论文缺失的最重要的一个消融实验。
3. 完整跑一个 decoder-only 架构。 哪怕是个 GPT-2 量级的小模型（2023 年完全可行），按相同协议训练+uptrain，就能直接验证论文隐含的“交叉注意力的发现能迁移到自注意力”——而后者才是真正影响整个领域的那个设定。
4. 直接实测 TP 切分下的吞吐量，而不是从第一原理推论 ”$G$ 是 TP 度数的倍数就能避免复制开销”。把 MQA、GQA-4、GQA-8、MHA 都切到 TP=2/4/8 上实测吐吐量和内存，就能把论文只断言的事情定下来。
5. 把 uptraining 预算扫描延伸到 10% 以上，专门针对 MQA，确定 MQA 的质量天花板到底是真实的架构极限（KV 头太少无论怎么训都不够表示信息），还是只是收敛慢、多训一会儿就能追上。这区分了“MQA 架构上就更差”和“MQA 只是需要更大 uptraining 预算”，而当前实验设计无法区分这两种情况。

复现性与工程实践

工程师今天想上 GQA，注意以下事项：

• PyTorch： 2.0 起 F.scaled_dot_product_attention 接受 query 形状 $[B, H_q, L, d_k]$、key/value 形状 $[B, H_{kv}, L, d_k]$（$H_q$ 是 $H_{kv}$ 的倍数），原生支持广播。FlashAttention 2 原生支持 GQA。
• vLLM、TGI、TensorRT-LLM： 所有主流推理框架开箱支持 GQA。基本只需设 num_key_value_heads = G。
• 转换脚本： PyTorch 里把 MHA 检查点转 GQA-$G$ 大致就是：

W_K_mha = state_dict["k_proj.weight"]            # [H * d_k, d_model]
W_K_mha = W_K_mha.view(H, d_k, d_model)          # [H, d_k, d_model]
W_K_gqa = W_K_mha.view(G, H // G, d_k, d_model).mean(dim=1)
state_dict["k_proj.weight"] = W_K_gqa.reshape(G * d_k, d_model)

V 同理。然后正常继续训练。

• $G$ 怎么选。 默认 $G = 8$（$H \geq 16$ 时）。要做 TP=$K$ 切分时确保 $G$ 是 $K$ 的倍数。常见生产配置：
  • LLaMA-2 70B / LLaMA-3 70B：$H = 64, G = 8$
  • Mistral 7B：$H = 32, G = 8$
  • LLaMA-3 8B：$H = 32, G = 8$
  • Qwen2.5 72B：$H = 64, G = 8$
• 长上下文： GQA 的重要性随上下文线性放大。128k 上下文下 MHA 70B 的 KV 缓存已无法部署在普通硬件；GQA-8 让它变得可行。

下面这些工程注意点论文没说但很重要：

KV 缓存内存布局。 MHA 自然布局是 $[L, H, d_k]$ 或 $[B, H, L, d_k]$；GQA 是 $[L, G, d_k]$。decode 时 Q 还是 $[B, H, L, d_k]$ 但 $L = 1$。内核需要在 $H$ 个 Q 头维度上广播 $G$ 个 KV 头。现代注意力内核（FlashAttention-2 的 GQA 支持、xformers）通过寄存器层或循环 tiling 处理这种广播，不会真的物化 $[B, H, L, d_k]$ 的 KV 张量——那样会抵消带宽节省。

混合精度。 现代服务系统经常把 KV 缓存存 FP8 甚至 INT4。GQA 可叠加：$G = 8$ + FP8 KV，每步 KV 带宽比 FP16 GQA-8 再小 2 倍。两个收益是相乘的：GQA-8 + FP8 = $8 \cdot 2 = 16$× 相对 MHA + FP16。

投机解码。 一次前向验证 $k$ 个投机 token 时，KV 缓存只读一次，$k$ 个新查询都对它做注意力。带宽成本主要在读 $L$ 个已缓存条目上，GQA 的节省继续生效。GQA 还增加了”每尝试一个 token 的成本”的弹性，给投机解码留出更多加速空间。

训练成本。 从头训 GQA 比训等大 MHA 更便宜：K、V 投影参数缩到 $G/H$，每步 FLOP 略减。70B（$H = 64, G = 8$）下每步训练 FLOP 大约低 3%——小但不为零。

参数节省。 GQA 略微减少参数总量。一层中 MHA 的 K+V 参数是 $2 \cdot d_{model} \cdot H \cdot d_k$，GQA 是 $2 \cdot d_{model} \cdot G \cdot d_k$。LLaMA-2 70B（$d_{model} = 8192, H = 64, G = 8, d_k = 128$）每层省 $2 \cdot 8192 \cdot (64-8) \cdot 128 = 117$M 参数，80 层共 9.4B。绝对值大，但相对 70B 不算多——参数主体在 MLP，不在注意力。

转换的工程细节。 在 HF Transformers 里转换 MHA 检查点除了对 $W^K, W^V$ 做均值池化，还要更新模型 config（num_key_value_heads = G），让运行时知道按 $G$ 广播。多数架构共用这一字段（LLaMA、Mistral、Mixtral、Qwen 都用 num_key_value_heads）。老一些的代码可能叫 n_kv_heads。

数值影响。 均值池化后 $W_g^K$ 的范数比原本每头的 $W_h^K$ 范数小约 $\sqrt{H/G}$（独立性假设下）。这把注意力分数幅度往下推，softmax 后分布更平。实操中影响很小，5% uptraining 会吸收掉；但极端转换（$G$ 极小、零 uptraining）下，分数幅度的下移会以微妙方式影响生成质量（注意力过于均匀 → 输出过度平滑）。

从带宽到资源占用：另一个视角

GQA 通常以”减少 KV 带宽”被讲述，但它同样在另一个维度上提供红利：KV 缓存占用——这进一步决定了单卡能装多大的 batch、能服务多少并发用户。

举例：A100 80 GB 卡部署 LLaMA-2 70B（fp16 权重约 140 GB，需要 TP=2 才能装下，本节假设 TP=2 后每卡 70 GB 权重）。剩余 HBM 约 10 GB 留给 KV 缓存与激活。

• MHA 70B：KV 缓存约 2.7 MB / token / 请求（按 $H = 64, d_k = 128$ 算）。剩余 10 GB 能装 $\sim 3700$ token·请求的 KV。如果上下文 $L = 4096$，并发数最多 $\sim 0.9$——根本撑不到 1 个完整请求。
• GQA-8 70B：KV 缓存约 0.34 MB / token / 请求。剩余 10 GB 能装 $\sim 29700$ token·请求。$L = 4096$ 下能容纳 $\sim 7$ 个并发请求。

差 8 倍——直接体现在能服务多少用户上。Paged Attention / continuous batching 等机制能进一步把这种容量收益放大成吞吐收益。换句话说：GQA 不仅让单请求 decode 更快，更让多请求并发的经济性出现质的提升。

我的看法

GQA 是近年 LLM 文献里少见的 “小想法、大影响” 论文。核心洞察——你不必在 $H$ 和 $1$ 个 KV 头之间二选一，可以选中间任何因数——事后看来显得太显然了。技术深度藏在 uptraining 配方里，特别是均值池化的论证：它给出了把已有 MHA 检查点低成本转 GQA 的有原则的方法。

它有外溢影响的几个原因：

• 改动量极小。 不引入新算子、新激活、新优化器。每层只改两个权重矩阵的形状。在现有代码库里 retrofit 进去几乎没成本。
• 与张量并行干净组合。 $G$ 个头在 $G$ 个 TP rank 上切分，零复制、零额外 all-reduce、零边角 case。
• 能和别的一切组合。 GQA + FlashAttention 行；GQA + Paged KV（vLLM）行；GQA + 投机解码行；GQA + 滑动窗口（Mistral）行；GQA + 量化 KV 行。
• 击中了设计空间的甜点。 $G = 8$、$H = 64$ 能拿回几乎全部 MHA 质量、同时拿到几乎全部 MQA 速度。“8” 这个魔数能在多个模型尺度上稳定胜出，提示它反映的是注意力的结构性属性，而非调参产物。

论文没完全回答、领域之后正在推进的几个方向：

• 非对称 KV。 为什么 K 和 V 必须用同样的分组数？DeepSeek-V2 的 MLA 用低秩潜变量联合压缩 KV，是比 GQA 更激进的下一步。
• 可学习分组。 没有强证据证明学习分组比朴素连续分组更好。这本身就有意思——它说明在”组级”粒度上头是大致可互换的。
• 超越”恢复预训练”。 均值池化论证在”组内头本来就相关”这个假设上最强。从头训的 GQA 不依赖这条假设，但 GQA 仍然 work。为什么 KV 头的边际价值这么低——这个更深的问题论文没回答。

经验事实是 GQA 在 $G = 8$ 处几乎是免费的，这是 LLM 推理领域的悄然革命。它是今天长上下文服务（32k、128k）在当前价格点上可行的核心原因之一。结合 FlashAttention、Paged KV、投机解码，GQA 是现代 LLM 推理的四根支柱之一。这篇论文短、目标明确、经得起时间检验。

参考文献

• Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., & Sanghai, S. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023. arXiv:2305.13245。
• Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. arXiv:1911.02150。
• Vaswani, A. et al. (2017). Attention Is All You Need. NeurIPS 2017。
• Raffel, C. et al. (2020). Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. JMLR。
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691。
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention（vLLM）. SOSP 2023。
• Touvron, H. et al. (2023). LLaMA 2: Open Foundation and Fine-Tuned Chat Models. arXiv:2307.09288。
• DeepSeek-AI (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arXiv:2405.04434。
• Jiang, A. et al. (2023). Mistral 7B. arXiv:2310.06825。
• Pope, R. et al. (2022). Efficiently Scaling Transformer Inference. MLSys 2023.（PaLM 推理论文——奠定了 GQA 立足的”带宽受限 decode”分析框架）。