Technische Samenvatting: CoSA - Versnelling van Long-Context Inferentie via Proxy-Kernel Co-Designed Sparse Attention

1. Probleemstelling

De kwadratische computationele kosten van self-attention mechanismen maken long-context inferentie voor moderne Large Language Models (LLMs) prohibitief duur. Hoewel sparse attention een oplossing biedt door selectief belangrijke query-key interacties te berekenen, kampen bestaande benaderingen met een fundamentele afruil tussen efficiëntie en getrouwheid (fidelity), met name onder strikte sparsity-budgetten.

Huidige methoden vallen over het algemeen uiteen in twee categorieën:

1. Proxy-gestuurde Benaderingen: Deze gebruiken een lichtgewicht proxy om een binaire sparse mask te voorspellen voordat de kernel wordt uitgevoerd. Hoewel deze proxies effectief zijn bij gematigde budgetten, missen ze onvermijdelijk saillante blokken wanneer het budget wordt beperkt, wat leidt tot significante nauwkeurigheidsdegradatie.
2. In-Kernel Skipping Benaderingen: Deze gebruiken exacte Softmax-statistieken binnen de kernel om verwaarloosbare blokken over te slaan. Hoewel accuraat, lijden ze aan twee beperkingen:
   • Dichte Proxy Vereiste: Ze moeten de volledige Query-Key (QK) interactie voor elk blok berekenen om een skip-beslissing te nemen, waardoor de dominante kwadratische QK-kosten niet worden geëlimineerd.
   • Conservatief Overslaan: Vanwege het "bucket effect" (waarbij een enkele uitschieter in een rij de skip blokkeert) en het gebruik van een running maximum (dat een ondergrens is van het globale maximum in plaats van een globaal maximum), passen deze methoden overly conservatieve skip-regels toe, waardoor potentieel snelheidswinsten onbenut blijven.

De kern van de geïdentificeerde bottleneck is de ontkoppeling tussen de proxy en de backend kernel. Proxies genereren statische binaire masks zonder rekening te houden met de executievolgorde, terwijl kernels in vaste volgordes worden uitgevoerd, wat een optimale sparsity in de weg staat.

2. Methodologie: CoSA

De auteurs stellen CoSA voor, een twee-fasen, training-vrije block-sparse attention mechanism gebaseerd op het principe van proxy-kernel co-design. CoSA koppelt een Kernel-Aware Proxy (KAP) aan een Ordered-Skipping Kernel (OSK) om zowel de blokselectie als de executievolgorde gezamenlijk te optimaliseren.

2.1 Kernprincipes van het Ontwerp

• Orde-invariantie: Door gebruik te maken van de eigenschap dat Online-Softmax (OSM) accumulatie onafhankelijk is van de volgorde waarin key-blokken worden bezocht, staat CoSA willekeurige herordening van blokverwerking toe zonder de uiteindelijke output te veranderen.
• Computation-Order Mask: In plaats van een binair mask genereert CoSA een mask dat de specifieke volgorde voorschrijft waarin de KV-pages bezocht moeten worden.

2.2 Fase 1: Kernel-Aware Proxy (KAP)

De KAP opereert onder een gematigd budget om blokken te selecteren en hun bezoekvolgorde te bepalen.

• Block Ranking: Het downsampled queries en keys, berekent benaderde scores, en past MaxPool toe op keys gevolgd door MaxNorm. Dit proces behoudt de "rowmax" (de maximale waarde in een rij) binnen blokken en zet scores om in relatieve ranking-bijdragen.
• HRM Flag Generatie: Het identificeert "Have Rowmax" (HRM) blokken—blokken die de maximale waarde bevatten voor een specifieke query-rij.
• Flag-Grouped Sorting: De KAP construeert een computation-order mask ($M_{KAP}$) die prioriteit geeft aan het eerst bezoeken van HRM-blokken, gevolgd door andere blokken die gesorteerd zijn op hun geschatte scores. Deze ordening zorgt ervoor dat de kernel de globale maximum vroegtijdig tijdens de traversatie tegenkomt.

2.3 Fase 2: Ordered-Skipping Kernel (OSK)

De OSK consumeert de $M_{KAP}$ mask om sparse attention uit te voeren met drie verschillende patronen:

