Kara: スライディングウィンドウKVキャッシュ圧縮による効率的な推論LLMサービス

推論言語モデル（特にチェーン・オブ・ソート推論を行うモデル）は、長い思考連鎖を生成する際、デコードフェーズで膨大なKVキャッシュが蓄積される。これにより、高いレイテンシと限られたスループットが生じる。KVキャッシュ圧縮は、重要でないKVペアを選択的に削除することでメモリ負荷を軽減する有望な技術である。しかし、既存手法には二つの主要な限界がある。第一に、閾値トリガー圧縮ポリシーはスループット改善が限定的であり、場合によっては低下させる。また、シーケンスの特定ブロックからKVペアを完全に除去し、情報損失を悪化させる可能性がある。第二に、既存手法は孤立したKVペアまたは固定サイズのチャンクのみを保持するため、任意のトークン位置にある重要な可変サイズのチャンクを保存できない。

これらの限界を克服するために、著者らはKaraを提案する。Karaはスライディングウィンドウ方式を採用し、最近生成されたコンテキストのみを対象にデコード時に圧縮を実行する。双方向アテンションを利用してウィンドウ内の各KVペアのアテンションスコアを計算し、情報価値の高いものを選択する。さらに、重要な意味情報を柔軟に保持するために、Token2Chunkモジュールを設計した。このモジュールは選択されたKVペアのサブセットを可変サイズのチャンクに拡張し、固定境界による情報の断片化を防ぐ。これにより、任意の位置にある可変サイズの意味ブロックを保持できるようになった。

KaraはさらにPagedAttentionに適応され、vLLMをベースにした推論フレームワークKvLLMが構築された。PagedAttentionはページング管理によりKVキャッシュの断片化を減らす効率的なメモリ管理技術である。KaraとPagedAttentionの組み合わせにより、メモリ使用量がさらに最適化される。実験は言語モデリング、質問応答、推論タスクなど多岐にわたって行われ、KaraとKvLLMが一貫してKVキャッシュメモリを削減し、出力スループットを大幅に向上することを示した。例えば、長いシーケンス生成タスクでは、Karaはベースライン手法と比較してKVキャッシュサイズを最大50%削減し、生成品質を維持または向上させた。

本研究の意義は、長いコンテキスト推論の効率的なデプロイメントに実用的なソリューションを提供する点にある。特に、数学的推論、コード生成、マルチステップの質問応答など、多くのCoT推論を必要とするアプリケーションに有用である。Karaの設計は圧縮効率と情報保持のバランスを考慮しており、将来の推論モデルのサービス最適化に道筋を示すものである。