DRTriton

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

DRTriton trains language models to convert PyTorch programs into equivalent, high-performance Triton GPU kernels. It combines three ingredients:

1. A synthetic data generator that builds valid PyTorch tensor-computation graphs with comprehensive operator coverage using constraint programming.
2. Curriculum reinforcement learning (built on verl) that rewards both conversion correctness and execution speedup over eager PyTorch.
3. A test-time search / subproblem decomposition procedure that splits a program into fusable subproblems, generates a kernel for each, and recombines them.

DRTriton-7B achieves a speedup over PyTorch on 92% of KernelBench Level 2 tasks, compared to 23% for GPT-5.2 and 19% for Claude-Sonnet-4.5 (see paper).

Reference

If you use this code, please cite:

Siqi Guo, Ming Lin, Tianbao Yang. DRTriton: Large-Scale Synthetic Data Driven Reinforcement Learning for Triton Kernel Generation. arXiv:2603.21465. https://arxiv.org/pdf/2603.21465

@article{guo2026drtriton,
  title={DRTriton: Large-Scale Synthetic Data Driven Reinforcement Learning for Triton Kernel Generation},
  author={Guo, Siqi and Lin, Ming and Yang, Tianbao},
  journal={arXiv preprint arXiv:2603.21465},
  year={2026}
}

Repository layout

Path	Purpose
generator/	Synthetic PyTorch computation-graph generator (CP-SAT based). See generator/README.md.
evaluation/	Rollout + correctness/speedup evaluation harness for models and APIs. See evaluation/README.md.
data/	Benchmark datasets and the evaluation harness config.
verl/	Vendored verl RL training framework used for curriculum RL.
install.sh	Environment bootstrap (conda + vLLM/SGLang + verl, editable install).
requirements.txt	Full development dependency set.

Datasets in data/

File	Description
data/synthetic_bench.jsonl	Synthetic benchmark produced by the generator, with per-sample subproblems.
data/kernelbench_level1.jsonl	KernelBench Level 1 tasks (single primitive ops), converted to the fused_operator / get_inputs format.
data/kernelbench_level2.jsonl	KernelBench Level 2 tasks (fused operator sequences).
data/kernelbench_level3.jsonl	KernelBench Level 3 tasks (full model architectures).
data/config.yaml	Templates the evaluation harness fills with generated Triton code to check correctness, NaN/Inf safety, dtype/shape match, and timing.

Each task is a JSON line containing a pytorch field with two functions:

def get_inputs():
    # returns a list of input tensors
    ...

def fused_operator(*inputs):
    # the computation to be turned into a Triton kernel
    return [outputs]

Installation

git clone <this-repo> lpc_new && cd lpc_new
bash install.sh

install.sh creates a local conda env at ./venv (Python 3.10), installs vLLM/SGLang (and optionally Megatron), installs this repo editable, and adds the training/eval extras (hydra-core, datasets, accelerate, wandb, ...). A CUDA GPU is required for kernel execution and validation.

For a data-generation-only environment:

pip install torch ortools

Usage

1. Generate synthetic training data

python generator/gen.py \
    --output_file data/random_torch_l1.jsonl \
    --level 1 \
    --target_count 10000

--level controls the number of fused operators per sample (0: single op, 1: pairs, 2/3/4: 3/5/10 randomly sampled ops). Generation is multi-GPU and each candidate is executed before being written out. See generator/README.md for the full option list, supported operators (47+), and FLOPS/shape constraints.

2. Curriculum RL training

Training uses the vendored verl framework (PPO / GRPO-style pipelines under verl/trainer/, with the batched Triton reward manager in verl/workers/reward_manager/batch.py). The reward combines a correctness term and a speedup term (speed_reward_type=log), and the DRPO loss (loss_type=drpo) trades the two off with a temperature tau and a weight Lambda that follow the curriculum. Correctness/speedup are scored by the custom reward function verl/utils/reward_score/triton_verify.py.

Starting from an SFT checkpoint, launch python3 -m verl.trainer.main_ppo with the DRPO overrides. Example (8 GPUs/node × 4 nodes, Qwen2.5-Coder-7B):

#!/bin/bash
set -x

export WANDB_API_KEY=<your-wandb-key>
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# Base model checkpoint from SFT training
SFT_CHECKPOINT=/path/to/checkpoints/triton-sft/qwen-2.5-coder-7b-instruct/global_step_615
EXP_NAME="qwen-2.5-coder-7b-drpo-level1-5-tau5-lambda0.1-log-dist"
PROJECT_NAME="verl_drpo_triton"