1. Mask-Skipped: Blokken die niet door de KAP zijn geselecteerd, worden fysiek overgeslagen via KV page remapping. Dit elimineert de control-flow en synchronisatiekosten die geassocieerd worden met overgeslagen blokken in standaard sparse kernels.
2. In-Kernel Skipped: Voor de blokken die door Fase 1 zijn behouden, voert de OSK een tweede fase van skipping uit met behulp van exacte online-softmax statistieken. Omdat de HRM-blokken als eerste zijn bezocht (door de ordening van de KAP), stijgt de running maximum snel naar de globale maximum. Dit maakt het mogelijk om de skip-conditie (gebaseerd op het verschil tussen de lokale en de running max) agressiever en veiliger te voldoen, waardoor het "bucket effect" wordt vermeden.
3. Computed: Blokken die beide filters passeren, worden volledig berekend.

3. Belangrijkste Bijdragen

• Het Blootleggen van de Ontkoppeling: Het artikel identificeert dat de isolatie tussen proxy-ontwerp en kernel-executie de fundamentele barrière is voor het bereiken van hoge sparsity zonder verlies van nauwkeurigheid.
• KAP Ontwerp: Een nieuwe proxy die een computation-order mask afgeeft in plaats van een binair mask. Het selecteert blokken onder een gematigd budget en schrijft een executievolgorde voor die prioriteit geeft aan blokken die rowmax-waarden bevatten.
• OSK Ontwerp: Een geoptimaliseerde kernel die pagina's fysiek overslaat op basis van de KAP mask en agressieve in-kernel skipping uitvoert door gebruik te maken van de verbeterde running maximum statistieken verkregen door de herordende traversatie.
• Empirische Validatie: Uitgebreide experimenten tonen aan dat CoSA een hogere nauwkeurigheid bereikt bij lagere budgetten vergeleken met state-of-the-art methoden (MInference, FlexPrefill, XAttention) over diverse LLM-backbones (Qwen3-8B, Llama-3.1-8B-Instruct) en benchmarks.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op RULER (synthetische taken) en LongBench-v2 (real-world reasoning) met contextlengtes variërend van 4K tot 128K.

• Nauwkeurigheid: CoSA presteert consequent beter dan bestaande sparse methoden. Op LongBench-v2 met Qwen3-8B behaalde CoSA de hoogste gemiddelde nauwkeurigheid (36,51%) terwijl het opereerde bij het laagste budget (15%). Het behield een robuuste prestatie in zowel "w/o CoT" (direct) als "w/ CoT" (reasoning) modi.
• Efficiëntie:
  • Attention Speedup: CoSA bereikt een 4,93× versnelling in attention computation bij een contextlengte van 128K.
  • End-to-End Latency: Het vermindert de Time-to-First-Token (TTFT) met 2,53× bij 128K context met verwaarloosbare prestatievermindering vergeleken met dense attention.
• Ablatie Studies: Het artikel demonstreert dat de combinatie van KAP (ordered selection) en OSK (page remapping + in-kernel skip) synergetisch is. Het gebruik van alleen KAP verbetert de nauwkeurigheid; het toevoegen van in-kernel skipping verhoogt de sparsity; en het toevoegen van page remapping herstelt en verbetert de nauwkeurigheid terwijl het budget verder wordt aangescherpt.

5. Betekenis en Claims

Het artikel claimt dat CoSA een paradigmaverschuiving vertegenwoordigt van het behandelen van de proxy en de kernel als geïsoleerde componenten naar een co-designed systeem. Door hen te verbinden met een computation-order mask, lost CoSA de efficiency-fidelity afruil op die huidige sparse attention methoden teisteren.

De auteurs stellen dat:

1. Agressieve sparsity is haalbaar zonder nauwkeurigheidsverlies: Door HRM-blokken naar voren te halen, kan de kernel meer blokken veilig overslaan, waardoor het sparsity-budget lager kan worden gemaakt dan voorheen mogelijk was.
2. Fysieke executie ertoe doet: Het simpelweg hebben van een goed mask is onvoldoende; de kernel moet de geheugentoegang fysiek herordenen (via page remapping) om de theoretische versnellingen te realiseren en redundante synchronisatiekosten te vermijden.
3. Schaalbaarheid: CoSA is bijzonder effectief voor lange contexten (128K+), waar de kwadratische kosten van attention domineren, en biedt substantiële real-world latency reducties voor toepassingen zoals Retrieval-Augmented Generation (RAG) en autonome agenten.

Het artikel concludeert dat hoewel CoSA zich momenteel richt op de prefilling fase, toekomstig werk de principes van dit co-design zal uitbreiden naar de decoding fase, die andere computationele patronen vertoont.