💬 NLP

CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection

CompactAttention は、2 次元のブロック疎マスクを GQA 対応のグループ別 KV ブロックテーブルに変換するブロック結合 KV 選択メカニズムを導入することで、長文脈大規模言語モデルにおけるチャンク化プレフィルを高速化し、明示的な圧縮を行わずにインプレースメモリアクセスを可能にしながら、ほぼ密な精度を維持し、最大 2.72 倍の速度向上を実現します。

原著者： Jiwon Song, Dongwon Jo, Beomseok Kang, Jae-Joon Kim

公開日 2026-05-19

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Jiwon Song, Dongwon Jo, Beomseok Kang, Jae-Joon Kim

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは図書館司書（AI）になりきって、膨大な数の本（コンテキスト）から質問に答える場面を想像してください。過去には、巨大な図書館があれば、正しい答えを見つけるために一冊一冊の本を読み通さなければならず、それは永遠に続くような時間がかかりました。これを加速させるため、研究者たちは「チャンク化されたプリフィル」システムを発明しました。これは、図書館全体を一度に読むのではなく、小さなバッチ（チャンク）ごとに読み、進みながらメモ帳（KV キャッシュ）にメモを追加していくというものです。

しかし、新たな問題が発生しました：新しい質問のバッチが来るたびに、メモ帳のすべての内容を再び読み直さずに、どのようにしてメモ帳の中の正しいページを素早く見つけられるのでしょうか？

この論文は、この問題を解決する新しい方法としてCompactAttentionを紹介しています。その仕組みを、簡単なアナロジーを用いて説明します。

旧来の方法の問題点

この論文は、人々がこの問題を解決しようとして試みた 2 つの主要な方法と、それらがなぜ失敗したかを特定しています。

「スパースカーネル」アプローチ（非効率なスキャナー）:
- アイデア: 図書館の地図を用意し、重要な本だけを赤い点でマークします。そして、白いスペース（重要でない部分）をスキップして、赤い点の箇所だけを見るようにします。
- 失敗: 巨大な図書館（長いコンテキスト）を読みながら、小さな質問（小さなチャンク）しかしていない場合、この方法は遅くなります。これは、壁一面のテキストをスキャンするには優れたスキャナーを持っているのに、たった一文しかない場合、スキャナーのセットアップと較正に時間がかかりすぎるようなものです。「白いスペースをスキップする」ためのオーバーヘッドが、実はすべてを読み通すよりも遅くさせてしまいます。
「クエリ部分サンプリング」アプローチ（怠慢な図書館司書）:
- アイデア: すべての質問をチェックする代わりに、バッチからいくつかのランダムな質問だけを選び、それらのために重要な本を見つけ、その本が全員にとって重要であると仮定します。
- 失敗: これはリスクがあります。間違った少数の質問を選んでしまうと、たった一つの特定の質問だけが必要とした重要な本を見逃してしまう可能性があります。また、一度本を選んだら、それらを棚から特別なテーブルへ物理的に運ばなければ読み始められません。この「運ぶこと」（データの転送）には、多くの時間とエネルギーを要します。

解決策：CompactAttention

CompactAttention は、本を見つけることと本を読むことを分離することでゲームを変えます。

ステップ 1: 「ユニオン」戦略（検索のグループ化）

複雑な「スキップリスト」（スパースカーネル）を実行したり、少数の質問に基づいて推測したりする代わりに、CompactAttention はスマートなグループ化のトリックを使用します。

事件を捜査する探偵チーム（クエリヘッド）がいると想像してください。各探偵は、自分にとって重要だと考える「容疑者」（KV ブロック）のリストを持っています。
CompactAttention は、各探偵に個別に作業させるのではなく、「チーム全体の容疑者を一つのマスターリストにまとめよう」と言います。
これは 2 つのステップで行われます。
1. Q-ブロックユニオン: 現在のバッチに含まれるすべての質問のリストを結合します。
2. グループ内ユニオン: 一緒に働く探偵たちのリストを結合します。
結果: 全員が必要とするものを網羅した、単一の最小限の「マスターリスト」が得られます。どの探偵も必要な本はリストに含まれるため、重要な本が置き去りにされることはありません。

