xKV : Décomposition SVD intercouches pour la compression du cache KV

xKV: Cross-Layer SVD for KV-Cache Compression

March 24, 2025

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

cs.AI

Résumé

Les modèles de langage à grande échelle (LLMs) avec des fenêtres de contexte longues permettent des applications puissantes, mais au prix d'une consommation de mémoire élevée pour stocker les états de Clé et de Valeur (KV-Cache). Des études récentes ont tenté de fusionner le KV-cache de plusieurs couches en représentations partagées, mais ces approches nécessitent soit un pré-entraînement coûteux, soit reposent sur des hypothèses de similarité cosinus élevée par token entre les couches, ce qui ne se vérifie généralement pas en pratique. Nous constatons que les vecteurs singuliers dominants sont remarquablement alignés à travers plusieurs couches du KV-Cache. Exploitant cette observation, nous proposons xKV, une méthode simple post-entraînement qui applique la Décomposition en Valeurs Singulières (SVD) sur le KV-Cache de couches groupées. xKV consolide le KV-Cache de plusieurs couches en un sous-espace partagé de faible rang, réduisant significativement la taille du KV-Cache. À travers des évaluations approfondies sur le benchmark de contexte long RULER avec des LLMs largement utilisés (par exemple, Llama-3.1 et Qwen2.5), xKV atteint des taux de compression jusqu'à 6,8 fois supérieurs à la technique inter-couches de pointe tout en améliorant la précision de 2,7%. De plus, xKV est compatible avec l'attention latente multi-têtes émergente (MLA) (par exemple, DeepSeek-Coder-V2), offrant un taux de compression notable de 3x sur les tâches de codage sans dégradation de performance. Ces résultats mettent en évidence la capacité et la polyvalence de xKV à résoudre les goulots d'étranglement de mémoire pour l'inférence des LLMs à contexte long. Notre code est disponible publiquement à l'adresse : 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.