Multi Query Attention
Multi-Query Attention
MQA 保留多个 Query heads,但让所有 Query heads 共享一组 K/V。它针对的是 decode 的具体瓶颈:每生成一个 token,模型都要读取不断增长的历史 K/V;如果每个 Q head 都有独立 KV,读取量和缓存容量会随 Hq 线性增加。
最小机制
1 | Q: [B,Hq,Tq,Dh] |
注意 MQA 只共享 K/V,不共享 Q,也不把所有 head 合成一个 attention。多个 Q head 仍可以提出不同查询,只是它们读取同一个 KV 表示。
为什么 decode 更受益
粗略地说,MHA 的每层 KV Cache 元素量与 Hq×Tk×Dh 成正比,MQA 则变为 1×Tk×Dh。因果链是:KV head 减少 → 每步需要搬运的历史数据减少 → 内存带宽压力下降 → 同显存可容纳更多并发/更长上下文 → decode 吞吐或 TPOT 可能改善。
这条链在 prefill 中不一定同样显著,因为 prefill 的大矩阵计算更容易成为主导;实际收益还依赖 batch、kernel、硬件和序列长度。不要把“KV 更小”直接等同于“端到端一定更快”。
质量与训练边界
共享 KV 限制了不同 Q head 读取独立历史子空间的能力,可能造成质量退化,尤其在需要多种长距离关系的任务上。MQA 不是对现有 MHA 权重零成本的推理开关;转换、uptraining 或从头训练都需要重新验证。GQA 把共享程度作为连续旋钮,常用于在质量与效率之间取得中间点。
| 机制 | Hkv |
KV Cache | 主要风险 |
|---|---|---|---|
| MHA | Hq |
最大 | 带宽与显存压力 |
| GQA | 1 < Hkv < Hq |
中间 | 共享程度需调参 |
| MQA | 1 |
最小 | 表达能力可能下降 |
关系与验证
MQA 是 GQA 在 Hkv=1 时的特例,向下依赖 KV Cache 的布局,向上影响 Autoregressive Generation 的 decode 阶段。验证时固定模型质量集、batch/长度分布和硬件,分别测 KV 峰值、TPOT、吞吐与长上下文任务,而不是只测单请求首 token。
参考资料
从 MHA 转成 MQA 不是改一个维度
已有 MHA checkpoint 若直接把多个 K/V 投影平均或挑一个 head,得到的是一种近似初始化,不是等价变换。不同 KV head 可能承载互补信息,压成一个之后模型需要继续 uptraining 才能恢复部分质量。转换实验应把“权重转换误差”和“结构本身的质量损失”分开:先固定 tokenizer、采样和评测,比较转换前后 logits,再短训若干步观察能否恢复。
什么时候 MQA 的收益不明显
若 batch 很大、prefill GEMM 已经占满 GPU,KV 读取不是主导项;若 kernel 为了适配矩阵乘而实际复制 K/V,理论上的缓存节省也不会转化成相同的带宽收益。多租户 serving 还会受到请求长度分布和 cache 碎片影响,因此要记录每层实际读取字节、显存峰值、TPOT 和并发曲线,而不是只比较单请求速度。
MQA 适合把“decode 历史读取”作为第一约束的场景;如果任务对长距离关系极其敏感或质量预算紧,先尝试 GQA 的中间 Hkv,并用长上下文和事实一致性集验证共享造成的退化。