CommVQ: Quantização Vetorial Comutativa para Compressão de Cache KV

CommVQ: Commutative Vector Quantization for KV Cache Compression

June 23, 2025

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

Resumo

Modelos de Linguagem de Grande Escala (LLMs) são cada vez mais utilizados em aplicações que exigem contextos longos, mas o cache de chave-valor (KV) frequentemente se torna um gargalo de memória em GPUs à medida que o contexto aumenta. Para resolver isso, propomos a Quantização Vetorial Comutativa (CommVQ) para reduzir significativamente o uso de memória na inferência de LLMs com contextos longos. Primeiro, introduzimos a quantização aditiva com um codificador leve e um codebook para comprimir o cache KV, que pode ser decodificado por meio de uma simples multiplicação de matrizes. Para reduzir ainda mais os custos computacionais durante a decodificação, projetamos o codebook para ser comutativo com o Embedding de Posição Rotacional (RoPE) e o treinamos usando um algoritmo de Expectation-Maximization (EM). Isso permite a integração eficiente da decodificação no mecanismo de auto-atenção. Nossa abordagem alcança alta precisão com quantização aditiva e baixa sobrecarga por meio do codebook comutativo com RoPE. Experimentos em benchmarks de contexto longo e GSM8K mostram que nosso método reduz o tamanho do cache KV FP16 em 87,5% com quantização de 2 bits, superando os métodos de quantização de cache KV mais avançados. Notavelmente, ele permite a quantização de 1 bit do cache KV com perda mínima de precisão, permitindo que um modelo LLaMA-3.1 8B funcione com um contexto de 128K em uma única GPU RTX 4090. O código-fonte está disponível em: 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.