FlashAttention ontleed: snellere berekening via GPU-geheugen

De klassieke implementatie van het transformer-aandachtsmechanisme kent een berucht schaalprobleem: de benodigde hoeveelheid geheugen groeit kwadratisch met de lengte van de context. Lange tijd werd aangenomen dat deze vertraging puur werd veroorzaakt door het aantal rekenkundige bewerkingen (FLOPs). In de praktijk blijkt de echte remmende factor echter niet de rekenkracht van de GPU-rekenkernen te zijn, maar de snelheid waarmee data van en naar het hoofdgeheugen wordt verplaatst. Dit hardwarematige knelpunt wordt geheugenbandbreedte genoemd.

FlashAttention, oorspronkelijk ontwikkeld door Tri Dao en collega's, herstructureert de berekening van self-attention fundamenteel. In plaats van tussenresultaten continu weg te schrijven naar het trage hoofdgeheugen van de GPU, maakt het algoritme optimaal gebruik van de kleine, ultrasnelle SRAM-chips die zich direct naast de rekenkernen bevinden. In dit artikel ontleden we stap voor stap hoe FlashAttention functioneert, welke wiskundige technieken dit mogelijk maken, en wat de praktische implicaties zijn voor trainings- en inferentietijd.

Het fysieke knelpunt: GPU-geheugenhiërarchie ontleed

Om te begrijpen waarom standaard attention traag is bij lange reeksen, moeten we kijken naar de fysieke architectuur van moderne videokaarten zoals de NVIDIA A100 of H100. Een moderne GPU bevat grofweg twee typen geheugen: het grote High Bandwidth Memory (HBM, ook wel het VRAM genoemd) en het interne Static RAM (SRAM, verdeeld over de Streaming Multiprocessors als Shared Memory en L1-cache).

HBM is ruim bemeten (bijvoorbeeld 80 GB op een A100), maar de doorvoersnelheid ligt rond de 1,5 tot 2,0 TB/s. Dat klinkt enorm, maar de rekenkernen (Tensor Cores) kunnen opgeteld honderden teraflops aan data verwerken. Het interne SRAM is vele malen kleiner (ongeveer 192 KB per Streaming Multiprocessor, wat over de hele chip neerkomt op enkele tientallen megabytes), maar haalt een bandbreedte van meer dan 19 TB/s. Wanneer een GPU-kernel data ophaalt uit HBM, staan de rekeneenheden feitelijk te wachten op data. We noemen zo'n bewerking memory-bound.

Standaard attention voert opeenvolgende bewerkingen uit: eerst vermenigvuldigt het de Query-matrix $Q$ met de getransponeerde Key-matrix $K^T$, schrijft het resultaat $S$ (de $N \times N$ aandachtsmatrix) weg naar HBM, leest deze matrix opnieuw in om de softmax-functie toe te passen tot matrix $P$, schrijft $P$ weer weg naar HBM, en leest $P$ ten slotte opnieuw in om deze te vermenigvuldigen met de Value-matrix $V$. Bij een context van 8.192 tokens telt de $N \times N$-matrix ruim 67 miljoen floating-point getallen per attention head. Dit leidt tot gigantische hoeveelheden onnodig lees- en schrijfverkeer over de tragere HBM-bus.

De wiskundige doorbraak: Tiling en Online Softmax

De centrale innovatie van FlashAttention is tiling: het opknippen van de grote $Q$-, $K$- en $V$-matrices in kleinere blokken die precies passen binnen het snelle SRAM van de GPU. Het probleem hierbij is echter de softmax-functie. Softmax normaliseert waarden over een volledige rij met de formule:

$$\text{softmax}(x)_i = \frac{e^{x_i}}{\sum_{j=1}^N e^{x_j}}$$

Om de noemer (de som van exponenten) te berekenen, lijkt het alsof we eerst de complete rij moeten kennen. Als we een rij in stukjes hakken, ontbreekt immers de globale som. FlashAttention lost dit op door gebruik te maken van het principe van online softmax. Hierbij worden tussenresultaten per blok berekend en dynamisch geschaald zodra een nieuw blok wordt verwerkt.

Stel dat we een deelsom hebben berekend over het eerste blok tokens met een lokaal maximum $m_1$ en een lokale noemer $l_1$. Wanneer we het tweede blok inladen met een lokaal maximum $m_2$, bepalen we het nieuwe globale maximum $m_{\text{nieuw}} = \max(m_1, m_2)$. We kunnen de oude som en de geaccumuleerde tussenvector simpelweg herschalen met de correctiefactor $e^{m_1 - m_{\text{nieuw}}}$. Hierdoor is het nooit nodig om de volledige $N \times N$-aandachtsmatrix in het globale HBM-geheugen op te slaan.

