GQA 分组查询注意力：KV Cache 与模型质量的折中

大语言模型逐 token 生成时，会把历史 token 的 Key 和 Value 保存在 KV Cache 中，避免每轮重复计算。这个缓存换来了更快的生成，也逐渐成为长上下文和高并发服务的主要显存开销。

Grouped-Query Attention（GQA）解决的就是这组矛盾：保留较多 Query heads，让模型仍能从多个角度检索信息；同时减少 Key/Value heads，让多个 Query heads 共用一组缓存。它处在传统 Multi-Head Attention（MHA）和共享最彻底的 Multi-Query Attention（MQA）之间。

理解 GQA 不能只记住“KV heads 更少”。还要回答三个问题：不同 head 在共享什么，显存究竟省在哪里，以及训练好的 MHA 权重为什么不能只改一个配置数字就变成 GQA。

多头注意力中的 Q、K、V

对输入隐藏状态 $X\in\mathbb{R}^{B\times S\times D_{model}}$，注意力层通过三个线性投影得到 Query、Key 和 Value：

多头注意力再将它们拆成多个 head。对第 $h$ 个 head：

Query 表示“当前 token 要寻找什么”，Key 表示历史 token 可被怎样匹配，Value 则是匹配后真正汇总的信息。多个 Query heads 可以学习不同检索方式，例如实体关系、位置和语法；这里的语义只是帮助理解，模型并没有被人工规定每个 head 必须负责某类知识。

传统 MHA 为每个 Query head 配置独立的 Key 和 Value head：

Q0 ──► K0, V0
Q1 ──► K1, V1
Q2 ──► K2, V2
Q3 ──► K3, V3

这给每个 head 最大自由度，但生成时也要为每一层、每个历史 token 保存全部 $K_h$ 和 $V_h$。

KV Cache 为什么会不断变大

假设提示词含 1000 个 token。Prefill 先处理这 1000 个输入，建立它们在每一层的 K/V，并用最后一个位置的 logits 采样第 1001 个 token。此时第 1001 个 token 已经成为输出，但它自己的 K/V 还没有经过模型计算。

若生成尚未结束，下一次 decode 把已采样的第 1001 个 token 作为输入，只为它计算新的 Query、Key、Value。它的 Query 读取前 1000 个 token 以及自身的 K/V，得到用于采样第 1002 个 token 的 logits。之后重复这一过程：每处理一个新输入，缓存追加一份状态，并预测下一个输出。Hugging Face 的缓存生成循环正是先调用模型，再把选出的 next_token_ids 作为下一轮输入。

单条序列、单层 KV Cache 的常见形状为：

K: [sequence_length, num_kv_heads, head_dim]
V: [sequence_length, num_kv_heads, head_dim]

若层数为 $L$、上下文长度为 $S$、KV head 数为 $H_{kv}$、head dimension 为 $d_h$，每个元素占 $b$ 字节，则缓存主体约为：

系数 2 来自 K 和 V。这个式子揭示了几个直接事实：

• 上下文翻倍，缓存近似翻倍；
• 并发序列数翻倍，所有序列的缓存总量近似翻倍；
• 在其他条件不变时，KV head 数从 32 降到 8，缓存主体降为四分之一。

最后一条只适用于 KV tensor，不等于整张 GPU 的占用也减少 75%。权重、激活、采样缓冲区、CUDA Graph、内存块元数据与碎片仍然存在。

MHA、MQA 与 GQA 的结构关系

设 Query head 数为 $H_q$，KV head 数为 $H_{kv}$：

以 $H_q=8$、$H_{kv}=2$ 为例，每 4 个 Query heads 组成一组：

Q0 ┐
Q1 ├──► K0, V0
Q2 ┤
Q3 ┘

Q4 ┐
Q5 ├──► K1, V1
Q6 ┤
Q7 ┘

组大小为：

常见实现要求 $H_q$ 能被 $H_{kv}$ 整除，第 $h$ 个 Query head 对应的 KV head 为：

注意，GQA 没有把 Query heads 也减少。Query 仍能保留不同投影，变化发生在 $W_K, W_V$ 的输出维度以及缓存布局上。

一次 GQA 计算如何展开

常见张量布局为：

Q: [batch, H_q,  sequence, head_dim]
K: [batch, H_kv, sequence, head_dim]
V: [batch, H_kv, sequence, head_dim]

为了复用普通 MHA 代码，教学实现可能先把 K/V 在 head 维逻辑扩展到 $H_q$：

def repeat_kv(x, num_groups):
    # x: [batch, H_kv, sequence, head_dim]
    b, h_kv, s, d = x.shape
    x = x[:, :, None, :, :].expand(b, h_kv, num_groups, s, d)
    return x.reshape(b, h_kv * num_groups, s, d)

expand 本身可以创建零步长 view，但随后的 reshape、kernel 或 contiguous 操作是否产生物理副本，要看具体实现。高性能 GQA kernel 会直接根据 head 映射读取对应 K/V，避免在热路径把整份 KV Cache 真正复制 $G$ 次；否则节省缓存的主要优势会被抵消。

对一个 Query head，数学计算仍是：

同组 Query heads 使用相同的 $K_{g(h)}, V_{g(h)}$，但因为各自的 $Q_h$ 不同，得到的注意力权重和输出仍然可以不同。“共享 KV”并不等于同组 head 产生相同结果。

用一个具体模型估算缓存

假设模型具有：

• 32 层；
• 4096 token 上下文；
• head dimension 为 128；
• KV Cache 使用 BF16，每元素 2 字节。

