AutoTriton: Автоматическое программирование на Triton с использованием обучения с подкреплением в крупных языковых моделях

AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs

July 8, 2025

Авторы: Shangzhan Li, Zefan Wang, Ye He, Yuxuan Li, Qi Shi, Jianling Li, Yonggang Hu, Wanxiang Che, Xu Han, Zhiyuan Liu, Maosong Sun

cs.AI

Аннотация

Разработка ядер в глубоком обучении требует оптимизации вычислительных блоков с учетом аппаратного обеспечения, балансировки управления памятью, параллелизма и аппаратно-специфичных оптимизаций через обширную эмпирическую настройку. Хотя предметно-ориентированные языки, такие как Triton, упрощают программирование для GPU, абстрагируя низкоуровневые детали, разработчикам всё равно приходится вручную настраивать критические параметры, такие как размеры блоков и шаблоны доступа к памяти, через итеративное экспериментирование, что создает значительные барьеры для достижения оптимальной производительности и более широкого внедрения. В данной работе мы представляем AutoTriton — первую модель, предназначенную для программирования на Triton, основанную на обучении с подкреплением (RL). AutoTriton проходит этап контролируемого тонкого настроя (SFT), чтобы быть оснащенной необходимыми знаниями в программировании на Triton с использованием высококачественного конвейера сбора данных, и выполняет RL с алгоритмом Group Relative Policy Optimization (GRPO), комбинируя правило-ориентированное вознаграждение и вознаграждение на основе выполнения для дальнейшего улучшения способностей в программировании на Triton. Эксперименты на пяти оценочных каналах TritonBench и KernelBench показывают, что наша 8B-модель AutoTriton достигает производительности, сопоставимой с основными крупными моделями, включая Claude-4-Sonnet и DeepSeek-R1-0528. Дополнительный экспериментальный анализ демонстрирует ключевую роль каждого модуля в AutoTriton, включая этап SFT, этап RL и стратегию проектирования вознаграждений. Эти результаты подчеркивают перспективность RL для автоматической генерации высокопроизводительных ядер, и поскольку высокопроизводительные ядра являются ключевыми компонентами AI-систем, этот прорыв устанавливает важную основу для создания более эффективных AI-систем. Модель и код будут доступны по адресу 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.