python3 -m verl.trainer.main_ppo \
    algorithm.adv_estimator=grpo \
    data.train_files=data/stage_merged_all.parquet \
    data.val_files=data/stage_merged_all.parquet \
    data.train_batch_size=256 \
    data.max_prompt_length=8192 \
    data.max_response_length=8192 \
    data.filter_overlong_prompts=True \
    data.truncation='error' \
    actor_rollout_ref.model.path=${SFT_CHECKPOINT} \
    actor_rollout_ref.rollout.max_num_batched_tokens=32000 \
    actor_rollout_ref.actor.optim.lr=1e-6 \
    actor_rollout_ref.model.use_remove_padding=True \
    actor_rollout_ref.actor.ppo_mini_batch_size=64 \
    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \
    actor_rollout_ref.actor.use_kl_loss=False \
    actor_rollout_ref.actor.kl_loss_coef=0.001 \
    actor_rollout_ref.actor.kl_loss_type=low_var_kl \
    +actor_rollout_ref.actor.ppo_kl_type=kl \
    +actor_rollout_ref.actor.delta=1e-4 \
    +actor_rollout_ref.actor.beta=1e3 \
    +actor_rollout_ref.actor.tau=5 \
    +actor_rollout_ref.actor.Lambda=0.1 \
    +actor_rollout_ref.actor.loss_type=drpo \
    +actor_rollout_ref.actor.speed_reward_type=log \
    actor_rollout_ref.actor.entropy_coeff=0 \
    actor_rollout_ref.model.enable_gradient_checkpointing=True \
    actor_rollout_ref.actor.fsdp_config.param_offload=True \
    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \
    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \
    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \
    actor_rollout_ref.rollout.name=vllm \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \
    actor_rollout_ref.rollout.n=8 \
    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \
    actor_rollout_ref.ref.fsdp_config.param_offload=True \
    actor_rollout_ref.actor.use_torch_compile=False \
    actor_rollout_ref.ref.use_torch_compile=False \
    algorithm.use_kl_in_reward=False \
    trainer.critic_warmup=0 \
    trainer.logger=['console','wandb'] \
    trainer.project_name=${PROJECT_NAME} \
    trainer.experiment_name=${EXP_NAME} \
    trainer.balance_batch=False \
    trainer.n_gpus_per_node=8 \
    trainer.nnodes=4 \
    trainer.save_freq=20 \
    trainer.test_freq=-1 \
    trainer.total_epochs=1 \
    trainer.max_actor_ckpt_to_keep=2 \
    trainer.max_critic_ckpt_to_keep=2 \
    trainer.val_before_train=False \
    custom_reward_function.path=verl/utils/reward_score/triton_verify.py $@

Adjust data.train_files, actor_rollout_ref.model.path, trainer.nnodes / trainer.n_gpus_per_node, and the tau / Lambda curriculum values for your setup.

3. Evaluation

The evaluation/ scripts run a model over a benchmark, extract the generated Triton code, and score correctness + speedup against eager PyTorch. Typical flow:

# Pass@n rollout with a local model via vLLM
python evaluation/vllm_gen_eval.py \
    --model_path /path/to/checkpoint \
    --data_path data/kernelbench_level2.jsonl \
    --output_dir evaluation_results/kb_l2/model.json \
    --max_tokens 8192 \
    --prompt_style original \
    --tensor_parallel_size 4

# Score a results file (correctness + speedup)
python evaluation/eval_from_json_fast.py --input evaluation_results/kb_l2/model.json

Test-time search via subproblem decomposition:

# 1. Split programs into subproblems
python evaluation/subproblem.py --input data/eval_v4.jsonl --output data/eval_v4_sub.jsonl --max-ops 6

# 2. Generate a Triton kernel per subproblem (vLLM / Claude / OpenAI variants exist)
python evaluation/vllm_gen_subproblem.py \
    --model_path /path/to/model \
    --data_path data/eval_v4_sub.jsonl \
    --output_path evaluation_results/tts/model.jsonl \
    --tensor_parallel_size 4 --max_tokens 8192 --prompt_style original

# 3. Stitch subproblem kernels back into runnable programs
python evaluation/subproblem_replace.py --input evaluation_results/tts/model.jsonl

# 4. Correctness + speedup evaluation
python evaluation/eval_from_subproblem.py \
    --input evaluation_results/tts/model_with_replacements.jsonl \
    --output evaluation_results/tts/model_annotated.jsonl

Additional helpers: evaluation/op_success_rate.py (per-operator pass rate), evaluation/eval_summary_v0.py (pass@k aggregation), evaluation/compare_eval_results.py. See evaluation/README.md for more.

How correctness is checked

For every task the harness runs the reference fused_operator and the generated triton_fused_operator on the same inputs (repeated with fresh random inputs) and requires: matching output count, type, shape, and dtype; no NaN/Inf; and value agreement (torch.norm(x - y) / (1e-6 + torch.norm(x)) < 1e-3 for floats, exact equality otherwise). Speedup is measured with CUDA events after warmup. The exact templates live in data/config.yaml.

Acknowledgements

The verl/ directory vendors verl (Apache-2.0). Benchmark tasks are adapted from KernelBench.