FlashAttention (FR)

FlashAttention est un algorithme révolutionnaire pour le calcul du mécanisme d'attention (attention), conçu pour accélérer de manière significative l'entraînement et l'inférence des grands modèles de langage (LLM) tout en conservant une précision de calcul exacte. L'algorithme a été présenté pour la première fois en 2022 par une équipe de chercheurs de l'Université de Stanford, dirigée par Tri Dao.

L'idée clé de FlashAttention consiste à réorganiser les calculs en tenant compte de la hiérarchie de la mémoire du GPU, ce qui permet de minimiser le nombre d'accès à la mémoire lente et d'éliminer le principal goulot d'étranglement du mécanisme d'attention standard.

Problématique de l'attention standard

Le mécanisme standard d'auto-attention dans les transformeurs est calculé selon la formule : $\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$ où Q, K, V sont les matrices de requêtes (queries), de clés (keys) et de valeurs (values).

Le principal problème de cette approche est sa complexité quadratique en temps et en mémoire (O(N²)) par rapport à la longueur de la séquence N. Dans une implémentation naïve, il est nécessaire de calculer et de stocker en mémoire GPU la matrice d'attention complète S de taille N×N, ce qui entraîne deux problèmes critiques :

1. Consommation de mémoire importante : Le stockage de la matrice N×N devient impossible lorsque l'on travaille avec des contextes longs.
2. Opérations d'entrée/sortie (E/S) : Le principal goulot d'étranglement n'est pas le nombre d'opérations arithmétiques, mais les accès constants à la mémoire lente du GPU.

Hiérarchie de la mémoire du GPU

