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

Together AI a lancé ThunderKittens — un langage de programmation (DSL) pour l'écriture de code optimisé sur GPU. Le projet se positionne comme une tentative de simplifier le développement des réseaux de neurones, qui nécessite actuellement une connaissance approfondie de l'architecture matérielle.

Pourquoi c'était nécessaire

L'écriture de code efficace pour GPU est de la magie noire pour la plupart des ingénieurs ML. Sur des puces comme NVIDIA H100, il existe des tensor cores spécialisés qui fournissent 94 % de toute la puissance de calcul. Mais pour les utiliser, vous devez écrire en CUDA — un langage bas niveau complexe que seule une poignée de professionnels comprend.

En 2023, la solution populaire FlashAttention2 a montré qu'il était possible de simplifier considérablement le travail avec l'attention dans les transformers. Mais ce n'était qu'un cas particulier — un algorithme pour un opérateur spécifique. Les développeurs avaient besoin d'un moyen universel d'écrire du code rapide sans ce CUDA infernal, qui fonctionnerait pour différents kernels.

Comment fonctionne ThunderKittens

ThunderKittens se situe entre deux extrêmes. D'un côté, il y a le CUDA pur — très rapide, mais très complexe, avec une barrière d'entrée élevée. De l'autre côté, il y a Triton, qui masque les détails du GPU et simplifie la vie, mais parfois ne peut pas extraire les performances maximales du matériel.

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

L'objet fondamental dans ThunderKittens est la tile (tuile), une matrice de taille qui s'ajuste au tensor core. C'est une unité de calcul qui permet de charger complètement le noyau spécialisé et d'extraire tout le potentiel du matériel moderne.

Performance et résultats

Sur les puces modernes A100 et RTX 4090, ThunderKittens correspond en vitesse à FlashAttention2 — tandis que le code est quelque peu plus court et plus clair. Sur H100, cette solution gagne : plus rapide que FlashAttention2 à la fois en avant et en rétropropagation. Autrement dit, il n'y a aucun compromis entre la propreté du code et la vitesse.

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

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

Un point intéressant : les ingénieurs qui ont suivi un cours CUDA de seulement deux heures ont pu écrire leur propre code en ThunderKittens. Cela montre que le langage simplifie réellement le développement, tout en ne masquant pas les détails du matériel.

Projet ouvert et formation

Les auteurs de Together AI disent honnêtement : c'est un art project. Ne vous attendez pas à des mises à jour régulières et à un support pour tous vos commentaires dans le suivi des bogues. Le projet a été publié en accès ouvert parce que les développeurs pensent qu'il est intéressant de partager des idées et des outils avec la communauté.

Avec ThunderKittens, ils ont également lancé NanoGPT-TK — une version du célèbre projet NanoGPT d'Andrej Karpathy, où les kernels de calcul principaux ont été réécrits en ThunderKittens. Cela a été fait spécifiquement pour l'enseignement et la démonstration des capacités.

NanoGPT a longtemps été reconnu comme l'un des meilleurs projets en IA pour comprendre comment fonctionne l'entraînement des transformers à partir de zéro.

Ce que cela signifie

ThunderKittens montre qu'il existe un véritable écart dans l'IA entre la commodité 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 légèrement plus complexe s'il leur donne du contrôle et une vraie vitesse en pratique.

Pour les ingénieurs ML, cela pourrait signifier que dans le futur, le portage d'un modèle entraîné d'une puce à une autre sera plus facile — il suffira de réécrire quelques kernels GPU au lieu 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 nécessiter un cours CUDA d'un mois.