Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

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.

Siqi Guo, Ming Lin, Tianbao Yang• 2026

Related benchmarks

TaskDatasetResultRank
Triton kernel generationSynthetic Benchmark Level 1
Accuracy86.8
7
Triton kernel generationSynthetic Benchmark Level 2
Accuracy96
7
Triton kernel generationSynthetic Benchmark (Level 5)
Acc99
7
Triton kernel generationSynthetic Benchmark Level 20
Accuracy99
7
Triton kernel generationSynthetic Benchmark Overall All Levels
Average Speedup1.57
7
Triton kernel generationKernelBench Level 1
Accuracy69
6
Triton kernel generationKernelBench Level 2
Accuracy96
6
Triton kernel generationKernelBench Level 3
Accuracy76
6
Showing 8 of 8 rows

Other info

Follow for update