AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs

2025-07-08

AI Summary

Developing deep learning kernel functions faces challenges in hardware-specific adaptation and high manual tuning costs. Method: This paper introduces reinforcement learning (RL) to Triton programming automation for the first time, proposing a staged reward mechanism and Group Relative Policy Optimization (GRPO) to jointly optimize supervised fine-tuning (SFT) and RL; it integrates rule-based rewards with real execution feedback to construct a high-quality training data pipeline. Contribution/Results: Evaluated across five benchmarks in TritonBench and KernelBench, our 8B model achieves performance on par with Claude-4-Sonnet and DeepSeek-R1-0528. The approach significantly improves both the efficiency and quality of automated kernel generation while substantially reducing reliance on expert domain knowledge.

Technology Category

• Machine Learning: Reinforcement Learning
• Natural Language Processing: Code Generation / Program Synthesis from Natural Language
• Search and Optimization: Learning to Search

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

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.

Problem

• Automates Triton GPU programming using reinforcement learning
• Optimizes kernel parameters like tile sizes and memory access
• Enhances AI system efficiency via high-performance kernel generation

Innovation

• Uses reinforcement learning for Triton programming
• Combines rule-based and execution-based rewards
• Achieves performance comparable to large models