PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Siqi Guo, Ming Lin, Tianbao Yang

arXiv:2603.21465v2cs.CLcs.LG

TL;DR

Automating efficient CUDA kernel development remains difficult because even Triton requires substantial GPU expertise and trial and error. DRTriton trains LLMs with large-scale synthetic PyTorch programs and curriculum reinforcement learning, substantially surpassing baselines on synthetic and real-world benchmarks, including 76% accuracy on KernelBench Level 3.

  • Problem

    Efficient CUDA kernel development remains difficult and labor-intensive, while writing optimized Triton kernels still requires substantial GPU-programming expertise and trial and error.

  • Method

    DRTriton combines CSP-DAG synthetic program generation with curriculum reinforcement learning that decouples conversion-accuracy and execution-speed rewards.

  • Results

    DRTriton-7B substantially surpasses baselines on synthetic and real-world benchmarks, achieving 76% accuracy on KernelBench Level 3.

  • Takeaways & Limitations

    Synthetic-data-driven curriculum reinforcement learning generalizes effectively to complex, human-authored GPU kernels.

  • Takeaways & Limitations

    The training process faces challenges from uncontrolled sample complexity and quality, since difficult samples introduced too early can produce zero rewards and poor gradient signals.

Abstract

from arXiv · show

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.

1 Introduction

DRTriton targets the difficult, expertise-intensive conversion of PyTorch programs into efficient Triton kernels, where existing LLMs and repository-based training approaches remain limited by weak performance, small datasets, and uncontrolled difficulty. It addresses these limitations with systematically generated programs of varying difficulty and curriculum reinforcement learning, achieving 99% synthetic-benchmark accuracy and outperforming PyTorch with 86% of generated kernels.

  • Motivation: Efficient CUDA and Triton kernel development remains difficult because manual optimization is costly, while Triton still demands GPU expertise and extensive trial and error.FlashAttention Dao et al. [2022] is cited as requiring years of development effort, and current LLMs still struggle with efficient PyTorch-to-Triton or CUDA translation [Ouyang et al., 2025, Li et al., 2025c].
  • Limitations: Existing repository-based approaches are constrained by limited training data, with Li et al. [2025d] using 14k samples and Woo et al. [2025] using 11k samples.The paper argues that these dataset sizes are insufficient for learning complex Triton or CUDA kernels.
  • Limitations: Uncontrolled sample complexity can cause sparse-reward training to yield zero rewards, poor gradient signals, wasted tokens, and misleading updates when difficult examples appear too early.This motivates controlling the order and difficulty of training samples.
  • Approach: DRTriton systematically generates PyTorch programs with varying difficulty and uses curriculum reinforcement learning to progress from simple to complex tasks.The approach begins with a representative subset of 61 widely used PyTorch operators, including operators used in KernelBench.
  • Results: 99% accuracy was achieved on a synthetic benchmark with 20 operators, and 86% of generated Triton kernels outperformed the original PyTorch implementations.These results were obtained by DRTriton-7B after warmup on 2,026 curated PyTorch–Triton pairs and curriculum RL on 100,000 synthetic PyTorch programs.

2 Related Work

