从MHA到GQA - Zhonghui

2017 年的 Transformer 用 Multi-Head Attention（MHA，多头注意力）取代循环结构时，最重要的问题是：怎样让模型从不同表示子空间、不同位置同时读取信息。到了大语言模型时代，注意力面对的另一个问题越来越突出：模型生成每个 token 时，都要重新读取此前所有 token 的 Key 和 Value。上下文越长、并发请求越多，这部分 KV Cache 越容易占满显存、耗尽内存带宽。

MQA 和 GQA 没有改变注意力的基本含义。它们保留多个 Query head，只是逐步减少与之配套的 Key/Value head。MLA 的目标同样是缩小 KV Cache，但做法不同：它不再缓存若干组完整的 K/V，而是缓存一份低维 latent representation（潜在表示）。

从 MHA 到 MQA、GQA 与 MLA 的论文脉络

这四种结构并不是一条严格的线性升级路线：

MHA、GQA、MQA 可以放在同一条 KV head 共享程度的连续谱上；
MLA 改变了 K/V 的参数化和缓存形式，属于 低秩压缩路线；
它们都主要面向自回归生成中的缓存与带宽问题，不等于解决了长序列 dense attention 的全部计算问题。

问题为什么从「注意力怎么算」变成「KV 怎么存」#

缩放点积注意力可以写成：

\operatorname{Attention}(Q,K,V) =\operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_h}}\right)V

Query 决定当前需要寻找什么，Key 用于和 Query 计算匹配程度，Value 则是匹配后真正汇聚的内容。对 decoder-only 语言模型来说，生成过程可以分成两个计算形态。

Prefill 处理已经完整给出的 prompt。所有位置都已知，只要加上 causal mask，各位置的 Q、K、V 可以并行计算。Decode 每次只生成一个新 token：当前 token 产生新的 Query，读取此前所有 token 的 K/V，得到输出后再把当前 token 的新 K/V 追加进缓存。

Prefill、Decode 与 KV Cache 的关系

过去 token 的 Query 用完就不会再被未来位置查询，因此没有必要保留；过去 token 的 K/V 却会在之后每个解码步骤中反复使用。KV Cache 正是用空间换计算：若不缓存它们，每生成一个 token 都要重新计算整个前缀的 K/V。

设 batch size 为 $B$ ，层数为 $N$ ，已经缓存的序列长度为 $L$ ，KV head 数为 $H_{kv}$ ，每个 head 的维度为 $d_h$ ，每个数占 $b$ bytes。在不考虑内存对齐、分页和量化元数据时，常规 MHA/GQA/MQA 的 KV Cache 可以近似写成：

\operatorname{KVCache} \approx 2BNLH_{kv}d_hb

式子开头的 2 来自 Key 和 Value。上下文长度与 batch size 都会线性放大缓存，而解码的每一步还要从显存读取相当一部分历史 K/V。模型可能因此不是算力不够，而是等待数据搬运。

MHA：每个 Query head 都有自己的 K/V#

Vaswani 等人在 2017 年的 Attention Is All You Need 中提出了 Transformer 使用的 Multi-Head Attention。第 $i$ 个 head 有独立的投影矩阵：

\operatorname{head}_i =\operatorname{Attention} \left(QW_i^Q,KW_i^K,VW_i^V\right)

\operatorname{MHA}(Q,K,V) =\operatorname{Concat}(\operatorname{head}_1,\ldots,\operatorname{head}_H)W^O

原论文给出的动机不是 KV Cache，而是单个 attention head 的加权平均会限制表示能力。多个 head 使用不同的投影，可以同时从不同表示子空间、不同位置组合信息。Transformer Base 使用 8 个 head；在保持总计算量近似不变的消融中，单 head 比论文采用的 8-head 设置低 0.9 BLEU，不过 head 并不是越多越好，32 个更窄的 head 同样出现了下降。

在标准 MHA 中，Query、Key、Value 的 head 数相同：

H_q=H_k=H_v=H

