CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation

Weinan Dai; Hanlin Wu; Qiying Yu; Huan-ang Gao; Jiahao Li; Chengquan Jiang; Weiqiang Lou; Yufan Song; Hongli Yu; Jiaze Chen; Wei-Ying Ma; Ya-Qin Zhang; Jingjing Liu; Mingxuan Wang; Xin Liu; Hao Zhou

Source-linked AI summary

CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation

Weinan Dai, Hanlin Wu, Qiying Yu, Huan-ang Gao, Jiahao Li, Chengquan Jiang, Weiqiang Lou, Yufan Song, Hongli Yu, Jiaze Chen, Wei-Ying Ma, Ya-Qin Zhang, Jingjing Liu, Mingxuan Wang, Xin Liu, Hao Zhou

arXiv:2602.24286v1Feb 27, 2026cs.LGcs.AI

TL;DR

CUDA Agent targets the limited CUDA optimization ability of LLMs, which leaves existing approaches uncompetitive with torch.compile. It combines scalable data synthesis, a skill-augmented execution-driven environment, and stable agentic reinforcement learning. On KernelBench, it achieves state-of-the-art results, including 100%, 100%, and 92% faster rate over torch.compile across Levels 1–3.

Problem
Existing CUDA generation approaches remain uncompetitive with torch.compile, while training-free refinement and fixed feedback loops do not fundamentally improve models’ CUDA optimization ability.
Method
CUDA Agent combines scalable data synthesis, a skill-augmented CUDA development environment with automated verification and profiling, and reinforcement learning techniques for stable training.
Results
CUDA Agent achieves state-of-the-art KernelBench performance, with 100%, 100%, and 92% faster rate over torch.compile on Levels 1, 2, and 3, respectively.
Takeaways & Limitations
The results suggest that structured environments and reliable execution-based rewards can move LLMs from passive code generation toward hardware-aware performance optimization.
Takeaways & Limitations
The study does not compare CUDA Agent with more sophisticated compiler frameworks such as TVM and relies on a large GPU pool, limiting accessibility.

Abstract

from arXiv · show

GPU kernel optimization is fundamental to modern deep learning but remains a highly specialized task requiring deep hardware expertise. Despite strong performance in general programming, large language models (LLMs) remain uncompetitive with compiler-based systems such as torch.compile for CUDA kernel generation. Existing CUDA code generation approaches either rely on training-free refinement or fine-tune models within fixed multi-turn execution-feedback loops, but both paradigms fail to fundamentally improve the model's intrinsic CUDA optimization ability, resulting in limited performance gains. We present CUDA Agent, a large-scale agentic reinforcement learning system that develops CUDA kernel expertise through three components: a scalable data synthesis pipeline, a skill-augmented CUDA development environment with automated verification and profiling to provide reliable reward signals, and reinforcement learning algorithmic techniques enabling stable training. CUDA Agent achieves state-of-the-art results on KernelBench, delivering 100\%, 100\%, and 92\% faster rate over torch.compile on KernelBench Level-1, Level-2, and Level-3 splits, outperforming the strongest proprietary models such as Claude Opus 4.5 and Gemini 3 Pro by about 40\% on the hardest Level-3 setting.

1 Introduction

CUDA Agent addresses the gap between general software competence and CUDA kernel optimization by combining scalable data, a structured development environment, and stable agentic reinforcement learning. It reports state-of-the-art KernelBench performance, exceeding torch.compile across difficulty levels.

Motivation: CUDA kernel optimization remains challenging because it requires deep knowledge of GPU microarchitecture and sophisticated profiling tools.
Motivation: Existing CUDA generation methods remain uncompetitive with torch.compile despite LLMs’ strong general software development abilities.
Motivation: Training-free refinement and fixed multi-turn fine-tuning improve kernels without fundamentally resolving base models’ limited CUDA optimization abilities.
Results: 100%, 100%, and 92% faster rate over torch.compile is achieved on KernelBench Levels 1, 2, and 3, respectively.CUDA Agent also supports 128k-token contexts and up to 200 interaction turns.
Approach: CUDA Agent combines scalable data synthesis, a skill-augmented CUDA environment, and reinforcement learning techniques for stable long-context, multi-turn training.
Results: CUDA Agent is presented as a competitive and often superior alternative to compiler-driven CUDA kernel optimization.

