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

In plain words

Multiplying two large grids of numbers is the core job of a neural network, and GPU programmers spend a lot of effort making that job fast. The best results usually come from NVIDIA's own math library, which already tries many internal strategies. This paper asks whether a language model can write faster versions automatically.

The system generates candidate GPU programs, compiles them, runs them, and measures how fast they are. Programs that are faster and still give correct answers earn a higher score, and the model is trained to produce better programs over time. The authors also train it in stages, first on general GPU code and then on matrix multiplication only, and they give it profiler statistics as hints.

The main claim is that the generated kernels beat the standard library baselines on average across a thousand matrix shapes, with the largest gains on small problems. The authors report these numbers themselves.

TL;DR

• CUDA-L2 is an LLM plus reinforcement-learning pipeline that writes half-precision matrix-multiply (HGEMM) CUDA kernels §1 §3.
• The reward is the measured speedup over a reference kernel, minus a penalty for numerical deviation from an FP32 CPU result, minus a penalty for code length eq. 3 §3.2.3.
• It is evaluated on 1,000 (M, N, K) shapes built from ten sizes, and it reports average speedups over torch.matmul, cuBLAS and two cuBLASLt setups §1 §4.
• All numbers are author-reported. The official repository ships A100 kernels for the 16-bit-accumulator release, plus 3090 and H100 kernels, but not the training pipeline benchmarking_offline.py.

Problem

Hand-tuned HGEMM kernels must be tuned per matrix shape, GPU architecture and accumulator precision, and tuning rarely transfers across them §1. Earlier LLM-plus-RL kernel work was evaluated on KernelBench, with one fixed shape per task, so it is unclear how those kernels behave on realistic production shapes §1. The paper's stated target is to beat NVIDIA's cuBLAS and cuBLASLt on large grids of shapes, not on a single shape §1.

The evaluation assumes two timing scenarios: offline, where kernels run back to back, and server, where kernels run at random intervals. Results differ between them §2.4.

Method

Overview. CUDA-L2 extends CUDA-L1 §3.1 in four ways: (1) continued pretraining on more diverse CUDA code, (2) multi-stage RL from general kernels to matmul-only kernels, (3) NCU profiler metrics in the prompt context, and (4) retrieval-augmented context for knowledge the base model lacks §1 §3.2. The base model is DeepSeek 671B §3.2.1.

Step 1: continued pretraining §3.2.1.

1. Collect CUDA code from web sources, cleaned with rules and an LLM, and from libraries such as PyTorch, ATen and CUTLASS §3.2.1.
2. For web snippets, generate a natural-language instruction per snippet with Claude Sonnet 4 §3.2.1.
3. Use each instruction as a search query; the retrieved documents and examples are added as context §3.2.1.
4. Train on the resulting instruction, context and code triplets §3.2.1.

Step 2: general-kernel RL §3.2.2. Train on about 1K kernels from PyTorch, ATen and CUTLASS, each paired with a reference implementation. The reward is the average speedup over all test iterations. Contrastive RL and GRPO are used, and rewards are smoothed and clipped against reward hacking §3.2.2 §3.1.

Step 3: HGEMM RL §3.2.3. Restrict the tasks to HGEMM with varying (M, N, K). The reward is given in eq. 3 below. The context also includes NCU metrics such as memory throughput, compute utilisation, warp occupancy and cache hit rates §3.2.3. The output must be CUDA C++, CuTe, inline PTX or CUTLASS templates, compiled with nvcc. Python DSLs such as Triton are excluded §3.2.3.

Key equations, as written in the paper:

s(\text{custom}) = \frac{t_{\text{ref}}}{t_{\text{custom}}} - 1 eq. 2

r(\text{custom}) = \frac{1}{N} \sum_{i=1}^{N} \left[\frac{t^{i}_{\text{ref}}}{t^{i}_{\text{custom}}} - \alpha \cdot \text{diff}^i \right] - \beta \, L(\text{custom}) eq. 3

where diff^i is the maximum element-wise absolute difference from the FP32 CPU reference in iteration i, L is the code length, and α > 0 and β > 0 are penalty coefficients §3.2.3 eq. 3.

The reference correctness value is the FP32 CPU product eq. 1 §2.3.

Reward (one kernel), following eq. 3:

