CommVQ：KVキャッシュ圧縮のための可換ベクトル量子化

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

要旨

大規模言語モデル（LLMs）は、長いコンテキスト長を必要とするアプリケーションでますます使用されていますが、コンテキストが長くなるにつれて、キー・バリュー（KV）キャッシュがGPU上のメモリボトルネックとなることがよくあります。この問題に対処するため、我々はCommutative Vector Quantization（CommVQ）を提案し、長いコンテキストのLLM推論におけるメモリ使用量を大幅に削減します。まず、軽量なエンコーダとコードブックを用いた加法量子化を導入し、KVキャッシュを圧縮します。これは単純な行列乗算によってデコード可能です。デコード時の計算コストをさらに削減するため、コードブックをRotary Position Embedding（RoPE）と可換になるように設計し、Expectation-Maximization（EM）アルゴリズムを用いて訓練します。これにより、デコードをセルフアテンションメカニズムに効率的に統合することが可能になります。我々のアプローチは、加法量子化による高精度と、RoPE可換コードブックによる低オーバーヘッドを実現します。長いコンテキストのベンチマークとGSM8Kでの実験により、我々の方法は2ビット量子化でFP16 KVキャッシュサイズを87.5%削減し、最先端のKVキャッシュ量子化手法を上回ることを示しました。特に、最小限の精度損失で1ビットKVキャッシュ量子化を可能にし、LLaMA-3.1 8Bモデルを単一のRTX 4090 GPUで128Kのコンテキスト長で実行できるようにします。ソースコードは以下で公開されています：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.