AutoTriton : Programmation automatique de Triton avec apprentissage par renforcement dans les LLM

AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs

July 8, 2025

papers.authors: Shangzhan Li, Zefan Wang, Ye He, Yuxuan Li, Qi Shi, Jianling Li, Yonggang Hu, Wanxiang Che, Xu Han, Zhiyuan Liu, Maosong Sun

cs.AI

papers.abstract

Le développement de noyaux en apprentissage profond nécessite l'optimisation des unités de calcul sur le matériel tout en équilibrant la gestion de la mémoire, le parallélisme et les optimisations spécifiques au matériel grâce à un réglage empirique approfondi. Bien que les langages spécifiques à un domaine comme Triton simplifient la programmation GPU en masquant les détails de bas niveau, les développeurs doivent encore ajuster manuellement des paramètres critiques tels que les tailles de tuiles et les modèles d'accès à la mémoire par le biais d'expérimentations itératives, ce qui crée des obstacles importants à l'obtention de performances optimales et à une adoption plus large. Dans ce travail, nous présentons AutoTriton, le premier modèle dédié à la programmation Triton alimenté par l'apprentissage par renforcement (RL). AutoTriton effectue un réglage supervisé (SFT) pour acquérir une expertise essentielle en programmation Triton en utilisant un pipeline de collecte de données de haute qualité, et mène un apprentissage par renforcement avec l'algorithme Group Relative Policy Optimization (GRPO), combinant une récompense basée sur des règles et une récompense basée sur l'exécution pour améliorer séquentiellement les capacités de programmation Triton. Les expériences menées sur cinq canaux d'évaluation de TritonBench et KernelBench montrent que notre modèle AutoTriton de 8B atteint des performances comparables aux grands modèles dominants, notamment Claude-4-Sonnet et DeepSeek-R1-0528. Une analyse expérimentale approfondie démontre le rôle crucial de chaque module au sein d'AutoTriton, y compris l'étape SFT, l'étape RL et la stratégie de conception des récompenses. Ces résultats soulignent le potentiel de l'apprentissage par renforcement pour générer automatiquement des noyaux haute performance, et puisque les noyaux haute performance sont des composants essentiels des systèmes d'IA, cette percée établit une base importante pour la construction de systèmes d'IA plus efficaces. Le modèle et le code seront disponibles à l'adresse https://github.com/AI9Stars/AutoTriton.

English

Kernel development in deep learning requires optimizing computational units across hardware while balancing memory management, parallelism, and hardware-specific optimizations through extensive empirical tuning. Although domain-specific languages like Triton simplify GPU programming by abstracting low-level details, developers must still manually tune critical parameters such as tile sizes and memory access patterns through iterative experimentation, creating substantial barriers to optimal performance and wider adoption. In this work, we introduce AutoTriton, the first model dedicated to Triton programming powered by reinforcement learning (RL). AutoTriton performs supervised fine-tuning (SFT) to be equipped with essential Triton programming expertise using a high-quality data gathering pipeline, and conducts RL with Group Relative Policy Optimization (GRPO) algorithm, combining a rule-based reward and an execution-based reward to further improve Triton programming ability, sequentially. Experiments across five evaluation channels of TritonBench and KernelBench illustrate that our 8B model AutoTriton achieves performance comparable to mainstream large models, including Claude-4-Sonnet and DeepSeek-R1-0528. Further experimental analysis demonstrates the crucial role of each module within AutoTriton, including the SFT stage, the RL stage, and the reward design strategy. These findings underscore the promise of RL for automatically generating high-performance kernels, and since high-performance kernels are core components of AI systems, this breakthrough establishes an important foundation for building more efficient AI systems. The model and code will be available at https://github.com/AI9Stars/AutoTriton.