Research map / Large language models / Agents, code and math

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

NeurIPS 2026 · Large language models › Agents, code and math · approach: LLMs for optimization and heuristics · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Trains Qwen-2.5-Coder-7B to convert PyTorch programs into Triton kernels via SFT on 2,026 pairs, then curriculum RL with decoupled rewards (DRPO) on 100k CSP-DAG synthetic programs, plus test-time fragment search.

Problem

LLMs such as GPT-5.2 and Claude-Sonnet-4.5 struggle to convert PyTorch reference programs into efficient Triton (and hence CUDA) kernels, and existing LLM kernel-training approaches are limited by small, uncontrolled, real-world training data and sparse rewards.

Approach

Generates valid, difficulty-controlled PyTorch programs by formulating DAG construction plus tensor-shape constraints as a CP-SAT problem (CSP-DAG), then trains with curriculum RL using decoupled correctness and log-speed rewards (DRPO), and composes the fastest verified Triton fragments at test time.

Key results

• On the synthetic benchmark, DRTriton with test-time search reaches 99% accuracy at Level 20 and average speedup 1.57x, versus 1.34x for GPT-5.2 (which gets 0% at Level 20).
• On KernelBench Level 2, DRTriton with test-time search achieves 96% accuracy with 92% of kernels faster than Torch Eager, versus 32% accuracy and 23% faster for GPT-5.2.
• On KernelBench Level 3, DRTriton reaches 76% accuracy with 54% faster than Torch Eager and 34% faster than torch.compile, though AutoTriton with search has higher Level 3 accuracy (86%).

Surprising findings

• Without test-time search, DRTriton's Level 20 accuracy is 0%, and test-time search is what lifts harder programs to high accuracy.
• Longer and harder programs achieve higher accuracy and speedup after test-time search, attributed to more fusion opportunities.

Limitations

• Operator set limited to 61 PyTorch operators; no sparse ops or custom CUDA extensions.
• Without test-time search, performance on complex (Level 5 and 20) programs is low (15% and 0% accuracy).
• On KernelBench Level 3, AutoTriton with test-time search achieves higher accuracy than DRTriton.

Experimental setup

Datasets: KernelBench CSP-DAG synthetic PyTorch programs DeepSeek-R1 and GPT-5.2 generated PyTorch-Triton pairs

Benchmarks: KernelBench synthetic benchmark (Levels 1, 2, 5, 20)

Metrics: Acc (pass@1) Faster1 geometric mean speedup

Baselines: GPT-5.2 Claude-Sonnet-4.5 DeepSeek-R1 Qwen-3-Coder-480B AutoTriton torch.compile

Builds on: DRPO GRPO KernelBench torch.export CP-SAT

Closest papers by meaning

• $\mathcal{B}$-Coder: Value-Based Deep Reinforcement Learning for Program Synthesis (ICLR 2024 · similarity 0.85)
• Breaking the SFT Plateau: Multimodal Structured Reinforcement Learning for Chart-to-Code Generation (ICLR 2026 · similarity 0.85)
• TorchRL: A data-driven decision-making library for PyTorch (ICLR 2024 · similarity 0.85)
• UFO-RL: Uncertainty-Focused Optimization for Efficient Reinforcement Learning Data Selection (NeurIPS 2025 · similarity 0.85)
• EfficientRollout: System-Aware Self-Speculative Decoding for RL Rollouts (NeurIPS 2026 · similarity 0.85)
• AREAL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning (NeurIPS 2025 · similarity 0.85)