SparseEngine、長文脈エージェント向けに疎な推論を中心とした設計を提案

長文脈 LLM エージェントでは、会話、ツール呼び出し、タスクの中間状態が履歴として蓄積されます。これらを KV キャッシュに保持すると、コンテキストが伸びるほど GPU メモリの消費と注意計算のコストが増えます。疎な注意や KV の削除は負荷を抑える手段ですが、方式ごとにキャッシュの配置、保持ルール、実行手順が異なるため、既存の汎用推論エンジンへ統合するのは容易ではありません。

SparseEngine は、この問題を個別の最適化ではなく、サービス基盤の設計問題として扱います。最初から疎な推論を前提にエンジンを構成し、共通のライフサイクル契約を導入しました。各方式は KV をどのように表現するか、どの履歴を残すか、どのように計算するかを自分で制御できます。一方、リクエスト処理に必要な状態遷移は共通のサービス基盤が調整します。これにより、方式固有のロジックと、スケジューリングやリクエスト管理を分離しやすくなります。

主なポイント

• 4 つのカテゴリにまたがる 15 種類の方式をサポートし、異なるキャッシュ表現やワークフローを扱います。
• Chain Cache はリクエスト間の状態管理を可能にします。KV 削除方式は保持された履歴から処理を再開でき、毎回完全なコンテキストを作り直す必要を減らします。
• 制御可能な Prefix-Cache Pruning は、履歴の指定領域から KV を削除しながら、論理的なプレフィックス一致を維持します。
• 論文の要約では、KV 削除時のスループットが 10 倍超、同じ並列度でのデコードが vLLM より 2.5 倍超、エージェントベンチマークのエンドツーエンド速度が 2 倍超向上したと報告されています。ただし、提示された素材にはハードウェア、モデル、詳細な測定条件がないため、数値は評価設定と合わせて解釈する必要があります。

この研究の意義は、疎な推論を単なるカーネル高速化から、状態を扱うサービスアーキテクチャへ広げた点にあります。長文脈エージェントでは、一度の注意計算を軽くするだけでなく、複数の関連リクエストにまたがって状態を保存し、復元し、削除し、再利用しなければなりません。SparseEngine のような契約型の設計は、方式ごとに別の実行基盤を用意する負担を減らす可能性があります。

もっとも、報告された高速化がすべての環境で再現するとは限りません。実際の効果は、モデル、履歴の分布、キャッシュヒット率、並列度、許容できる品質変化などに左右されます。それでも本プロジェクトは、長文脈エージェントの推論では、注意カーネルだけでなく KV 状態のライフサイクルも主要な設計対象になることを示しています。実装は https://github.com/CURRENTF/SparseEngine で公開されています。

出典：Hugging Face Daily Papers