🤖 machine learning

CUDA-L2: Surpassing cuBLAS Performance for Matrix Multiplication through Reinforcement Learning

The paper introduces CUDA-L2, a system that leverages large language models and reinforcement learning to automatically optimize Half-precision General Matrix Multiply (HGEMM) kernels, achieving significant performance gains over established baselines like torch.matmul, cuBLAS, and cuBLASLt by systematically exploring configuration spaces at scales impractical for human engineers.

Original authors: Songqiao Su, Xiaoya Li, Albert Wang, Guoyin Wang, Jiwei Li, Chris Shum

Published 2026-07-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Songqiao Su, Xiaoya Li, Albert Wang, Guoyin Wang, Jiwei Li, Chris Shum

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to bake the perfect cake. For years, the world's best bakers (human experts) have been tweaking the recipe for a specific type of cake called "Matrix Multiplication." This cake is the secret ingredient in almost every modern AI brain. These bakers have spent years making it faster, but they hit a wall: every time the cake size changes (from a tiny cupcake to a massive wedding tier), they have to start from scratch, and the tricks that worked for one size often fail on another. It's like trying to use a tiny spoon to eat a giant soup; the tool just doesn't fit.

Enter CUDA-L2, a new team of "AI bakers" built by combining a super-smart language model with a video game-style learning system called Reinforcement Learning (RL). Instead of asking a human to guess the best recipe, CUDA-L2 plays a massive game of "try, fail, learn, and try again."

The Big Discovery: The AI Bakes Faster Cakes

The main finding of this paper is that CUDA-L2 can automatically design these matrix multiplication "recipes" (called kernels) and bake them faster than the best human-made recipes currently available.

The researchers tested this on 1,000 different cake sizes (combinations of dimensions M, N, and K ranging from 64 to 16,384). These sizes cover the exact dimensions used in famous open-source AI models like Qwen, Llama, and DeepSeek.

Here is how much faster the AI bakes compared to the current champions, measured on an A100 GPU:

Vs. The Standard Tool (torch.matmul): In a continuous baking session (offline mode), CUDA-L2 is 22.0% faster. In a busy kitchen where orders come in randomly (server mode), it's 28.7% faster.
Vs. The Gold Standard (cuBLAS): Even against NVIDIA's own highly optimized library, CUDA-L2 wins. It's 19.2% faster in continuous mode and 26.0% faster in the busy server mode.
Vs. The "Auto-Optimizer" (cuBLASLt-AutoTuning): This is the most important comparison. cuBLASLt-AutoTuning is a tool that tries up to 100 different recipes to find the best one. CUDA-L2 still beats it by 11.4% in continuous mode and 15.9% in server mode.

The paper is very sure about these numbers because they were measured directly by running the code thousands of times, not just guessed or simulated.

What the AI Did NOT Do

It is important to know what this system doesn't do. The paper explicitly states that the current version of CUDA-L2 is limited to the A100 architecture. However, the framework is designed for broad applicability, and the team is actively working to extend it to other GPU architectures, including older Ampere chips like the RTX 3090, as well as newer Hopper (e.g., H100) and Blackwell (e.g., B200) chips. The paper also argues against the idea that human experts have "solved" matrix multiplication; the fact that the AI found better recipes proves that human tuning is far from perfect.

How the AI Learned to Bake Better

The AI didn't just guess; it learned specific, clever tricks that even human experts might miss because they are too busy to check every single possibility.

Padding the Cake: Imagine you have a cake that is 8,192 inches wide, but your baking pan only fits perfectly if the width is divisible by 160. 8,192 isn't divisible by 160. A human might say, "I'll use a pan size of 128 because it fits." But the AI said, "I'll add a tiny bit of extra batter (padding) to make the cake 8,320 inches wide so I can use the 160-inch pan." This tiny bit of extra work actually made the whole process faster.
The Ping-Pong Strategy: Usually, a baker loads ingredients, mixes them, then loads the next batch. The AI figured out a "ping-pong" method: while the mixer is working on Batch A, the assistant is already loading ingredients for Batch B. This means the mixer never has to wait.
The "Staggered" Delivery: Sometimes, the AI realized that asking for two ingredients at once (A and B) causes a traffic jam in the kitchen. Instead, it asked for ingredient A, started mixing, and then asked for ingredient B. This kept the kitchen moving smoothly.
Choosing the Right Pan Size: The AI learned that for small cakes, small pans work best. But for giant cakes, it needs much larger pans. It figured out the perfect size for every single cake dimension, a task that would take a human a lifetime to calculate for 1,000 different sizes.