Prior work spans conventional kernel compilers, LLM-based kernel generation, multi-agent systems, synthetic code generation, and RLVR. These approaches motivate DRTriton’s focus on scalable data construction and decoupled rewards for sparse-reward kernel generation.

  • Kernel Compilers: Conventional compilers such as TVM [Chen et al., 2018], Ansor [Zheng et al., 2020], and TorchInductor optimize kernels through learned schedules, pattern matching, and loop fusion but remain bounded by predefined patterns.TorchInductor has been PyTorch’s default compiler since PyTorch 2.0 and maps PyTorch models to Triton kernels.
  • GPU Kernel Generation with LLMs: LLM-based kernel generation moved beyond rigid compilers, but KernelLLM trained by SFT on 25k PyTorch–Triton pairs suffers from hallucinated or “fake” kernels.Subsequent RL approaches use hand-curated data [Li et al., 2025d, Woo et al., 2025], limiting scalability because training samples must be curated from public datasets.
  • Multi-agent systems: Multi-agent kernel systems coordinate planning, generation, and verification [Hong et al., 2025, Wang et al., 2025a, Wei et al., 2025, Zhang et al., 2025, Liao et al., 2025, Li et al., 2025b] but remain bounded by foundation models and can exploit testing loopholes.Robust benchmarking [Lange et al., 2025] identified loopholes including redundant-operator elimination and hard-coding for specific input patterns.
  • Synthetic Code Generation: LLM-based synthetic code generation, including Magicoder [Wei et al., 2023] and WizardCoder [Luo et al., 2023], cannot guarantee correctness or coverage, while prior kernel-data synthesis relies on expensive LLM generation [Paliskara and Saroufim, 2025, Liao et al., 2025].The paper instead introduces a CSP-based framework for kernel-data generation.
  • Reinforcement Learning with Verifiable Rewards (RLVR): RLVR can produce abilities beyond training data [Guo et al., 2025], but code-generation methods such as AutoTriton [Li et al., 2025d] and TritonRL [Woo et al., 2025] face sparse early rewards when compiler feedback and kernel correctness are combined.DRPO [Li et al., 2025a] first proposed decoupled rewards; this work applies curriculum learning with decoupled rewards to address sparse rewards.

3 CSP-DAG for PyTorch Program Generation

CSP-DAG synthesizes PyTorch programs by constructing operator DAGs and solving tensor-shape constraints, producing valid programs with controlled difficulty. It guarantees coverage over all valid programs composable from the given operators while scaling to 100k programs in roughly 1.5 hours on 32 CPU cores.

  • DAG generation: CSP-DAG constructs PyTorch programs as directed acyclic graphs of operators and tensors, then fills missing inputs with tensor-creating operators during graph generation.OpCompute operators consume tensors, whereas OpCreate operators produce tensors without tensor inputs; the algorithm iteratively adds nodes using a candidate-tensor list.
  • Shape constraints: The method formulates tensor ranks and dimensions as constraint variables, enforcing operator-specific shape compatibility and bounds on dimension sizes, FLOPs, and tensor elements.Each tensor dimension is an integer between 1 and 2^15, while neighboring operators impose additional rank and dimension constraints.
  • Shape solving: The CP-SAT solver [Perron and Didier] finds feasible tensor shapes for the generated DAG and randomly selects among multiple feasible solutions.This constraint-solving stage converts the structurally generated graph into valid PyTorch code.
  • Coverage and scalability: CSP-DAG guarantees coverage over all valid PyTorch programs composable from the given operators, solves DAGs with up to 20 operators in 1–2 seconds per CPU core, and generates 100k programs in roughly 1.5 hours on 32 cores.Program difficulty is defined by the number of OpCompute operators, enabling controlled difficulty levels.
  • Program representation: The synthesized program format closely resembles PyTorch compiler IR and can represent standard PyTorch models lowered through torch.export.The representation is flat and functional, with each line corresponding to a single operator application.

4 Training Pipeline

