AI 大模型之 GQA

原创

ITPAPA
ITPAPA

ITPAPA

2026年10月13日 12:38
江苏

在小说阅读器读本章

去阅读

在公众号小说中沉浸阅读

Grouped-Query Attention   GQA

分组查询注意力     技术深化学习

研究范围：基本概念、技术演化、数学及计算机制、模块架构、主要配置参数、训练与推理流程、工程实现、性能评测、典型应用、风险与2026年最新公开进展。

1  技术定位与解决的问题

GQA让若干查询头共享同一组Key/Value头，使推理KV缓存与KV头数量而非Query头数量成正比。它位于MHA（每个Q头一个KV头）和MQA（全部Q头共享一个KV头）之间。

问题背景：大规模模型的结构、计算、显存、训练稳定性与生成能力之间存在不同的取舍。本技术对应的改进点，应通过与恰当的基线比较才能判断。它属于“模型结构或训练/生成机制”，不应把单篇论文的实验数字误解为通用硬件指标。

2  结构组成与数据流

图1  Grouped-Query Attention (GQA) 核心计算/数据流示意

•  步骤1：输入隐藏状态。

•  步骤2：分别投影 Q头与少量KV头。

•  步骤3：按每组Q头映射至同一KV头。

•  步骤4：计算缩放点积注意力。

•  步骤5：输出投影并缓存分组KV。

3  底层原理与关键技术

3.1  机制设计与公式

在同维同精度下，GQA相比MHA的KV缓存比例约为G/H。

3.2  复杂度与状态开销

如果H=32、G=8，每个KV头对应4个Q头，KV缓存约为MHA的1/4。

3.3  硬件执行与精度

减小G可降低KV读写，但可能影响注意力表达能力；精度和吞吐需联合评估。

架构识别原则：需要区分“数学形式”“训练目标”“GPU计算内核”“模型实例配置”“在线服务系统”。同名概念在这些层面的实现方式并不自动相同；评测必须说明层次、模型版本、数据精度、硬件与任务。

4  核心参数及工程含义

参数取值原则：除明确标注“论文示例”的值之外，本表列的是配置项与设计量，而非默认推荐值。实际有效范围受模型结构、训练方式、推理框架、显卡架构与可用显存约束。

4.1  定量公式与参数估算

KV缓存比：GQA/MHA = G/H；每序列KV字节 = 2×L×G×d×S×b。

4.2  可复核的示例计算

示例：H=32、G=8、L=32、d=128、S=4096、b=2，GQA约512 MiB；对应MHA约2048 MiB；节省75%是相同维度下的理论KV数据量比较。

4.3  核心伪代码与执行序列

reshape Q为[B,S,H,d]；reshape K/V为[B,S,G,d]；逻辑广播每G组到H/G个Q头；执行Attention。

4.4  易出错的工程边界

检查组数整除、KV缓存真实形状和Tensor并行分片；避免在模型计算前实际repeat形成无谓复制。

5  实施流程与关键技术细节

5.1  训练准备

在模型config中设置num\_attention\_heads和num\_key\_value\_heads。

5.2  实现验证

对已有MHA权重的KV头做分组聚合，并使用少量继续训练恢复质量。

5.3  测试优化

注意张量reshape/repeat\_kv只影响计算映射，不应复制实际缓存。

5.4  交付评估

按长度、批量和精度绘制质量—带宽—显存权衡曲线。

6  技术验证与量化评估设计

图2  机制—参数—实现—应用的验证路径

实验控制：固定模型权重/训练Token预算、Tokenization、输入与输出长度分布、温度与采样配置、数据类型、GPU与软件版本；每组记录至少一次预热、多轮重复以及波动区间。结果按实际运行配置报告，避免跨研究条件直接比倍率。

7  典型应用场景

•  低延迟线上大模型服务：先建立真实业务负载与基线，再决定是否引入该结构。

•  消费级GPU部署中大型语言模型：先建立真实业务负载与基线，再决定是否引入该结构。

•  针对带宽瓶颈的解码阶段优化：先建立真实业务负载与基线，再决定是否引入该结构。

8  局限 代价与容易误解的地方

•  极端分组可能伤害某些注意力头的专门能力。