这种一对一结构给每个 Query head 一套独立的 K/V 表示，但也意味着每个历史 token、每一层都要保存 $H$ 组 Key 和 $H$ 组 Value。原始 Transformer 主要研究机器翻译训练，当时的序列与模型都远小于今天的长上下文 LLM；MHA 的这项推理成本后来才成为主要矛盾。

MQA：保留多个 Query，只写一组 K/V#

Noam Shazeer 在 2019 年的 Fast Transformer Decoding: One Write-Head is All You Need 中提出 Multi-Query Attention（MQA，多查询注意力）。这里的 “one write-head” 指只有一组 K/V 被写入缓存，并不是整个注意力只剩一个 head。

MQA 仍然保留 $H_q$ 个不同的 Query 投影，但所有 Query head 共用同一组 Key 和 Value：

H_{kv}=1

不同 Query head 仍会产生不同的注意力分数，并形成多个输出 head；被删除的只是 K/V 张量上的 head 维度。若 MHA 原本有 $H$ 个 KV head，在 head dimension 不变的简化条件下，MQA 的 KV Cache 就约为它的 $1/H$ 。

MHA、GQA 与 MQA 的 Query/KV head 连接方式

MQA 论文首先从硬件瓶颈出发分析这个变化。当序列中的位置无法并行生成时，计算设备需要在每一步重新读取很大的 K/V 张量；现代 GPU/TPU 的算术吞吐远高于内存带宽，计算量没有饱和也可能被数据搬运拖慢。去掉 K/V 的 head 维度，正好减少这部分读取量。

论文的主要实验不是今天常见的百亿参数 decoder-only LLM，而是 WMT 2014 英德翻译和 Billion Word 语言模型。翻译模型约 2.11 亿参数、8 个 head，在 TPUv2 上以 1024 条序列组成的 batch 做增量 greedy decoding 时，Decoder 的摊销时间从 MHA 的每 token 46 微秒降到 MQA 的 3.8 微秒；beam search 则从 203 微秒降到 32 微秒。与此同时，dev BLEU 从 26.7 小幅降至 26.5，beam-4 的 test BLEU 分别为 28.4 和 28.5。

这些数字证明了内存带宽优化确实可以转化成速度，但不能直接套到任意模型和 GPU 上。它们依赖当时的模型、TPU、batch、序列长度和实现。论文也把 FFN 加宽以对齐总参数量，并把 Encoder self-attention、Decoder self-attention 和 cross-attention 全部换成 MQA，因此不是只改现代 LLM 的 causal self-attention 所得到的孤立对照。

MQA 还带来两个取舍。第一，所有 Query head 只能面对同一套 K/V 表示，压缩比较激进，模型容量可能下降。第二，只有一个 KV head 时，按 head 做 tensor parallelism 不够自然；单一 K/V 有时会在多个设备上复制，理论缓存优势不一定全部变成端到端收益。

GQA：把 Query head 分组#

MQA 在 MHA 的两个极端之间留下了很大的设计空间。2023 年的 GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints 把 Query head 分为 $G$ 组，每组共享一个 Key head 和一个 Value head：

H_{kv}=G,\qquad \text{每组 Query head 数}=\frac{H_q}{G}

于是三种结构可以用同一个参数 $G$ 表示：

\begin{aligned} G=H_q &\Rightarrow \text{MHA}\\ 1<G<H_q &\Rightarrow \text{GQA}\\ G=1 &\Rightarrow \text{MQA} \end{aligned}

例如 32 个 Query head 搭配 8 个 KV head 时，每 4 个 Query head 共享一组 K/V，缓存约为对应 MHA 的四分之一。相比 MQA，它保留了更多 K/V 表示子空间；相比 MHA，它仍然显著减少了缓存和解码时的读取量。

原论文不只提出了一种新结构#

GQA 论文还要解决一个现实问题：已经花费大量算力训练好的 MHA checkpoint，能否不用从头训练就转换成 MQA 或 GQA？论文把这个过程称为 uptraining：

把同一组内原有 Key 投影矩阵取平均，得到新的 Key head；Value 投影也做同样处理；
用原来的预训练配置和数据继续训练，让模型适应新的共享结构。

