CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection

Riepilogo Tecnico: CompactAttention

Enunciato del Problema
Man mano che i Large Language Models (LLM) scalano per supportare finestre di contesto da centinaia di migliaia a milioni di token, l'elaborazione di sequenze lunghe tramite un singolo passaggio di prefill "one-shot" è diventata impraticabile a causa dei costi computazionali quadratici e del rischio di bloccare le richieste di decodifica nei sistemi di servizio online. Di conseguenza, il prefill frammentato (chunked prefill) — ovvero l'elaborazione degli input in modo sequenziale in blocchi di dimensione fissa, prestando attenzione alle cache Key-Value (KV) accumulate — è diventata la strategia di servizio standard in framework come vLLM e SGLang.

Tuttavia, i metodi di accelerazione esistenti presentano limitazioni significative in questo regime:

1. Inefficienza dell'Attention Sparsa a Blocchi: I metodi tradizionali a blocchi sparsi (ad esempio XAttention, SeerAttention) si basano su kernel sparsi che richiedono lunghe query per ammortizzare l'overhead dell'accesso irregolare alla memoria e dell'interpretazione delle maschere. Nel prefill frammentato, la lunghezza della query è limitata dalla dimensione del blocco ($Q \ll KV$), causando la perdita di parallelismo ed efficienza da parte di questi kernel, che spesso non riescono a superare l'attention densa.
2. Overhead della Ricerca di Pattern: I metodi sparsi richiedono la ricerca di pattern sparsi dipendenti dall'input. Nel prefill frammentato, questa ricerca deve essere ripetuta sulla cache KV accumulata in crescita per ogni blocco, rendendo l'overhead cumulativo di ricerca un collo di bottiglia critico.
3. Limitazioni della Selezione Diretta KV: Metodi come QUOKA evitano i kernel sparsi selezionando un sottoinsieme di token KV basato su un insieme di query campionati e eseguendo un'attention densa. Tuttavia, questo approccio soffre di due problemi:
   • Lacune di Copertura: Le voci KV specifiche per una query, importanti per query non campionate, potrebbero essere perse, portando a un degrado della precisione in compiti che richiedono un accesso distribuito alle informazioni.
   • Overhead di Copia: La selezione a livello di token richiede una raccolta e una copia esplicita delle voci KV in un buffer compatto prima dell'esecuzione dell'attention, introducendo un overhead che scala con la lunghezza del contesto e la dimensione del batch.

Metodologia: CompactAttention
Gli autori propongono CompactAttention, un meccanismo che disaccoppia la selezione dei blocchi KV a livello di blocco dall'esecuzione dell'attention. Invece di utilizzare maschere sparse 2D come piani di esecuzione diretti per kernel sparsi, CompactAttention le tratta come segnali di selezione KV per costruire tabelle di blocchi KV per gruppo consapevoli di GQA.

Il metodo opera in due fasi:

1. Fase di Selezione (Costruzione della Tabella KV Unione di Blocchi):

   • Il sistema accetta una maschera sparsa a blocchi per testa 2D generata da qualsiasi metodo di ricerca di pattern leggero (ad esempio SeerAttention o FlashPrefill).
   • Applica una unione di blocchi Q (Q-block union) tra i blocchi di query all'interno del blocco corrente. Questo è necessario perché l'attention paged densa consuma un'unica lista di blocchi KV per un gruppo di blocchi di query eseguiti insieme, piuttosto che decisioni individuali per ogni piastrella $(Q\text{-block}, KV\text{-block})$.
   • Applica una unione intra-gruppo tra le teste di query che condividono un gruppo KV (impostazione predefinita sui gruppi di Grouped-Query Attention).
   • Il risultato è una tabella minima di blocchi KV per gruppo ($T_{b,g}$) che preserva tutti i blocchi KV selezionati dalle maschere di input originali. Questa costruzione garantisce la preservazione della copertura: qualsiasi blocco KV selezionato da qualsiasi blocco di query o testa nel gruppo rimane selezionato.

2. Fase di Esecuzione (Attention Paged Zero-Copy):

   • A differenza dei metodi di selezione a livello di token, CompactAttention opera a granularità di blocco.
   • Utilizza un layout di memoria KV-head-major $[B, H_{kv}, L, D]$ invece del layout sequence-major standard. Ciò consente di trattare ogni tripla $(batch, KV\text{-head}, block)$ come una regione di memoria contigua.
   • Le tabelle di blocchi costruite vengono passate come metadati (specificamente kv_indptr e kv_indices) a un kernel standard di attention paged densa.
   • Il kernel accede ai blocchi KV selezionati in-place dalla cache KV originale senza copia esplicita o compattazione in un nuovo buffer.

Contributi Chiave

• Disaccoppiamento di Selezione ed Esecuzione: Il documento introduce un design che separa la logica di quali blocchi KV selezionare da come vengono eseguiti, consentendo il riutilizzo di metodi leggeri di ricerca di pattern sparsi evitando al contempo le inefficienze dei kernel sparsi.
• Costruzione dell'Unione di Blocchi: L'uso innovativo delle unioni di blocchi Q e intra-gruppo per convertire maschere sparse 2D in tabelle di blocchi eseguibili e consapevoli di GQA, che sono minime ma preservano la copertura.
• Esecuzione Paged Zero-Copy: Sfruttando un layout KV-head-major e metadati di attention paged, il metodo elimina l'overhead di copia KV esplicito intrinseco nei metodi di selezione a livello di token come QUOKA.

Risultati Sperimentali
Valutati su LLaMA-3.1-8B-Instruct e Qwen3-30B-A3B-Instruct-2507 sotto prefill frammentato:

• Accelerazione: Su GPU H200 con una lunghezza di contesto di 128K, CompactAttention raggiunge fino a 2.72× di accelerazione nell'attention e 1.96× di accelerazione end-to-end rispetto alla baseline di attention densa. Questo supera significativamente le baseline sparse a blocchi (che spesso performano peggio dell'attention densa nel regime $Q \ll KV$) e QUOKA (che è limitata dall'overhead di copia).
• Precisione: CompactAttention mantiene una precisione vicina all'attention densa sui benchmark RULER e LongBench V2. Supera QUOKA, in particolare su compiti che richiedono un accesso distribuito alle informazioni (ad esempio Multi-key NIAH), dove il campionamento delle query di QUOKA porta a un degrado della precisione.
• Efficienza: Gli studi di ablazione confermano che la strategia di esecuzione paged zero-copy è più efficiente sia dell'esecuzione con kernel sparsi che della compattazione esplicita KV, anche tenendo conto dell'overhead di costruzione dei metadati.

Significato
Il documento afferma che il collo di bottiglia principale per l'attention sparsa nel prefill frammentato non è semplicemente quali blocchi KV selezionare, ma come i blocchi selezionati vengono eseguiti. Rendendo la selezione sparsa compatibile con efficienti kernel di attention paged densa attraverso la costruzione di unioni di blocchi e l'esecuzione zero-copy, CompactAttention offre una via pratica verso un servizio LLM a contesto lungo più veloce. Affronta le specifiche inefficienze del regime $Q \ll KV$ che hanno ostacolato l'adozione dei metodi di attention sparsa esistenti nei sistemi di prefill frammentato.