FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao; Daniel Y. Fu; Stefano Ermon; Atri Rudra; Christopher Ré

Articles · arXiv 2205.14135

FlashAttention : Attention exacte rapide et économe en mémoire grâce à la prise en compte des E/S

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré

Publié le 2022-05-27cs.LG arXiv PDF

Lire avec traduction et IA

Résumé en trois lignes

🎯 L'attention exacte standard matérialise une matrice d'attention $O (N^{2})$ dans la HBM lente du GPU, ce qui la rend limitée par la mémoire et difficile à adapter aux contextes longs.
🔧 FlashAttention calcule l'attention exacte au sein d'un seul noyau CUDA fusionné en découpant en tuiles les entrées sur la SRAM rapide sur puce, en mettant à jour le softmax de manière incrémentale et en recalculant l'attention sur puce pendant la passe arrière.
📈 FlashAttention apporte jusqu'à 3x d'accélération de bout en bout à l'entraînement sur GPT-2, surpasse de 15 % le record MLPerf 1.1 sur BERT-large et étend le contexte à 64K jetons.

Résumé

Les Transformers sont lents et gourmands en mémoire sur les séquences longues, car la complexité en temps et en mémoire de l'auto-attention est quadratique par rapport à la longueur de la séquence. Les méthodes d'attention approchée ont tenté de résoudre ce problème en faisant un compromis sur la qualité du modèle pour réduire la complexité de calcul, mais elles n'obtiennent souvent aucun gain de temps d'exécution réel (wall-clock speedup). Nous soutenons qu'un principe manquant consiste à rendre les algorithmes d'attention conscients des E/S (IO-aware) — en prenant en compte les lectures et écritures entre les niveaux de mémoire du GPU. Nous proposons FlashAttention, un algorithme d'attention exacte conscient des E/S qui utilise le pavage (tiling) pour réduire le nombre de lectures/écritures mémoire entre la mémoire à large bande passante du GPU (HBM) et la SRAM sur puce du GPU. Nous analysons la complexité d'E/S de FlashAttention, en montrant qu'il nécessite moins d'accès à la HBM que l'attention standard, et qu'il est optimal pour une gamme de tailles de SRAM. Nous étendons également FlashAttention à une attention bloc-creuse (block-sparse attention), produisant un algorithme d'attention approchée plus rapide que toute méthode d'attention approchée existante. FlashAttention entraîne les Transformers plus rapidement que les références existantes : accélération de 15 % du temps réel de bout en bout sur BERT-large (longueur de séquence 512) par rapport au record de vitesse d'entraînement MLPerf 1.1, accélération de 3 $\times$ sur GPT-2 (longueur de séquence 1K) et accélération de 2,4 $\times$ sur Long-Range Arena (longueur de séquence 1K-4K). FlashAttention et Block-Sparse FlashAttention permettent d'augmenter la taille du contexte dans les Transformers, ce qui améliore la qualité des modèles (perplexité améliorée de 0,7 sur GPT-2 et gain de 6,4 points sur la classification de documents longs) ainsi que des capacités entièrement nouvelles : les premiers Transformers à dépasser les performances du hasard sur le défi Path-X (longueur de séquence 16K, exactitude de 61,4 %) et Path-256 (longueur de séquence 64K, exactitude de 63,1 %).

Résumé original (anglais)

Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem by trading off model quality to reduce the compute complexity, but often do not achieve wall-clock speedup. We argue that a missing principle is making attention algorithms IO-aware -- accounting for reads and writes between levels of GPU memory. We propose FlashAttention, an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. We analyze the IO complexity of FlashAttention, showing that it requires fewer HBM accesses than standard attention, and is optimal for a range of SRAM sizes. We also extend FlashAttention to block-sparse attention, yielding an approximate attention algorithm that is faster than any existing approximate attention method. FlashAttention trains Transformers faster than existing baselines: 15% end-to-end wall-clock speedup on BERT-large (seq. length 512) compared to the MLPerf 1.1 training speed record, 3 $\times$ speedup on GPT-2 (seq. length 1K), and 2.4 $\times$ speedup on long-range arena (seq. length 1K-4K). FlashAttention and block-sparse FlashAttention enable longer context in Transformers, yielding higher quality models (0.7 better perplexity on GPT-2 and 6.4 points of lift on long-document classification) and entirely new capabilities: the first Transformers to achieve better-than-chance performance on the Path-X challenge (seq. length 16K, 61.4% accuracy) and Path-256 (seq. length 64K, 63.1% accuracy).

Guide de lecture

Problème et motivation

Les Transformers présentent une complexité temporelle et mémoire quadratique par rapport à la longueur de séquence $N$ . Bien que les méthodes d'attention approchée réduisent le nombre théorique de FLOPs, elles échouent souvent à obtenir des gains de vitesse réels car les calculs modernes sur GPU sont principalement limités par les transferts mémoire (E/S) plutôt que par les opérations arithmétiques.

