CUDA Agent Trains LLMs to Write GPU Kernels Faster Than torch.compile

Introduction

ByteDance Seed and Tsinghua University’s Institute for AI Industry Research (AIR) have introduced CUDA Agent, a large-scale agentic reinforcement learning system designed to teach language models how to generate high-performance CUDA kernels.

The work targets a persistent weakness in AI-generated low-level code. Frontier language models can often produce CUDA that compiles and runs correctly, but correctness alone is not enough for production GPU workloads. A generated kernel must also compete with highly optimized compiler-generated implementations.

CUDA Agent is built around that second problem: not simply generating valid CUDA code, but learning how to iteratively profile, verify, and optimize kernels until they deliver meaningful runtime improvements.

Why Correct CUDA Code Was Not Enough

Before agentic reinforcement learning, the Seed1.6 base model already showed reasonably strong CUDA coding ability on KernelBench.

Across the 250 Level 1–3 KernelBench tasks used by the researchers, Seed1.6 achieved a 74.0% pass rate. In other words, it could generate functionally correct solutions for a substantial majority of the benchmark.

Performance was a different story.

Only 27.2% of the base model’s generated kernels were faster than torch.compile, and its geometric-mean speed relative to the compiler baseline was only 0.69×.

Model	Pass Rate	Faster Than `torch.compile`	Geomean Speed vs. `torch.compile`
Seed1.6 base model	74.0%	27.2%	0.69×
CUDA Agent	98.8%	96.8%	2.11×

This gap captures the main challenge of automatic GPU kernel generation.

A language model can understand CUDA syntax and still produce inefficient memory access, excessive kernel launches, poor tiling choices, unnecessary intermediate operations, or other implementations that lose to a compiler.

For latency-sensitive workloads, a correct but slower kernel is not enough.

CUDA Agent Learns Inside a Real CUDA Development Loop

To address that problem, the researchers place the model inside a realistic CUDA development environment rather than asking it to produce a final kernel in one shot.

The agent can repeatedly inspect performance, modify code, verify correctness, and profile the resulting implementation.

The project’s SKILL.md describes the optimization loop in roughly the same order a CUDA engineer would use:

Profile the native PyTorch implementation to identify bottlenecks and optimization opportunities.
Rewrite the performance-critical operators as custom CUDA extensions.
Compile and test the implementation for correctness.
Measure runtime performance against the baseline.
Continue iterating when further optimization is possible.

The agent is therefore trained to behave less like a one-shot code generator and more like an engineer working through an optimization cycle.

A Skill-Augmented Environment

CUDA Agent gives the model CUDA-specific instructions and tools rather than relying only on general coding ability.

The environment includes:

Kernel compilation
Correctness verification
Performance profiling
CUDA-specific optimization guidance
A restricted sandbox
Reward signals tied to actual execution results

The SKILL.md file also restricts shortcuts that could make benchmark results misleading. For example, the agent cannot simply fall back to arbitrary PyTorch operations inside the custom kernel implementation.

This matters because reinforcement learning systems can otherwise discover ways to maximize a reward without solving the intended optimization problem.

PPO Training Scales the Agent to Long Optimization Trajectories

The researchers train CUDA Agent with Proximal Policy Optimization (PPO).

A major challenge is that GPU optimization is naturally a long-horizon task. The model may need many cycles of code editing, compilation, debugging, profiling, and further optimization before reaching a strong result.

The paper therefore uses different context lengths and training stages to stabilize learning.

The base model is Seed1.6, a Mixture-of-Experts model with 230 billion total parameters and 23 billion active parameters.

For agentic reinforcement learning:

The context window is 131,072 tokens.
Training rollouts allow up to 150 agent turns.
Evaluation allows up to 200 agent turns.
The model is trained for 150 RL training steps.
Actor and critic models use a multi-stage warm-up strategy before full long-horizon optimization.

This is more involved than simply fine-tuning a model on pairs of CUDA prompts and answers.

The aim is to teach the model how to improve its own kernel through repeated interaction with execution feedback.

The Sandbox Separates Compilation From GPU Profiling

Reliable performance measurement is essential because runtime speed is part of the reward.

The researchers therefore built a CPU–GPU resource-decoupled sandbox architecture.

A Docker-based terminal sandbox handles CPU-oriented tasks such as compilation, while verification and performance profiling are dispatched to a dedicated GPU sandbox pool containing 128 NVIDIA H20 GPUs.

This separation serves two purposes.

First, it keeps compilation and agent interaction isolated from GPU benchmarking. Second, exclusive GPU allocation reduces interference between concurrent workloads, which makes latency measurements more stable.

That is important for reinforcement learning: if timing measurements are noisy, the model receives unreliable reward signals and may learn the wrong optimization behavior.

CUDA Agent Reaches 98.8% Pass Rate and Beats the Compiler on 96.8% of Tasks

The final system was evaluated on 250 tasks from KernelBench Levels 1 through 3.

CUDA Agent achieved:

98.8% overall pass rate
98.4% faster rate versus PyTorch eager execution
96.8% faster rate versus torch.compile
2.60× geometric-mean speedup versus eager execution
2.11× geometric-mean speedup versus torch.compile

Those results represent a substantial improvement over the Seed1.6 starting point.

The model did not merely become better at producing code that passed correctness tests. Its generated kernels became much more likely to outperform the compiler baseline.

Results by KernelBench Difficulty

The gains remained strong across all three KernelBench levels.

