GQA 深入理解：从 MHA、MQA 到 KV Cache 与内存带宽瓶颈

核心结论：GQA（Grouped Query Attention，分组查询注意力）本质上是让多个 Query Head 共享一组 Key/Value Head。它在尽量保留多头注意力表达能力的同时，显著减少 KV Cache 的容量和读取流量，因此尤其适合 算力很强，但显存容量或显存带宽相对紧张 的推理设备。

1. GQA 解决的到底是什么问题？

标准 Transformer Attention 为：

\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V

在自回归 LLM 推理中，模型每生成一个新 token，都需要让当前 Query 与历史 token 的 Key、Value 做 Attention。

如果每次都重新计算历史 token 的 K/V，计算会被大量重复。因此推理框架会把历史 K/V 保存下来，这就是 KV Cache。

历史 token
   ↓
K / V 投影
   ↓
KV Cache ─────────────┐
                     │
当前 token → Query ──┼→ Attention → 新 token
                     │
下一步继续读取 KV Cache

这带来一个非常关键的事实：

在长上下文和高并发场景下，Attention 的瓶颈往往不只是计算，而是“需要保存多少 KV，以及每一步需要从显存读取多少 KV”。

GQA 就是直接针对这个问题进行模型结构级优化。

2. 从 MHA 到 MQA，再到 GQA

2.1 MHA：每个 Query Head 都有独立的 K/V Head

传统 Multi-Head Attention（MHA）中：

Q1 ↔ K1, V1
Q2 ↔ K2, V2
Q3 ↔ K3, V3
Q4 ↔ K4, V4
...

如果 Query Head 数量为 $H_q$ ，那么通常：

H_q = H_k = H_v

例如 32 个 Query Head，对应 32 个 KV Head。

优点是表达能力强，各个 Attention Head 可以学习不同的表示空间。

缺点则是 KV Cache 很大。

2.2 MQA：所有 Query Head 共享同一组 K/V

Multi-Query Attention（MQA）走向另一个极端：

Q1 ─┐
Q2 ─┤
Q3 ─┼→ 共享 K1, V1
Q4 ─┤
... │

即：

H_{kv}=1

如果原来是 32 个 KV Head，那么理论上 KV Cache 可以缩小到约原来的 $1/32$ 。

但问题也很明显：所有 Query Head 共用同一套 K/V 表示，可能损失一定模型表达能力和效果。

2.3 GQA：在 MHA 与 MQA 之间折中

Grouped Query Attention（GQA）把多个 Query Head 分组，每一组共享一个 KV Head。

例如：

Group 1:
Q1 Q2 Q3 Q4  → K1, V1

Group 2:
Q5 Q6 Q7 Q8  → K2, V2

...

假设：

Query Heads = 32
KV Heads = 8

那么每 4 个 Query Head 共享一组 K/V。

可写成：

\text{Group Size}=\frac{H_q}{H_{kv}}=\frac{32}{8}=4

于是得到一个非常实用的折中：

Attention 类型	Query Heads	KV Heads	KV Cache	模型表达能力
MHA	32	32	最大	强
GQA	32	8	约 MHA 的 1/4	通常接近 MHA
MQA	32	1	约 MHA 的 1/32	共享程度最高

因此可以把三者理解为一条连续谱：

表达能力优先                                      内存效率优先

MHA ---------------- GQA ---------------- MQA
32 KV Heads           8 KV Heads            1 KV Head

3. 为什么 KV Cache 会成为推理瓶颈？

3.1 KV Cache 容量公式

忽略一些实现细节，一层 Transformer 的 KV Cache 大小可以近似写成：

\mathrm{KV\ Cache}\approx B\times S\times H_{kv}\times D_h\times 2\times \mathrm{Bytes}

其中：

$B$ ：batch size / 并发序列数
$S$ ：上下文长度
$H_{kv}$ ：KV Head 数量
$D_h$ ：每个 Head 的维度
$2$ ：分别保存 K 和 V
Bytes：FP16/BF16 通常为 2 Bytes

整个模型还要再乘 Transformer 层数 $L$ ：

\mathrm{Total\ KV\ Cache}\approx L\times B\times S\times H_{kv}\times D_h\times 2\times \mathrm{Bytes}

这里最值得注意的是：

KV Cache 大小与 $H_{kv}$ 成正比。

因此，如果从：

32 KV Heads
↓
8 KV Heads

那么其他条件不变时，KV Cache 理论容量直接降到约原来的：

\frac{8}{32}=\frac14

这正是 GQA 的核心收益。

4. “高算力、低内存设备”到底是什么意思？

这个说法容易产生歧义，因为这里实际上有两个完全不同的资源约束。

4.1 显存容量不足：Capacity Bound

假设一块加速卡 Tensor Core / Matrix Core 算力很强，但只有较小的 HBM 容量。

LLM 服务时显存大致被以下部分占用：

模型权重
+ KV Cache
+ Activation / Workspace
+ Runtime Buffer

其中模型权重在加载后相对固定，而 KV Cache 会随着以下因素持续增长：

上下文长度 ↑
并发请求数 ↑
KV Cache ↑

所以经常出现：

算力核心还有余量，但显存已经没有空间继续增加 batch 或 sequence。

GQA 通过减少 KV Head 数量，使 KV Cache 显著下降，于是设备可以支持：

更长上下文
更高并发
更大的 continuous batching
更多同时驻留的请求

因此它对 低显存容量设备 很友好。

4.2 显存带宽不足：Bandwidth Bound

另一个更重要的问题是 Memory Bandwidth。

现代 GPU/NPU 的矩阵计算单元非常快，但数据必须先从 HBM 搬进计算单元附近：

HBM
 ↓
L2 / Cache
 ↓
Shared Memory / SRAM
 ↓
