Flash Attention

Flash Attention (FlashAttention) est un algorithme et une famille de noyaux GPU qui calculent exactement l’attention des modèles Transformer sans jamais écrire la matrice d’attention complète dans la mémoire principale du GPU. L’auto-attention standard sur N jetons construit une matrice de scores N × N, l’écrit dans la mémoire à large bande passante (HBM), puis la relit pour appliquer le softmax et calculer le produit avec les valeurs ; la majeure partie du temps d’exécution de l’attention est donc consacrée au transfert des données. Tri Dao et ses collègues de Stanford ont présenté cette méthode dans l’article de 2022 FlashAttention : une attention exacte rapide et économe en mémoire, tenant compte des entrées-sorties. Elle divise le calcul en blocs qui tiennent dans la mémoire SRAM rapide intégrée à la puce ; l’utilisation de la mémoire augmente ainsi de manière linéaire avec la longueur de la séquence, et le résultat correspond à celui de l’attention standard, à l’arrondi en virgule flottante près. Ses noyaux s’exécutent dans PyTorch et dans des moteurs d’inférence pour grands modèles de langage (LLM), comme vLLM.

Fonctionnement de Flash Attention#

Un GPU dispose d’une petite quantité de mémoire très rapide près de ses unités de calcul et d’une grande mémoire plus lente hors puce. L’article sur FlashAttention décrit un NVIDIA A100 équipé de 40 à 80 Go de HBM à 1,5–2,0 To/s et de 192 Ko de SRAM dans chacun de ses 108 processeurs multiprocesseurs de flux, avec un débit estimé à 19 To/s. L’attention effectue peu d’opérations par octet lu ; sa vitesse dépend donc de la fréquence à laquelle les données passent entre ces deux niveaux. La méthode comptabilise ces transferts, d’où le terme « tient compte des entrées-sorties ».

FlashAttention charge un bloc de requêtes, puis fait défiler des blocs de clés et de valeurs dans la SRAM. Pour chaque bloc, il calcule les scores, met à jour un maximum et une somme cumulés pour chaque ligne de requête, redimensionne la sortie partielle chaque fois que le maximum augmente et ajoute la contribution du nouveau bloc. Ce softmax en ligne fournit le résultat exact du softmax sans qu’il soit nécessaire de traiter une ligne entière d’un seul coup ; seules la sortie finale et deux valeurs par ligne sont renvoyées dans la HBM. Le passage arrière recalcule les blocs de scores à partir de Q, K et V au lieu de stocker la matrice N × N ; cela ajoute des opérations, mais reste plus rapide, car il lit beaucoup moins de mémoire. Voici la même boucle en PyTorch, qui traite toutes les requêtes avec un seul bloc de clés et de valeurs à la fois :

import torch

torch.manual_seed(0)
n, d, block = 4096, 64, 256
q, k, v = torch.randn(3, n, d).unbind(0)

# Standard attention builds the full n x n score matrix in memory
reference = torch.softmax(q @ k.T / d**0.5, dim=-1) @ v

# FlashAttention-style: stream key/value tiles, keeping a running max and sum per query (online softmax)
row_max, row_sum, out = torch.full((n, 1), -torch.inf), torch.zeros(n, 1), torch.zeros(n, d)
for start in range(0, n, block):
    scores = q @ k[start : start + block].T / d**0.5  # only an n x block tile exists at once
    new_max = torch.maximum(row_max, scores.max(dim=-1, keepdim=True).values)
    rescale, probs = torch.exp(row_max - new_max), torch.exp(scores - new_max)
    row_sum = row_sum * rescale + probs.sum(dim=-1, keepdim=True)
    out = out * rescale + probs @ v[start : start + block]
    row_max = new_max
out = out / row_sum

print(f"Largest score tile: {tuple(scores.shape)} vs full matrix: ({n}, {n})")
print(f"Max difference from standard attention: {(out - reference).abs().max().item():.1e}")

Le plus grand bloc de scores mesure 4 096 × 256, soit un seizième de la matrice 4 096 × 4 096, et le résultat diffère de l’attention standard d’au plus 4.3e-07, soit l’erreur d’arrondi FP32. Un véritable noyau divise également les requêtes en blocs, à raison d’un bloc par bloc de threads GPU, et fusionne toute la boucle dans un seul noyau afin que les blocs ne quittent jamais la SRAM.

FlashAttention-2, -3 et -4#

Chaque version vise une génération de GPU plus récente. FlashAttention-2 (2023) a réorganisé la répartition du travail entre les threads et s’exécute environ deux fois plus vite que la version originale, atteignant 50 à 73 % du débit maximal d’un A100. FlashAttention-3 (2024) utilise les Tensor Cores asynchrones du H100 et prend en charge FP8, atteignant 740 TFLOPs/s en FP16 et près de 1,2 PFLOPs/s en FP8. FlashAttention-4 (2026) cible les GPU Blackwell et annonce jusqu’à 1 613 TFLOPs/s en BF16 sur un B200, soit 71 % du débit maximal. Les noyaux sont fournis dans la bibliothèque open source flash-attention, qui fonctionne avec des entrées FP16 et BF16 ; la fonction scaled_dot_product_attention de PyTorch sélectionne automatiquement un noyau FlashAttention-2 lorsque le GPU et les entrées le prennent en charge, et utilise sinon des implémentations économes en mémoire ou classiques. Le SAM 2 de Meta, disponible dans le package Ultralytics, calcule l’attention dans son encodeur d’images Hiera et son attention mémoire au moyen de cette fonction ; PyTorch utilise donc FlashAttention pour ces couches dès que le matériel et la précision des entrées le permettent.

Flash Attention et notions connexes#

• Attention standard : même opération mathématique, calculée en écrivant les matrices complètes de scores et de probabilités dans la HBM.
• Attention creuse : ne calcule les scores que pour certaines paires requête-clé. L’article sur FlashAttention a également introduit une variante creuse par blocs qui ignore les blocs vides.
• Attention linéaire : remplace le softmax par une carte de caractéristiques afin que le coût augmente de manière linéaire avec la longueur, ce qui modifie les sorties du modèle.
• PagedAttention : gestionnaire de mémoire qui stocke le cache KV dans des pages de taille fixe pendant le service. Les noyaux d’attention, y compris FlashAttention, lisent les clés et les valeurs dans ces pages.
• Attention en anneau : répartit une très longue séquence sur plusieurs appareils, chacun calculant l’attention par blocs sur sa partition tandis que les blocs de clés et de valeurs circulent dans un anneau.

Effets sur l’entraînement et les longs contextes#

Comme l’utilisation de la mémoire augmente de manière linéaire avec la longueur de la séquence, FlashAttention a rendu possible l’entraînement avec de longues fenêtres de contexte. L’article original a fait état d’une accélération de bout en bout de 15 % pour BERT-large avec une longueur de séquence de 512, par rapport au record d’entraînement MLPerf 1.1, d’une accélération de 3× pour GPT-2 avec une longueur de séquence de 1 024 et de 2,4× sur le benchmark Long Range Arena. Les mêmes économies s’appliquent à la vision, où un Transformer de vision (ViT) ou un Transformer de diffusion convertit des images et des vidéos haute résolution en longues séquences de jetons. Les séquences plus longues ont également amélioré la précision : l’article a présenté les premiers Transformers à dépasser le hasard sur Path-X, une tâche qui consiste à déterminer si deux points d’une image de 128 × 128 pixels, aplatie en une séquence de 16 384 pixels, sont reliés par un chemin.