多头潜在注意力（MLA）

Multi-Head Latent Attention（MLA）是 DeepSeek 提出的一种高效注意力机制，通过对 KV 缓存进行低秩联合压缩，在显著降低推理阶段显存占用的同时，保持与标准多头注意力（MHA）相当的模型性能。其核心思想是将键（Key）和值（Value）投影到低维潜在空间，仅缓存压缩后的隐向量，并在计算时通过上投影恢复。以下是 MLA 的详细数学推导。

约定：所有计算使用行向量，即 $y = x W$ 。

1. Q 的计算

\begin{aligned} c_{t}^{Q} & = h_{t} W^{D Q}, \\ [q_{t, 1}^{C}; q_{t, 2}^{C}; \dots; q_{t, n_{h}}^{C}] = q_{t}^{C} & = c_{t}^{Q} W^{U Q}, \\ [q_{t, 1}^{R}; q_{t, 2}^{R}; \dots; q_{t, n_{h}}^{R}] = q_{t}^{R} & = RoPE (c_{t}^{Q} W^{Q R}), \\ q_{t, i} & = [q_{t, i}^{C}; q_{t, i}^{R}], \end{aligned}

其中：

$c_{t}^{Q} \in R^{d_{c}^{'}}$ 是查询的压缩隐向量；
$d_{c}^{'} (≪ d_{h} n_{h})$ 表示查询压缩维度；
$W^{D Q} \in R^{d \times d_{c}^{'}}$ 、 $W^{U Q} \in R^{d_{c}^{'} \times d_{h} n_{h}}$ 分别是查询的下投影和上投影矩阵；
$W^{Q R} \in R^{d_{c}^{'} \times d_{h}^{R} n_{h}}$ 用于生成携带 RoPE 的解耦查询。

2. KV 的计算

\begin{aligned} c_{t}^{K V} & = h_{t} W^{D K V}, \\ [k_{t, 1}^{C}; k_{t, 2}^{C}; \dots; k_{t, n_{h}}^{C}] = k_{t}^{C} & = c_{t}^{K V} W^{U K}, \\ k_{t}^{R} & = RoPE (h_{t} W^{K R}), \\ k_{t, i} & = [k_{t, i}^{C}; k_{t}^{R}], \\ [v_{t, 1}^{C}; v_{t, 2}^{C}; \dots; v_{t, n_{h}}^{C}] = v_{t}^{C} & = c_{t}^{K V} W^{U V}, \end{aligned}

其中：

$c_{t}^{K V} \in R^{d_{c}}$ 是键值的压缩隐向量；
$d_{c} (≪ d_{h} n_{h})$ 表示 KV 压缩维度；
$W^{D K V} \in R^{d \times d_{c}}$ 是下投影矩阵；
$W^{U K}, W^{U V} \in R^{d_{c} \times d_{h} n_{h}}$ 是键和值的上投影矩阵；
$W^{K R} \in R^{d \times d_{h}^{R}}$ 用于生成携带 RoPE 的解耦键；
$RoPE (\cdot)$ 表示应用旋转位置编码的操作。

注意：对于 MLA，仅需缓存蓝色框中的向量（ $c_{t}^{K V}$ 和 $k_{t}^{R}$ ），从而显著减少 KV 缓存大小，同时保持与标准多头注意力（MHA）相当的性能。

最终，注意力查询（ $q_{t, i}$ ）、键（ $k_{j, i}$ ）和值（ $v_{j, i}^{C}$ ）组合得到最终输出 $u_{t}$ ：

\begin{aligned} o_{t, i} & = \sum_{j = 1}^{t} {Softmax}_{j} (\frac{q_{t, i} k_{j, i}^{⊤}}{\sqrt{d_{h} + d_{h}^{R}}}) v_{j, i}^{C}, \\ u_{t} & = [o_{t, 1}; o_{t, 2}; \dots; o_{t, n_{h}}] W^{O}, \end{aligned}

其中 $W^{O} \in R^{d_{h} n_{h} \times d}$ 是输出投影矩阵。

3. 实际参数配置

$d = hidden_size = 7168$
$d_{c} = kv_lora_rank = 512$
$d_{c}^{'} = q_lora_rank = 1536$
$n_{h} = num_heads = 128$
$d_{h} = qk_nope_head_dim = 128$
$d_{h}^{R} = qk_rope_head_dim = 64$

此外：

$W^{U Q}$ 和 $W^{Q R}$ 可合并， $q_h e a d_d i m = q k_n o p e_h e a d_d i m + q k_r o p e_h e a d_d i m = 192$ 。
$W^{D K V}$ 和 $W^{K R}$ 可合并， $k v_l o r a_r a n k + q k_r o p e_h e a d_d i m = 576$ 。

4. 矩阵吸收（Absorb）

考虑如下计算：

Y = X A B, C = A B

其中：

$X \in R^{m \times d}$ 是输入隐状态（hidden states）；
$A \in R^{d \times d_{c}}$ 、 $B \in R^{d_{c} \times n}$ 是权重矩阵；
$C \in R^{d \times n}$ 是 absorb 后的等效权重矩阵。

直接计算的 FLOPs 为：

2 m d d_{c} + 2 m n d_{c} = 2 m d_{c} (d + n)

合并权重后计算的 FLOPs 为：

2 m d n

当 $d_{c}$ 较小时，通常有：

d n > d_{c} (d + n)

因此不一定需要合并两个权重矩阵。

不考虑 RoPE 部分，仅从 $c^{Q}$ 和 $c^{K V}$ 计算 $q_{i} k_{i}^{⊤}$ （第 $i$ 个 head）：

\begin{aligned} q_{i} k_{i}^{⊤} & = c^{Q} W_{i}^{U Q} (c^{K V} W_{i}^{U K})^{⊤}, \\ = c^{Q} W_{i}^{U Q} (W_{i}^{U K})^{⊤} (c^{K V})^{⊤}, \\ = q_{i} (W_{i}^{U K})^{⊤} (c^{K V})^{⊤}, & (Absorb) \\ = q_{i} (c^{K V} W_{i}^{U K})^{⊤}, & (Normal) \end{aligned}

警告：此处 "Absorb" 的真实含义是利用矩阵乘法结合律，优先将 $q$ 与 $W^{U K}$ 结合，并缓存压缩隐向量 $c^{K V}$ 。它并非合并权重矩阵，"Absorb" 这一命名具有一定误导性。

4.1 为什么计算时不把 $W_{i}^{U Q} (W_{i}^{U K})^{⊤}$ 合并

对单个 token、单个 head，FLOPs 分别为：

分开计算： $2 d_{h} (d_{c}^{'} + d_{c}) = 524,288$
合并计算： $2 d_{c}^{'} d_{c} = 1,572,864 = 3 \times 524,288$

合并后计算量反而是原来的 3 倍。

4.2 为什么 Prefill 阶段显式计算 k 和 v，而 Decode 阶段不需要

假设输入 shape 如下：

$q : (b, n_{h}, s_{q}, d_{h})$
$c^{K V} : (b, 1, s_{k v}, d_{c})$
$W^{U K} : (d_{c}, n_{h} d_{h})$

4.2.1 Prefill 阶段（ $s_{q} = s_{k v} = s$ ）

FLOPs 对比：

\begin{aligned} T_{Normal} & = 2 b n_{h} d_{h} s (d_{c} + s), \\ T_{Absorb} & = 2 b n_{h} d_{c} s (d_{h} + s), \\ \frac{T_{Normal}}{T_{Absorb}} & = \frac{d_{h} (d_{c} + s)}{d_{c} (d_{h} + s)} = \frac{s + 512}{4 s + 512} \in (\frac{1}{4}, 1) \end{aligned}

Prefill 阶段 Normal 更快，且此阶段是计算瓶颈，故显式计算 $q$ 和 $k$ 。

4.2.2 Decode 阶段（ $s_{q} = 1, s_{k v} = s$ ）

FLOPs 对比：

\begin{aligned} T_{Normal}^{K} & = 2 b n_{h} d_{h} (d_{c} + s), & （缓存 k） \\ T_{Normal}^{L} & = 2 b n_{h} d_{h} (d_{c} s + s), & （缓存 latent） \\ T_{Absorb} & = 2 b n_{h} d_{c} (d_{h} + s), \\ \frac{T_{Normal}^{K}}{T_{Absorb}} & = \frac{d_{h} (d_{c} + s)}{d_{c} (d_{h} + s)} = \frac{s + 512}{4 s + 512} \in (0.25, 1), \\ \frac{T_{Normal}^{L}}{T_{Absorb}} & = \frac{513 s}{4 s + 512} \in (0.99, 128.25) \end{aligned}

虽然缓存 k 的计算量最小（极限为 Absorb 的 1/4），但 Decode 阶段瓶颈是显存带宽。

4.2.3 内存读取量对比（bfloat16 精度）

MLA（Absorb）： $(b, n_{h}, 1, d_{c}) \times (b, 1, s, d_{c})$
$M_{MLA} = 2 b d_{c} (n_{h} + s)$
标准 MHA： $(b, n_{h}, 1, d_{h}) \times (b, n_{h}, s, d_{h})$
$M_{MHA} = 2 b d_{h} n_{h} (1 + s)$

内存读取比例：

\frac{M_{MLA}}{M_{MHA}} = \frac{d_{c} (n_{h} + s)}{d_{h} n_{h} (1 + s)} = \frac{128 + s}{32 (1 + s)}

当 $s = 20$ 时，比值 $\approx 0.22$ ；
极限情况（ $s \to \infty$ ）：比值 $\to 1 / 32$ 。

因此，Decode 阶段采用 Absorb 方式，可大幅降低显存带宽压力，并复用 MQA（Multi-Query Attention）实现。

5. 矩阵吸收问题总结

"矩阵吸收"的本质是如何应用矩阵乘法结合律：

\begin{aligned} Y & = (X A) B = X (A B), \\ Z & = (X W) Y = X (W Y), \end{aligned}

其中 $A, B, W$ 均为权重矩阵。

决策依据应综合权衡：

计算量（FLOPs）
显存读写量（Memory Traffic）
当前阶段瓶颈（计算 or 带宽）

可借助 Roofline Model 进行系统性分析。

多头潜在注意力（MLA） ​

1. Q 的计算 ​

2. KV 的计算 ​

3. 实际参数配置 ​

4. 矩阵吸收（Absorb） ​

4.1 为什么计算时不把 WiUQ(WiUK)⊤ 合并 ​

4.2 为什么 Prefill 阶段显式计算 k 和 v，而 Decode 阶段不需要 ​

4.2.1 Prefill 阶段（sq=skv=s） ​

4.2.2 Decode 阶段（sq=1,skv=s） ​

4.2.3 内存读取量对比（bfloat16 精度） ​

5. 矩阵吸收问题总结 ​