ByteDance Seed and Tsinghua AIR Introduces CUDA Agent: A Large-Scale Agentic RL System for CUDA Kernel Generation

ByteDance Seed and Tsinghua AIR have released CUDA Agent, an agentic reinforcement learning system that trains a large language model to write GPU kernels that beat a compiler. The gap it targets is narrow but stubborn: frontier models already produce correct CUDA, they just produce slow CUDA. On KernelBench, the base model Seed1.6 passes 74.0% of tasks yet outruns torch.compile on only 27.2% of them, at a 0.69× geometric-mean speedup which means its kernels are, on average, slower than what the compiler generates on its own. CUDA Agent closes that gap by putting the model inside a real CUDA development environment with profiling, correctness checks and a permission-locked sandbox, then training it with PPO for 150 steps at a 131,072-token context. The result is a 98.8% pass rate and a 96.8% faster-than-torch.compile rate across the 250-task benchmark, at 2.11× geomean over compile — roughly 40 points ahead of Claude Opus 4.5 and Gemini 3 Pro on the hardest Level-3 split.

Is it deployable?

Partly, but the trained agent is not released. It is built on Seed1.6, a proprietary MoE model with 23B active and 230B total parameters, and the paper ships no weights. Public: the CUDA-Agent-Ops-6K dataset, the SKILL.md spec and the reward and warm-up recipes.

Which companies: The profiling sandbox alone used 128 NVIDIA H20 GPUs, which puts full replication inside frontier labs, GPU clouds and large infrastructure teams. Mid-size teams can still adopt the parts — dataset, milestone reward, anti-reward-hacking constraints, skill spec — on top of an open base model.

Industries and applications: AI infrastructure and inference serving, GPU cloud, autonomous driving, quantitative trading, medical imaging and recommendation systems — anywhere fused kernels sit on a latency-critical path. Uses include fusing operator sequences torch.compile handles poorly, cutting cost per token, and re-tuning kernels across GPU generations.

Data synthesis

The research team crawls reference operators from the torch and transformers libraries. An LLM then samples up to five torch operator classes and stacks them into one fused layer. A filter keeps only operators that execute in both eager and compile modes, are deterministic, produce non-constant outputs, and run between 1 ms and 100 ms in eager mode. Samples with AST similarity above 0.9 to any KernelBench task are removed. The result is CUDA-Agent-Ops-6K: 6,000 samples, 83.77% of them two-operator compositions.

Environment and reward

The agent loop mirrors OpenHands tooling — Bash, Read/Write, Edit/MultiEdit, Glob, Grep, NotebookEdit, BashOutput, KillBash — under a ReAct pattern. CUDA instructions ship in the Agent Skills format. SKILL.md tells the model to profile the PyTorch model, rewrite model_new.py with custom kernels, compile in a GPU sandbox, and iterate until the kernel is at least 5% faster than torch.compile at atol=1e-2, rtol=1e-2.

Reward hacking gets five countermeasures: permission-locked verification and profiling scripts, context managers that forbid torch.nn.functional fallbacks, checks against five random inputs, profiling with device synchronization and warm-up, and no web search tool.

The reward is discrete rather than a raw speedup ratio. r ∈ {−1, 1, 2, 3}: −1 on correctness failure, 3 if the kernel clears both eager and torch.compile by more than 5%, 2 if it clears eager only, 1 otherwise.

Results

Table 1, overall: 98.8% pass rate, 98.4% faster than eager, 96.8% faster than torch.compile, at 2.60× and 2.11× geomean respectively. Level 2 (operator sequences) is the strongest split: 100% pass, 100% faster rate, 2.80× over torch.compile. Level 3 lands at 94.0% pass, 90.0% faster rate and 1.52×, roughly 40 points above Claude Opus 4.5 (50.0%) and Gemini 3 Pro (52.0%) on faster rate versus compile.

One inconsistency: the abstract and introduction state 100% / 100% / 92% faster rates for Levels 1–3, while Table 1 reports 97.0% / 100.0% / 90.0%. Table 1 is the main results table.

Ablations are blunt. Removing the agent loop drops faster rate versus compile from 96.8% to 14.1%. A raw speedup reward gives 60.4%, no RFT gives 49.8% plus reward collapse, no value pretraining gives 50.9% plus runaway trajectories.

Case studies show what the policy learns. A diagonal matmul rewritten as row-wise scaling: 73.31× over torch.compile. A matmul-divide-sum-scale chain reordered and fused: 24.04×. A ResNet BasicBlock with BatchNorm folded into convolution andcudnnConvolutionBiasActivationForward: 3.59×.

Key Takeaways

• CUDA Agent hits 98.8% pass rate and 96.8% faster-than-torch.compile rate on KernelBench, at 2.11× geomean.
• Level 2 fusion is the standout: 100% faster rate and 2.80× over torch.compile.
• The discrete milestone reward beats a raw speedup ratio by 36.4 points on faster rate.
• RFT plus value pretraining is what turns a 17-step collapse into 150 stable steps.
• Weights are closed; the 6,000-sample dataset, SKILL.md and the recipe are public.

Check out the Paper, the Project Page and the Dataset. Also, feel free to follow us on Twitter and don’t forget to join our 150k+ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.

Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us