CommVQ: Quantizzazione Vettoriale Commutativa per la Compressione della Cache KV

CommVQ: Commutative Vector Quantization for KV Cache Compression

June 23, 2025

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

Abstract

I modelli linguistici di grandi dimensioni (LLM) sono sempre più utilizzati in applicazioni che richiedono contesti lunghi, ma la cache chiave-valore (KV) spesso diventa un collo di bottiglia della memoria sulle GPU man mano che il contesto cresce. Per affrontare questo problema, proponiamo la Quantizzazione Vettoriale Commutativa (CommVQ) per ridurre significativamente l'uso della memoria durante l'inferenza di LLM con contesti lunghi. Introduciamo innanzitutto una quantizzazione additiva con un encoder leggero e un codebook per comprimere la cache KV, che può essere decodificata tramite una semplice moltiplicazione di matrici. Per ridurre ulteriormente i costi computazionali durante la decodifica, progettiamo il codebook in modo che sia commutativo con l'Embedding Posizionale Rotazionale (RoPE) e lo addestriamo utilizzando un algoritmo Expectation-Maximization (EM). Ciò consente un'integrazione efficiente della decodifica nel meccanismo di self-attention. Il nostro approccio raggiunge un'elevata accuratezza con la quantizzazione additiva e un basso overhead grazie al codebook commutativo con RoPE. Esperimenti su benchmark di contesti lunghi e su GSM8K dimostrano che il nostro metodo riduce la dimensione della cache KV FP16 dell'87,5% con una quantizzazione a 2 bit, superando i metodi di quantizzazione della cache KV all'avanguardia. In particolare, consente una quantizzazione della cache KV a 1 bit con una perdita di accuratezza minima, permettendo a un modello LLaMA-3.1 8B di funzionare con una lunghezza di contesto di 128K su una singola GPU RTX 4090. Il codice sorgente è disponibile all'indirizzo: 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.