FlashAttention-3 : une attention plus rapide et plus précise grâce à l’asynchronisme et à la basse précision

• Le goulot d’étranglement de l’attention dans les Transformers limitait la vitesse d’entraînement et d’inférence des LLM à long contexte ; FlashAttention-3 est une nouvelle version qui vise à le réduire davantage en exploitant les fonctionnalités des GPU Hopper
• L’idée centrale consiste à utiliser l’asynchronisme des Tensor Cores et de TMA pour chevaucher calcul et déplacement de données, et à entrelacer matmul et softmax afin de réduire les temps d’inactivité du GPU
• Sur H100, FlashAttention-2 n’exploitait que 35 % des FLOPS théoriques maximaux, tandis que FlashAttention-3 atteint jusqu’à 740 TFLOPS en FP16, soit 75 % du maximum théorique du H100
• En basse précision FP8, les valeurs aberrantes des activations de LLM amplifient les erreurs ; un incoherent processing fondé sur la transformée de Hadamard réduit l’erreur de 2,6× par rapport à l’attention FP8 de référence
• FlashAttention-3 est disponible sur GitHub ; à mesure que le coût de l’attention diminue, les marges de progression augmentent pour traiter des contextes plus longs et améliorer l’efficacité de l’entraînement et de l’inférence des LLM

---

Objectifs et performances de FlashAttention-3

• L’attention est une couche essentielle des Transformers, mais elle devient un goulot d’étranglement majeur dans les grands modèles de langage et les applications à long contexte
• FlashAttention et FlashAttention-2 ont ouvert la voie à une approche qui accélère l’attention en réduisant les lectures et écritures en mémoire GPU ; la plupart des bibliothèques l’utilisent aujourd’hui pour accélérer l’entraînement et l’inférence des Transformers
• Cette approche a contribué à l’augmentation de la longueur de contexte des LLM ces deux dernières années, de 2–4K pour GPT-3 et OPT à 128K pour GPT-4, et jusqu’à 1M pour Llama 3
• FlashAttention-2 ne parvenait qu’à 35 % d’utilisation des FLOPS théoriques maximaux sur GPU H100, mais FlashAttention-3 améliore ce résultat en exploitant les nouvelles fonctionnalités des GPU Hopper
• Les performances publiées de FlashAttention-3 sont les suivantes
  • 1,5 à 2,0× plus rapide que FlashAttention-2 en FP16
  • Jusqu’à 740 TFLOPS en FP16
  • 75 % d’utilisation des FLOPS théoriques maximaux du H100
  • Près de 1,2 PFLOPS en FP8
  • Une erreur 2,6× plus faible que l’attention FP8 de référence

Rappel du fonctionnement de FlashAttention

• FlashAttention réorganise l’ordre de calcul de l’attention et accélère le traitement avec du tiling et du recalcul, tout en réduisant l’usage mémoire par rapport à la longueur de séquence, de quadratique à linéaire
• Les blocs d’entrée sont chargés de la HBM vers la SRAM, l’attention est calculée sur ces blocs, puis la sortie est mise à jour en HBM
• Comme la grande matrice d’attention intermédiaire n’est pas écrite en HBM, les lectures et écritures mémoire diminuent, ce qui permet une accélération réelle de 2 à 4× à l’exécution
• En combinant tiling et rescaling du softmax, on peut traiter les données par blocs tout en obtenant une sortie correcte sans approximation

Fonctionnalités des GPU Hopper : WGMMA, TMA, FP8

• FlashAttention-2 peut atteindre jusqu’à 70 % des FLOPS théoriques maximaux sur les GPU Ampere A100, mais n’exploite pas pleinement les nouvelles fonctionnalités des GPU Hopper
• FlashAttention-3 utilise trois fonctionnalités de Hopper
  • WGMMA : fonctionnalité warpgroup matrix multiply-accumulate qui exploite les nouveaux Tensor Cores de Hopper, avec un débit supérieur à mma.sync sur Ampere
  • TMA : unité matérielle dédiée qui accélère les transferts de données entre global memory et shared memory, en prenant en charge le calcul des index et la prédication out-of-bound afin de réduire l’utilisation des registres
  • FP8 : peut doubler le débit des Tensor Cores par rapport au FP16, mais représente les valeurs en virgule flottante avec moins de bits, ce qui implique un compromis avec la précision
