Calculateur de mémoire d'attention
Données
| Longueur de séquence | 4 096 |
|---|---|
| Têtes d'attention | 32 |
| Taille de lot | 1 |
| Précision | FP16 / BF16 (2 octets) |
Calculateur de mémoire d'attention
Estimez la mémoire de la matrice de scores d'attention matérialisée dans l'attention standard d'un transformeur, à partir de la longueur de séquence, du nombre de têtes, de la taille de lot et des octets par élément — le terme quadratique que FlashAttention élimine.
Données
Charge de travail
Résultats
Saisissez une valeur pour afficher les résultats.
La mémoire d'attention
L'attention standard d'un transformeur compare chaque jeton avec tous les autres, construisant une matrice de scores avant le softmax. Cette matrice est carrée en longueur de séquence, donc sa mémoire croît avec le carré du contexte — le fameux coût quadratique de l'attention. Ce calculateur estime la taille de cette matrice matérialisée à partir de la longueur de séquence, du nombre de têtes d'attention, de la taille de lot et des octets utilisés par score. C'est précisément la mémoire que FlashAttention évite de stocker.
Le terme quadratique
Pour une séquence de jetons, l'attention forme une matrice de scores : la ligne , colonne indique à quel point le jeton prête attention au jeton . La stocker demande une mémoire proportionnelle à . C'est différent des poids du modèle ou du cache clé-valeur, qui croissent linéairement avec le contexte — la matrice de scores croît avec son carré, de sorte qu'à long contexte elle devient le tampon intermédiaire dominant et la raison pour laquelle l'attention naïve épuise la mémoire.
La formule
Chaque tête construit sa propre matrice de scores et chaque séquence du lot porte sa propre copie, donc la mémoire en octets vaut
Aoctets=B⋅H⋅s2⋅eoù est la taille de lot, le nombre de têtes d'attention, la longueur de séquence et les octets par score stocké. La division par donne des gigaoctets. Le carré sur est ce qui rend le long contexte coûteux : le nombre de têtes et le lot ne multiplient que linéairement.
Ce qu'est cette mémoire
La valeur mesure la matrice de scores matérialisée que l'attention standard écrit en mémoire entre le produit requête-clé et la somme sur les valeurs pondérée par le softmax. C'est le tampon d'activations qu'une implémentation naïve doit conserver. Sa persistance dépend du contexte : dans une simple passe avant, il est transitoire et peut être réutilisé entre les couches, tandis que les frameworks d'entraînement peuvent en garder une copie par couche pour la passe arrière, sauf si le checkpointing ou un noyau fusionné intervient.
FlashAttention
FlashAttention calcule la sortie identique sans jamais écrire la matrice complète. Il parcourt les clés et les valeurs par petits blocs, en maintenant des statistiques de softmax courantes, de sorte qu'il n'a besoin que d'un bloc en mémoire à la fois, ce qui fait passer le stockage de quadratique à linéaire en longueur de séquence. La matrice que ce calculateur dimensionne est exactement ce que FlashAttention refuse de stocker — le résultat obtenu ici est donc un bon indicateur de la mémoire qu'un noyau d'attention fusionné économise à une longueur de contexte donnée.
Exemple chiffré
Prenons une seule séquence de 4 096 jetons avec 32 têtes d'attention en précision 16 bits :
Aoctets=1×32×40962×2=1073741824soit environ 1,07 Go pour la matrice de scores d'une seule couche. Doublez le contexte à 8 192 jetons et le terme au carré prend le dessus : la même expression donne environ 4,29 Go, quatre fois plus pour deux fois la longueur. Le versant clé-valeur du même budget d'attention est traité dans le Calculateur de taille du cache KV, et la part des poids du modèle dans le Calculateur de VRAM pour l'inférence LLM.
Questions fréquentes (FAQ)
Pourquoi la mémoire d'attention est-elle quadratique en longueur de séquence ?
L'attention compare chaque jeton avec tous les autres, produisant une matrice de scores dont les deux dimensions sont la longueur de séquence. Stocker cette matrice complète demande donc une mémoire proportionnelle au carré de la longueur de séquence.
Doubler le contexte quadruple la matrice de scores, ce qui explique pourquoi l'attention standard devient limitée par la mémoire à long contexte — le terme quadratique dépasse le coût linéaire des poids et des activations.
Comment FlashAttention réduit-il cette mémoire ?
FlashAttention calcule le même résultat sans jamais matérialiser la matrice de scores complète. Il diffuse l'attention sur de petits blocs de clés et de valeurs, en conservant des statistiques de softmax courantes, de sorte qu'il ne garde qu'un bloc à la fois.
La matrice quadratique que ce calculateur mesure est exactement la mémoire que FlashAttention évite de stocker, ce qui lui permet de passer à des séquences bien plus longues. L'arithmétique est inchangée ; seul le stockage intermédiaire passe de quadratique à linéaire.
Cette mémoire est-elle par couche ou pour tout le modèle ?
Cela dépend de l'implémentation. La valeur donnée ici est la taille de la matrice de scores d'une seule couche. Dans une passe avant naïve, ce tampon peut être libéré et réutilisé d'une couche à l'autre, il est donc transitoire plutôt que multiplié par le nombre de couches.
Pendant l'entraînement, en revanche, les frameworks peuvent conserver les tenseurs d'attention de chaque couche pour la passe arrière, sauf si le checkpointing de gradient ou FlashAttention est utilisé, auquel cas le total peut croître avec le nombre de couches.
Mentions légales
Cette estimation ne couvre que la matrice de scores matérialisée à la précision indiquée ; elle exclut les tenseurs de requêtes, de clés et de valeurs, la projection de sortie et les autres activations. Elle décrit l'attention standard — des noyaux comme FlashAttention ne stockent pas du tout cette matrice.