FlashAttention-3: 2.6x speedup su FA2 ottimizzando per H100 Hopper con wgmma, TMA e FP8
In una frase Tri Dao e NVIDIA pubblicano FlashAttention-3: ottimizzato per H100 Hopper con overlapping compute/memory tramite wgmma e TMA, supporto FP8 low-precision, 2.6x speedup su FA2 e 75% del picco H100.
Ogni volta che NVIDIA lancia una nuova generazione di GPU, le vecchie implementazioni non sfruttano le nuove funzionalità hardware. H100 (architettura Hopper) ha introdotto istruzioni specializzate completamente nuove — wgmma per le moltiplicazioni di matrice e TMA per il trasferimento asincrono dei dati — che FA2 non usava affatto.
FlashAttention-3 è una riscrittura da zero di FlashAttention per sfruttare al massimo H100. Il trucco principale: sovrapporre le operazioni di calcolo e di memoria invece di eseguirle in sequenza. Mentre H100 sta facendo le moltiplicazioni di matrici per un blocco, FA3 carica già in SRAM i dati del blocco successivo.
Il risultato: 2.6x più veloce di FA2 su H100, raggiungendo il 75% del picco teorico della GPU. Supporta anche FP8, il nuovo formato a bassa precisione di H100 che dimezza la memoria e quasi raddoppia il throughput rispetto a FP16.
Aziende
Tri Dao Research, NVIDIA
Tool
FlashAttention-3, CUDA, PyTorch, cuDNN
Tag
Fonti