Flash Attention

Flash Attention (FlashAttention) è un algoritmo e una famiglia di kernel GPU che calcolano l'attenzione esatta per i modelli Transformer senza mai scrivere l'intera matrice di attenzione nella memoria principale della GPU. La self-attention standard su N token costruisce una matrice di punteggi N × N, la scrive nella memoria ad ampia larghezza di banda (HBM) e la rilegge per calcolare la softmax e il prodotto con i valori; perciò la maggior parte del tempo di esecuzione dell'attenzione è dedicata allo spostamento dei dati. Tri Dao e i suoi colleghi di Stanford hanno presentato il metodo nell'articolo del 2022 FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Suddivide il lavoro in tile che entrano nella veloce SRAM on-chip, così l'uso della memoria cresce linearmente con la lunghezza della sequenza e l'output corrisponde a quello dell'attenzione standard, a meno degli arrotondamenti in virgola mobile. I suoi kernel vengono eseguiti in PyTorch e nei motori di inferenza dei modelli linguistici di grandi dimensioni (LLM), come vLLM.

Come funziona Flash Attention#

Una GPU ha una piccola quantità di memoria molto veloce accanto alle unità di calcolo e una memoria più grande e lenta all'esterno del chip. L'articolo su FlashAttention descrive una NVIDIA A100 con 40-80 GB di HBM a 1.5-2.0 TB/s e 192 KB di SRAM in ciascuno dei suoi 108 multiprocessori streaming, con una velocità stimata di 19 TB/s. L'attenzione esegue poche operazioni per ogni byte letto, quindi la sua velocità dipende dalla frequenza con cui i dati passano da un livello all'altro; il metodo conteggia questi trasferimenti, da cui il termine «consapevole dell'I/O».

FlashAttention carica un tile di query e poi trasmette in streaming tile di chiavi e valori attraverso la SRAM. Per ogni tile calcola i punteggi, aggiorna il massimo e la somma progressivi per ogni riga di query, riscalando l'output parziale ogni volta che il massimo aumenta e aggiungendo il contributo del nuovo tile. Questa softmax online produce il risultato esatto della softmax senza dover esaminare una riga intera in una volta; solo l'output finale e due numeri per riga vengono riscritti in HBM. Il passaggio backward ricalcola i tile dei punteggi a partire da Q, K e V invece di memorizzare la matrice N × N: aggiunge operazioni, ma è comunque più veloce perché legge molta meno memoria. Lo stesso ciclo in PyTorch, che elabora tutte le query rispetto a un tile chiave-valore alla volta:

import torch

torch.manual_seed(0)
n, d, block = 4096, 64, 256
q, k, v = torch.randn(3, n, d).unbind(0)

# Standard attention builds the full n x n score matrix in memory
reference = torch.softmax(q @ k.T / d**0.5, dim=-1) @ v

# FlashAttention-style: stream key/value tiles, keeping a running max and sum per query (online softmax)
row_max, row_sum, out = torch.full((n, 1), -torch.inf), torch.zeros(n, 1), torch.zeros(n, d)
for start in range(0, n, block):
    scores = q @ k[start : start + block].T / d**0.5  # only an n x block tile exists at once
    new_max = torch.maximum(row_max, scores.max(dim=-1, keepdim=True).values)
    rescale, probs = torch.exp(row_max - new_max), torch.exp(scores - new_max)
    row_sum = row_sum * rescale + probs.sum(dim=-1, keepdim=True)
    out = out * rescale + probs @ v[start : start + block]
    row_max = new_max
out = out / row_sum

print(f"Largest score tile: {tuple(scores.shape)} vs full matrix: ({n}, {n})")
print(f"Max difference from standard attention: {(out - reference).abs().max().item():.1e}")

Il tile dei punteggi più grande misura 4,096 × 256, un sedicesimo della matrice 4,096 × 4,096, e il risultato differisce dall'attenzione standard al massimo di 4.3e-07, un valore pari all'errore di arrotondamento FP32. Un kernel reale suddivide anche le query in tile, uno per ogni blocco di thread della GPU, e fonde l'intero ciclo in un unico kernel, così i tile non lasciano mai la SRAM.

FlashAttention-2, -3 e -4#

Ogni versione è pensata per una generazione più recente di GPU. FlashAttention-2 (2023) ha riorganizzato la suddivisione del lavoro tra i thread ed è circa due volte più veloce dell'originale, raggiungendo il 50-73% del throughput massimo di una A100. FlashAttention-3 (2024) usa i Tensor Core asincroni di H100 e supporta FP8, raggiungendo 740 TFLOPs/s in FP16 e quasi 1.2 PFLOPs/s in FP8. FlashAttention-4 (2026) è progettato per le GPU Blackwell e dichiara fino a 1,613 TFLOPs/s in BF16 su B200, il 71% del picco. I kernel sono distribuiti nella libreria open source flash-attention, compatibile con input FP16 e BF16; PyTorch seleziona automaticamente un kernel FlashAttention-2 tramite scaled_dot_product_attention quando la GPU e gli input lo supportano, ricorrendo altrimenti a implementazioni efficienti in termini di memoria o standard. SAM 2 di Meta, incluso nel pacchetto Ultralytics, calcola l'attenzione nel suo encoder di immagini Hiera e nella sua attenzione della memoria tramite questa funzione, perciò PyTorch instrada questi livelli a FlashAttention ogni volta che l'hardware e la precisione degli input lo consentono.

Flash Attention e termini correlati#

• Attenzione standard: la stessa operazione matematica, calcolata scrivendo in HBM le matrici complete dei punteggi e delle probabilità.
• Attenzione sparsa: calcola i punteggi solo per coppie selezionate di query e chiavi. L'articolo su FlashAttention ha introdotto anche una variante sparsa a blocchi che salta i tile vuoti.
• Attenzione lineare: sostituisce la softmax con una mappa di feature, così il costo cresce linearmente con la lunghezza, modificando però gli output del modello.
• PagedAttention: gestore della memoria che archivia la cache KV in pagine di dimensione fissa durante il serving. I kernel di attenzione, incluso FlashAttention, leggono chiavi e valori da queste pagine.
• Attenzione ad anello: suddivide una sequenza molto lunga tra più dispositivi; ciascuno calcola l'attenzione a blocchi sulla propria partizione mentre i blocchi chiave-valore passano lungo un anello.

Impatto sull'addestramento e sui contesti lunghi#

Poiché la memoria cresce linearmente con la lunghezza della sequenza, FlashAttention ha reso pratico addestrare con finestre di contesto lunghe. L'articolo originale ha riportato un'accelerazione complessiva del 15% per BERT-large con sequenze di lunghezza 512 rispetto al record di addestramento MLPerf 1.1, un'accelerazione di 3× per GPT-2 con sequenze di lunghezza 1,024 e di 2.4× sul benchmark Long Range Arena. Gli stessi risparmi si applicano alla visione, dove un Vision Transformer (ViT) o un Transformer di diffusione converte immagini e video ad alta risoluzione in lunghe sequenze di token. Le sequenze più lunghe hanno migliorato anche la precisione: l'articolo ha presentato i primi Transformer capaci di superare il caso nel test Path-X, che chiede se due punti in un'immagine 128 × 128, appiattita in una sequenza di 16,384 pixel, siano collegati da un percorso.