AsyncTLS: Efficient Generative LLM Inference with Asynchronous Two-level Sparse Attention

Probleemstelling

De inferentie van Large Language Models (LLMs) in lange contexten (honderdduizenden tokens) wordt beperkt door twee fundamentele uitdagingen:

1. Kwadratische complexiteit: Het zelf-attentie-mechanisme heeft een rekencomplexiteit van $O(n^2)$, wat onhoudbaar wordt bij zeer lange sequenties.
2. KV-cache geheugen: De Key-Value (KV) cache groeit lineair met de sequentielengte en vult vaak het beperkte, hoogwaardige GPU-geheugen (VRAM). Dit dwingt tot het verplaatsen (offloading) van data naar langzamere CPU-geheugenlagen, wat latency introduceert.

Bestaande oplossingen hebben een afweging nodig tussen nauwkeurigheid en efficiëntie:

• Token-niveau sparse attention (bijv. H2O, SnapKV) biedt hoge nauwkeurigheid door individuele belangrijke tokens te selecteren, maar de indexering (het vinden van deze tokens) is computatief duur en kostbaar.
• Block-niveau sparse attention (bijv. Quest) is efficiënter door blokken van tokens te selecteren, maar verliest precisie omdat het onrelevante tokens binnen een geselecteerd blok meeneemt en kritieke tokens in niet-geselecteerde blokken kan verliezen.

Methodologie: AsyncTLS

AsyncTLS is een hiërarchisch systeem dat de voordelen van beide benaderingen combineert en een asynchroon offloading-mechanisme introduceert. Het bestaat uit twee hoofdcomponenten:

1. Hiërarchische Twee-Niveau Sparse Attention

In plaats van direct naar individuele tokens te zoeken (duur) of hele blokken te kiezen (onprecies), gebruikt AsyncTLS een gefaseerde aanpak:

• Stap 1: Grofkorrelige Blokkering (Block-Level Filtering):
  De KV-cache wordt opgedeeld in blokken. Eerst wordt een snelle, grofkorrelige score berekend voor elk blok om blokken te identificeren die waarschijnlijk relevante tokens bevatten. Dit verkleint de zoekruimte drastisch. De auteurs hebben de berekening van deze scores geoptimaliseerd naar standaard GEMM-bewerkingen (matrixvermenigvuldiging) om compatibel te zijn met moderne GPU-architecturen (zoals MQA en MLA), wat de rekenkosten verlaagt.
• Stap 2: Fijnkorrelige Token-selectie (Token-Level Selection):
  Binnen de geselecteerde blokken wordt vervolgens een fijnkorrelige selectie uitgevoerd om de meest cruciale individuele tokens te kiezen. Hierbij wordt gebruikgemaakt van technieken zoals "Double Sparsity" (kanalen selecteren) en kwantisatie om de indexering te versnellen.
• Resultaat: Dit behoudt de nauwkeurigheid van token-niveau selectie, maar vermindert de indexeringsoverhead aanzienlijk door de zoekruimte eerst te beperken tot relevante blokken.

2. Asynchroon Offloading Engine

Om het geheugenprobleem op te lossen zonder de snelheid te verliezen door wachttijden bij data-overdracht:

• Asynchroon Prefetching: Het systeem maakt gebruik van de temporele localiteit van attentiepatronen (de selectie van het volgende tijdstip lijkt sterk op die van het huidige tijdstip).
  • Terwijl de GPU de attentie berekent voor het huidige tijdstip (gebaseerd op de blokken van de vorige stap), wordt er asynchroon prefetching uitgevoerd voor de blokken die nodig zijn voor de volgende stap.
  • Hierdoor overlappen de geheugenoverdrachten (CPU $\leftrightarrow$ GPU) met de berekeningen, waardoor de latency van PCIe-overdrachten wordt verborgen.
• Incrementele Overdracht: In plaats van hele blokken te verplaatsen bij elke stap, wordt alleen het verschil (de delta) tussen de huidige en de vorige selectie overgedragen. Dit minimaliseert de bandbreedtevereisten op de PCIe-verbinding.

Belangrijkste Bijdragen

1. Hiërarchische Architectuur: Een nieuw tweelaags selectiemechanisme dat de precisie van token-niveau sparsiteit combineert met de efficiëntie van block-niveau filtering.
2. AsyncTLS Offloading Engine: Een innovatieve engine die token-niveau sparsiteit toepasbaar maakt in offloading-scenario's door asynchroon prefetching en incrementele data-overdracht te gebruiken.
3. Brede Architecturale Compatibiliteit: De methode is getest en werkt effectief op diverse architecturen, waaronder Multi-Head Attention (MHA), Grouped-Query Attention (GQA) en de nieuwere Multi-Head Latent Attention (MLA) (zoals gebruikt in DeepSeek-modellen).

Resultaten

De auteurs hebben AsyncTLS geëvalueerd op state-of-the-art modellen (Qwen3-8B/14B en GLM-4.7-Flash) met contextlengtes tot 128k tokens.

• Nauwkeurigheid: AsyncTLS bereikt nauwkeurigheid die bijna niet te onderscheiden is van volledige attentie (Full Attention) en presteert significant beter dan bestaande block-niveau methoden (zoals Quest) en vergelijkbaar met token-niveau methoden (zoals Double-Sparsity) op benchmarks zoals LongBench en RULER.
• Snelheid (Operator Level):
  • GQA: 1,2x tot 4,0x sneller dan token-niveau methoden en 1,7x tot 6,2x sneller dan volledige attentie.
  • MLA: Tot wel 10x sneller dan volledige attentie.
• End-to-End Throughput:
  • Door het gebruik van offloading kan AsyncTLS grotere batchgroottes verwerken (bijv. batch size 6) waar volledige attentie vastloopt (batch size 1) vanwege geheugenbeperkingen.
  • Bij sequentielengtes van 96k tokens wordt een 1,3x tot 4,7x verbetering in end-to-end throughput behaald ten opzichte van volledige attentie.

Betekenis

AsyncTLS lost de fundamentele tegenstelling op tussen nauwkeurigheid en efficiëntie bij lange context-inferentie. Het bewijst dat token-niveau sparsiteit, die historisch gezien te duur was voor praktische implementatie door indexeringsoverhead, nu haalbaar is door middel van hiërarchische filtering en slimme geheugenbeheerstrategieën. Dit maakt het mogelijk om ultra-lange sequenties (96k-128k tokens) kostenefficiënt en snel te verwerken op bestaande hardware, zonder dat er extra training van het model nodig is (training-free). Het biedt een schaalbare oplossing voor de volgende generatie LLM-toepassingen die diepe context vereisen.