AutoTriton: Automatische Triton-Programmierung mit Reinforcement Learning in LLMs

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

Die Entwicklung von Kerneln im Bereich des Deep Learnings erfordert die Optimierung von Recheneinheiten über verschiedene Hardwareplattformen hinweg, wobei ein Gleichgewicht zwischen Speicherverwaltung, Parallelität und hardware-spezifischen Optimierungen durch umfangreiche empirische Abstimmung hergestellt werden muss. Obwohl domänenspezifische Sprachen wie Triton die GPU-Programmierung durch die Abstraktion von Low-Level-Details vereinfachen, müssen Entwickler dennoch kritische Parameter wie Tile-Größen und Speicherzugriffsmuster manuell durch iterative Experimente anpassen, was erhebliche Hindernisse für optimale Leistung und breitere Akzeptanz darstellt. In dieser Arbeit stellen wir AutoTriton vor, das erste Modell, das sich der Triton-Programmierung widmet und durch Reinforcement Learning (RL) angetrieben wird. AutoTriton führt ein Supervised Fine-Tuning (SFT) durch, um mit essenzieller Triton-Programmierungsfähigkeit ausgestattet zu werden, und führt RL mit dem Group Relative Policy Optimization (GRPO)-Algorithmus durch, der eine regelbasierte Belohnung und eine ausführungsbasierte Belohnung kombiniert, um die Triton-Programmierungsfähigkeit schrittweise weiter zu verbessern. Experimente über fünf Evaluationskanäle von TritonBench und KernelBench zeigen, dass unser 8B-Modell AutoTriton eine Leistung erzielt, die mit gängigen großen Modellen wie Claude-4-Sonnet und DeepSeek-R1-0528 vergleichbar ist. Weitere experimentelle Analysen demonstrieren die entscheidende Rolle jedes Moduls innerhalb von AutoTriton, einschließlich der SFT-Phase, der RL-Phase und der Belohnungsdesignstrategie. Diese Ergebnisse unterstreichen das Potenzial von RL für die automatische Generierung von Hochleistungskerneln, und da Hochleistungskernel Kernkomponenten von KI-Systemen sind, legt dieser Durchbruch eine wichtige Grundlage für den Aufbau effizienterer KI-Systeme. Das Modell und der Code werden unter https://github.com/AI9Stars/AutoTriton verfügbar sein.

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.