ステップ 2: 「ゼロコピー」実行（その場で読む）

ここが魔法のパートです。

旧来の方法: マスターリストができたら、素早く読むために、それらの本を棚から特別なテーブルへ物理的に移動させなければなりませんでした。この「移動」には時間がかかります。
CompactAttention の方法: 本を移動させません。代わりに、図書館司書に地図（メタデータ）を手渡します。「棚 A、3 列、5 番の本へ行き、次に棚 B、1 列、2 番の本へ」というようなものです。
図書館司書（コンピュータカーネル）は、棚のその場所へ直接行き、そこにある本を読みます。これは**「ゼロコピーページドアテンション」**と呼ばれます。データを移動させるために費やされていた時間とエネルギーをすべて節約します。

なぜ画期的なのか

この論文は、128,000 語のコンテキスト（非常に長い文書）を持つ巨大な AI モデル（LLaMA-3.1-8B）でこれをテストしました。

精度: 図書館全体を読み通すこと（密アテンション）と同等の賢さでした。重要な見落としはありませんでした。
速度: 標準的な方法に比べて最大2.72 倍高速でした。

結論

CompactAttention を、図書館の配置換えを諦め、代わりに完璧な統合された索引カードを使う賢い図書館司書だと考えてください。

「検索」（重要なブロックを見つけること）と「実行」（それらを読むこと）を分離すべきだと認識し、何一つ見逃さないようにする「グループ化」のトリックを用いることで、知性を失うことなく長文書の AI 処理を劇的に高速化することに成功しました。ボトルネックは、どの本を選ぶかという点だけでなく、それらをどのように取り出すかという点にあったことを証明しました。

技術サマリー：CompactAttention

問題定義
大規模言語モデル（LLM）が数十万から数百万トークンのコンテキストウィンドウをサポートするようにスケールするにつれ、単一の「ワンショット」プリフィルパスによる長系列の処理は、二次的な計算コストとオンライン serving システムにおけるデコード要求のストールリスクにより、非現実的となっています。その結果、蓄積されたキー・バリュー（KV）キャッシュに注意を向けながら、固定サイズのチャンクで入力順次処理を行う「チャンクドプリフィル」が、vLLM や SGLang などのフレームワークにおける標準的な serving 戦略となっています。

しかし、既存の加速手法はこの領域において重大な制限に直面しています：

ブロックスパース注意の非効率性: 従来のブロックスパース手法（XAttention、SeerAttention など）は、不規則なメモリアクセスとマスク解釈のオーバーヘッドを償却するために長いクエリ長を必要とするスパースカーネルに依存しています。チャンクドプリフィルでは、クエリ長がチャンクサイズによって制限され（ $Q \ll KV$ ）、これらのカーネルは並列性と効率性を失い、しばしば密な注意（dense attention）を上回ることができません。
パターン探索のオーバーヘッド: スパース手法は入力依存のスパースパターンを探索する必要があります。チャンクドプリフィルでは、この探索が各チャンクに対して成長する蓄積 KV キャッシュ全体で繰り返される必要があり、累積的な探索オーバーヘッドが重要なボトルネックとなります。
直接 KV 選択の限界: QUOKA などの手法は、スパースカーネルを回避するために、サブサンプリングされたクエリのセットに基づいて KV トークンのサブセットを選択し、密な注意を実行します。しかし、このアプローチは以下の 2 つの問題に悩まされています：
- カバレッジの欠落: サンプル化されなかったクエリにとって重要なクエリ固有の KV エントリが見逃され、分散情報アクセスを必要とするタスクにおいて精度が低下します。
- コピーオーバーヘッド: トークンレベルの選択は、注意実行前に KV エントリを明示的に収集してコンパクトなバッファにコピーすることを必要とし、コンテキスト長とバッチサイズに比例してオーバーヘッドが増大します。

手法：CompactAttention
著者は、ブロックレベルの KV 選択と注意実行を分離する「CompactAttention」を提案します。2 次元ブロックスパースマスクをスパースカーネルの直接実行計画として使用するのではなく、CompactAttention はそれらをKV 選択信号として扱い、GQA 対応のグループ別 KV ブロックテーブルを構築します。

