StitchCUDA: An Automated Multi-Agents End-to-End GPU Programing Framework with Rubric-based Agentic Reinforcement Learning
summary
The gist
StitchCUDA proposes a multi-agent framework integrated with rubric-based agentic reinforcement learning to automate end-to-end GPU program generation, addressing limitations in prior work that
In short
StitchCUDA automates creating entire GPU programs by using three specialized agents—Planner, Coder, and Verifier—in an iterative loop to generate CUDA code. It uses rubric-based reinforcement learning to train the Coder to handle complex system requirements and optimize performance across multiple kernels simultaneously. This framework achieves near 100% success on difficult programming tasks.
Key concepts
- Multi-Agent Framework
- This involves coordinating three distinct AI agents—a Planner, a Coder, and a Verifier—that work together in an iterative loop to solve a complex problem. They share information about the generated code and profiling results to ensure the final GPU program is correct and performs well across all parts.
- Rubric-Based Agentic Reinforcement Learning
- This training method uses structured rules (rubrics) alongside standard rewards to guide an AI agent, specifically the Coder. The rubric evaluates performance across four dimensions: anti-hacking, CUDA engineering quality, operator coverage, and skill compliance. This prevents the agent from exploiting weaknesses in the reward system.
- Atomic Skills Decomposition
- Instead of training one large model for all tasks at once, this method splits the complex learning process into two simpler skills: 'from-scratch generation' and 'feedback-driven optimization.' This decomposition makes training more efficient by collecting specific types of data for each skill separately, allowing the model to learn advanced techniques like kernel fusion effectively.
Terminology used across episodes
This episode discusses
- StitchCUDA: An Automated Multi-Agents End-to-End GPU Programing Framework with Rubric-based Agentic Reinforcement Learning · Paper Radio
- Kevin: Multi-Turn RL for Generating CUDA Kernels
- cuPilot: A Strategy-Coordinated Multi-agent Framework for CUDA Kernel Evolution
- An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
- Qwen2.5-Coder Technical Report
- Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning
- CUDABeaver: Benchmarking LLM-Based Automated CUDA Debugging
- CUDAHercules: Benchmarking Hardware-Aware Expert-level CUDA Optimization for LLMs · Paper Radio
- CUDA-L1: Improving CUDA Optimization via Contrastive Reinforcement Learning
- OpenRubrics: Towards Scalable Synthetic Rubric Generation for Reward Modeling and LLM Alignment
- CUDA-L2: Surpassing cuBLAS Performance for Matrix Multiplication through Reinforcement Learning
- Qwen3 Technical Report
- RAGEN: Understanding Self-Evolution in LLM Agents via Multi-Turn Reinforcement Learning
- CudaForge: An Agent Framework with Hardware Feedback for CUDA Kernel Optimization
- CUDABench: Benchmarking LLMs for Text-to-CUDA Generation
- QiMeng-Kernel: Macro-Thinking Micro-Coding Paradigm for LLM-Based High-Performance GPU Kernel Generation
The paper
StitchCUDA: An Automated Multi-Agents End-to-End GPU Programing Framework with Rubric-based Agentic Reinforcement Learning · Read on arXiv
University of Minnesota-Twin Cities
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.
Transcript
Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.
Tom: Today's paper: "StitchCUDA: An Automated Multi-Agents End-to-End GPU Programing Framework with Rubric-based Agentic Reinforcement Learning".
Jane: StitchCUDA proposes a multi-agent framework integrated with rubric-based agentic reinforcement learning to automate end-to-end GPU program generation, addressing limitations in prior work that focused only on single kernels.
Tom: First, who's behind it and why it matters.
Title and authors: Tom: So, to recap, we’re talking about "StitchCUDA: An Automated Multi-Agents End-to-End GPU Programing Framework with Rubric-based Agentic Reinforcement Learning." The authors are Li, Zhang, Chen, Luo, and Hong. They’re proposing this multi-agent setup to solve the issue where existing methods only focus on optimizing individual kernels instead of the whole program.
Jane: Right. And their core idea is to use this multi-agent system—a Planner for design, a Coder for writing code, and a Verifier for checking correctness and performance—to build complete GPU programs automatically. It’s about orchestrating the whole thing rather than just tackling one piece at a time.
Lu: I found it interesting how they structured the agents to handle system-level design alongside kernel implementation, which suggests they are thinking about data movement and host settings in a holistic way.
Meng: That holistic view is what concerns me practically; if the Planner designs something that’s fundamentally flawed from a data flow perspective, does the Coder just dutifully implement it, or can it catch those big system bottlenecks?
Lalam: The idea of an end-to-end framework means we aren't just getting a snippet of code; we are getting a fully functional piece of software that runs on the GPU. That level of automation is something I think could really streamline our entire development pipeline for AI applications.
The paper's summary: Tom: Now, let’s look at the summary of what StitchCUDA actually does. Basically, it takes a reference PyTorch code and uses this iterative "plan–code–profile–refine" loop to create a working GPU program that performs well.
Jane: The core mechanism they use is this loop where the agents talk to each other in stages: the Planner figures out what needs to be done, the Coder writes the code for a part, and then the Verifier profiles it and sends feedback back so the Coder can improve.
Lu: What really caught my eye was how they tackled making this agentic process better by using rubric-based reinforcement learning over two specific skills: generating from scratch and then optimizing based on feedback.
Meng: That part about decomposing the RL into atomic skills sounds clever for reducing the computational load, but I'm still wondering if that decomposition actually helps the Coder learn more complex optimization strategies than a standard prompt might achieve.
Lalam: I think the rubric-based reinforcement learning is where it gets really smart because they designed rewards to stop it from just guessing or doing something technically correct but useless, which is what we call reward hacking.
The paper's improvements: Tom: Let’s talk about what they claim StitchCUDA actually improved over prior work. They show that this framework achieves a nearly one hundred percent success rate on KernelBench Level three tasks, which is a big jump from earlier attempts like CUDAForge that struggled there <ref:2603.02637#pg2>.
Jane: The authors highlight that StitchCUDA doesn't just write code; it manages cross-kernel optimization and host-side orchestration, which was a major gap in previous single-kernel focused tools. This suggests it can handle the complex system-level requirements of larger models.
Lu: They specifically point out that by using rubric rewards, they manage to steer the Coder toward advanced techniques like tiling or tensor cores rather than settling for simpler fixes when feedback is given.
Meng: So, if we look at the practical gains mentioned in their results, they claim a one point five times average speedup over PyTorch eager mode code and two point seven three times better performance compared to RL model baselines on H200 GPUs <ref:2603.02637#pg2>. That’s a significant number when you look at real-world execution speedups.
Lalam: I think the improvement in correctness, going from three out of ten success on Level three tasks to nine out of ten, shows that this structured approach actually builds more reliable and powerful code than just relying on trial and error in prompting <ref:2603.02637#pg1>.
Conclusion: Tom: So, to wrap up StitchCUDA: it’s a multi-agent system with rubric-based RL that successfully automates end-to-end GPU program generation, achieving nearly one hundred percent success on Level three benchmarks and showing substantial speedups <ref:2603.02637#pg2>.
Jane: Essentially, this paper shows how breaking down the complex task into distinct roles for planning, coding, and verifying allows an AI to generate complete GPU programs that perform much better than previous single-kernel approaches.
Lu: I think the real weight of this work is in demonstrating that integrating structured feedback through rubric rewards can actually guide the Coder toward implementing sophisticated CUDA engineering techniques when dealing with complex system-level tasks.
Meng: From a practical standpoint, if we see these kinds of performance gains on H200 GPUs, it means we could deploy much larger and more complex AI models without needing a massive team to manually optimize every single line of code.
Lalam: I'm really optimistic about this; if this framework can consistently deliver high-quality GPU programs with such reliability, it could significantly accelerate the pace at which we can bring truly advanced, performant AI applications into production environments.
More episodes
- 2610.15391-DC-Flow: Target-Conditioned Cyclic Peptide Design via Delayed-Coupling Flow Matching
- 2610.15379-LiveAnon: Improving Privacy and Utility in Streaming Speaker Anonymization with Codec Language Models
- 2610.15097-Anchors Enhance Controllability and Reasoning of Diffusion Language Models
- 2610.15398-GA-SVD: Generation-aware Singular Value Decomposition for LLM Compression
- 2610.15106-RELAGRAM: Recursive Latent Graph Model for the Capacitated Vehicle Routing Problem
- 2610.15159-Beyond Adjacent Layers: Graph-Guided Layer Fusion for Compressing Large Language Models
- 2610.15185-DeltaNet-PC: Test-Time Memory Without Backpropagation Through Time
- 2610.15246-ObsSpec: Accelerating Agentic Reinforcement Learning with Observation Speculation
- 2610.15388-Approximation in Contrastive Representation Learning
- 2610.15427-Riemannian Graph Tokenizer for Structural Knowledge Transfer