FlashAttention

FlashAttention is een techniek die AI-taalmodellen sneller en zuiniger met geheugen maakt, zonder hun antwoorden te veranderen. Het herschikt één zware berekening, attention, zodat de videokaart veel minder vaak gegevens uit zijn trage hoofdgeheugen hoeft te halen. Daardoor passen langere documenten op dezelfde hardware en gaat trainen sneller. Het zit ingebouwd in PyTorch, een veelgebruikt framework voor AI-modellen.

Welk probleem lost het op?

Attention vergelijkt elk token in een tekst met elk ander token, dus tijd en geheugen groeien met het kwadraat van de tekstlengte. Een standaardimplementatie schrijft de hele tabel met die vergelijkingen naar het hoofdgeheugen van de GPU (HBM) en leest hem daarna weer in. Op een Nvidia A100 bevat dat geheugen 40 tot 80 GB met 1,5 tot 2,0 TB/s; het SRAM op de chip bevat maar 192 KB per processorblok, maar is ongeveer een orde van grootte sneller (Dao e.a., 2022). Volgens het paper worden de meeste bewerkingen in een transformer opgehouden door geheugentoegang.

Hoe werkt het?

FlashAttention, van Tri Dao en collega's (mei 2022), is IO-aware: het behandelt het lezen en schrijven tussen de geheugenlagen van de GPU als de echte kostenpost.

• Tiling: het verdeelt de queries, keys en values in blokken die in het SRAM passen, berekent attention blok voor blok en bouwt de softmax stap voor stap op, zodat de volledige tabel nooit naar het hoofdgeheugen gaat.
• Herberekening: voor training bewaart het alleen kleine normalisatiewaarden en berekent het de attention-blokken opnieuw tijdens de backward pass, wat sneller is dan een opgeslagen tabel inlezen.
• Eén kernel: alle attention-stappen draaien in één GPU-programma.

Het resultaat is exacte attention, met extra geheugen dat lineair groeit met de tekstlengte. Het aantal berekeningen groeit nog steeds met het kwadraat. In het paper trainde FlashAttention GPT-2 tot drie keer zo snel als de implementatie van Hugging Face en 1,8 keer zo snel als Megatron, BERT-large 15% sneller dan het record van MLPerf 1.1, en gebruikte het tot 20 keer minder geheugen dan de exacte attention-implementaties waarmee het werd vergeleken.

Welke versies zijn er?

• FlashAttention-2 (juli 2023) verdeelt het werk beter over de GPU: ongeveer twee keer zo snel, op 50 tot 73% van het theoretische maximum van de A100. PyTorch 2.2 (januari 2024) bouwde het in zijn standaardfunctie voor attention in.
• FlashAttention-3 (juli 2024) is gemaakt voor de H100 van Nvidia: 1,5 tot 2 keer zo snel als versie 2 in de forward pass, tot 740 TFLOPs/s in FP16.
• FlashAttention-4 (maart 2026) is gemaakt voor Blackwell-GPU's zoals de B200: tot 1.613 TFLOPs/s in BF16.

De repository biedt FlashAttention-2 ook voor GPU's van AMD, via ROCm. Fine-tuningtools zoals LLaMA-Factory kunnen het aanzetten met een instelling.

Voorbeeld

In maart 2026 beschreven de makers van de vLLM Semantic Router, die verzoeken indeelt voordat ze een taalmodel bereiken, het probleem op een AMD MI300X (Liu e.a.). Met standaard-attention had één classifier ongeveer 1,5 GB nodig voor de attention-tabel bij 8K tokens, ongeveer 6 GB bij 16K en ongeveer 24 GB bij 32K: elke verdubbeling van de tekst verviervoudigde het geheugen. Naast het taalmodel op dezelfde GPU paste dat niet. Een FlashAttention-operator liet het attention-geheugen lineair groeien en maakte teksten van 8K tot 32K tokens mogelijk.