从 MHA checkpoint 转换并继续训练为 GQA

平均池化比“只保留第一个 head”或“随机初始化新 head”更好，因为它在结构压缩时保留了更多原 checkpoint 的信息。消融还显示，GQA 在刚转换后就有一定可用性，MQA 则更依赖后续训练；两者从相当于原预训练计算量 5% 的 uptraining 中都获得了明显收益，继续增加到 10% 时收益趋缓。

主实验基于 T5.1.1，而不是从头训练的 decoder-only 模型。研究者把 MQA/GQA 用在 Decoder self-attention 和 cross-attention 中，没有修改可以并行计算的 Encoder self-attention。在 8 个 TPUv4、每个芯片尽可能使用最大 batch，并为各模型分别优化并行策略的设置下，T5-XXL 的平均单样本推理时间为：

模型	单样本推理时间	说明
MHA-XXL	1.51 s	原始 T5-XXL
MQA-XXL	0.24 s	经过 5% uptraining
GQA-8-XXL	0.28 s	8 组 KV head，经过 5% uptraining

GQA-8 在论文的摘要、翻译和问答任务上取得接近 MHA-XXL 的整体质量，同时速度接近 MQA。这里的“接近”来自这组 T5 实验，不是关于所有模型的定理；分组数量仍然需要结合模型容量、训练方式和部署硬件选择。

GQA 后来成为 decoder-only LLM 很常见的工程方案。Llama 2 在 70B 模型中使用 GQA；Llama 3 又把 GQA 扩展到 8B 和 70B 两种尺寸。它受欢迎的原因并不是结构最简单，而是落在一个实用的中间点：压缩显著，又没有把所有 Query head 都压到同一份 K/V 上。

三种共享方式放进同一个缓存公式#

在常规注意力中，只要确定 $H_{kv}$ ，就能比较 MHA、GQA 和 MQA。假设一个模型有 32 层、32 个 Query head，head dimension 为 128，使用 BF16，batch size 为 1，上下文长度为 4096：

结构	$H_{kv}$	KV Cache
MHA	32	2 GiB
GQA-8	8	512 MiB
MQA	1	64 MiB

同一配置下 MHA、GQA 与 MQA 的 KV Cache 大小

如果只把上下文从 4K 拉长到 128K，三者分别变成约 64 GiB、16 GiB 和 2 GiB；batch size 再扩大几倍，缓存也跟着等比例增加。这个例子只计算 K/V tensor 本身，不包含模型参数、activation、临时 workspace、内存碎片和框架管理开销，但已经足以说明为什么 KV head 数会直接影响可服务的上下文和并发量。

结构	Query head	KV head	共享程度	主要优势	主要代价
MHA	$H$	$H$	不共享	每个 Query head 有独立 K/V 表示	KV Cache 最大
GQA	$H$	$G$	组内共享	容量与推理效率之间可调	需要选择分组数，仍有一定信息共享
MQA	$H$	1	全部共享	常规结构中 KV Cache 最小	压缩最激进，质量和并行策略可能受影响

缓存越小不代表模型必然越好，也不保证端到端速度按同样比例增加。Prefill 阶段仍要计算多个 Query head 与序列的注意力；短上下文、小 batch 时，模型权重读取或其他模块可能才是瓶颈；部署框架是否有高效 kernel、KV head 能否自然分片，也会改变实际结果。

MLA：不再缓存完整的多头 K/V#

Multi-head Latent Attention（MLA，多头潜在注意力）由 2024 年的 DeepSeek-V2 技术报告提出。它没有一篇与 MQA、GQA 完全对应的独立原始论文，而是 DeepSeek-V2 架构中的 attention 模块；同一模型的 FFN 还使用了 DeepSeekMoE，二者不应混为一项技术。

MLA 的核心是 low-rank key-value joint compression（低秩 Key/Value 联合压缩）。对第 $t$ 个 token 的 hidden state $h_t$ ，先得到低维表示：