The training pipeline combines Level 1 SFT cold-starting, curriculum DRPO with decoupled correctness and speed rewards, and test-time search over verified kernel compositions. A Triton verifier enforces syntactic validity, genuine kernel use, and exact agreement with PyTorch outputs.

  • Triton Verification: The Triton verifier requires a kernel to be syntactically valid, genuinely used, and exactly output-equivalent to PyTorch across five random test cases.Validation combines linting and compilation, no-op monkey-patch testing to detect copied PyTorch implementations, and precise output comparison.
  • Supervised Fine-Tuning: Level 1 SFT provides a cold start by teaching basic PyTorch-to-Triton conversion before RL increases task difficulty.Direct RL would otherwise produce sparse, uninformative rewards because most generated programs would fail to compile; SFT programs each contain exactly one OpCompute operator.
  • Curriculum RL: DRPO jointly favors correct Triton translations and faster implementations by decoupling correctness and speed rewards after SFT.Correct outputs are weighted by speed, incorrect outputs are penalized, and the logarithmic speed-reward function performs best among evaluated choices.
  • Curriculum RL: Curriculum learning advances to the next difficulty level when held-out Pass@1 accuracy exceeds 50% and stops when performance plateaus, reaching Level 5 in the experiments.The curriculum is intended to improve learning efficiency and stability by gradually increasing task difficulty.
  • Test-Time Search: Test-time search decomposes an n-operator PyTorch program into contiguous fragments of length at most 5, verifies generated kernels, and selects the fastest correct hybrid implementation.The search evaluates fusion strategies by replacing verified fragments while leaving other operators unchanged, yielding the optimal strategy among the verified options.

5 Experiments

DRTriton substantially outperforms all baselines on synthetic and real-world benchmarks in accuracy and execution speed, with curriculum RL, test-time search, and decoupled rewards driving these gains. On KernelBench, it reaches 96% accuracy and 92% Faster1 at Level 2, while ablations confirm DRPO and logarithmic speed rewards are effective design choices.

  • Synthetic benchmark: 87% accuracy at synthetic Level 1 and 75% at Level 2, versus Claude’s 68% and 49%, while Level 5 retains 15% accuracy as most baselines nearly fail.DRTriton outperforms all baselines in accuracy across difficulty levels.
  • Synthetic benchmark: 54% of Level 2 synthetic kernels exceed 1× speedup, versus 17% for GPT-5.2 and 14% for Claude at Level 1, with DRTriton’s curves above all baselines.The speed-based reward weighting ω(o|q) guides generation toward optimized rather than merely correct kernels.
  • Curriculum RL: Stage 2 delivers peak Level 1 and Level 2 performance, Stage 3 slightly regresses there, and Level 5 improves steadily throughout curriculum RL.DRPO training substantially improves over SFT.
  • Test-time search: Average speedup rises from 1.20× to 1.57× with test-time search, which particularly improves complex programs by optimizing inefficient code fragments.Longer and harder programs can achieve higher accuracy and speedup after search because they contain more optimization opportunities.
  • KernelBench results: 96% accuracy and 92% Faster1 at KernelBench Level 2, while Level 3 reaches 76% accuracy and 54% Faster1 over Torch Eager, exceeding baseline LLMs on challenging real-world implementations.At Level 2, DRTriton also achieves 56% Faster1 over torch.compile; at Level 3, AutoTriton has higher accuracy but DRTriton is substantially faster.
  • Ablations: DRPO consistently outperforms GRPO across all metrics, while the logarithmic speed reward consistently achieves the best Acc and Faster1 among tested reward forms.The compared speed rewards include log(ttorch/ttriton) and power forms (ttorch/ttriton)^α for α ∈ {0.25, 0.5, 0.75, 1.0}; the logarithmic form is used thereafter.

6 Conclusion · A Full set of operators · B Constraint Details

