Prefix Sliding、プロンプトのプレフィックスと直近の推論を保持し、超長時間思考のKVメモリを定額化

推論モデルで完全なattentionを使用すると、新しいtokenを1つ生成するたびに、増大し続けるKV cacheを読み取る必要がある。思考の軌跡が長くなるほど、1ステップ当たりのコストとメモリ要件は高くなる。Prefix Slidingは、Qwen3-1.7Bのattention分析に基づき、コンテキストを常に参照可能な2つの領域に分割する。1つはシステム指示、ツール定義、元の質問を含むプレフィックスで、もう1つは直近数千個の推論tokenからなるsliding windowだ。その間にある古いドラフトは破棄される。100-tokenのプレフィックスと4,096-tokenのwindowを使用する場合、attentionが処理するtokenは最大4,196個となり、生成コストは思考の軌跡全体の長さに伴って線形に増加しなくなる。

この手法は単調増加するposition IDをそのまま使用するため、windowの移動時にも、すでにRoPEを適用したKVエントリを再計算する必要がない。チームはさらにFlashAttentionを改変し、tile内maskingとblock全体のskipという2段階のフィルタリングによって、プレフィックスと直近のwindowのみを計算するようにした。単体の80GB H100上でvLLMを使って1,024本のsequenceを生成したテストでは、windowが定常状態に入った後もthroughputは約5,000 token/秒を維持した一方、完全なattentionはさらに低速化し続けた。論文が示す約3倍の高速化は、同じ時間内により多くの推論tokenを生成できるというend-to-endの結果であり、tokenごとの推論品質が向上するという意味ではない。

研究では、この構造を強化学習にも導入している。100,000-tokenのrollout全体をtrainerへ渡す必要はなく、例では末尾の8,192 tokenのみを返し、先頭の6,144 tokenをコンテキストとして使用し、最後の2,048 tokenだけでlossを計算する。著者らは7Bモデルを用いた制御実験で、完全なattentionに近い性能を達成できると報告している。ただし、gradientはtruncated approximationであり、frontier規模のモデルまで拡張できることはまだ実証されていない。

制約にもエンジニアリング上の重要な意味がある。LiveCodeBenchでは、モデルが数千token後に以前のコードの続きを書く場合があり、windowが小さすぎると関数の冒頭部分が失われる。また、大量のツール出力によって、有効なコンテキストが一気にwindow外へ押し出される可能性もある。公開コードは依然として旧バージョンのTorch、vLLM、Prime-RL、FlashAttentionを使用しており、カスタムkernelのインストールには約10時間かかる。エンジニアリングチームは、本番のinference engineへの統合を判断する前に、ワークロードに応じたwindow size、長い出力の比率、過去の情報を参照し直す必要性を検証すべきだ。