CompactAttention: Versnelling van blokgewijze prefill met blok-unie KV-selectie

Samenvatting

Chunked prefill is een algemeen toegepaste serving-strategie voor grote taalmodellen met lange contexten geworden, maar efficiënte aandachtsberekening in dit regime blijft uitdagend. Bestaande methoden voor schaarse aandacht zijn primair ontworpen voor eenmalige prefill en zijn niet efficiënt overdraagbaar naar chunked prefill: blok-schaarse kernels verliezen efficiëntie wanneer de querylengte beperkt wordt door de chunkgrootte, terwijl fijnmazige patroonzoektochten kostbaar worden wanneer ze bij elke chunk over de geaccumuleerde KV-cache worden herhaald. QUOKA, een recente methode die direct is gericht op chunked prefill, vermijdt de overhead van schaarse kernels, maar vertrouwt op query-gesubsampleerde, token-niveau KV-selectie, wat kan leiden tot het missen van query-specifieke KV-items en expliciete KV-kopieeroverhead introduceert. Om deze beperkingen aan te pakken, stellen we CompactAttention voor, een chunked-prefill aandachtsmechanisme gebaseerd op Block-Union KV-selectie. CompactAttention behandelt 2D blok-schaarse maskers als KV-selectiesignalen in plaats van directe uitvoeringsplannen voor schaarse kernels, en converteert deze naar GQA-bewuste per-groep KV-bloktabellen via Q-blok-unie en intra-groepsunie. Deze constructie produceert de minimale bloktabellen die alle KV-blokken behouden die door de invoermaskers zijn geselecteerd, onder beperkingen van gepagineerde uitvoering, waardoor geselecteerde KV-blokken op hun plaats kunnen worden benaderd zonder expliciete KV-compactie. Op LLaMA-3.1-8B-Instruct behoudt CompactAttention nauwkeurigheid dicht bij dichte aandacht op de RULER-benchmark, terwijl het tot 2,72× aandachtssnelheidswinst levert bij een contextlengte van 128K onder chunked prefill.

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.