L'article présente une nouvelle recherche intitulée « CODA : Rewriting Transformer Blocks as GEMM-Epilogue Programs », dont l'objectif principal est d'optimiser l'efficacité de l'entraînement des modèles Transformer, en particulier en résolvant des opérations « gourmandes en mémoire » qui semblent isolées mais s'accumulent pour entraîner des délais significatifs.

Auteur de l'article, source : Machine Heart

Le 22 mai, Tri Dao a partagé sur les réseaux sociaux un tweet de Han Guo. Il a également écrit : « Après une réécriture mathématique, il s'est avéré que tout ce qui constitue les Transformer n'est qu'une série de GEMM + epilogue (multiplication matricielle plus épilogue). Avec quelques primitives optimisées, les LLM (et les débutants) peuvent écrire des noyaux à vitesse lumineuse pour toutes les opérations Transformer ! »

Tri Dao est l'un des auteurs principaux de la série FlashAttention, et ce tweet pointe vers un article qu'ils ont publié ce jour-là : CODA.

Titre de l'article : CODA : Réécriture des blocs Transformer comme programmes GEMM-Epilogue
Adresse du papier : https://arxiv.org/abs/2605.19269
Adresse du code : https://github.com/HanGuo97/coda-kernels

Ce nom se lit comme « finale » et se prononce comme « CUDA ». Des chercheurs du MIT, de Princeton, de Together AI et de Meta tentent d’éliminer systématiquement les calculs épars, peu remarqués mais constamment gourmands en temps, dans l’entraînement des Transformers, à l’aide d’une nouvelle abstraction de programmation.

Le « impôt de la paresse » pour l'entraînement des grands modèles

Pour comprendre ce que CODA résout, il faut d'abord savoir où passe le temps lors de l'entraînement de grands modèles.

En entraînant un modèle de 1 milliard de paramètres de style LLaMA-3 sur une seule NVIDIA H100, la plupart des gens auraient l'intuition que le temps est principalement consacré aux multiplications matricielles et au calcul d'attention, car ce sont là les « véritables calculs ». Cette intuition est globalement correcte : les multiplications matricielles (GEMM) et l'attention occupent effectivement la majeure partie de la puissance de calcul.

Mais si vous ouvrez le profileur et que vous examinez attentivement, vous constaterez qu’un ensemble d’« opérateurs mineurs » consomment silencieusement du temps : la normalisation (RMSNorm), les fonctions d’activation (SwiGLU, RoPE), l’addition résiduelle, la réduction entre couches… Ils ont chacun une charge de calcul modeste, mais ils déplacent fréquemment de grands tenseurs intermédiaires entre la mémoire vidéo et la mémoire principale.

C'est ce qu'on appelle un « goulot d'étranglement de la bande passante mémoire » : comme un chef exceptionnel qui doit aller chercher les ingrédients dans un entrepôt éloigné, les ramener, puis les rapporter après chaque plat, au lieu de les avoir directement sur son plan de travail. Même si le chef est extrêmement rapide, le temps attendu pour le transport reste un gaspillage réel.

Pire encore, à mesure que les formats de précision réduite d’NVIDIA, tels que FP8 et FP4, accélèrent les calculs matriciels, le coût relatif de ces opérations de « transfert » augmente : les multiplications matricielles s’accélèrent, mais le coût de transfert des tenseurs vers et depuis la mémoire ne diminue pas proportionnellement.

Un ensemble de données dans l'article est très clair : lors de l'entraînement d'un modèle de 1 milliard de paramètres sur H100 avec TorchTitan, les opérations autres que les multiplications matricielles représentent une part significative du temps d'exécution end-to-end, et cette proportion devient encore plus marquée avec l'introduction de la précision FP8.

Les frameworks de programmation existants sont presque impuissants face à cela. PyTorch exprime le calcul Transformer comme une séquence d'opérateurs, avec des frontières claires entre eux. Ces frontières sont très favorables à la dérivation automatique (autograd), mais empêchent précisément l'optimisation par fusion entre opérateurs : chaque frontière d'opérateur correspond souvent à un écriture inutile en mémoire vidéo.

CODA : « La fin » cache un trésor

L'origine de CODA repose sur une observation simple.

Sur GPU, un noyau de multiplication matricielle高性能 (GEMM) est structuré en deux parties : la boucle principale (mainloop), qui effectue les calculs de multiplication et d'addition par blocs matriciels, et l'épilogue (epilogue), qui effectue des traitements finaux avant l'écriture des résultats dans la mémoire vidéo, tels que l'ajout de biais, la conversion de type et une mise à l'échelle simple.

