MQA: Multi-Query Attention (2019, Shazeer et al.)

9. MQA: Multi-Query Attention (2019, Shazeer et al.)

arxiv: 1911.02150

核心问题：Transformer 的推理瓶颈在哪里？不是在计算（前向传播很快），而是在内存带宽——每个 token 生成时都要把巨大的 Key 和 Value 张量从显存加载到计算单元。多头注意力（MHA）中，每个 head 有自己的 K 和 V，KV cache 大小 = n_heads × d_head × seq_len。当模型有 96 个 head、128 维 head、4K 上下文时，这个 cache 是 GB 级的。怎么在不严重损害质量的前提下砍掉它？

方法创新：
MQA 的方案极其激进：所有 query head 共享同一组 K 和 V。

在 MHA 中：

• Q: [batch, n_heads, seq_len, d_head]
• K: [batch, n_heads, seq_len, d_head] ← n_heads 份
• V: [batch, n_heads, seq_len, d_head] ← n_heads 份

在 MQA 中：

• Q: [batch, n_heads, seq_len, d_head] ← 保持不变
• K: [batch, 1, seq_len, d_head] ← 只有 1 份！
• V: [batch, 1, seq_len, d_head] ← 只有 1 份！

KV cache 从 n_heads 份降到 1 份，内存带宽需求 drastic 降低。解码速度大幅提升。

代价：所有 head 看同一个 "记忆"，失去了不同 head 关注不同子空间的能力。质量下降。

关键数字：

• "much faster to decode"
• "only minor quality degradation from the baseline"
• 作者：Noam Shazeer（Transformer 作者之一，后来去了 Character.AI）

影响评估：
MQA 是注意力"瘦身"的第一步。它证明了一个重要原则：KV cache 是推理瓶颈，不是计算。后续所有注意力优化（GQA、MLA、SWA）都围绕"减少 KV cache"展开。但 MQA 的质量下降让它没有被主流模型直接采用——它更像一个"思想实验"，证明了瘦身的可能性，同时也证明了瘦太多会伤脑子。

费曼点评：

MQA 的真正价值是教会你识别真正的瓶颈。大多数人以为 Transformer 推理慢是因为 attention 计算量大——不，慢的是内存带宽。注意力矩阵的计算是 O(n²)，但每个 token 只算一行，实际很快。慢的是把 KV cache 从显存搬到计算单元的过程。MQA 不是"更好的注意力"，是"把注意力中不需要复制的部分删掉"。就像搬家时不是扔家具，是把 96 把一模一样的椅子合成 1 把。

arxiv: 1911.02150

#论文深度研究 #小凯