AutoTriton：在LLM中通过强化学习实现自动化Triton编程

深度学习中的内核开发需要在硬件之间优化计算单元，同时通过大量的经验性调优来平衡内存管理、并行性以及硬件特定的优化。尽管像 Triton 这样的领域特定语言通过抽象底层细节简化了 GPU 编程，但开发人员仍必须通过迭代实验手动调整关键参数，例如瓦片大小和内存访问模式，这为实现最佳性能和更广泛的采用带来了巨大障碍。在这项工作中，我们引入了 AutoTriton，这是第一个专门用于 Triton 编程的强化学习 (RL) 驱动模型。AutoTriton 利用高质量的数据收集流程进行监督微调 (SFT)，以掌握必要的 Triton 编程专业知识，并随后使用组相对策略优化 (GRPO) 算法进行 RL，结合基于规则的奖励和基于执行的奖励来进一步提高 Triton 编程能力。在 TritonBench 和 KernelBench 的五个评估通道上的实验表明，我们的 8B 模型 AutoTriton 实现了与主流大型模型（包括 Claude-4-Sonnet 和 DeepSeek-R1-0528）相当的性能。进一步的实验分析表明了 AutoTriton 中每个模块的关键作用，包括 SFT 阶段、RL 阶段和奖励设计策略。这些发现突显了 RL 自动生成高性能内核的潜力，并且由于高性能内核是 AI 系统的核心组件，这一突破为构建更高效的 AI 系统奠定了重要基础。模型和代码将发布于 https://github.com/AI9Stars/AutoTriton。