中文

TritonRL: Training LLMs to Think and Code Triton Without Cheating

软件工程 2026-02-10 v2 机器学习

摘要

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 resolve the credit assignment problem in long-sequence generation. Comprehensive evaluations on KernelBench 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.

关键词

引用

@article{arxiv.2510.17891,
  title  = {TritonRL: Training LLMs to Think and Code Triton Without Cheating},
  author = {Jiin Woo and Shaowei Zhu and Allen Nie and Zhen Jia and Yida Wang and Youngsuk Park},
  journal= {arXiv preprint arXiv:2510.17891},
  year   = {2026}
}