• FlashAttention-3 utilise les abstractions de NVIDIA CUTLASS pour exploiter les fonctionnalités de Hopper
• Le simple fait de réécrire FlashAttention pour utiliser ces nouvelles fonctionnalités fait passer les performances du forward pass FP16 d’environ 350 TFLOPS avec FlashAttention-2 à 540–570 TFLOPS

Chevauchement de GEMM et softmax grâce à l’asynchronisme

• Les principales opérations de l’attention sont les GEMM entre Q-K et P-V, ainsi que le softmax
• Sur les accélérateurs modernes, les opérations autres que matmul sont beaucoup plus lentes que les matmul, et les fonctions spéciales comme l’exponentielle du softmax sont traitées par des unités distinctes de celles du floating point multiply-add ou du matrix multiply-add
• Le H100 SXM5 fournit 989 TFLOPS en matrix multiply FP16, mais le débit des fonctions spéciales n’est que de 3,9 TFLOPS, soit 256 fois moins
• Avec une dimension de tête de 128, même si les FLOPS de matmul sont 512 fois plus nombreux que ceux de l’exponentielle, l’exponentielle peut représenter 50 % du temps par rapport au matmul
• En FP8, les FLOPS de matmul doublent, mais la vitesse de l’exponentielle reste inchangée ; exécuter matmul et softmax en parallèle devient donc encore plus important

• Ping-pong scheduling entre warpgroups

  • Le warp scheduler du GPU effectue automatiquement une partie de l’overlap en exécutant d’autres warps pendant que certains attendent les résultats de GEMM
  • FlashAttention-3 utilise des barrières de synchronisation pour mieux chevaucher manuellement le GEMM et le softmax de deux warpgroups
  • Le warpgroup 1 exécute d’abord le GEMM1 d’une itération et le GEMM0 de l’itération suivante
  • Ensuite, pendant que le warpgroup 2 exécute le GEMM, le warpgroup 1 traite le softmax
  • Ce ping-pong schedule consiste à cacher le softmax derrière le temps d’exécution du GEMM de l’autre warpgroup
  • En pratique, le scheduling n’est pas aussi parfaitement net que dans le schéma, mais il fait passer le forward pass d’attention FP16 d’environ 570 TFLOPS à 620 TFLOPS pour une dimension de tête de 128 et une longueur de séquence de 8K

• Overlap au sein d’un warpgroup

  • Au sein d’un même warpgroup, il est également possible d’exécuter une partie du softmax pendant que le GEMM de ce warpgroup est en cours
  • Ce pipelining fait passer le débit du forward d’attention FP16 d’environ 620 TFLOPS à 640–660 TFLOPS
  • En contrepartie, il faut conserver simultanément les accumulateurs GEMM et les entrées/sorties du softmax, ce qui augmente la pression sur les registres
  • Globalement, cette technique offre un compromis avantageux

Basse précision FP8 et incoherent processing

• Les activations des LLM peuvent contenir des outliers dont la magnitude est beaucoup plus élevée que celle des autres features
• Ces outliers compliquent la quantification et augmentent fortement l’erreur de quantification
• FlashAttention-3 exploite l’incoherent processing utilisé dans la littérature sur la quantification, comme QuIP
• Il multiplie les queries et les keys par une matrice orthogonale aléatoire afin de disperser les outliers et de réduire les erreurs de quantification
• L’implémentation utilise une transformée de Hadamard avec des signes aléatoires
  • Si l’on note d la dimension de tête, elle peut être appliquée à chaque tête d’attention en O(d log d) au lieu de O(d²)
  • La transformée de Hadamard étant memory-bandwidth bound, elle peut être fusionnée avec l’opération précédente également memory-bandwidth bound, comme le rotary embedding, sans coût supplémentaire
• Dans une expérience où Q, K et V sont générés selon une loi normale standard et où de grandes magnitudes sont ajoutées à 0,1 % des entrées pour simuler des outliers, l’incoherent processing réduit l’erreur de quantification de 2,6×

Benchmarks et disponibilité

• FlashAttention-3 est comparé non seulement à FlashAttention-2, mais aussi aux implémentations Triton et cuDNN qui utilisent déjà les nouvelles fonctionnalités matérielles des GPU Hopper
• En FP16, il affiche une accélération d’environ 1,6 à 1,8× par rapport à FlashAttention-2
• En FP8, il atteint près de 1,2 PFLOPS
• Le dépôt GitHub de FlashAttention-3 est public
• L’article de recherche est également disponible dans le même dépôt flash-attention

