TileLang est un langage Python de haut niveau spécifique à un domaine (DSL) destiné à concevoir et compiler, via TVM, des noyaux GPU axés sur les performances. Alors que les grands modèles de langage et d’autres architectures de type transformer accroissent les exigences de calcul, la capacité à écrire des noyaux GPU personnalisés exploitant pleinement les Tensor Cores — les accélérateurs dédiés de NVIDIA pour la multiplication matricielle — est devenue de plus en plus essentielle pour l’efficacité de l’entraînement et de l’inférence. Toutefois, l’écriture de tels noyaux exige traditionnellement une expertise approfondie en CUDA C++, en programmation au niveau des warps et en hiérarchies mémoire propres au matériel. TileLang comble cet écart en permettant aux développeurs d’exprimer les calculs au niveau des tuiles, tandis que le compilateur prend en charge les détails de bas niveau. Ce tutoriel propose un parcours complet des capacités de TileLang, en commençant par la validation de l’environnement et la mise en place d’utilitaires réutilisables de benchmarking et de vérification numérique. Il implémente ensuite progressivement l’addition de vecteurs, la multiplication matricielle tuilée sur Tensor Cores, l’exploration d’ordonnancements, les épilogues GEMM fusionnés, le softmax par ligne et FlashAttention.

Tout au long du tutoriel, les développeurs travaillent directement avec les tuiles en mémoire partagée de TileLang, les fragments de registres, les boucles pipelinées, les primitives d’itération parallèle, les réductions et les opérateurs GEMM sur Tensor Cores. Le compilateur gère automatiquement le mappage des threads, les dispositions mémoire, la synchronisation, la vectorisation et la génération d’instructions CUDA de bas niveau. Les noyaux sont évalués face à des références PyTorch et cuBLAS, avec inspection du code source CUDA généré, évaluation du débit mémoire et de calcul, et autotuning afin d’identifier des configurations de noyaux dépendantes de l’architecture. TileLang est disponible sur GitHub.

Le tutoriel commence par la configuration de l’environnement CUDA de Google Colab, l’installation de TileLang avec une solution de repli nightly, puis l’importation des modules PyTorch et TileLang requis. Des utilitaires réutilisables de benchmarking, de validation et de reporting sont définis afin de mesurer la latence des noyaux et de comparer les sorties numériques à l’aide de l’erreur relative. Un noyau TileLang d’addition de vecteurs est ensuite implémenté, exécuté sur le GPU, comparé à la bande passante de PyTorch, puis inspecté à travers le code source CUDA généré par le compilateur.

Ensuite, un noyau de multiplication matricielle tuilée sur Tensor Cores est implémenté, en déplaçant les tuiles d’entrée à travers la mémoire globale, la mémoire partagée et les fragments de registres. Les dimensions des tuiles, les étapes du pipeline, le nombre de threads et le swizzling L2 sont contrôlés manuellement, tandis que TileLang génère les instructions Tensor Core, la synchronisation et la logique de transfert mémoire. Plusieurs configurations d’ordonnancement sont benchmarkées, leur exactitude numérique est vérifiée, et la configuration de noyau dépendante de l’architecture la plus performante est identifiée. Le code complet du tutoriel est disponible ici.

Le noyau de multiplication matricielle est ensuite étendu en fusionnant l’ajout de biais et l’activation GELU directement dans l’accumulateur résidant dans les registres. Cette approche réduit le trafic intermédiaire vers la mémoire globale en terminant l’épilogue avant l’écriture du tenseur de sortie final. Une telle fusion de noyaux est une technique bien établie pour réduire les goulets d’étranglement liés à la bande passante mémoire, qui dominent souvent la latence dans les charges de travail de réseaux neuronaux à grande échelle. L’implémentation fusionnée est comparée à une exécution PyTorch eager. En outre, un noyau softmax par ligne est implémenté à l’aide de réductions de maximum et de somme au niveau des fragments, en conservant le processus de normalisation largement dans les registres.

Un noyau forward FlashAttention fusionné est implémenté afin de traiter les tuiles de requêtes, de clés et de valeurs sans matérialiser la matrice complète des scores d’attention en mémoire globale. Les mises à jour du softmax en ligne sont appliquées au moyen de maxima courants, de sommes de normalisation, de facteurs de remise à l’échelle et de multiplications matricielles tuilées sur Tensor Cores. FlashAttention, introduit initialement par des chercheurs de Stanford, est devenu une technique largement adoptée pour réduire l’empreinte mémoire de la self-attention d’un espace quadratique à un espace linéaire, et il est désormais intégré dans de grands frameworks, notamment l’API d’attention par produit scalaire mis à l’échelle de PyTorch. Les attentions causale et non causale sont toutes deux validées face à l’attention par produit scalaire mis à l’échelle de PyTorch, avec comparaison de la latence et du débit de calcul.

Un espace de recherche d’autotuning est défini sur les tailles de tuiles matricielles, les dimensions des blocs K, les profondeurs de pipeline et le nombre de threads, en filtrant les configurations qui dépassent le budget de mémoire partagée. Le décorateur d’autotuning de TileLang est utilisé pour compiler, benchmarker, valider et mettre en cache plusieurs ordonnancements de noyaux pour la même charge de travail de multiplication matricielle. Le noyau sélectionné est exécuté, vérifié face à PyTorch, puis évalué en termes de latence obtenue et de débit Tensor Core. Cette recherche automatisée répond à un défi pratique du développement de noyaux GPU : les tailles de tuiles et profondeurs de pipeline optimales varient selon les architectures GPU telles qu’Ampere et Hopper, ce qui rend le réglage manuel fragile lorsqu’il vise plusieurs générations de matériel.

Le flux de débogage et d’introspection de TileLang est présenté au moyen de l’impression côté device, de l’inspection du CUDA généré et du profileur de noyau intégré. Les repères émis par le compilateur sont examinés, notamment les opérations Tensor Core, les copies asynchrones, les barrières de synchronisation et les instructions de chargement matriciel. Toutes les sections du tutoriel sont organisées dans un runner tolérant aux pannes qui enregistre l’état d’exécution, rapporte les informations de temps d’exécution et imprime une référence compacte de programmation TileLang.

Le tutoriel montre comment TileLang traduit des programmes Python au niveau des tuiles en noyaux GPU optimisés, sans nécessiter la gestion manuelle des indices de threads, des dispositions de données au niveau des warps, des instructions Tensor Core ou des barrières mémoire asynchrones. Les noyaux implémentés et validés couvrent des opérations élément par élément limitées par la bande passante, des charges de travail GEMM intensives en calcul, des épilogues de réseaux neuronaux fusionnés, des réductions résidant dans les registres et l’attention avec softmax en ligne. L’exploration montre également comment les dimensions de blocs, la consommation de mémoire partagée, la profondeur de pipeline, le nombre de threads, la forme des tuiles et le swizzling L2 influencent les performances sur différentes architectures GPU. L’inspection du code source généré, le débogage côté device, le profilage et la recherche automatisée d’ordonnancements établissent ensemble un flux complet pour développer, vérifier, benchmarker et affiner des noyaux TileLang personnalisés.