Flash Attention

Flash Attention est un algorithme optimisé pour calculer le mécanisme d'attention au cœur des modèles transformers. Introduit par Tri Dao et al. de Stanford dans un article de 2022, il produit des résultats numériquement identiques à l'attention standard mais réorganise le calcul pour minimiser le mouvement des données entre la mémoire à large bande passante du GPU (HBM) et le SRAM rapide sur puce. Il est classé comme un algorithme conscient des E/S—sa conception est explicitement motivée par le goulot d'étranglement de la bande passante mémoire du matériel GPU moderne plutôt que par la réduction de la complexité arithmétique de l'opération d'attention elle-même.

L'attention standard matérialise la matrice d'attention complète N×N en HBM, où N est la longueur de la séquence, résultant en une utilisation mémoire O(N²). Flash Attention utilise une stratégie de pavage : il traite l'attention par blocs suffisamment petits pour tenir dans le SRAM, en appliquant l'algorithme softmax en ligne numériquement stable pour accumuler les résultats partiels de manière incrémentale. La matrice d'attention complète n'est jamais écrite en HBM. Flash Attention 2 (2023) a amélioré le parallélisme dans la dimension de séquence et réduit les opérations en virgule flottante non-matmul. Flash Attention 3 (2024) a introduit la spécialisation asynchrone de warp et le support de précision FP8 ciblant les GPUs NVIDIA H100 et Hopper, poussant le débit plus près des pics théoriques du matériel.

L'attention est le goulot d'étranglement dominant en mémoire et calcul dans l'entraînement et l'inférence des transformers à longueurs de séquence longues. Flash Attention réduit la mémoire de O(N²) à O(N) et fournit des accélérations en temps mural de 2–4x par rapport aux implémentations non optimisées, se traduisant directement par des coûts d'entraînement plus bas ou plus d'itérations expérimentales par budget de calcul. Ces gains se composent à des longueurs de séquence plus longues, faisant de Flash Attention une infrastructure essentielle pour les modèles pratiques à long contexte.

En 2026, Flash Attention est intégré nativement dans PyTorch via torch.nn.functional.scaled_dot_product_attention, Hugging Face Transformers, vLLM et pratiquement tous les principaux cadres d'entraînement et d'inférence LLM. C'est une dépendance d'entraînement pour LLaMA 3, Mistral et la plupart des autres modèles frontière. Ring attention, qui distribue le pavage Flash Attention sur plusieurs appareils via une communication all-to-all, étend ses avantages aux séquences qui dépassent la mémoire d'un seul GPU, permettant les fenêtres de contexte d'un million de tokens maintenant proposées commercialement.

Exemple

“Une équipe de recherche entraînant un transformer de 7 milliards de paramètres sur des séquences de 32 000 tokens passe de l'attention par défaut de PyTorch à Flash Attention 2, réduisant l'utilisation mémoire GPU par étape d'environ 60% et réduisant le temps d'entraînement par étape de 3x, permettant des tailles de lot plus grandes et divisant par deux le temps d'entraînement total en temps mural.