KV 재구축을 넘어: 추론 기반 디코딩 모델에서 기능적 재구축을 통한 MLA 최적화

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

멀티 헤드 잠재 어텐션(MLA)은 긴 문맥 LLM 추론에 있어 점점 더 중요해지고 있습니다. MLA는 증가하는 키-값(KV) 캐시를 대체하고 디코딩 메모리 트래픽을 줄이는 작고 효율적인 잠재 상태를 사용하기 때문입니다. 하지만 대부분의 성능이 뛰어난 공개 모델은 멀티 헤드 또는 그룹 쿼리 어텐션(MHA/GQA)을 사용하므로, 처음부터 다시 학습하지 않고 MLA의 캐시 효율성을 얻기 위해서는 변환이 필요합니다. 추론 기반 디코딩은 상호 보완적인 가속화를 제공하지만, 그 속도 향상은 초안 제안과 대상 검증 간의 일치성에 따라 달라집니다. 우리는 MHA/GQA에서 MLA로 직접 변환하면 이러한 일치성이 크게 감소할 수 있다는 것을 발견했습니다. 저차원 분해 및 RoPE 처리는 어텐션 기능 오류를 발생시키는데, 이는 독립적인 생성에서는 허용될 수 있지만 초안 토큰 채택률을 상당히 낮출 수 있습니다. 따라서 우리는 MLA 초안 구성을 캐시 압축이 아닌 기능적 재구축으로 정의합니다. 우리의 엔드-투-엔드(E2E) 방법은 각 변환된 MLA 어텐션 모듈을 최적화하여, 교정된 히든 상태에서 원래의 MHA/GQA 모델이 출력 투영 후 보이는 응답을 재현하도록 합니다. 이 변환기 독립적인 사후 변환 절차는 변환된 캐시와 추론 그래프를 보존하며, 검증자 로짓이나 검증자 감독 신호가 필요하지 않습니다. 우리는 네 가지 Llama/Qwen 초안-대상 쌍, TransMLA 및 MHA2MLA, HF 및 vLLM, 그리고 네 가지 200개 프롬프트 작업에 이르는 총 192개의 모델-변환기-백엔드-방법-작업 구성을 평가했습니다. 0.5%의 보고 오차 허용 범위 내에서, 기능적 재구축은 64개의 일치된 작업 영역 중 37개에서 채택률을 크게 향상시키고, 26개에서는 거의 변화가 없으며, 1개에서는 감소시켰습니다. 코드 및 평가 결과는 https://github.com/swyhahaha/FunctionalMLA 에서 확인할 수 있습니다.

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.