BLASST: Dynamic BLocked Attention Sparsity via Softmax Thresholding

Probleemstelling

De groeiende vraag naar inferentie met lange contexten in Large Language Models (LLM's) heeft de computationele en geheugenbottlenecks van het zelf-attentie-mechanisme (self-attention) verergerd.

• Quadratische Complexiteit: Voor een sequentie van lengte $n$ vereist de standaard attentie $O(n^2)$ bewerkingen en geheugentoegang. Dit maakt het verwerken van zeer lange sequenties (bijv. 128K tot 1M tokens) computationeel onhaalbaar, zelfs met state-of-the-art hardware.
• Beperkingen van Bestaande Oplossingen: Bestaande methoden voor sparse attention (zoals MInference, XAttention, of methoden die fine-tuning vereisen) hebben aanzienlijke nadelen:
  1. Ze vereisen vaak dure pre-computatie die de theoretische snelheidswinst tenietdoet.
  2. Ze vereisen model-finetuning of het trainen van nieuwe architecturen.
  3. Ze focussen vaak alleen op de 'prefill'-fase (het verwerken van de input) en niet op de 'decode'-fase (het genereren van tokens), of andersom.
  4. Ze missen geoptimaliseerde ondersteuning voor moderne GPU-architecturen (zoals NVIDIA Blackwell en Hopper).
  5. Ze zijn moeilijk te integreren in bestaande frameworks zonder ingrijpende wijzigingen.

Methodologie: BLASST

BLASST (BLocked Attention Sparsity via Softmax Thresholding) is een dynamische, trainingsvrije methode die irrelevante attentieblokken overslaat tijdens de inferentie.

Kerninzicht:
Tijdens de berekening van attentiescores in FlashAttention (specifiek bij de online softmax), kunnen blokken worden geïdentificeerd die een verwaarloosbare bijdrage leveren aan de uiteindelijke output.

• Dynamische Overslaan: Het algoritme onderhoudt een 'running maximum' ($m_i$) van de attentiescores over de blokken heen. Voor elk blok $j$ wordt het lokale maximum ($\tilde{m}_i$) berekend.
• Drempelbeslissing: Als het verschil tussen het lokale maximum en het running maximum kleiner is dan een bepaalde drempelwaarde ($\ln(\lambda)$), wordt aangenomen dat de exponentiële waarde van dit blok na normalisatie verwaarloosbaar klein is ($\approx 0$).
• Gevolg: Voor deze blokken worden drie dure operaties overgeslagen:
  1. Berekening van de exponentiële functie voor de softmax.
  2. Laden van het bijbehorende 'Value'-blok uit het HBM (High Bandwidth Memory).
  3. Matrixvermenigvuldiging tussen de attentiewaarden en de values.

Calibratie en Thresholding:

• De optimale drempelwaarde $\lambda$ hangt omgekeerd evenredig af van de contextlengte $L$ ($\lambda = a/L$).
• Een geautomatiseerde calibratieprocedure bepaalt de parameter $a$ zodat een vaste sparsiteitsratio (bijv. 70%) behouden blijft, ongeacht de lengte van de inputsequentie. Dit voorkomt dat een vaste drempel leidt tot onvoorspelbare prestaties bij verschillende contextlengtes.

Kernel-ontwerp:
BLASST is geïmplementeerd als geoptimaliseerde CUDA-kernels voor Blackwell en Hopper GPU's:

• Prefill-fase (Compute-bound): De kernels overslaan zowel de softmax-berekeningen als de tensor-core-matrixvermenigvuldigingen (MMA) voor gesprongen blokken.
• Decode-fase (Memory-bound): De focus ligt op het overslaan van het laden van de Value-matrix uit het geheugen, wat het geheugendebiet (bandwidth) bespaart. De pipeline is herontworpen om "pipeline bubbles" te voorkomen door batched loading van blokken die de drempeltest passeren.

Belangrijkste Bijdragen

1. BLASST-algoritme: Een "drop-in" methode die geen pre-computatie, geen proxy-scores en geen training vereist, maar toch aanzienlijke versnelling biedt.
2. Geautomatiseerde Calibratie: Een procedure die de relatie tussen drempelwaarde en contextlengte kwantificeert, waardoor betrouwbare implementatie mogelijk is zonder handmatige tuning.
3. Sparsity-Aware Training: Een extensie waarbij modellen worden getraind met BLASST actief, waardoor ze leren om informatie te concentreren in de blokken met de hoogste scores, wat de nauwkeurigheids-sparsiteitsgrens verder duwt.
4. Geoptimaliseerde Kernels: Implementaties voor zowel prefill als decode die compatibel zijn met MHA, GQA, MQA en MLA (Multi-head Latent Attention) en geïntegreerd zijn in frameworks zoals TensorRT-LLM en FlashInfer.

Resultaten

De prestaties zijn getest op moderne GPU's (NVIDIA B200 en H200) met modellen zoals Llama-3.1-8B en Qwen3-8B.

• Snelheidswinst:
  • Prefill: Tot 1.52x versnelling bij een sparsiteit van 71,9%.
  • Decode: Tot 1.48x versnelling bij een sparsiteit van 73,2%.
  • Bij lagere sparsiteit (~50%) is er al een versnelling van ~1.33x (prefill) en ~1.25x (decode) met minimaal verlies aan nauwkeurigheid.
• Nauwkeurigheid:
  • BLASST behoudt de benchmarknauwkeurigheid op lange contexttaken (RULER, LongBench) en redeneertaken (MATH500, AIME, GPQA).
  • Opvallend genoeg presteert BLASST op sommige taken (zoals MATH500) zelfs beter dan de dichte baseline, waarschijnlijk omdat het filteren van ruis (lage attentiescores) de concentratie op relevante tokens verbetert.
• Stabiliteit: De calibratiemethode zorgt ervoor dat de sparsiteit consistent blijft over verschillende contextlengtes (van 4K tot 64K+ tokens), terwijl een vaste drempelwaarde leidt tot grote variatie in de daadwerkelijke sparsiteit.
• Extremen: Zelfs bij zeer lange sequenties (200K tokens) en hoge sparsiteitsniveaus (tot 90%) blijft de degradatie in nauwkeurigheid beperkt en voorspelbaar.

Betekenis

BLASST is een doorbraak in de praktische implementatie van lange-context LLM's. Het lost de belangrijkste adoptiebarrières op van eerdere sparse attention-methoden:

• Het vereist geen training en geen dure voorbereiding.
• Het versnelt zowel de prefill- als de decode-fase, wat essentieel is voor totale inferentielatentie.
• Het is hardware-agnostisch in concept maar hardware-geoptimaliseerd in implementatie, specifiek voor de nieuwste GPU-generaties.
• Het is naadloos integreerbaar in bestaande frameworks.

Door de computationele complexiteit van de attentie-effectief te reduceren zonder de modelkwaliteit te offeren, maakt BLASST het haalbaar om LLM's in productie te zetten voor toepassingen die extreem lange contexten vereisen, zoals het analyseren van volledige codebases of uitgebreide documenten.