L'attention fusionnée

Dix-septième article de la série Smallem. Précédent : À durée égale.

Deux problèmes sont restés ouverts dans les articles précédents.

Le premier est un mur. À l'article 15, allonger le contexte à 4 096 positions faisait manquer de mémoire une carte de 24 Gio.

Le second est un écart. À l'article 11, PyTorch faisait un pas d'entraînement en 75 ms avec les mêmes opérations que nous, et en 45,5 ms avec son attention fusionnée. Nous étions à 157 ms.

Les deux ont la même cause, et le même remède : ne plus écrire la matrice des scores en mémoire.

Ce remède s'appelle FlashAttention (Dao et al., 2022 puis FlashAttention-2, en 2023). Il ne change pas le modèle, seulement la façon de le calculer.

Ce que l'attention écrit en mémoire

Une couche d'attention calcule, pour chaque tête, les scores de toutes les requêtes face à toutes les clés, puis leur softmax, puis le mélange des valeurs :

S = Q · Kᵀ × échelle        [T, T]
P = softmax_causal(S)       [T, T]
O = P · V                   [T, D]

Q, K et V sont des matrices [T, D], avec D = 64. S et P sont des matrices [T, T] : avec 512 positions, huit fois plus grosses que Q. Elles passent toutes deux par la mémoire de la carte, et la tape les garde pour la passe arrière, couche après couche.

Leur taille croît comme le carré du contexte. À 4 096 positions, un tenseur de scores pour deux séquences et huit têtes occupe 1 Gio, et il y en a plusieurs par couche. Le mur vient de là.

Une opération plutôt que trois

Jusqu'ici, l'attention était trois opérations sur la tape : un produit, le softmax causal de l'article 8, un second produit. Pour qu'un kernel puisse tout faire d'un coup, il faut d'abord une seule opération :

pub struct FusedAttention {
    pub scale: f64,
}

Trois entrées, q, k, v, une sortie. Comme toutes les opérations de l'article 12, elle a une définition écrite avec les primitives : les trois étapes de toujours. C'est cette définition qui tourne sur le CPU, et c'est contre elle que le GPU est testé.

Il manquait une chose au contrat des opérations. La passe arrière aura besoin, pour chaque requête, d'un nombre que la passe avant calcule en passant : nous y reviendrons. Jusqu'ici, une opération rendait son résultat et rien d'autre. Elle peut maintenant rendre en plus ce que sa rétropropagation voudra relire, et la tape le lui rend :

fn run_saving(&self, op: &O, inputs: &[&Self::Tensor])
    -> (Self::Tensor, Option<Self::Tensor>) {
    (self.run(op, inputs), None)
}

Par défaut, rien. Aucune autre opération n'a eu à changer.

Le softmax en ligne

Un kernel fusionné prend un bloc de requêtes et parcourt les clés par tuiles. Il n'écrit jamais S en entier en mémoire : il garde seulement la tuile du moment, dans la mémoire partagée du bloc.

L'obstacle est le softmax. Pour normaliser une ligne, il faut connaître son maximum et la somme de ses exponentielles, donc avoir vu toute la ligne. Or on la découvre tuile après tuile.

