CUDA-L2: Surpassing cuBLAS Performance for Matrix Multiplication through Reinforcement Learning (2512.02551v1)

Abstract: In this paper, we propose CUDA-L2, a system that combines LLMs and reinforcement learning (RL) to automatically optimize Half-precision General Matrix Multiply (HGEMM) CUDA kernels. Using CUDA execution speed as the RL reward, CUDA-L2 automatically optimizes HGEMM kernels across 1,000 configurations. CUDA-L2 systematically outperforms major matmul baselines to date, from the widely-used {\it torch.matmul} to state-of-the-art Nvidia's closed-source libraries, i.e., {\it cuBLAS}, {\it cuBLASLt}. In offline mode, where kernels are executed consecutively without time intervals, CUDA-L2 yields +22.0\% over {\it torch.matmul} on average; +19.2\% over {\it cuBLAS} using the optimal layout configuration (normal-normal NN and transposed-normal TN); +16.8\% over {\it cuBLASLt-heuristic}, which queries {\it cuBLASLt} library and selects the algorithm based on the heuristic's suggestion; and +11.4\% over the most competitive {\it cuBLASLt-AutoTuning} model, which selects the fastest algorithm from up to 100 candidates from {\it cuBLASLt}'s suggestions. In server mode, where kernels are executed at random intervals simulating real-time inference, the speedups further increase to +28.7\%, +26.0\%, +22.4\%, and +15.9\% for {\it torch.matmul}, {\it cuBLAS}, {\it cuBLASLt-heuristic}, and {\it cuBLASLt-AutoTuning} respectively. CUDA-L2 shows that even the most performance-critical, heavily-optimized kernels like HGEMM can be improved through LLM-guided RL automation by systematically exploring configuration spaces at scales impractical for humans. Project and code can be found at github.com/deepreinforce-ai/CUDA-L2

• The paper demonstrates that an RL-guided LLM pipeline can auto-optimize CUDA HGEMM kernels, achieving up to 22% speedup over cuBLAS.
• It integrates multi-stage reinforcement learning and contrastive methods to balance speed, correctness, and code brevity in kernel generation.
• Benchmarking on 1,000 configurations shows improved performance in both offline and server settings, highlighting robust generalization capabilities.

CUDA-L2: Reinforcement Learning-Guided Optimization of Matrix Multiplication Kernels

Overview

This work introduces CUDA-L2, an automated kernel optimization system that combines LLMs and reinforcement learning (RL) to systematically surpass the performance of NVIDIA's cuBLAS and related libraries for Half-precision General Matrix Multiply (HGEMM) across a comprehensive set of matrix sizes. The approach addresses the documented limitations in manual tuning and static code generation, utilizing RL-guided large models for kernel code generation and selection, achieving robust average speedup—up to 22% over torch.matmul and 19.2% over cuBLAS—over 1,000 configurations, with even higher gains in server mode. The implication is a new paradigm for performance-critical kernel development, portending wide applicability in deep learning workloads on novel hardware targets (2512.02551).

Motivation and Problem Scope

Matrix multiplication (matmul), and specifically HGEMM, underpins computational workloads in LLMs and transformer-based models, constituting a performance-critical kernel in both training and inference paths. While NVIDIA’s cuBLAS and cuBLASLt libraries offer high performance via expert-crafted, architecture-specific kernel schedules, the vast combinatorial space dictated by $(M, N, K)$ dimensions, layout constraints, and underlying hardware implies that manual or heuristic-based approaches leave significant performance untapped. This challenge is amplified by divergent optimal strategies across architecture generations (Ampere, Hopper, Blackwell, etc.) and tile sizes, as well as by input-dependent optimal parameterization.

The key technical gap lies in discovering optimized strategies that generalize across this space without sacrificing correctness or incurring reward hacking during RL or code generation. Prior LLM-based systems and RL-augmented frameworks lack systematic, configuration-wide coverage or only operate on narrow benchmarks unsuited for production-level sweep over real model workloads.

CUDA-L2 Methodology

CUDA-L2 advances the state-of-the-art through a tightly integrated pipeline:

• Continued Pretraining: The LLM backbone is further pretrained on a highly diverse corpus of CUDA code, spanning established libraries (e.g., ATen, CUTLASS) and web-scraped examples, with instruction synthesis for RL prompt engineering. Retrieval-augmented context is used to incorporate architectural and API updates in the foundation model.
• Multi-Stage RL Training: Optimization is staged from general-kernel RL (on varied CUDA primitives) to HGEMM-specific RL, utilizing execution time as reward, and integrating comprehensive NVIDIA Nsight Compute profiling metrics to inform reward shaping.
• Contrastive RL with Correctness Regularization: Speed, numerical deviation (bounded by established kernel outputs), and code brevity are jointly optimized. Hacking opportunities (multiple streams, lazy execution) are suppressed by enforcing compilation and direct execution of generated ${.cu}$ files via nvcc.
• Coverage and Generalization: CUDA-L2 is evaluated on $10^3$ configurations, including those deployed in commercial LLMs, demonstrating adaptability across input size, tiling strategies, and both offline (back-to-back kernels) and server (random interval, cache-cold) inference settings.

