KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache

要約

大規模言語モデル (LLM) を効率的に処理するには、リクエストあたりのコストを削減するために多くのリクエストをバッチ処理する必要があります。
しかし、バッチ サイズが大きくなり、コンテキストの長さが長くなると、再計算を避けるためにアテンション キーと値を格納するキーバリュー (KV) キャッシュがメモリ需要を大幅に増加させ、速度とメモリ使用量の新たなボトルネックになります。
さらに、KV キャッシュのロードにより計算コアがアイドル状態になり、推論速度が制限されます。
KV キャッシュ サイズを削減する簡単で効果的な解決策は、KV キャッシュが使用する総バイト数を減らす量子化です。
ただし、KV キャッシュの量子化の難しさと制限を理解するために、KV キャッシュの要素の分布を調査する詳細な研究は不足しています。
このギャップを埋めるために、一般的な LLM の KV キャッシュ内の要素の分布に関する包括的な調査を実施しました。
私たちの調査結果は、キー キャッシュはチャネルごとに量子化する必要がある、つまりチャネルの次元に沿って要素をグループ化し、それらを一緒に量子化する必要があることを示しています。
対照的に、値キャッシュはトークンごとに量子化する必要があります。
この分析から、KIVI というチューニング不要の 2 ビット KV キャッシュ量子化アルゴリズムを開発しました。
ハードウェアに優しい実装により、KIVI を使用すると、Llama、Falcon、および Mistral モデルがほぼ同じ品質を維持しながら、$\mathbf{2.6\times}$ 少ないピーク メモリ (モデルの重みを含む) を使用することができます。
このメモリ使用量の削減により、バッチ サイズを最大 $\mathbf{4\times}$ 大きくすることができ、実際の LLM 推論ワークロードで $\mathbf{2.35\times \sim 3.47\times}$ のスループットが実現します。
ソース コードは https://github.com/jy-yuan/KIVI で入手できます。

要約(オリジナル)

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\times}$ less peak memory (including model weight). This reduction in memory usage enables up to $\mathbf{4\times}$ larger batch size, bringing $\mathbf{2.35\times \sim 3.47\times}$ throughput on real LLM inference workload. The source code is available at https://github.com/jy-yuan/KIVI.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google