第1章：Multi-Query Attention (MQA) 与 Grouped-Query Attention (GQA)
各位同学，今天我们来聊聊推理加速里一个非常实在的话题——怎么把KV缓存这块硬骨头啃下来。
我在做模型部署的时候，最头疼的就是显存。你想想看，一个7B的模型，光KV缓存就能吃掉好几个G的显存。尤其是做长序列推理时，显存直接爆掉。后来我接触到MQA和GQA，才发现原来业界早就有了解决方案。
1.1 标准Attention的痛点
先回顾一下标准的多头注意力（MHA）。每个注意力头都有自己的K、V矩阵。假设有32个头，每个头的维度是128，那KV缓存就是32 × 128 × 2 × 序列长度。这个数字随着序列长度线性增长，非常恐怖。
我遇到过最夸张的情况：一个8K长度的序列，KV缓存占了将近4GB显存。而模型参数本身才占6GB。你说这合理吗？
核心问题：标准MHA中，每个头都独立存储K和V，导致KV缓存随头数线性增长。推理时，这成了显存和带宽的双重瓶颈。
1.2 MQA：让所有头共享KV
MQA的思路很简单：既然每个头都存一份K和V太浪费，那让所有Query头共享同一组K和V不就行了？
说白了，就是把KV头的数量降到1。所有Query头都去访问这同一组K和V。
# 标准MHA：每个头独立KV
K = [K_head1, K_head2, ..., K_head32] # 32个KV头
V = [V_head1, V_head2, ..., V_head32]
# MQA：所有头共享KV
K = K_shared # 1个KV头
V = V_shared这样做的好处很明显：KV缓存直接降到原来的1/32。显存占用大幅降低，带宽需求也小了。
但代价呢？我刚开始用MQA时也担心质量会掉。毕竟所有头都用同一组KV，表达能力会不会受限？
实验结果证明，MQA在大多数任务上质量损失很小，甚至有些任务上还能持平。原因在于，Query头之间的差异已经足够表达不同的注意力模式了。
我的经验：在部署对话模型时，MQA能让我把batch size从4提到16，吞吐量直接翻倍。质量上几乎看不出差别。
1.3 GQA：在MQA和MHA之间找平衡
MQA虽然省显存，但有些人觉得它太激进了。所有头共享一组KV，万一某些任务需要更细粒度的注意力呢？
于是GQA出现了。它把KV头分成几个组，每个组内的Query头共享KV。比如8个KV头，每个组有4个Query头共享。
# GQA：分组共享KV
# 假设32个Query头，8个KV头
K = [K_group1, K_group2, ..., K_group8] # 8个KV头
V = [V_group1, V_group2, ..., V_group8]
# 每个KV头被4个Query头共享GQA相当于在MHA和MQA之间做了个折中。KV头数可以自由调节，比如8、16、32等。你想想看，这给了我们很大的灵活性。
| 方法 | KV头数 | 显存节省 | 质量损失 |
|---|---|---|---|
| MHA | 32 | 1x | 无 |
| GQA-8 | 8 | 4x | 极小 |
| GQA-4 | 4 | 8x | 较小 |
| MQA | 1 | 32x | 可接受 |
我建议在实际项目中，先试试GQA-8。如果显存还是紧张，再降到GQA-4或MQA。这样能最大程度保证质量。
1.4 核心逻辑图
下面这张图展示了MHA、GQA、MQA三者的核心区别。你可以直观地看到KV头数的变化。
1.5 实际部署中的选择
那么问题来了：到底该用MQA还是GQA？
我个人习惯这样判断：
- 显存极度紧张（比如部署在边缘设备）：直接用MQA，省到极致
- 需要平衡质量和速度（比如云端服务）：用GQA-8或GQA-4
- 对质量要求极高（比如医疗、金融场景）：还是用MHA吧，但要做好显存规划
避坑指南：我曾经在部署一个对话模型时，直接用了MQA，结果发现模型在长文本理解任务上质量下降明显。后来换成GQA-8，问题就解决了。所以一定要根据实际任务来测试，不要盲目追求极致的显存节省。
1.6 实现要点
实现MQA和GQA其实不复杂。关键是在模型定义时，把KV头的数量设成1（MQA）或某个较小的数（GQA）。
以PyTorch为例：
class GQAAttention(nn.Module):
def __init__(self, d_model, n_heads, n_kv_heads):
super().__init__()
self.n_heads = n_heads
self.n_kv_heads = n_kv_heads
self.head_dim = d_model // n_heads
# Q投影：头数不变
self.wq = nn.Linear(d_model, n_heads * self.head_dim)
# K和V投影：头数减少
self.wk = nn.Linear(d_model, n_kv_heads * self.head_dim)
self.wv = nn.Linear(d_model, n_kv_heads * self.head_dim)
def forward(self, x):
batch_size, seq_len, _ = x.shape
q = self.wq(x).view(batch_size, seq_len, self.n_heads, self.head_dim)
k = self.wk(x).view(batch_size, seq_len, self.n_kv_heads, self.head_dim)
v = self.wv(x).view(batch_size, seq_len, self.n_kv_heads, self.head_dim)
# 关键步骤：将KV头扩展到与Q头数一致
# 每个KV头被 (n_heads // n_kv_heads) 个Q头共享
k = k.repeat_interleave(self.n_heads // self.n_kv_heads, dim=2)
v = v.repeat_interleave(self.n_heads // self.n_kv_heads, dim=2)
# 后续就是标准的attention计算
# ...嗯，这里要注意repeat_interleave这一步。它把KV头复制到和Q头一样多，这样就能直接做矩阵乘法了。虽然计算时KV头数变多了，但缓存里只存了原始的少量KV头，所以显存还是省下来了。
1.7 小结
MQA和GQA的核心思想就是减少KV头的数量。MQA把所有头合并成一个，GQA分成几个组。两者都能显著降低显存占用和带宽需求，而且质量损失很小。
我在实际项目中，90%的情况都用GQA-8。既省显存，又不用担心质量。如果你也在做推理部署，不妨试试这个方案。
记住：没有银弹。具体用哪个，还是要根据你的模型、任务和硬件来定。多测试，多对比，总能找到最适合的方案。