AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs
作者: Shangzhan Li, Zefan Wang, Ye He, Yuxuan Li, Qi Shi, Jianling Li, Yonggang Hu, Wanxiang Che, Xu Han, Zhiyuan Liu, Maosong Sun
分类: cs.LG, cs.CL
发布日期: 2025-07-08
🔗 代码/项目: GITHUB
💡 一句话要点
AutoTriton：利用强化学习在LLM中自动生成高性能Triton程序
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: Triton编程 强化学习 自动优化 Kernel生成 深度学习编译器
📋 核心要点
- 深度学习kernel开发需要手动调整tile大小和内存访问模式等参数，过程繁琐且依赖经验。
- AutoTriton利用强化学习自动生成高性能Triton程序，无需人工调优，降低了开发门槛。
- 实验表明，AutoTriton的8B模型性能可与Claude-4-Sonnet等主流大模型媲美。
📝 摘要（中文）
深度学习中的Kernel开发需要在硬件上优化计算单元，同时平衡内存管理、并行性和硬件特定优化，这通常需要大量的经验性调整。尽管Triton等领域特定语言通过抽象底层细节简化了GPU编程，但开发者仍然需要手动调整诸如tile大小和内存访问模式等关键参数，这阻碍了最佳性能和更广泛的应用。本文提出了AutoTriton，这是第一个基于强化学习（RL）的Triton编程模型。AutoTriton首先通过监督微调（SFT）和一个高质量的数据收集流程来获得Triton编程的专业知识，然后使用Group Relative Policy Optimization (GRPO)算法进行强化学习，结合基于规则的奖励和基于执行的奖励，以进一步提高Triton编程能力。在TritonBench和KernelBench的五个评估通道上的实验表明，我们的8B模型AutoTriton实现了与主流大型模型（包括Claude-4-Sonnet和DeepSeek-R1-0528）相当的性能。进一步的实验分析表明了AutoTriton中每个模块的关键作用，包括SFT阶段、RL阶段和奖励设计策略。这些发现强调了RL在自动生成高性能内核方面的潜力，并且由于高性能内核是AI系统的核心组件，因此这一突破为构建更高效的AI系统奠定了重要的基础。
🔬 方法详解
问题定义：论文旨在解决深度学习kernel开发中手动调整Triton程序参数的难题。现有方法需要开发者具备丰富的经验，并通过大量的实验来寻找最佳的tile大小、内存访问模式等参数，效率低下且成本高昂。这种手动调优的方式阻碍了Triton语言的广泛应用，也限制了深度学习模型的性能优化。
核心思路：论文的核心思路是利用强化学习（RL）来自动搜索最佳的Triton程序参数配置。通过将Triton程序生成过程建模为一个RL问题，AutoTriton能够学习到如何根据不同的硬件和计算任务，自动生成高性能的Triton kernel。这种方法避免了人工调优的繁琐过程，并有望实现更优的性能。
技术框架：AutoTriton的技术框架主要包含两个阶段：监督微调（SFT）和强化学习（RL）。首先，通过SFT，AutoTriton学习Triton编程的基本知识和技巧，使其具备生成Triton程序的能力。然后，利用RL，AutoTriton进一步优化生成的Triton程序，使其在目标硬件上达到最佳性能。RL阶段使用Group Relative Policy Optimization (GRPO)算法，并结合基于规则的奖励和基于执行的奖励。
关键创新：AutoTriton的关键创新在于将强化学习应用于Triton程序的自动生成。与传统的基于规则或模板的方法不同，AutoTriton能够通过学习自动适应不同的硬件和计算任务，生成定制化的Triton kernel。此外，AutoTriton还提出了GRPO算法，并设计了有效的奖励函数，以提高RL的训练效率和性能。
关键设计：AutoTriton的关键设计包括：1) 高质量的数据收集流程，用于SFT阶段的训练数据；2) GRPO算法，用于RL阶段的策略优化；3) 结合基于规则和基于执行的奖励函数，用于指导RL的学习过程。基于规则的奖励可以提供一些先验知识，而基于执行的奖励则可以根据实际的性能反馈来调整策略。具体的参数设置和网络结构等细节在论文中没有详细描述，属于未知信息。
🖼️ 关键图片
📊 实验亮点
AutoTriton在TritonBench和KernelBench的五个评估通道上进行了实验，结果表明其8B模型实现了与Claude-4-Sonnet和DeepSeek-R1-0528等主流大型模型相当的性能。实验还验证了SFT阶段、RL阶段和奖励设计策略对AutoTriton性能的关键作用，表明了强化学习在自动生成高性能kernel方面的潜力。
🎯 应用场景
AutoTriton可应用于各种需要高性能kernel的深度学习任务，例如图像识别、自然语言处理和推荐系统。通过自动生成优化的Triton程序，AutoTriton可以显著提高这些任务的计算效率，降低计算成本，并加速AI模型的部署和应用。未来，AutoTriton有望成为深度学习编译器和自动优化工具的重要组成部分。
📄 摘要（原文）
Kernel development in deep learning requires optimizing computational units across hardware while balancing memory management, parallelism, and hardware-specific optimizations through extensive empirical tuning. Although domain-specific languages like Triton simplify GPU programming by abstracting low-level details, developers must still manually tune critical parameters such as tile sizes and memory access patterns through iterative experimentation, creating substantial barriers to optimal performance and wider adoption. In this work, we introduce AutoTriton, the first model dedicated to Triton programming powered by reinforcement learning (RL). AutoTriton performs supervised fine-tuning (SFT) to be equipped with essential Triton programming expertise using a high-quality data gathering pipeline, and conducts RL with Group Relative Policy Optimization (GRPO) algorithm, combining a rule-based reward and an execution-based reward to further improve Triton programming ability, sequentially. Experiments across five evaluation channels of TritonBench and KernelBench illustrate that our 8B model AutoTriton achieves performance comparable to mainstream large models, including Claude-4-Sonnet and DeepSeek-R1-0528. Further experimental analysis demonstrates the crucial role of each module within AutoTriton, including the SFT stage, the RL stage, and the reward design strategy. These findings underscore the promise of RL for automatically generating high-performance kernels, and since high-performance kernels are core components of AI systems, this breakthrough establishes an important foundation for building more efficient AI systems. The model and code will be available at https://github.com/AI9Stars/AutoTriton.