Benchmark Results

CUDA-L2's results on A100 GPUs, benchmarked against torch.matmul, cuBLAS, cuBLASLt-heuristic, and cuBLASLt-AutoTuning, validate both mean speedup and win rates:

Figure 1: CUDA-L2 speedup against torch.matmul, cuBLAS, cuBLASLt-heuristic, and cuBLASLt-AutoTuning on 1,000 HGEMM configurations in offline and server settings. Offline: consecutive execution, Server: randomized intervals simulating live inference.

Key findings:

• Offline (back-to-back execution): CUDA-L2 exceeds torch.matmul by +22.0%, cuBLAS-optimal by +19.2%, cuBLASLt-heuristic by +16.8%, and cuBLASLt-AutoTuning by +11.4%.
• Server (randomized intervals, inference-mimicking): Gains increase to +28.7%, +26.0%, +22.4%, and +15.9% respectively, confirming improved cache utilization and hardware scheduling adaptation by RL-discovered kernels.
• Max selection scenarios (choosing the best of CUDA-L2 and a baseline for each configuration) further raise speedups, supporting ensemble deployment for maximal throughput.

The strongest relative improvements materialize for small and mid-sized matrix shapes, where manual tile and pipeline schedules are insufficient and the search space of possible tilings remains underexplored by heuristics.

Technical Analysis of Optimizations

Inspection of CUDA-L2-generated kernels reveals that the RL/LLM pipeline not only finds canonical best practices but extends them with context-adaptive variations:

• Abstraction and Tiling Adaptation: The system selects between raw WMMA and CuTe abstractions based on problem scale, tuning tiling/block sizes (BM, BN, BK) non-uniformly to maximize compute or occupancy, and even applies input padding with zeros to flexibly exploit more efficient tile shapes.
• Advanced Memory Scheduling: CUDA-L2 discovers and discriminates between single/double-buffered register usage, aggressive multi-step prefetching (for deep pipelines with large K), direct wide-type (e.g., uint128_t) register-to-shared-memory copies to minimize memory traffic, and staggered fetch schedules for deeper compute/memory overlap.
• Block Swizzle Parameterization: The usage and stride of block swizzle are modulated in response to problem scale, with high-stride swizzling employed for large configurations where L2 cache thrashing is a limiting factor.
• Stage Numbering in Pipelining: The number of pipeline stages increases with the K dimension, in line with deeper matrix products requiring multi-stage latency hiding; the RL policy uncovers these empirical breakpoints without domain-specific hardcoding.

Additional analyses demonstrate that the system learns to avoid, or penalize, solutions that would incur correctness degradation or reward hacking, as validated by matching CPU-based FP32 ground truths with appropriately tight tolerances.

Practical and Theoretical Implications

Practically, CUDA-L2 demonstrates that RL-guided LLM code generation can outperform closed-source, vendor-optimized libraries even for the most scrutinized and performance-critical CUDA workloads. The framework enables configuration-wide, architecture-general code synthesis, capable at scale and within software deployment cycles unsustainable for manual tuning.

Theoretically, the paper underscores the value of combined reward sources (speed, correctness, brevity), advanced reward shaping, and contextual profiling incorporation in RL for code generation. It illustrates the emergence of latent domain knowledge, including workload-specific tiling, pipelining, and hardware-aware scheduling, from a pretraining + RL pipeline.

For future research, CUDA-L2 signals the plausibility of entirely automated kernel optimization for new hardware targets (e.g., Ada Lovelace, Hopper, Blackwell), and for other key primitives (convolution, reduction, attention). The modular system design will facilitate both architecture migration and extension to mixed-precision, sparse, or low-rank GEMM variants, as well as non-GEMM operators.

Conclusion

CUDA-L2 establishes a new data-driven, RL-guided, LLM-based paradigm for high-performance kernel generation, achieving consistent speedups over NVIDIA's state-of-the-art closed-source libraries for HGEMM across substantial real-world matrix configuration distributions (2512.02551). The system's design, encompassing diverse pretraining, multi-stage RL, correctness targeting, and adaptive scheduling, constitutes an effective template for generalized, architecture-resilient code optimization in modern AI computing. Future work will extend this framework across emerging GPU generations and broader operator classes.