TritonRL: Training LLMs to Think and Code Triton Without Cheating Article Swipe

MAKE IMPACT

Jin Seok Woo , Shaowei Zhu , Allen Nie , Zhen Jia , Yida Wang , Youngsuk Park · YOU? · · 2025 · Open Access · · DOI: https://doi.org/10.48550/arxiv.2510.17891

With the rapid evolution of large language models (LLMs), the demand for automated, high-performance system kernels has emerged as a key enabler for accelerating development and deployment. We introduce TritonRL, a domain-specialized LLM for Triton kernel generation, trained with a novel training framework that enables robust and automated kernel synthesis. Unlike general-purpose programming languages, Triton kernel generation faces unique challenges due to data scarcity and incomplete evaluation criteria, vulnerable to reward hacking. Our approach addresses these challenges end-to-end by distilling Triton-specific knowledge through supervised fine-tuning on curated datasets, and further improving code quality via reinforcement learning (RL) with robust, verifiable rewards and hierarchical reward assignment. Our RL framework robustly detects reward hacking and guides both reasoning traces and code tokens through fine-grained verification and hierarchical reward decomposition, enabling the model to generate high-quality Triton kernels that can truly replace existing modules. With robust and fine-grained evaluation, our experiments on KernelBench demonstrate that TritonRL achieves state-of-the-art correctness and speedup, surpassing all other Triton-specific models and underscoring the effectiveness of our RL-based training paradigm.

Type

preprint

Landing Page

http://arxiv.org/abs/2510.17891

PDF

https://arxiv.org/pdf/2510.17891

OA Status

green

OpenAlex ID

https://openalex.org/W4416054711

Raw OpenAlex JSON

OpenAlex ID

https://openalex.org/W4416054711

Canonical identifier for this work in OpenAlex

DOI

https://doi.org/10.48550/arxiv.2510.17891

Digital Object Identifier

Title

TritonRL: Training LLMs to Think and Code Triton Without Cheating

Work title

Type

preprint

OpenAlex work type

Publication year

2025

Year of publication

Publication date

2025-10-18

Full publication date if available

Authors

Jin Seok Woo, Shaowei Zhu, Allen Nie, Zhen Jia, Yida Wang, Youngsuk Park

List of authors in order

Landing page

https://arxiv.org/abs/2510.17891

Publisher landing page

PDF URL

https://arxiv.org/pdf/2510.17891

Direct link to full text PDF

Open access

Yes

Whether a free full text is available

OA status

green

Open access status per OpenAlex

OA URL

https://arxiv.org/pdf/2510.17891

Direct OA link when available

Cited by

0

Total citation count in OpenAlex

Full payload