ByteDance Seed and Tsinghua AIR introduced CUDA Agent, an agentic reinforcement learning system that trains a large language model to write GPU kernels that are both correct and faster than those generated by torch.compile. The core problem for AI infrastructure teams is that frontier LLMs can produce syntactically valid CUDA but the kernels often run slower than the compiler’s output, wasting compute and increasing latency in latency‑critical paths such as inference serving, autonomous driving, quantitative trading, medical imaging and recommendation systems.
CUDA Agent solves this by placing the model inside a real CUDA development environment that includes profiling, correctness checks and a permission‑locked sandbox. The agent iterates through a ReAct loop: it reads a SKILL.md specification that tells it to profile a PyTorch model, write a new kernel, compile it in a GPU sandbox and verify that the kernel is at least five percent faster than torch.compile on multiple random inputs. Training uses PPO for 150 steps with a 131 072‑token context, a discrete reward that penalizes correctness failures and rewards milestone speedups, and anti‑reward‑hacking measures such as locked verification scripts, disabled fallbacks and no web search.
The released artifacts— the CUDA‑Agent‑Ops‑6K dataset, the SKILL.md spec and the reward and warm‑up recipes—allow mid‑size teams to adopt the methodology on top of an open base model, while large labs can replicate the full sandbox using 128 H20 GPUs. Results on KernelBench show a jump from a base model’s 27.2 % faster‑than‑compile rate to 96.8 % for the trained agent, with a 2.11× geometric‑mean speedup over torch.compile and a 98.8 % pass rate. This closes the stubborn gap between correctness and performance, letting teams deploy fused kernels that cut cost per token and retune across GPU generations without manual tuning.
#AI #GPU #CUDA #LLM #PerformanceOptimization #Inference