# Reward for one generated HGEMM kernel  [§3.2.3] [eq. 3]
# custom_times, ref_times: per-iteration times  [§2.4]
# diffs: per-iteration max |FP32 reference - custom| output difference
# code_length: length of the generated code  [§3.2.3]
def hgemm_reward(custom_times, ref_times, diffs, code_length):
    # Speed term minus a numerical penalty, per iteration  [eq. 3]
    # alpha and beta are positive; their values are not stated
    terms = []
    for t_ref, t_custom, diff in zip(ref_times, custom_times, diffs):
        terms.append(t_ref / t_custom - alpha * diff)
    # Average over iterations, then subtract the code-length penalty  [eq. 3]
    mean_term = sum(terms) / len(terms)
    return mean_term - beta * code_length

Timing protocol for one shape and one mode §2.4:

# Timing protocol for one (M, N, K) shape  [§2.4]
# Offline: kernels run back to back. Server: idle gaps are not timed.
def evaluate_kernel(custom, reference, shape):
    # Warm up for 10 seconds, untimed  [§2.4]
    warm_up(custom, reference, seconds=10)
    speedups = []
    # Keep measuring until at least 30 seconds have passed  [§2.4]
    while elapsed_seconds() < 30:
        # Randomise which kernel runs first  [§2.4]
        order = shuffle([custom, reference])
        t_custom, t_ref = time_both(order, shape)
        # Single-run speedup, a ratio minus one  [eq. 2]
        speedups.append(t_ref / t_custom - 1)
    # Reported score is the mean over runs  [§2.4]
    return mean(speedups)

Correctness check with binary inputs §2.3 eq. 1:

# Exact-match check with random 0/1 matrices A and B  [§2.3]
# Reference: FP32 CPU product, exact for integers below 2048  [§2.3]
def exact_match_ok(custom_out, ref_out):
    # Compare only entries whose reference value is below 2048  [§2.3]
    mask = ref_out < 2048
    # Those entries must equal the reference exactly  [§2.3]
    return all(custom_out[mask] == ref_out[mask])

Repeat with several random inputs. Any single failure marks the kernel incorrect §2.3. A second criterion bounds the deviation by the spread among NVIDIA baseline kernels on the same input §2.3.

Anti-hacking rules §2.4.1 §3.2.3. Extra CUDA streams are disallowed, and kernels must be .cu files, which avoids Python lazy evaluation §2.4.1. Timing uses CUDA events around the call with a synchronize §2.4.1. Length of generated code is penalised §3.2.3.

Discovered techniques §5. The authors describe, without a full recipe: swizzled shared-memory layouts with bank-conflict avoidance (bits, base, shift), multi-stage shared-memory pipelining (n_stage), asynchronous 128-bit global-to-shared copies, register accumulation, block swizzle (swizzle_stride), and epilogue optimisation §5.3. Specific variants include ping-pong double-buffered register fragments §5.3.1, multi-step register prefetch §5.3.2, direct register-to-shared epilogue copies with wide atoms §5.3.3, and staggered A/B prefetch §5.3.4. Abstraction choice shifts from raw WMMA for small shapes to CuTe for large ones §5.1. Padding M with zeros so that a non-divisor tile size can be used is also found §5.2.

Hyperparameters

Only values stated in the paper or repository are listed. Everything else is "not stated".

Results

Author-reported. Setup: 1,000 HGEMM shapes; A100 GPU; baselines are torch.matmul, cuBLAS (NN, TN, max), cuBLASLt-heuristic and cuBLASLt-AutoTuning (NN, TN, max) §2.2 §4. "Mean" is the mean of per-shape speedups, where a value of +22.0% means a time ratio of about 1.22 §2.4 eq. 2. Values from tab. 1 (rendered as an image in the source).

Source: tab. 1. The paper's win rates "79.3% to 95.7%" match the wins column §4.1. Standard deviations are listed in tab. 1 but are not repeated here.

Headline comparison, CUDA-L2 against each baseline with max over layouts tab. 1 §4.1:

Taking the best of CUDA-L2 and the baseline per shape tab. 2 §4.2:

Speedup against cuBLASLt-AutoTuning-max falls as problem size grows: about 1.4× for small problems (log2(MNK) about 18–20) and about 1.0× above about 38. The authors attribute this to larger matrices saturating the GPU §4.3 fig. 3.

Figure 1 plots the same benchmark over NN, TN and max for offline and server modes fig. 1.

Paper vs code

