AutoTriton: Programación Automática de Triton con Aprendizaje por Refuerzo en Modelos de Lenguaje de Gran Escala

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

Resumen

El desarrollo de núcleos (kernels) en el aprendizaje profundo requiere optimizar unidades computacionales en diferentes hardware, equilibrando la gestión de memoria, el paralelismo y optimizaciones específicas del hardware mediante un ajuste empírico extensivo. Aunque lenguajes específicos de dominio como Triton simplifican la programación en GPU al abstraer detalles de bajo nivel, los desarrolladores aún deben ajustar manualmente parámetros críticos, como los tamaños de los bloques (tile sizes) y los patrones de acceso a memoria, a través de experimentación iterativa, lo que crea barreras significativas para alcanzar un rendimiento óptimo y una adopción más amplia. En este trabajo, presentamos AutoTriton, el primer modelo dedicado a la programación en Triton impulsado por aprendizaje por refuerzo (RL). AutoTriton realiza un ajuste fino supervisado (SFT) para adquirir conocimientos esenciales de programación en Triton utilizando una canalización de recopilación de datos de alta calidad, y lleva a cabo RL con el algoritmo de Optimización de Política Relativa Grupal (GRPO), combinando una recompensa basada en reglas y una recompensa basada en la ejecución para mejorar secuencialmente la capacidad de programación en Triton. Los experimentos en cinco canales de evaluación de TritonBench y KernelBench muestran que nuestro modelo de 8B, AutoTriton, alcanza un rendimiento comparable a modelos grandes predominantes, como Claude-4-Sonnet y DeepSeek-R1-0528. Un análisis experimental adicional demuestra el papel crucial de cada módulo dentro de AutoTriton, incluyendo la etapa de SFT, la etapa de RL y la estrategia de diseño de recompensas. Estos hallazgos resaltan el potencial del RL para generar automáticamente núcleos de alto rendimiento, y dado que estos núcleos son componentes centrales de los sistemas de IA, este avance establece una base importante para construir sistemas de IA más eficientes. El modelo y el código estarán disponibles en 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.