Accueil › Glossaire › Triton

Triton (OpenAI Triton)

Triton est un langage de programmation et un compilateur open source créé par Philippe Tillet (OpenAI) qui permet d’écrire des kernels GPU haute performance directement en Python, en abstraisant les complexités de la programmation CUDA (gestion de la mémoire partagée, coalescence des accès mémoire, synchronisation des threads) tout en atteignant des performances proches du code CUDA écrit à la main.

Si vous utilisez torch.compile dans PyTorch, c’est Triton qui génère les kernels GPU en coulisse. Si vous exécutez Flash Attention, de nombreuses implémentations reposent sur des kernels Triton. Si vous servez un modèle via vLLM, des kernels Triton accélèrent certaines opérations. Triton est devenu une brique fondamentale de l’infrastructure IA moderne, rendant la programmation GPU accessible à des développeurs Python qui n’auraient jamais écrit de CUDA.

Le problème que Triton résout

Écrire des kernels GPU performants en CUDA est notoirement difficile. Il faut gérer manuellement la coalescence des accès mémoire DRAM, le placement des données en mémoire partagée (SRAM), la synchronisation entre threads au sein d’un bloc, et l’ordonnancement des opérations sur les Tensor Cores. Un expert CUDA peut passer des semaines à optimiser un seul kernel de multiplication matricielle.

Les alternatives existantes (Numba, CUDA Python) simplifient l’accès au GPU mais restent proches du modèle SIMT (Single Instruction, Multiple Thread) de CUDA, qui expose toutes ces complexités. Les frameworks de haut niveau (PyTorch, TensorFlow) fournissent des opérateurs pré-optimisés, mais ne permettent pas de créer facilement des opérations custom fusionnées.

Triton prend une approche radicalement différente. Au lieu d’exposer le modèle SIMT (un programme par thread), Triton utilise un modèle par blocs (un programme opère sur un bloc de données). Le développeur déclare des opérations sur des blocs (petits tableaux dont les dimensions sont des puissances de deux), et le compilateur Triton gère automatiquement la répartition entre threads, le placement en mémoire partagée, la coalescence des accès, et la synchronisation.

Comment fonctionne Triton

Le modèle de programmation par blocs

Dans Triton, chaque instance d’un kernel traite un bloc de données. Le parallélisme inter-blocs (entre instances) est géré par la grille de lancement. Le parallélisme intra-bloc (au sein d’une instance) est géré automatiquement par le compilateur. Le développeur n’a jamais à penser en termes de threads individuels.

Voici un kernel Triton de base pour l’addition de deux vecteurs :

Le décorateur @triton.jit marche en parcourant l’arbre syntaxique abstrait (AST) de la fonction Python pour générer du Triton-IR (une représentation intermédiaire basée sur LLVM/MLIR). Ce code est ensuite optimisé, parallélisé, et converti en LLVM-IR puis en PTX (pour NVIDIA) ou AMDGCN (pour AMD) pour l’exécution.

Pipeline de compilation

La compilation d’un kernel Triton passe par plusieurs étapes :

Python AST → Triton-IR : le code Python est transformé en représentation intermédiaire Triton, non optimisée et indépendante du matériel. Les opérations sur les blocs sont des citoyens de première classe.

Triton-IR → Triton-GPU IR : les optimisations GPU sont appliquées. Les données sont automatiquement placées en mémoire partagée en analysant les opérandes des opérations coûteuses (comme tl.dot). Le parallélisme intra-bloc est résolu.

Triton-GPU IR → LLVM-IR : la représentation est convertie en LLVM-IR standard.

LLVM-IR → PTX / AMDGCN : le code final est généré pour le GPU cible. Les deux étapes IR Triton sont implémentées comme des dialectes MLIR (sous-projet de LLVM pour la compilation hétérogène).

Autotuning