c_t^{KV}=W^{DKV}h_t

再分别上投影为各个 head 使用的 content Key 和 Value：

k_{t,i}^{C}=W_i^{UK}c_t^{KV},\qquad v_{t,i}^{C}=W_i^{UV}c_t^{KV}

K 和 V 共享压缩来源 $c_t^{KV}$ ，但有各自的上投影，仍然承担不同角色。训练时这些式子可以按通常方式计算；自回归推理时，关键是只保存压缩后的 $c_t^{KV}$ ，而不是为每个 head 保存展开后的 $k_{t,i}^{C}$ 和 $v_{t,i}^{C}$ 。

MLA 的低秩 KV 联合压缩、矩阵吸收与 RoPE 分支

为什么缓存 latent 后不需要反复把 K/V 展开#

只缓存 $c_j^{KV}$ 看起来会产生一个新问题：每次生成时，难道要为所有历史 token 重新计算完整 K/V？MLA 利用矩阵乘法结合律做 matrix absorption（矩阵吸收）。忽略位置编码，某个 head 的注意力分数包含：

q_{t,i}^{\top}k_{j,i}^{C} =q_{t,i}^{\top}W_i^{UK}c_j^{KV} =\left((W_i^{UK})^{\top}q_{t,i}\right)^{\top}c_j^{KV}

$W_i^{UK}$ 可以吸收到 Query 一侧。Value 的上投影同样可以与最终输出投影合并，使加权聚合直接围绕压缩 latent 进行。这里的“吸收”不是删掉训练参数，而是改变推理时矩阵乘法的执行顺序，避免先把所有历史 latent 还原成完整多头 K/V。

RoPE 为什么必须单独处理#

直接在 content Key 上应用 RoPE 会把一个随位置变化的旋转矩阵夹在低秩上投影与 Query 之间。这个矩阵对每个历史位置不同，不能再简单吸收到固定权重中；如果强行这样做，就要在每一步重新生成所有历史 Key，失去压缩缓存的意义。

DeepSeek-V2 因此使用 decoupled RoPE（解耦旋转位置编码）：

content 分量 $q^C,k^C$ 负责主要语义匹配，沿用可被矩阵吸收的低秩路径；
额外的 $q^R,k^R$ 专门承载 RoPE 位置信息；
每个 Query head 有自己的 $q^R$ ，Key 侧的 $k^R$ 在 head 之间共享。

最终 Query 和 Key 分别由 content 与 RoPE 分量拼接。推理时除了 $c_t^{KV}$ ，还要缓存共享的 $k_t^R$ 。因此 MLA 每层、每 token 的缓存元素数不是只有压缩维度 $d_c$ ，而是：

\operatorname{Cache}_{MLA}=d_c+d_h^R

相应地，四种结构每层、每 token 的简化缓存量为：

结构	缓存元素数
MHA	$2H_qd_h$
GQA	$2Gd_h$
MQA	$2d_h$
MLA	$d_c+d_h^R$

DeepSeek-V2 使用 128 个 attention head、每 head 128 维， $d_c=512$ ， $d_h^R=64$ 。代入后，MHA、GQA-8、MQA 和 MLA 每层每 token 分别需要 32768、2048、256 和 576 个元素。MLA 在这个配置下甚至比 MQA 多，但远小于 MHA 和常见组数的 GQA；论文把 576 个元素描述为等价于 2.25 组 GQA。它追求的不是绝对最小缓存，而是用低秩 latent 保留比激进 KV 共享更丰富的表示。

DeepSeek-V2 的实验应该怎样理解#

DeepSeek-V2 报告中的受控消融比“整体吞吐提升多少”更能说明 MLA 本身的效果。作者分别训练了约 16B 和 250B 总参数的 MoE 模型，每一对主要区别是 MHA 或 MLA。在四项困难 benchmark 上，MLA 版本整体更强；KV Cache 则分别只有对应 MHA 版本的约 14% 和 4%。不过两组模型的总参数和激活参数并非完全相同，单项结果也并非全部由 MLA 获胜，因此这是一组支持 MLA 取舍有效的实证结果，不是“MLA 必然优于 MHA”的普遍证明。

