DRTriton: Large-Scale Synthetic Data Driven Reinforcement Learning for Triton Kernel Generation

Siqi Guo; Ming Lin; Tianbao Yang

cs.CLMar 23, 2026

DRTriton: Large-Scale Synthetic Data Driven Reinforcement Learning for Triton Kernel Generation

Authors: Siqi Guo, Ming Lin, Tianbao Yang

Organizations: Texas A&M University · Oracle

Abstract

Developing efficient CUDA kernels is a fundamental yet challenging task in the generative AI industry. Recent research leverages Large Language Models (LLMs) to automatically convert PyTorch reference implementations to CUDA kernels, significantly reducing engineering effort. State-of-the-art LLMs, such as GPT-5.2 and Claude-Sonnet-4.5, still struggle with this task. To address this challenge, we propose DRTriton, a scalable learning framework for training LLMs to convert PyTorch programs into highly optimized Triton kernels, which are then compiled to CUDA kernels at runtime. DRTriton consists of three key components: (i) a data synthetic algorithm CSP-DAG that guarantees full coverage and unbiased uniform sampling over the operator space with controlled difficulty; (ii) a curriculum RL framework with decoupled rewards that jointly optimizes conversion success rate and execution speed; and (iii) a test-time search algorithm that further improves the execution speed of the generated Triton kernels. With a warmup stage of SFT on limited PyTorch-Triton pairs curated using existing LLMs, DRTriton trained by RL on synthesized PyTorch programs generalizes effectively to real-world CUDA kernels that are challenging even for human experts. Experimental results show that DRTriton-7B achieves speedup over PyTorch on 92% of KernelBench Level 2 tasks, compared to 23% for GPT-5.2 and 19% for Claude-Sonnet-4.5.

Figures & tables

Figure 1: Left: A DAG with 5 nodes and 10 edges. Each node represents an operator, and each edge represents a tensor. Right: The synthesized PyTorch program. Blue part indicates the tensors and the red part indicates the operators.

Model	Level 1		Level 2		Level 5		Level 20		Avg. speedup
Model	Acc	Faster1	Acc	Faster1	Acc	Faster1	Acc	Faster1	Avg. speedup
GPT-5.2	54	17	43	35	5	5	0	0	1.34
Claude-Sonnet-4.5	68	14	49	29	7	3	0	0	0.68
DeepSeek-R1	33	13	24	16	9	3	0	0	0.91
Qwen-3-Coder-480B	48	6	30	23	1	1	0	0	0.97
AutoTriton	36	7	15	8	2	1	0	0	1.08

Table 1: Main results on synthetic benchmark. Acc (%): Pass@1 accuracy breakdown by difficulty level. Faster1 (%): percentage of correct kernels achieving

>1\times

speedup. Avg. speedup: geometric mean speedup over PyTorch computed across all correct kernels.

Figure 2: Speedup distribution across difficulty levels for DRTriton, DRTriton (with test-time search), and other baselines, on synthetic benchmarks. Each curve shows speedups of generated Triton kernels over PyTorch programs across the percentages of testing data. The red dashed flat line is the reference line for 1x speedup, and the orange dashed flat line is the reference line for 2x speedup.

Model	Level 1			Level 2			Level 3
Model	Acc	TE	TC	Acc	TE	TC	Acc	TE	TC
AutoTriton	44	12	6	45	14	12	32	14	6
GPT-5.2	40	23	23	32	23	11	30	18	10
Claude-Sonnet-4.5	46	8	8	36	19	12	36	12	12
DeepSeek-R1	25	8	7	22	15	8	8	2	2
Qwen-3-Coder-480B	21	5	1	18	14	1	24	0	0

Table 2: Results on KernelBench real-world benchmark. We report accuracy (Acc) and Faster1 scores (percentage of correct kernels achieving

>1\times

speedup) against two baselines: Torch Eager (TE) and torch.compile (TC).

Panel of Figure 3 — Figure 3: Training dynamics. (a) Performance progression across training stages on three difficulty levels, showing Acc and Faster1 metrics for the base model, after SFT, and across three DRPO stages. (b) Comparison of DRPO and GRPO trained on Stage 1 data (20k samples) from the same SFT checkpoint; DRPO consistently outperforms GRPO across all metrics.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Category	Operator Class	PyTorch Function	#In
OpCreate	Randn	torch.randn	0
	Ones	torch.ones	0
	Zeros	torch.zeros	0
Elementwise Binary	Add	torch.add	2
	Mul	torch.mul	2
	Sub	torch.sub	2

Appendix

