FlashAttention (SV)

FlashAttention — är en revolutionerande algoritm för beräkning av uppmärksamhetsmekanismen (attention), utvecklad för att avsevärt påskynda träning och inferens av stora språkmodeller (LLM) med bibehållen full beräkningsnoggrannhet. Algoritmen presenterades första gången 2022 av ett forskarteam från Stanford University under ledning av Tri Dao.

Nyckelidén med FlashAttention är en omorganisering av beräkningarna med hänsyn till GPU:ns minneshierarki, vilket minimerar antalet åtkomster till det långsamma minnet och eliminerar den viktigaste flaskhalsen i den standardiserade uppmärksamhetsmekanismen.

Problematik med standard-attention

Den standardiserade självuppmärksamhetsmekanismen i transformers beräknas enligt formeln: $\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$ där Q, K, V är matriser för frågor, nycklar och värden.

Huvudproblemet med detta tillvägagångssätt är kvadratisk komplexitet avseende tid och minne (O(N²)) i förhållande till sekvensens längd N. Vid en naiv implementering måste hela uppmärksamhetsmatrisen S av storlek N×N beräknas och lagras i GPU:ns minne, vilket leder till två kritiska problem:

1. Stort minnesförbrukning: Lagring av en N×N-matris blir omöjlig vid arbete med långa kontexter.
2. In-/utdata-operationer (IO): Den viktigaste flaskhalsen är inte antalet aritmetiska operationer, utan de ständiga åtkomsterna till GPU:ns långsamma minne.

GPU:ns minneshierarki

För att förstå problemet är det viktigt att skilja på två typer av minne i GPU:n (med NVIDIA A100 som exempel):

• SRAM (statiskt minne): Snabbt on-chip-minne med liten kapacitet (~20 MB) och enorm bandbredd (upp till 19 TB/s).
• HBM (högbandbreddsminne): Långsamt minne med stor kapacitet (40–80 GB) och betydligt lägre bandbredd (cirka 1,5 TB/s).

Denna asymmetri gör standardalgoritmen för attention minnesbandbreddsbegränsad (memory-bound), eftersom den ständigt läser och skriver stora matriser från det långsamma HBM-minnet, vilket är den främsta källan till fördröjningar.

FlashAttentions viktigaste innovationer

FlashAttention är en IO-medveten (IO-aware) algoritm som löser problemet genom att minimera åtkomsterna till HBM. Detta uppnås med hjälp av tre grundläggande tekniker.

Tiling och blockbearbetning

Istället för att bearbeta hela matrisen på en gång delar FlashAttention upp ingångsmatriserna Q, K, V i mindre block (tiles), som får plats i det snabba SRAM-minnet. Algoritmen laddar successivt dessa block, utför alla uppmärksamhetsberäkningar för dem och uppdaterar slutresultatet, utan att spara hela uppmärksamhetsmatrisen i det långsamma HBM-minnet.

Online-beräkning av Softmax

Den viktigaste tekniska genombrotten var den "online"-baserade beräkningen av Softmax. Standard-Softmax kräver kännedom om alla element i indatavektorn för normalisering. FlashAttention använder en modifierad algoritm som möjliggör beräkning av Softmax i delar. Den håller reda på två mellanliggande värden (aktuellt maximum och summan av exponenter), som uppdateras allteftersom nya block bearbetas, vilket ger ett exakt resultat utan att hela matrisen behöver vara tillgänglig på en gång.

Sammanslagning av operationer till en enda CUDA-kärna

Alla uppmärksamhetsoperationer (matrismultiplikation QKᵀ, maskering, Softmax, multiplikation med V) kombineras till en enda sammanslagen CUDA-kärna (fused kernel). Detta minskar drastiskt antalet läs-/skrivoperationer till HBM: istället för upprepade genomgångar av hela matrisen laddar algoritmen ett block till SRAM en gång, utför alla beräkningar och skriver endast slutresultatet.

Teoretisk och praktisk effektivitet

Komplexitet och optimalitet

