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

summary

Video file (mp4)

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

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

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

← Home