CompactAttention: Beschleunigung des Chunked Prefill durch Block-Union KV-Auswahl

Zusammenfassung

Chunked Prefill hat sich zu einer weit verbreiteten Serving-Strategie für Large Language Models mit langen Kontexten entwickelt, doch die effiziente Attention-Berechnung in diesem Bereich bleibt herausfordernd. Bestehende Methoden für sparse Attention sind primär für einmaliges Prefill ausgelegt und lassen sich nicht effizient auf Chunked Prefill übertragen: Blocksparse-Kernel verlieren an Effizienz, wenn die Query-Länge durch die Chunk-Größe begrenzt ist, während feinkörnige Mustersuche aufwändig wird, wenn sie bei jedem Chunk über den angesammelten KV-Cache wiederholt werden muss. QUOKA, eine neuere Methode, die direkt auf Chunked Prefill abzielt, vermeidet den Overhead von Sparse-Kerneln, basiert jedoch auf einer Query-Subsampling-basierten, tokenweisen KV-Auswahl, die query-spezifische KV-Einträge übersehen und expliziten KV-Kopieraufwand verursachen kann. Um diese Einschränkungen zu adressieren, schlagen wir CompactAttention vor, einen Attention-Mechanismus für Chunked Prefill, der auf Block-Union-KV-Auswahl basiert. CompactAttention behandelt 2D-Blocksparse-Masken als KV-Auswahlsignale anstatt als direkte Ausführungspläne für Sparse-Kernel und wandelt sie durch Q-Block-Union und Intra-Gruppen-Union in GQA-bewusste KV-Blocktabellen pro Gruppe um. Diese Konstruktion erzeugt die minimalen Blocktabellen, die unter den Einschränkungen der seitenbasierten Ausführung alle durch die Eingabemasken ausgewählten KV-Blöcke erhalten, sodass auf ausgewählte KV-Blöcke direkt ohne explizite KV-Kompaktierung zugegriffen werden kann. Auf LLaMA-3.1-8B-Instruct bleibt CompactAttention im RULER-Benchmark nahe an der Genauigkeit von dichter Attention und erzielt bei einer Kontextlänge von 128K unter Chunked Prefill eine bis zu 2,72-fache Attention-Beschleunigung.

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.