Official code: github.com/ornith-ai/CUDA-L2, MIT License, cloned at code_commit (main branch, shallow clone) [code: LICENSE.txt]. The repository is linked from the paper's title block.

• Hardware scope. The paper's v4 text says CUDA-L2 is limited to A100 and that Ampere 3090, Ada, Hopper and Blackwell are future work §1 §7. The repository ships A100 kernels and also 3090 and H100 kernels, with accumulator-precision variants README.md benchmarking_offline.py. This is a later release than the paper, not a contradiction of its A100 results.
• Accumulator type. The paper does not say explicitly which accumulator the main A100 table uses. The repository's Dec 2, 2025 release (the paper date) is a100_F16F16F16F16 README.md. The baselines use FP16 compute §2.2.
• Training code is not published. The repository has kernels, benchmarks, defense checks and a pybind wrapper, but no RL, reward or pretraining code benchmarking_offline.py. Equations 3 and the RL settings cannot be checked against code.
• Timing scripts take their durations as required arguments. Warm-up and benchmark lengths are CLI inputs with no defaults benchmarking_offline.py benchmarking_server.py. The paper's 10 s and 30 s values are therefore not confirmed as the repository defaults.
• Anti-hacking checks. The paper describes two rules §2.4.1. The repository's defense module checks more: stream injection, lazy evaluation, precision downgrade and elapsed-time monkey-patching defense.py defense.py defense.py.
• Correctness threshold. The paper's 2048 exact-match threshold §2.3 is mirrored by a comment in the checker zero_one_correctness_check.py.
• Ownership check. The repository's LICENSE names "deepreinforce.ai" as the copyright holder LICENSE.txt. The paper's authors list ornith.ai addresses and the paper links the ornith-ai GitHub organisation §1. We accept the repository as the authors' because the paper links it, but the copyright line and organisation names differ.

Pitfalls

• Speedup is a mean of ratios minus one. "+22.0%" averages t_ref/t_custom − 1 over runs and shapes §2.4 eq. 2. It is not a ratio of total times, and it is not a geometric mean. Compare numbers only within the same definition.
• Server numbers are noisier. The paper attributes this to GPU thermal behaviour: idle gaps cause boosting and throttling §4.1. Server standard deviations are about 0.18 to 0.28 against about 0.14 to 0.21 offline tab. 1. The interval time is excluded from timing §2.4.
• Exactness only holds below 2048. The binary-input test ignores outputs at 2048 or above §2.3. Those entries are checked only indirectly, by the baseline-bounded deviation criterion §2.3. A kernel can pass the exact test and still be wrong on large outputs. Implementers should add their own tolerance checks.
• The sampling probability is not given. The paper says it is adjusted per matrix size so that a useful fraction of outputs lies below 2048 and above 0 §2.3. The values are not stated.
• Zero padding changes the problem. Padding M to a multiple of BM adds computed rows §5.2. Output buffers must be sized for the padded shape, and the padded rows must be dropped. The paper's example adds about 1.6%.
• Tile sizes are not a formula. BK correlates only weakly with K (ρ = 0.256) §6.1. Stage count and swizzle stride are described as correlations, not rules §6.2 §6.3.
• Prefetch depth trades registers for latency. Double buffering and multi-step prefetch help when registers are free and K is large, and hurt when register pressure is high §5.3.1 §5.3.2. Which to use is decided per shape by RL, so there is no single setting to copy.
• Figure numbering in the source. The two plots referenced as fig. 3 and fig. 4 (speedup vs size, and tile/stage selection) are placed in table floats in the LaTeX source tab. 3 tab. 4. The compiled PDF may label them "Table". The text calls them Figure, so we keep fig. 3 and fig. 4 here.
• Tables are images in the source. tab. 1 and tab. 2 are PNG files in the arXiv source. Numbers above were read from the images and cross-checked against the text §4.1 §4.2. A reader who needs the exact numbers should check the paper.
• Only FP16 inputs with the configured accumulator are covered. Other dtypes and shapes outside the ten-size grid are not evaluated in the paper.

Lineage

• Builds on: CUDA-L1, an LLM plus contrastive-RL pipeline for KernelBench §3.1. CUDA-L2 changes four things §3.2.
• Compared against: torch.matmul, cuBLAS, cuBLASLt-heuristic and cuBLASLt-AutoTuning §2.2. Other kernel-generation work named in the paper: Sakana AI's AI CUDA Engineer and KernelBench §1.
• Superseded by: not stated in this paper.