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

TL;DR

StitchCUDA is proposed, 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.

Abstract

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.

Figures (5)

• Figure 1: The performance of various automated GPU program generation methods on KernelBench Level 1/2/3. Achieved average speedup is relative to PyTorch eager mode code, using NVIDIA H200 GPU.
• Figure 2: StitchCUDA: ① The Planner profiles the reference implementation using Nsys to identify performance bottlenecks and generates a structured to-do list of optimization tasks. ② The Coder executes one subtask per iteration, ③ compiles the modified code with nvcc, and produces an executable for unit testing. ④ The Verifier selects the appropriate profiling tool (Nsys, NCU, or both) to analyze the target kernel. ⑤ If compilation/correctness fails, the profiling stage is skipped. ⑥ Relevant Nvidia's official documentation is retrieved via a RAG module, and ⑦ Construct a single task feedback to Coder to initiate the next coding-feedback iteration.
• Figure 3: Rubric-based Agentic RL process for Coder. We train the Coder on two atomic skills via GRPO. Besides the rule-based reward function, we introduce a rubric reward to avoid reward hacking and encourage more kernel optimizations.
• Figure 4: Comparison of success rate and average speedup between StitchCUDA and StitchCUDA-A across Level 1/2/3. StitchCUDA achieves higher success rates and speedups at all levels, with the rubric-based reward.
• Figure 5: Illustration of online documents processed into the database for RAG