BeaconKV: ビーコンクエリによる KV キャッシュ圧縮手法

大規模推論モデル（LRM）は、拡張された思考連鎖（CoT）生成を通じて優れた問題解決能力を発揮しますが、その結果生じるキーバリュー（KV）キャッシュはシーケンス長に比例して増大し、長い推論トレースでは GPU の容量を超える深刻なメモリボトルネックを引き起こします。既存の KV キャッシュ圧縮手法は、直近のクエリを用いて将来のトークンの重要度を推定しますが、これはこれらのクエリが将来のアテンションパターンの信頼できる代理変数であると暗黙に仮定しています。

しかし、私たちはこの仮説が長期的な推論では成り立たないことを示しました。特定のデコーディングステップでは「思考再訪トークン（TRT）」が発生し、トレースの初期段階で策定されたタスク解決計画など、遠くの過去の文脈に再度アテンションを向けることがあります。体系的な分析を通じて、私たちは TRT に対応するクエリが埋め込み空間内で少数の類似度グループにクラスタリングされることを発見しました。

この知見に基づき、本研究ではトレーニングフリーの KV キャッシュ圧縮手法「BeaconKV」を提案します。これは各グローバルクエリのクラスターに対するコンパクトな代表値である「ビーコンクエリ」を維持し、クエリ履歴全体を保存することなく、どの KV ペアが再訪されるかを予測するものです。

4 つのオープンソース LRM と多様な推論ベンチマークにわたる評価では、BeaconKV は既存の圧縮手法を上回る性能を示しました。最大で 5.8 倍のメモリ削減を実現しながら、キャッシュ精度をほぼ維持し、スループットは 4.3 倍以上向上させています。