超越KV重建：投机解码中MLA草稿模型的功能重构

# 投机解码中 MLA 草稿模型的功能性重建
来源：https://arxiv.org/html/2607.27269
###### 摘要

多头潜在注意力（MLA）是长上下文 LLM 推理中日益重要的架构，因为它将不断增长的键值（KV）缓存替换为紧凑的潜在状态，从而减少解码过程中的内存流量。然而，大多数能力较强的开放检查点都是使用多头注意力或分组查询注意力（MHA/GQA）预训练的，因此在不进行昂贵从头训练的情况下，转换为 MLA 对于实现这些缓存效率优势至关重要。与此同时，投机解码已成为加速自回归生成的一种补充方法。然而，其实践中的加速效果在很大程度上取决于草稿提案与目标验证之间的一致性：随机采样由分布重叠决定，而贪婪解码则需要最高排名 token 的一致性。我们发现，直接的 MHA/GQA 到 MLA 转换会严重损害这种一致性。低秩分解和 RoPE 处理引入的注意力函数误差在独立生成中或许可以容忍，但在投机解码中却会导致 token 接受率急剧下降。因此，我们将 MLA 草稿构建表述为一个功能性重建问题，而非缓存压缩问题。我们的端到端（E2E）重建方法直接优化每个转换后的 MLA 注意力模块，使其在校准隐藏状态上重现其原始 MHA/GQA 对应模块在输出投影之后（post-\(W_O\)）的响应。这是一个与转换器无关的转换后流程，它保留转换后的缓存结构，并且不使用验证器模型的 logits 或监督信号。我们评估了 192 种模型-转换器-后端-方法-任务配置，涵盖四个 Llama/Qwen 草稿-目标对、TransMLA 和 MHA2MLA、HF 和 vLLM 后端，以及四个包含 200 个提示词的任务。使用 0.5 个百分点的报告容差，功能性重建在 64 个匹配任务单元中的 37 个中显著改善了接受率，26 个实际上没有变化，1 个显著下降。我们已在 https://github.com/swyhahaha/FunctionalMLA 发布我们的实现和完整评估工件。

## 引言

