大模型为什么越聊越慢？为什么同样是 Transformer，有的模型能撑长上下文，有的很容易被 KV Cache 卡住？MHA、MQA、GQA、FlashAttention 这几个词经常一起出现，但它们解决的并不是同一个问题。

简单讲：MHA / MQA / GQA 关注的是 “K/V 头到底存几份”；FlashAttention 关注的是 “Attention 计算时怎么少搬数据”。前者偏模型结构，后者偏 GPU 计算内核。

提示：不要把 FlashAttention 理解成 MQA 或 GQA 的替代品。它们经常是组合使用：例如模型结构用 GQA，底层计算用 FlashAttention。

一、MHA 为什么会慢？

MHA 的慢，不是一个单点问题，而是两类成本叠加。第一类是注意力矩阵的计算成本：序列长度越长，token 两两之间的相关性矩阵越大。第二类是生成阶段的 KV Cache：模型每生成一个新 token，都要读取历史 token 的 K/V。上下文越长、层数越多、KV 头越多，显存和带宽压力越大。

训练时，Attention 的 N×N 矩阵和反向传播中间结果会造成很高显存压力；推理时，尤其是自回归解码阶段，KV Cache 读取会成为吞吐瓶颈。Hugging Face 文档也强调，KV cache 通过复用已计算的 key/value 来减少重复计算，从而优化生成速度；NVIDIA TensorRT-LLM 文档则说明每个 Transformer layer 都有对应的 KV cache。

二、MHA：效果强，但每个头都要存 K/V

MHA，全称 Multi-Head Attention。它把注意力拆成多个头，每个头都从不同角度看上下文。有的头可能更关注主语和谓语，有的头关注长距离依赖，有的头关注局部短语。最后再把所有头的结果拼起来。

MHA 的好处是表达能力强，坏处也很直接：每个头都要有自己的 K 和 V。推理时，历史 token 的 K/V 需要缓存下来，头数越多，KV Cache 越大。

提示：可以把 MHA 理解成“每个专家都有一套独立资料库”。专家越多，看问题越丰富，但资料库也越占空间。

三、KV Cache：推理阶段真正容易卡住的地方

大模型生成文本是一个 token 一个 token 往外吐。如果每次生成新 token 都重新计算前面所有 token 的 K/V，那会非常浪费。所以推理框架会把历史 K/V 缓存下来，下一步直接复用。

KV Cache 的估算公式可以写成：

这里最关键的是 `num_kv_heads`。MHA 的 KV 头数通常等于注意力头数；MQA 只有 1 组 KV；GQA 则是介于二者之间。也就是说，减少 KV 头数，可以直接压低 KV Cache 的显存和带宽需求。

四、MQA：所有 Query 头共享一组 K/V

MQA，全称 Multi-Query Attention。它的思路很激进：Query 头可以有很多个，但所有 Query 头共享同一组 Key 和 Value。这样一来，KV Cache 的头数从很多组降到 1 组。

好处是显存和带宽明显下降，解码阶段更快。但代价是不同 Query 头看到的 K/V 信息变得更相似，模型表达细节可能下降。Noam Shazeer 的 MQA 论文提出的核心动机就是：增量解码时反复加载巨大的 key/value 张量会拖慢推理，而共享 K/V 可以显著降低内存带宽压力。

五、GQA：在 MHA 和 MQA 之间折中

GQA，全称 Grouped-Query Attention。它不是让所有 Query 头共享一组 K/V，而是把 Query 头分成几组，每一组共享一组 K/V。这样既能减少 KV Cache，又不会像 MQA 那样共享得太狠。

GQA 可以理解成一个连续谱：如果 KV 组数等于 Query 头数，它退化成 MHA；如果 KV 组数等于 1，它退化成 MQA；中间状态就是常见的 GQA。GQA 论文也强调，它是 MQA 的泛化形式，目标是在推理速度和模型质量之间取得更好的平衡。

六、FlashAttention：少搬数据，比少算更关键

FlashAttention 和 MQA、GQA 不是同一类东西。MHA / MQA / GQA 改的是模型结构里的 K/V 头分配；FlashAttention 改的是 attention 的计算方式。

普通 attention 往往会产生很大的中间注意力矩阵。这个矩阵在长上下文下非常大，频繁在 GPU HBM 和计算单元之间读写。FlashAttention 的思路是把 Q/K/V 分块搬到更快的片上 SRAM，块内完成计算和 softmax 的在线更新，不显式落完整 N×N 矩阵，从而减少 HBM 读写。

FlashAttention 论文把这种思路称为 IO-aware：不是只盯着数学计算量，而是把 GPU 内存层级之间的数据搬运也算进去。官方实现也说明，FlashAttention 支持 MQA/GQA 这种 Q 头数多、KV 头数少的输入形状。

七、它们到底是什么关系？

这几个概念最容易混的地方，就是把“结构优化”和“内核优化”混在一起。结构优化决定 KV Cache 存多少份；内核优化决定 attention 这件事怎么算得更省。

常见组合是：模型用 GQA 减少 KV Cache，再用 FlashAttention 或类似 fused kernel 降低 attention 计算中的显存读写。这样既能控制缓存，又能提升长上下文时的计算效率。

八、用 PyTorch 看一下代码长什么样

下面是一个简化版示例。重点不是拿它直接训练模型，而是看清楚 Q 头数和 K/V 头数的关系。PyTorch 的 `scaled_dot_product_attention` 支持 `enable_gqa=True`，但具体能否走高效内核，还要看设备、dtype、head shape 等条件。

如果是 MHA，通常 `num_q_heads == num_kv_heads`；如果是 MQA，`num_kv_heads == 1`；如果是 GQA，`1 < num_kv_heads < num_q_heads`。

九、生产环境怎么选？

图 8：选型判断树

如果你在训练一个新模型，MHA 的表达能力更稳，但成本高。如果你更关心推理成本，尤其是长上下文、多并发、显存紧张，GQA 往往是更平衡的选择。MQA 更激进，适合特别强调解码速度和缓存压缩的场景，但要认真评估质量。

如果你已经确定了模型结构，仍然可以继续优化 attention kernel。这时 FlashAttention、FlashAttention-2、框架内置的 SDPA、推理引擎里的 fused attention，都属于需要关注的方向。

十、最后总结：别再把它们混为一谈

• MHA：每个注意力头都有自己的 K/V，表达能力强，但 KV Cache 大。

• MQA：所有 Query 头共享一组 K/V，显存和带宽最省，但可能损失细节。

• GQA：Query 头分组共享 K/V，是 MHA 和 MQA 的折中方案。

• FlashAttention：不改变头的结构，而是用分块和在线 softmax 减少 HBM 读写。

• 生产中常见组合是 GQA + FlashAttention：一个省缓存，一个省计算 IO。