xKV: SVD Transcamadas para Compressão de Cache KV

xKV: Cross-Layer SVD for KV-Cache Compression

March 24, 2025

Autores: Chi-Chih Chang, Chien-Yu Lin, Yash Akhauri, Wei-Cheng Lin, Kai-Chiang Wu, Luis Ceze, Mohamed S. Abdelfattah

cs.AI

Resumo

Modelos de Linguagem de Grande Escala (LLMs) com janelas de contexto longas permitem aplicações poderosas, mas vêm com o custo de alto consumo de memória para armazenar os estados de Chave e Valor (KV-Cache). Estudos recentes tentaram mesclar o KV-cache de múltiplas camadas em representações compartilhadas, mas essas abordagens ou exigem um pré-treinamento dispendioso ou dependem de suposições de alta similaridade de cosseno por token entre as camadas, o que geralmente não se sustenta na prática. Descobrimos que os vetores singulares dominantes estão notavelmente bem alinhados em múltiplas camadas do KV-Cache. Aproveitando essa percepção, propomos o xKV, um método simples de pós-treinamento que aplica a Decomposição em Valores Singulares (SVD) no KV-Cache de camadas agrupadas. O xKV consolida o KV-Cache de múltiplas camadas em um subespaço compartilhado de baixa dimensão, reduzindo significativamente o tamanho do KV-Cache. Por meio de avaliações extensas no benchmark de contexto longo RULER com LLMs amplamente utilizados (por exemplo, Llama-3.1 e Qwen2.5), o xKV alcança taxas de compressão até 6,8x maiores do que a técnica inter-camadas mais avançada, enquanto melhora a precisão em 2,7%. Além disso, o xKV é compatível com a emergente Atenção Latente Multi-Cabeça (MLA) (por exemplo, DeepSeek-Coder-V2), resultando em uma notável taxa de compressão de 3x em tarefas de codificação sem degradação de desempenho. Esses resultados destacam a forte capacidade e versatilidade do xKV em abordar gargalos de memória para inferência de LLMs de contexto longo. Nosso código está publicamente disponível em: https://github.com/abdelfattah-lab/xKV.

English

Large Language Models (LLMs) with long context windows enable powerful applications but come at the cost of high memory consumption to store the Key and Value states (KV-Cache). Recent studies attempted to merge KV-cache from multiple layers into shared representations, yet these approaches either require expensive pretraining or rely on assumptions of high per-token cosine similarity across layers which generally does not hold in practice. We find that the dominant singular vectors are remarkably well-aligned across multiple layers of the KV-Cache. Exploiting this insight, we propose xKV, a simple post-training method that applies Singular Value Decomposition (SVD) on the KV-Cache of grouped layers. xKV consolidates the KV-Cache of multiple layers into a shared low-rank subspace, significantly reducing KV-Cache sizes. Through extensive evaluations on the RULER long-context benchmark with widely-used LLMs (e.g., Llama-3.1 and Qwen2.5), xKV achieves up to 6.8x higher compression rates than state-of-the-art inter-layer technique while improving accuracy by 2.7%. Moreover, xKV is compatible with the emerging Multi-Head Latent Attention (MLA) (e.g., DeepSeek-Coder-V2), yielding a notable 3x compression rates on coding tasks without performance degradation. These results highlight xKV's strong capability and versatility in addressing memory bottlenecks for long-context LLM inference. Our code is publicly available at: https://github.com/abdelfattah-lab/xKV.