大语言模型推理日益受限于内存移动而非算术运算。在自回归解码过程中，每个生成的 token 都会读取不断增长的键值（KV）缓存，因此服务延迟和吞吐量在很大程度上取决于该缓存的表示和访问效率。分组查询注意力（GQA）相对于多头注意力减少了 KV 头的数量，但它仍然存储显式的键和值。DeepSeek 风格模型中使用的多头潜在注意力（MLA）更进一步：它缓存一个低维潜在向量和一个较小的位置分量，然后通过吸收投影计算注意力（DeepSeek-AI 2024a (https://arxiv.org/html/2607.27269#bib.bib5)）。最近的转换方法使得这种设计无需从头预训练 MLA 模型即可实现：TransMLA 转换现有的 GQA/MHA 检查点，而 MHA2MLA 提供了相关的 MHA 到 MLA 迁移路径（Meng et al. 2025 (https://arxiv.org/html/2607.27269#bib.bib6)；Ji 等 2025 (https://arxiv.org/html/2607.27269#bib.bib7)）。

本文研究了一个如果只衡量独立解码就很容易被忽略的失败模式。转换后的 MLA 模型可能是一个高效的生成器，但仍然是一个糟糕的*草稿模型*。在投机解码中，草稿提出多个 token，目标模型并行验证它们（Leviathan et al. 2023 (https://arxiv.org/html/2607.27269#bib.bib8)；Chen et al. 2023 (https://arxiv.org/html/2607.27269#bib.bib9)）。当草稿提案与目标验证一致时，许多 token 会被接受，目标在每次验证步骤中前进多个位置。当转换改变了提案概率或排名时，验证会拒绝更多 token，草稿就变成了开销。

在本文中，我们介绍一种新的功能性重建方法，以提高投机解码过程中草稿的接受率和输出 token 速度。这是一个与转换器无关、仅在训练时进行的转换后阶段，不需要验证器的 logits、验证器的激活、基准标签或生成目标，并且不会改变转换后的缓存结构或推理图。从 TransMLA 或 MHA2MLA 转换的检查点开始，它训练转换器引入的 query 和 KV 投影，使完整的 MLA 注意力模块在校准隐藏状态上重现冻结的原始 MHA/GQA 模块在 \(W_O\) 之后的响应。在表1 (https://arxiv.org/html/2607.27269#Sx6.T1) 中的 64 个匹配任务单元中，我们的方法产生了 37 个显著的接受率改进，26 个实际上没有变化的结果，以及 1 个显著下降；37 个改进单元中有 12 个还显示出吞吐量的显著提升。这些结果涵盖了 TransMLA 和 MHA2MLA、HF 和 vLLM，以及 Llama 和 Qwen 模型家族。

我们的贡献有三方面：

(1) 转换和草稿是不同的目标。结构转换可以成功，而草稿接受所需的提案一致性却可能丢失。

(2) 草稿质量可以完全在训练时优化。我们的功能性目标使用校准状态和冻结的原始注意力模块，无需验证器监督或推理时更改。

(3) 训练时的功能性重建与转换器无关。相同的目标在 64 个匹配任务单元中的 63 个中有所改进或保持在报告容差之内，涵盖 TransMLA/MHA2MLA、Llama/Qwen 和 HF/vLLM。

## 相关工作

KV 缓存高效注意力。自回归 Transformer 在解码过程中会为每个先前的 token 和层保留键和值状态（Vaswani et al. 2017 (https://arxiv.org/html/2607.27269#bib.bib1)）。MQA 和 GQA 通过在查询头之间共享 KV 头来减少这种状态（Shazeer 2019 (https://arxiv.org/html/2607.27269#bib.bib3)；Ainslie et al. 2023 (https://arxiv.org/html/2607.27269#bib.bib2)）。其他方法保留原始注意力表示但减少其运行时占用：PagedAttention 对显式 KV 块进行虚拟化（Kwon et al. 2023 (https://arxiv.org/html/2607.27269#bib.bib10)），H2O 和 StreamingLLM 为长序列保留选定的缓存条目（Zhang et al. 2023b (https://arxiv.org/html/2607.27269#bib.bib19)；Xiao et al. 2023 (https://arxiv.org/html/2607.27269#bib.bib20)），KIVI 对缓存进行量化（Liu et al. 2024 (https://arxiv.org/html/2607.27269#bib.bib21)）。FlashAttention 和 FlashAttention-2 则在不改变缓存表示的情况下优化精确注意力内核（Dao et al. 2022 (https://arxiv.org/html/2607.27269#bib.bib11)；Dao 2023 (https://arxiv.org/html/2607.27269#bib.bib12)）。MLA 在架构上有所不同：它用学习到的潜在缓存替换显式 KV 缓存，因此可以减少表示本身，而不仅仅是其存储或执行成本（DeepSeek-AI 2024a (https://arxiv.org/html/2607.27269#bib.bib5), b (https://arxiv.org/html/2607.27269#bib.bib22)）。

MLA 与转换。MLA 在 DeepSeek-V2 中引入，并在 DeepSeek-V3 中保留为面向推理的注意力架构（DeepSeek-AI 2024a (https://arxiv.org/html/2607.27269#bib.bib5), b (https://arxiv.org/html/2607.27269#bib.bib22)）。从预训练的 MHA/GQA 迁移到 MLA 并非易事，因为低秩 KV 压缩和旋转位置处理都可能扰动原始注意力图（Su et al. 2024 (https://arxiv.org/html/2607.27269#bib.bib4)）。TransMLA 使用 RoPE 分解与低秩 QKV 转换（Meng et al. 2025 (https://arxiv.org/html/2607.27269#bib.bib6)），而 MHA2MLA 则结合了部分 RoPE 移除与联合低秩近似（Ji 等 2025 (https://arxiv.org/html/2607.27269#bib.bib7)）。这些方法通过转换器局部重建目标解决了初始化问题。我们的 E2E 流程是互补的：它从任一转换后的检查点开始，并使用原始 MHA/GQA 注意力在校准状态上的输出作为对投机草稿至关重要的功能性目标。

投机解码。投机解码使用廉价的草稿提出多个 token，并通过精确的目标端修正来保持目标分布（Leviathan et al. 2023 (https://arxiv.org/html/2607.27269#bib.bib8)；Chen et al. 2023 (https://arxiv.org/html/2607.27269#bib.bib9)）。基于树的方法通过固定或硬件感知的草稿树来扩大并行候选集（Miao et al. 2023 (https://arxiv.org/html/2607.27269#bib.bib13)；Chen et al. 2024 (https://arxiv.org/html/2607.27269#bib.bib25)）；Medusa 使用辅助解码头（Cai et al. 2024 (https://arxiv.org/html/2607.27269#bib.bib14)）；EAGLE/EAGLE-2 在特征空间中使用静态或上下文自适应树进行草稿（Li et al. 2024b (https://arxiv.org/html/2607.27269#bib.bib15), a (https://arxiv.org/html/2607.27269#bib.bib26)）。自投机方法通过层跳过或提前退出重用目标模型（Zhang et al. 2023a (https://arxiv.org/html/2607.27269#bib.bib23)；Elhoushi et al. 2024 (https://arxiv.org/html/2607.27269#bib.bib24)），而 Lookahead 解码使用 Jacobi 式并行更新（Fu et al. 2024 (https://arxiv.org/html/2607.27269#bib.bib16)）。这些工作改变了提案生成或验证调度。我们则改进 MLA 转换草稿的提案保真度，这与上述设计互补，并直接提高了它们的加速所依赖的接受率。

## 预备：为什么转换后的 MLA 不能保留草稿

### MHA/GQA 注意力

设 \(h_t \in \mathbb{R}^D\) 为位置 \(t\) 处的隐藏状态，其中 \(D\) 表示模型隐藏维度，有 \(H\) 个查询头和 \(G\) 个 KV 组。我们下面使用 GQA 表示；标准 MHA 是 \(G=H\) 且每个查询头一个 KV 组的特例。对于组 \(g(i)\) 中的查询头 \(i\)，GQA 计算

\[
\begin{aligned}
q_{t,i} &= W^Q_i h_t, \\
k_{j,g(i)} &= W^K_{g(i)} h_j, \\
v_{j,g(i)} &= W^V_{g(i)} h_j.
\end{aligned}
\tag{1}
\]
设 \(R_t\) 为位置 \(t\) 处的 RoPE 旋转，并定义相对旋转 \(R_{tj}=R_t^\top R_j\)（Su et al. 2024 (https://arxiv.org/html/2607.27269#bib.bib4)）。设键维度为 \(d_k\)，缩放因果注意力为

\[
s^{\mathrm{GQA}}_{tij} = \frac{q_{t,i}^\top R_{tj} k_{j,g(i)}}{\sqrt{d_k}},
\tag{2}
\]
\[
a^{\mathrm{GQA}}_{tij} = \frac{\exp(s^{\mathrm{GQA}}_{tij})}{\sum_{u \le t} \exp(s^{\mathrm{GQA}}_{tiu})},
\tag{3}
\]
\[
o^{\mathrm{GQA}}_{t,i} = \sum_{j \le t} a^{\mathrm{GQA}}_{tij} v_{j,g(i)}.
\tag{4}
\]

### 潜在 KV 瓶颈

一个 MLA 层存储潜在缓存 \(c_j = W^{DKV} h_j \in \mathbb{R}^r\) 和一个较小的位置键 \(k^{PE}_j\)。所有 KV 组的非位置键和值被堆叠起来并从潜在状态产生：

\[
\begin{bmatrix} \hat{k}^{nope}_j \\ \hat{v}_j \end{bmatrix} = W^{UKV} W^{DKV} h_j.
\tag{5}
\]
这里 \(W^{DKV} \in \mathbb{R}^{r \times D}\) 且 \(W^{UKV} \in \mathbb{R}^{m \times r}\)，其中 \(m=G(d_k^{nope}+d_v)\)。对所有隐藏状态精确重现这个堆叠线性映射需要

\[
\begin{bmatrix} W^{K,nope} \\ W^V \end{bmatrix} = W^{UKV} W^{DKV}.
\tag{6}
\]
右手边的秩最多为 \(r\)。因此，精确等价需要

\[
\mathrm{rank}\!\left(\begin{bmatrix} W^{K,nope} \\ W^V \end{bmatrix}\right) \le r.
\tag{7}
\]
无论何时这个秩超过 \(r\)，所选的压缩因子分解都无法精确重现堆叠的 K/V 映射，并且必须留下残差。这个条件语句涉及在所选秩下的直接映射重建：它既不排除每个功能等价的注意力参数化，也不与足够大秩下的精确 MLA 重写相矛盾。它说明了为什么一个激进压缩且缓存有效的转换通常是原始 K/V 投影的近似。

图 1：与转换器无关的 E2E 功能性重建。TransMLA 或 MHA2MLA 转换首先提供一个具备缓存能力的 MLA 初始化。对于每个校准隐藏状态 \(h^\ell\)，冻结的原始 GQA 块和可训练的转换 MLA 块产生 \(W_O\) 之后的输出。E2E 最小化它们的掩蔽均方误差，停止通过原始路径的梯度，并且只更新转换器引入的 query 和 KV 投影。下方的示意图展示了经验验证的路径：减少残差可以保留提案排名并提高接受率。

### RoPE 处理与接受率

秩瓶颈并不是唯一的不匹配。转换流程还必须选择如何处理 RoPE。为了明确该误差，将原始键拆分为 \(k_j = k_j^{keep} + k_j^{drop}\)，其中 \(k_j^{keep}\) 保留在旋转路径上，而 \(k_j^{drop}\) 被移到非位置路径。设转换后的查询为 \(\hat{q}_t = q_t + \delta q_t\)，重建的保留键为 \(\hat{k}_j^{keep} = k_j^{keep} + \delta k_j\)。在这种代表性路由下，原始 logit 使用 \(q_t^\top R_{tj}(k_j^{keep}+k_j^{drop})/\sqrt{d_k}\)，而转换后的 logit 使用 \((q_t+\delta q_t)^\top [R_{tj}(k_j^{keep}+\delta k_j)+k_j^{drop}]/\sqrt{d_k}\)。它们的精确差值为

\[
\hat{s}_{tj} - s_{tj} = \frac{1}{\sqrt{d_k}}\Big[
q_t^\top R_{tj} \delta k_j + q_t^\top (I - R_{tj}) k_j^{drop}
\tag{8}
\]
\[
+ \delta q_t^\top R_{tj}(k_j^{keep}+\delta k_j) + \delta q_t^\top k_j^{drop}\Big].
\]
这个分解揭示了保留键重建误差、位置路径误差和一阶查询扰动；第三项还包含它们的交互。Softmax 将 logit 残差映射为注意力权重误差，输出残差还继承了值重建误差。

投机解码使这个残差在操作上可见。在固定上下文 \(x\) 下，记 \(p_D^x(y)=p_D(y \mid x)\)，类似地记 \(p_T^x\)。随机拒绝采样验证的单 token 期望接受率为（Leviathan et al. 2023 (https://arxiv.org/html/2607.27269#bib.bib8)；Chen et al. 2023 (https://arxiv.org/html/2607.27269#bib.bib9)）

\[
\alpha_{\mathrm{sample}}(x) = \sum_y \min\!\left(p_D^x(y), p_T^x(y)\right)
\tag{9}
\]
\[
= 1 - \mathrm{TV}\!\left(p_D^x, p_T^x\right).
\]
我们的实验改为使用贪婪解码，其单 token 接受事件为

\[
\alpha_{\mathrm{greedy}}(x) = \mathbf{1}[\arg\max_y p_D^x(y) = \arg\max_y p_T^x(y)].
\]
\al