CommVQ: KV 캐시 압축을 위한 교환 가능 벡터 양자화

CommVQ: Commutative Vector Quantization for KV Cache Compression

June 23, 2025

저자: Junyan Li, Yang Zhang, Muhammad Yusuf Hassan, Talha Chafekar, Tianle Cai, Zhile Ren, Pengsheng Guo, Foroozan Karimzadeh, Colorado Reed, Chong Wang, Chuang Gan

cs.AI

초록

대규모 언어 모델(LLM)은 긴 문맥 길이가 필요한 애플리케이션에서 점점 더 많이 사용되고 있지만, 문맥이 길어질수록 키-값(KV) 캐시가 GPU에서 메모리 병목 현상을 일으키는 경우가 많습니다. 이를 해결하기 위해, 우리는 긴 문맥 LLM 추론을 위한 메모리 사용량을 크게 줄이는 교환 가능한 벡터 양자화(CommVQ)를 제안합니다. 먼저, 가벼운 인코더와 코드북을 사용한 가법 양자화를 도입하여 KV 캐시를 압축하며, 이를 간단한 행렬 곱셈으로 디코딩할 수 있도록 합니다. 디코딩 과정에서의 계산 비용을 더욱 줄이기 위해, 우리는 코드북이 회전 위치 임베딩(RoPE)과 교환 가능하도록 설계하고, 이를 기대값 최대화(EM) 알고리즘을 사용하여 학습시킵니다. 이를 통해 디코딩을 자기 주의 메커니즘에 효율적으로 통합할 수 있습니다. 우리의 접근 방식은 가법 양자화를 통해 높은 정확도를 달성하고, RoPE-교환 가능한 코드북을 통해 낮은 오버헤드를 유지합니다. 긴 문맥 벤치마크와 GSM8K에 대한 실험 결과, 우리의 방법은 2비트 양자화로 FP16 KV 캐시 크기를 87.5% 줄이면서도 최신 KV 캐시 양자화 방법을 능가하는 성능을 보여줍니다. 특히, 최소한의 정확도 손실로 1비트 KV 캐시 양자화를 가능하게 하여, LLaMA-3.1 8B 모델이 단일 RTX 4090 GPU에서 128K 문맥 길이로 실행될 수 있도록 합니다. 소스 코드는 https://github.com/UMass-Embodied-AGI/CommVQ에서 확인할 수 있습니다.

English

Large Language Models (LLMs) are increasingly used in applications requiring long context lengths, but the key-value (KV) cache often becomes a memory bottleneck on GPUs as context grows. To address this, we propose Commutative Vector Quantization (CommVQ) to significantly reduce memory usage for long-context LLM inference. We first introduce additive quantization with a lightweight encoder and codebook to compress the KV cache, which can be decoded via simple matrix multiplication. To further reduce computational costs during decoding, we design the codebook to be commutative with Rotary Position Embedding (RoPE) and train it using an Expectation-Maximization (EM) algorithm. This enables efficient integration of decoding into the self-attention mechanism. Our approach achieves high accuracy with additive quantization and low overhead via the RoPE-commutative codebook. Experiments on long-context benchmarks and GSM8K show that our method reduces FP16 KV cache size by 87.5% with 2-bit quantization, while outperforming state-of-the-art KV cache quantization methods. Notably, it enables 1-bit KV cache quantization with minimal accuracy loss, allowing a LLaMA-3.1 8B model to run with a 128K context length on a single RTX 4090 GPU. The source code is available at: https://github.com/UMass-Embodied-AGI/CommVQ.