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

• URL: http://arxiv.org/abs/2603.21465v1
• Date: Mon, 23 Mar 2026 00:59:35 GMT
• Title: DRTriton: Large-Scale Synthetic Data Reinforcement Learning for Triton Kernel Generation
• Abstract summary: Large Language Models (LLMs) automatically convert PyTorch reference implementations to kernels, significantly reducing engineering efforts.<n>We propose DRTriton, a scalable learning framework for training LLMs to convert PyTorch codes into highly optimized Triton kernels.<n> DRTriton consists of three key components: (i) a synthetic algorithm CSP-DAG that guarantees full coverage and unbiased uniform sampling over the operator space with difficulty; (ii) a curriculum reinforcement learning with decoupled reward efficiently optimize conversion success rate and 19% speed simultaneously; and (iii) a test-time search algorithm that further improves inference speed.
• Score: 35.68487613748886
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• RealisticTritonBench: A Benchmark for Triton-Kernel Generation in Real-World AI Frameworks [7.371299443009948]
  We introduce RealisticTritonBench, the first benchmark to derive Triton kernel generation tasks from real-world pull requests in popular AI frameworks.<n> RealisticTritonBench integrates PRs that modify Triton kernels from popular open-source AI frameworks and transforms them into generation tasks with concrete engineering contexts.<n>We evaluate leading LLMs on RealisticTritonBench and find that they still struggle with real-world Triton kernel generation tasks.
  arXiv  Detail & Related papers  (2026-08-12T12:38:17Z)
• BWTA: Accurate and Efficient Binarized Transformer by Algorithm-Hardware Co-design [71.97035034203275]
  We analyze zero-point distortion in binarization and propose a Binary Weights & Ternary Activations (BWTA) quantization scheme.<n>For training, we propose Smooth Multi-Stage Quantization, combining a Levelwise Degradation Strategy and a MagnitudeAlignment Projection Factor.<n>Experiments show that BWTA approaches full-precision performance for Transformer-based models, with an average 3.5% drop on GLUE and less than 2% drop on five tasks.
  arXiv  Detail & Related papers  (2026-04-05T04:25:07Z)
• Towards Cold-Start Drafting and Continual Refining: A Value-Driven Memory Approach with Application to NPU Kernel Synthesis [68.7701048879757]
  EvoKernel is a self-evolving agentic framework that automates the lifecycle of kernel synthesis.<n>It learns stage-specific Q-values that prioritize experiences based on their contribution to the current objective.<n>It improves models' correctness from 11.0% to 83.0% and a median speedup of 3.60x over initial drafts.
  arXiv  Detail & Related papers  (2026-03-11T14:57:06Z)
• CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation [51.72529978689561]
  Agent is a large-scale agentic reinforcement learning system that develops kernel expertise through three components.<n>Agent delivers 100%, 100%, and 92% faster rate over torchcompile on KernelBench.
  arXiv  Detail & Related papers  (2026-02-27T18:58:05Z)
• Fine-Tuning GPT-5 for GPU Kernel Generation [5.109141377873154]
  We present Makora's environment and tools for reinforcement learning finetuning of frontier models.<n>In the single-attempt setting, our fine-tuned model improves kernel correctness from 43.7% to 77.0%.<n>When integrated into a full coding agent, it is able to solve up to 97.4% of problems in an expanded KernelBench suite.
  arXiv  Detail & Related papers  (2026-02-11T16:22:54Z)
• Dr. Kernel: Reinforcement Learning Done Right for Triton Kernel Generations [32.98036846113632]
  We study reinforcement learning (RL) for kernel generation.<n>We propose Turn-level Reinforce-Leave-One-Out (TRLOO) to provide unbiased advantage estimation.<n>We introduce Profiling-based Rewards (PR) and Profiling-based Rejection Sampling (PRS) to overcome the issue.
  arXiv  Detail & Related papers  (2026-02-05T17:01:09Z)
• AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs [87.8306870967343]
  We introduce AutoTriton, the first model dedicated to Triton programming powered by reinforcement learning (RL)<n>AutoTriton performs supervised fine-tuning (SFT) to be equipped with essential Triton programming expertise using a high-quality data gathering pipeline.<n> Experiments across five evaluation channels of TritonBench and KernelBench illustrate that our 8B model AutoTriton achieves performance comparable to mainstream large models.
  arXiv  Detail & Related papers  (2025-07-08T05:38:24Z)
• Trinity: A General Purpose FHE Accelerator [17.213234642867537]
  We present the first multi-modal FHE accelerator based on a unified architecture, which efficiently supports CKKS, TFHE, and their conversion scheme within a single accelerator. We propose a novel FHE accelerator named Trinity, which incorporates algorithm optimizations, hardware component reuse, and dynamic workload scheduling.
  arXiv  Detail & Related papers  (2024-10-17T10:02:38Z)
• Point-to-Voxel Knowledge Distillation for LiDAR Semantic Segmentation [74.67594286008317]
  This article addresses the problem of distilling knowledge from a large teacher model to a slim student network for LiDAR semantic segmentation. We propose the Point-to-Voxel Knowledge Distillation (PVD), which transfers the hidden knowledge from both point level and voxel level.
  arXiv  Detail & Related papers  (2022-06-05T05:28:32Z)
• PolyScientist: Automatic Loop Transformations Combined with Microkernels for Optimization of Deep Learning Primitives [55.79741270235602]
  We develop a hybrid solution to the development of deep learning kernels. We use the advanced polyhedral technology to automatically tune the outer loops for performance.
  arXiv  Detail & Related papers  (2020-02-06T08:02:34Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.