AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs
Shangzhan Li, Zefan Wang, Ye He, Yuxuan Li, Qi Shi, Jianling Li, Yonggang Hu, Wanxiang Che, Xu Han, Zhiyuan Liu, Maosong Sun
2025-07-10
Summary
This paper talks about AutoTriton, a system that uses artificial intelligence to automatically write and optimize code for GPUs, which are specialized computer chips used for fast calculations in AI and graphics.
What's the problem?
The problem is that writing efficient GPU code is very hard and requires experts to manually adjust settings like memory use and calculation order to get the best speed. This process is slow and difficult, which slows down AI development.
What's the solution?
The researchers created AutoTriton, which first learns the basics of a GPU programming language called Triton by studying lots of existing code. Then, it uses reinforcement learning, where it tries different programs and gets feedback on how well they perform, to improve its code generation. This two-step training helps AutoTriton write faster and more efficient GPU programs automatically.
Why it matters?
This matters because it makes it easier and faster to create good GPU programs, speeding up AI research and applications. By automating this complex task, more people can develop advanced AI systems without needing deep expertise in GPU programming.
Abstract
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.