Triton fournit un mécanisme d’autotuning intégré via le décorateur @triton.autotune. Vous définissez un ensemble de configurations (tailles de blocs, nombre de warps, nombre de stages du pipeline), et Triton teste chaque combinaison pour trouver la plus rapide sur votre matériel :

Triton et torch.compile

L’intégration la plus impactante de Triton est avec torch.compile de PyTorch. Quand vous compilez un modèle PyTorch avec torch.compile(model, backend="inductor"), le backend TorchInductor capture le graphe de calcul du modèle et génère automatiquement des kernels Triton optimisés.

Ce processus est transparent pour le développeur. Vous ajoutez une seule ligne de code (model = torch.compile(model)) et PyTorch transforme automatiquement votre modèle en une séquence de kernels Triton fusionnés, éliminant les transferts mémoire intermédiaires et réduisant les lancements de kernels. Les gains de performance typiques vont de 1,5× à 3× selon le modèle et le hardware.

Depuis PyTorch 2.7, torch.library.triton_op permet aussi de créer des opérateurs custom backed par des kernels Triton définis par l’utilisateur, que torch.compile peut inspecter et optimiser dans le graphe global.

Support multi-matériel

Triton est conçu pour être portable entre différents fabricants de GPU :

NVIDIA (natif) : le backend principal. Triton compile vers PTX/CUBIN. Le support est le plus mature, incluant les Tensor Cores (via tl.dot) et le pipelining mémoire. Depuis début 2025, NVIDIA collabore directement avec OpenAI sur le support Blackwell, avec un nouveau backend CUDA Tile IR qui cible les Tensor Cores de 5ème génération. Les benchmarks montrent jusqu’à 1,5× d’accélération sur Flash Attention FP16 par rapport à Hopper, sans modification du code kernel.

AMD (ROCm) : Triton compile vers AMDGCN pour les GPU AMD. Le support est fonctionnel et en développement actif. Les kernels Triton fonctionnent sur les GPU Instinct (datacenter) et RDNA (grand public), bien que les performances ne soient pas toujours au niveau de CUDA natif.

Intel (XPU) : le backend intel-xpu-backend-for-triton permet de cibler les GPU Intel Arc et Data Center Max. Il supporte les Tensor Descriptors pour des opérations matmul optimisées via les opérations DPAS et le block IO 2D matériel d’Intel.

Cette portabilité multi-vendor est un argument clé de Triton : vous écrivez un kernel une fois, et il fonctionne (en théorie) sur NVIDIA, AMD et Intel, le compilateur gérant les optimisations spécifiques à chaque architecture.

Où Triton est utilisé dans l’écosystème IA

torch.compile / TorchInductor : le cas d’usage le plus répandu. PyTorch utilise Triton pour générer les kernels GPU de modèles compilés. TorchInductor analyse le graphe de calcul, fusionne les opérations adjacentes (par exemple, une matmul suivie d’un bias add et d’un ReLU devient un seul kernel), et génère le code Triton correspondant. Cette fusion élimine les allers-retours vers la HBM entre chaque opération, ce qui est souvent le facteur limitant des performances.

Flash Attention : l’implémentation originale de Flash Attention 2 par Tri Dao et les variantes communautaires utilisent des kernels Triton pour l’attention optimisée en mémoire. C’est le composant qui rend les longs contextes viables dans les LLM.

vLLM : des kernels Triton sont utilisés pour certaines opérations d’inférence, y compris des kernels de quantification pour GPTQ.

Unsloth : la bibliothèque d’optimisation du fine-tuning utilise des kernels Triton custom pour accélérer l’entraînement LoRA jusqu’à 2× tout en réduisant la consommation mémoire.

SageAttention : une implémentation d’attention optimisée basée sur des kernels Triton, utilisée pour l’inférence rapide de modèles de diffusion.

Recherche ML : les chercheurs écrivent des kernels Triton pour prototyper rapidement de nouvelles opérations (attention sparse, nouveaux schémas de quantification, opérateurs custom) sans devoir maîtriser CUDA.

