新上线今天0 投票
超越KV重建：推测解码中MLA草稿模型的功能重构方法
背景：MLA与推测解码的碰撞
多潜变量注意力（MLA）正成为长上下文大语言模型推理的关键技术。它通过紧凑的潜状态替代不断增长的键值（KV）缓存，显著降低解码时的内存带宽压力。然而，目前最强大的开源模型大多采用多头注意力（MHA）或分组查询注意力（GQA），若想获得MLA的缓存效率，通常需要将MHA/GQA转换为MLA，而无需从头训练。
推测解码（Speculative Decoding）是另一种加速策略：先由草稿模型快速生成候选token，再由目标模型验证。其加速效果高度依赖于草稿与目标输出的一致性。
核心问题：直接转换为何失效？
论文作者发现，直接将MHA/GQA转换为MLA会大幅降低草稿token的接受率。原因在于：低秩分解和旋转位置编码（RoPE）处理会引入注意力函数误差。这种误差在独立生成时或许可以容忍，但在推测解码中，却会显著削弱草稿与目标模型之间的对齐，导致验证频繁拒绝，加速效果大打折扣。
新方法：功能重构而非缓存压缩
为此，研究者提出将MLA草稿构建视为功能重构问题，而非简单的缓存压缩。他们设计了一种端到端（E2E）方法：在转换后，针对每个MLA注意力模块，优化其输出，使其在校准隐藏状态上，能够再现原始MHA/GQA对应模块经过输出投影后的响应。
这一方法具有以下特点：
- 转换器无关：适用于任何MHA/GQA到MLA的转换方法（如TransMLA、MHA2MLA）；
- 无需验证器监督：不需要验证器的logits或额外监督信号；
- 保持推理图不变：转换后的缓存和推理结构得以保留。
实验效果：多数场景显著提升
研究团队在192种模型-转换器-后端-任务组合上进行了评估，涵盖Llama/Qwen草稿-目标对、HF和vLLM后端，以及4个200条提示的任务。在0.5个百分点的容差下，功能重构在64个匹配任务单元中的37个显著提升了接受率，26个基本不变，仅1个轻微下降。
小结
这项研究揭示了推测解码中MLA转换的隐性瓶颈，并提供了实用的后处理优化方案。对于希望在长上下文场景中兼顾内存效率与解码速度的开发者而言，功能重构不失为一种值得尝试的轻量级增强手段。