クラスタリングに基づくKVCache残差量子化による大規模言語モデル推論の高速化

近年，LLMの汎用性能を活かした応用が拡大している．一方でモデルの大規模化に伴い，推論時の計算・メモリ要求が増大し，効率的推論が重要な課題となっている．多くのLLMはTransformer[1]をベースとしており，Attentionの計算が必要となる．この計算を高速化するため，各層のKey/Valueを保存して再利用するKVCacheが一般に用いられる．しかしKVCacheは層数・ヘッド数・隠れ次元・生成長に比例して増加し，深刻なメモリボトルネックとなる．特にGPU環境ではKVCacheがメモリを圧迫し，バッチサイズや同時処理数を制限してスループット/レイテンシや運用コストに影響する[2]．この課題に対し，低ビット量子化[3,4]，重要度の低いトークンの削除・要約[5,6]，Attention構造や保持方式の改変[7,8,9]などが提案されているが，タスク依存の性能劣化[10]，既存実装との互換性といったトレードオフが残る．またKeyとValueは役割・統計特性が異なり，一様な量子化は品質劣化の要因となり得る．さらにアルゴリズム上の工夫が必ずしも実スループット向上に直結せず，メモリアクセスや計算効率を含む検討が必要である．本研究ではKVCacheの統計的構造を活用した量子化を提案する．事前サンプルからKVCacheをクラスタリングして重心を得て，推論時は各Key/Valueを最近傍重心で近似し，残差のみを量子化して保持することで，高圧縮と低誤差を両立する．さらにKeyには次元ごとのダイナミックレンジ差を考慮した次元単位量子化を適用し，Valueにはランダム回転による分散均一化[11]後にLloyd–Max量子化[12,13]を行うことで誤差分散を低減する．加えてGPU上で量子化・復号とAttentionを統合した最適化CUDAカーネルを実装し，メモリ削減に加えて推論スループット向上を達成した．実験ではメモリ使用量を大幅に削減しつつ，ベースライン比で最大約3.5倍のスループット向上を確認した．