The "Server" vs. "Offline" Difference

The paper also noticed something interesting about the kitchen environment.

Offline Mode: Imagine a factory line where cakes are baked one after another without stopping. The oven stays hot, and the workers are in a rhythm. Here, the AI was 11.4% to 22.0% faster than the competition.
Server Mode: Imagine a busy restaurant where orders come in at random times. The oven might cool down between orders, and the workers have to warm up again. In this "real-world" chaos, the AI's advantage actually grew, beating the competition by 15.9% to 28.7%. The paper suggests this is because the AI's recipes are better at handling these "cold starts" and unpredictable pauses.

The Bottom Line

The paper concludes that even for the most critical, heavily optimized tasks like matrix multiplication, LLM-guided Reinforcement Learning can find better solutions than humans by exploring a space of possibilities that is too huge for a person to check. While this specific version of CUDA-L2 is currently limited to A100 GPUs, the framework is designed to be expanded to other chips in the future.

In short: The AI didn't just tweak the recipe; it discovered entirely new ways to bake the cake that humans hadn't thought of, proving that even in a field as mature as GPU optimization, there is still plenty of room for improvement.

Technical Summary: CUDA-L2

Problem Statement

Matrix multiplication (matmul) is a fundamental operation in Large Language Models (LLMs), accounting for a significant portion of computation time in both training and inference. While manually optimized CUDA kernels exist, achieving optimal performance across diverse matrix dimensions ( $M, N, K$ ) and GPU architectures remains a significant challenge. Different dimension triplets often require distinct optimization strategies that do not transfer well across hardware. Furthermore, even for identical dimensions, optimal strategies vary based on setup details such as accumulator precision (FP16 vs. FP32).

Existing automated approaches, such as those using Large Language Models (LLMs) for kernel generation, have primarily focused on benchmarks with single, fixed configurations (e.g., KernelBench). It remains unclear how these methods translate to production environments where they must compete against highly optimized, closed-source libraries like NVIDIA's cuBLAS and cuBLASLt. There is a lack of automated systems capable of systematically exploring the vast configuration space of Half-precision General Matrix Multiply (HGEMM) to surpass human-expert-optimized baselines.

Methodology: CUDA-L2

The authors propose CUDA-L2, a system that combines Large Language Models (LLMs) and Reinforcement Learning (RL) to automatically optimize HGEMM CUDA kernels. The system is designed to handle 1,000 distinct configurations representing all combinations of $M, N, K$ values from $\{64, 128, 256, 512, 1024, 2048, 4096, 8192, 12288, 16384\}$ .

System Architecture and Training

CUDA-L2 extends the previous CUDA-L1 framework with four key technical enhancements:

Continued Pretraining: The model is pre-trained on a diverse corpus of CUDA code beyond standard benchmarks, including sources from PyTorch, ATen, CUTLASS, and NVIDIA tutorials. Instruction descriptions are generated for raw code snippets, and retrieval-augmented context is added to incorporate new knowledge and architectural characteristics.
Multi-Stage RL Training: The training progresses from broad general kernel types to specialized domains (specifically matmul).
- General Kernel RL: Trains the model on a wide range of operations (linear algebra, convolution, reduction, etc.) using execution speed as the reward.
- HGEMM RL: Focuses specifically on HGEMM kernels with varying $M, N, K$ configurations.
Enhanced Context with Profiling Metrics: The RL context includes comprehensive NVIDIA Nsight Compute (NCU) profiling metrics (e.g., memory throughput, SM occupancy, cache efficiency) to enable better optimization decisions beyond simple end-to-end execution time.
Reward Function Design: The reward for a generated kernel is based on the average speedup score across test iterations. To ensure robustness, the reward is penalized for:
- Numerical deviation from an FP32 CPU ground truth.
- Excessive code length (to encourage conciseness).
- The system enforces strict correctness criteria, including exact matches for binary inputs and bounded deviation against reliable baselines (cuBLAS, cuBLASLt).