Dans l'attention standard, le calcul de $S = Q K^{⊤}$ , $P = softmax (S)$ et $O = PV$ écrit des matrices intermédiaires de taille $N \times N$ dans la mémoire lente High Bandwidth Memory (HBM) et les relit plusieurs fois. Cela entraîne $Θ (N d + N^{2})$ accès HBM et une empreinte mémoire de $O (N^{2})$ , ce qui restreint les fenêtres de contexte en pratique.

Idées maîtresses

Prise en compte des E/S (IO-Awareness) : Prendre explicitement en compte la hiérarchie mémoire du GPU : la SRAM sur puce, rapide mais petite (par exemple, 192 Ko/SM avec une bande passante d'environ 19 To/s sur A100), face à la HBM, grande mais lente (40–80 Go avec une bande passante de 1,5–2,0 To/s).
Pavage / Tiling (Calcul du Softmax par blocs) : Découper les entrées $Q, K, V$ en blocs qui tiennent entièrement dans la SRAM. En suivant de manière glissante les statistiques de maximum par ligne et de normalisation du softmax, la réduction softmax est calculée de manière incrémentale à travers les blocs sans matérialiser la matrice complète $N \times N$ en HBM.
Recalcul lors de la passe arrière : Éviter de stocker la matrice d'attention intermédiaire $N \times N$ pour la rétropropagation. À la place, stocker uniquement la sortie $O$ et les statistiques du softmax $(ℓ, m)$ , puis recalculer $S$ et $P$ sur puce, bloc par bloc dans la SRAM, pendant la passe arrière.
Fusion de noyaux (Kernel Fusion) : Fusionner les multiplications matricielles, le masquage, le softmax et le dropout en un seul noyau CUDA pour éliminer les allers-retours redondants en HBM.

Méthode (avec les équations clés)

Étant donnés $Q, K, V \in R^{N \times d}$ , FlashAttention définit des tailles de blocs $B_{c} = ⌈ M / (4 d)⌉$ et $B_{r} = min (⌈ M / (4 d)⌉, d)$ , où $M$ est la capacité de la SRAM.

Suivi incrémental du Softmax

Pour un vecteur décomposé en sous-vecteurs $x^{(1)}$ et $x^{(2)}$ , la mise à jour des statistiques glissantes s'effectue via : $m (x) = max (m (x^{(1)}), m (x^{(2)})),$ $ℓ (x) = e^{m (x^{(1)}) - m (x)} ℓ (x^{(1)}) + e^{m (x^{(2)}) - m (x)} ℓ (x^{(2)}) .$

Pour le bloc $j$ de clés/valeurs et le bloc $i$ de requêtes :

Calculer les scores non normalisés sur la SRAM : $S_{ij} = Q_{i} K_{j}^{⊤} \in R^{B_{r} \times B_{c}}$ .
Calculer les statistiques locales au bloc : $\tilde{m}_{ij} = rowmax (S_{ij})$ , $\tilde{P}_{ij} = exp (S_{ij} - \tilde{m}_{ij})$ et $\tilde{ℓ}_{ij} = rowsum (\tilde{P}_{ij})$ .
Mettre à jour les normalisateurs globaux glissants : $m_{i}^{new} = max (m_{i}, \tilde{m}_{ij}), ℓ_{i}^{new} = e^{m_{i} - m_{i}^{new}} ℓ_{i} + e^{\tilde{m}_{ij} - m_{i}^{new}} \tilde{ℓ}_{ij} .$
Mettre à jour le bloc de sortie glissant $O_{i} \in R^{B_{r} \times d}$ : $O_{i} \leftarrow diag (ℓ_{i}^{new})^{- 1} (diag (ℓ_{i}) e^{m_{i} - m_{i}^{new}} O_{i} + e^{\tilde{m}_{ij} - m_{i}^{new}} \tilde{P}_{ij} V_{j}) .$

Analyse de complexité

Complexité des E/S : L'attention standard nécessite $Θ (N d + N^{2})$ accès HBM. FlashAttention nécessite $Θ (N^{2} d^{2} M^{- 1})$ accès HBM. Les auteurs prouvent qu'aucun algorithme d'attention exacte ne peut asymptotiquement atteindre $o (N^{2} d^{2} M^{- 1})$ accès pour l'ensemble des $M \in [d, N d]$ .
FLOPs et mémoire : FlashAttention effectue $O (N^{2} d)$ FLOPs (avec un nombre de FLOPs légèrement supérieur dans la passe arrière en raison du recalcul) et n'exige que $O (N)$ de mémoire auxiliaire au-delà des entrées et des sorties.

Extension : Block-Sparse FlashAttention

Restreindre l'évaluation aux blocs non nuls selon un masque de parcimonie prédéfini $M \in {0, 1}^{N / B_{r} \times N / B_{c}}$ permet d'atteindre une complexité d'E/S de $Θ (N d + N^{2} d^{2} M^{- 1} s)$ , où $s$ est la fraction de blocs non nuls.

Résultats (avec les chiffres les plus importants)

Accélération de l'entraînement :
- BERT-large (longueur de séquence 512) : Surpasse le record de vitesse d'entraînement MLPerf 1.1 de 15 % (17,4 min contre 20,0 min sur $8 \times$ A100).
- GPT-2 (longueur de séquence 1K) : Atteint jusqu'à 3,5x d'accélération par rapport à HuggingFace et jusqu'à 1,8x par rapport à Megatron-LM.
- Long-Range Arena (LRA) : Offre une accélération de 2,4x par rapport à l'attention standard ; Block-Sparse FlashAttention atteint 2,8x d'accélération.
Capacités accrues sur les contextes longs :
- GPT-2 avec un contexte de 4K : S'entraîne plus rapidement que Megatron-LM avec un contexte de 1K tout en améliorant la perplexité de validation de 0,7 point (17,5 contre 18,2).
- Classification de documents longs : Étendre le contexte de 512 à 16K apporte un gain de 4,3 points de micro- $F_{1}$ sur MIMIC-III (52,8 à 57,1) et jusqu'à 8,5 points sur ECtHR (72,2 à 80,7 à 8K).
- Path-X et Path-256 : FlashAttention est le premier Transformer à résoudre Path-X (longueur de séquence 16K) avec une exactitude de 61,4 %. Block-Sparse FlashAttention atteint 63,1 % d'exactitude sur Path-256 (longueur de séquence 64K).

Limites et questions ouvertes

Complexité d'ingénierie CUDA : L'implémentation de noyaux conscients des E/S nécessite un développement CUDA spécifique de bas niveau plutôt que des abstractions PyTorch standard ou de compilateurs, ce qui complique la portabilité entre architectures matérielles.
Attention multi-GPU : L'ordonnancement optimal des E/S sur un seul nœud ne s'étend pas automatiquement à l'attention distribuée multi-GPU via des interconnexions (NVLink, InfiniBand).
Borne inférieure de l'attention exacte : Bien qu'optimal pour une gamme générale de $M$ , prouver des bornes inférieures de complexité paramétrée fines en fonction de $M$ reste un problème ouvert.

Termes clés

IO-awareness (prise en compte des E/S): Le principe algorithmique consistant à concevoir des opérations en tenant compte des lectures et écritures entre les niveaux de mémoire matérielle (par exemple, la SRAM rapide sur puce par rapport à la HBM plus lente).
High Bandwidth Memory (HBM): La mémoire principale volumineuse du GPU qui offre une grande capacité totale mais une bande passante nettement inférieure à celle du cache sur puce.
SRAM: Mémoire tampon (scratchpad) rapide sur puce sur les GPU modernes, située au sein de chaque multiprocesseur de flux (Streaming Multiprocessor, SM).
Tiling (pavage): Restructuration d'un calcul en divisant les matrices d'entrée en blocs plus petits qui s'insèrent dans la SRAM et en accumulant les résultats au fil des itérations.
Selective gradient checkpointing (checkpointing sélectif de gradient): Une technique qui évite de sauvegarder la grande matrice d'attention intermédiaire pendant la passe avant en la recalculant sur puce pendant la passe arrière à l'aide des statistiques de softmax mises en cache.
Kernel fusion (fusion de noyaux): Combinaison de plusieurs opérations distinctes (multiplication matricielle, masquage, softmax, dropout) en une seule exécution de noyau GPU afin d'éliminer les lectures/écritures répétées vers la HBM.
Block-sparse FlashAttention: Une extension de FlashAttention qui ne calcule que les blocs de tuiles non nuls définis par un masque de parcimonie par blocs, réduisant ainsi la complexité d'E/S proportionnellement à la fraction de blocs non nuls.
Path-X / Path-256: Benchmarks de raisonnement visuel sur séquences longues issus de Long-Range Arena (longueurs de séquence de 16K et 64K) sur lesquels les Transformers précédents n'ont pas réussi à obtenir des performances supérieures au hasard.

Questions

FlashAttention fait-il une approximation de la matrice d'attention ?

Non. FlashAttention calcule une auto-attention exacte avec des sorties mathématiques identiques à celles de l'attention standard, à la précision des nombres à virgule flottante près.

Pourquoi FlashAttention est-il plus rapide s'il calcule davantage de FLOPs dans la passe arrière ?

Parce que l'attention standard est limitée par les transferts mémoire ; déplacer la matrice $N \times N$ entre la HBM et la SRAM prend beaucoup plus de temps réel que d'effectuer des opérations arithmétiques supplémentaires sur la SRAM rapide.

Quelle mémoire auxiliaire FlashAttention conserve-t-il pour la passe arrière ?

FlashAttention ne stocke que la sortie avant $O$ , les facteurs de normalisation du softmax $ℓ$ et $m$ (ce qui nécessite un espace mémoire supplémentaire de $O (N)$ ), ainsi que l'état du générateur pseudo-aléatoire pour reproduire les masques de dropout sur puce.

Comment FlashAttention se compare-t-il en termes de complexité d'E/S à l'attention standard ?

L'attention standard nécessite $Θ (N d + N^{2})$ accès HBM, tandis que FlashAttention nécessite $Θ (N^{2} d^{2} M^{- 1})$ accès, où $M$ est la taille de la SRAM et $d$ est la dimension de la tête d'attention.

La présentation de cette page a été rédigée par une IA à partir de l'article et peut contenir des erreurs. Consultez l'article pour les détails.

Lire avec traduction et IA