KernelBench Split	Pass Rate	Faster Than `torch.compile`	Geomean Speedup vs. `torch.compile`
Level 1	100.0%	97.0%	1.87×
Level 2	100.0%	100.0%	2.80×
Level 3	94.0%	90.0%	1.52×

Free · Use We0 AI

Build a showcase site and grow leads in minutes

Describe your idea once, and We0 AI can generate a showcase site, pages, and CMS, then help you attract customers and traffic after launch.

One complete project generation for free registration

Best for trying one complete generation flow and seeing a first project draft quickly.

Untitled-1

Layout

Fill

Radius

Generating

Level 2 was particularly strong.

These tasks involve operator sequences, where fusion and memory movement create opportunities that static compiler heuristics may not always exploit.

The authors argue that an iterative learned optimizer can search a broader space of implementation strategies, including hardware-specific memory access, tiling, and fusion choices.

The Training Data Is Being Released

The full trained model weights are not currently part of the public release described by the project.

However, the researchers have released several important pieces of the training stack.

CUDA-Agent-Ops-6K

The CUDA-Agent-Ops-6K dataset contains 6,000 synthesized operator-level training tasks.

Its construction pipeline includes:

Collecting seed operators from libraries such as torch and transformers.
Using an LLM to combine multiple operators into more complex fused tasks.
Filtering generated tasks through execution-based checks.

The filtering stage removes cases that are stochastic, trivial, invalid, or too similar to evaluation tasks.

The dataset is available on Hugging Face and is released under a CC BY 4.0 license.

SKILL.md and Agent Environment

The GitHub repository also includes the CUDA-specific SKILL.md instructions and the agent working environment.

These materials document the optimization workflow, available tools, restrictions, and execution setup used to guide the agent.

Reward and Training Recipes

The paper and repository describe the reward design, warm-up stages, critic initialization, and PPO-based training recipe used to stabilize long-horizon optimization.

This is particularly useful for researchers interested in training agents for systems programming rather than only reproducing the final benchmark score.

Why This Matters Beyond KernelBench

CUDA kernel optimization sits underneath a large part of modern AI infrastructure.

Faster kernels can improve:

Model training throughput
LLM inference latency
GPU utilization
Serving cost
Real-time AI pipelines
High-performance numerical workloads

The original AIBase article points to possible applications in AI infrastructure, large-model inference services, autonomous driving, and quantitative trading—areas where small latency differences can matter.

The CUDA Agent result does not mean traditional compilers are obsolete.

torch.compile remains a strong automatic baseline, and the CUDA Agent evaluation itself depends on a controlled benchmark and specific GPU environment.

What the research shows is narrower but still important: with enough execution feedback and specialized reinforcement learning, a language model can learn optimization strategies that outperform a static compiler baseline on a large majority of the tested GPU-kernel tasks.

That turns low-level performance engineering into a plausible domain for agentic learning rather than only one-shot code generation.

常见问题

What is CUDA Agent?

CUDA Agent is a large-scale agentic reinforcement learning system developed by ByteDance Seed, Tsinghua AIR, and their joint SIA-Lab. It trains a language model to iteratively write, verify, profile, and optimize CUDA kernels.

Which model is CUDA Agent based on?

The research uses Seed1.6 as the base model. The paper describes it as a Mixture-of-Experts model with 230B total parameters and 23B active parameters.

What is KernelBench?

KernelBench is an open benchmark for evaluating whether language models can generate correct and efficient GPU kernels for PyTorch programs. CUDA Agent is evaluated on 250 tasks covering KernelBench Levels 1 through 3.

How much faster is CUDA Agent than `torch.compile`?

Across the full benchmark, CUDA Agent reports a 2.11× geometric-mean speedup relative to torch.compile. Its generated kernels are faster than the compiler baseline on 96.8% of the evaluated tasks.

Does CUDA Agent use reinforcement learning?

Yes. The system uses PPO together with multi-stage warm-up, value-model pretraining, robust reward design, and long multi-turn agent trajectories.

How long can one CUDA Agent optimization trajectory run?

Training rollouts allow up to 150 agent turns, while evaluation allows up to 200 turns. The agentic training context window is 131,072 tokens.

Is CUDA Agent open source?

The project has released its GitHub repository, agent environment, SKILL.md, training recipes, and the CUDA-Agent-Ops-6K dataset. The complete trained model weights are not listed as part of the current public release.

What is CUDA-Agent-Ops-6K?

CUDA-Agent-Ops-6K is a dataset of 6,000 synthesized operator-level CUDA training tasks. It is designed for large-scale agentic RL and includes filtering intended to keep tasks executable, deterministic, non-trivial, and separated from the KernelBench evaluation set.

Summary

CUDA Agent addresses a specific weakness in LLM-generated systems code: producing CUDA that is not only correct, but actually faster than compiler-generated alternatives.

Starting from Seed1.6, the researchers use a CUDA-specific agent environment, execution feedback, robust reward design, and long-horizon PPO training to raise the pass rate from 74.0% to 98.8% and the faster-than-torch.compile rate from 27.2% to 96.8%.

The project also releases a 6,000-task training dataset, SKILL.md, agent-environment materials, and training recipes, giving researchers a concrete foundation for further work on learned GPU optimization.

CUDA Agent’s main contribution is showing that performance engineering itself can be learned through an iterative agent loop—not just approximated with one-shot code generation.

CUDA Agent Trains LLMs to Write GPU Kernels Faster Than torch.compile

Introduction

Why Correct CUDA Code Was Not Enough