Comment dompter Blackwell sans devenir fou à cause de CUDA

Lorsqu'il s'agit d'écrire des noyaux GPU, la plupart d'entre nous imaginent immédiatement des pages interminables de code CUDA C++, jonglant avec des pointeurs et le débogage fastidieux des fuites mémoire ou des conditions de course. La situation devient encore plus complexe lorsque des monstres comme l'architecture NVIDIA Blackwell (sm_100a) arrivent sur le marché. Il y a tellement d'asynchronisme et de nouveaux moteurs de données à l'intérieur que les anciennes approches « lancez simplement un thread vers un cœur » ne permettent plus d'extraire même la moitié des performances.

Récemment, je suis tombé sur le projet modern-gpu-programming-for-mlsys de l'équipe MLC-AI. En essence, ce n'est pas juste un dépôt, mais un tutoriel interactif vivant sur la programmation GPU moderne. Les gars ont mis le paquet : ils vous apprennent à écrire des noyaux SOTA (state-of-the-art) en utilisant Python et une nouvelle DSL appelée TIRx.

Pourquoi ce n'est pas qu'un simple tutoriel

D'habitude, les cours GPU restent coincés sur la multiplication matricielle classique des années 2010. Ici, l'accent est mis sur le matériel de demain. Le centre de l'attention est l'architecture Blackwell, sa hiérarchie mémoire, ses cœurs tensoriels et ses mécanismes de mouvement de données asynchrones.

Les auteurs proposent un chemin depuis la compréhension du matériel jusqu'à l'écriture de noyaux qui fonctionnent réellement dans les systèmes ML modernes. Au lieu de se battre avec du C++ bas niveau, vous écrivez en TIRx (Tensor IR next). C'est une extension d'Apache TVM qui vous permet de décrire la logique du noyau en Python tout en maintenant le contrôle sur les registres, le cache et les moteurs tensoriels.

Quoi de neuf dans ce tutoriel

Tout le contenu est divisé en blocs logiques qui élèvent progressivement la barre de difficulté.

D'abord, ils couvrent le modèle d'exécution et de mémoire. Pas de définitions ennuyeuses ici, mais beaucoup de pratique avec TMA (Tensor Memory Accelerator) et la compréhension de comment les données atteignent réellement les unités de calcul. Curieusement, les auteurs s'attardent en détail sur le modèle Roofline — cela aide à comprendre sur quoi votre code est étranglé : la bande passante mémoire ou la puissance brute des cœurs.

Ensuite vient l'immersion profonde dans TIRx. Vous apprenez à travailler avec les tuiles (TileLayout), les axes et le mécanisme de swizzle (mélange astucieux des données en mémoire pour éviter les conflits de banks). À la fin, vous obtenez un noyau GEMM (multiplication matricielle) fonctionnel.

Après cela vient la vraie « viande ». Les auteurs montrent comment transformer un GEMM ordinaire en un monstre haute performance. Les pipelines TMA, la spécialisation warp et le regroupement des CTA en clusters entrent en jeu. Si ces termes vous semblent être des incantations, le tutoriel les rangera joliment sur des étagères.

Dans le finale, vous créez Flash Attention 4. C'est exactement ce sur quoi fonctionnent les grands modèles de langage modernes. Le noyau inclut deux étapes MMA (Matrix-Matrix Accumulation), un softmax en ligne et le support GQA (Grouped Query Attention).

Stack technique et fonctionnalités

Le projet est étroitement lié à l'écosystème Apache TVM. Pour exécuter les exemples, vous aurez besoin de :

1. Bibliothèque apache-tvm (TIRx est inclus dedans).
2. PyTorch pour vérifier les résultats et préparer les données.
3. Un GPU Blackwell (par exemple, B200) si vous voulez exécuter le code sur du matériel réel.

Si vous n'avez pas de Blackwell sous la main (ce qui est logique pour la plupart d'entre nous en ce moment), le dépôt est toujours extrêmement utile comme base théorique. La logique d'optimisation mémoire et d'ordonnancement des calculs s'applique aussi aux architectures plus anciennes, bien que certaines fonctionnalités comme TMA soient spécifiques aux générations NVIDIA plus récentes.

Pour construire le livre localement et le lire dans un format pratique, Sphinx standard suffit :

pip install -r requirements-docs.txt
sphinx-build -b html . _build/html

Est-ce que ça vaut votre temps

Je rencontre souvent des développeurs d'infrastructure ML qui ont peur de se plonger dans l'écriture de noyaux, la considérant comme la prérogative des « dieux C++ ». Ce projet prouve qu'avec les bonnes abstractions (comme TIRx) et la compréhension de l'architecture, vous pouvez écrire du code de classe mondiale tout en restant dans l'écosystème Python.

Le projet sera particulièrement utile pour ceux qui travaillent sur l'optimisation des réseaux de neurones « sous le capot » ou qui veulent comprendre pourquoi Flash Attention est plus rapide que l'approche habituelle. Même si vous ne prévoyez pas de publier votre propre bibliothèque Blackwell demain, comprendre comment la mémoire tensorielle et les files d'attente asynchrones fonctionnent éclaircit vraiment l'esprit et change votre perspective sur le calcul haute performance.

Le seul inconvenient est que la barrière à l'entrée reste élevée. Ce n'est pas « GPU pour les nuls », mais plutôt « GPU pour ceux qui veulent devenir des pros ». Mais si vous êtes prêt à comprendre comment les octets se déplacent réellement à l'intérieur d'une carte graphique, il est difficile de trouver meilleur matériel sur le sujet en ce moment.