维度压缩 MLA
不减 entry 数量，而是减小每个 token 的 KV——共享 head 或压到低秩 latent
核心要点：
- 维度压缩减每 token 的 KV 大小，与序列压缩正交
- MQA/GQA 减少 KV head 数（head 共享）
- MLA 把 KV 联合压到低秩 latent，不减 head
- 权重吸收让推理直接在 latent 上算，不展开 KV
- decoupled RoPE 绕开 RoPE 对权重吸收的阻断
本篇是注意力机制章实现卷的**压缩化（维度）**篇，是 MLA / GQA / MQA 的唯一主词条。它与09-序列压缩注意力同属压缩化机制族但作用轴正交：序列压缩减 entry 数量，本篇减每 entry 维度。
名词定义
| 名词 | 定义 |
|---|---|
| 维度压缩 | 减小每个 token 缓存的 KV 大小（head 数或 head 维度） |
| KV head 共享 | 多个 query head 共用一组 K/V head（MQA/GQA） |
| MQA (Multi-Query Attention) | 所有 query head 共享 1 组 KV head |
| GQA (Grouped-Query Attention) | query head 分组，每组共享一组 KV head |
| MLA (Multi-head Latent Attention) | 把 K/V 联合投到低秩 latent 存储，推理时上投影恢复 |
| 低秩 latent | MLA 缓存的低维向量 $c^{KV}$，$d_c \ll n_h d_h$ |
| 权重吸收 | 把上投影矩阵融进 Q/O 投影，推理直接在 latent 算 attention |
| Decoupled RoPE | 把位置编码拆到独立维度，绕开 RoPE 对权重吸收的阻断 |
@tbl-attn-dimcomp-glossary 本篇名词定义
维度压缩怎么减小每 token 的 KV？
维度压缩不动 entry 数量，而是减小单个 token 缓存的 KV 体积——这是它与序列压缩（09）的根本分野。单 token 单层 KV 大小是 $2 \cdot n_{\text{kv}} \cdot d_h$，维度压缩从这两个因子下手：
- 减 KV head 数 $n_{\text{kv}}$:多个 query head 共享一组 K/V（MQA/GQA），head-level sharing。
- 压 head 维度:把 K/V 联合投到低秩 latent $d_c \ll n_h d_h$（MLA），latent-level joint compression。
两条路本质不同：前者减少"KV 的份数"，后者把"每份 KV"压成低秩表示。下面先讲简单的 head 共享，再讲 MLA。
MQA 和 GQA 怎么靠共享 KV head 省显存？
MQA 让所有 query head 共享一组 KV、GQA 让 query 分组各享一组，用减少 KV head 数换显存[1][2]。它们不改 attention 算子，只改 KV head 的数量。
- MQA:KV head 数降到 1，所有 query head 共享。KV 体积降到 $1/n_h$，但只剩"一种聚合视角"，表达力损失明显——大模型已少用。
- GQA:query head 分 $G = n_q / n_{\text{kv}}$ 组，每组共享一组 KV。KV 体积降到 $1/G$（典型 $n_{\text{kv}}=8$），质量接近 MHA。
现状:GQA 是当前 dense 模型的事实标准（LLaMA 3 / Qwen / Mistral / Gemma 默认），因为它实现简单、与 FlashAttention 兼容、质量损失可控。MQA 是它的极端情形($G=n_h$)。
head 共享的局限：它只能整组砍 KV head，粒度粗。MLA 走另一条更细的路——不减 head，把 KV 压到低秩。
MLA 怎么把 KV 压到低秩 latent？
MLA 把 K 和 V 联合投影到一个低维 latent 向量，推理时只 cache 这个 latent、用时上投影恢复[3]。它不减 head 数（保留 MHA 的表达力），只压维度。
下投影把隐状态压到 $d_c$ 维 latent:
$$\begin{equation} c_t^{KV} = W^{DKV} h_t \in \mathbb{R}^{d_c} \label{eq:attn-dimcomp-mla-down} \end{equation}$$每个 head 的 K/V 由 latent 上投影恢复：
$$\begin{equation} k_{t,i}^C = W^{UK}_i\, c_t^{KV}, \qquad v_{t,i}^C = W^{UV}_i\, c_t^{KV}, \qquad i = 1, \ldots, n_h \label{eq:attn-dimcomp-mla-up} \end{equation}$$推理时只 cache $c_t^{KV}$，不 cache 展开后的 $k, v$。DeepSeek-V2 配置：$d_c = 512$、$n_h=128$、$d_h=128$,KV latent 维度约 $4 d_h$——所有 head 的 K/V 共享这一个 $c^{KV}$。
把 KV 看成一个秩为 $d_c$ 的低秩矩阵对（$W^{UK}, W^{UV}$ 是解码基）——这是 MLA 与 head 共享的根本不同：head 共享减少基的数量,MLA 压低每基的秩。
权重吸收为什么能让推理不展开 KV？
K/V 的上投影是线性的，可以"吸收"进相邻的 Q 投影和输出投影，让推理直接在 $d_c$ 维 latent 上算 attention、不显式恢复 KV[3]。这是 MLA 推理高效的关键——否则每步都要上投影展开全 KV，省了存储却没省计算。
K 端吸收:attention 分数原本是 $k^\top q$，展开 latent:
$$\begin{equation} S_h = (W^{UK}_h c^{KV})^\top (W^{UQ}_h c^Q) = (c^{KV})^\top \underbrace{(W^{UK}_h)^\top W^{UQ}_h}_{W^{KQ}_h}\, c^Q \label{eq:attn-dimcomp-absorb-k} \end{equation}$$$W^{KQ}_h = (W^{UK}_h)^\top W^{UQ}_h$ 是 $d_c \times d_c$ 的常数矩阵，可离线预融合——分数直接在 latent 上算，不展开 $k$。
V 端吸收:输出原本是 $W^O \sum_h \text{attn}_h \cdot v_h$，同理把 $W^{UV}$ 吸进 $W^O$:
$$\begin{equation} W^{OV} = W^O W^{UV} \label{eq:attn-dimcomp-absorb-v} \end{equation}$$直接把 attention 加权后的 $c^{KV}$ 映射到输出,不展开 $v$。
陷阱:权重吸收只在推理时做（$W^{KQ}, W^{OV}$ 离线预融合）;训练仍保持 $W^{UK}, W^{UQ}$ 等分开，以便独立求梯度。融合是推理优化，不改训练。
为什么 MLA 需要 decoupled RoPE？
RoPE 的位置旋转矩阵会夹在上投影矩阵中间，让权重吸收失效——MLA 把位置编码拆到独立维度绕开这个冲突[3]。这是 MLA 设计里最不直观但必须的一环。
问题在于：若对 $k$ 直接施 RoPE（位置 $i$ 的旋转 $R_i$），分数变成：
$$\begin{equation} S_{h,ij} = (c_i^{KV})^\top (W^{UK}_h)^\top R_i^\top R_j\, W^{UQ}_h\, c_j^Q \label{eq:attn-dimcomp-rope-block} \end{equation}$$$R_i^\top R_j$ 随位置变化，夹在 $(W^{UK})^\top$ 和 $W^{UQ}$ 之间,无法合并成常数 $W^{KQ}$，权重吸收完全失效。
decoupled RoPE 的解法:把每个 Q/K 拆成不含 RoPE 的内容部分 + 专门承载位置的 RoPE 部分，拼回：
$$\begin{equation} q_{t,h} = [q_{t,h}^C;\; q_{t,h}^R], \qquad k_{t,h} = [k_{t,h}^C;\; k_t^R] \label{eq:attn-dimcomp-decouple} \end{equation}$$其中 $k_t^R = \mathrm{RoPE}(W^{KR} h_t)$ 是所有 head 共享的单一向量（MQA-style），直接从 $h_t$ 投影、不经过 KV latent，因此绕开权重吸收路径。分数拆成两项之和（下式大写 $K/Q$ 表示各 head 堆叠后的矩阵）:
$$\begin{equation} S_h = (K^C)^\top Q^C + (K^R)^\top Q^R \label{eq:attn-dimcomp-decouple-score} \end{equation}$$内容项在 latent 上用权重吸收算，位置项 RoPE 独立算再叠加。推理 cache 存 $(c_t^{KV}, k_t^R)$，维度 $d_c + d_h^R$（V2 取 $d_h^R=64$），与权重吸收兼容。
MLA 和 GQA 的压缩本质差在哪？
GQA/MQA 是 head-level sharing（减 KV 份数），MLA 是 latent-joint-compression（压每份的秩）——在压缩空间的定义上根本不同。
| 维度 | GQA / MQA | MLA |
|---|---|---|
| 压缩轴 | 减 KV head 数 $n_{\text{kv}}$ | 压 KV 到低秩 latent $d_c$ |
| 缓存形式 | 完整 K/V（头数少） | latent $c^{KV}$ + 解耦 RoPE key $k^R$ |
| 头间关系 | 多 Q head 共享一组 KV head | 所有 head 共享一个 $c^{KV}$ latent |
| 表达力 | 低于 MHA（头数少损表达） | 论文报告优于 MHA（低秩未必损表达） |
| RoPE | 直接兼容 | 需 decoupled RoPE |
| 推理优化 | 不适用 | 权重吸收省 KV 展开 |
@tbl-attn-dimcomp-mla-vs-gqa MLA 与 GQA/MQA 的压缩本质对比
KV cache 大小（DeepSeek-V2，每 token 单层元素数）:MHA $2 n_h d_h = 32768$、MQA $2 d_h = 256$、MLA $d_c + d_h^R = 576$——MLA 等价于只有约 2.25 个 KV head 的 GQA，比 MHA 压约 98.3%;V2 报告 KV cache 减 93.3%、生成吞吐提升 5.76×[3]。
演进:MLA 是 DeepSeek V2/V3 的标志[4]。但 V4 转向序列压缩（CSA/HCA，见 09）替代 MLA，改用 compressor + MQA 共享 KV——序列方向的压缩比维度方向更激进。两条压缩轴的取舍仍在演化。
Takeaway
| 知识点 | 核心结论 |
|---|---|
| 维度压缩定位 | 减每 token KV 大小，与序列压缩正交（减 head 数 / 压 head 维度） |
| MQA/GQA | 减 KV head 数（共享），GQA-8 是 dense 模型默认 |
| MLA 低秩 | K/V 联合压到 latent $c^{KV}$（$d_c$=512），不减 head |
| 权重吸收 | $W^{UK}$ 吸进 Q、$W^{UV}$ 吸进 O，推理在 latent 算不展开 KV（仅推理） |
| decoupled RoPE | RoPE 旋转阻断吸收，故拆独立维度、$k^R$ 全 head 共享绕开 |
| MLA vs GQA | head 共享（减份数）vs latent 压缩（压秩），MLA 压约 98.3% |
@tbl-attn-dimcomp-takeaway 维度压缩核心知识点
参考资料
- Shazeer, Fast Transformer Decoding: One Write-Head is All You Need, arXiv:1911.02150, 2019. https://arxiv.org/abs/1911.02150
- Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, arXiv:2305.13245, 2023. https://arxiv.org/abs/2305.13245
- DeepSeek-AI, DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model, arXiv:2405.04434, 2024. https://arxiv.org/abs/2405.04434
- DeepSeek-AI, DeepSeek-V3 Technical Report, arXiv:2412.19437, 2024. https://arxiv.org/abs/2412.19437