FlashAttention-1 versus FlashAttention-2 en latere iteraties

FlashAttention heeft sinds de eerste publicatie aanzienlijke iteratieve verbeteringen doorgemaakt. In de eerste versie (FlashAttention-1) werd de buitenste lus over de sleutel- en waardeblokken ($K$ en $V$) gelegd en de binnenste lus over de queryblokken ($Q$). Dit zorgde voor extra synchronisatiestappen tussen verschillende GPU-threads.

FlashAttention-2 draaide deze logica om: de buitenste lus itereert nu over de rijen van de $Q$-matrix. Omdat verschillende thread-blocks nu onafhankelijk aan verschillende rijen kunnen werken, is er veel minder communicatie en synchronisatie tussen de Streaming Multiprocessors nodig. Daarnaast werd het aantal niet-gemoduleerde matrixvermenigvuldigingen geminimaliseerd, waardoor de bezettingsgraad van de Tensor Cores steeg van circa 35% naar meer dan 70% van het theoretische maximum.

FlashAttention-3 richt zich specifiek op modernere hardware zoals NVIDIA Hopper-chips. Deze chips introduceren Tensor Memory Accelerator (TMA) en asynchrone instructies, waarmee data rechtstreeks van HBM naar shared memory kan worden gekopieerd zonder tussenkomst van de reguliere rekenregisters. Hierdoor kunnen dataoverdracht en matrixberekeningen tegelijkertijd plaatsvinden zonder vertraging.

Geheugencomplexiteit: van O(N²) naar O(N)

In een standaard transformer-architectuur vergt de opslag van de activaties tijdens de forward pass voor aandachtsmechanismen $O(N^2)$ geheugenruimte, waarbij $N$ de contextlengte in tokens is. Bij het trainen van een model moeten deze activaties bewaard blijven in het HBM om tijdens de backward pass de gradiënten te kunnen uitrekenen.

FlashAttention brengt deze geheugenvoetafdruk terug naar $O(N)$. Omdat de $N \times N$-aandachtsmatrix nooit in zijn geheel wordt geconstrueerd in HBM, hoeft deze ook niet te worden opgeslagen voor de backward pass. In plaats daarvan berekent FlashAttention de benodigde aandachtsblokken tijdens de backward pass razendsnel opnieuw (recomputation) vanuit de opgeslagen vectoren $Q$, $K$ en $V$ in combinatie met de bewaarde softmax-statistieken (het maximum $m$ en de noemer $l$).

Hoewel herberekening rekenkundig extra bewerkingen kost, is het in de praktijk veel sneller dan het inladen van gigantische matrices uit HBM. Het elimineren van de geheugendoorvoer compenseert ruimschoots de extra rekenstappen.

Interactie met de KV-cache tijdens inferentie

Tijdens de inferentiefase, wanneer een model token voor token antwoord genereert, verandert de rol van het aandachtsmechanisme. Hier hoeven geen gradiënten bewaard te worden, maar moeten eerdere tokens snel geraadpleegd worden. Om te begrijpen hoe eerdere invoer bewaard blijft, bekijken we de opbouw van KV-caching in transformer-architecturen, waar de allocatie van sleutels en waarden per laag wordt behandeld.

Bij het genereren van een nieuw token heeft de Query-vector een lengte van 1 ($N_q = 1$), terwijl de Key- en Value-matrices groeien tot de totale historische lengte van het gesprek ($N_{kv}$). In deze situatie is sprake van FlashDecoding, een gespecialiseerde variant van FlashAttention voor autoregressieve inferentie. FlashDecoding splitst de $N_{kv}$-sequentie over meerdere GPU-kernen op om ook bij $N_q=1$ een maximale parallelle bezetting te behouden.

FlashAttention werkt naadloos samen met moderne architecturen die de geheugendruk van de KV-cache verkleinen. Zie voor een diepere analyse van deze geheugenbesparende techniek ook het artikel over grouped-query attention en geheugengebruik, waarin meerdere query-heads dezelfde key- en value-heads delen.

Algoritmische kern in pseudocode

Onderstaande vereenvoudigde representatie toont het kernconcept van de FlashAttention forward pass met bloksgewijze tiling en online herschaling van de softmax-waarden.

# Invoer: Q, K, V matrices in HBM (formaat N x d)
# Blokgroottes: B_r (rijen van Q), B_c (kolommen van K, V)
# Uitvoer: O matrix in HBM (formaat N x d)