报告摘要还给出“相对 DeepSeek 67B，KV Cache 减少 93.3%、最大生成吞吐提升到 5.76 倍”。这是 DeepSeek-V2 与上一代完整模型的系统级比较，两者还存在 MoE、模型规模、训练和部署等差异，不能把整个数字都当作 MLA 单模块的因果效果。

四种结构是两条路线#

MHA、GQA、MQA 改变的是 KV head 的数量；MLA 改变的是 K/V 的表示空间。把它们放进同一张表，关系会更清楚：

结构	K/V 的组织方式	实际缓存	与 MHA 的关系
MHA	每个 Query head 独立 K/V	所有 head 的 K/V	基准结构
GQA	一组 Query head 共享 K/V	$G$ 组 K/V	减少 KV head
MQA	所有 Query head 共享 K/V	1 组 K/V	把共享推到极端
MLA	K/V 来自共享低维 latent	KV latent + RoPE key	改变 K/V 参数化

注意力效率技术分别改变了什么

这张分类也能避免把几种常见技术混在一起：

Sliding Window / Local Attention 改变每个 Query 能读取哪些历史位置，主要减少被访问的 token 数；
Sparse Attention 只选择一部分位置或块，主要降低长序列 attention 的计算量；
FlashAttention 不改变注意力的数学结果，而是重新组织分块计算和显存读写；
KV Cache quantization 减少每个缓存元素占用的 bytes；
PagedAttention 改善 KV Cache 的分配、分页和复用方式；
GQA/MQA 减少完整 KV head 数，MLA 则压缩每个 token 被缓存的表示。

这些方法位于不同设计轴上，可以组合出现。GQA 可以配合 FlashAttention，也可以同时使用滑动窗口和量化 KV Cache。MQA 原论文也专门训练了 local attention 版本，用来说明“减少 K/V head”和“减少被关注的位置”是正交的。

MHA、GQA、MQA、MLA 同样没有自动把 dense attention 的二次复杂度变成线性。对长度为 $L$ 的完整 causal attention，prefill 仍然要处理数量级为 $O(L^2)$ 的 query-key 关系；decode 时每个新 token 仍要读取随历史长度线性增长的缓存。它们首先解决的是缓存容量和带宽，而不是所有长上下文问题。

应该记住的四件事#

第一，MHA、GQA、MQA 的 Query head 都可以是多个。三者的关键差异是 KV head 数，不要把 MQA 误解成 single-head attention。

第二，GQA 是 MHA 与 MQA 的一般化表达。 $H_{kv}=H_q$ 时是 MHA， $H_{kv}=1$ 时是 MQA，中间的组数构成可调的质量—缓存折中。

第三，MLA 不是 GQA 继续减少 group 后的别名。它通过低秩联合压缩、矩阵吸收和解耦 RoPE，让推理围绕缓存的 latent representation 进行。

第四，KV Cache 小不等于所有阶段都等比例加速。真正的端到端表现还取决于上下文长度、batch、模型宽度、精度、并行策略、kernel 和硬件。论文中的加速数字必须连同实验条件一起阅读。

从 2017 年关注“不同 head 能学到什么”，到 2019 年开始追问“每步生成要从内存读多少”，再到 GQA 的分组共享和 MLA 的低秩压缩，attention 的演变反映出一个很实际的变化：现代大模型的瓶颈不仅是参数和 FLOPs，保存、搬运历史信息的成本同样会决定模型能否高效运行。

参考资料#

Vaswani et al. Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762
Shazeer. Fast Transformer Decoding: One Write-Head is All You Need. 2019. arXiv:1911.02150
Ainslie et al. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023. arXiv:2305.13245
DeepSeek-AI. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. 2024. arXiv:2405.04434
Touvron et al. Llama 2: Open Foundation and Fine-Tuned Chat Models. 2023. arXiv:2307.09288
Meta AI. Introducing Meta Llama 3. 2024. 官方介绍