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 프로그래밍을 단순화하지만, 개발자는 여전히 타일 크기와 메모리 접근 패턴과 같은 핵심 매개변수를 반복적인 실험을 통해 수동으로 조정해야 하며, 이는 최적의 성능과 광범위한 채택에 상당한 장벽으로 작용합니다. 본 연구에서는 강화 학습(RL) 기반의 Triton 프로그래밍 전용 모델인 AutoTriton을 소개합니다. AutoTriton은 고품질 데이터 수집 파이프라인을 통해 필수적인 Triton 프로그래밍 전문성을 습득하기 위해 지도 미세 조정(SFT)을 수행하고, 규칙 기반 보상과 실행 기반 보상을 결합한 Group Relative Policy Optimization(GRPO) 알고리즘을 사용하여 RL을 진행함으로써 Triton 프로그래밍 능력을 순차적으로 향상시킵니다. TritonBench와 KernelBench의 다섯 가지 평가 채널에서의 실험은 8B 모델인 AutoTriton이 Claude-4-Sonnet 및 DeepSeek-R1-0528을 포함한 주류 대형 모델과 비슷한 성능을 달성함을 보여줍니다. 추가 실험 분석은 SFT 단계, RL 단계, 보상 설계 전략을 포함한 AutoTriton 내 각 모듈의 중요한 역할을 입증합니다. 이러한 결과는 고성능 커널을 자동으로 생성하기 위한 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.