ByteDance and Tsinghua Introduce CUDA Agent for Kernels
ByteDance Seed and Tsinghua AIR introduce CUDA Agent, a large-scale agentic reinforcement learning system for CUDA kernel generation.
Stock photo for illustration only, not from the actual event
- ByteDance Seed and Tsinghua AIR launch CUDA Agent for automated kernel generation
- Built on Seed1.6 proprietary MoE model with 23B active and 230B total parameters
- Profiling sandbox utilizes 128 NVIDIA H20 GPUs achieving a 98.8% overall pass rate
- Releases CUDA-Agent-Ops-6K dataset, SKILL.md spec, and reward recipes to the public
ByteDance Seed and Tsinghua AIR have jointly introduced CUDA Agent, a large-scale agentic reinforcement learning system engineered specifically for CUDA kernel generation. While the trained agent itself is not publicly released—having been built upon Seed1.6, a proprietary Mixture-of-Experts model featuring 23 billion active and 230 billion total parameters without ship weights—the research team has made the CUDA-Agent-Ops-6K dataset, the SKILL.md specification, along with reward and warm-up recipes, fully accessible to the public.
Regarding the required infrastructure, the profiling sandbox alone utilizes 128 NVIDIA H20 GPUs, restricting full replication to frontier research laboratories, GPU clouds, and large infrastructure teams. Nevertheless, mid-sized teams can still adopt individual components such as the dataset, milestone rewards, anti-reward-hacking constraints, and skill specifications on top of an open base model.
The technology targets various industries and latency-critical applications, including AI infrastructure and inference serving, GPU clouds, autonomous driving, quantitative trading, medical imaging, and recommendation systems. Use cases involve fusing operator sequences that torch.compile handles poorly, reducing costs per token, and re-tuning kernels across different GPU generations.
Stock photo for illustration only, not from the actual event
The agent loop mirrors OpenHands tooling—incorporating Bash, Read/Write, Edit/MultiEdit, Glob, Grep, NotebookEdit, BashOutput, and KillBash—operating under a ReAct pattern. CUDA instructions are delivered within the Agent Skills format. Specifically, SKILL.md directs the model to profile the PyTorch model, rewrite model_new.py using custom kernels, compile within a GPU sandbox, and iterate continuously until the resulting kernel is at least 5% faster than torch.compile at atol=1e-2 and rtol=1e-2.
To address security and performance integrity, the system implements five countermeasures against reward hacking: permission-locked verification and profiling scripts, context managers prohibiting torch.nn.functional fallbacks, consistency checks against five random inputs, device synchronization and warm-up profiling, and the exclusion of any web search tool. Rewards are discrete rather than continuous speedup ratios, categorized as r ∈ {-1, 1, 2, 3}: -1 for correctness failures, 3 if clearing both eager and torch.compile by over 5%, 2 for clearing eager only, and 1 otherwise.
Kernel optimization has traditionally been one of the most meticulous bottlenecks in high-performance deep learning, requiring engineers to write low-level code to extract maximum hardware efficiency. Applying reinforcement learning via autonomous agents represents a major milestone in automating this complex workflow. However, the heavy hardware reliance on 128 NVIDIA H20 GPUs underlines that full-scale execution remains tightly bound to well-funded enterprise and research infrastructure.
Main benchmark results demonstrate a 98.8% overall pass rate, running 98.4% faster than eager execution and 96.8% faster than torch.compile, achieving 2.60x and 2.11x geometric means respectively. Case studies further illustrate policy learnings, such as a diagonal matrix multiplication rewritten into row-wise scaling achieving 73.31x over torch.compile.
Source: MarkTechPost
Found something wrong in this article? Report an issue with this article
Comments
Leave a Comment