CommVQ: Коммутативное векторное квантование для сжатия кэша ключей и значений

CommVQ: Commutative Vector Quantization for KV Cache Compression

June 23, 2025

Авторы: 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

Аннотация

Крупные языковые модели (LLM) всё чаще применяются в задачах, требующих работы с длинными контекстами, однако кэш ключей и значений (KV) часто становится узким местом в памяти GPU при увеличении длины контекста. Для решения этой проблемы мы предлагаем метод коммутативного векторного квантования (CommVQ), который значительно сокращает использование памяти при выводе LLM с длинными контекстами. Сначала мы вводим аддитивное квантование с использованием легковесного кодера и кодовой книги для сжатия KV-кэша, который может быть декодирован с помощью простого матричного умножения. Чтобы дополнительно снизить вычислительные затраты при декодировании, мы проектируем кодовую книгу так, чтобы она была коммутативна с ротационным позиционным кодированием (RoPE), и обучаем её с использованием алгоритма максимизации ожиданий (EM). Это позволяет эффективно интегрировать декодирование в механизм самовнимания. Наш подход обеспечивает высокую точность благодаря аддитивному квантованию и низкие накладные расходы за счёт RoPE-коммутативной кодовой книги. Эксперименты на бенчмарках с длинными контекстами и GSM8K показывают, что наш метод сокращает размер FP16 KV-кэша на 87,5% при 2-битном квантовании, превосходя современные методы квантования KV-кэша. Примечательно, что он позволяет использовать 1-битное квантование KV-кэша с минимальной потерей точности, что позволяет модели LLaMA-3.1 8B работать с длиной контекста 128K на одном GPU RTX 4090. Исходный код доступен по адресу: 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.