用合成数据训练大模型，自动高效生成Triton内核。
DRTriton: Large-Scale Synthetic Data Driven Reinforcement Learning for Triton Kernel Generation
- 用合成数据和课程强化学习训练模型转换PyTorch到Triton。
- 在92%的基准任务上比PyTorch快，远超GPT-5.2和Claude-Sonnet-4.5。
- 适合需要高效编译CUDA内核的AI工程团队使用。
在生成式AI领域，开发高效的CUDA内核是一项基础但极具挑战的任务。近期研究利用大语言模型（LLMs）将PyTorch参考实现自动转换为CUDA内核，显著降低工程成本。然而，当前最先进的LLMs如GPT-5.2和Claude-Sonnet-4.5仍难以胜任此任务。为此，我们提出DRTriton，一个可扩展的学习框架，用于训练LLMs将PyTorch程序转化为高度优化的Triton内核，并在运行时编译为CUDA内核。DRTriton包含三个核心组件：(i) CSP-DAG数据合成算法，确保在算子空间中全覆盖且无偏均匀采样，同时控制难度；(ii) 解耦奖励的课程强化学习框架，联合优化转换成功率与执行速度；(iii) 测试时搜索算法，进一步提升生成Triton内核的执行效率。通过在少量人工标注的PyTorch-Triton对上进行SFT预热，DRTriton在合成的PyTorch程序上通过RL训练后，能有效泛化至真实世界中甚至对人类专家也极具挑战的CUDA内核。实验表明，DRTriton-7B在KernelBench Level 2任务中实现了92%的加速率，而GPT-5.2和Claude-Sonnet-4.5分别仅为23%和19%。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。