În atenția standard, timpul și memoria cresc pătratic cu lungimea secvenței, pentru că se construiește matricea completă de scoruri N × N, scrisă și citită din memoria principală a GPU-ului (HBM).
FlashAttention (Dao, Fu, Ermon, Rudra, Ré, 2022) este un algoritm exact, „IO-aware”: împarte calculul în blocuri (tiling) care încap în memoria SRAM de pe cip și calculează softmax-ul incremental, fără să scrie matricea N × N în HBM. Memoria suplimentară devine liniară în lungimea secvenței.
FlashAttention-2 (Dao, 2023) împarte mai bine munca între blocurile de fire de execuție și este de aproximativ 2 ori mai rapid decât prima versiune. FlashAttention-3 (Shah et al., 2024) folosește asincronia Tensor Cores și TMA de pe GPU-urile Hopper și este de 1,5–2 ori mai rapid pe H100 în FP16. FlashAttention-4 este scris în CuTeDSL pentru Hopper și Blackwell.
Concepte cheie și caracteristici
| Concept / Parametru | Descriere practică |
|---|---|
Tiling |
Q, K și V sunt împărțite în blocuri procesate în SRAM. |
Softmax online |
Normalizarea softmax se calculează pas cu pas, pe blocuri, cu rezultat identic. |
Complexitate I/O |
Atenția standard: Θ(Nd + N²) accesări HBM; FlashAttention: Θ(N²d²/M), unde M este mărimea SRAM (Teorema 2 din lucrare). |
Hardware FA2 |
GPU-uri Ampere, Ada sau Hopper (A100, RTX 3090, RTX 4090, H100), CUDA 12.0+; pentru Turing (T4, RTX 2080) există repository-ul separat flash-attention-turing. |
Hardware FA3 / FA4 |
FA3: H100 / H800, CUDA 12.3+ (beta). FA4: Hopper și Blackwell (de exemplu H100, B200). |
Tipuri de date |
FA2: fp16 și bf16 (bf16 cere Ampere sau mai nou). FA3 adaugă FP8 la forward. |
Exemple practice și scenarii reale
Exemplul 1: FlashAttention-2 în Hugging Face Transformers
Argumentul attn_implementation din documentația Transformers; dtype înlocuiește vechiul torch_dtype.
import torch
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3.2-1B",
dtype=torch.bfloat16,
attn_implementation="flash_attention_2",
)Exemplul 2: Schimbarea implementării fără reîncărcarea modelului
model.set_attn_implementation("sdpa")Greșeli frecvente și capcane
- Crezi că FlashAttention aproximează atenția. Calculul este exact; câștigul vine din mai puține transferuri de memorie.
- Crezi că FlashAttention face atenția liniară ca timp. Memoria devine liniară, dar numărul de operații rămâne pătratic în lungimea secvenței.
- Încerci FA2 pe un GPU Turing sau în FP32. Proiectul principal cere Ampere, Ada sau Hopper și fp16/bf16.