Grouped Query Attention (GQA)

分组查询注意力（GQA）是一种用于 Transformer 的多头注意力形式，其中多个查询头组成一组，共享一个键头和一个值头。例如，Llama 3.1 8B 的每一层都将 32 个查询头与 8 个键值头配对。它介于多头注意力（MHA）和多查询注意力（MQA）之间：在 MHA 中，每个查询头都有自己的键和值；在 MQA 中，所有查询头共享单个键头和值头。Google Research 的 Joshua Ainslie 及其同事在 EMNLP 2023 上发表的论文 GQA：从多头检查点训练通用多查询 Transformer 模型中介绍了 GQA，并展示了 GQA 模型能够以接近 MHA 的质量实现接近 MQA 的解码速度。Llama 2 70B、Llama 3 系列、Mistral 7B 和 Gemma 2 等开源大语言模型（LLM）都使用了 GQA；基于这些模型构建的视觉语言模型（VLM）也会在语言解码器中继承 GQA，例如基于 Gemma 2 的 Google PaliGemma 2。

GQA 的工作原理#

在自注意力中，每个词元都会针对每个头投影为查询、键和值；每个头会将查询与键进行比较。GQA 保留完整的查询头集合，但会减少键和值头的数量，并让每组查询头使用该组共享的键和值进行注意力计算。若有 8 个查询头和 2 个键值头，则每个键值头为 4 个查询头服务：

每个头内部的注意力计算仍是原样的缩放点积注意力。只有一组时，该层就成为 MQA；组数与查询头数相同时，就成为 MHA。PyTorch通过 scaled_dot_product_attention 的 enable_gqa 参数直接支持这种布局：

import torch
import torch.nn.functional as F

torch.manual_seed(0)
q = torch.randn(1, 8, 16, 64)  # 8 query heads, 16 tokens, 64 dimensions per head
k = torch.randn(1, 2, 16, 64)  # only 2 key heads
v = torch.randn(1, 2, 16, 64)  # and 2 value heads

out = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)

# Equivalent multi-head attention: copy each key-value head to its group of 4 query heads
k_full, v_full = k.repeat_interleave(4, dim=1), v.repeat_interleave(4, dim=1)
reference = F.scaled_dot_product_attention(q, k_full, v_full, is_causal=True)

print(f"Output {tuple(out.shape)}, same as copied heads: {torch.allclose(out, reference, atol=1e-6)}")
print(f"Cached values per token per layer: {k[0, :, 0].numel() * 2} with GQA, {k_full[0, :, 0].numel() * 2} with MHA")

它会打印 Output (1, 8, 16, 64), same as copied heads: True 和 256 with GQA, 1024 with MHA：全部 8 个查询头都能产生输出，而所存储的键和值只有 MHA 所需数量的四分之一。

GQA 如何缩小 KV 缓存#

在文本生成过程中，模型会在KV 缓存中保留每个先前词元的键和值，并在生成每个新词元时读取整个缓存，因此解码速度更多受内存带宽而非运算能力限制。缓存大小与键值头数量成正比。Llama 3.1 8B 有 32 层，每层包含 8 个维度为 128 的键值头，因此在16 位精度下，每个词元占用 128 KiB；在完整的 131,072 词元上下文窗口下，总量为 16 GiB。若使用 MHA，同一模型每个词元需要 512 KiB，完整长度下需要 64 GiB。节省下来的内存可用于支持更长上下文、增大连续批处理的批次，以及降低每个词元的推理延迟。Ultralytics LLM 接口接受适用于任何兼容 OpenAI 的服务器的 base_url，例如自托管的 Llama 3.1 8B；由于有 8 个键值头，在给定的缓存内存中，服务器可容纳的对话词元数是 32 头 MHA 布局的四倍。

将模型转换为 GQA#

GQA 无需从头训练。论文作者通过对每组中的键和值头求平均值，将现有 MHA 检查点转换为 GQA，然后使用约 5% 的原始计算量继续预训练；他们将这一步称为后续训练。MQA 是共享单个头的极端形式，由 Noam Shazeer 于 2019 年提出，也是速度最快的布局，但它会降低质量并可能导致训练不稳定；采用中间数量的组则能恢复大部分质量。

GQA 与相关技术的比较#

• 多头潜在注意力（MLA）：随 DeepSeek-V2 一同推出，会将每个词元的键和值压缩为一个较小的共享潜在向量。
• PagedAttention：以固定大小的块分配 KV 缓存，从而避免内存碎片化，不受填充缓存的注意力头数量影响。
• 滑动窗口注意力：限制每层缓存的词元数量；Mistral 7B 将其与 GQA 结合使用。
• Flash Attention：一种能够更快计算精确注意力、并原生支持 GQA 布局的 GPU 内核。