Table 3: Complete List of Operators

Figure 4: Speedup distribution across difficulty levels for DRTriton (with test-time search), and other baselines, on KernelBench. Each curve shows speedups of generated Triton kernels over PyTorch programs across the percentages of testing data. The red dashed flat line is the reference line for 1x speedup, and the orange dashed flat line is the reference line for 2x speedup.

Reward Function	Level 1
Reward Function	Acc	Faster1
Power ( $\alpha=0.25$ )	38.1	14.4
Power ( $\alpha=0.50$ )	30.9	5.2
Power ( $\alpha=0.75$ )	30.9	2.1
Power ( $\alpha=1.0$ )	32.0	3.1
Logarithmic	42.3	18.6

Appendix

Table 4: Ablation study on speed reward function design. All models are trained with DRPO on Stage 1 data using Qwen-2.5-Coder-1.5B as the base model, starting from the SFT checkpoint.

Benchmark	Setting	Fragments evaluated	Generation (h)	Validation (h)	Total (h)
Synthetic	w/o test-time search	400	0.01	0.11	0.12
Synthetic	w/ test-time search	10,612	0.13	2.72	2.85
KernelBench	w/o test-time search	250	0.01	0.07	0.08
KernelBench	w/ test-time search	16,859	0.23	3.30	3.53

Appendix

Table 5: Test-time search overhead on synthetic and KernelBench benchmarks. Generation and validation times are reported in hours.

GPU Kernel Generation

GPU (graphics processing unit) kernel generation automates writing and optimizing device code for faster computation across accelerators.

See all 67 papers

CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation
Feb 27, 2026Weinan Dai, Hanlin Wu, Qiying Yu +13GPU Kernel Generation GPU Kernel Optimization
ByteDance Seed · Institute for AI Industry Research (AIR), Tsinghua University · SIA-Lab of Tsinghua AIR and ByteDance Seed
RealisticTritonBench: A Benchmark for Triton-Kernel Generation in Real-World AI Frameworks
Aug 12, 2026Jinjun Huang, Zhongzhen Wen, Tongtong Xu +3LLM Evaluation GPU Kernel Generation
College of Computer Science and Technology and the State Key Laboratory of Blockchain and Data Security, Zhejiang University · State Key Lab for Novel Software Technology, Nanjing University · Software Engineering Application Technology Laboratory, Huawei +1
KernelZero: Co-Evolving Proposer and Coder for Continuously Improved GPU Kernel Generation
Sep 27, 2026Changxin Ke, Rui Zhang, Zixiang Fang +8GPU Kernel Generation LLM-Based Program Synthesis
State Key Lab of Processors, Institute of Computing Technology, CAS · University of Chinese Academy of Sciences · Intelligent Software Research Center, Institute of Software, CAS +1

Feb 27, 2026J/K move · Enter open · S save

CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation

Weinan Dai, Hanlin Wu, Qiying Yu, Huan-ang Gao+12

ByteDance Seed · Institute for AI Industry Research (AIR), Tsinghua University · SIA-Lab of Tsinghua AIR and ByteDance Seed

GPU Kernel Generation GPU Kernel Optimization RL for Code Generation Agentic RL

Open paper AskPDF arXiv

GPU kernel optimization is fundamental to modern deep learning but remains a highly specialized task requiring deep hardware expertise. Despite strong performance in general programming, large language models (LLMs) remain uncompetitive with compiler-based systems such as torch.compile for CUDA kernel generation. Existing CUDA code generation approaches either rely on training-free refinement or fine-tune models within fixed multi-turn execution-feedback loops, but both paradigms fail to fundamentally improve the model's intrinsic CUDA optimization ability, resulting in limited performance gains. We present CUDA Agent, a large-scale agentic reinforcement learning system that develops CUDA kernel expertise through three components: a scalable data synthesis pipeline, a skill-augmented CUDA development environment with automated verification and profiling to provide reliable reward signals, and reinforcement learning algorithmic techniques enabling stable training. CUDA Agent achieves state-of-the-art results on KernelBench, delivering 100%, 100%, and 92% faster rate over torch.compile on KernelBench Level-1, Level-2, and Level-3 splits, outperforming the strongest proprietary models such as Claude Opus 4.5 and Gemini 3 Pro by about 40% on the hardest Level-3 setting.

Abstract

GPU Kernel Generation

Explore similar work

CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation

RealisticTritonBench: A Benchmark for Triton-Kernel Generation in Real-World AI Frameworks

KernelZero: Co-Evolving Proposer and Coder for Continuously Improved GPU Kernel Generation