QJL: 1-Bit Quantized JL Transform for KV Cache Quantization with Zero Overhead

要約

LLM を提供するには、KV キャッシュ内の Key-Value (KV) 埋め込みのストレージ要件により、大量のメモリが必要になります。これは、シーケンスの長さに応じて増大します。
KV キャッシュを圧縮する効果的なアプローチは量子化です。
ただし、従来の量子化方法では、量子化定数 (少なくともゼロ点とスケール) をデータ ブロックごとに完全な精度で保存する必要があるため、重大なメモリ オーバーヘッドに直面します。
ブロック サイズに応じて、このオーバーヘッドにより量子化された数値ごとに 1 ビットまたは 2 ビットが追加される可能性があります。
QJL は、ジョンソン リンデンシュトラウス (JL) 変換とそれに続く符号ビット量子化で構成される新しい量子化アプローチです。
既存の方法とは対照的に、QJL では量子化定数を保存する必要がなくなるため、メモリのオーバーヘッドがなくなります。
我々は 2 つのベクトルの内積に対する非対称推定器を提案し、一方のベクトルに QJL を適用し、もう一方のベクトルに量子化なしの標準的な JL 変換を適用すると、歪みが最小限に抑えられた不偏推定器が提供されることを実証します。
私たちは、最適化された計算のための軽量 CUDA カーネルを組み込んだ、QJL スケッチとそれに対応する内積推定器の効率的な実装を開発しました。
QJL をさまざまな LLM および NLP タスクに適用して KV キャッシュをわずか 3 ビットに量子化すると、精度を損なうことなく KV キャッシュ メモリ使用量が 5 倍以上削減され、同時に実行時間の高速化が実現します。
コードは \url{https://github.com/amirzandieh/QJL} で入手できます。

要約(オリジナル)

Serving LLMs requires substantial memory due to the storage requirements of Key-Value (KV) embeddings in the KV cache, which grows with sequence length. An effective approach to compress KV cache is quantization. However, traditional quantization methods face significant memory overhead due to the need to store quantization constants (at least a zero point and a scale) in full precision per data block. Depending on the block size, this overhead can add 1 or 2 bits per quantized number. We introduce QJL, a new quantization approach that consists of a Johnson-Lindenstrauss (JL) transform followed by sign-bit quantization. In contrast to existing methods, QJL eliminates memory overheads by removing the need for storing quantization constants. We propose an asymmetric estimator for the inner product of two vectors and demonstrate that applying QJL to one vector and a standard JL transform without quantization to the other provides an unbiased estimator with minimal distortion. We have developed an efficient implementation of the QJL sketch and its corresponding inner product estimator, incorporating a lightweight CUDA kernel for optimized computation. When applied across various LLMs and NLP tasks to quantize the KV cache to only 3 bits, QJL demonstrates a more than fivefold reduction in KV cache memory usage without compromising accuracy, all while achieving faster runtime. Codes are available at \url{https://github.com/amirzandieh/QJL}.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google