AutoTriton: Automatisch Triton Programmeren met Reinforcement Learning in LLM's

AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs

July 8, 2025

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

cs.AI

Samenvatting

Kernelontwikkeling in deep learning vereist het optimaliseren van rekenkundige eenheden over hardware heen, terwijl geheugenbeheer, parallellisme en hardware-specifieke optimalisaties in balans worden gebracht door uitgebreide empirische afstemming. Hoewel domeinspecifieke talen zoals Triton GPU-programmering vereenvoudigen door laag-niveau details te abstraheren, moeten ontwikkelaars nog steeds kritieke parameters zoals tilegroottes en geheugentoegangspatronen handmatig afstemmen via iteratief experimenteren, wat aanzienlijke barrières creëert voor optimale prestaties en bredere adoptie. In dit werk introduceren we AutoTriton, het eerste model gewijd aan Triton- programmering aangedreven door reinforcement learning (RL). AutoTriton voert supervised fine-tuning (SFT) uit om uitgerust te worden met essentiële Triton-programmerings- expertise met behulp van een hoogwaardige datainzamelingspijplijn, en voert RL uit met het Group Relative Policy Optimization (GRPO) algoritme, waarbij een op regels gebaseerde beloning en een op uitvoering gebaseerde beloning worden gecombineerd om de Triton-programmerings- vaardigheid verder te verbeteren, sequentieel. Experimenten over vijf evaluatiekanalen van TritonBench en KernelBench illustreren dat ons 8B-model AutoTriton prestaties bereikt die vergelijkbaar zijn met mainstream grote modellen, waaronder Claude-4-Sonnet en DeepSeek-R1-0528. Verdere experimentele analyse toont de cruciale rol aan van elke module binnen AutoTriton, inclusief de SFT-fase, de RL-fase, en de beloningsontwerpstrategie. Deze bevindingen onderstrepen de belofte van RL voor het automatisch genereren van hoogwaardige kernels, en aangezien hoogwaardige kernels kerncomponenten zijn van AI-systemen, legt deze doorbraak een belangrijke basis voor het bouwen van efficiëntere AI-systemen. Het model en de code zijn beschikbaar op 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.