•  并非所有模型均支持不重训地修改KV头数。

•  理论缓存节省与实际TPS增益不同，应以真实引擎实测。

9  技术进展 截至2026年10月

2023年EMNLP论文系统提出GQA以及从MHA检查点上训的方案；2026年仍是稠密Transformer解码效率的基础选项。

10  推荐实验记录模板

11  原始论文与权威资料

•  [1] https://aclanthology.org/2023.emnlp-main.298/

说明：所有实验数值应以所引文献的相应版本、配置与对照组为准。本文的示意图为重新绘制的知识结构图，不能作为作者官方架构图。

技术深化 分组查询的张量映射与权重转换

本章在原有概览基础上展开推导与实现。文中的尺寸和计算例为教学配置，实验代码验证明确标注的简化核心；论文方法、工程建议和生产实测分别说明。

分组不是把查询头合并

设查询头H=8，键值头G=2，每组有4个查询头。Q保持[B,H,Sq,d]，K和V保持[B,G,Sk,d]；第i个查询头读取第floor(i/4)个键值头。每个查询头仍计算独立softmax，所以GQA主要减少KV表示与读取，不把注意力配对数量缩减为G/H。[1]

M为加性掩码，可见位置为0，不可见位置为负无穷。预填充时Sq与Sk通常相同；增量解码时Sq=1而Sk为完整历史长度。因果遮罩要用查询的绝对位置判断，直接对一个1×Sk矩阵使用左上角三角遮罩可能只允许读取第一个键。

从MHA检查点转为GQA

把原始K投影按头重排为[H,d,D]，分为G组后在组内求均值；V投影同样处理，Q与输出投影通常保持原结构。若含偏置，对偏置按相同组映射聚合。均值只是初始化方法，不能保证原模型函数不变，需要继续训练并验证质量。[1]

权重的实际存储方向取决于框架Linear约定；错误地在输入维上平均，会得到尺寸合法但语义错误的矩阵。旋转位置编码应作用在正确的键头维度；查询与键的旋转参数一致并不意味着两者头数必须相等。

缓存省下多少以及何时不省

总缓存字节为2BLGSdb，其中L为层数、b为元素字节数。H=32、G=8时，持久化KV量是MHA的四分之一。理论上单步KV读取也减少，但内核是否把共享键值重用到寄存器或共享内存，决定真实带宽收益。

如果先repeat\_interleave并把展开后的K/V写入缓存，实际上又保存了H份头。教学参考代码可以展开来做数值对照，生产内核应直接索引共享头。张量并行还要求检查G与并行度：G小于设备数时，某些实现会复制KV头，单卡缓存节省会小于理想比例。

最小参考实现

随包02\_gqa.py使用显式头映射和逻辑展开两种方法对照，输入为单条序列。两种方法相等只能证明索引实现正确，不能证明GQA和MHA质量相同。

out = []

for i in range(H):

g = i // (H // G)

scores = Q[i] @ K[g].T / np.sqrt(d)

out.append(softmax(scores) @ V[g])

