MQA、GQA、MLA 怎么减小 KV cache？
一个 token 的 KV cache 有多大？

英文问法How do MQA and GQA reduce the KV cache? · Must Q and K have the same dimension? · Implement grouped-query attention.

考过面经：NVIDIA TRT-LLM 实习（"Q和K的dim必须一样吗？什么情况下不一样，提示是MQA和GQA的区别"，知乎）；他司：阿里 AI Infra 实习（「MHA、GQA与MLA」，牛客）、百度 AI Infra 提前批（「了解GQA、MLA、Hybrid架构吗？」，牛客）、滴滴（「GQA 结构下，一次推理需要多大的 KV Cache 显存？」，牛客）、Datadog MLE "Implement Grouped-Query Attention (GQA)"（PracHub）

先弄清楚：每个概念是什么

KV head: MHA 里每个 head 都有自己的 K、V，decode 时都要存进 KV cache，每步还要整个读一遍（memory-bound）。KV head 越少，decode 越快。
MQA（multi-query attention）: 所有 query head 共用 1 组 K、V，cache 省 h 倍，quality 略降。例：原论文的 decoder 每个 token 从 46 µs 降到 3.8 µs。
GQA（grouped-query attention）: query head 分成 g 组，每组共用一组 K、V。例：LLaMA 3 8B 有 32 个 query head、8 个 KV head。
MLA（multi-head latent attention）: 不存 K、V，只存一个低维 latent 向量。例：DeepSeek-V2 每层每 token 只存 512 + 64 = 576 个数。

30 秒回答

In multi-head attention every query head has its own K and V, and decoding caches all of them: 2 × layers × KV heads × head dim × bytes per token. Decoding is memory-bound, so this cache limits speed, batch size and context length. MQA shares one K/V head across all query heads, shrinking the cache by the head count at some quality cost. GQA shares one per group and stays close to MHA quality. MLA in DeepSeek-V2 caches a low-rank latent per token, about as small as GQA with 2.25 groups.

关键公式与手算

\[ \text{KV cache / token} = 2 \cdot L \cdot n_{kv} \cdot d_{head} \cdot \text{bytes} \]

\(n_{kv}\) 在 MHA 是 \(h\)，GQA 是 \(g\)，MQA 是 1。总量再乘 context 长度和 batch。

bf16（2 bytes），\(d_{head} = 128\)：

LLaMA 2 7B（MHA，32 层，32 个 head）：2 × 32 × 32 × 128 × 2 = 512 KiB/token
LLaMA 3 8B（GQA，8 个 KV head）：2 × 32 × 8 × 128 × 2 = 128 KiB/token，小 4 倍。128K context 一条序列 16 GiB，和 weights（8B × 2 bytes）差不多大
同样形状换成 MQA：16 KiB/token

	每层每 token 存几个数	谁在用	quality
MHA	\(2\,h\,d_h\)	LLaMA 2 7B	baseline
MQA	\(2\,d_h\)（÷ \(h\)）	PaLM	略降，可能训练不稳
GQA	\(2\,g\,d_h\)（÷ \(h/g\)）	LLaMA 2 34B、70B；LLaMA 3 全系 \(g = 8\)	接近 MHA
MLA	\(d_c + d_R\)	DeepSeek-V2	论文称强于 MHA

MLA 怎么做

\[ c_t = W^{DKV} h_t \in \mathbb{R}^{d_c},\qquad k_t = W^{UK} c_t,\qquad v_t = W^{UV} c_t \]

只缓存 \(c_t\)；\(W^{UK}\) 可以并进 \(W^Q\)，\(W^{UV}\) 并进 \(W^O\)，推理时不用还原 K、V。DeepSeek-V2 比同形状的 MHA（128 个 head）小约 57 倍。

手撕 GQA（NumPy，省略 mask）

def gqa(Q, K, V):
    # Q: (h, n, dk)；K、V: (g, n, dk)，h 个 query head 分成 g 组
    r = Q.shape[0] // K.shape[0]          # 每组几个 query head
    K = np.repeat(K, r, axis=0)           # head i 用第 i // r 组
    V = np.repeat(V, r, axis=0)
    S = Q @ K.transpose(0, 2, 1) / np.sqrt(Q.shape[-1])
    P = np.exp(S - S.max(-1, keepdims=True))
    return (P / P.sum(-1, keepdims=True)) @ V

cache 里只存 \(g\) 组，repeat 只在算的时候做。完整练习见 ml-coding。

必答点

每 token 存 2 × 层数 × KV head 数 × d_head × bytes；decode 是 memory-bound
MQA：1 组 K、V，小 \(h\) 倍，quality 略降
GQA：\(g\) 组，小 \(h/g\) 倍，quality 接近 MHA；LLaMA 3 用 \(g = 8\)
MLA：缓存低维 latent，up-projection 并进 \(W^Q\)、\(W^O\)；DeepSeek-V2 相当于 2.25 组的 GQA
同一个 head 里 Q、K 维度必须相同；head 数可以不同

追问

Q 和 K 的 dim 必须一样吗？（NVIDIA 原题）

同一个 head 里必须一样，因为要做 dot product，但 head 数可以不同，就是 GQA、MQA。V 的 head dim 也可以不同，DeepSeek-V2 的 q、k 是 128 + 64 = 192 维，v 是 128 维。

MQA 更省，为什么大模型多用 GQA？

GQA 论文的 T5-XXL 每个样本推理 MHA 1.51 s、MQA 0.24 s、GQA-8 0.28 s，平均分 47.2、46.6、47.1，GQA 速度接近 MQA、分数接近 MHA。tensor parallel 时 KV head 少于 GPU 数就没法按 head 切，只能每张卡复制一份，KV cache 和 GQA 一样大。

MLA 和 RoPE 有什么冲突？

RoPE 矩阵随位置变化，夹在 \(W^Q\) 和 \(W^{UK}\) 中间，\(W^{UK}\) 并不进 \(W^Q\)，每步得重算前缀的 K。DeepSeek-V2 另加 \(d_R = 64\) 维、所有 head 共用的 key 带 RoPE，所以 cache 是 \((d_c + d_R)\cdot L\)。

来源Shazeer 2019, MQA §2.4.1（p.4）、§3–3.1（pp.5–6）、Table 1–2（p.8）· Ainslie et al. 2023, GQA §1–2.2（pp.1–2）、§3.1–3.2、Table 1（pp.2–3）· DeepSeek-V2 §2.1.1–2.1.4（pp.6–9，eq. 9–19，Table 1）、§3.1.2（p.12）· Touvron et al. 2023a, LLaMA Table 2（p.3）· Touvron et al. 2023b, Llama 2 §2.2、Table 1（pp.5–6）、Appendix A.2.1（p.47）· Llama 3 Herd §3.2 Table 3（p.7）

MQA、GQA、MLA 怎么减小 KV cache？一个 token 的 KV cache 有多大？