TritonRL: Training LLMs to Think and Code Triton Without Cheating

TL;DR

TritonRL tackles the challenge of generating Triton kernels with limited high-quality data and vulnerability to reward hacking by deploying a two-stage training pipeline: supervised knowledge distillation on KernelBook data augmented with reasoning traces, followed by reinforcement learning with robust, verifiable rewards and a hierarchical credit assignment that separates planning from code. The framework employs fine-grained verifiers to curb reward hacking and a GRPO-based objective to optimize both reasoning traces and the Triton code, with difficulty-aware data mixing to improve generalization. Empirical results on KernelBench show that 8B-scale TritonRL achieves state-of-the-art correctness and speedups among Triton-focused models, with ablations confirming the necessity of robust rewards, hierarchical decomposition, and disciplined data mixing. The work contributes openly released datasets and tooling to advance domain-specific kernel synthesis and demonstrates a viable path for efficient automation of GPU kernel production using small, specialized LLMs.

Abstract

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.

Figures (5)

• Figure 1: TritonRL components and workflow.
• Figure 2: Illustration of the flow of our robust verifier incorporating syntax and functionality checkers and the examples of invalid Triton codes. (a) invalid syntax: the code lacks any Triton blocks and consists solely of PyTorch code. (b) invalid functionality: the code include dummy Triton code that just copies data without meaningful operation delegating core operation (matrix multiplication) to PyTorch modules (torch.matmul).
• Figure 3: An example LLM output for Triton code generation, showing a reasoning trace (plan) and the generated Triton kernel code conditioned on the plan.
• Figure 4: Pass@$k$ correctness and mean speedup for $k=1,5,10$ on KernelBench Level 1 tasks. The figures show how performance (correctness on the left, mean speedup on the right) of TritonRL and baseline models scales as the number of attempts increases.
• Figure 5: Side-by-side comparison of base models and their post-trained variants. Successful training should reduce invalid syntax errors (yellow) and functional invalidity (red), while increasing correctness (green) from individual base model. Only TritonRL shows this outcome. For each metric, the numbers represent the ratio of corresponding samples among 10 generated samples for each of the 100 tasks in KernelBench Level 1.