arXiv:2507.05687cs.LGcs.CL2025-07被引 40

用强化学习自动优化Triton代码,提升AI计算效率。

AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs

  • 基于强化学习自动调优Triton代码参数
  • 性能媲美主流大模型,达80%以上最优率
  • 适合想提升GPU算子性能的工程师

深度学习内核开发需在硬件上优化计算单元,同时平衡内存管理、并行性与硬件特异性优化,依赖大量实验调参。尽管领域专用语言Triton通过抽象底层细节简化编程,开发者仍需手动调整分块大小、内存访问模式等关键参数,阻碍高性能实现与广泛应用。本文提出AutoTriton,首个基于强化学习的Triton编程模型。它先通过监督微调(SFT)获取高质量编程知识,再采用组相对策略优化(GRPO)算法进行强化学习,结合规则奖励与执行反馈奖励逐步提升能力。在TritonBench和KernelBench五个评估通道上的实验表明,其80亿参数模型性能可媲美主流大模型(如Claude-4-Sonnet和DeepSeek-R1-0528)。分析显示,SFT阶段、强化学习阶段及奖励设计均对性能提升至关重要。该研究证明了强化学习在自动生成高性能内核方面的潜力,而高性能内核是人工智能系统的核心组件,为构建更高效AI系统奠定基础。代码与模型将公开于https://github.com/AI9Stars/AutoTriton。

原文摘要 · Abstract (English)

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.

强化学习Triton代码生成算子优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。