2 Related Works

Related work uses test-time search or model training to improve CUDA kernels, but these approaches remain constrained by base-model capability, limited scale, and scarce high-quality data.

Test-time search: Test-time methods such as STARK and ReGraphT search CUDA optimization trajectories using compilation, correctness checks, timing feedback, or retrieval.
Test-time search: These test-time scaling approaches rely heavily on the base model’s CUDA coding capability and are orthogonal to CUDA Agent.
Model training: Training-based approaches include multi-turn reinforcement learning and contrastive reinforcement learning with execution-based rewards.
Open limitations: Existing approaches are constrained by scarce high-quality training data, limited training scale, and hand-designed optimization loops.

3 Method

CUDA Agent combines scalable task synthesis, a skill-integrated coding loop, robust reward scheduling, and stabilization techniques to train CUDA kernel expertise with agentic RL.

Method overview: CUDA Agent combines scalable data collection, a skill-integrated environment with robust reward scheduling, and RL techniques for stable training.These components address data scale, structured CUDA development, reliable feedback, and optimization stability.
Data synthesis: 6,000 samples form CUDA-Agent-Ops-6K2, a curated operator-level training dataset produced by crawling seeds, synthesizing combinations, and filtering execution behavior.Filtering requires Eager and Compile execution, excludes stochastic or hackable operators, restricts eager runtime to 1–100 ms, and removes similar KernelBench cases.
Skill-integrated agent loop: The agent loop follows a ReAct-style workflow that interleaves reasoning, tool execution, and observations for iterative CUDA coding, debugging, and optimization.CUDA-specific skills guide profiling, custom operator implementation, compilation, evaluation, and repeated refinement against correctness and performance requirements.
Robust reward scheduling: The reward r ∈ {−1, 1, 2, 3} jointly reflects correctness and whether the generated kernel exceeds 5% speedup over eager and compile baselines.The indicator b(t,t0) is one when runtime improvement over baseline t0 exceeds 5%; evaluation scripts are protected against reward hacking.

4 Experiments

CUDA Agent is evaluated on KernelBench using correctness, faster-rate, and speed-up metrics against strong model and compiler baselines. Results show that agentic RL, interactive feedback, milestone-based rewards, RFT, and Value Pretraining materially affect optimization quality and training stability.

4.1 Experimental Setup: Evaluations use 250 KernelBench operator tasks across Level 1 to Level 3, with Pass Rate, Faster Rate, and Speed-up as the principal metrics.Baselines include Claude Opus 4.5, Gemini 3 Pro, GLM 4.6, and Kimi K2 under the same agent loop.
4.2 Main Results: CUDA Agent attains a 98.8% Pass Rate and a 96.8% faster rate against proprietary model baselines.Claude Opus 4.5 and Gemini 3 Pro achieve 91.2%–95.2% Pass Rates but only 66%–69% faster rates.
4.2 Main Results: 100% faster rate and 2.80× speed-up over torch.compile are achieved on Level 2 Operator Sequences.The reported advantage is attributed to learned optimization policies exploring hardware-specific fusion, memory-access, and tiling strategies.
4.3 Ablations: Removing the interactive agent loop substantially reduces correctness and optimization quality, with compilation, runtime, and profiler feedback enabling iterative refinement.The single-turn model predicts the final kernel without execution feedback, whereas CUDA Agent operates through multi-turn interactions.
4.3 Ablations: Replacing the robust reward schedule with raw Speed-up Reward preserves comparable functional correctness but substantially weakens optimization performance.The paper reports that normalized, milestone-based rewards align better with consistently faster kernels than noisy runtime-ratio regression.
4.3.3 Impact of Multi-Stage Training: Removing either RFT or Value Pretraining substantially degrades optimization and produces training instability despite largely preserved pass rates.Without RFT, reward collapses as actor entropy rises; without Value Pretraining, low explained variance accompanies inefficient exploration and excessively long trajectories.

