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

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

cs.MAMar 3, 2026

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

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

Organizations: University of Minnesota-Twin Cities, Minnesota, USA.

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

Figures & tables

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.

		Level 1 (20 Tasks)			Level 2 (20 Tasks)			Level 3 (10 Tasks)
Hardware	Method	Correctness	Avg Speedup	$\rm{Fast}_{1}$	Correctness	Avg Speedup	$\rm{Fast}_{1}$	Correctness	Avg Speedup	$\rm{Fast}_{1}$
	GPT5.2	10/20	0.73x	25%	8/20	0.66x	30%	2/10	0.47x	10%
	Claude-4-sonnet	12/20	1.02x	25%	8/20	0.76x	30%	3/10	0.67x	20%
	Qwen3-32B	2/20	0.02x	5%	0	0	0%	0	0	0%
	Kevin32B	9/20	0.65x	10%	11/20	0.49x	30%	4/10	0.22x	0%
	CUDAForge	18/20	2.74x	50%	18/20	1.28x	80%	6/10	0.75x	30%

Table 1: The success rate and E2E mean speedup relative to PyTorch Eager of various methods on our testing set. The best results are highlighted in bold , the second best is underline . StitchCUDA-Q is our multi-agent framework, with the original Qwen3-32B as the Coder. Similarly, StitchCUDA-K is a backend with Kevin32B, StitchCUDA-G is a backend with GPT-5.2 as Coder, and StitchCUDA-C is a backend with Claude-4.6-Sonnet.

Variant	Corre. (/10)	Speedup	Hacking (/50)
StitchCUDA	9.0	1.50×	7.7
w/o Anti-Hacking	5.7	0.75×	15.7
w/o CUDA Engineering	8.0	1.18×	9.3
w/o Operator Coverage	8.0	1.08×	10.0
w/o Skill Compliance	7.3	1.33×	8.7
StitchCUDA-A	5.0	0.46×	16.3

Table 2: Ablation results of training rubrics

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.

	Level 1		Level 2		Level 3		Sum
Model	partial	total	partial	total	partial	total	partial	total
StitchCUDA-K	10	3	8	0	4	1	22/50	4/50
StitchCUDA-Q	4	0	4	0	2	0	10/50	0
StitchCUDA-A	8	0	5	0	3	0	16/50	0
StitchCUDA	3	0	3	0	2	0	8/50	0

Table 3: Hacking detection with prompts hint.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Illustration of online documents processed into the database for RAG

Topic	URL
CUDA-12.9	https://docs.nvidia.com/cuda/archive/12.9.0/
cuBLAS-12.9	https://docs.nvidia.com/cuda/archive/12.9.0/cublas/index.html
cuSPARSE-12.9	https://docs.nvidia.com/cuda/archive/12.9.0/cusparse/index.html
CUTLASS-4.3.5	https://docs.nvidia.com/cutlass/latest/
cuSPARSELt-0.8.1	https://docs.nvidia.com/cuda/cusparselt/index.html
Nvidia Blackwell Architecture Whitepaper	https://images.nvidia.com/aem-dam/Solutions/geforce/blackwell/nvidia-rtx-blackwell-gpu-architecture.pdf

Appendix

Table 4: Example official documentation topics and links

Method	Turns per trajectory	Total turns	H200-Hour
Multi-turn agentic RL	$15$	$144000$	$9600-12000$
Our method (w/o rubric reward)	$1$	-	$128$
Rubric-based agentic RL (ours)	$1$	-	$160$

Appendix

Table 5: Computation overhead comparison measured in H200-Hour under the same dataset and training hyperparameters (400 samples, group=8, epoch=3). We estimate the average environment interaction takes 4–5 minutes.

Program Synthesis

Program synthesis automatically constructs executable programs from examples, specifications, demonstrations, or natural-language descriptions.

See all 179 papers

CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation
Feb 27, 2026Weinan Dai, Hanlin Wu, Qiying Yu +13GPU Kernel Generation GPU Kernel Optimization
ByteDance Seed · Institute for AI Industry Research (AIR), Tsinghua University · SIA-Lab of Tsinghua AIR and ByteDance Seed
CuTeGen: An LLM-Based Agentic Framework for Generation and Optimization of High-Performance GPU Kernels using CuTe
Apr 1, 2026Tara Saba, Zhiyang Chen, Jikai Jason Li +3GPU Kernel Generation GPU Kernel Optimization
Department of Computer Science, University of Toronto · Standard Kernel
Optimizing CUDA like a Human: Micro-Profiling Tools as Expert Surrogates for LLM-Based GPU Kernel Optimization
Jun 24, 2026Jiading Gai, Shuai Zhang, Kaj Bostrom +6GPU Kernel Optimization High-Performance Computing
Amazon · Siemens · University of Minnesota

Feb 27, 2026J/K move · Enter open · S save

CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation

Weinan Dai, Hanlin Wu, Qiying Yu, Huan-ang Gao+12

ByteDance Seed · Institute for AI Industry Research (AIR), Tsinghua University · SIA-Lab of Tsinghua AIR and ByteDance Seed

GPU Kernel Generation GPU Kernel Optimization RL for Code Generation Agentic RL

Open paper AskPDF arXiv

GPU kernel optimization is fundamental to modern deep learning but remains a highly specialized task requiring deep hardware expertise. Despite strong performance in general programming, large language models (LLMs) remain uncompetitive with compiler-based systems such as torch.compile for CUDA kernel generation. Existing CUDA code generation approaches either rely on training-free refinement or fine-tune models within fixed multi-turn execution-feedback loops, but both paradigms fail to fundamentally improve the model's intrinsic CUDA optimization ability, resulting in limited performance gains. We present CUDA Agent, a large-scale agentic reinforcement learning system that develops CUDA kernel expertise through three components: a scalable data synthesis pipeline, a skill-augmented CUDA development environment with automated verification and profiling to provide reliable reward signals, and reinforcement learning algorithmic techniques enabling stable training. CUDA Agent achieves state-of-the-art results on KernelBench, delivering 100%, 100%, and 92% faster rate over torch.compile on KernelBench Level-1, Level-2, and Level-3 splits, outperforming the strongest proprietary models such as Claude Opus 4.5 and Gemini 3 Pro by about 40% on the hardest Level-3 setting.

Abstract

Program Synthesis

Explore similar work

CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation

CuTeGen: An LLM-Based Agentic Framework for Generation and Optimization of High-Performance GPU Kernels using CuTe

Optimizing CUDA like a Human: Micro-Profiling Tools as Expert Surrogates for LLM-Based GPU Kernel Optimization