AutoTriton: Automatic Triton Programming with Reinforcement Learning in
  LLMs
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.
686e2c00a5f0f70d9de40c8c
Shangzhan Li
686e2c00a5f0f70d9de40c8d
Zefan Wang
686e2c00a5f0f70d9de40c8e
Ye He
686e2c00a5f0f70d9de40c8f
Yuxuan Li
686e2c00a5f0f70d9de40c90
Qi Shi
claimed_verified
2025-07-09T08:49:03.237Z
62ccd26d376917c022420a46
/avatars/629858b1c7419dbcdbead3484a36abd1.svg
Qi Shi
user
qshi
686e2c00a5f0f70d9de40c91
Jianling Li
686e2c00a5f0f70d9de40c92
Yonggang Hu
686e2c00a5f0f70d9de40c93
Wanxiang Che
686e2c00a5f0f70d9de40c94
Xu Han
686e2c00a5f0f70d9de40c95
Zhiyuan Liu
686e2c00a5f0f70d9de40c96
Maosong Sun
AutoTriton: Automatic Triton Programming with Reinforcement Learning in
  LLMs
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.
AutoTriton, a reinforcement learning model, automates the generation of high-performance kernels for GPU programming by combining supervised fine-tuning and Group Relative Policy Optimization.
reinforcement learning supervised fine-tuning Group Relative Policy Optimization Triton programming high-performance kernels AI systems