将MLA转换为功能重建,提升推测解码的采纳率。
Beyond KV Reconstruction: Functional Reconstruction for MLA Draft Models in Speculative Decoding

- 通过功能重建而非缓存压缩转换MHA/GQA到MLA
- 在192种配置中37组任务采纳率显著提升
- 无需验证器监督,兼容主流推理框架
多头潜在注意力(MLA)因紧凑的潜在状态替代不断增长的键值(KV)缓存,成为长上下文大模型推理的关键技术,可降低解码内存流量。然而,多数高性能开源模型采用多头或分组查询注意力(MHA/GQA),需转换才能获得MLA的缓存效率,且无需从头训练。推测解码可提供额外加速,但其提速依赖草稿生成与目标验证的一致性。我们发现,直接将MHA/GQA转为MLA会显著降低一致性:低秩分解和旋转位置编码处理引入了注意力函数误差,虽对独立生成可容忍,却大幅降低草稿令牌采纳率。因此,我们将草稿构建重新定义为功能重建而非缓存压缩。提出的端到端(E2E)方法优化每个转换后的MLA注意力模块,在校准隐藏状态上复现原始MHA/GQA模块的输出投影后响应。该转换器无关的后转换流程保留转换后的缓存与推理图,既不依赖验证器的输出也不需要验证器监督。我们在4个Llama/Qwen草稿-目标对、TransMLA和MHA2MLA转换器、HF与vLLM后端,以及4项共200个提示的任务中评估了192种配置。在0.5个百分点的报告容差下,功能重建在64个匹配任务单元中使采纳率显著提升37组,26组基本不变,仅1组显著下降。代码与评估数据可在https://github.com/swyhahaha/FunctionalMLA获取。
原文摘要 · Abstract (English)
Multi-head latent attention (MLA) is increasingly important for long-context LLM inference because compact latent states replace the growing key-value (KV) cache and reduce decoding memory traffic. Yet most capable open checkpoints use multi-head or grouped-query attention (MHA/GQA), so conversion is needed to obtain MLA's cache efficiency without retraining from scratch. Speculative decoding offers complementary acceleration, but its speedup depends on agreement between draft proposals and target verification. We find that direct MHA/GQA-to-MLA conversion can sharply reduce this agreement: low-rank factorization and RoPE handling introduce attention-function errors that may be tolerable for standalone generation but substantially lower draft-token acceptance. We therefore formulate MLA draft construction as functional reconstruction rather than cache compression. Our end-to-end (E2E) method optimizes each converted MLA attention module to reproduce the post-output-projection response of its original MHA/GQA counterpart on calibration hidden states. This converter-agnostic post-conversion procedure preserves the converted cache and inference graph and requires neither verifier logits nor verifier supervision. We evaluate 192 model-converter-backend-method-task configurations spanning four Llama/Qwen draft-target pairs, TransMLA and MHA2MLA, HF and vLLM, and four 200-prompt tasks. With a 0.5-percentage-point reporting tolerance, Functional Reconstruction materially improves acceptance in 37 of 64 matched task cells, leaves 26 practically unchanged, and materially decreases one. Code and evaluation artifacts are available at https://github.com/swyhahaha/FunctionalMLA.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。