initialiseer O in HBM met nullen
initialiseer l = [0] * N  (som van exponenten)
initialiseer m = [-oneindig] * N  (maxima per rij)

splits Q in blokken Q_1, ..., Q_Tr van grootte B_r x d
splits K in blokken K_1, ..., K_Tc van grootte B_c x d
splits V in blokken V_1, ..., V_Tc van grootte B_c x d

voor elk blok Q_i in SRAM:
  voor elk blok K_j, V_j in SRAM:
    # 1. Bereken lokale matrixvermenigvuldiging
    S_ij = (Q_i * K_j^T) / sqrt(d)
    
    # 2. Bereken lokaal rij-maximum
    m_lokaal = max_per_rij(S_ij)
    m_nieuw = max(m_i, m_lokaal)
    
    # 3. Bereken herschaalde exponenten
    P_ij = exp(S_ij - m_nieuw)
    l_nieuw = exp(m_i - m_nieuw) * l_i + som_per_rij(P_ij)
    
    # 4. Werk de uitvoer O_i bij met correctiefactoren
    O_i = diag(exp(m_i - m_nieuw)) * O_i + P_ij * V_j
    
    # 5. Werk statistieken bij voor volgende iteratie
    m_i = m_nieuw
    l_i = l_nieuw

  # Normaliseer uiteindelijke rij met de totale noemer l_i
  O_i = diag(1 / l_i) * O_i
  schrijf O_i weg naar HBM

Praktische integratie en hardware-eisen

FlashAttention is een low-level C++/CUDA-implementatie en is niet standaard beschikbaar op alle architecturen. De implementatie vereist specifieke hardware-instructies die alleen aanwezig zijn op modernere GPU's.

Op NVIDIA-hardware wordt FlashAttention-2 ondersteund vanaf de Turing- en Ampere-architecturen (bijvoorbeeld RTX 30-serie, RTX 40-serie, A10, A100 en nieuwer met Compute Capability 8.0 of hoger). Oudere architecturen zoals Pascal (GTX 1080) missen de instructiesets om shared memory op deze manier asynchroon aan te sturen. Bij het inrichten van een lokaal systeem of een clusterserver is het cruciaal om vooraf te controleren of de hardware en contextlengte binnen de grenzen van de videokaart passen. Raadpleeg de VRAM-rekenhulp: hoeveel geheugen heeft een model nodig? om precies te berekenen hoe modelparameters en contextlengtes zich verhouden tot het beschikbare GPU-geheugen.

Binnen moderne deep learning frameworks zoals PyTorch is FlashAttention vaak al geïntegreerd via torch.nn.functional.scaled_dot_product_attention (SDPA). PyTorch kiest hierbij automatisch voor de FlashAttention-backend wanneer de invoertensors voldoen aan de juiste vereisten (zoals datatypes FP16 of BF16 en geschikte tensorvormen).

Beperkingen, zwakke punten en randgevallen

Hoewel FlashAttention aanzienlijke voordelen biedt, zijn er duidelijke technische beperkingen en trade-offs:

Ten eerste biedt FlashAttention weinig tot geen snelheidswinst bij zeer korte reeksen (bijvoorbeeld minder dan 256 tokens). Bij dergelijke sequenties is de $N \times N$-matrix dermate klein dat deze al volledig in de L2-cache of SRAM past, waardoor de overhead van het opdelen in blokken niet opweegt tegen een standaard matrixoperatie.

Ten tweede is FlashAttention hardware-specifiek. Omdat de kernel nauwkeurig is afgestemd op de exacte registergroottes, shared memory-verdelingen en cache-structuren van specifieke GPU-generaties, vraagt ondersteuning voor andere hardware (zoals AMD ROCm of Apple Silicon) om compleet herschreven kernels. Hoewel projecten zoals FlashAttention voor ROCm bestaan, lopen deze implementaties vaak achter op de officiële NVIDIA-versies.

Ten derde kan er sprake zijn van lichte numerieke afrondingsverschillen. Hoewel FlashAttention wiskundig exact dezelfde formule berekent als standaard attention (in tegenstelling tot benaderingsmethoden zoals sparse attention of low-rank projecties), kan de gewijzigde volgorde van drijvende-kommaberekeningen door afronding in FP16 subtiel afwijkende uitkomsten opleveren. Dit leidt in de praktijk niet tot kwaliteitsverlies, maar zorgt er wel voor dat testresultaten tussen verschillende backends niet bit-voor-bit identiek zijn.