单条序列的缓存主体为：

结果约为：

如果同时服务 16 条满长度序列，MHA 的理论 KV 主体会达到约 32 GiB，而 8 KV heads 的 GQA 约为 8 GiB。真实引擎采用分页分配，不一定每条请求都填满最大长度；但这个计算足以说明，减少 KV heads 会直接提高可容纳的 token 数和并发量。

缓存减小还带来第二项收益：decode 每轮都要读取历史 K/V，较少的 KV heads 能降低显存带宽压力。对只生成一个新 token 的阶段，读取缓存经常比 Query 投影的计算更关键。

为什么不全部改成 MQA

如果只看缓存公式，$H_{kv}=1$ 总是最省。但 Key/Value 投影承担着对历史信息进行编码的职责，过度共享可能限制模型质量。不同 Query heads 虽然提出不同问题，却只能面对同一套 Key/Value 表示。

GQA 提供连续的设计空间：

• KV heads 多一些，表示容量更接近 MHA；
• KV heads 少一些，缓存和读取成本更接近 MQA。

最合适的分组数不能从公式直接推出。它与模型规模、训练数据、上下文长度、目标任务和服务硬件有关，需要在训练质量与推理 workload 上共同验证。

从已有 MHA checkpoint 转成 GQA

MHA 的 $W_K, W_V$ 包含 $H_q$ 组输出，GQA 只需要 $H_{kv}$ 组。把配置文件中的 num_key_value_heads 改小，会造成权重形状不匹配；即使强行裁剪，也会丢掉模型已经学习到的信息。

GQA 论文提出的 uptraining 流程，先把同组的 MHA K/V heads 做 mean pooling，作为新 GQA 权重的初始化：

然后继续预训练，让模型适应共享后的表示。这个过程可以理解为：平均只是把多份已有知识压到一份起点，后续训练才让 Query 与新的 K/V 分工重新协调。

论文实验使用了原始预训练计算量 5% 的 uptraining 配方。这个数字属于论文的特定模型和实验条件，不是“任何 MHA 模型只需训练 5% 就能无损转换”的通用定律。分组方式、数据混合、学习率与训练预算都会影响结果。

与位置编码和 KV Cache 的关系

使用 RoPE 时，Query 和 Key 会根据 token 位置做旋转，Value 通常不做。GQA 减少的是 Key heads 数量，并没有取消位置编码；写入缓存的 Key 必须与推理框架采用的 RoPE 约定一致。

缓存还涉及位置推进与请求状态：

1. prefill 为提示词批量计算 K/V；
2. 将上一步已经选出的 token 作为下一次 forward 的输入，计算并追加它的 K/V；
3. attention 用这个输入的新 Query 读取截止当前位置的缓存，再产生用于选择下一个 token 的 logits；
4. beam search、推测解码或前缀共享可能复制、共享、提交或回滚缓存块。

GQA 只改变每个 token 的 K/V 形状，不替代引擎对这些生命周期的管理。

与 Tensor Parallel 的交互

Tensor Parallel（TP）通常把 attention heads 分到多个 GPU rank。若 $H_q=32, H_{kv}=8, TP=4$，每个 rank 可以自然持有 8 个 Query heads 和 2 个 KV heads。

当 TP rank 数大于 KV head 数，例如 $H_{kv}=2, TP=8$，就无法让每个 rank 都拿到不同 KV head。框架可能在多个 rank 之间复制 KV heads，或采用更特殊的切分与通信方式。由此会出现两个后果：

• 单 rank 的缓存下降不再严格等于 $H_{kv}/H_q$；
• kernel 必须正确处理本地 Query head 到复制 KV head 的映射。

因此，模型配置不能脱离部署拓扑讨论。计划使用较大 TP 时，应先检查引擎对该 num_key_value_heads 的约束和复制策略。

训练与部署时应验证什么

从模型角度，需要比较转换前后的 perplexity 和目标任务质量，尤其关注长上下文检索、代码、数学与多轮对话，而不是只看平均分。

从系统角度，至少分别记录：

• prefill latency；
• decode 的 time per output token；
• 每 token 的 KV Cache 字节数；
• 给定显存下可容纳的并发 token 数；
• attention backend 是否真正支持该 GQA shape；
• TP 后是否发生 KV head 复制。

配置检查可以从模型文件开始：

from transformers import AutoConfig

config = AutoConfig.from_pretrained("your-model")
print("query heads:", config.num_attention_heads)
print("kv heads:", config.num_key_value_heads)
print("group size:",
      config.num_attention_heads // config.num_key_value_heads)

如果 num_key_value_heads 缺省，有些架构会按 num_attention_heads 处理，即退化为 MHA；具体行为应以该模型的配置类与实现为准。

小结

GQA 的核心不是“少几个 head”，而是把 Query 的表达自由度与 Key/Value 的存储成本拆开。多个 Query heads 仍然学习不同的检索方式，同组 Query 只共享对历史内容的 K/V 表示，由此显著减少 KV Cache 和 decode 读取带宽。

它也不是免费的推理开关。分组会改变投影权重和模型容量，从 MHA 转换需要合理初始化与继续训练；部署时还要考虑 kernel 支持、RoPE、缓存布局和 Tensor Parallel。只有把模型质量与真实 serving 成本放在一起测量，才能找到合适的 $H_{kv}$。

参考资料

• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
• Fast Transformer Decoding: One Write-Head is All You Need
• Llama 2: Open Foundation and Fine-Tuned Chat Models
• Hugging Face Transformers 中的 Llama Attention 实现