CommVQ: Kommutative Vektorquantisierung zur KV-Cache-Kompression

CommVQ: Commutative Vector Quantization for KV Cache Compression

June 23, 2025

Autoren: 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

Zusammenfassung

Große Sprachmodelle (LLMs) werden zunehmend in Anwendungen eingesetzt, die lange Kontextlängen erfordern, doch der Key-Value (KV)-Cache wird oft zu einem Speicher-Engpass auf GPUs, wenn der Kontext wächst. Um dies zu lösen, schlagen wir Commutative Vector Quantization (CommVQ) vor, um den Speicherverbrauch für die Inferenz von LLMs mit langem Kontext erheblich zu reduzieren. Zunächst führen wir additive Quantisierung mit einem leichtgewichtigen Encoder und Codebook ein, um den KV-Cache zu komprimieren, der durch einfache Matrixmultiplikation dekodiert werden kann. Um die Rechenkosten während der Dekodierung weiter zu reduzieren, entwerfen wir das Codebook so, dass es mit Rotary Position Embedding (RoPE) kommutativ ist, und trainieren es mit einem Expectation-Maximization (EM)-Algorithmus. Dies ermöglicht eine effiziente Integration der Dekodierung in den Self-Attention-Mechanismus. Unser Ansatz erreicht hohe Genauigkeit durch additive Quantisierung und geringen Overhead durch das RoPE-kommutative Codebook. Experimente auf Benchmarks für lange Kontexte und GSM8K zeigen, dass unsere Methode die Größe des FP16 KV-Caches um 87,5 % bei 2-Bit-Quantisierung reduziert und dabei state-of-the-art KV-Cache-Quantisierungsmethoden übertrifft. Bemerkenswerterweise ermöglicht sie eine 1-Bit-KV-Cache-Quantisierung mit minimalem Genauigkeitsverlust, sodass ein LLaMA-3.1 8B-Modell mit einer Kontextlänge von 128K auf einer einzelnen RTX 4090 GPU ausgeführt werden kann. Der Quellcode ist verfügbar unter: 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.