ThunderKittens de Together AI : un nouveau langage pour des kernels GPU efficaces

Together AI a publié ThunderKittens, un langage dédié (DSL) pour écrire du code GPU optimisé. Le projet se présente comme une tentative de simplifier le développement de réseaux de neurones, qui exige aujourd'hui une connaissance approfondie de l'architecture matérielle.

Pourquoi c'était nécessaire

Écrire du code GPU efficace relève de la magie noire pour la plupart des ingénieurs ML. Les puces comme la NVIDIA H100 disposent de tensor cores spécialisés, qui fournissent 94% de la puissance de calcul totale. Mais pour les exploiter, il faut écrire en CUDA, un langage complexe et de bas niveau que seule une poignée de personnes maîtrise dans l'industrie.

En 2023, le populaire FlashAttention2 a montré qu'on pouvait nettement simplifier le travail sur l'attention dans les transformers. Mais ce n'est qu'un cas particulier : un algorithme pour un opérateur précis. Les développeurs ont besoin d'une méthode universelle pour écrire du code rapide sans cet enfer de CUDA, valable pour différents kernels.

Comment fonctionne ThunderKittens

ThunderKittens se place entre deux extrêmes. D'un côté, le CUDA pur : très rapide, mais très complexe, avec un ticket d'entrée élevé. De l'autre, Triton, qui masque les détails du GPU et facilite la vie, mais ne parvient parfois pas à tirer le maximum de performances du matériel.

ThunderKittens propose une troisième voie. Son API ressemble à PyTorch, ce qui parle à tous les développeurs ML habitués aux opérations sur les tenseurs. En même temps, il reste assez transparent pour que le développeur comprenne ce qui se passe au niveau matériel. Les auteurs affirment que si vous connaissez CUDA, vous pouvez « compiler » ThunderKittens dans votre tête.

L'objet fondamental de ThunderKittens est le tile (tuile) : une matrice dont la taille tient dans un tensor core. C'est l'unité de calcul qui permet de charger entièrement le cœur spécialisé et d'exploiter tout le potentiel du matériel moderne.

Performances et résultats

Sur les puces A100 et RTX 4090, ThunderKittens égale FlashAttention2 en vitesse, avec un code un peu plus court et plus lisible. Sur H100, il l'emporte : il est plus rapide que FlashAttention2, aussi bien en passe avant qu'en rétropropagation. Autrement dit, il n'y a aucun compromis entre propreté du code et vitesse.

Les auteurs ont déjà écrit en ThunderKittens plusieurs kernels pour d'autres algorithmes :

• Based, une version optimisée de l'attention linéaire
• Hedgehog et d'autres kernels spécialisés pour les transformers
• Plusieurs solutions qui se distinguent favorablement des versions Triton en vitesse

Un point intéressant : des ingénieurs qui n'avaient suivi qu'un cours de CUDA de deux heures ont réussi à écrire leur propre code en ThunderKittens. Cela suggère que le langage simplifie réellement le développement, même s'il ne masque pas les détails du matériel.

Projet ouvert et pédagogie

Les auteurs de Together AI le disent franchement : c'est un art project. N'attendez pas de mises à jour régulières ni de prise en charge de toutes vos plaintes dans le bug tracker. Le projet est publié en accès libre parce que les développeurs trouvent intéressant de partager idées et outils avec la communauté.

Avec ThunderKittens, ils ont publié NanoGPT-TK, une version du projet culte NanoGPT d'Andrej Karpathy dont les principaux kernels de calcul ont été réécrits en ThunderKittens. Elle a été conçue spécifiquement pour l'enseignement et pour démontrer les capacités du langage. NanoGPT est depuis longtemps considéré comme l'un des meilleurs projets d'IA pour comprendre comment fonctionne l'entraînement des transformers à partir de zéro.

Ce que cela signifie

ThunderKittens montre qu'il existe dans l'IA un véritable fossé entre le confort des abstractions (PyTorch, Triton) et le contrôle du matériel réel (CUDA). Il s'avère que les développeurs sont prêts à écrire un code un peu plus complexe si cela leur donne du contrôle et une vitesse réelle en pratique.

Pour les ingénieurs ML, cela peut signifier qu'il sera plus facile à l'avenir de porter un modèle entraîné d'une puce à une autre : il suffirait de réécrire quelques kernels GPU, plutôt que de refaire la moitié de l'infrastructure. Pour les chercheurs, c'est un outil pour expérimenter rapidement avec des algorithmes spécialisés, sans exiger un cours de CUDA de plusieurs mois.