AutoTriton: Programmazione Automatica di Triton con Apprendimento per Rinforzo nei Modelli Linguistici di Grande Dimensione

AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs

July 8, 2025

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

cs.AI

Abstract

Lo sviluppo di kernel nell'ambito del deep learning richiede l'ottimizzazione delle unità computazionali su diversi hardware, bilanciando la gestione della memoria, il parallelismo e le ottimizzazioni specifiche per l'hardware attraverso un'ampia sintonizzazione empirica. Sebbene linguaggi specifici per dominio come Triton semplifichino la programmazione GPU astraendo i dettagli di basso livello, gli sviluppatori devono comunque sintonizzare manualmente parametri critici come le dimensioni dei tile e i pattern di accesso alla memoria attraverso sperimentazioni iterative, creando ostacoli significativi alle prestazioni ottimali e a una più ampia adozione. In questo lavoro, introduciamo AutoTriton, il primo modello dedicato alla programmazione Triton basato sull'apprendimento per rinforzo (RL). AutoTriton esegue un fine-tuning supervisionato (SFT) per acquisire competenze essenziali nella programmazione Triton utilizzando una pipeline di raccolta dati di alta qualità, e conduce l'RL con l'algoritmo Group Relative Policy Optimization (GRPO), combinando una ricompensa basata su regole e una ricompensa basata sull'esecuzione per migliorare ulteriormente, in modo sequenziale, le capacità di programmazione Triton. Esperimenti condotti su cinque canali di valutazione di TritonBench e KernelBench dimostrano che il nostro modello da 8B, AutoTriton, raggiunge prestazioni paragonabili ai principali modelli di grandi dimensioni, inclusi Claude-4-Sonnet e DeepSeek-R1-0528. Un'ulteriore analisi sperimentale evidenzia il ruolo cruciale di ciascun modulo all'interno di AutoTriton, inclusa la fase SFT, la fase RL e la strategia di progettazione delle ricompense. Questi risultati sottolineano il potenziale dell'RL per la generazione automatica di kernel ad alte prestazioni, e poiché i kernel ad alte prestazioni sono componenti fondamentali dei sistemi di intelligenza artificiale, questa svolta stabilisce una base importante per la costruzione di sistemi AI più efficienti. Il modello e il codice saranno disponibili su 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.