Register
 ↓
Tensor Core / Matrix Core

如果计算单元每秒能处理的数据远多于 HBM 每秒能喂进去的数据，就会出现：

算力很强，但“吃不饱”。

这就是 memory-bandwidth bound。

LLM decode 阶段恰恰非常容易出现这种情况。

5. 为什么 Decode 阶段尤其受益于 GQA？

LLM 推理通常分成两个阶段：

Prefill

一次性处理整个 Prompt：

Prompt: 4096 tokens
        ↓
一次处理大量 token
        ↓
大规模矩阵乘

Prefill 通常具有较高并行度，更容易把 GPU/NPU 的计算单元利用起来，因此相对更偏 compute-bound。

Decode

之后模型每次只生成一个或少量 token：

token 4097
   ↓
读取前 4096 token 的 KV Cache
   ↓
生成 token 4098
   ↓
再次读取更大的 KV Cache

此时每个新 token 都要扫描历史 KV Cache。

因此 decode 的一个典型特征是：

计算量相对有限
但需要读取大量历史 K/V

Arithmetic Intensity 较低，于是瓶颈往往从 FLOPS 转移到：

HBM Bandwidth

GQA 将 KV Head 从 32 减到 8，不只是把“存储容量”减少到约 1/4，还意味着 Attention decode 时需要读取的 K/V 数据量也显著减少。

这就是为什么 GQA 对“高 FLOPS、但 HBM 容量或带宽相对紧张”的芯片尤其有价值。

6. 一个硬件视角下的直觉

假设存在两块推理芯片：

芯片 A
计算能力：很强
HBM：很大、带宽很高

芯片 B
计算能力：同样很强
HBM：更小、带宽更低

如果模型使用 MHA，大量 KV Cache 会让芯片 B 很快遇到：

HBM容量不足
或
HBM带宽打满

此时增加 Tensor Core 数量并不能解决问题，因为瓶颈根本不在计算单元。

而采用 GQA 后：

KV Cache ↓
HBM占用 ↓
HBM读取流量 ↓
可容纳Batch ↑
有效吞吐 ↑

因此所谓“GQA 亲和高算力低内存设备”，更准确的说法应该是：

当设备的计算能力相对于内存容量/带宽更富余时，GQA 能减少 Attention 对内存系统的压力，从而让已有算力更容易被真正利用起来。

7. GQA 并不是“免费优化”

需要注意，GQA 并不是简单把一个已经训练好的 MHA 模型在推理时强行改成更少 KV Head，就一定不会损失效果。

它本质上属于 模型架构设计。

减少 KV Head 意味着多个 Query Head 共享 K/V 表示，因此共享程度越高，理论上的表示自由度越小。

所以存在一个 trade-off：

更多 KV Heads
→ 更强表示自由度
→ 更大 KV Cache

更少 KV Heads
→ 更高内存效率
→ 更强共享约束

GQA 的价值就在于找到 MHA 与 MQA 之间较好的 Pareto 点。

8. GQA、FlashAttention 和 PagedAttention 有什么区别？

这三个概念经常一起出现，但它们优化的是不同层次。

技术	优化层次	核心目标
GQA	模型结构	减少 KV Head，从源头减少 KV Cache
FlashAttention	Attention Kernel / IO	通过 tiling 等方式减少 HBM 中间读写
PagedAttention	推理系统 / KV Cache 管理	降低 KV Cache 内存碎片，提高动态分配效率

可以理解为：

GQA
↓
让需要保存的 KV 本身变少

FlashAttention
↓
让 Attention 计算过程减少无效 HBM IO

PagedAttention
↓
让已经存在的 KV Cache 被更高效地组织和管理

三者并不冲突，而是可以叠加。

9. 从 AI Infra 视角应该建立什么认知？

学习 GQA 的意义不只是记住一个 Attention 变体，而是建立一个重要的推理性能模型：

模型架构
   ↓
决定 KV Head 数量
   ↓
决定 KV Cache 大小
   ↓
决定显存容量压力 + HBM流量
   ↓
影响可支持的并发与上下文长度
   ↓
最终影响推理吞吐和成本

这体现了典型的 Algorithm-System-Hardware Co-design：

算法结构：GQA
        ↓
系统层：KV Cache / Continuous Batching / PagedAttention
        ↓
Kernel层：FlashAttention / Fused Attention
        ↓
硬件层：HBM Bandwidth / SRAM / Tensor Core

一个成熟的 LLM Inference 工程师不能只问：

“这个模型有多少 FLOPS？”

还应该继续问：

每生成一个 token 要读取多少字节？
KV Cache 每个 token 占多少空间？
当前是 compute-bound 还是 memory-bound？
Prefill 与 Decode 的瓶颈是否相同？
GQA ratio 对 batch size 和吞吐有什么影响？
HBM 容量与 HBM bandwidth 哪一个先成为瓶颈？

10. 最后的记忆模型

可以用一个非常简单的类比记住三种 Attention：

MHA：
32个人，每个人自己带一套资料。

MQA：
32个人，所有人共享一套资料。

GQA：
32个人分成8组，每4个人共享一套资料。

这里：

人       = Query Head
资料     = Key / Value
资料仓库 = KV Cache

因此：

GQA 的本质不是减少 Query 的数量，而是减少独立 K/V 的数量。

这也是理解它为什么能显著降低 KV Cache、减少 HBM 流量，并改善 LLM decode 效率的关键。

11. 一句话总结

GQA 是一种模型结构级的 Attention 优化：通过让多组 Query Head 共享更少的 KV Head，在模型表达能力与 KV Cache 成本之间取得折中；由于 LLM decode 常受显存容量和内存带宽限制，因此它尤其适合计算能力强、但内存系统相对紧张的 GPU/NPU 推理平台。