L'idée clé (Dao et al., 2022) : l'attention standard matérialise la matrice complète N×N dans la HBM du GPU (high bandwidth memory), ce qui est à la fois gourmand en mémoire (quadratique par rapport à la longueur de la séquence) et lent (la bande passante HBM est le goulot d'étranglement). Flash Attention ne matérialise jamais cette matrice. Elle calcule plutôt l'attention par tuiles, chargeant de petits blocs de Q, K, V dans la SRAM rapide sur puce, calculant des résultats partiels et les accumulant — une technique appelée « tiling » ou « fusion de noyaux ».
L'attention standard stocke la matrice N×N, nécessitant O(N²) de mémoire. Pour un contexte de 128K avec 128 têtes d'attention, ça représente des centaines de gigaoctets. Flash Attention utilise O(N) de mémoire en calculant le softmax de manière incrémentale et sans jamais stocker la matrice complète. C'est ce qui a rendu les fenêtres de contexte de 128K à 1M faisables sur le matériel existant. FlashAttention-2 a encore amélioré le débit en parallélisant mieux à travers les blocs de threads GPU.
Flash Attention illustre un principe plus large : sur le matériel moderne, le goulot d'étranglement est souvent la bande passante mémoire, pas le calcul. Les GPU peuvent effectuer des trillions d'opérations par seconde mais ne peuvent lire/écrire que quelques centaines de gigaoctets par seconde depuis la HBM. Les algorithmes qui minimisent le trafic mémoire (même au prix de calculs supplémentaires) gagnent souvent. Cette approche « consciente des I/O » influence la façon dont tout le domaine pense la conception d'algorithmes pour les charges de travail IA.