CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection

Technische Zusammenfassung: CompactAttention

Problemstellung
Da Large Language Models (LLMs) skaliert werden, um Kontextfenster von Hunderttausenden bis zu Millionen von Tokens zu unterstützen, ist die Verarbeitung langer Sequenzen über einen einzigen „One-Shot"-Prefill-Durchlauf aufgrund quadratischer Rechenkosten und des Risikos, Decode-Anfragen in Online-Service-Systemen zu blockieren, unpraktisch geworden. Folglich hat sich chunked prefill (Prefill in Abschnitten) – die sequenzielle Verarbeitung von Eingaben in festen Größenabschnitten unter Berücksichtigung der akkumulierten Key-Value (KV)-Caches – zur Standard-Service-Strategie in Frameworks wie vLLM und SGLang entwickelt.

Jedoch stoßen bestehende Beschleunigungsmethoden in diesem Regime auf erhebliche Einschränkungen:

1. Ineffizienz block-sparser Aufmerksamkeit: Traditionelle block-sparse Methoden (z. B. XAttention, SeerAttention) verlassen sich auf sparse Kernel, die lange Query-Längen benötigen, um den Overhead irregulärer Speicherzugriffe und der Maskeninterpretation zu amortisieren. Beim chunked prefill ist die Query-Länge durch die Chunk-Größe begrenzt ($Q \ll KV$), was dazu führt, dass diese Kernel Parallelismus und Effizienz verlieren und oft keine bessere Leistung als dichte Aufmerksamkeit erzielen.
2. Overhead der Muster-Suche: Sparse Methoden erfordern das Finden eingangsabhängiger sparse Muster. Beim chunked prefill muss diese Suche für jeden Chunk über den wachsenden akkumulierten KV-Cache wiederholt werden, wodurch der kumulative Such-Overhead zu einem kritischen Engpass wird.
3. Einschränkungen der direkten KV-Auswahl: Methoden wie QUOKA vermeiden sparse Kernel, indem sie eine Teilmenge von KV-Tokens basierend auf einer unterprobenen Menge von Queries auswählen und dichte Aufmerksamkeit durchführen. Dieser Ansatz leidet jedoch unter zwei Problemen:
   • Deckungslücken: Query-spezifische KV-Einträge, die für nicht unterprobenen Queries wichtig sind, können übersehen werden, was zu Genauigkeitsverlusten bei Aufgaben führt, die einen verteilten Informationszugriff erfordern.
   • Kopier-Overhead: Die Token-Ebene-Auswahl erfordert das explizite Zusammenführen und Kopieren von KV-Einträgen in einen kompakten Puffer vor der Ausführung der Aufmerksamkeit, was einen Overhead einführt, der mit der Kontextlänge und der Batch-Größe skaliert.

Methodik: CompactAttention
Die Autoren schlagen CompactAttention vor, einen Mechanismus, der die Block-Level-KV-Auswahl von der Ausführung der Aufmerksamkeit entkoppelt. Anstatt 2D block-sparse Masken als direkte Ausführungspläne für sparse Kernel zu verwenden, behandelt CompactAttention sie als KV-Auswahlsignale, um GQA-bewusste pro-Gruppe KV-Block-Tabellen zu konstruieren.

Die Methode arbeitet in zwei Stufen:

1. Auswahlstufe (Konstruktion der Block-Union-KV-Tabelle):

   • Das System akzeptiert eine 2D pro-Head block-sparse Maske, die von einer beliebigen leichten Muster-Suchmethode generiert wird (z. B. SeerAttention oder FlashPrefill).
   • Es wird eine Q-Block-Union über Query-Blöcke innerhalb des aktuellen Chunks angewendet. Dies ist notwendig, da dichte paginierte Aufmerksamkeit eine einzelne KV-Blockliste für eine Gruppe von gemeinsam ausgeführten Query-Blöcken verbraucht, anstatt individuelle Entscheidungen für jedes $(Q\text{-Block}, KV\text{-Block})$-Tile zu treffen.
   • Es wird eine intra-gruppale Union über Query-Heads angewendet, die eine KV-Gruppe teilen (standardmäßig Grouped-Query-Attention-Gruppen).
   • Das Ergebnis ist eine minimale pro-Gruppe KV-Block-Tabelle ($T_{b,g}$), die alle KV-Blöcke bewahrt, die von den ursprünglichen Eingabemasken ausgewählt wurden. Diese Konstruktion gewährleistet die Erhaltung der Abdeckung: Jeder KV-Block, der von einem beliebigen Query-Block oder Head in der Gruppe ausgewählt wurde, bleibt ausgewählt.