Kr = np.repeat(K, H // G, axis=0)

Vr = np.repeat(V, H // G, axis=0)

reference = softmax(Q @ Kr.transpose(0, 2, 1) / np.sqrt(d)) @ Vr

assert np.allclose(np.stack(out), reference)

再增加padding、非方形查询、跨页缓存和张量分片四组边界样例。测试时应包含各组数值明显不同的KV头，避免错误映射在随机小样本中被低敏感度输出掩盖。

配置与排障

配置文件、投影矩阵行数、推理缓存布局三者必须一致。仅改num\_key\_value\_heads而不转换权重，会加载失败或错误解释维度。速度不升时，先确认内核分派是否退回通用实现，再查看是否由MLP、权重读取或通信主导。长上下文质量下降时，同时检查权重转换训练与RoPE外推，不能全部归因于键值共享。

消融实验与判定

在相同模型宽度下设置G为1、2、4、8、H，记录训练损失、任务准确率与单位有效输出成本。架构实验应从训练或继续训练阶段控制预算；执行实验只比较同一个GQA检查点的广播和直接索引路径。短序列批量1用于暴露启动开销，长序列高并发用于暴露缓存带宽，最终选择满足质量要求的最小G。

图3 教学配置为32层 4096词元 BF16 不含权重和管理开销

本册配套实验验证记录

配套文件为 实验代码/02\_gqa.py。本次使用NumPy CPU固定随机输入运行，执行通过。结果验证本章所标注的公式、索引、容量或事务边界，不代表完整模型训练效果或GPU性能实测。

max\_abs\_error 0.0

工程实操 非方形因果遮罩与头分片检查

固定输入与检查目标

查询来自绝对位置5和6，键包含位置0至6。设置8个查询头、2个KV头，每个KV头服务4个查询头。这个例子同时检查长前缀增量输入与分组索引。

逐步计算与预期结果

1  因果条件是键位置不大于查询位置。位置5的查询能读6个键，位置6的查询能读7个键。

2  若直接对2×7矩阵取左上角下三角，只得到1个和2个可见键，尺寸虽然正确，语义已经改变。

3  查询头0至3对应KV头0，查询头4至7对应KV头1。测试数据应让两个KV头明显不同，避免映射错误被相同数值掩盖。

4  加入padding后，应取因果条件与键有效性条件的交集；padding位置编号不能替代有效长度。

接口与实现约束

模型配置与缓存接口应分别记录查询头数和KV头数。执行端要明确KV头在设备上分片还是复制，返回的布局必须与投影权重行数相容。

先在无padding场景验证绝对位置，再加入左padding、右padding和多请求不同前缀长度。非方形causal的对齐语义应随所选内核版本记录，不能以接口参数名推断结果。

故障核查表

扩展验证组合

使用Sq=1、2和Sk=7、17，测试G=1、2、H及不支持的H/G组合；不支持的配置应明确拒绝，不能静默映射。

配套代码与结果

进阶实验/boundary\_cases.py 中的 c02 函数对应本节。可在进阶实验目录运行 python boundary\_cases.py，或执行 python -c \"from boundary\_cases import c02; print(c02())\" 单独检查。本节固定样例执行通过，结果保存在边界实验结果.json。

验证范围为正文明确标注的简化机制与错误反例。工程接口、扩展测试和故障注入属于本资料提出的实施建议；完整模型训练、GPU数值与吞吐、跨进程恢复仍须在目标环境验证。

训练与部署 共享键值的梯度归约

目标与验证范围

GQA让多个查询头读取同一个KV头，因此共享K/V的梯度是所有使用者贡献的总和。教学检查使用4个查询头、2个KV头，每个共享头对应2个查询头。

解析关系与梯度路径

先为每个查询头独立计算attention反向，再按g(i)把dK和dV累加到共享头。查询梯度仍属于各自查询头，不能一并合并。

本例共享K梯度范数约0.879371。如果在共享头归约时错误求平均，范数变为0.439686；损失中原本是否包含头平均，必须在上游梯度处决定，不能在缓存共享层再平均一次。

训练配置与误用边界

训练损失归一化、张量并行通信归约和KV共享归约是三个不同位置。先在单设备验证总和语义，再核对各rank是否重复累加或漏掉贡献。

MHA到GQA的均值初始化改变权重结构，并非这里的梯度均值操作。继续训练时保存转换映射、优化器新参数状态及训练预算，避免两种平均被混为一谈。

图4 共享键梯度需要汇总全部使用头的贡献

部署验收与恢复

缓存只在有效KV头维存储；教学repeat参考允许展开，但反向归约和部署缓存布局分别检查。变换维度时也需保留原头顺序。

验收以共享头梯度与逻辑展开后反向的归约结果对照，覆盖不同组数、mask和分片布局。执行收益与模型转换质量分开报告。

配套检查与本次数值结果

训练部署实验/training\_checks.py 中的 t02 函数已执行通过。本例解析梯度与中心差分的最大绝对误差为 2.720e-10。完整输入、解析实现和实际输出见代码及训练与状态验证结果.json。

可在训练部署实验目录执行 python training\_checks.py，或执行 python -c \"from training\_checks import t02; print(t02())\"。所示结果来自NumPy CPU教学验证；目标模型的自动微分、低精度和分布式执行需另做对应验收。