BeaconKV: ビーコンクエリに導かれるKVキャッシュ圧縮による大規模推論モデルの効率的な推論

要旨

大規模推論モデル（LRM）は、拡張されたChain-of-Thought (CoT) 生成を通じて優れた問題解決を達成するが、その結果生じるキー・バリュー（KV）キャッシュは系列長に比例して線形に増大し、深刻なメモリボトルネックを引き起こし、長い推論トレースにおいてしばしばGPU容量を超える。既存のKVキャッシュ圧縮手法は、最近のクエリに依存して将来のトークンの重要度を推定し、暗黙のうちにこれらが将来のアテンションパターンの信頼できる代理指標となると仮定している。我々は、この仮定が長期的な推論において失敗することを示す。すなわち、特定のデコーディングステップが、トレースの初期に定式化されたタスク解決計画など、遠く離れた以前のコンテキストに再アテンションする思考再訪トークン（TRT）を生成するのである。系統的な分析を通じて、我々はTRTに対応するクエリが埋め込み空間において少数の類似グループにクラスタリングされることを発見した。この洞察に基づき、我々はBeaconKVを提案する。これはトレーニング不要のKVキャッシュ圧縮手法であり、各グローバルクエリクラスタのコンパクトな代表であるビーコンクエリを保持することにより、クエリ履歴全体を保存することなく、どのKVペアが再訪されるかを予測する。4つのオープンソースLRMと多様な推論ベンチマークにおいて、BeaconKVは一般に既存の圧縮手法を上回り、フルキャッシュの精度をほぼ維持しつつ、最大5.8倍のメモリ削減を達成し、スループットを4.3倍以上向上させる。

English

Large Reasoning Models (LRMs) achieve superior problem-solving through extended Chain-of-Thought (CoT) generation, but the resulting key-value (KV) cache grows linearly with sequence length and creates severe memory bottlenecks, often exceeding GPU capacity for long reasoning traces. Existing KV cache compression methods rely on recent queries to estimate future token importance, implicitly assuming these serve as reliable proxies for future attention patterns. We demonstrate that this assumption fails in long-horizon reasoning: certain decoding steps generate Thought Revisiting Tokens (TRT) that re-attend to distant previous context, such as task-solving plans formulated early in the trace. Through systematic analysis, we discover that queries corresponding to the TRT cluster into a small number of similarity groups in the embedding space. Based on this insight, we propose BeaconKV, a training-free KV cache compression method that maintains beacon queries, compact representatives for each global query cluster, to anticipate which KV pairs will be revisited without storing the entire query history. Across four open-source LRMs and diverse reasoning benchmarks, BeaconKV generally outperforms existing compression methods, achieving up to 5.8times memory reduction while nearly preserving full cache accuracy and improving throughput by over 4.3times.