ByteDance Seed and Tsinghua AIR's CUDA Agent Generates Kernels That Beat torch.compile

The system targets a narrow but persistent gap: frontier models already produce correct CUDA code, but the code is often slow. The base model Seed1.6 passes 74.0% of KernelBench tasks but outruns torch.compile on only 27.2% of them, with a 0.69x geometric-mean speedup, meaning its kernels are on average slower than compiler-generated code. CUDA Agent closes this gap by placing 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.

Across the full benchmark, the trained agent achieved a 98.8% pass rate, was faster than eager execution on 98.4% of tasks and faster than torch.compile on 96.8%, with geometric-mean speedups of 2.60x and 2.11x respectively. Level 2, containing operator sequences, was the strongest split at 100% pass, 100% faster than compile and 2.80x speedup. On the hardest Level-3 split, the system reached 94.0% pass, 90.0% faster than compile and 1.52x speedup, roughly 40 points ahead of Claude Opus 4.5 and Gemini 3 Pro on the faster-than-compile rate. The paper contains an inconsistency: the abstract and introduction state 100%, 100% and 92% faster rates for Levels 1 through 3, while Table 1 reports 97.0%, 100.0% and 90.0%; Table 1 is described as the main results table.

To build the training data, the researchers crawled reference operators from the torch and transformers libraries. An LLM sampled up to five torch operator classes and stacked them into a single fused layer. Samples were kept only if they executed in both eager and compile modes, were deterministic, produced non-constant outputs and ran between 1 ms and 100 ms in eager mode; samples with AST similarity above 0.9 to any KernelBench task were removed. The resulting CUDA-Agent-Ops-6K dataset contains 6,000 samples, 83.77% of them two-operator compositions.

The agent loop mirrors OpenHands tooling, including Bash, Read/Write, Edit/MultiEdit and Grep, under a ReAct pattern. SKILL.md instructs 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 set tolerances. The reward is discrete, r ∈ {−1, 1, 2, 3}, with 3 given when the kernel clears both eager and torch.compile by more than 5%. Five countermeasures address reward hacking, including permission-locked verification 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.

Ablations show the agent loop is critical: removing it drops the faster-than-compile rate from 96.8% to 14.1%. A raw speedup reward gives 60.4%, skipping reference-based fine-tuning gives 49.8% plus reward collapse, and omitting value pretraining gives 50.9% plus runaway trajectories. Case studies illustrate the learned policies: a diagonal matmul rewritten as row-wise scaling reached 73.31x over torch.compile, a matmul-divide-sum-scale chain reordered and fused reached 24.04x, and a ResNet BasicBlock with BatchNorm folded into convolution reached 3.59x.

The trained agent is not released. It is built on Seed1.6, a proprietary mixture-of-experts model with 23B active parameters and 230B total parameters, and the paper ships no weights. Publicly released components include the CUDA-Agent-Ops-6K dataset, the SKILL.md specification and the reward and warm-up recipes. Full replication requires substantial compute; the profiling sandbox alone used 128 NVIDIA H20 GPUs. Mid-size teams can still adopt individual components, such as the dataset, milestone reward, anti-reward-hacking constraints and skill spec, on top of an open base model. Potential applications include AI infrastructure and inference serving, GPU cloud, autonomous driving, quantitative trading, medical imaging and recommendation systems, wherever fused kernels sit on a latency-critical path.