AutoTriton: Programação Automática em Triton com Aprendizado por Reforço em LLMs

AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs

July 8, 2025

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

cs.AI

Resumo

O desenvolvimento de kernels em aprendizado profundo requer a otimização de unidades computacionais em hardware, equilibrando o gerenciamento de memória, paralelismo e otimizações específicas de hardware por meio de ajustes empíricos extensivos. Embora linguagens específicas de domínio, como Triton, simplifiquem a programação de GPUs ao abstrair detalhes de baixo nível, os desenvolvedores ainda precisam ajustar manualmente parâmetros críticos, como tamanhos de blocos e padrões de acesso à memória, por meio de experimentação iterativa, criando barreiras significativas para o desempenho ideal e uma adoção mais ampla. Neste trabalho, apresentamos o AutoTriton, o primeiro modelo dedicado à programação Triton impulsionado por aprendizado por reforço (RL). O AutoTriton realiza ajuste fino supervisionado (SFT) para ser equipado com conhecimentos essenciais de programação Triton usando um pipeline de coleta de dados de alta qualidade, e conduz RL com o algoritmo Group Relative Policy Optimization (GRPO), combinando uma recompensa baseada em regras e uma recompensa baseada em execução para melhorar sequencialmente a capacidade de programação Triton. Experimentos em cinco canais de avaliação do TritonBench e KernelBench ilustram que nosso modelo de 8B, AutoTriton, alcança desempenho comparável a modelos grandes mainstream, incluindo Claude-4-Sonnet e DeepSeek-R1-0528. Análises experimentais adicionais demonstram o papel crucial de cada módulo dentro do AutoTriton, incluindo a etapa de SFT, a etapa de RL e a estratégia de design de recompensa. Esses achados destacam a promessa do RL para a geração automática de kernels de alto desempenho, e como kernels de alto desempenho são componentes centrais de sistemas de IA, esse avanço estabelece uma base importante para a construção de sistemas de IA mais eficientes. O modelo e o código estarão disponíveis em 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.