Flash Attention

Définition

Flash Attention est un algorithme d'optimisation du mécanisme d'attention des transformers, conçu pour exploiter l'architecture mémoire hiérarchique des GPU, réduisant drastiquement la mémoire nécessaire au calcul de l'attention sans modifier son résultat exact, contrairement aux approches d'approximation sacrifiant la précision. Le mécanisme d'attention standard nécessite de matérialiser explicitement en mémoire GPU une matrice de dimension quadratique par rapport à la longueur de la séquence traitée, opération dont la complexité mémoire croît en O(n²), ce qui limite structurellement la longueur de contexte traitable et engendre de nombreux transferts coûteux entre la mémoire haute capacité mais lente (HBM) et la mémoire ultra-rapide mais limitée en taille (SRAM) du GPU. Flash Attention résout ce goulot d'étranglement en appliquant deux techniques combinées : le tiling, découpant le calcul en blocs suffisamment petits pour tenir entièrement dans la mémoire rapide SRAM du GPU, évitant ainsi la matérialisation complète de la matrice d'attention en mémoire lente ; et la recomputation, recalculant certaines valeurs intermédiaires à la volée durant la rétropropagation plutôt que de les stocker, réduisant encore l'empreinte mémoire au prix d'un calcul additionnel largement compensé par le gain d'accès mémoire. Cette optimisation, publiée en 2022 par des chercheurs de Stanford, accélère l'entraînement et l'inférence des grands modèles de langage tout en permettant de traiter des contextes plus longs à ressources matérielles constantes.