从 MHA、GQA 到 MLA：大模型注意力机制 KV Cache 极致压缩原理与工程实践

在大语言模型（LLM）向超长上下文（128K 乃至 1M Tokens）与高并发生产服务演进的过程中，推理系统的核心物理瓶颈已从模型权重显存转移至 KV Cache（键值缓存）的显存占用与访存带宽。当上下文长度与并发批次（Batch Size）同步扩大时，庞大的 KV Cache 会迅速耗尽高带宽显存（HBM），限制系统最大吞吐上限。

为了在保持模型表征能力的同时压减 KV Cache 体积，注意力机制架构经历了从 MHA、MQA、GQA 到 MLA（Multi-Head Latent Attention，多头潜在注意力） 的深度演进。本文系统梳理注意力机制的显存压缩脉络，重点推导 MLA 的低秩联合压缩、推理期“矩阵吸收”加速以及解耦 RoPE 的数学原理，并给出生产级代码实现与工程部署避坑指南。

一、 注意力机制显存演进：MHA、MQA 与 GQA 的权衡

在标准自回归解码中，单个 Token 在每层产生的 KV Cache 显存占用（以 FP16/BF16 精度，每元素 2 字节计算）为：

$$\text{Bytes}_{\text{KV}} = 2 \times 2 \times n_{\text{kv}} \times d_h = 4 \cdot n_{\text{kv}} \cdot d_h$$

其中 $n_{\text{kv}}$ 为键值注意力头数，$d_h$ 为单头维度。

围绕削减 $n_{\text{kv}}$，业界先后演化出三代经典变体：

二、 MLA（多头潜在注意力）核心数学原理

以 DeepSeek-V2/V3 为代表的现代架构提出了 MLA（Multi-Head Latent Attention）。其核心洞察在于：传统 KV 矩阵在注意力头维度上存在大量冗余，可以通过低秩投影将其联合压缩为一个极紧凑的“潜在向量（Latent Vector）”进行缓存，在保持甚至超越 MHA 精度的同时，将 KV Cache 压缩至比 MQA 更低的水平。

1. 低秩联合压缩（Low-Rank Joint Compression）

对于第 $t$ 个位置的输入隐状态 $h_t \in \mathbb{R}^d$，MLA 不再直接投影出庞大的多头 $K_t, V_t$，而是通过一个下投影矩阵 $W^{DKV} \in \mathbb{R}^{d_c \times d}$ 将其压缩为低维潜在向量 $c_t^{KV}$：

$$c_t^{KV} = W^{DKV} h_t, \quad c_t^{KV} \in \mathbb{R}^{d_c}$$

其中压缩维度 $d_c \ll n_h d_h$（例如 $d_c = 512$，而 $n_h d_h = 128 \times 128 = 16384$）。

在需要计算注意力时，再通过两个上投影矩阵 $W^{UK}, W^{UV} \in \mathbb{R}^{(n_h d_h) \times d_c}$ 还原出多头键向量与值向量：

$$k_t^C = W^{UK} c_t^{KV}, \quad v_t^C = W^{UV} c_t^{KV}$$

显存收益：在推理阶段，系统只需在 KV Cache 中缓存紧凑的 $c_t^{KV}$（而非还原后的 $k_t^C$ 与 $v_t^C$），单层 KV Cache 显存占用直接从 $2 n_h d_h$ 骤降至 $d_c$，实现超过 90% 的显存削减。

2. 推理期“矩阵吸收（Matrix Absorption）”加速

如果每次解码都先将历史所有缓存的 $c_s^{KV}$ ($s \le t$) 通过 $W^{UK}, W^{UV}$ 解压缩为完整的多头矩阵再计算注意力，依然会消耗大量显存带宽与计算量。

利用矩阵乘法的结合律，MLA 在推理阶段将上投影矩阵直接“吸收”进相邻的线性层中：