Evaluation Protocol

The system is evaluated under two scenarios to reflect different deployment realities:

Offline Mode: Kernels are executed consecutively without time intervals, simulating peak throughput.
Server Mode: Kernels are executed at random intervals to simulate real-time inference, accounting for GPU thermal dynamics (cooling periods and clock boosting).

The evaluation compares CUDA-L2 against four baselines:

torch.matmul (PyTorch's default dispatch).
cuBLAS (using optimal layout selection between NN and TN).
cuBLASLt-heuristic (selecting the top-ranked algorithm from NVIDIA's heuristic API).
cuBLASLt-AutoTuning (exhaustively testing up to 100 candidates to find the fastest).

Key Contributions and Discovered Techniques

CUDA-L2 demonstrates that LLM-guided RL can discover superior implementations by systematically exploring configuration spaces impractical for manual tuning. The system autonomously discovers and applies several advanced optimization techniques:

Abstraction Selection: The model automatically selects between raw WMMA intrinsics for small matrices and CuTe abstractions for larger, complex tiled operations.
Input Padding: The system discovers strategies where input matrices are padded with zeros to non-divisible dimensions (e.g., padding $M=8192$ to $8320$) to enable more efficient tile sizes ( $B_M$ ) that would otherwise be invalid, balancing the overhead of extra computation against performance gains.
Advanced Memory and Compute Overlap:
- Double-Buffered Register Fragments: Alternating between two sets of register fragments to overlap data movement with computation (ping-pong execution).
- Aggressive Prefetching: Moving beyond single-step prefetching to multi-step prefetching (loading multiple iterations ahead) when register pressure allows.
- Staggered Prefetch Scheduling: Interleaving the prefetching of matrices A and B around the matrix multiply-accumulate (MMA) operation to maximize instruction-level parallelism.
- Direct Register-to-Shared-Memory Copy: Eliminating intermediate temporary tensors in the epilogue when register layouts match shared memory layouts, using wide data types (e.g., uint128_t) to reduce copy operations.
- Block Swizzling: Reordering thread block execution to improve L2 cache hit rates, particularly for large problem sizes.

Results

CUDA-L2 systematically outperforms all major baselines across 1,000 configurations on A100 GPUs.

Offline Mode Performance (Average Speedup over Baselines):

+22.0% over torch.matmul.
+19.2% over cuBLAS (using optimal layout).
+16.8% over cuBLASLt-heuristic.
+11.4% over cuBLASLt-AutoTuning (the strongest baseline, which exhaustively tests 100 candidates).

Server Mode Performance:
Speedups increase in the server scenario, likely due to the system's ability to better handle thermal dynamics and cache cooling effects:

+28.7% over torch.matmul.
+26.0% over cuBLAS.
+22.4% over cuBLASLt-heuristic.
+15.9% over cuBLASLt-AutoTuning.

The win rates range from 79.3% to 95.7% across all baselines, indicating that improvements are systematic rather than driven by outliers. The system also shows that speedups are generally higher for smaller matrix sizes where GPU utilization is lower, while larger matrices saturate the GPU's throughput, narrowing the gap with baselines.

Significance

The paper claims that CUDA-L2 demonstrates that even the most performance-critical and heavily optimized kernels, such as HGEMM, can be improved through LLM-guided RL automation. By scaling optimization across configuration spaces that are difficult for humans to tune manually, the system achieves performance gains that surpass state-of-the-art closed-source libraries.

While the current version is limited to A100 (Ampere) architectures, the framework is designed for broad applicability. The authors note ongoing work to extend the system to other GPU architectures, including Ada Lovelace, Hopper, and Blackwell. The work suggests that automated methods can systematically explore vast configuration spaces to find optimizations that manual tuning or heuristic-based libraries might miss.