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

TL;DR

CUDA Agent is a Seed1.6-based agentic RL system from ByteDance Seed + Tsinghua AIR for generating high-performance CUDA kernels. It pairs a synthesized 6K-task operator dataset (CUDA-Agent-Ops-6K) with an OpenHands-style ReAct agent loop, a sandboxed GPU verifier pool (128 H20s) with anti-reward-hacking guards, and a multi-stage PPO training recipe (single-turn RL warm-up → RFT actor init → Value Pretraining critic init → 150-step PPO). On KernelBench it lifts the 23B-active / 230B-total Seed1.6 base from ~27% to 96.8% faster rate over torch.compile across Level-1/2/3, outperforming Claude Opus 4.5 and Gemini 3 Pro by roughly 40% on the hardest split — and beats torch.compile itself in the vast majority of operator tasks.

Key claims

CUDA Agent attains a 98.8% overall pass rate and a 96.8% faster-than-torch.compile rate on KernelBench, vs 66.4% / 69.6% for Claude Opus 4.5 / Gemini 3 Pro under the same agent loop [Table 1].
Per-level faster-rate over torch.compile: 97% (L1), 100% (L2), 90% (L3); per-level geomean speedup over torch.compile: 1.87× / 2.80× / 1.52× [Table 1].
LLM-generated CUDA kernels can systematically beat static compiler heuristics in operator-fusion settings (L2: 100% faster rate, 2.80× geomean over compile), suggesting the iterative agent loop discovers tiling and memory-access patterns inaccessible to rule-based fusion [§4.2].
The agent loop ablation removes 51%+ of the gains: w/o agent loop = 14.1% faster rate vs 96.8% for the full system, isolating multi-turn compilation+profiler feedback as the load-bearing component [Table 2].
Robust milestone-based reward (discrete credit for hitting 1.0× / 1.05× / 1.1× / … of torch.compile) outperforms the raw speedup reward used in prior work (Kevin / CUDA-L1): w/o robust reward = 60.4% faster rate vs 96.8% full [Table 2, §3.2].
Removing Rejection-Fine-Tuning actor initialization (w/o RFT) causes catastrophic policy collapse with a sharp entropy spike around the same training step [§4.3.3, Fig. 4]; removing Value Pretraining causes trajectory-length explosion as the uninitialized critic fails to penalize redundant search [§4.3.3, Fig. 5].
Training instability traced to a precision-mismatch failure mode: CUDA-coding tokens make up <0.01% of pretraining, so the base policy emits many low-probability tokens whose importance-sampling ratios explode under BF16/FP16 mismatch between training and inference engines [§3.3].
CUDA-Agent-Ops-6K is built by mining torch/transformers operator classes as seeds, prompting an LLM to compose ≤5 operators into fused tasks, then filtering for (a) eager+compile correctness, (b) determinism, (c) anti-trivial outputs, (d) eager runtime in [1ms, 100ms], and (e) low similarity to KernelBench test cases for contamination control [§3.1].
Reward-hacking countermeasures are codified: file-permission–protected verification scripts, banning torch.nn.functional fallbacks via execution-time context managers, output checks against 5 random inputs, careful warm-up/synchronization in profiling, and no web-search tool in the agent’s toolset [§3.2].
Scales to a 128K-token context and up to 200 agent turns per task at evaluation (150 during training); base model is Seed1.6, MoE with 23B active / 230B total parameters [§4.1].
ChatGPT-5/5.1/5.2 declined to respond to CUDA-related prompts and were therefore not evaluated — a footnote, but a notable data point on safety-training side-effects in code-generation evaluations [§1, footnote 1].

Method

The system is three components glued by PPO. Data: a scalable synthesis pipeline mines real PyTorch operator classes from torch and transformers, prompts an LLM to combinatorially compose up to 5 of them into a single fused layer, and filters the results through execution-feedback gates (correctness, determinism, runtime band, anti-trivial-output, KernelBench-similarity exclusion) to produce 6,000 operator-level training problems (CUDA-Agent-Ops-6K, released on HF). Agent loop: an OpenHands-aligned ReAct loop with BashTool/GlobTool/MultiEditTool/TodoWriteTool plus a CUDA-specific profiler skill, governed by a SKILL.md that codifies the optimize loop (profile native → write model_new.py + CUDA kernels + bindings → compile + verify + iterate until ≥5% over torch.compile). Rewards are milestone-based — discrete credit for hitting performance bands relative to eager / compile baselines — rather than raw speedup, to avoid outliers and bias toward easy kernels. Anti- reward-hacking is enforced at the system level: protected verifier scripts, banned fallbacks, repeated profiling with synchronization, no web tool. RL: PPO on Seed1.6, but only after a multi-stage warm-up: (1) single-turn RL on the base model to lift CUDA coding, (2) RFT on agent trajectories sampled from the warm-up model, filtered by positive reward and pattern-validity, to initialize the actor, (3) Value Pretraining of the critic via GAE targets on the same trajectories to initialize it before PPO. The sandbox uses 128 H20 GPUs with process-level isolation between CPU (compilation) and GPU (verification + profiling) workloads.

Results

Main result on KernelBench, 250 operator tasks (Levels 1–3), agent loop fixed across baselines [Table 1]:

Model	Overall pass rate	Overall faster-rate vs compile	L3 faster-rate vs compile
Seed1.6 (base)	74.0%	27.2%	2.0%
GLM 4.6	75.6%	19.2%	10.0%
Kimi K2	66.8%	22.8%	6.0%
Gemini 3 Pro	91.2%	69.6%	52.0%
Claude Opus 4.5	95.2%	66.4%	50.0%
CUDA Agent	98.8%	96.8%	90.0%

L2 (operator sequences) is the cleanest signal for “learned > static compiler heuristics”: 100% faster rate, 2.80× geomean over compile vs 1.60× for Claude Opus 4.5. Ablations confirm each component is load-bearing [Table 2]: full = 96.8% faster rate; w/o agent loop = 14.1%; w/o robust reward = 60.4%; w/o RFT = 49.8%; w/o Value Pretraining = 50.9%. Beyond the headline, RFT specifically prevents policy collapse (Fig. 4), and Value Pretraining specifically prevents trajectory-length explosion (Fig. 5) — distinct failure modes with distinct fixes.

Why it’s interesting

Three angles. (1) CUDA Agent is the first filed agentic-SWE entry in a kernel-optimization domain rather than general repo-level SWE, which makes it a useful test of the Agentic Software Engineering recipe-set: the same shape (large-scale verifiable-task synthesis + sandboxed execution + multi-format tool loop + multi-stage post- training) survives the domain shift, but the reward design has to change — raw speedup is a bad reward, milestones are needed, and contamination control against KernelBench is a non-trivial gate that the SWE-Bench-style work doesn’t face. (2) The RFT-actor + Value-Pretrained-critic recipe for stabilizing PPO on a domain that is <0.01% of pretraining is a directly transferable result for any Reasoning RL setup where the target distribution is far from the base — sharper than the no-KL prescriptions in Human detectors are surprisingly powerful reward models and complementary to the experience-library framing of Training-Free Group Relative Policy Optimization (whose context- space variant cannot inject capabilities the base model lacks — which is exactly CUDA Agent’s regime). (3) The fused-operator synthesis pipeline is a new datapoint for Synthetic Training Data: unlike InfTool (Close the Loop: Synthesizing Infinite Tool-Use Data via Multi-Agent Role-Playing), which synthesizes trajectories from APIs, CUDA Agent synthesizes problem distributions over real APIs by compositional sampling, with the verification gate doing both correctness and contamination duty.