2. Ausführungsstufe (Zero-Copy paginierte Aufmerksamkeit):

   • Im Gegensatz zu Methoden der Token-Ebene-Auswahl arbeitet CompactAttention auf Block-Granularität.
   • Es wird ein KV-Head-major Speicherlayout $[B, H_{kv}, L, D]$ verwendet, anstatt des standardmäßigen sequenz-major Layouts. Dies ermöglicht, dass jedes $(Batch, KV\text{-Head}, Block)$-Triple als zusammenhängender Speicherbereich behandelt wird.
   • Die konstruierten Block-Tabellen werden als Metadaten (speziell kv_indptr und kv_indices) an einen Standard-paginierten dichten Attention-Kernel übergeben.
   • Der Kernel greift auf die ausgewählten KV-Blöcke in-place aus dem ursprünglichen KV-Cache zu, ohne explizites Kopieren oder Komprimieren in einen neuen Puffer.

Hauptbeiträge

• Entkopplung von Auswahl und Ausführung: Die Arbeit stellt ein Design vor, das die Logik, welche KV-Blöcke ausgewählt werden sollen, von der Art und Weise, wie sie ausgeführt werden, trennt. Dies ermöglicht die Wiederverwendung leichter sparse-Muster-Suchmethoden, während die Ineffizienzen sparse Kernel vermieden werden.
• Block-Union-Konstruktion: Die neuartige Verwendung von Q-Block- und intra-gruppalen Unions, um 2D sparse Masken in ausführbare, GQA-bewusste Block-Tabellen umzuwandeln, die minimal, aber abdeckend sind.
• Zero-Copy paginierte Ausführung: Durch die Nutzung eines KV-Head-major Layouts und paginierter Attention-Metadaten eliminiert die Methode den expliziten KV-Kopier-Overhead, der Methoden der Token-Ebene-Auswahl wie QUOKA inhärent ist.

Experimentelle Ergebnisse
Evaluiert an LLaMA-3.1-8B-Instruct und Qwen3-30B-A3B-Instruct-2507 unter chunked prefill:

• Beschleunigung: Auf H200 GPUs mit einer Kontextlänge von 128K erreicht CompactAttention bis zu 2,72-fache Beschleunigung der Aufmerksamkeit und 1,96-fache End-to-End-Beschleunigung gegenüber der Basislinie dichter Aufmerksamkeit. Dies übertrifft block-sparse Baselines (die im Regime $Q \ll KV$ oft schlechter abschneiden als dichte Aufmerksamkeit) und QUOKA (die durch Kopier-Overhead begrenzt ist) erheblich.
• Genauigkeit: CompactAttention hält die Genauigkeit auf den Benchmarks RULER und LongBench V2 nahe an der dichten Aufmerksamkeit. Es übertrifft QUOKA, insbesondere bei Aufgaben, die einen verteilten Informationszugriff erfordern (z. B. Multi-key NIAH), wo die Query-Untersampling von QUOKA zu Genauigkeitsverlusten führt.
• Effizienz: Ablationsstudien bestätigen, dass die Zero-Copy paginierte Ausführungsstrategie effizienter ist als sowohl sparse-Kernel-Ausführung als auch explizite KV-Komprimierung, selbst unter Berücksichtigung des Overheads für den Aufbau der Metadaten.

Bedeutung
Die Arbeit behauptet, dass die primäre Engstelle für sparse Aufmerksamkeit beim chunked prefill nicht lediglich welche KV-Blöcke ausgewählt werden sollen, sondern wie die ausgewählten Blöcke ausgeführt werden. Indem CompactAttention durch Block-Union-Konstruktion und Zero-Copy-Ausführung eine sparse Auswahl mit effizienten dichten paginierten Attention-Kerneln kompatibel macht, bietet es einen praktischen Weg zu einem schnelleren Long-Context LLM-Service. Es adressiert die spezifischen Ineffizienzen des Regimes $Q \ll KV$, die die Adoption bestehender sparse-Aufmerksamkeitsmethoden in chunked-prefill-Systemen behindert haben.