L’architecture Transformer, au cœur du développement des modèles de langage génératifs (LLM) comme GPT, DeepSeek ou Llama, repose sur le mécanisme d’attention auto-attentionnelle, permettant un traitement parallèle de séquences entières plutôt que mot par mot. Cette capacité à capter des dépendances à longue portée est essentielle, mais elle entraîne une complexité computationnelle et mémoire quadratique en fonction de la longueur de la séquence, créant un goulot d’étranglement critique pour les contextes longs. Pour surmonter ces limites, FlashAttention a été conçu comme une avancée algorithmique optimisée pour l’entrée/sortie (IO-aware), offrant le même résultat mathématique que l’attention standard, mais avec une efficacité accrue. Cette approche réduit considérablement la consommation mémoire et accélère à la fois l’entraînement et l’inférence, permettant aux modèles de gérer des séquences plus longues — une nécessité pour des applications comme l’analyse d’images haute résolution ou les conversations prolongées. FlashAttention-4 (FA4), la dernière itération, est le fruit d’une co-conception logicielle et matérielle spécifiquement adaptée à l’architecture NVIDIA Blackwell, notamment les GPU HGX B200. FA4 atteint un pic de performance de 1 605 TFLOPS/s, soit 71 % du maximum théorique du matériel. Grâce à une redéfinition du mécanisme d’attention, FA4 surmonte les limitations de l’architecture Blackwell, notamment son échelle asymétrique où la puissance de calcul augmente bien plus vite que la bande passante mémoire. Il bat ainsi les références standard, offrant jusqu’à 1,3x de vitesse en plus par rapport à cuDNN et 2,4x par rapport à Triton Inference Server. Ces gains s’étendent à la phase de rétropropagation, où FA4 utilise la mémoire Tensor Memory (TMEM) — une mémoire on-chip de 256 Ko par unité de traitement (SM) — pour stocker les intermédiaires (S, P, dP, dS, dQ), évitant ainsi la surcharge de la mémoire partagée (SMEM) et réduisant les débordements de registres. FA4 s’appuie sur plusieurs innovations : une réduction de la charge sur les unités MUFU (exponentielles) via des approximations polynomiales basées sur des opérations FMA, des tuiles d’opérations matricielles plus grandes (jusqu’à 128×128), une nouvelle planification des tâches (LPT) pour le masquage causal, et des pipelines entièrement asynchrones qui maximisent le chevauchement entre calcul, softmax et accès mémoire. Ces optimisations permettent de tirer pleinement parti de la double capacité de calcul des nouveaux Tensor Cores, sans être freinées par les contraintes mémoire. En outre, FA4 profite de nouvelles fonctionnalités logicielles comme CUDA 13 et les outils CUDA-X, ainsi que du DSL CuTe en Python, qui accélèrent les temps de compilation de 20 à 30 fois par rapport à FA3, tout en maintenant une expressivité élevée des noyaux. Des frameworks d’inférence comme SGLang et vLLM supportent déjà le pré-remplissage (prefill) avec FA4, et les techniques ont été intégrées dans cuDNN 9.14. Cette avancée est particulièrement bénéfique dans les configurations distribuées multi-GPU et multi-nœuds, où elle améliore à la fois la performance et l’efficacité énergétique. En résumé, FlashAttention-4 incarne une nouvelle génération d’algorithmes conçus pour exploiter pleinement les capacités des accélérateurs modernes. Son succès repose sur une compréhension fine des contraintes matérielles et une optimisation logicielle poussée, positionnant FA4 comme un pilier clé du développement futur des LLM à très grande échelle.