KVCache 架构压缩 — GQA、MQA、MLA
核心要点:
- KV head 共享按 MHA→GQA→MQA 递减
- GQA-8 省 8 倍 KV,质量损失可控
- MLA 投影到低秩潜空间存储
- 三者均为架构层手段,训练时定型
- GQA 是默认,MLA 是 DeepSeek 标志
本文回答 02-第一性挑战 中 KV $O(n)$ 显存问题的架构层缓解方案。本类方法改模型结构、降单 token KV 大小,与改 attention 计算复杂度(04-注意力机制变体)、训练后动态压缩(07-推理-KV 管理)三者正交,工业模型同时用。
为什么是架构层
架构层唯一可压缩的量是单 token KV 大小 $H_{\text{kv}} \cdot d_h$。回顾 KVCache 占用:
$$\begin{equation} M_{\text{KV}} = 2 \cdot L \cdot H_{\text{kv}} \cdot d_h \cdot n \cdot b \label{eq:lc-kvarch-kv-memory} \end{equation}$$序列长度 $n$ 和层数 $L$ 不能动,bytes $b$ 由量化决定(→ 07-推理-KV 管理)。剩下可在架构层下手的就是 $H_{\text{kv}} \cdot d_h$ —— "KV 多大"。
| 思路 | 操作 | 代表方法 |
|---|---|---|
| 减少 KV head 数 | $H_{\text{kv}}$ 从 query head 数变小,多个 query 共享同一份 KV | MQA、GQA |
| 降维度 | $d_h$ 通过低秩投影变小 | MLA(潜空间) |
| 两者结合 | 共享 + 低秩 | 部分 hybrid 设计 |
@tbl-longctx-kvarch-paths KV 架构压缩的三类思路
MHA:原版
核心问题:KV 体积的 baseline 来自哪里?
| 项 | 值 |
|---|---|
| Query head 数 $H_q$ | $H$(如 LLaMA 3 70B 是 64) |
| KV head 数 $H_{\text{kv}}$ | $H$(与 query 一致) |
| 单 token 单层 KV 大小 | $2 \cdot H \cdot d_h$ |
@tbl-longctx-kvarch-mha MHA 配置
每个 query head 有独立的 K 和 V。模型表达力最强,但 KV 体积最大。
MQA:所有 query 共享一份 KV
核心问题:MQA 用什么代价换 KV 体积?
Shazeer, 2019[1] 提出 Multi-Query Attention:
| 项 | 值 |
|---|---|
| Query head 数 | $H$ |
| KV head 数 | 1(所有 query head 共享) |
| KV 体积相对 MHA | $1/H$ |
@tbl-longctx-kvarch-mqa MQA 配置
优势:KV 直接降到 $1/H$(LLaMA 3 70B 量级降到 ~1.5%)。decode 阶段读 KV 极快,TPOT 大幅改善。
代价:所有 query 共享 KV,模型表达力损失——只有"一种聚合视角"。早期 PaLM 用 MQA,但 LLaMA 等大模型逐步退化到 GQA 平衡。
GQA:分组共享
核心问题:GQA 如何在 MQA 与 MHA 之间取平衡?
Ainslie et al., 2023[2] 提出 Grouped-Query Attention,是 MQA 与 MHA 的折中:
| 项 | 值 |
|---|---|
| Query head 数 | $H_q$ |
| KV head 数 | $H_{\text{kv}}$($1 < H_{\text{kv}} < H_q$,典型 8) |
| 分组比 | $G = H_q / H_{\text{kv}}$,每 $G$ 个 query 共享一份 KV |
@tbl-longctx-kvarch-gqa GQA 配置
LLaMA 3 70B 实例:$H_q = 64, H_{\text{kv}} = 8, G = 8$,KV 体积是 MHA 的 $1/8$。
优势:
- KV 直接降到 $1/G$
- 比 MQA 保留更多"聚合视角",质量接近 MHA
- 实现简单,与 FlashAttention 等 kernel 兼容
业界现状:LLaMA 2 70B 开始引入 GQA,LLaMA 3 系列、Qwen2.5、Mistral、Gemma 等几乎所有现代开源模型默认 GQA。GQA 是当前 dense 架构的事实标准。
MLA:低秩潜空间投影
核心问题:不减 head 数,还能怎样降 KV 体积?
DeepSeek-V2 / DeepSeek-AI, 2024[3] 提出 Multi-head Latent Attention:不减 head 数,把 KV 整体投到低秩潜空间存储。
核心机制(机制详见 10,本篇只对比 KV 体积)
MLA 把 K/V 联合投到低秩潜空间 $c^{KV}_t$(DeepSeek-V2 取 $d_c=512$,V3[4] 沿用),推理时只存这个潜向量、用时上投影恢复;靠权重吸收让推理直接在潜空间算 attention、不展开 KV;RoPE 需解耦处理(decoupled RoPE)。
单 token 单层 KV 实际存 $(d_c + d_h^R) \cdot b = (512 + 64) \cdot 2 = 1152$ bytes,相对 MHA baseline 约 3.5%,即 V2 报告"KV cache 减少 93.3%"[3](等效 GQA group≈2.25)。同报告定性断言质量排序 MLA > MHA > GQA,但精确 delta 未在报告中拆出。本篇只取此体积用于下方架构对比。
完整机制(低秩投影公式、权重吸收 $W^{KQ}=(W^{UK})^\top W^{UQ}$、为何需要 decoupled RoPE)见 06-大模型解构 / 10-维度压缩 MLA——MLA 的算子实现已迁那里做唯一主词条,本篇(长上下文视角)只对比它的 KV 体积。V4 起改用序列压缩替代 MLA,见 07-前沿模型追踪 / DeepSeek-V4。
三方法对比
| 维度 | MHA | MQA | GQA | MLA |
|---|---|---|---|---|
| KV head 数 | $H$ | 1 | $H_{\text{kv}}$ | 等效低秩 $d_c$ |
| KV 体积(相对 MHA) | $1$ | $1/H$ | $H_{\text{kv}}/H$ | $\sim (d_c + d_h^R) / (H \cdot d_h)$ |
| 模型质量 | baseline | 退化明显 | 略逊于 MHA(同报告定性)[3] | 优于 MHA(V2 报告定性;PPL/MMLU 精确 delta 未在报告中独立拆出)[3] |
| 实现复杂度 | 标准 | 简单 | 简单 | 中等(需权重吸收) |
| 与 RoPE 兼容性 | 直接兼容 | 直接兼容 | 直接兼容 | 需特别处理(V2 论文有专门设计) |
| 工业采用 | 早期模型 | 已少用 | 主流 dense 模型默认 | DeepSeek 系标志 |
@tbl-longctx-kvarch-comparison MHA / MQA / GQA / MLA 对比
实际 KV 节省对比(LLaMA 3 70B 配置基准)
以单 batch、$L=80$ 全模型、序列长度 1M token(按 $1\text{M} = 2^{20}$ 二进制约定,与 Doc 02 § KV 数量级直观 口径一致)为例:
| 方案 | 单 token KV (bytes) | 全模型 1M token KV (GiB),$L=80$ |
|---|---|---|
| MHA(虚构对照) | $2 \cdot 64 \cdot 128 \cdot 2 = 32768$ | 2560 (~2.5 TiB) |
| GQA-8(LLaMA 3 70B 实际) | $2 \cdot 8 \cdot 128 \cdot 2 = 4096$ | 320 |
| MQA-1(早期方案) | $2 \cdot 1 \cdot 128 \cdot 2 = 512$ | 40 |
| MLA(DeepSeek-V2 风格,$d_c{=}512, d_h^R{=}64$) | 1152 | 90 |
如 所示,长上下文场景下即使是 GQA-8 在 80 层全模型 1M 单 batch 仍达 320 GiB,远超任何单卡乃至超节点显存。这是长上下文必须叠加推理层 KV 管理(→ 07-推理-KV 管理)和并行切分(→ 上下文并行 CP)的根本原因。
@tbl-longctx-kvarch-savings 1M token、$L=80$ 全模型 KV 大小:不同架构方案对比
与 GQA / MQA / MLA 的实际选型
| 场景 | 推荐方案 | 理由 |
|---|---|---|
| 新训 dense 模型 | GQA-8 | 工业标准,性价比最高,兼容性最好 |
| 极致 KV 压缩 + 自研栈 | MLA | DeepSeek 路线,KV 体积最小,需配套残差 / 位置编码设计 |
| 已有 MHA 模型转长上下文 | 蒸馏到 GQA | 见 Ainslie 2023[2] 的"up-training"方案 |
| 边缘 / 端侧极致省 KV | MQA / GQA-2 | 质量妥协换 KV 极小 |
@tbl-longctx-kvarch-selection 架构层 KV 压缩的实际选型
Takeaway
| 知识点 | 核心结论 |
|---|---|
| KV head 共享 | MHA→GQA→MQA,KV head 数递减,体积随之降 |
| GQA-8 | 省 8 倍 KV,质量接近 MHA,当前 dense 模型默认 |
| MQA | KV 降到 $1/H$,但表达力损失明显,大模型已少用 |
| MLA | 低秩潜空间存储,KV 体积最小,DeepSeek 系标志 |
| 架构层定位 | 训练时定型,与推理层动态压缩、并行切分正交 |
| 选型 | 端侧 / 极致省 KV → MQA / GQA-2;通用 dense → GQA-8;自研栈极致压缩 → MLA |
@tbl-longctx-kvarch-takeaway 全文要点
参考资料
- Shazeer, Fast Transformer Decoding: One Write-Head is All You Need (MQA), 2019. https://arxiv.org/abs/1911.02150
- Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, 2023. https://arxiv.org/abs/2305.13245
- DeepSeek-AI, DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model (含 MLA),2024. https://arxiv.org/abs/2405.04434
- DeepSeek-AI, DeepSeek-V3 Technical Report, 2024. https://arxiv.org/abs/2412.19437