CommVQ : Quantification Vectorielle Commutative pour la Compression du Cache KV

CommVQ: Commutative Vector Quantization for KV Cache Compression

June 23, 2025

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

Résumé

Les modèles de langage de grande taille (LLMs) sont de plus en plus utilisés dans des applications nécessitant des contextes longs, mais le cache clé-valeur (KV) devient souvent un goulot d'étranglement mémoire sur les GPU à mesure que le contexte s'allonge. Pour résoudre ce problème, nous proposons la Quantisation Vectorielle Commutative (CommVQ) afin de réduire significativement l'utilisation de la mémoire pour l'inférence des LLMs à contexte long. Nous introduisons d'abord une quantification additive avec un encodeur léger et un codebook pour compresser le cache KV, qui peut être décodé via une simple multiplication matricielle. Pour réduire davantage les coûts de calcul lors du décodage, nous concevons le codebook pour qu'il soit commutatif avec l'Embedding de Position Rotatif (RoPE) et l'entraînons à l'aide d'un algorithme d'Espérance-Maximisation (EM). Cela permet une intégration efficace du décodage dans le mécanisme d'auto-attention. Notre approche atteint une haute précision avec la quantification additive et un faible surcoût grâce au codebook commutatif avec RoPE. Les expériences sur des benchmarks à contexte long et GSM8K montrent que notre méthode réduit la taille du cache KV FP16 de 87,5 % avec une quantification à 2 bits, tout en surpassant les méthodes de quantification de cache KV les plus avancées. Notamment, elle permet une quantification à 1 bit du cache KV avec une perte de précision minimale, permettant à un modèle LLaMA-3.1 8B de fonctionner avec une longueur de contexte de 128K sur un seul GPU RTX 4090. Le code source est disponible à l'adresse suivante : 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.