arXiv 2603.026372026-08-11cs.MAcs.CLcs.PL

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

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

Department of Computer Science\&Engineering, University of Minnesota-Twin Cities, Minnesota, USA；Department of Electrical Engineering, University of Minnesota-Twin Cities, Minnesota, USA

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

发表信息：: Proceedings of the 43rd International Conference on Machine Learning (ICML 2026), PMLR 306:72290-72314

阅读 PDF

arXiv 页面 ↗HTML 原文 ↗

AI总结：

StitchCUDA通过基于规则的多智能体强化学习，实现了端到端GPU编程的高效生成与验证，显著提升了性能和准确性。

摘要：

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

英文摘要：

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. Code of the STITCHCUDA framework is avalaible at https://github.com/UMN-APEX-Lab/StitchCUDA.

机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。