DRTriton combines synthetic-data training with curriculum reinforcement learning to produce Triton kernels that generalize to complex real-world GPU kernels. Its appendices specify the operator coverage and shape, broadcasting, and size constraints used in synthesis.

  • 6 Conclusion: DRTriton-7B surpasses state-of-the-art commercial and specialized models on synthetic and real-world benchmarks, showing effective generalization from synthetic-data SFT warmup and RL to human-authored GPU kernels.The framework trains PyTorch-to-Triton LLMs through large-scale synthetic data and curriculum reinforcement learning.
  • A Full set of operators: The operator set includes normalization operators, one- to three-dimensional pooling, one- to three-dimensional convolution, and one- to three-dimensional transposed convolution.The listed normalization operators are BatchNorm, LayerNorm, GroupNorm, and InstanceNorm; each pooling operator has one input, while convolutional operators have two.
  • B Constraint Details: Broadcasting permits dimensions that are equal or have size 1 to align, including tensors with different orders by right-aligning their dimensions.For example, shape (5, ) is treated as (1, 1, 5) when broadcast with (3, 4, 5).
  • B Constraint Details: Constraint checking uses N = max(in1.n, in2.n, . . . , inm.n) and pads leading dimensions with size 1 to express constraints uniformly.Unless otherwise stated, constraints apply to these right-aligned inputs rather than the original tensors.
  • B Constraint Details: Elementwise operators allow broadcast-compatible inputs, while reduction operators constrain output dimensions according to dim and keepdim.Elementwise outputs use the maximum aligned dimension sizes; reductions remove or retain the reduced dimension depending on keepdim.
  • B Constraint Details: Matmul requires at least two dimensions per input, matches the contracting dimensions, and broadcasts compatible batch dimensions.Transpose swaps the final two dimensions while preserving tensor order and preceding dimensions.
  • B Constraint Details: ConvNd constrains channel divisibility by groups and output channels, while PoolNd and ConvTransposeNd are handled similarly.Normalization operators preserve input shape, with additional type-specific requirements such as GroupNorm channel divisibility by group count.
  • B Constraint Details: The experiments constrain operator programs to FLOPSmin = 234, FLOPSmax = 235, SIZEmin = 32, and SIZEmax = 232.These bounds define the FLOPs and size ranges used in the experiments.

C SFT Dataset Construction Details

The SFT dataset is built through uniform synthetic generation, execution-based filtering, and targeted augmentation, producing 2,026 PyTorch–Triton pairs covering all 61 operators. Iterative evaluation concentrates additional data on operators with low model success rates.

  • Dataset Construction: The multi-stage construction process is designed to provide coverage and quality across diverse PyTorch operators.
  • Initial Data Generation: 12,200 synthetic programs were generated by sampling 200 programs for each of 61 fundamental operators, then filtering out kernels that failed execution-based correctness validation.DeepSeek-R1 generated the Triton implementations, and the filtering stage yielded 1,464 valid samples.
  • Iterative Refinement: An initial SFT-plus-RL model was evaluated on 100 test samples per operator to identify operators with extremely low success rates and insufficient training-data coverage.These operator-level results guided the targeted augmentation stage.
  • Final Dataset: 2,026 PyTorch–Triton pairs form the final dataset, combining 1,464 initial samples with 562 targeted augmentations across all 61 operators.The final distribution preserves broad coverage from uniform sampling while strengthening representation of operators that challenged the base model.

D Additional Experiment Results · D.1 Speedup Curves on KernelBench

Figure 4 compares DRTriton with test-time search against baseline models across all three KernelBench difficulty levels. It displays the distribution of generated-kernel speedups over PyTorch, with 1× and 2× reference lines.

  • D.1 Speedup Curves on KernelBench: The figure spans all three KernelBench difficulty levels.Its curves show how speedup distributions vary across the benchmark’s difficulty settings.
  • D.1 Speedup Curves on KernelBench: Figure 4 compares DRTriton with test-time search against all baseline models on KernelBench.The comparison covers speedup distributions across the benchmark’s three difficulty levels.
  • D.1 Speedup Curves on KernelBench: Each curve reports generated Triton-kernel speedups over the corresponding PyTorch programs.The distributions are shown across percentages of the testing data.
  • D.1 Speedup Curves on KernelBench: The curves represent speedup distributions across percentages of testing data.This presentation shows the distribution rather than a single aggregate speedup value.
  • D.1 Speedup Curves on KernelBench: A red dashed flat line marks the 1× speedup reference.This line provides the baseline for parity with PyTorch speed.
  • D.1 Speedup Curves on KernelBench: An orange dashed flat line marks the 2× speedup reference.This line provides a higher reference point for interpreting the plotted speedup distributions.

