BitDecoding: Unlocking Tensor Cores for Long-Context LLMs Decoding with Low-Bit KV Cache

要約

長いコンテキストの大規模な言語モデル（LLMS）の採用の増大は、キー価値（kV）キャッシュの拡大により、自己回帰デコードに重要なメモリと計算上の課題をもたらしました。
KVキャッシュの量子化は有望なソリューションとして浮上しており、4ビットまたは2ビットの量子化でさえ、メモリコストを削減しながらモデルの精度を維持できることを示しています。
ただし、これらの利点にもかかわらず、低ビットKVキャッシュの予備的実装は、量子化と定量化オーバーヘッド、およびテンソルコアの利用の欠如による予想されるスピードアップを提供するのに苦労しています。
この作業では、低ビットKVキャッシュで効率的なデコードのためにテンソルコアのロックを解除するGPU最適化されたフレームワークであるBitDecodingを提案します。
低ビットKVキャッシュのテンソルコアを効率的に活用することは、各デコードステップでのKVキャッシュ生成の動的な性質のために困難です。
ビットデコードは、テンソルコアの高い利用を可能にするデータレイアウトの互換性を保証するテンソルコア中心のビットフュージョンスキームでこれらの課題に対処します。
さらに、BitDecodingには、ワープ効率の高い並列デコードカーネルと細粒の非同期パイプラインが組み込まれているため、脱量化のオーバーヘッドを最小限に抑え、計算効率が改善されます。
実験では、FP16 FlashDecoding-V2と比較して、BitDecodingがRTX 4090で最大7.5倍のスピードアップを達成し、H100で4.8x、H100で8.9xを達成することが示されています。
また、最先端の低ビットKVキャッシュ実装（QSERVE）を最大4.3倍上回ることも優れています。
128Kシーケンス長のllama-3.1-8bでは、ビットデコードはシングルバッチデコードレイテンシを3倍に減らし、長いコンテキスト生成シナリオでの有効性を示します。
このコードは、https：//github.com/dd-duda/bitdecodingで入手できます。

要約(オリジナル)

The growing adoption of long-context Large Language Models (LLMs) has introduced significant memory and computational challenges in autoregressive decoding due to the expanding Key-Value (KV) cache. KV cache quantization has emerged as a promising solution, with prior work showing that 4-bit or even 2-bit quantization can maintain model accuracy while reducing memory costs. However, despite these benefits, preliminary implementations for the low-bit KV cache struggle to deliver the expected speedup due to quantization and dequantization overheads and the lack of Tensor Cores utilization. In this work, we propose BitDecoding, a GPU-optimized framework that unlocks Tensor Cores for efficient decoding with low-bit KV cache. Efficiently leveraging Tensor Cores for low-bit KV cache is challenging due to the dynamic nature of KV cache generation at each decoding step. BitDecoding addresses these challenges with a Tensor Cores-Centric BitFusion Scheme that ensures data layout compatibility to enable high utilization of Tensor Cores. Additionally, BitDecoding incorporates a warp-efficient parallel decoding kernel and a fine-grained asynchronous pipeline, minimizing dequantization overhead and improving computational efficiency. Experiments show that BitDecoding achieves up to 7.5x speedup on RTX 4090, 4.8x on A100, and 8.9x on H100, compared to FP16 FlashDecoding-v2. It also outperforms the state-of-the-art low-bit KV cache implementation (QServe) by up to 4.3x. On LLaMA-3.1-8B with a 128K sequence length, BitDecoding reduces single-batch decoding latency by 3x, demonstrating its effectiveness in long-context generation scenarios. The code is available at https://github.com/DD-DuDa/BitDecoding.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google