La raison d’être de la fin réside dans le fait que la sortie de la multiplication matricielle est encore « vivante » dans les registres sur puce et n’a pas encore été écrite dans la mémoire vidéo globale. C’est une fenêtre dorée brève : si l’on peut effectuer davantage de calculs à ce moment-là, on peut entièrement éviter un aller-retour d’écriture et de lecture de la mémoire vidéo.

L'insight central de CODA est que de nombreuses opérations gourmandes en mémoire dans les Transformer peuvent être réparamétrées algébriquement et exécutées dans cette fenêtre de « conclusion ».

Cela nécessite un peu de compétence mathématique. Prenons le modèle le plus courant GEMM-RMSNorm-GEMM : le résultat d'une multiplication matricielle est soumis à une addition résiduelle, puis à une normalisation RMS, avant d'effectuer une autre multiplication matricielle. La méthode traditionnelle consiste à exécuter séquentiellement trois opérateurs indépendants, avec deux écritures intermédiaires en mémoire GPU.

L'équipe CODA a constaté que le facteur d'échelle de ligne r dans la normalisation RMS, étant un scalaire partagé par ligne, commute avec la multiplication matricielle suivante : l'application de r peut être reportée du « moment avant le deuxième GEMM » au « terme du deuxième GEMM ». Après ce report, la fin du premier GEMM ne nécessite plus que le calcul d'une « racine moyenne quadratique partielle » (partial RMS), regroupée par un noyau de réduction extrêmement léger, supprimant ainsi le calcul complet de RMSNorm.

Une telle réparamétrisation s'applique également aux opérations telles que SwiGLU, RoPE (Positional Encoding rotative) et la perte d'entropie croisée, et s'étend même à la rétropropagation. Un théorème dans l'article original démontre que tant que la phase finale de la propagation avant est « localement par blocs », la rétropropagation hérite automatiquement de la même structure. Consultez l'article original pour plus de détails.

Cinq « blocs » et un « langage Lego »

CODA n'est pas un noyau de fusion spécifique, mais un ensemble d'abstractions de programmation.

Il fixe la boucle principale GEMM optimisée par des experts, puis expose à la fin cinq types de primitives composites :