D.2 Ablation on Speed Reward Function

This ablation compares alternative functional forms for the speed reward component r_s(o) in DRPO, using Qwen-2.5-Coder-1.5B initialized from an SFT checkpoint and trained on Stage 1 data.

  • Ablation setup: The study evaluates different functional forms for the DRPO speed reward component r_s(o).Table 4 reports this speed reward function ablation.
  • Ablation setup: The comparison isolates the design of the speed reward function while holding the base model and training initialization fixed.The supplied passages specify the shared model, data stage, and SFT initialization, but do not provide outcome values.
  • Ablation setup: All variants are trained with DRPO on Stage 1 data using Qwen-2.5-Coder-1.5B as the base model.Training starts from the SFT checkpoint.

D.3 Test-Time Search Overhead · E PyTorch Code Rewriting Example

Test-time search adds modest absolute overhead that is negligible relative to training cost, while the code-rewriting example converts an MLP into a functional representation aligned with DRTriton’s training distribution for optimization.

  • D.3 Test-Time Search Overhead: Test-time search incurs modest absolute overhead on both synthetic and KernelBench benchmarks, remaining negligible relative to overall training cost.Table 5 reports generation and validation times in hours despite evaluating more fragments.
  • E PyTorch Code Rewriting Example: The example starts from a KernelBench multi-layer perceptron with two hidden linear-ReLU layers and a final linear output layer.The test configuration uses batch_size 128, input_size 16384, hidden layer sizes [16384, 16384], and output_size 8192.
  • E PyTorch Code Rewriting Example: torch.export traces the MLP into a functional graph that explicitly materializes weights and biases and expresses execution through fused_operator.The representation uses functional operators for the forward pass rather than the original module-level structure.
  • E PyTorch Code Rewriting Example: The rewritten representation aligns with DRTriton’s training-data distribution, enabling the model to identify optimization opportunities.This alignment is the stated reason for converting the traced graph into the functional form.
  • E PyTorch Code Rewriting Example: The functional implementation initializes and exposes parameters for two 16384×16384 hidden-layer matrices, their biases, and an 8192×16384 output matrix with bias.Weights use Kaiming-uniform initialization, while biases use uniform initialization over [-0.0078125, 0.0078125].
  • E PyTorch Code Rewriting Example: The fused operator applies the first linear layer and ReLU, then the second linear layer and ReLU, followed by the output linear layer.These operations correspond to the original MLP’s sequential forward pass.

F Evaluation Prompts · G Case Study on KernelBench

The section specifies evaluation prompts for synthetic benchmarks and KernelBench, then illustrates DRTriton’s complete pipeline on a Level 3 LeNet-5 task. The case study traces transformations from object-oriented PyTorch through functional IR to an optimized Triton kernel with test-time search.

  • F Evaluation Prompts: KernelBench evaluation uses the Triton backend and the default prompt provided by the KernelBench pipeline.This differs from the detailed synthetic-benchmark prompt template.
  • F Evaluation Prompts: Synthetic-benchmark prompts require converting PyTorch into an efficient Triton kernel with identical numerical results while maximizing performance.The template includes detailed requirements for kernel structure, memory efficiency, and performance optimization, plus a concrete guiding example.
  • F Evaluation Prompts: The synthetic prompt requires exactly one @triton.jit kernel, a triton_fused_operator entrypoint, and fusion of all operations into one kernel.It also requests coalesced accesses, correct stride handling, masked loads and stores, and minimized redundant memory transactions.
  • F Evaluation Prompts: The prompt emphasizes handling input tensor shapes with masks to prevent out-of-bounds access and prioritizing GPU execution efficiency.It also asks for optimal block sizes and maximized GPU utilization.
  • F Evaluation Prompts: Synthetic evaluation prompts require returning only Triton code wrapped in <triton_code> tags and provide PyTorch input, operator, and output-code examples.The example uses two length-128 tensors and an elementwise addition implemented with masked loads and stores.
  • G Case Study on KernelBench: DRTriton’s LeNet-5 case study shows three stages: object-oriented PyTorch, torch.export functional IR, and an optimized Triton kernel produced with test-time search.The demonstration uses a Level 3 KernelBench task.