5 Conclusion

CUDA Agent combines scalable data synthesis, structured agent environments, and stability-oriented RL training to generate and optimize CUDA kernels through execution-driven workflows. The paper reports consistent gains over torch.compile and strong proprietary models on KernelBench, suggesting that structured environments and reliable execution-based rewards can broaden foundation-model coding beyond syntax.

5 Conclusion: CUDA Agent jointly scales data synthesis, agent environments, and stability-oriented RL training for execution-driven CUDA kernel optimization.The system is presented as moving LLMs beyond syntactic code generation toward hardware-aware performance optimization.
5 Conclusion: The system achieves consistent gains over torch.compile and strong proprietary models on KernelBench.The conclusion frames these results as evidence that structured environments and reliable execution-based rewards can transform passive code generators into active optimizers.

A Details of Collected Training Data

The training data consists of executable PyTorch operator tasks, including primitive, composite, and transformers-library examples. The dataset is synthesized and filtered to balance complexity while reducing structural overlap with evaluation samples, and the agent uses tools for controlled code development.

Operator Format: Each training sample is a self-contained executable PyTorch task comprising an operator class plus initialization and runtime input generators.The operator class subclasses torch.nn.Module and defines __init__ and forward methods.
Data Contamination Check: Training samples are decontaminated by removing examples whose maximum AST similarity to any evaluation program exceeds 0.9.Similarity is computed between extracted Model classes using PythonASTSimilarity.
Final Dataset Composition: Training examples include transformers operators and combinatorial torch operators formed by sequentially stacking one to five primitive classes.The final composition balances simple operators, moderate compositions, and higher-level transformers modules.
Agent Tooling: The agent interacts with the local execution environment through controlled tools for inspection, editing, compilation, debugging, and background-process management.The toolset includes Bash, file read/write, Edit, MultiEdit, Glob, Grep, NotebookEdit, BashOutput, and KillBash.

B.2 Original SKILL.md Content

The CUDA Agent skill specifies a constrained workflow for implementing, testing, and optimizing custom CUDA extensions. It emphasizes correctness verification, profiling-driven iteration, algorithmic optimizations before parameter tuning, and cleanup of intermediate implementations.

Restrictions and Allowed Operations: The workflow restricts C++ to raw CUDA kernels, cuBLAS for GEMM, and cuDNN for convolution, while Python may use tensor creation and custom extension operations.Torch operators, third-party libraries beyond the stated exceptions, and modifications to fixed infrastructure are forbidden.
Implementation: CUDA implementations use paired .cu and binding files, with launchers handling grid configuration, shared memory, streams, and selectable kernel variants.The examples include templated kernels, grid-stride loops, dynamic shared-memory allocation, and configuration-dependent launches.
Optimization Strategy: Optimization prioritizes kernel fusion, shared-memory tiling, and memory coalescing before hardware-utilization and fine-tuning techniques.The prescribed order places algorithmic changes ahead of vectorization, warp primitives, occupancy tuning, instruction-level parallelism, mixed precision, and prefetching.
Parameter Tuning: Parameter tuning is a last resort after algorithmic options are exhausted and only when performance is within 1.2× of the target.The skill also requires benchmarking configurations with warmups, repeated timing, synchronization, and selection of the best configuration.
Iteration Requirements: Correctness failures require iteration until verification passes, including debugging boundary conditions, synchronization, data types, precision, and memory alignment.Fixes are applied in kernel and binding files before recompilation and retesting.
Final Cleanup: Before completion, the kernels directory must contain only the final optimized implementation, with intermediate, old, test, and backup files removed.The final workspace preserves the optimized kernel and binding files while leaving fixed infrastructure and utility directories unchanged.

D.1 Common Optimization Patterns