L’intérêt principal : la fusion de kernels

La raison pour laquelle Triton est si important dans l’IA moderne tient en un mot : la fusion. Dans un modèle de deep learning classique exécuté avec PyTorch eager mode, chaque opération lance un kernel GPU séparé. Entre deux kernels, les résultats intermédiaires sont écrits en HBM (mémoire globale du GPU, relativement lente) puis relus par le kernel suivant. Pour un Transformer avec des dizaines de couches, ces allers-retours mémoire sont le vrai goulet d’étranglement.

Triton permet de fusionner plusieurs opérations en un seul kernel : les données restent en SRAM (mémoire partagée, très rapide) entre les opérations, éliminant les écritures/lectures HBM intermédiaires. C’est exactement ce que fait Flash Attention : au lieu de calculer l’attention en plusieurs passes (Q×K, softmax, ×V), tout est fusionné en un seul kernel qui ne matérialise jamais la matrice d’attention complète en mémoire. Le gain est double : moins de mémoire consommée ET plus de vitesse.

Des benchmarks récents montrent que cette fusion réduit les lancements de kernels de 67 à 88 % sur certains modèles, ce qui se traduit directement en gains de latence et de throughput.

Triton vs CUDA

En résumé : CUDA reste le choix quand chaque cycle compte (kernels de bibliothèque ultra-optimisés type cuBLAS, cuDNN). Triton est le choix pour les opérations custom, les fusions de kernels, le prototypage rapide, et les cas où la portabilité multi-vendor est importante. Les deux coexistent dans l’écosystème : les opérations standards utilisent souvent cuBLAS/cuDNN, les opérations fusionnées et custom utilisent Triton.

Support Blackwell et CUDA Tile IR

En février 2026, NVIDIA a annoncé l’intégration de CUDA Tile comme backend pour Triton. CUDA Tile est un nouveau modèle de programmation NVIDIA (introduit dans CUDA 13.1) qui remplace le modèle SIMT par un modèle natif par tuiles (tiles), aligné avec la philosophie de Triton. L’intégration CUDA Tile → Triton permet de cibler les Tensor Cores de 5ème génération des GPU Blackwell sans réécrire les kernels existants.

Les benchmarks montrent des améliorations significatives en FP16 et FP8 pour les GEMM et Flash Attention sur Blackwell. Le support reste en développement actif (certaines opérations comme les tensor-of-pointers ne sont pas encore optimales), mais il établit les fondations pour que les kernels Triton tirent automatiquement parti du matériel Blackwell.

Limites

Performance pas toujours au niveau de CUDA expert. Pour les opérations critiques comme les GEMM de grande taille, les kernels Triton atteignent typiquement 80 à 95 % des performances de cuBLAS/cuDNN optimisé. L’écart se réduit à chaque version, mais il existe encore pour les cas les plus optimisés.

Kernels de petite taille (GEMM_K petit). Les multiplications matricielles avec un K petit montrent une utilisation plus faible du GPU avec Triton. La mitigation passe par un sub-tiling manuel dans le kernel, en attendant que le compilateur gère automatiquement le warp-specialization.

Support Windows. Le support Windows de Triton est assuré via un fork communautaire (triton-windows). Le support natif progresse mais n’est pas encore au niveau de Linux. Triton >= 3.3 est requis avec PyTorch >= 2.7 et CUDA >= 12.8.

Courbe d’apprentissage. Bien que plus accessible que CUDA, Triton demande une compréhension des concepts GPU (grilles, blocs, mémoire partagée, bande passante) pour écrire des kernels performants. Le modèle par blocs simplifie l’implémentation mais ne dispense pas de la réflexion sur l’architecture matérielle.

GPU Turing retiré. Depuis Triton 3.3, le support des GPU Turing (RTX 20xx, compute capability 7.5) a été supprimé. Ces GPU restent supportés avec Triton 3.2, mais les nouvelles fonctionnalités ne sont plus testées dessus.