FlashAttention-4 : comment Together AI a accéléré l'attention sur GPU Blackwell

FlashAttention-4 est une refonte de l'algorithme d'attention spécialement conçue pour les GPU modernes, où les noyaux tensoriels augmentent en performance beaucoup plus rapidement que la mémoire et les autres ressources.

Pourquoi l'approche standard ne fonctionne plus

À première vue, la performance de l'attention (attention) est contrôlée par la vitesse de deux multiplications matricielles : S = Q × K^T et O = P × V. Mais l'analyse de Blackwell B200 montre quelque chose d'inattendu : le goulot d'étranglement ne se trouve pas dans les noyaux tensoriels, mais dans les blocs fonctionnels spécialisés (SFU) pour le calcul de l'exponentielle dans softmax (forward pass) et dans le trafic de mémoire partagée (backward pass).

De Hopper (H100) à Blackwell (B200), la performance BF16 des noyaux tensoriels a augmenté de 1 à 2,25 pétaflops, tandis que le nombre de SFU et la bande passante de la mémoire partagée sont restés inchangés. Cette asymétrie casse l'optimisation standard — on ne peut pas simplement supposer que les noyaux contrôlent entièrement la performance. En réalité, ils s'interfèrent les uns avec les autres.

Comment Together AI a résolu le problème

L'équipe de recherche, avec les ingénieurs de NVIDIA, Meta et Princeton, a proposé trois idées clés :

• Nouveau pipelining — des pipelines logiciels qui chevauchent au maximum le travail des noyaux tensoriels, des SFU et de la mémoire sans temps d'inactivité.
• Imitation de l'exponentielle via des polynômes — au lieu du bloc SFU lent, la forward pass calcule l'exponentielle sur des unités FMA rapides (fused multiply-add).
• TMEM et 2-CTA MMA — utilisation de la nouvelle tensor memory (256 KB par SM) plus un mode où deux blocs de threads travaillent sur une seule opération matricielle, réduisant le trafic de mémoire partagée.

Les capacités de Blackwell qui l'ont rendu possible

Blackwell a spécialement ajouté plusieurs fonctionnalités pour ces optimisations :

Tensor memory (TMEM) — stockage rapide sur puce (256 KB par SM), câblé directement aux noyaux tensoriels. Les résultats intermédiaires peuvent rester dans TMEM sans accéder à la mémoire partagée plus lente, ce qui réduit drastiquement la latence d'accès.

Noyaux tensoriels asynchrones de 5e génération — chaque noyau est lancé par un seul thread et accumule les résultats dans TMEM. Le tile maximal pour BF16 est 128×256×16 (environ 2 fois plus grand que Hopper), ce qui permet un pipelining plus profond sans débordement de registres.

2-CTA MMA — un nouveau mode où deux blocs de threads travaillent simultanément sur une seule opération matricielle. Cela réduit la charge sur la mémoire partagée de moitié et diminue le nombre d'opérations atomiques.

Chiffres et résultats

FlashAttention-4 sur Blackwell B200 avec BF16 atteint 1605 TFLOPs/s (71% d'utilisation). C'est 1,3× plus rapide que cuDNN 9.13, et 2,7× plus rapide que Triton. Pour le contexte : 1605 pétaflops, c'est presque la moitié de la performance de pointe de Blackwell, mais c'est précisément ce qui est extrait du noyau d'attention complexe.

Ce que cela signifie

FlashAttention-4 montre comment il faut travailler à l'ère de la mise à l'échelle asymétrique des GPU — ne pas peaufiner les anciens algorithmes, mais les refondre avec les nouvelles capacités matérielles. Le résultat ressemble à une réussite, mais c'est simplement le début de l'adaptation à la nouvelle réalité du matériel.

*Meta est reconnue comme une organisation extrémiste et interdite en Russie.