Optimisations restantes et intégration future

• L’article inclut, en plus des éléments abordés dans le billet de blog, des optimisations comme les séquences de longueur variable, les persistent kernels et la transposition in-kernel pour FP8
• Concevoir l’algorithme en fonction du matériel d’exécution permet d’importants gains d’efficacité et ouvre de nouvelles capacités de modèles, comme les longs contextes
• Les travaux futurs incluent l’optimisation de l’inférence des LLM et la généralisation des techniques à d’autres architectures matérielles
• FlashAttention-3 devrait être intégré dans une future version de PyTorch

[GN⁺]

Avis de Hacker News

• À en juger par les commentaires dans le code, Tri Dao semble travailler sur FA3 depuis avril 2022, juste après l’annonce de Hopper/H100
  Il est assez intéressant qu’il ait fallu plus de deux ans pour que le code soit publié aujourd’hui ; c’est peut-être parce que de meilleures solutions sont en préparation
  Les publications récentes de Tri penchent vers les architectures de type SSM et Mamba. FlashAttention a une complexité temporelle quadratique par rapport à la longueur de séquence, mais les algorithmes récents sont sous-quadratiques : ils ne se contentent donc pas de réaliser le même calcul plus efficacement, ils réduisent fortement la quantité de calcul elle-même
  Dao et Gu ont montré cette année, dans un long article, que Mamba/SSM peuvent eux aussi être formulés de façon à bien s’accélérer avec les mêmes primitives matérielles dont profitent les Transformers

  • Tant que la Strong Exponential Time Hypothesis (SETH) n’est pas prouvée ou réfutée, il faudra soit payer un coût quadratique, soit renoncer à quelque chose. Au fond, c’est le coût d’une recherche exhaustive
    Prouver ou réfuter SETH résoudrait aussi le problème P vs NP, donc il est difficile de s’attendre à ce que cela arrive bientôt
    La question clé est de savoir si un cas d’usage donné peut supporter ce coût

• Je me demande à quel point l’algorithme FlashAttention est lié au matériel
  Par exemple, cette annonce dit exploiter les fonctionnalités asynchrones du GPU H100 ; cela semble donc signifier que les cartes qui ne sont pas de la série H ne bénéficient pas de cette accélération
  Par ailleurs, la bibliothèque FlashAttention réelle nécessite CUDA, mais l’algorithme semble avoir été porté vers Metal[^0]. Si l’algorithme est proche d’une fonction pure, on pourrait penser qu’il devrait être implémentable sur n’importe quel GPU/framework de machine learning
  [0]: https://github.com/philipturner/metal-flash-attention

  • Il y a beaucoup de bonnes réponses, mais en bref : « en pratique, assez fortement » lié au matériel. L’exemple ci-dessous est parlant

    https://github.com/karpathy/nanoGPT/blob/master/model.py#L45
    Le nanoGPT de Karpathy vérifie l’existence de torch.nn.functional.scaled_dot_product_attention pour appeler FlashAttention
    https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
    Si l’on regarde la documentation, en pratique on voudra surtout appeler FA2 ; FA2 optimise les kernels du périphérique pour découper l’opération Softmax sur la matrice triangulaire et réduire les allers-retours inutiles de lots de flottants entre le GPU et le CPU
    https://arxiv.org/pdf/2307.08691
    L’article sur FA2 est presque entièrement rédigé du point de vue du matériel sur lequel il s’exécute

  • L’amélioration algorithmique de FlashAttention consiste principalement à découper puis recombiner la partie Softmax de l’attention, ce qui n’est pas en soi une idée totalement nouvelle. Sa contribution écrasante tient à la mise en œuvre efficace de cette méthode et de ses détails sur le matériel Nvidia
  • À l’origine, FlashAttention dépendait très peu du matériel
    Pour les versions récentes, cela dépend du niveau d’abstraction. ThunderKittens[0] offre, comme ce qui est décrit dans l’article, un gain d’environ 1,3× à 2× par rapport à FA2 tout en restant relativement applicable de manière générale sur les GPU
    Chaque nouveau matériel peut avoir des fonctionnalités spécifiques permettant d’en tirer des performances supplémentaires. En général, les fournisseurs adoptent les fonctionnalités qui les placent devant les autres, mais comme on le voit déjà avec CUDA, les API et bibliothèques se fragmentent
    [0]: https://hazyresearch.stanford.edu/blog/2024-05-12-tk
  • Conceptuellement, il l’est un peu ; du point de vue de l’implémentation pratique, beaucoup. Même l’implémentation Python standard compile en interne des kernels adaptés à un matériel spécifique
  • Pour ajouter un point pratique, le matériel AMD manque encore d’une bonne implémentation de flash-attention-2. ROCm devient lentement utilisable, mais il n’est pas encore au niveau de CUDA

