Beyond KV Reconstruction: Functional Reconstruction for MLA Draft Models in Speculative Decoding

2026-07-29

AI Summary

Directly converting multi-head attention (MHA) or grouped-query attention (GQA) to multi-head latent attention (MLA) for speculative decoding introduces attention function discrepancies that substantially reduce draft token acceptance rates. This work proposes an architecture-agnostic, post-conversion functional reconstruction method that treats MLA draft model construction as a functional approximation of the original MHA/GQA module’s output influence, rather than merely cache compression. By performing end-to-end optimization on calibrated hidden states, the method reproduces the original model’s behavior without requiring verifier logits or additional supervision. Evaluated across 192 configurations, the approach significantly improves acceptance rates on 37 tasks, matches baseline performance on 26, and incurs only a minor degradation on a single task, demonstrating its effectiveness and strong generalization capability.

Abstract

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.

Problem

• speculative decoding
• multi-head latent attention
• attention-function error
• draft-token acceptance
• MHA/GQA-to-MLA conversion

Innovation

• Functional Reconstruction
• Speculative Decoding
• Multi-head Latent Attention
• MHA-to-MLA Conversion
• Draft Model Optimization