# Using Span Queries to Optimize Cache and Attention Locality > [!info] Talk metadata > - **会議:** [[MLSys2026]] Day 3 (May 20 / Wed)、Research Track Oral "LLM Serving 3" セッション（13:30 PDT 開始） > - **登壇者:** Paul Castro, Nick Mitchell, Nathan Ordonez, Thomas Parnell, Mudhakar Srivatsa, Antoni Viros i Martin（全員 IBM Research。連絡先: Nick Mitchell <[email protected]>） > - **URL:** https://mlsys.org/virtual/2026/oral/3747 > - **OpenReview:** https://openreview.net/forum?id=qcGGSXpFcM > - **スライド:** https://mlsys.org/media/mlsys-2026/Slides/3747_FfJdljm.pdf > [!abstract] 概要（論文アブストラクト忠実訳） > LLM のクライアントはチャット補完を超えて進化しつつあり、推論時スケーリングや深い推論技法など多様な手法を取り入れている。一方、推論サーバは依然としてチャット補完に重度に最適化されており、線形的な KV キャッシュ戦略を主に採用している。先行研究は、推論サーバがこうしたより非線形なユースケースへ進化すれば KV キャッシュヒット率の大幅な改善が可能であることを示してきたが、単一ユースケース（RAG など）向けの解にとどまっていた。本論文では、チャット、RAG、推論時スケーリング、エージェント的ワークロードをすべて、より一般的な構造であるスパンクエリの特殊ケースとして表現できることを示す。先行研究が暗黙に仮定していた決定的な区別は、入力の順序が重要か否か——すなわち入力が可換であるか否か——にある。チャットでは順序は重要である。RAG では多くの場合重要ではない。スパンクエリは推論呼び出しの式木であり、可換性制約を伴って結合される。スパンクエリの構文と意味論を定義し、それらが KV キャッシュの局所性を改善するために自動的に最適化できることを示す。vLLM へのわずかな変更（492 行のみ）で、スパンクエリの高性能実行が可能になることを示す。このスタックを用いて、2 つの異なる非チャットユースケースにおいて TTFT で 10〜20 倍の削減を達成できることを実証する。最後に、スパンクエリはアテンションの局所性の改善——いわゆる Lost-in-the-Middle 問題の回避——にも最適化できることを示す。アテンション最適化されたスパンクエリにより、2b パラメータモデルが 8b モデルの精度を大幅に上回ることを実証する。 ## 背景: vLLM と KV キャッシュの局所性 ### ページドアテンションとプレフィクスキャッシング vLLM はトークン列を固定サイズのブロック（ページ）に分割し、ページテーブルを用いて不連続な物理ページを連続的な仮想ページにマッピングする。KV キャッシュの操作（挿入・検索・退避等）はすべてブロック粒度で行われる。キャッシュ検索では各ブロックのハッシュが前方ブロックのハッシュに依存するチェイン構造が使われ、プレフィクス一致のみを探索する。最初のキャッシュミスが発生した時点で走査を打ち切る仕様のため、プレフィクスベースの再利用パターンに特化した設計である。 ### チャット補完ではなぜ高キャッシュヒット率が達成されるか チャット補完では会話が追記的に蓄積される。ターンが進むにつれて先行メッセージのブロックはキャッシュに残り続けるため、再利用が前方から単調に拡大し、キャッシュヒット率は 100% に漸近する。 ### 双対出力パラドックス vLLM のチャット補完 API はクライアントに送るトークン列とキャッシュに挿入するトークン列が異なる場合がある。生成出力ブロックは暗黙的に入力ブロックのハッシュチェインに依存するが、クライアントへは出力トークンのみが返却される。この不一致を論文は「双対出力パラドックス」と呼ぶ。チャットではクライアントが前回出力を次回入力のプレフィクスとして付加することで解消されるが、非チャットユースケースではサーバ側の対処が必要になる。 ## 非チャットユースケースとキャッシュの問題 ### RAG（検索拡張生成） RAG ではコーパスをインデキシングした断片を検索し、ユーザメッセージとともにモデルに入力する。しかしリクエストごとに検索される断片の集合と順序が変わるため、再利用される断片であっても先行プレフィクスが異なるとキャッシュミスとなる。断片が長くなるほどキャッシュヒット率は 0% に漸近する。 ### ネスト生成（推論時スケーリング） judge/generator パターンに代表されるネスト生成では、内側のモデル呼び出しが複数の候補を生成し、外側の呼び出しがそれらを評価・選択する。内側と外側ではシステムプロンプトが異なり、内側の出力が外側の入力にネストされる。共通プレフィクスが存在しないため、キャッシュヒット率は出力長の増大とともに 0% に漸近する。 ### 両ユースケースの本質的問題 チャットの追記的・線形的な再利用パターンとは異なり、RAG とネスト生成はブロックの再利用はあっても位置が変化するため vLLM のプレフィクスベースのキャッシュ戦略では対処できない。 ## ページとスパンの定義 ### ページ（ブロック） ページドアテンションサーバにおけるトークンの連続列。物理メモリ上は不連続でも仮想的に連続として扱われる。 ### スパン ページドアテンションサーバにおけるページの連続列。ページテーブルの上位にスパンテーブルという間接層を導入することで、スパン単位での順序変更を可能にする。 ## 可換性仮説 入力の順序が結果に影響しない場合——すなわち AB = BA が成立する場合——最適化の余地が生まれる。これが本研究の中心的な概念である可換性（commutativity）である。可換性が成立する場合に得られる恩恵は二つある: 1. **キャッシュの局所性**: ブロックが元々異なるコンテキストで使われていたとしてもページを再利用できる。プレフィクスの一致に依存しないキャッシュ戦略が可能になる 2. **アテンションの局所性**: Lost-in-the-Middle 問題に対して分割統治アプローチを適用できる。可換な入力を再配置して重要な情報がアテンションの届きやすい位置に来るよう最適化できる ## スパンクエリの設計 ### 式木としてのスパンクエリ スパンクエリは以下の演算子集合の上に定義される式木（expression tree）である（Table 1）: | 区分 | 演算子 | 意味 | |---|---|---| | 糖衣構文 | C | チャット補完 | | 糖衣構文 | R | コーパス検索 | | 糖衣構文 | F | コーパス断片 | | メッセージ結合 | + | 可換結合（順序不変） | | メッセージ結合 | $\bowtie$ | 非可換結合（順序依存） | | コア演算子 | S, A, U | System / Assistant / User メッセージ | | コア演算子 | G | トークン生成 | C, R, F, G$^1$ はコア演算子（$\bowtie$, +, S, A, U, G）に脱糖できる。 ### メッセージリストからメッセージツリーへ 従来のチャット補完 API はメッセージのリスト（S, A, U の線形列）を入力とする。スパンクエリはこれをメッセージのツリーに一般化し、generate（G）ノードを根に持つ木構造でモデル呼び出しの階層的関係を宣言的に表現する。`query.execute()` により木が実行され出力トークンが生成される。 ### ユースケース別のクエリ構造 **チャット**: G ノードの下に $\bowtie$（非可換結合）で S, A, U を結合する。順序がアテンションに影響するため可換にできない。 **RAG**: チャットと類似の構造だが、検索された断片 F$_1$, F$_2$, ... が +（可換結合）で結合される。断片の順序はアテンション品質に影響しないと宣言することで最適化の余地を生む。 **RAG（サーバサイド検索付き）**: R ノードと C（コーパス仕様）ノードを用いてサーバ側で検索を実行する構造も表現できる。 **judge/generator**: 内側の生成（候補メール作成等）と外側の判定（最良候補の選択）がネストした構造を表現する。内側の生成結果は +（可換結合）で結合され、候補の提示順序は結果に影響しないと宣言できる。ファンアウト数が増えても（2-way、4-way 等）同一の代数的構造で表現できる。 ## KV キャッシュ局所性の改善 ### 高レベル最適化（式木変換） 高レベルオプティマイザは式木に対して固定点反復で 4 種の書換え規則を適用する: 1. **C の脱糖**: C 演算子を G + $\bowtie$ に展開する 2. **R の脱糖**: R 演算子を G$^1$（`max_tokens=1` の生成）+ 可換結合に展開する。これにより断片を本クエリに先立って「準備」できる 3. **+ の簡約**: 可換結合のチェインを不要な入れ子なく平坦化する 4. **+ の分配**: 可換結合 + をトークン生成 G にわたって分配し、双対出力パラドックスを「先回り」で解消する。G の内部に + を配置することで、生成の出力が将来再利用される際にプレフィクスが一致するようにする ### スパンクエリのトークン化 最適化されたクエリを vLLM が消費できる形式にエンコードする。特殊トークン（パディング用 $\square$、独立部分木の開始 `(`、兄弟境界 `)(` 、部分木終端 `)n`）を用いて式木を括弧化し、トークン列に変換する。 ### 低レベル最適化 トークン化されたスパンクエリに対して二つの最適化を行う: 1. **ブロック整列**: 特殊トークンがブロック境界に揃うようパディングを挿入する。これによりブロック先頭のトークンのみを走査すれば特殊トークンの有無を判定でき、全トークンの走査が不要になる 2. **末尾部分ブロックの切り詰め**: 内側生成の出力がブロックを満たさない場合、最終トークンを切り詰めてキャッシュヒット率を維持する。vLLM は部分充填ブロックをキャッシュしないため、切り詰めないとプレフィクス走査が早期に停止してしまう ### vLLM スケジューラ層の変更 ブロックのキャッシュヒット走査において、ハッシュチェインの蓄積ロジックを修正する。ブロック先頭トークンが `(` の場合はチェインを中断し、`)` の場合はチェインを再開する。これによりスパン部分木のブロックが独立にキャッシュ検索される。 ### vLLM GPU ランナー層の変更と CIDRA キャッシュヒット時、再利用ブロックの位置が以前の使用時と異なる場合がある。各 KV ベクトルは RoPE による位置符号化を含むため、位置変更時には逆 RoPE を適用してから新位置の RoPE を再適用する再配置（repositioning、ReRoPE）が必要になる。 論文は Concurrent In-place Duplicating ReRoPE Algorithm（CIDRA）を提案する。CIDRA はブロック再配置の依存グラフ（ブロック A がブロック B の位置に移動、B が C へ……）を構築し、強連結成分分析を行う。出次数が 1 より大きいノード（複数リクエストが同一ブロックを異なる位置に再配置する場合）はブロックを複製する。独立な部分グラフはビンパッキングによりバッチ化して並列 GPU 実行される。CIDRA の再配置スループットは最大約 500 トークン/ms である。 ### vLLM への変更量 変更は合計 492 行に収まる（Table 2）: | vLLM ソースファイル | 変更行数 | |---|---| | `core/block_pool.py` | 68 | | `core/kv_cache_manager.py` | 89 | | `core/kv_cache_utils.py` | 73 | | `core/sched/output.py` | 4 | | `core/sched/scheduler.py` | 19 | | `core/single_type_kv_cache_manager.py` | 4 | | `worker/gpu_model_runner.py` | 232 | | **合計** | **492** | さらにクライアント側ライブラリ `spnl` を提供し、クライアントがスパンクエリを発行すると vLLM が解析・計画・最適化を行う。vLLM 側の主要変更は (1) RoPE を書き込み時ではなく読み出し時に適用する変更、(2) プレフィクス走査時にスパン領域のブロックハッシュチェインを選択的に無効化する変更の 2 点である。 ## RAG での精度評価 ### 実験設定 Block Attention (Ma et al., ICLR 2025) が使用した `ldsjmdy/Tulu3-Block-FT` モデルで 3 つの RAG データセットを評価した。実験は NVIDIA H100 上で実行された。 | データセット | 文書数 | 文書特性 | |---|---|---| | 2Wiki | 3,000 | 短文書（数百トークン） | | NQ | 12,000 | 短文書（数百トークン） | | LongBenchV2 | 86 | 長文書（数万トークン） | ### RAG 精度結果（Table 4） | 実装 | バッチサイズ | 2Wiki | NQ | LongBenchV2 | |---|---|---|---|---| | Block Attention | 1 | 68.36 $\pm$ 0.09% | 62.36 $\pm$ 0.24% | 20.69 $\pm$ 2.37% | | Span Queries | 1 | 69.39 $\pm$ 0.07% | 61.58 $\pm$ 0.28% | 20.57 $\pm$ 2.95% | | Span Queries | 2 | - | - | 21.16 $\pm$ 1.80% | | Span Queries | 16 | 68.14 $\pm$ 0.80% | 61.63 $\pm$ 0.45% | - | 精度は Block Attention と同等であり、バッチサイズの増大による劣化も見られない。LongBenchV2 は GPU メモリ制約によりバッチサイズ 2 に制限された。 ### MSMARCO での精度（スライド） スライドでは MSMARCO データセットでのカテゴリ別精度が示されている: - 全体: +2.6 pp - DESCRIPTION (n=6645): -1.4 pp - NUMERIC (n=1587): +13.2 pp - ENTITY (n=942): +6.7 pp - LOCATION (n=504): +11.9 pp - PERSON (n=322): +5.6 pp HotpotQA では全体 -2.9 pp（bridge: -3.7 pp、comparison: +0.4 pp）とやや劣るケースもある。 ## RAG での TTFT 高速化 ### マイクロベンチマーク 平均 2857 トークンのランダム生成文書を用い、検索文書数を 1〜32 まで変化させて 3 構成の TTFT を比較した: 1. **標準 vLLM**: 各トークンが全先行トークンにアテンドするため、文書数に対して最悪二次的に増大 2. **Span Queries（キャッシュミス）**: スパンの存在によりアテンションが疎になり（文書内トークンは同文書内の先行トークンのみにアテンド）、定数係数が小さい。やはり二次的だが標準 vLLM より大幅に高速 3. **Span Queries（キャッシュヒット）**: プリフィルをスキップし再配置コストのみで済むため、最悪でも線形的に増大 実験結果として、キャッシュヒット時は文書 30 件で標準 vLLM の 25 秒超に対して約 2〜3 秒と、10 倍以上の TTFT 削減を達成した。キャッシュミス時でも約 3 倍の高速化が得られている。 ### バルク実行による追加高速化 バルククエリ実行 API を用いて複数スパンクエリをまとめて実行すると、貪欲法によるクラスタリングで時間的局所性が向上する: | バルクサイズ | 2Wiki | NQ | |---|---|---| | 1 | 1.21x | 1.03x | | 1024 | 1.31x | 1.05x | | コーパス全体 | 1.59x | 1.13x | ## ネスト生成での精度評価 ネスト生成ではクエリ実行時に内側生成出力の末尾部分ブロックを切り詰めるため、重要情報が末尾に位置する場合に精度への影響が懸念される。既存ベンチマークではこのシナリオを分離できないため、needle-in-a-haystack 変種のベンチマークを構築して評価した。内側生成がランダムな名前（needle）をランダム生成内容（hay）に混ぜて出力し、外側の judge がその名前を抽出する。needle の位置（先頭・中間・末尾）と文書長（10, 100, 1000, 10000 文字）を変化させて F1 スコアを測定した。 切り詰めなしの場合、短文書ではスパンクエリの精度は良好で位置依存性も小さい。長文書でもモデルのファインチューニング（Block-FT）により + トークンによるアテンション疎性を適切に活用できている。切り詰めありの場合、短文書では影響が小さいが、長文書で needle が末尾に位置する場合に精度が低下する。これは切り詰めにより needle 自体が除去される可能性があるためであり、予想通りの結果である。 ## ネスト生成での TTFT 高速化 ### 温度依存性 内側生成の温度が 0 の場合は出力が定数となり通常のプレフィクスキャッシングで十分なため、スパンクエリの恩恵は小さい。温度が 0 でない場合、TTFT の改善はファンアウト数に依存する: - 内側ファンアウト 1（温度 0 以外）: 標準 vLLM 比で TTFT を 7〜32% 削減 - 内側ファンアウト 24: 標準 vLLM 比で TTFT を 12〜13 倍高速化 ### ファンアウト数との関係 内側生成数を 1〜32 まで変化させたマイクロベンチマーク（温度 0.5）では、p50 TTFT の高速化がファンアウト数にほぼ線形に増大し、32 生成で最大約 16 倍の高速化を達成する。p99 でも安定した高速化が得られるが、p99 付近で安定性がやや低下し始める。 CIDRA の再配置オーバーヘッドは p50 で最大約 24%（ファンアウト 9〜17 付近）、p99 で最大約 28% だが、TTFT 高速化の恩恵がこのオーバーヘッドを大幅に上回る。 ## アテンション局所性の改善（Lost-in-the-Middle 対策） ### Lost-in-the-Middle 問題 小規模モデルは長い入力の中間に位置する重要情報への注意が弱まる（U 字型の精度分布）。先行研究ではアテンション行列の疎性を高めることでアテンション局所性を改善できることが示されている。 ### 分割統治による最適化 スパンクエリの式木変換を拡張し、8-way judge/generator クエリ（8 候補を一度に判定）を 3 段の 2-way 判定のトーナメント構造に変換する。各段では 2 候補のみを評価するため、judge が注意を向けるべき範囲が狭まりアテンション局所性が改善される。 ### 実験結果 granite3.3 モデルの 2b と 8b の二つの変種で、hay の長さを 0〜800 文字まで変化させて needle 抽出の成功率を測定した: - **未最適化クエリ**: 2b モデルは hay が約 100 文字を超えると急激に精度が低下する。8b モデルはより長い hay に耐えるが、やはり中間位置の情報で精度が落ちる - **アテンション最適化クエリ（k=2 のトーナメント）**: 2b モデルが 800 文字の hay でも高精度を維持し、**未最適化の 8b モデルの精度を大幅に上回る** これはモデルサーバへの変更なしに、クライアント側のクエリ構造変換のみで達成される。 ## 技術的貢献のまとめ 1. **スパンクエリ**: チャット、RAG、推論時スケーリング、エージェント的ワークロードを統一的に表現する宣言的中間表現。可換性制約をクエリに明示的にエンコードすることで自動最適化を可能にする 2. **KV キャッシュ局所性**: vLLM への 492 行の変更で、非チャットユースケースの TTFT を 10〜20 倍改善。キャッシュミス時でもアテンション疎性により約 3 倍の高速化 3. **アテンション局所性**: 式木変換による分割統治最適化で Lost-in-the-Middle を回避し、2b モデルが 8b モデルを上回る精度を達成。モデルサーバへの変更は不要 4. **CIDRA アルゴリズム**: 複数リクエストが重複する KV ブロックセットを再利用する場面で、インプレースの並行 ReRoPE を安全に実行するアルゴリズム ## 今後の展望 - ネスト生成構造が KV キャッシュを経由すべきかどうかの検討。gather/scatter 構造に配置して最適化スケジューラを構築する可能性 - DeepServe、Mooncake、llm-d 等のデータセンタ規模推論エンジンとの統合。コーパス全体を事前にプリフィル・格納し、検索リクエストを先回りで処理する構想 - 切り詰めに代わるアプローチとして、カスタムロジットプロセッサで内側生成出力をブロック整列にバイアスする手法の検討