[論文レビュー] KIVI : Plug-and-play 2bit KV Cache Quantization with Streaming Asymmetric Quantization

Zirui Liu, Jiayi Yuan+6 more|arXiv (Cornell University)|Jan 1, 2023

Quantum-Dot Cellular Automata被引用数 6

ひとこと要約

KIVIは、大規模言語モデルの推論におけるキー・バリュー（KV）キャッシュに、微調整不要のプラグアンドプレイ型2ビット非対称量子化手法を提案する。キーキャッシュにはチャネルごとの量子化、バリューキャッシュにはトークンごとの量子化を適用する。これにより、ピークメモリ使用量が2.6倍削減され、バッチサイズを最大4倍に拡大可能となり、Llama、Falcon、Mistralモデルにおいて最小限の精度損失で2.35倍～3.47倍の高いスループットを達成する。

ABSTRACT

Efficiently serving large language models (LLMs) requires batching of many requests to reduce the cost per request. Yet, with larger batch sizes and longer context lengths, the key-value (KV) cache, which stores attention keys and values to avoid re-computations, significantly increases memory demands and becomes the new bottleneck in speed and memory usage. Additionally, the loading of the KV cache causes the computational core to be idle, which limits the inference speed. A straightforward and effective solution to reduce KV cache size is quantization, which decreases the total bytes taken by KV cache. However, there is a lack of in-depth studies that explore the element distribution of KV cache to understand the hardness and limitation of KV cache quantization. To fill the gap, we conducted a comprehensive study on the element distribution in KV cache of popular LLMs. Our findings indicate that the key cache should be quantized per-channel, i.e., group elements along the channel dimension and quantize them together. In contrast, the value cache should be quantized per-token. From this analysis, we developed a tuning-free 2bit KV cache quantization algorithm named KIVI. With hardware-friendly implementation, KIVI can enable Llama, Falcon, and Mistral models to maintain almost the same quality while using $\mathbf{2.6 imes}$ less peak memory (including model weight). This reduction in memory usage enables up to $\mathbf{4 imes}$ larger batch size, bringing $\mathbf{2.35 imes \sim 3.47 imes}$ throughput on real LLM inference workload. The source code is available at https://github.com/jy-yuan/KIVI.

研究の動機と目的

• 高いバッチサイズと長いコンテキスト長さの下で、KVキャッシュに起因する増大するメモリと速度のボトルネックに対処すること。
• 代表的な大規模言語モデル（LLM）におけるKVキャッシュの要素分布を分析し、量子化の課題と制限を理解すること。
• モデルの精度を維持しながら、メモリ使用量を大幅に削減できる、ハードウェアにやさしい微調整不要の2ビット量子化手法を開発すること。
• 効率的なKVキャッシュ圧縮を通じて、実世界のLLM推論ワークロードにおけるより大きなバッチサイズと高いスループットを実現すること。

提案手法

• 微調整やモデル再トレーニングを必要としない、プラグアンドプレイ型の2ビット非対称量子化手法KIVIを提案する。
• キーキャッシュにチャネルごとの量子化を適用し、チャネル次元に沿って要素をグループ化することで、各チャネルごとの量子化誤差を制限する。
• バリューキャッシュにトークンごとの量子化を適用し、自己回帰的生成のストリーミング特性に適合させ、誤差を各トークンごとに隔離する。
• KVキャッシュをグループ部と残差部に分割：グループ部にはグループ単位の量子化を適用し、残差部はフル精度で保持する。
• アテンション計算中にグループ部と残差部を組み合わせるためにタイル型行列乗算を用い、精度を保持する。
• プリフィルおよびデコード段階における量子化オーバーヘッドを最小限に抑えるハードウェアにやさしい実装を採用する。

実験結果

リサーチクエスチョン

• RQ1 LLMにおけるキーキャッシュとバリューキャッシュの要素分布はどのように異なり、これらの差異が量子化戦略にどのような影響を及えるか？
• RQ2 なぜキーキャッシュにはチャネルごとの量子化が効果的であり、バリューキャッシュにはトークンごとの量子化がより適しているのか？
• RQ3 ストリーミング的かつ動的な性質を持つKVキャッシュにおいて、微調整なしに2ビット量子化スキームが高精度を達成可能か？
• RQ4 異なる量子化スキームが、実際のLLM推論ワークロードにおけるメモリ使用量、バッチサイズ、スループットに与える影響は何か？
• RQ5 KIVIにおけるグループ化と残差分割機構は、大規模でストリーミング的なKVキャッシュの効率的量子化を可能にしつつ、精度をどのように維持しているか？

主な発見

• キーキャッシュのチャネルごとの量子化は、チャネル間での誤差伝播を低減し、特に高マグニチュードのアウリヤーを持つチャネルが存在するため不可欠である。
• バリューキャッシュのトークンごとの量子化は不可欠である。なぜなら、バリューキャッシュはアテンション計算におけるミキサーの役割を果たし、誤差を各トークンごとに隔離する必要があるからである。
• KIVIは、Llama-2-7Bにおいてピークメモリ使用量を2.6倍削減しながら、複数のベンチマークでほぼ同一のモデル品質を維持する。
• この手法により、バッチサイズを最大4倍に拡大でき、実際のLLM推論ワークロードで2.35倍～3.47倍の高いスループットを達成する。
• アブレーションスタディの結果、グループサイズと残差長は性能に測定可能なが、管理可能な影響を及ぼすことが判明し、最適な設定は経験的に特定された。
• KIVIは、Llama、Falcon、Mistralなど多様なモデルにおいて、2ビット精度ですら最小限の精度損失で優れた性能を発揮する。