StitchCUDA: An Automated Multi-Agents End-to-End GPU Programing Framework with Rubric-based Agentic Reinforcement Learning

📄 arXiv: 2603.02637v1 📥 PDF

作者: Shiyang Li, Zijian Zhang, Winson Chen, Yuebo Luo, Mingyi Hong, Caiwen Ding

分类: cs.MA, cs.CL, cs.PL

发布日期: 2026-03-03

💡 一句话要点

StitchCUDA：一种基于规则的Agent强化学习的自动化多智能体端到端GPU编程框架

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: GPU编程 自动化代码生成 多智能体系统 强化学习 CUDA 端到端程序 性能优化

📋 核心要点

  1. 现有方法在自动GPU编程中主要关注单内核优化，无法有效扩展到端到端程序，限制了实际应用。
  2. StitchCUDA提出一个多智能体框架，包含Planner、Coder和Verifier，实现端到端GPU程序自动生成。
  3. 实验表明，StitchCUDA在端到端GPU编程任务上成功率接近100%，性能优于现有基线方法。

📝 摘要（中文）

现代机器学习工作负载越来越依赖GPU，但由于GPU内核效率和主机端设置的依赖性，实现高端到端性能仍然具有挑战性。虽然基于LLM的方法在自动GPU内核生成方面显示出前景，但先前的工作主要集中在单内核优化，并没有扩展到端到端程序，从而阻碍了实际部署。为了解决这一挑战，本文提出StitchCUDA，一个用于端到端GPU程序生成的多智能体框架，包含三个专门的智能体：Planner用于协调整个系统设计，Coder专门用于逐步实现它，Verifier用于使用Nsys/NCU进行正确性检查和性能分析。为了从根本上提高Coder在端到端GPU编程中的能力，StitchCUDA集成了基于规则的Agent强化学习，该学习基于两个原子技能：任务到代码生成和反馈驱动的代码优化，结合了来自真实执行的规则奖励和基于规则的奖励。因此，Coder学习如何实现高级CUDA编程技术（例如，自定义内核融合，cublas epilogue），并且我们还有效地防止了Coder在基准测试期间的奖励攻击（例如，仅复制PyTorch代码或硬编码输出）。在KernelBench上的实验表明，StitchCUDA在端到端GPU编程任务上实现了近100％的成功率，与多智能体基线相比，加速提高了1.72倍，比RL模型基线提高了2.73倍。

🔬 方法详解

问题定义：现有基于LLM的GPU程序生成方法主要集中于单个kernel的优化，无法处理端到端的GPU程序生成，这限制了它们在实际应用中的价值。同时，如何保证生成代码的正确性和高效性也是一个挑战。

核心思路：StitchCUDA的核心思路是利用多智能体协作的方式，将端到端GPU程序生成任务分解为多个子任务，并由不同的智能体负责。通过强化学习，使智能体能够学习到高效的CUDA编程技巧，并避免奖励攻击。

技术框架：StitchCUDA框架包含三个主要智能体：Planner、Coder和Verifier。Planner负责整体系统设计和任务分解；Coder负责根据Planner的指令逐步生成CUDA代码；Verifier负责代码的正确性检查和性能分析，并提供反馈给Coder。Coder使用基于规则的Agent强化学习，学习任务到代码生成和反馈驱动的代码优化。

关键创新：StitchCUDA的关键创新在于其多智能体架构和基于规则的Agent强化学习方法。多智能体架构能够有效地处理端到端GPU程序生成任务的复杂性，而基于规则的Agent强化学习方法能够提高Coder的学习效率和代码质量，并防止奖励攻击。

关键设计：Coder的强化学习基于两个原子技能：任务到代码生成和反馈驱动的代码优化。奖励函数结合了规则奖励和基于规则的奖励，其中规则奖励来自真实执行结果，基于规则的奖励则根据代码是否符合CUDA编程规范进行评估。这种奖励机制能够引导Coder学习高效且正确的CUDA编程技巧。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

StitchCUDA在KernelBench基准测试中表现出色，在端到端GPU编程任务上实现了近100%的成功率，并且相比于多智能体基线，性能提升了1.72倍，相比于RL模型基线，性能提升了2.73倍。这些结果表明StitchCUDA在端到端GPU程序自动生成方面具有显著优势。

🎯 应用场景

StitchCUDA具有广泛的应用前景，可以应用于各种需要高性能GPU计算的领域，例如深度学习、科学计算、图像处理等。它可以帮助开发者快速生成高效的GPU程序，提高开发效率，并降低开发成本。未来，该框架可以进一步扩展到支持更多的GPU架构和编程语言。

📄 摘要（原文）

Modern machine learning (ML) workloads increasingly rely on GPUs, yet achieving high end-to-end performance remains challenging due to dependencies on both GPU kernel efficiency and host-side settings. Although LLM-based methods show promise on automated GPU kernel generation, prior works mainly focus on single-kernel optimization and do not extend to end-to-end programs, hindering practical deployment. To address the challenge, in this work, we propose StitchCUDA, a multi-agent framework for end-to-end GPU program generation, with three specialized agents: a Planner to orchestrate whole system design, a Coder dedicated to implementing it step-by-step, and a Verifier for correctness check and performance profiling using Nsys/NCU. To fundamentally improve the Coder's ability in end-to-end GPU programming, StitchCUDA integrates rubric-based agentic reinforcement learning over two atomic skills, task-to-code generation and feedback-driven code optimization, with combined rubric reward and rule-based reward from real executions. Therefore, the Coder learns how to implement advanced CUDA programming techniques (e.g., custom kernel fusion, cublas epilogue), and we also effectively prevent Coder's reward hacking (e.g., just copy PyTorch code or hardcoding output) during benchmarking. Experiments on KernelBench show that StitchCUDA achieves nearly 100% success rate on end-to-end GPU programming tasks, with 1.72x better speedup over the multi-agent baseline and 2.73x than the RL model baselines.