Transformation élément par élément (addition résiduelle, fonction d'activation, RoPE)
Chargement et stockage des vecteurs (diffusion des poids RMSNorm)
Chargement et stockage par blocs de matrices (enregistrement des activations intermédiaires pour la rétropropagation)
Réduction par blocs (racine moyenne quadratique locale, log-sum-exp par blocs)
État de transformation (statistiques max et sum-exp nécessaires pour la normalisation en ligne)

Avec ces cinq types de blocs, presque toutes les opérations d'une architecture Transformer standard, hors l'attention, peuvent être couvertes.

Ce qui est plus intéressant, c’est la tolérance de cette abstraction envers « qui écrit le code ». Dans les expériences, l’article évalue deux modes d’implémentation : l’un réalisé par des programmeurs humains, l’autre généré par Claude Code — en fournissant les primitives de CODA, quelques exemples et les journaux d’implémentation, l’IA effectue la majeure partie du code noyau, avec une supervision humaine légère.

Les performances des deux modes ont atteint un niveau élevé. Tri Dao a déclaré dans un tweet : « Un LLM et un débutant peuvent écrire un noyau à vitesse lumineuse », ce qui correspond exactement à la résonance pratique des résultats de l'expérience du papier.

Résultats de l'expérience

Les tests de référence de CODA choisissent des adversaires exigeants : cuBLAS avec torch.compile, ainsi que les noyaux Liger et FlashInfer optimisés pour les LLM.

Le papier évalue deux implémentations pour chaque noyau : CODA (LLM), générée par Claude Code, avec des instructions primitives, plusieurs exemples et un journal de techniques d'implémentation mis à jour en continu fournis par les chercheurs ; l'IA génère le code principal, tandis qu'une supervision humaine légère est appliquée ; CODA (Human), développée indépendamment par des programmeurs humains, utilisant les mêmes idées de réparamétrisation de haut niveau, mais sans dépendre du jeu d'primitives CODA. Les résultats des deux groupes sont comparés à des bibliothèques optimisées telles que cuBLAS + torch.compile, Liger Kernel et FlashInfer.

Au niveau du seul opérateur, avec le modèle typique GEMM-RMSNorm-GEMM, CODA dépasse la base cuBLAS + PyTorch pour les dimensions cachées des trois tailles de modèles correspondantes : 1B, 7B et 70B. Des combinaisons finales telles que SwiGLU, RoPE et l'entropie croisée présentent des performances similaires.

Les noyaux générés par LLM sont à la hauteur des versions écrites à la main sur la plupart des benchmarks, et dépassent légèrement ces dernières dans certaines configurations. Il s'agit d'une conclusion assez rare dans le domaine historiquement très exigeant de l'optimisation des noyaux GPU.

Les gains de la rétropropagation sont particulièrement marquants : le noyau de rétropropagation GEMM-Residual-PartialRMS-GEMM offre une accélération de 1,6 à 1,8 fois par rapport à la base, tandis que la rétropropagation de SwiGLU présente une amélioration d’environ 1,4 à 1,6 fois. Dans ce domaine, l’écart entre les LLM et les implémentations humaines est également minime. Cela n’est pas surprenant : la rétropropagation implique naturellement un accès plus important aux tenseurs intermédiaires, ce qui amplifie les bénéfices de la fusion en fin de chaîne ; de plus, la conception des primitives de CODA est suffisamment claire pour permettre aux modèles d’IA d’effectuer correctement les combinaisons.

Dans les benchmarks end-to-end de couches Transformer complètes, l'accélération en avant de CODA varie entre environ 5 % et 20 % selon les tailles, avec des effets plus marqués sur les grands modèles (correspondant à des dimensions cachées de 70B).

En ce qui concerne la précision numérique, la réparamétrisation de CODA ajuste le moment d'application du facteur d'échelle RMSNorm, mais les expériences montrent que son erreur numérique est comparable à celle de l'implémentation de référence PyTorch, et même inférieure dans certaines configurations — grâce à l'accumulateur de précision supérieure du boucle principale GEMM.

Ce que CODA peut faire : une fiche de référence pour clarifier les limites des capacités de CODA avant d'aborder une perspective plus large.

Couverture : presque tous les calculs, hors attention et embeddings de mots, dans la propagation avant et arrière des Transformers standards (comme l'architecture LLaMA), incluant RMSNorm, addition résiduelle, activation SwiGLU, codage de position rotatoire RoPE, perte d'entropie croisée, ainsi que les calculs de gradients inverses pour ces opérations.
Effet d'accélération : à l'échelle des dimensions cachées de 1B à 70B, une amélioration variable est observée au niveau des opérateurs par rapport à la base cuBLAS + torch.compile, avec un gain particulièrement marqué lors de la rétropropagation (certains noyaux atteignent plus de 1,6 fois la performance) ; l'accélération end-to-end de la couche Transformer complète se situe entre 5 % et 20 %, avec des effets plus prononcés sur les modèles de plus grande taille.
Qui peut utiliser : CODA, implémenté sur la base de CuTeDSL (DSL Python de NVIDIA CUTLASS), qui prend en charge deux méthodes d'écriture de noyaux : pour les programmeurs humains et pour les modèles d'IA, toutes deux atteignant des performances élevées.
Limitations actuelles : Seul le scénario à un seul GPU est actuellement pris en charge, sans entraînement distribué ; la réparamétrisation cible principalement l'architecture Transformer standard, et son applicabilité à d'autres architectures reste à vérifier.

Conclusion

CODA n'est pas un travail isolé. C'est une mise en œuvre concrète d'une catégorie d'idées : sur GPU, le véritable espace d'optimisation réside souvent moins dans « ce qui est calculé » que dans « comment les données sont déplacées ».

FlashAttention permet de stocker les calculs d'attention dans la mémoire intégrée au silicium, tandis que CODA tente d'y intégrer également la normalisation et les fonctions d'activation. Triton abaisse la barrière à l'écriture de noyaux personnalisés, et des projets comme ThunderKittens et TileLang explorent davantage cet espace à différents niveaux. Ces travaux convergent tous vers un même objectif : unifier véritablement, au sein d'un cadre programmable, la facilité d'expression du graphe d'opérateurs PyTorch et l'efficacité d'exécution proche du CUDA écrit à la main.

La dernière phrase du tweet de Tri Dao mérite d'être réfléchie : « Les LLM et les débutants peuvent écrire des noyaux à vitesse lumineuse pour toutes les opérations Transformer. » Derrière cela se cache une logique plus profonde : lorsque l'abstraction de la programmation est suffisamment bien conçue, les modèles d'IA peuvent eux-mêmes participer à l'optimisation de leur propre infrastructure d'entraînement. Ce cycle est ce qui rend le CODA le plus fascinant.

Du point de vue de cette perspective, le nom « CODA » pourrait avoir une signification plus profonde. Dans la musique classique, la coda est le passage final qui conclut une œuvre. Ici, elle représente la « conclusion » du noyau GEMM — et rédiger cette conclusion pourrait bien être le prochain chapitre essentiel pour améliorer l’efficacité des systèmes d’entraînement Transformer.