Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Siqi Guo, Ming Lin, Tianbao Yang

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.

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

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.
Paper Structure (30 sections, 8 equations, 6 figures, 3 tables)

This paper contains 30 sections, 8 equations, 6 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. CSP-DAG for PyTorch Program Generation
  4. Generating DAGs
  5. Constraint Programming
  6. Training Pipeline
  7. Triton Verifier
  8. Supervised Fine-Tuning
  9. RLVR with Decoupled Rewards
  10. Curriculum Reinforcement Learning
  11. Test-time Search for Optimal Compositional Kernels
  12. Experiments
  13. Experimental Setup
  14. Results on Synthetic Benchmarks
  15. Results on KernelBench
  16. ...and 15 more sections

Figures (6)

  • 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.
  • Figure 2: Frequency of collected SFT PyTorch operators.
  • Figure 3: 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.
  • Figure 4: Performance progression across training stages on three difficulty levels. We show Acc and Faster1 metrics for the base model, after SFT, and across three DRPO training stages.
  • Figure 5: Comparison of DRPO and GRPO performance when trained on Stage 1 data (20k samples) from the same SFT checkpoint.
  • ...and 1 more figures