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

论文针对多头潜在注意力（MLA）草稿模型在投机解码中的构建问题，提出功能性重构（Functional Reconstruction）方法。该方法优化转换后的MLA注意力模块，使其复现原MHA/GQA模型经输出投影后的响应，无需验证器监督。作者在192种配置（含Llama/Qwen草稿-目标对、TransMLA和MHA2MLA转换器、HF和vLLM后端）上评估，发现相比直接转换，该方法在37个匹配任务单元中显著提升草稿接受率，26个基本不变，其余下降。

此前MHA/GQA转MLA仅被视为缓存压缩，未考虑重构后注意力函数误差对投机解码草稿接受率的负面影响。该方法将构建目标从缓存压缩转为功能复现，显著提升草稿接受率，为长上下文推理的加速与内存效率兼得提供新路径。

MLA通过低秩因子分解和RoPE处理压缩KV缓存，但直接转换会引入注意力函数误差，降低草稿与目标的一致率。功能性重构在转换后，用校准隐藏状态优化每个MLA注意力模块，使其输出接近原MHA/GQA模块经输出投影后的结果，从而在不改变推理图的前提下提升草稿接受率。

对使用MHA/GQA开源模型但希望获得MLA缓存效率的团队，该方法可避免从头训练而直接转换并保持推理加速效果。在vLLM等推理后端中集成该后转换流程，可能提升投机解码的吞吐、降低长上下文推理的内存带宽成本，但需增加校准开销。

• 在vLLM中为MLA转换模型提供可选的post-conversion微调模块，面向长上下文推理场景优化。
• 开发自动化工具，根据目标模型自动选择转换器并应用功能重构，降低手动调优成本。
• 将此方法集成到模型部署流水线，作为MHA/GQA转MLA的标准后处理步骤。

• 功能重构所需的校准隐藏状态数量和计算开销是多少？
• 该方法在更大规模模型或更多任务上的效果是否稳定？
• 重构后模型的长期稳定性和泛化能力如何？