FlashAttention minskar minnesförbrukningen från O(N²) till O(N), vilket ger linjär skalning. Det har bevisats att algoritmens IO-komplexitet är teoretiskt optimal för beräkning av attention i en tvånivås minneshierarki, det vill säga det är omöjligt att utföra exakt attention snabbare utan att ändra hårdvaran.

Empiriska resultat

Den första versionen av FlashAttention visade på betydande förbättringar:

• Acceleration:
  • BERT-large (längd 512): 15% snabbare träning.
  • GPT-2 (längd 1K): 3-faldig acceleration.
  • Long-Range Arena-uppgifter (1K–4K): 2,4-faldig acceleration.
• Minnesbesparing: Upp till 20-faldig minnesbesparing jämfört med exakta referensimplementationer.
• Förbättrad modellkvalitet: Tack vare möjligheten att arbeta med längre kontexter förlorar FlashAttention inte bara inte i kvalitet, utan förbättrar den också. Till exempel förbättrades GPT-2:s perplexitet med 0,7 enheter, och noggrannheten i uppgifter för klassificering av långa dokument ökade med 6,4 enheter.

Utveckling och vidare forskning

Framgången med FlashAttention inledde en hel serie hårdvaruorienterade algoritmer.

FlashAttention-2 (2023)

Den andra versionen syftade till att utnyttja GPU:ns resurser mer fullständigt. I den ursprungliga FlashAttention uppgick effektiviteten på NVIDIA A100 till endast 25–40 % av maximalt möjligt. FlashAttention-2 introducerade förbättringar i beräkningsparallellisering, vilket möjliggjorde:

• En tvåfaldig acceleration jämfört med den första versionen.
• Ökad GPU-utnyttjning till 50–73 % av det teoretiska maximumet.
• Utökat stöd för attention-huvuden av storlek 256, samt för Multi-Query Attention (MQA)-arkitekturer.

FlashAttention-3 (2024)

Den tredje versionen optimerades specifikt för GPU-arkitekturen NVIDIA Hopper (H100). Den utnyttjar nya hårdvarufunktioner, såsom asynkronitet i Tensor Cores och stöd för FP8, vilket möjliggjorde:

• En ytterligare 1,5–2-faldig acceleration jämfört med FlashAttention-2.
• Prestanda upp till 740 TFLOPS på FP16 och nära 1,2 PFLOPS på FP8.

Specialiserade lösningar

Idéerna bakom FlashAttention vidareutvecklades i andra projekt:

• FlashInfer (2025): En anpassningsbar attention-motor, optimerad specifikt för inferens av LLM. Den fokuserar på effektiv hantering av KV-cache i strömningsgenerering.
• FlashMLA (2024): En implementering av attention med komprimering av kontextuell cache (latent attention), som möjliggör minnesbesparing vid mycket långa sekvenser med minimal informationsförlust.

Påverkan på industrin och ekosystemet

FlashAttention har blivit ett grundläggande genombrott och snabbt blivit industristandard för effektiv träning och inferens av LLM. Det har integrerats i viktiga bibliotek som PyTorch och Hugging Face och används i de flesta stora språkmodeller (LLaMA, MPT, Falcon, Claude m.fl.).

Närmast FlashAttention och dess efterföljande versioner har spelat en avgörande roll i utökningen av kontextfönster för språkmodeller: från 2–4 tusen token (GPT-3) till 128 tusen token (GPT-4) och till och med miljontals token i experimentella modeller. Algoritmen undanröjde ett av de viktigaste hindren för skalning av transformers och öppnade nya möjligheter för AI-applikationer, från analys av långa dokument till multimodal förståelse.

Länkar

• Officiellt FlashAttention-förråd på GitHub

Litteratur

• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision. arXiv:2407.08608.
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
• Hong, K. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
• Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
• Dege, P. et al. (2025). FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs. arXiv:2506.01969.
• Wang, G. et al. (2025). FlashMask: Efficient and Rich Mask Extension of FlashAttention. OpenReview wUtXB43Chi.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (OpenReview version). OpenReview H4DqfPSibmx.
• Gholami, A. et al. (2024). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. OpenReview pF2ukh7HxA.