この手法は 2 つの段階で動作します：

選択段階（ブロック結合 KV テーブルの構築）:
- システムは、軽量なパターン探索手法（SeerAttention や FlashPrefill など）によって生成された 2 次元のヘッド別ブロックスパースマスクを受け入れます。
- 現在のチャンク内のクエリブロック間でQ-ブロック結合を適用します。これは、密なページド注意が、個々の $(Q\text{-block}, KV\text{-block})$ タイルごとの決定ではなく、一緒に実行されるクエリブロックのグループに対して単一の KV ブロックリストを消費するため、必要となります。
- KV グループを共有するクエリヘッド間でグループ内結合を適用します（デフォルトでは Grouped-Query Attention グループ）。
- その結果、元の入力マスクによって選択されたすべての KV ブロックを保持する最小限のグループ別 KV ブロックテーブル（ $T_{b,g}$ ）が得られます。この構築により、カバレッジの保持が保証されます。つまり、グループ内の任意のクエリブロックまたはヘッドによって選択された KV ブロックは、すべて選択されたままとなります。
実行段階（ゼロコピーページド注意）:
- トークンレベルの選択手法とは異なり、CompactAttention はブロック粒度で動作します。
- 標準的なシーケンス優先レイアウトではなく、KV-ヘッド優先メモリラayout $[B, H_{kv}, L, D]$ を利用します。これにより、各 $(batch, KV\text{-head}, block)$ の組が連続したメモリ領域として扱われます。
- 構築されたブロックテーブルは、標準的なページド密注意カーネルにメタデータ（具体的には kv_indptr と kv_indices）として渡されます。
- カーネルは、明示的なコピーや新しいバッファへの圧縮を行わず、元の KV キャッシュから選択された KV ブロックにその場でアクセスします。

主要な貢献

選択と実行の分離: どの KV ブロックを選択するかという論理と、それらをどのように実行するかという論理を分離する設計を導入し、軽量なスパースパターン探索手法の再利用を可能にしながら、スパースカーネルの非効率性を回避します。
ブロック結合構築: 2 次元スパースマスクを実行可能で GQA 対応のブロックテーブルに変換するための、Q-ブロック結合とグループ内結合の新たな利用法。これらは最小限でありながらカバレッジを保持します。
ゼロコピーページド実行: KV-ヘッド優先レイアウトとページド注意メタデータを活用することで、QUOKA などのトークンレベル選択手法に固有の明示的な KV コピーオーバーヘッドを排除します。

実験結果
チャンクドプリフィル条件下でLLaMA-3.1-8B-InstructおよびQwen3-30B-A3B-Instruct-2507を評価しました：

高速化: H200 GPU で 128K のコンテキスト長において、CompactAttention は密な注意のベースラインに対して最大2.72 倍の注意速度向上と1.96 倍のエンドツーエンド速度向上を達成しました。これは、 $Q \ll KV$ の領域で密な注意を下回ることが多いブロックスパースベースラインや、コピーオーバーヘッドに制限される QUOKA を大幅に上回ります。
精度: CompactAttention は、RULERおよびLongBench V2ベンチマークにおいて、密な注意に近い精度を維持します。特に分散情報アクセスを必要とするタスク（例：Multi-key NIAH）において、QUOKA のクエリサブサンプリングによる精度低下を回避し、QUOKA を上回ります。
効率性: 除去実験により、メタデータ構築のオーバーヘッドを考慮しても、ゼロコピーページド実行戦略がスパースカーネル実行および明示的な KV 圧縮の両方よりも効率的であることが確認されました。

意義
本論文は、チャンクドプリフィルにおけるスパース注意の主要なボトルネックは、単にどの KV ブロックを選択するかではなく、選択されたブロックをどのように実行するかであると主張しています。ブロック結合構築とゼロコピー実行を通じて、スパース選択を効率的な密なページド注意カーネルと互換性のあるものにする CompactAttention は、より高速な長コンテキスト LLM serving への実用的な道筋を提供します。これは、既存のスパース注意手法のチャンクドプリフィルシステムへの採用を妨げてきた、 $Q \ll KV$ 領域固有の非効率性に対処するものです。