CompactAttention: ブロックユニオンKV選択によるチャンク化プリフィルの高速化

要旨

チャンク化プリフィルは、長コンテキスト大規模言語モデルにおいて広く採用されるサービング戦略となっているが、この領域での効率的なアテンション計算は依然として困難である。既存のスパースアテンション手法は主にワンショットプリフィル向けに設計されており、チャンク化プリフィルに効率的に適用できない。ブロックスパースカーネルは、クエリ長がチャンクサイズに制限されると効率が低下し、一方、細粒度パターン検索は、各チャンクで累積されたKVキャッシュに対して繰り返し行われるとコストが高くなる。チャンク化プリフィルを直接対象とした最近の手法であるQUOKAは、スパースカーネルのオーバーヘッドを回避するが、クエリサブサンプリングによるトークンレベルのKV選択に依存しており、その結果、クエリ固有のKVエントリを見逃したり、明示的なKVコピーのオーバーヘッドが生じる可能性がある。これらの制限に対処するため、我々はブロックユニオンKV選択に基づくチャンク化プリフィルアテンション機構であるCompactAttentionを提案する。CompactAttentionは、2次元ブロックスパースマスクを直接的なスパースカーネル実行計画ではなくKV選択信号として扱い、Qブロックユニオンとグループ内ユニオンを通じて、それらをGQAを考慮したグループごとのKVブロックテーブルに変換する。この構成により、ページング実行制約の下で入力マスクによって選択されたすべてのKVブロックを保持する最小限のブロックテーブルが生成され、選択されたKVブロックを明示的なKVコンパクションなしにその場でアクセスできるようになる。LLaMA-3.1-8B-Instructにおいて、CompactAttentionはRULERベンチマークで密なアテンションに近い精度を維持しつつ、チャンク化プリフィル下の128Kコンテキスト長において最大2.72倍のアテンション高速化を実現する。

English

Chunked prefill has become a widely adopted serving strategy for long-context large language models, but efficient attention computation in this regime remains challenging. Existing sparse attention methods are primarily designed for one-shot prefill and do not translate efficiently to chunked prefill: block-sparse kernels lose efficiency when the query length is limited by the chunk size, while fine-grained pattern search becomes costly when repeated over the accumulated KV cache at every chunk. QUOKA, a recent method that directly targets chunked prefill, avoids sparse-kernel overhead but relies on query-subsampled, token-level KV selection, which can miss query-specific KV entries and introduce explicit KV-copy overhead. To address these limitations, we propose CompactAttention, a chunked-prefill attention mechanism based on Block-Union KV Selection. CompactAttention treats 2D block-sparse masks as KV-selection signals rather than direct sparse-kernel execution plans, and converts them into GQA-aware per-group KV block tables through Q-block union and intra-group union. This construction produces the minimal block tables that preserve all KV blocks selected by the input masks under paged execution constraints, enabling selected KV blocks to be accessed in place without explicit KV compaction. On LLaMA-3.1-8B-Instruct, CompactAttention maintains accuracy close to dense attention on the RULER benchmark while delivering up to 2.72times attention speedup at 128K context length under chunked prefill.