• J’aimerais poser la question aux spécialistes des compilateurs. Y a-t-il une chance qu’un compilateur découvre de lui-même des optimisations comme FlashAttention ? TVM et tinygrad semblent aller dans cette direction, mais j’ai du mal à croire que ce soit possible

  • En théorie, c’est possible. Grâce aux propriétés algébriques des mathématiques, de vastes réordonnancements sont possibles, auxquels on peut ajouter un tiling de boucles polyédrique relativement structuré
    Cela dit, le coût est élevé, donc il faut mettre en cache les résultats de cette exploration
    L’optimisation par e-graphs semble bien adaptée à ce domaine. Mais elle exige un changement de paradigme important dans la façon de traiter les passes d’optimisation, si bien qu’elle n’a quasiment pas été déployée en dehors de quelques outils de niche. Par exemple, elle s’accorde mal avec les graphes d’appels traditionnels : pour déployer des e-graphs au-delà et entre les blocs de base et les boucles for, il faut fortement modifier le flot de contrôle, et break ainsi que return ne sont pas pris en charge
  • Cela me semble extrêmement difficile, mais pas impossible
    Je ne sais pas vraiment jusqu’où va l’état de l’art en optimisation de compilateurs en matière de placement des données et de maximisation de l’utilisation du processeur
    J’ai vu autrefois une vidéo sur l’optimisation qui expliquait qu’une petite optimisation avait bien accéléré le programme, mais que son effet restait minime par rapport aux variations de vitesse dues aux différences de placement mémoire causées par cette optimisation, voire par des modifications aléatoires
    Cette présentation portait davantage sur la distinction du signal dans le bruit, mais ce bruit lui-même indique que les compilateurs ne gèrent pas encore bien des formes bien plus simples que le problème évoqué ici
    Rien que l’architecture CPU/mémoire, où les caches et les schémas d’accès influencent les performances, est complexe ; si l’on y ajoute l’architecture GPU, cela semble être un terrain encore largement inexploré
    Ce sera peut-être possible un jour. Comme on parle d’IA, la question se pose aussi de savoir si une IA suffisamment intelligente pourrait le faire, mais tout dépend du sens de « suffisamment »
    Comme test de très haut niveau pour un modèle d’IA, on pourrait imaginer lui donner quelque chose comme micrograd et lui demander de créer, en conservant la même interface, quelque chose de plus rapide que torch. On n’en est pas encore du tout là, mais ce serait intéressant si cela devenait possible
  • Je ne pense pas. Il faut plutôt voir cela comme un autre algorithme. Au lieu de ne considérer que les mathématiques, on conçoit l’algorithme en tenant compte de la forme du matériel
    TVM, je peux le comprendre. À strictement parler, il fait autre chose, mais c’est un domaine assez proche
    En revanche, je ne vois pas pourquoi tinygrad donne cette impression
  • https://github.com/uwplse/tensat
  • Appeler des opérateurs de haut niveau depuis un langage enveloppant comme Python est assez délicat

• Si quelqu’un veut porter cela vers ROCm / AMD MI300x, contactez [e-mail]. Je ne vous enverrai absolument aucun spam
  Je peux fournir du temps de calcul pour ce travail

  • Vous êtes donc une société de serveurs avec accélérateurs AMD ! Beau travail, j’espère que quelqu’un acceptera :)
  • Sans vouloir être impoli, je me demande quel est l’objectif de cette proposition. Quelqu’un va-t-il vraiment faire ce portage gratuitement en échange d’un simple accès au matériel ? Quel est son intérêt ?

