CommVQ: Cuantización Vectorial Conmutativa para la Compresión de la Caché 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

Resumen

Los Modelos de Lenguaje de Gran Escala (LLMs, por sus siglas en inglés) se utilizan cada vez más en aplicaciones que requieren contextos largos, pero la caché de clave-valor (KV) suele convertirse en un cuello de botella de memoria en las GPU a medida que el contexto crece. Para abordar este problema, proponemos la Cuantización Vectorial Conmutativa (CommVQ) para reducir significativamente el uso de memoria en la inferencia de LLMs con contextos largos. Primero, introducimos la cuantización aditiva con un codificador ligero y un libro de códigos para comprimir la caché KV, la cual puede decodificarse mediante una simple multiplicación de matrices. Para reducir aún más los costos computacionales durante la decodificación, diseñamos el libro de códigos para que sea conmutativo con la Incrustación de Posición Rotatoria (RoPE) y lo entrenamos utilizando un algoritmo de Expectation-Maximization (EM). Esto permite una integración eficiente de la decodificación en el mecanismo de auto-atención. Nuestro enfoque logra alta precisión con cuantización aditiva y bajo sobrecosto gracias al libro de códigos conmutativo con RoPE. Los experimentos en benchmarks de contextos largos y GSM8K muestran que nuestro método reduce el tamaño de la caché KV en FP16 en un 87.5% con cuantización de 2 bits, superando a los métodos de cuantización de caché KV más avanzados. Destacablemente, permite la cuantización de 1 bit en la caché KV con una pérdida mínima de precisión, permitiendo que un modelo LLaMA-3.1 8B funcione con un contexto de 128K en una sola GPU RTX 4090. El código fuente está disponible en: 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.