超越KV重构:推测解码中MLA草案模型的功能重构
Beyond KV Reconstruction: Functional Reconstruction for MLA Draft Models in Speculative Decoding
摘要
多头潜在注意力 (MLA) 对于长上下文 LLM 推理越来越重要,因为紧凑的潜在状态取代了不断增长的键值 (KV) 缓存并减少了解码内存流量。然而,大多数有能力的开放检查点都使用多头或分组查询注意力(MHA/GQA),因此需要进行转换才能获得 MLA 的缓存效率,而无需从头开始重新训练。 推测解码 提供补充加速,但其加速取决于提案草案和目标验证之间的协议。我们发现直接 MHA/GQA 到 MLA 的转换可以大大降低这种一致性:低秩分解和 RoPE 处理引入了注意力函数错误,这些错误对于独立生成来说可能是可以容忍的,但会大大降低草稿词元的接受度。因此,我们将 MLA 草稿构造制定为功能重建而不是缓存压缩。我们的端到端 (E2E) 方法优化每个转换后的 MLA 注意力模块,以重现其原始 MHA/GQA 对应物在校准隐藏状态上的输出后投影响应。这种与转换器无关的后转换过程保留了转换后的缓存和推理图,并且既不需要验证者 logits 也不需要验证者监督。我们评估了 192 个模型-转换器-后端-方法-任务配置,涵盖四个 Llama/Qwen 草案目标对、TransMLA 和 MHA2MLA、HF 和 vLLM,以及四个 200 提示任务。凭借 0.5 个百分点的报告容差,功能重建显着提高了 64 个匹配任务单元中 37 个任务单元的接受度,使 26 个任务单元几乎保持不变,并大幅减少了 1 个任务单元。代码和评估工件可在 https://github.com/swyhahaha/FunctionalMLA 获取。