• FlashAttention-3 is optimized for Hopper GPUs (e.g. H100).
  Quelles performances FA3 obtient-il sur des GPU grand public comme les 3090 et 4090 ?

  • C’est spécifique à Hopper. Les améliorations sont fortement liées aux fonctionnalités Hopper comme les warp groups et TMA
    Sur une 4090, l’implémentation Triton de l’attention en FP8 pourrait apporter un gain de vitesse : https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html

• TMA (Tensor Memory Accelerator). This is a special hardware unit that accelerates the transfer of data between global memory and shared memory, taking care of all index calculation and out-of-bound predication. This frees up registers, which is a valuable resource to increase tile size and efficiency.
  D’après ce que je comprends, il est vrai que TMA réduit l’usage des registres, mais le point le plus important est qu’il permet de confier la génération des adresses au matériel. À mesure que les opérations autour s’accélèrent, la génération d’adresses peut devenir un goulot d’étranglement

• C’est l’une des améliorations les plus importantes pour l’IA dans son ensemble. Elle permet de faire plus, plus vite, avec le même matériel, et apporte un bénéfice presque sans compromis à la plupart des utilisateurs d’IA

  • Pour les utilisateurs qui ont un H100, oui

• Je me demande pourquoi FlashAttention est environ 5 fois plus lent avec un masquage variable que sans. Si la prise en charge du masquage n’est pas bonne, les gains d’optimisation disparaissent presque

  • Où voyez-vous ce benchmark ?

• Ce serait bien qu’un expert réponde à quelques questions :)
  FlashAttention est-il un remplacement drop-in de l’opération d’attention des LLM ? Peut-on l’utiliser partout où une opération d’« attention » est employée, ou faut-il entraîner séparément le LLM pour utiliser FA ?
  Quel rapport FA entretient-il avec des stratégies comme GQA ou l’attention à fenêtre glissante ? S’agit-il de concepts orthogonaux, ou faut-il une implémentation FA distincte pour chaque stratégie ?
  Récemment, llama.cpp a ajouté la prise en charge de FlashAttention : cela signifie-t-il qu’il a commencé à utiliser des choses comme les kernels CUDA fournis par FlashAttention ?
  Enfin, cet article compare FlashAttention et Triton. Triton n’est-il pas une sorte de couche d’abstraction ? Ne peut-on pas implémenter FA avec Triton ? Je ne comprends pas bien l’expression « FlashAttention contre Triton »

  • 1. C’est à peu près ça. Mathématiquement, c’est équivalent. Côté logiciel, les seuls problèmes sont des choses comme la gestion des versions des dépendances ou le format des données en mémoire, et FlashAttention 2 est déjà intégré à HuggingFace et à plusieurs bibliothèques populaires. FlashAttention 3 a de fortes chances d’y arriver bientôt aussi, mais il faut un GPU H100 pour l’exécuter
    2. FlashAttention 2 a ajouté la prise en charge de GQA dans une mise à jour précédente :
       https://github.com/Dao-AILab/flash-attention
    3. Ici, la comparaison porte sur cette implémentation de FlashAttention écrite en CUDA C++ pur, et sur une implémentation Triton d’un algorithme similaire écrite en Triton : https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html
  • FlashAttention peut remplacer en drop-in l’opération d’attention des LLM
    FlashAttention est une méthode pour calculer la partie Softmax(QK^T)V de l’attention, tandis que GQA est une méthode pour calculer les matrices Q, K et V. Pour l’attention à fenêtre glissante, j’en suis moins sûr, mais il s’agit de modifier le masque d’attention qui contrôle quelles requêtes peuvent voir quelles clés
    Je n’ai pas utilisé llama.cpp, mais l’explication selon laquelle il aurait commencé à utiliser des kernels CUDA semble globalement correcte
    La dernière question fait référence à une implémentation de FlashAttention précédemment écrite en Triton

• Je suis curieux, car l’article dit que des opérations comme sigmoid sont très lentes
  Les LLM modernes utilisent beaucoup de fonctions d’activation intégrant sigmoid ou Softmax, comme SiLU, Swish ou SOLU
  ReLU entraîne-t-il moins de perte de performances ? Si oui, vaudrait-il mieux revenir au bon vieux ReLU ?

  • ReLU est littéralement une fonction linéaire tronquée à zéro à partir d’un certain point, donc elle demande beaucoup moins de calcul que celles qui impliquent des fonctions exponentielles. Cela dit, il paraît difficile d’obtenir des résultats compétitifs avec une fonction d’activation aussi simple