• 吸收 $W^{UK}$ 到 Query 投影中：
  对于第 $i$ 个注意力头的查询向量 $q_{t,i}^C$ 与历史第 $s$ 个键向量 $k_{s,i}^C$ 的点积：

  $$(q_{t,i}^C)^T k_{s,i}^C = (q_{t,i}^C)^T \left(W_i^{UK} c_s^{KV}\right) = \left((W_i^{UK})^T q_{t,i}^C\right)^T c_s^{KV}$$

  推理时只需对当前单个 Query 向量预先左乘 $(W_i^{UK})^T$，即可直接与缓存的压缩向量 $c_s^{KV}$ 做点积，完全省去解压历史 $K$ 的操作。
• 吸收 $W^{UV}$ 到输出投影 $W^O$ 中：
  注意力加权求和输出为：

  $$o_{t,i} = \sum_{s=1}^t \alpha_{t,s,i} v_{s,i}^C = \sum_{s=1}^t \alpha_{t,s,i} \left(W_i^{UV} c_s^{KV}\right) = W_i^{UV} \left(\sum_{s=1}^t \alpha_{t,s,i} c_s^{KV}\right)$$

  最终的多头输出拼接投影 $u_t = \sum_{i=1}^{n_h} W_i^O o_{t,i} = \sum_{i=1}^{n_h} (W_i^O W_i^{UV}) \left(\sum_{s=1}^t \alpha_{t,s,i} c_s^{KV}\right)$。直接将 $W_i^O W_i^{UV}$ 离线合并为单个等价矩阵，彻底消除 $V$ 的解压缩开销。

3. 冲突破解：解耦旋转位置编码（Decoupled RoPE）

上述“矩阵吸收”成立的前提是 $q_{t,i}^C$ 与 $k_{s,i}^C$ 之间没有位置相关的非线性算子阻隔。然而，若直接对 $q_t^C$ 和 $k_s^C$ 施加旋转位置编码（RoPE），由于旋转矩阵 $R_{t-s}$ 依赖相对位置 $(t-s)$，会破坏矩阵结合律：

$$(R_t W_i^{UQ} c_t^Q)^T (R_s W_i^{UK} c_s^{KV}) = (c_t^Q)^T (W_i^{UQ})^T R_{t-s} W_i^{UK} c_s^{KV}$$

$R_{t-s}$ 夹在中间导致无法离线合并 $(W_i^{UQ})^T$ 与 $W_i^{UK}$。

MLA 的解耦设计（Decoupled RoPE）：
将 Query 和 Key 拆分为两部分——不带位置编码的内容分支（NoPE）与专门携带位置信息的旋转分支（RoPE）：

1. 压缩潜在向量 $c_t^{KV}$ 生成的 $k_t^C$ 不施加 RoPE，保持矩阵吸收能力；
2. 额外开辟极小维度的投影路径（维度 $d_h^R = 64$），生成多头查询旋转分量 $q_{t,i}^R = \text{RoPE}(W_i^{QR} c_t^Q)$，以及所有注意力头共享的单个键旋转分量 $k_t^R = \text{RoPE}(W^{KR} h_t)$；
3. 最终拼接两者计算注意力打分：

   $$\text{Score}_{t,s,i} = \frac{(q_{t,i}^C)^T k_{s,i}^C + (q_{t,i}^R)^T k_s^R}{\sqrt{d_h + d_h^R}}$$

推理时，KV Cache 仅需保存 $[c_t^{KV}; k_t^R] \in \mathbb{R}^{d_c + d_h^R}$（如 $512 + 64 = 576$ 维），兼顾了 RoPE 相对位置建模与无损矩阵吸收。

三、 核心代码实现：支持矩阵吸收的 MLA 模块原型

以下代码基于 PyTorch 演示了 MLA 的完整前向计算及推理期矩阵吸收机制的等价实现：

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """标准 RoPE 旋转位置编码应用 [Batch, Heads, Seq_Len, Rope_Dim]"""
    half = x.shape[-1] // 2
    x1, x2 = x[..., :half], x[..., half:]
    rotated = torch.cat((-x2, x1), dim=-1)
    return (x * cos) + (rotated * sin)

