DRTriton: Large-Scale Synthetic Data Driven Reinforcement Learning for Triton Kernel Generation
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Triton kernel generation | Synthetic Benchmark Level 1 | Accuracy86.8 | 7 | |
| Triton kernel generation | Synthetic Benchmark Level 2 | Accuracy96 | 7 | |
| Triton kernel generation | Synthetic Benchmark (Level 5) | Acc99 | 7 | |
| Triton kernel generation | Synthetic Benchmark Level 20 | Accuracy99 | 7 | |
| Triton kernel generation | Synthetic Benchmark Overall All Levels | Average Speedup1.57 | 7 | |
| Triton kernel generation | KernelBench Level 1 | Accuracy69 | 6 | |
| Triton kernel generation | KernelBench Level 2 | Accuracy96 | 6 | |
| Triton kernel generation | KernelBench Level 3 | Accuracy76 | 6 |