Le softmax en ligne (Milakov et Gimelshein, 2018) tient à jour un maximum m et une somme l. Quand une nouvelle tuile fait monter le maximum de m à m', tout ce qui a été accumulé l'a été avec des exp(s - m). Il suffit de le multiplier par exp(m - m') pour le remettre à jour :

pour chaque tuile de clés :
    s  = q · kᵀ × échelle            le futur reçoit moins l'infini
    m' = max(m, max(s))
    p  = exp(s - m')
    l  = l × exp(m - m') + Σ p
    o  = o × exp(m - m') + p · v
    m  = m'
o = o / l

La division attend la toute fin. Dans le kernel, chaque thread tient deux lignes, et les seize threads d'une même ligne réunissent leur maximum et leur somme par échanges de registres, comme les réductions de l'article 11 :

let new_max0 = max0.max(half_warp_max(s00.max(s01)));
let (p00, p01) = ((s00 - new_max0).exp(), (s01 - new_max0).exp());
let fix0 = (max0 - new_max0).exp();
sum0 = sum0 * fix0 + half_warp_sum(p00 + p01);
o00 *= fix0;

Deux économies viennent en prime. Une tuile qui ne contient que du futur n'est jamais parcourue : pour une séquence entière, c'est la moitié du travail. Et aucun masque n'existe : chaque thread sait, par sa position, si une clé est dans le futur.

La passe arrière : recalculer plutôt que garder

La passe arrière a besoin des poids P, qu'on n'a pas gardés. On les recalcule.

C'est là que sert le nombre gardé par la passe avant. Pour chaque requête, le kernel a écrit en mémoire L = m + ln(l), le logarithme de la somme de la ligne. Avec lui, un poids se retrouve sans refaire le softmax : p = exp(s - L). On recalcule les scores, tuile par tuile, et on n'a jamais besoin de la ligne entière.

Le reste est la règle du softmax de l'article 6, réécrite pour qu'aucune matrice [T, T] n'apparaisse :

dV = Pᵀ · dO
dS = P × (dO · Vᵀ - Δ)        Δ = Σ dO × O, une valeur par ligne
dQ = dS · K × échelle
dK = dSᵀ · Q × échelle

La somme Σ dP × P de la règle du softmax demanderait toute la ligne de P. Mais O vaut justement P · V, donc cette somme vaut Σ dO × O : une valeur par ligne, qu'un kernel de réduction calcule avant tout le reste.

FlashAttention-2 calcule les trois gradients dans un seul kernel. Chaque bloc prend une tuile de clés, accumule dK et dV dans ses registres, et ajoute sa part de dQ en mémoire globale par des additions atomiques. Nous avons préféré deux kernels : l'un pour dK et dV, l'autre pour dQ, qui parcourt les tuiles dans l'ordre de la passe avant. Il refait trois produits, mais il n'a aucune addition atomique. Le résultat ne dépend donc pas de l'ordre dans lequel la carte exécute les blocs, et deux entraînements identiques restent identiques.

Un premier kernel plus lent que l'ancien

Les tests sont passés du premier coup, sur des formes choisies pour ne pas tomber juste : 70 positions, des têtes de 5 et de 3 composantes, une seule requête face à 70 positions comme en génération. Puis la première mesure, en précision mixte, est tombée :

Plus lente. La raison est l'article 11. L'ancienne attention passait ses produits par la multiplication de matrices sur Tensor Cores, en bf16. Le nouveau kernel faisait tout en f32, sur les unités ordinaires. Économiser des allers et retours en mémoire ne rattrapait pas un calcul bien moins rapide.

Il fallait donc une seconde version, sur les Tensor Cores.

Sur les Tensor Cores

L'instruction mma.sync de l'article 11 multiplie une tuile 16 × 16 par une tuile 16 × 8, chaque thread du warp tenant dans ses registres quelques éléments de chacune. Le nouveau kernel donne seize requêtes à chaque warp, qui les garde en registres pendant tout le parcours des clés.

Une coïncidence de disposition rend le reste élégant. Le résultat d'une multiplication laisse à chaque thread les éléments des lignes g et g + 8, colonnes 2t et 2t + 1. C'est exactement ce qu'attend un fragment de l'opérande de gauche pour la multiplication suivante. Les poids passent donc des scores au produit par les valeurs sans toucher la mémoire partagée :

pub fn fragment_a(left: [f32; 4], right: [f32; 4]) -> [u32; 4] {
    [
        bf16_pair(left[0], left[1]),
        bf16_pair(left[2], left[3]),
        bf16_pair(right[0], right[1]),
        bf16_pair(right[2], right[3]),
    ]
}

Le compilateur ne déroule pas

La première version écrite à la main tenait en quelques boucles sur des tableaux : huit tuiles de scores, quatre pas de seize. Elle passait les tests, et l'attention prenait 36 ms par pas.

Le PTX a donné la raison. Le compilateur de kernels n'avait pas déroulé les boucles : il restait 8 instructions mma.sync au lieu de 64, dans une boucle, et les tableaux indexés par un compteur étaient placés en mémoire locale, bien plus lente que les registres. L'article 11 l'avait déjà noté pour la multiplication de matrices, qui écrit pour cette raison seize variables séparées.

Plutôt que seize variables, une macro déroule la boucle à la compilation :

macro_rules! unroll {
    ($var:ident in [$($value:literal),*] $body:block) => {
        $({
            const $var: usize = $value;
            $body
        })*
    };
}

Chaque indice devient une constante, chaque case de tableau une variable. Le PTX montre alors les 64 instructions, et plus aucun accès à la mémoire locale.

Les mesures

Le pas d'entraînement entier, en précision mixte (banc d'essai, modèle modern, 16 séquences de 512 positions) :

L'attention elle-même passe de 38 à 15 ms par pas. Le modèle GPT-2 de l'article 11, que l'on comparait à PyTorch, passe de 157 à 127 ms : l'écart avec PyTorch et son attention fusionnée (45,5 ms) passe de 3,5 à 2,8 fois. Ce qui reste n'est plus dans l'attention.

Le tableau de l'article 15, refait (8 192 tokens par pas, temps mesurés au profileur) :

La mémoire ne dépend plus du contexte : à nombre de tokens égal, elle reste à 7 Gio. Le mur a disparu, et 8 192 positions tiennent sans peine. Le temps croît encore, parce que le calcul des scores, lui, reste quadratique : la fusion supprime les allers et retours en mémoire, pas les multiplications.

Une remarque d'honnêteté sur la dernière colonne. En trois opérations, l'attention tient désormais à 4 096 positions, alors que l'article 15 y manquait de mémoire. Elle n'est plus tout à fait l'ancienne : regroupée en une seule opération, elle ne garde plus que les poids, quand l'ancienne tape gardait aussi les scores. Le seul fait d'avoir une opération au lieu de trois a donc déjà repoussé le mur d'un cran.

La génération, elle, ne change presque pas : 165 tokens par seconde contre 161. Avec un cache, chaque nouveau token ne compare qu'une requête aux clés passées, et le temps part ailleurs. Et en mode glouton, les deux chemins écrivent exactement le même texte.

Entraîner avec

Trois entraînements de 6 000 pas, avec 8 192 tokens par pas comme toujours, découpés en séquences de plus en plus longues :

La première ligne est le contrôle. Le même entraînement avec l'ancienne attention donnait 0,9592 : l'écart est dans le bruit. Les arrondis en bf16 des poids et des gradients ne coûtent rien de mesurable, et l'entraînement prend 14 minutes au lieu de 16.

La troisième ligne est la validation de l'article : un entraînement à 4 096 positions, qui tient dans 7 Gio là où il ne tenait pas dans 24.

Il apprend pourtant moins bien. À nombre de tokens égal, la perte monte avec la longueur des séquences : +0,009 à 2 048, +0,033 à 4 096. Deux explications sont plausibles, et aucune n'a été vérifiée. Si les pages de FineWeb-Edu sont le plus souvent plus courtes que 4 096 tokens, ce qui n'a pas été mesuré, une longue séquence enchaîne plusieurs documents sans rapport, et le contexte supplémentaire n'apprend rien. Et deux séquences par lot au lieu de seize, ce sont des tokens moins variés à chaque pas : le gradient est plus bruité. Les normes du gradient vont dans ce sens : des pics à 12 et 15 en seconde moitié d'entraînement, quand ceux de l'entraînement à 512 positions restaient près de 1.

C'est la réponse à une question laissée à l'article 15 : l'attention clairsemée de DeepSeek valait-elle d'être écrite à 2 048 positions ? Avec l'attention fusionnée, l'attention ne pèse plus qu'un quart du pas à cette longueur : en divisant son coût par huit, on gagnerait au mieux un cinquième du temps. Et un contexte plus long n'aide pas encore notre modèle. Il n'y a donc toujours rien à gagner.

Ce qui reste

L'attention latente de l'article 13 n'en profite pas. Ses requêtes et ses clés ont 96 composantes, et nos kernels s'arrêtent à 64 : au-delà, la mémoire partagée d'un bloc ne suffit plus pour la passe arrière. Elle passe donc encore par les trois opérations.

L'écart avec PyTorch, ensuite. Il ne tient plus à l'attention, mais au reste : PyTorch garde aussi ses activations en 16 bits, ce qui divise par deux le trafic mémoire de toutes les autres opérations.

Ce qu'il faut retenir

• Le mur de mémoire de l'attention venait de matrices [T, T] écrites en mémoire alors qu'elles ne sont qu'un intermédiaire. Un calcul réorganisé n'en écrit aucune.
• Le softmax en ligne normalise une ligne qu'on découvre par morceaux : il suffit de corriger ce qui a été accumulé à chaque nouveau maximum.
• Recalculer coûte souvent moins cher que garder en mémoire : la passe arrière refait les scores à partir d'un nombre par ligne.
• Un kernel fusionné n'est pas rapide par nature. Le nôtre, en f32, perdait contre trois opérations sur Tensor Cores.
• Lire le PTX : une boucle non déroulée y coûtait un facteur 2,4.

Il reste à prendre du recul sur l'ensemble : Bilan.