xKV: KV 캐시 압축을 위한 크로스 레이어 SVD

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

초록

긴 컨텍스트 윈도우를 갖춘 대형 언어 모델(LLMs)은 강력한 애플리케이션을 가능하게 하지만, 키와 값 상태(KV-Cache)를 저장하기 위한 높은 메모리 소비라는 비용이 따른다. 최근 연구들은 여러 레이어의 KV-Cache를 공유된 표현으로 병합하려는 시도를 했으나, 이러한 접근법들은 비용이 많이 드는 사전 학습을 요구하거나, 실제로는 일반적으로 성립하지 않는 레이어 간 높은 토큰별 코사인 유사성을 가정에 의존한다. 우리는 KV-Cache의 여러 레이어에서 지배적인 특이 벡터(singular vectors)가 놀랍도록 잘 정렬되어 있음을 발견했다. 이러한 통찰을 활용하여, 우리는 그룹화된 레이어의 KV-Cache에 특이값 분해(SVD)를 적용하는 간단한 사후 학습 방법인 xKV를 제안한다. xKV는 여러 레이어의 KV-Cache를 공유된 저차원 부분 공간으로 통합하여 KV-Cache 크기를 크게 줄인다. 널리 사용되는 LLMs(예: Llama-3.1 및 Qwen2.5)을 사용한 RULER 장문 컨텍스트 벤치마크에서의 광범위한 평가를 통해, 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.