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

TL;DR

DRTriton is a scalable framework that trains large language models to convert PyTorch code into optimized Triton kernels, achieving significant speedups on real-world CUDA kernels despite training solely on synthetic data.

Contribution

It introduces a synthetic data generation algorithm, a curriculum reinforcement learning approach, and a test-time search method for effective CUDA kernel generation.

Findings

DRTriton-7B outperforms GPT-5.2 and Claude-Sonnet-4.5 on KernelBench speedup.

Synthetic training data enables effective generalization to real-world CUDA kernels.

The framework achieves speedups on 92% of benchmark kernels.

Abstract

Developing efficient CUDA kernels is a fundamental yet challenging task in the generative AI industry. Recent researches leverage Large Language Models (LLMs) to automatically convert PyTorch reference implementations to CUDA kernels, significantly reducing the engineering efforts. State-of-the-art LLMs, such as GPT-5.2 and Claude-Sonnet-4.5, still struggle in this specific task. To address this challenge, we propose DRTriton, a scalable learning framework for training LLMs to convert PyTorch codes 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 reinforcement learning with decoupled reward efficiently optimizes conversion success rate and inference speed simultaneously; and (iii) a test-time search algorithm that further improves the inference speed of the generated Triton kernels. Notably, despite being trained exclusively on synthetic data, DRTriton generalizes effectively to real-world CUDA kernels that are challenging even for human experts. Experimental results show that DRTriton-7B achieves speedup on 92% of the KernelBench Level 2, compared to 23% for GPT-5.2 and 19% for Claude-Sonnet-4.5.