Pour comprendre le problème, il est important de distinguer deux types de mémoire dans un GPU (en prenant l'exemple du NVIDIA A100) :

• SRAM (mémoire statique) : Mémoire sur puce rapide de faible capacité (~20 Mo) avec une bande passante énorme (jusqu'à 19 To/s).
• HBM (mémoire à large bande passante) : Mémoire lente de grande capacité (40–80 Go) avec une bande passante beaucoup plus faible (environ 1,5 To/s).

Cette asymétrie rend l'algorithme d'attention standard limité par la bande passante mémoire (memory-bound), car il lit et écrit constamment de grandes matrices depuis la HBM lente, ce qui est la principale source de latence.

Innovations clés de FlashAttention

FlashAttention est un algorithme conscient des E/S (IO-aware) qui résout ce problème en minimisant les accès à la HBM. Cet objectif est atteint grâce à trois techniques principales.

Tiling et traitement par blocs

Au lieu de traiter la matrice entière en une seule fois, FlashAttention divise les matrices d'entrée Q, K, V en petits blocs (tiles) qui peuvent tenir dans la SRAM rapide. L'algorithme charge séquentiellement ces blocs, effectue tous les calculs d'attention pour eux et met à jour le résultat final, sans jamais stocker la matrice d'attention complète dans la HBM lente.

Calcul en ligne du Softmax

Une avancée technique clé a été le calcul « en ligne » (online) du Softmax. Le Softmax standard nécessite de connaître tous les éléments du vecteur d'entrée pour la normalisation. FlashAttention utilise un algorithme modifié qui permet de calculer le Softmax par morceaux. Il maintient deux valeurs intermédiaires (le maximum actuel et la somme des exponentielles), qui sont mises à jour à mesure que de nouveaux blocs sont traités, permettant d'obtenir un résultat exact sans accéder à la matrice entière en une seule fois.

Fusion des opérations en un seul noyau CUDA

Toutes les opérations d'attention (multiplication matricielle QKᵀ, masquage, Softmax, multiplication par V) sont combinées en un unique noyau CUDA fusionné (fused kernel). Cela réduit considérablement le nombre d'opérations de lecture/écriture sur la HBM : au lieu de multiples passages sur la matrice entière, l'algorithme charge un bloc en SRAM une seule fois, effectue tous les calculs et n'écrit que le résultat final.

Efficacité théorique et pratique

Complexité et optimalité

FlashAttention réduit la consommation de mémoire de O(N²) à O(N), ce qui permet une mise à l'échelle linéaire. Il a été démontré que la complexité en E/S de l'algorithme est théoriquement optimale pour le calcul de l'attention dans une hiérarchie de mémoire à deux niveaux, ce qui signifie qu'il est impossible d'exécuter une attention exacte plus rapidement sans modifier le matériel.

Résultats empiriques

La première version de FlashAttention a démontré des améliorations significatives :

• Accélération :
  • BERT-large (longueur de séquence 512) : accélération de l'entraînement de 15 %.
  • GPT-2 (longueur de séquence 1K) : accélération 3x.
  • Tâches de la Long-Range Arena (1K-4K) : accélération 2,4x.
• Économie de mémoire : Jusqu'à 20x d'économie de mémoire par rapport aux implémentations de base exactes.
• Amélioration de la qualité des modèles : En permettant de travailler avec des contextes plus longs, FlashAttention non seulement préserve mais améliore la qualité des modèles. Par exemple, la perplexité de GPT-2 s'est améliorée de 0,7 point, et la précision sur les tâches de classification de documents longs a augmenté de 6,4 points.

Évolution et développements ultérieurs

Le succès de FlashAttention a initié toute une série d'algorithmes orientés matériel.

FlashAttention-2 (2023)

La deuxième version visait à utiliser plus pleinement les ressources du GPU. Dans le FlashAttention original, l'efficacité sur un NVIDIA A100 n'était que de 25 à 40 % du maximum. FlashAttention-2 a introduit des améliorations dans la parallélisation des calculs, ce qui a permis de :

• Atteindre une accélération 2x par rapport à la première version.
• Augmenter l'utilisation du GPU jusqu'à 50–73 % du maximum théorique.
• Étendre la prise en charge aux têtes d'attention de taille 256, ainsi qu'aux architectures Multi-Query Attention (MQA).

FlashAttention-3 (2024)

La troisième version a été spécifiquement optimisée pour l'architecture GPU NVIDIA Hopper (H100). Elle tire parti des nouvelles capacités matérielles, telles que l'asynchronisme des Tensor Cores et la prise en charge du FP8, ce qui a permis de :

• Obtenir une accélération supplémentaire de 1,5 à 2x par rapport à FlashAttention-2.
• Atteindre des performances allant jusqu'à 740 TFLOPS en FP16 et près de 1,2 PFLOPS en FP8.

Solutions spécialisées

Les idées de FlashAttention ont été développées dans d'autres projets :

• FlashInfer (2025) : Un moteur d'attention personnalisable, optimisé spécifiquement pour les tâches d'inférence des LLM. Il se concentre sur la gestion efficace du cache KV en mode de génération en continu.
• FlashMLA (2024) : Une implémentation de l'attention avec compression du cache de contexte (latent attention), permettant d'économiser de la mémoire sur de très longues séquences avec une perte d'information minimale.

Impact sur l'industrie et l'écosystème

FlashAttention est devenu une avancée fondamentale et s'est rapidement imposé comme un standard de l'industrie pour l'entraînement et l'inférence efficaces des LLM. Il a été intégré dans des bibliothèques clés telles que PyTorch et Hugging Face, et est utilisé dans la plupart des grands modèles de langage (LLaMA, MPT, Falcon, Claude, etc.).

C'est FlashAttention et ses versions ultérieures qui ont joué un rôle décisif dans l'augmentation des fenêtres de contexte des modèles de langage : de 2 000 à 4 000 tokens (GPT-3) à 128 000 tokens (GPT-4) et même jusqu'à des millions de tokens dans les modèles expérimentaux. L'algorithme a levé l'un des principaux obstacles à la mise à l'échelle des transformeurs, ouvrant de nouvelles possibilités pour les applications d'IA, de l'analyse de documents longs à la compréhension multimodale.

Liens

• Dépôt officiel de FlashAttention sur GitHub

Bibliographie

• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision. arXiv:2407.08608.
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
• Hong, K. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
• Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
• Dege, P. et al. (2025). FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs. arXiv:2506.01969.
• Wang, G. et al. (2025). FlashMask: Efficient and Rich Mask Extension of FlashAttention. OpenReview wUtXB43Chi.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (OpenReview version). OpenReview H4DqfPSibmx.
• Gholami, A. et al. (2024). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. OpenReview pF2ukh7HxA.