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

📄 arXiv: 2603.21465v1 📥 PDF

作者: Siqi Guo, Ming Lin, Tianbao Yang

分类: cs.CL, cs.LG

发布日期: 2026-03-23

💡 一句话要点

DRTriton：利用大规模合成数据强化学习生成Triton内核，显著提升CUDA内核效率。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: CUDA内核生成 强化学习 合成数据 大型语言模型 Triton 性能优化 深度学习 代码转换

📋 核心要点

  1. 现有LLM在将PyTorch代码转换为CUDA内核方面表现不佳，难以满足生成式AI行业对高效CUDA内核的需求。
  2. DRTriton通过合成数据强化学习训练LLM，生成优化的Triton内核，进而编译为CUDA内核，提升转换效率和推理速度。
  3. 实验表明，DRTriton在KernelBench Level 2上相比GPT-5.2和Claude-Sonnet-4.5，显著提高了CUDA内核的加速比例。

📝 摘要（中文）

本文提出DRTriton，一个可扩展的学习框架，用于训练大型语言模型（LLMs）将PyTorch代码转换为高度优化的Triton内核，这些内核在运行时被编译为CUDA内核。DRTriton包含三个关键组件：(i) 数据合成算法CSP-DAG，保证在算子空间上的完全覆盖和无偏均匀采样，并控制难度；(ii) 具有解耦奖励的课程强化学习，高效地同时优化转换成功率和推理速度；(iii) 测试时搜索算法，进一步提高生成的Triton内核的推理速度。值得注意的是，尽管仅在合成数据上训练，DRTriton能够有效地泛化到真实世界的CUDA内核，这些内核甚至对人类专家来说也具有挑战性。实验结果表明，与GPT-5.2（23%）和Claude-Sonnet-4.5（19%）相比，DRTriton-7B在KernelBench Level 2上实现了92%的加速。

🔬 方法详解

问题定义：论文旨在解决将PyTorch代码高效转换为CUDA内核的问题。现有方法，特别是直接使用大型语言模型（LLMs）如GPT-5.2和Claude-Sonnet-4.5，在这一特定任务上表现不佳，无法生成足够优化的CUDA内核，导致工程效率低下。

核心思路：论文的核心思路是利用强化学习，通过大规模合成数据训练LLM，使其能够生成高度优化的Triton内核。Triton内核随后被编译为CUDA内核。这种方法的核心在于，通过精心设计的合成数据和强化学习策略，使模型能够学习到从PyTorch到优化CUDA内核的有效转换规则。

技术框架：DRTriton框架包含三个主要模块：(1) 数据合成算法CSP-DAG，用于生成高质量的合成训练数据；(2) 课程强化学习模块，使用解耦奖励函数，同时优化转换成功率和推理速度；(3) 测试时搜索算法，用于进一步优化生成的Triton内核的推理速度。整体流程是从PyTorch代码开始，通过训练好的LLM生成Triton内核，然后使用测试时搜索进行优化，最终编译为CUDA内核。

关键创新：该论文的关键创新在于结合了数据合成、课程强化学习和测试时搜索，形成一个完整的、可扩展的训练框架。特别是CSP-DAG算法，能够保证算子空间的全面覆盖和无偏采样，为强化学习提供高质量的训练数据。此外，解耦奖励函数的设计，使得模型能够同时关注转换成功率和推理速度，从而生成更优的CUDA内核。

关键设计：CSP-DAG算法的具体设计细节（例如，如何控制算子空间的难度，如何保证无偏采样）是关键。课程强化学习中，解耦奖励函数的设计（如何平衡转换成功率和推理速度的权重）以及具体的强化学习算法选择（例如，PPO、SAC等）也是重要的技术细节。测试时搜索算法的具体实现（例如，搜索空间的设计，搜索策略的选择）也会影响最终的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DRTriton-7B在KernelBench Level 2测试中，相比GPT-5.2和Claude-Sonnet-4.5，实现了显著的性能提升。具体来说，DRTriton-7B在92%的测试用例上实现了加速，而GPT-5.2和Claude-Sonnet-4.5的加速比例分别为23%和19%。这表明DRTriton在生成高性能CUDA内核方面具有显著优势，即使面对真实世界的复杂内核，也能表现出良好的泛化能力。

🎯 应用场景

DRTriton具有广泛的应用前景，尤其是在生成式AI领域。它可以用于加速各种深度学习模型的推理速度，降低计算成本。此外，该方法还可以应用于其他需要高性能计算的领域，例如科学计算、金融建模等。通过自动化CUDA内核的生成，DRTriton可以显著降低开发成本，提高开发效率，并推动相关领域的发展。

📄 摘要（原文）

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.