Across KernelBench difficulty levels, CUDA Agent repeatedly combines algebraic simplification, fusion, memory optimization, hardware awareness, and library use to improve diverse workloads.

Algebraic Simplification and Operator Reduction: Algebraic simplification reduces generic operators by exploiting mathematical structure, such as replacing diagonal matrix multiplication with row-wise scaling.
Kernel Fusion: Kernel fusion combines related operations to avoid intermediate tensors and reduce launch overhead, including arithmetic pipelines and residual addition with activation.
Memory Access Optimization: Coalesced accesses, smaller intermediate footprints, shared memory, and vectorized loads improve memory efficiency across generated kernels.
Hardware-Aware Optimization: TF32 enables Tensor Core use for matrix multiplications and convolutions when lower precision is numerically acceptable.
Library-Aware Optimization: CUDA Agent also invokes fused cuDNN primitives when they better match high-level model semantics and provide hardware-tuned implementations.

D.2 Example of Level 1: Model that performs matrix multiplication of a diagonal matrix with another matrix

For diagonal matrix multiplication, CUDA Agent replaces explicit diagonal construction and general GEMM with direct row-wise scaling, producing a large speedup over Torch Compile.

73.31× speed-up versus Torch Compile results from the custom operator for diagonal matrix multiplication.
The method exploits the equivalence between multiplying by a diagonal matrix and scaling each row of the dense matrix.
Direct row-wise scaling avoids constructing the diagonal matrix and reduces time complexity from O(N 2M) to O(NM).
The specialized kernel fuses diagonal construction with multiplication, reducing kernel launches and global memory traffic.
This case exemplifies replacing a generic operator with a specialized kernel that directly implements the underlying mathematical operation.

D.3 Example of Level 2: Model that performs matrix multiplication, division, summation, and scaling

For a multi-operation matrix computation, CUDA Agent algebraically rearranges the workload and fuses reductions, dot products, division, and scaling into efficient CUDA kernels.

24.04× speed-up versus Torch Compile is achieved by the custom operator for matrix multiplication, division, summation, and scaling.
Linearity transforms matrix multiplication followed by reduction into weight reduction followed by a dot product.
Two custom kernels compute the column-wise weight sum and then fuse dot product, division, and scaling to avoid intermediate global-memory writes.
Float4 loads maximize memory bandwidth, while shared-memory tree reduction accumulates partial sums without expensive global atomics.
The case combines algebraic simplification, kernel fusion, and low-level CUDA optimization to collapse a multi-operator graph into few efficient kernels.

D.4 Example of Level 3: ResNet BasicBlock

For a ResNet BasicBlock, CUDA Agent combines graph transformations, fused cuDNN operations, Tensor Core use, and custom kernels to optimize a realistic neural-network workload.

The reference BasicBlock contains separate convolution, normalization, bias, activation, and residual-addition stages, creating multiple kernel launches.
BatchNorm parameters are folded into convolution weights and bias, eliminating the BatchNorm operator while preserving numerical equivalence.
cudnnConvolutionBiasActivationForward fuses convolution, bias addition, and ReLU, while TF32 enables Tensor Core use on Hopper GPUs.
NHWC layout conversion introduced enough overhead to offset potential gains, so the final implementation retained NCHW with compatible cuDNN APIs.
Residual addition and final ReLU are fused into one custom kernel instead of two separate element-wise operators.
The case integrates graph transformations, library fusion, and custom kernel design for complex neural-network blocks.

E Limitations

The study identifies two main limitations: it omits comparisons with sophisticated compiler frameworks and requires substantial GPU resources for training.

CUDA Agent is compared with torch.compile, but not with more sophisticated compiler frameworks such as TVM.The authors cite tuning overhead and complex deployment requirements as reasons these systems are difficult to integrate into large-scale RL training.
The training pipeline relies on a large GPU pool with process-level isolation, creating considerable computational and engineering costs.The authors note that this resource requirement may limit accessibility and motivate more resource-efficient training strategies.