ByteDance Seed team and the Institute for Intelligent Manufacturing at Tsinghua University (AIR) recently officially released a new system called CUDA Agent. This is an intelligent agent reinforcement learning system specifically designed for training large language models to write high-performance GPU kernels, marking a significant advancement in AI's low-level code optimization field.
In previous AI development, although cutting-edge large models could generate grammatically correct CUDA code, the generated code often had low efficiency and struggled to outperform traditional compilers in actual runtime speed. For example, using the KernelBench benchmark test, the basic model Seed1.6 achieved a task success rate of 74.0%, but the proportion of its runtime speed surpassing torch.compile was only 27.2%, with the average geometric speed actually lagging behind the compiler itself.
To break through this bottleneck, CUDA Agent places the large model in a real CUDA development environment, equipped with performance analysis tools, correctness verification mechanisms, and a secure sandbox with restricted permissions. The system uses the Proximal Policy Optimization (PPO) algorithm for deep training over 150 steps, ultimately achieving a total success rate of 98.8% on a benchmark test with 250 tasks, and 96.8% of the generated kernels ran faster than traditional torch.compile.
In terms of practical applications and open-source plans, although the complete model weights of the system are not yet publicly available, the research team has opened up to the public a CUDA-Agent-Ops-6K dataset containing 6,000 samples, related SKILL.md specification instructions, and reward and warm-up training recipes. This achievement is expected to be widely applied in industries such as AI infrastructure, large model inference services, autonomous driving, and quantitative trading, which are highly sensitive to latency.
Join Now