WUSH-KV、データ適応型変換で2-bit KVキャッシュ量子化を実現

論文「WUSH-KV: KV Cache Quantization with Data-Adaptive Transforms」は、低ビットKVキャッシュ量子化手法「WUSH-KV」を提案しています。WUSH-KVはキャリブレーションデータからキー用・バリュー用の変換を別々に構築し、2-bitの評価では、評価した全モデルと全下流タスクでOSCAR変換と同等か上回りました。

何が新しいか

WUSH-KVは、行列積の両因子の二次統計からデータ認識型変換を構築して量子化誤差を減らす手法「WUSH」を、KVキャッシュ向けに適応したものです。論文は背景として、KVキャッシュのメモリと帯域幅のコストがコンテキスト長とバッチサイズに伴って増大し、長文脈推論の効率を制限している点を挙げています。

仕組みと評価結果

項目	内容
変換の構築	キャリブレーションデータから、キーとバリューで別々に構築
バリュー変換	モデル重みに折り込む
キー変換	RoPEの後に適用
組み合わせる量子化器	クリップ付き量子化器(例: QuEST INT)
理論的性質	QuEST INT使用時、軽微な仮定の下でWUSH変換はnear-optimal
エンドツーエンド評価	SGLangに統合、OSCAR-style percentile-clipped affine quantizationを使用
評価ビット幅	2-bit

QuEST INTと組み合わせた場合、WUSH-KVは層単位の再構成誤差を低減し、テストした他の変換の中で最も低いエンドツーエンドのperplexityを達成しました。SGLangに統合した2-bitのエンドツーエンド評価では、評価対象の全モデルと全下流タスクでOSCAR変換と同等か、それを上回っています。

適用手順

1. キャリブレーションデータから、キー用とバリュー用の変換を構築します。
2. バリュー変換をモデル重みに折り込み、キー変換はRoPE後に適用します。
3. 変換をクリップ付き量子化器と組み合わせます。
4. 論文のエンドツーエンド評価では、WUSH-KVをSGLangに統合し、OSCAR方式のパーセンタイルクリップ付きアフィン量子化を使っています。

SGLangは、LMSYSのもとで公開されている言語モデル・マルチモーダルモデル向けの推論フレームワークで、単体GPUから分散環境までを対象としています。