G.1 Original PyTorch Code … H GRPO Objective

The appendix traces LeNet-5 from object-oriented PyTorch through a functional rewrite to a test-time-searched Triton/PyTorch implementation, then defines GRPO using group-relative rewards and a reference-policy KL term.

  • G.1 Original PyTorch Code: The original LeNet-5 uses two convolutional layers, three fully connected layers, ReLU activations, max pooling, and flattening in its forward pass.The implementation is written in standard object-oriented PyTorch with nn.Module and class-based layer definitions.
  • G.1 Original PyTorch Code: The example tests LeNet-5 with batch_size = 4096, num_classes = 20, and random inputs shaped [4096, 1, 32, 32].Initialization inputs provide the number of output classes, while runtime inputs are generated as random tensors.
  • G.2 Functionally Rewritten PyTorch Code: The object-oriented model is transformed into a functional PyTorch program that explicitly passes initialized weights and biases into fused_operator.Its functional forward sequence calls conv2d, ReLU, max pooling, linear layers, and returns the final output tensor.
  • G.3 Optimized Triton Kernel with Test-Time Search: Test-time search fuses the first F.conv2d and F.relu into one Triton kernel while retaining the remaining operations in PyTorch.The optimized composition applies max pooling, the second convolution and activation, subsequent pooling, flattening, and linear layers after the fused kernel.
  • G.3 Optimized Triton Kernel with Test-Time Search: The generated Triton kernel computes convolution outputs with masked input loads, accumulates weighted values, adds bias, applies ReLU, and stores results.The wrapper derives tensor shapes, allocates output, uses BLOCK_SIZE = 1024, and launches the kernel over the output grid.
  • H GRPO Objective: GRPO [Shao et al., 2024] generates multiple outputs per input and optimizes a group-relative objective for policy πθ against a fixed reference policy.The objective includes a Kullback-Leibler divergence term involving πref, while the advantage normalizes each output’s reward relative to its group.

NeurIPS Paper Checklist

The checklist finds that the paper’s claims, limitations, experimental disclosures, and LLM usage are appropriately documented, while open-access release and statistical uncertainty reporting remain incomplete. The work is algorithmic and empirical rather than theoretical.

  • Claims and scope: The paper accurately states its three contributions and supports them with quantitative results on synthetic and KernelBench benchmarks.These contributions are CSP-DAG, curriculum RL with decoupled rewards, and test-time search.
  • Limitations: The paper acknowledges that its operator set is limited to 61 PyTorch operators and that it targets Triton as the compilation backend.Sparse operations, custom CUDA extensions, and native CUDA generation are identified as future directions.
  • Reproducibility: The experimental setup provides the model, hyperparameters, curriculum configuration, evaluation benchmarks, and evaluation prompt needed to reproduce the reported results.Reported details include DRPO parameters (β0, τ, λ) = (100, 5, 0.1) and KL constraint bound δ = 0.001.
  • Open access: The paper does not provide open access to its data and code because release is deferred until after acceptance.The checklist marks this criterion [No].
  • Statistical significance: The paper reports no error bars or confidence intervals because its metrics are computed on fixed held-out benchmark sets and repeated training runs are computationally prohibitive.The primary metrics are Pass@1 accuracy, Faster1 percentage, and average speedup.
  • LLM usage: LLMs are a core methodological component, with DeepSeek-R1 and GPT-5.2 used to generate SFT training data.The checklist marks disclosure of this usage [Yes].
Loading 2603.21465v2…