StitchCUDA用多智能体自动生成端到端的高效GPU程序,提升编译性能。
StitchCUDA: An Automated Multi-Agents End-to-End GPU Programing Framework with Rubric-based Agentic Reinforcement Learning
- 三智能体协作:规划、编码、验证,实现全流程自动化。
- 在KernelBench上成功率近100%,速度比基线快1.72倍。
- 通过评分机制强化学习,防止代码抄袭或硬编码漏洞。
现代机器学习工作负载日益依赖GPU,但实现高性能端到端程序仍具挑战,受限于内核效率与主机配置。尽管基于大模型的方法在自动内核生成方面展现潜力,但现有研究主要聚焦单内核优化,难以扩展至端到端程序,制约实际部署。为此,本文提出StitchCUDA,一个用于端到端GPU程序生成的多智能体框架,包含三个专用智能体:规划器负责系统设计协调,编码器逐步实现代码,验证器利用Nsys/NCU进行正确性检查与性能分析。为显著提升编码器在端到端编程中的能力,StitchCUDA引入基于评分的智能体强化学习,融合任务到代码生成与反馈驱动优化两种基础技能,采用组合评分奖励与真实执行规则奖励。由此,编码器学会实现高级CUDA技术(如自定义内核融合、cublas后处理),并有效避免基准测试中奖励劫持问题(如复制PyTorch代码或硬编码输出)。在KernelBench上的实验表明,StitchCUDA在端到端GPU编程任务中达到近100%成功率,较多智能体基线提速1.72倍,较强化学习基线提速2.73倍。StitchCUDA代码已开源:https://github.com/UMN-APEX-Lab/StitchCUDA。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。