ByteDance Seed and Tsinghua AIR present CUDA Agent, an agentic reinforcement‑learning framework that fine‑tunes a large language model to produce CUDA kernels that are not only correct but also faster than the output of torch.compile. The system starts from Seed1.6, a proprietary mixture‑of‑experts model with 23 B active and 230 B total parameters, and places the model in a permission‑locked sandbox equipped with a CUDA compiler, profiler and correctness checker. Training proceeds with proximal policy optimization for 150 steps, providing the model with a 131 072‑token context that lets it inspect and edit source files, run build commands and view profiling results in a ReAct‑style loop.
On the KernelBench benchmark (250 tasks) the base Seed1.6 model solves 74.0 % of tasks but only outperforms torch.compile on 27.2 % of them, yielding a 0.69× geometric‑mean speedup. After CUDA Agent training the pass rate rises to 98.8 % and the fraction of kernels beating torch.compile reaches 96.8 %, giving a 2.11× geometric‑mean speedup over the compiler—about 40 points ahead of Claude Opus 4.5 and Gemini 3 Pro on the hardest Level‑3 split. The accompanying CUDA‑Agent‑Ops‑6K dataset contains 6 000 sampled operator pairs, 83.77 % of which are two‑operator compositions, and the SKILL.md specification, reward shaping recipes and anti‑reward‑hacking measures are released publicly, while the trained agent weights remain private.
- Context window: 131 072 tokens
- PPO training steps: 150
- Sandbox GPU resources: 128 NVIDIA H20 GPUs
- Reward discretization: r ∈ {−1, 1, 2, 3}
- KernelBench results: 98.8 % pass, 96.8 % faster‑than‑torch.compile, 2.11× geomean speedup
- Licensing: dataset, SKILL.md, reward recipes public; agent weights not released
Why this matters
The work shows that aligning a language model with a real‑world compilation and profiling loop can close the persistent correctness‑performance gap in GPU code generation. By rewarding kernels that exceed torch.compile by a modest 5 % margin, the system shifts the model’s output from merely functional to performance‑competitive, which could lower inference latency and operational cost in latency‑critical services. However, because the trained agent is not released, reproducibility is limited to organizations that can replicate the 128‑GPU sandbox and PPO pipeline, restricting immediate adoption to well‑resourced labs. The public dataset and skill specification nevertheless enable smaller teams to experiment with similar reward‑shaping strategies on open base models.