TritonRL: Training LLMs to Think and Code Triton Without Cheating
Jiin Woo · Shaowei Zhu · Allen Nie · Zhen Jia · Yida Wang · Youngsuk Park
TritonRL (8B) achieves frontier-level GPU kernel generation by using a multi-layered verification system and hierarchical RL to prevent "reward hacking" and solve complex credit assignment in long-sequence code synthesis.
摘要
The rapid evolution of Large Language Models (LLMs) has driven a growing demand for automated, high-performance system kernels to accelerate machine learning workloads. We introduce TritonRL, a domain-specialized 8B-scale LLM for Triton programming, trained via a novel reinforcement learning (RL) framework. While Triton synthesis faces unique challenges, including data scarcity and a high susceptibility to reward hacking, our approach enables robust kernel generation through two primary innovations. First, we implement a multi-layered verification system that provides high-fidelity reward signals, ensuring that generated kernels are both syntactically and functionally valid. Second, we propose Hierarchical Reward Decomposition (HRD), which decouples reinforcement for high-level reasoning and low-level implementation to improve credit assignment in plan-code structured generation. Comprehensive evaluations on KernelBench and TritonBench-T demonstrate that TritonRL achieves state-of-the-art correctness and runtime speedup, outperforming concurrent Triton-specific models and matching the performance of frontier models with over 100B parameters. Our results highlight the effectiveness of hardware-aware RL paradigms in specialized domain adaptation.