xKV: Межслойное SVD для сжатия KV-кэша

xKV: Cross-Layer SVD for KV-Cache Compression

March 24, 2025

Авторы: Chi-Chih Chang, Chien-Yu Lin, Yash Akhauri, Wei-Cheng Lin, Kai-Chiang Wu, Luis Ceze, Mohamed S. Abdelfattah

cs.AI

Аннотация

Крупные языковые модели (LLM) с длинными контекстными окнами открывают возможности для мощных приложений, но требуют значительного объема памяти для хранения состояний ключей и значений (KV-Cache). Недавние исследования пытались объединить KV-cache из нескольких слоев в общие представления, однако эти подходы либо требуют дорогостоящего предварительного обучения, либо основываются на предположениях о высокой косинусной схожести между токенами в разных слоях, что на практике обычно не выполняется. Мы обнаружили, что доминирующие сингулярные векторы удивительно хорошо согласованы между несколькими слоями KV-Cache. Используя это наблюдение, мы предлагаем xKV — простой метод пост-обучения, который применяет сингулярное разложение (SVD) к KV-cache сгруппированных слоев. xKV объединяет KV-cache нескольких слоев в общее низкоранговое подпространство, значительно уменьшая размеры KV-cache. В ходе обширных оценок на бенчмарке RULER для длинных контекстов с широко используемыми LLM (например, Llama-3.1 и Qwen2.5) xKV достигает сжатия до 6.8 раз выше, чем современные межслойные методы, при этом повышая точность на 2.7%. Более того, xKV совместим с развивающейся технологией Multi-Head Latent Attention (MLA) (например, DeepSeek-Coder-V2), обеспечивая заметное 3-кратное сжатие на задачах кодирования без потери производительности. Эти результаты подчеркивают мощные возможности и универсальность xKV в решении проблем с памятью при выводе LLM с длинными контекстами. Наш код доступен по адресу: 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.