class MultiHeadLatentAttention(nn.Module):
    def __init__(
        self,
        d_model: int = 2048,
        num_heads: int = 16,
        d_head_nope: int = 128,
        d_head_rope: int = 64,
        kv_lora_rank: int = 512,
        q_lora_rank: int = 1024
    ):
        super().__init__()
        self.num_heads = num_heads
        self.d_head_nope = d_head_nope
        self.d_head_rope = d_head_rope
        self.kv_lora_rank = kv_lora_rank
        self.scale = 1.0 / math.sqrt(d_head_nope + d_head_rope)

        # 1. Query 低秩压缩与升维投影
        self.w_dq = nn.Linear(d_model, q_lora_rank, bias=False)
        self.q_norm = nn.RMSNorm(q_lora_rank)
        self.w_uq = nn.Linear(q_lora_rank, num_heads * d_head_nope, bias=False)
        self.w_qr = nn.Linear(q_lora_rank, num_heads * d_head_rope, bias=False)

        # 2. KV 联合低秩压缩与升维投影 (含共享 RoPE Key 分支)
        # 输出维度: kv_lora_rank (c_kv) + d_head_rope (共享 k_rope)
        self.w_dkv_kr = nn.Linear(d_model, kv_lora_rank + d_head_rope, bias=False)
        self.kv_norm = nn.RMSNorm(kv_lora_rank)
        
        # 上投影矩阵: 将潜在向量映射回多头 NoPE Key 与 Value
        self.w_uk = nn.Linear(kv_lora_rank, num_heads * d_head_nope, bias=False)
        self.w_uv = nn.Linear(kv_lora_rank, num_heads * d_head_nope, bias=False)

        # 3. 输出投影矩阵
        self.w_o = nn.Linear(num_heads * d_head_nope, d_model, bias=False)

    def forward_decode_absorbed(
        self,
        h_t: torch.Tensor,
        cached_c_kv: torch.Tensor,
        cached_k_rope: torch.Tensor,
        cos_q: torch.Tensor,
        sin_q: torch.Tensor,
        cos_k: torch.Tensor,
        sin_k: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        推理 Decode 阶段 (启用矩阵吸收，直接与压缩潜在缓存计算，不还原完整 KV)
        :param h_t: 当前步输入隐状态 [Batch, 1, d_model]
        :param cached_c_kv: 历史缓存的潜在向量 [Batch, Past_Len, kv_lora_rank]
        :param cached_k_rope: 历史缓存的共享位置键 [Batch, Past_Len, d_head_rope]
        """
        bsz, seq_len, _ = h_t.shape
        assert seq_len == 1, "矩阵吸收模式专为单步自回归 Decode 阶段优化"

        # --- Query 侧投影 ---
        c_q = self.q_norm(self.w_dq(h_t))
        q_nope = self.w_uq(c_q).view(bsz, 1, self.num_heads, self.d_head_nope).transpose(1, 2)
        q_rope = self.w_qr(c_q).view(bsz, 1, self.num_heads, self.d_head_rope).transpose(1, 2)
        q_rope = apply_rope(q_rope, cos_q, sin_q)

        # --- 当前步 KV 压缩与缓存拼接 ---
        kv_compressed = self.w_dkv_kr(h_t)
        c_kv_new, k_rope_new = torch.split(
            kv_compressed, [self.kv_lora_rank, self.d_head_rope], dim=-1
        )
        c_kv_new = self.kv_norm(c_kv_new)
        
        # 对新生成的共享 k_rope 施加旋转位置编码
        k_rope_new = apply_rope(
            k_rope_new.unsqueeze(1), cos_k, sin_k
        ).squeeze(1)  # [Batch, 1, d_head_rope]

        # 更新紧凑型 KV Cache
        full_c_kv = torch.cat([cached_c_kv, c_kv_new], dim=1)         # [Batch, Total_Len, rank]
        full_k_rope = torch.cat([cached_k_rope, k_rope_new], dim=1)   # [Batch, Total_Len, rope_dim]

        # --- 核心优化: 矩阵吸收计算 ---
        # 1. 将 W_uk [Heads * d_nope, rank] 吸收进 q_nope: q_absorbed [Batch, Heads, 1, rank]
        w_uk_reshaped = self.w_uk.weight.view(self.num_heads, self.d_head_nope, self.kv_lora_rank)
        q_absorbed = torch.einsum("bhqd,hdr->bhqr", q_nope, w_uk_reshaped)

        # 2. 直接用吸收后的 Query 与压缩缓存 full_c_kv 点积，加上 RoPE 分支打分
        scores_nope = torch.matmul(q_absorbed, full_c_kv.unsqueeze(1).transpose(-1, -2))
        scores_rope = torch.matmul(q_rope, full_k_rope.unsqueeze(1).transpose(-1, -2))
        attn_weights = F.softmax((scores_nope + scores_rope) * self.scale, dim=-1)

        # 3. 潜在空间加权聚合: latent_out [Batch, Heads, 1, rank]
        latent_out = torch.matmul(attn_weights, full_c_kv.unsqueeze(1))

        # 4. 将 W_uv 吸收进输出投影，恢复多头特征并映射回 d_model
        w_uv_reshaped = self.w_uv.weight.view(self.num_heads, self.d_head_nope, self.kv_lora_rank)
        head_out = torch.einsum("bhqr,hdr->bhqd", latent_out, w_uv_reshaped)
        output = self.w_o(head_out.transpose(1, 2).contiguous().view(bsz, 1, -1))

        return output, full_c_kv, full_k_rope

四、 生产环境部署与工程避坑指南

1. Prefill 阶段与 Decode 阶段的计算路径动态切换

• 常见误区：在所有推理阶段都无脑启用“矩阵吸收”模式。
• 算力开销分析：在矩阵吸收模式下，注意力点积的特征维度从 $d_{\text{nope}}$（128）变为 $d_c$（512），且不再共享单一 KV 头（相当于所有注意力头独立与共享的潜在池计算）。
  • 在 Decode 阶段（单步输入长度为 1），系统属于访存密集型（Memory-bound），省去解压显存读写带来的收益远大于多出的矩阵乘浮点开销，必须开启矩阵吸收；
  • 在 Prefill 阶段（长文本输入，算力密集型 Compute-bound），直接将 $c_t^{KV}$ 显式上投影展开为标准多头 $K, V$ 并调用 FlashAttention-2/3 内核，其总 FLOPs 更低且 Tensor Core 吞吐效率最高。生产推理引擎（如 vLLM、SGLang）均实现了 Prefill 展开与 Decode 吸收的动态双路径切换。

2. 低比特 KV Cache 量化（FP8 / INT8）的数值敏感性

• 踩坑现象：在传统 MHA/GQA 中，直接对 KV Cache 开启 Per-Tensor 或 Per-Token INT8/FP8 量化通常精度无损；但直接对 MLA 缓存的潜在向量 $c_t^{KV}$ 执行粗粒度量化，往往导致长文本困惑度（PPL）显著恶化。
• 根因与对策：$c_t^{KV}$ 经过了高倍率的低秩信息压缩，其每一个通道承载的信息密度极高，且通道间数值方差差异显著。若需对 MLA 的 KV Cache 开启 FP8/INT8 量化，必须采用细粒度的逐通道（Per-Channel）或分组（Group-wise）缩放校准，同时保持共享位置向量 $k_t^R$ 为高精度 BF16/FP16 不量化（因其仅占 64 维，量化收益微乎其微但对位置旋转相位极为敏感）。

3. 分布式推理中的张量并行（TP）切分陷阱

• 架构细节：在标准 GQA 张量并行中，当 TP 并行度大于 KV 头数（例如 $n_{\text{kv}} = 8$，但使用 16 卡 TP）时，需要将 KV 权重在多卡间冗余复制。
• MLA 优势与注意点：在 MLA 架构中，缓存的潜在向量 $c_t^{KV}$ 与共享 $k_t^R$ 在物理上等效于只有 1 个全局注意力头。在多卡张量并行（TP）下，各卡上的 Query 头可以独立切分并行计算，但共享的潜在 KV Cache 会在各张量并行组内完整保留一份（仅 576 维，显存开销极小），无需复杂的跨卡 All-Gather 通信拼接 KV，大幅降低了多卡推理的通信延迟。