Une attention exacte, accélérée en n'écrivant jamais la gigantesque grille.

SRC·16 Source
Les mêmes calculs d'attention — simplement jamais écrits.

Les mêmes calculs d'attention — simplement jamais écrits.

Le vrai coût de l'attention, ce n'est pas la multiplication — c'est de trimballer une table géante de scores vers la mémoire lente et retour. Flash attention ne change rien aux calculs et tout sur l'endroit où ils se font. Même réponse, en une fraction du temps. Comme résoudre un cube à toute vitesse : même casse-tête, mêmes mouvements — l'avantage, c'est la pure méthode.
Le goulot d'étranglement, ce ne sont pas les calculs. C'est le transport.

Le goulot d'étranglement, ce ne sont pas les calculs. C'est le transport.

L'attention classique construit la table complète de chaque mot évalué face à tous les autres — puis écrit cette grille géante en mémoire lente, la relit pour normaliser, et la relit encore pour pondérer les valeurs. Trois allers-retours avec le plus gros objet du modèle, pendant que la puce termine les calculs et attend. Comme un tapis déroulé : trop grand pour la pièce, alors il reste dans le couloir — et chaque coup d'œil exige un nouveau trajet.
Une puce rapide, étranglée par un tuyau lent.

Une puce rapide, étranglée par un tuyau lent.

Tmax ⁣(FLOPscompute rate, bytes movedbandwidth)T \approx \max\!\left(\dfrac{\text{FLOPs}}{\text{compute rate}},\ \dfrac{\text{bytes moved}}{\text{bandwidth}}\right)
Un processeur fait de l'arithmétique bien plus vite qu'il ne peut aller chercher les nombres à traiter. Le temps total est donc le plus lent des deux : faire les calculs, ou déplacer les octets. L'attention est limitée par la mémoire — les calculs sont prêts, pas les données — alors réduire les trajets, et non les calculs, voilà ce qui l'accélère. Comme une pompe puissante sur un tuyau fin : elle pourrait déplacer un océan, mais c'est le tuyau qui décide du débit.
Alors ne chargez pas tout. Travaillez un bloc à la fois.

Alors ne chargez pas tout. Travaillez un bloc à la fois.

Sij=QiKjd,Qi,KjRB×dS_{ij} = \dfrac{Q_i K_j^{\top}}{\sqrt{d}}, \qquad Q_i,\, K_j \in \mathbb{R}^{B\times d}
La solution commence simplement : ne jamais rapprocher toute la grille. Découpez les Keys et les Values en blocs assez petits pour tenir sur le minuscule bloc-notes rapide de la puce. Amenez un bloc, faites sa part du travail, lâchez-le, allez chercher le suivant. En clair : chaque bloc de B lignes tient en mémoire rapide, donc vous ne traitez qu'une paire de blocs à la fois. Comme porter du bois de chauffage : on ne traîne pas tout le tas jusqu'à l'âtre — on en apporte une brassée.
Mais softmax a besoin de toute la ligne d'un coup.

Mais softmax a besoin de toute la ligne d'un coup.

softmax(x)i=eximj=1nexjm,m=maxjxj\mathrm{softmax}(x)_i = \dfrac{e^{\,x_i - m}}{\sum_{j=1}^{n} e^{\,x_j - m}}, \qquad m = \max_{j} x_j
Voilà le hic. Pour transformer les scores en poids, softmax divise chacun par le total sur chaque key — et par sécurité, il soustrait d'abord le plus grand score de la ligne, pour que les exponentielles n'explosent jamais. Le max comme la somme ont besoin de toute la ligne. Mais vous n'avez qu'un seul bloc. Comment normaliser face à des nombres que vous n'avez pas vus ? Comme verser des parts égales d'un seul pichet : impossible de remplir le premier verre avant de savoir combien il y en a.
L'astuce : garder un maximum courant, et tout recaler au fil de l'eau.

L'astuce : garder un maximum courant, et tout recaler au fil de l'eau.

mnew=max(mold,m~)new=emoldmnewold+ieximnew\begin{aligned} m_{\text{new}} &= \max(m_{\text{old}},\, \tilde m) \\ \ell_{\text{new}} &= e^{\,m_{\text{old}} - m_{\text{new}}}\,\ell_{\text{old}} + \sum_{i} e^{\,x_i - m_{\text{new}}} \end{aligned}
Pas besoin de toute la ligne d'abord — il faut un cumul courant. Gardez le plus grand score vu jusqu'ici et le total en cours. Quand un nouveau bloc apporte un score plus grand, ramenez l'ancien total à la nouvelle échelle avec un seul facteur de réduction, puis intégrez le bloc ; la sortie se recale de la même façon. Le résultat est identique à tout faire d'un coup. Comme la barre du saut en hauteur : relevez-la quand un meilleur sauteur arrive, et remesurez chaque marque précédente face à la nouvelle hauteur.
Même réponse. La grille géante n'a jamais existé.

Même réponse. La grille géante n'a jamais existé.

memory: O(n2)  O(n),oflash=oexact\text{memory: } O(n^2)\ \longrightarrow\ O(n), \qquad \mathbf{o}_{\text{flash}} = \mathbf{o}_{\text{exact}}
Mettez tout ensemble et la table n×n n'est jamais stockée — seulement un maximum courant, une somme courante et la sortie, tous petits. La mémoire passe de quadratique à linéaire, et l'accélération vient de bien moins de trajets vers la mémoire lente, pas d'une réduction des calculs. Et ce n'est pas une approximation : le résultat est une attention exacte. Comme plier une grue en origami : une feuille, la forme achevée, aucune chute sur la table.
🌱 Ce n'était peut-être jamais la pensée qui était lente.

🌱 Ce n'était peut-être jamais la pensée qui était lente.

Flash attention n'ajoute aucune idée nouvelle — elle respecte simplement où vivent les calculs, réorganisant les mêmes pensées pour le corps qui les abrite. Alors, quelle part de la vitesse d'un esprit tient à la pensée elle-même — et quelle part tient seulement à la rapidité avec laquelle il atteint ce qu'il sait déjà ?
touchez →balayez ↑ pour approfondirbalayez ↓ pour quitter