KernelMosaic Open-Sources CUDA Kernel Dataset and Training Trajectories

Summary:KernelMosaic introduces an open dataset and agentic framework for training models to generate CUDA kernels. It combines heterogeneous task specifications, verified solutions, and iterative implementation, repair, and optimization trajectories into supervision for supervised fine-tuning. After training on KernelMosaic, a small open-source model improves its correctness on KernelBench Levels 1, 2, and 3 from 4%, 0%, and 0% to 88%, 82%, and 16%, respectively. The fraction of kernels that outperfor

---

From One-Shot Generation to a Verified CUDA Kernel Training Loop

KernelMosaic addresses a persistent engineering problem: a large language model may produce syntactically valid CUDA code without simultaneously satisfying API constraints, numerical correctness, compiler compatibility, and runtime performance. Rather than adding more static code samples, the work introduces an open dataset of heterogeneous tasks, verified solutions, and optimization trajectories, together with an agentic framework for constructing those trajectories.

The data represents tasks through natural-language descriptions, reference programs, mathematical specifications, and combinations of these forms. This allows the model to learn how requirements map to implementations instead of merely copying answers from fixed templates. Each solution also records implementation, compilation testing, correctness repair, and performance optimization, exposing the behavioral adjustments required after an initial failure.

How the Dataset and Agentic Trajectories Are Built

Tasks undergo human-guided review and diversity-aware selection, reducing the dominance of similar operators, duplicated templates, and a single specification style. During generation, a strong model drives an agentic framework that repeatedly:

1. Produces a candidate kernel from the task or specification;
2. Compiles and executes it to capture compiler and numerical errors;
3. Repairs the implementation using execution feedback;
4. Optimizes the kernel with reusable CUDA engineering guidance and performance feedback;
5. Stores intermediate trajectories together with the verified final solution.

The resulting supervision is therefore not limited to final answers. It includes tool calls, tests, and revision records. Training experiments show that this trajectory-level supervision substantially outperforms direct-response supervision and encourages the model to continue testing and revising its code.

Correctness Gains from Supervised Fine-Tuning

The researchers fine-tune a small open-source model and report the following changes in correctness on KernelBench:

| KernelBench level | Before fine-tuning | After fine-tuning | Change | |---|---:|---:|---:| | L1 | 4% | 88% | +84 percentage points | | L2 | 0% | 82% | +82 percentage points | | L3 | 0% | 16% | +16 percentage points |

The gains on L1 and L2 suggest that relatively explicit fundamental operators can benefit greatly from a compact set of high-quality trajectories. The 16% result on L3 indicates that complex, real-world optimization problems remain far beyond the dataset's current capability envelope.

Performance also improves in the reported comparison: the fraction of kernels that run faster than PyTorch eager execution increases at every level. However, the fraction of faster tasks is not a substitute for a complete performance study. It does not expose the acceleration distribution, absolute latency, memory consumption, or gap to the best hand-written implementation.

CUDA Generation Through the GPU Execution Stack

CUDA kernel generation is not ordinary code completion. A model must understand thread hierarchy and index spaces, use threadIdx.x and blockIdx.x to partition global data, and handle blockDim.x boundaries. Common operators also involve shared-memory tiling, bank conflicts, coalesced access, and warp-level synchronization.

For dense matrix multiplication, if each thread computes an output element, global floating-point operation complexity is typically approximately O(n³), while memory traffic depends on caching and tile reuse. Whether a kernel approaches peak performance depends not only on arithmetic count but also on:

• Useful floating-point instructions executed per thread;
• Data movement between global and shared memory;
• Whether tile sizes fit the target GPU's streaming multiprocessors and occupancy targets;
• Whether bank conflicts, boundary checks, and synchronization branches obscure compute utilization;
• Support for asynchronous copies, double buffering, and vectorized loads.

Including execution feedback in training is therefore well motivated. At the same time, data generation and evaluation need fixed hardware, compiler versions, and timing protocols. Otherwise, a model may learn the quirks of one GPU architecture or toolchain rather than transferable parallel optimization principles.

Engineering Trade-offs and Evaluation Boundaries

The principal cost of trajectory supervision is moving online exploration expense into dataset construction. A strong model must compile, run, and repair code across multiple rounds, so producing a high-quality trajectory can cost substantially more than generating a conventional code response. Open data does not automatically remove this cost or make full-scale trajectory reproduction affordable for every team.

KernelBench correctness is better suited to functional validation than production kernel quality. A comprehensive evaluation should include at least:

• Speedups over PyTorch eager execution and mature library implementations;
• Robustness across shapes, data types, and boundary sizes;
• First-attempt compilation success and average repair rounds;
• Numerical error, random-seed sensitivity, and race conditions;
• Peak memory use, occupancy, and energy consumption;
• Transfer across A100, H100, and consumer-grade GPUs.

The 16% L3 correctness rate also shows that current coverage is insufficient for complex composed operators, cross-file dependencies, and aggressive optimization. If most trajectories are generated and filtered by a strong teacher model, researchers should also examine teacher bias and whether familiarity with established CUDA library interfaces has skewed the task distribution.

Deployment Recommendations

For teams seeking to generate CUDA kernels, KernelMosaic is best treated as a training-data resource or research baseline, not as a guarantee of production readiness. Integration should use three layers of validation:

1. Compilation and API checks: Compile with the target CUDA Toolkit and GPU architecture, then validate boundary sizes, index types, and synchronization legality.
2. Numerical regression: Compare against a PyTorch reference implementation across multiple precisions, input sizes, and random seeds using explicit error thresholds.
3. Performance gates: Fix warm-up counts, timing windows, and GPU clocks; report median and percentile speedups rather than only a binary faster-or-slower result.

Training teams should also stratify sampling by operator category, specification source, failure mode, and performance gain. Otherwise, large numbers of simple L1 and L2 examples may overwhelm L3 coverage. Independent test sets should preserve unseen GPU architectures to measure transfer rather than memorization.

Research Value of the Open Dataset

KernelMosaic's central contribution is not proof that a model can already reliably write high-performance CUDA kernels. It provides a more realistic software-engineering data paradigm in which task definition, execution validation, error repair, and performance optimization form continuous supervision. The metric shifts from whether a model succeeds once toward whether it can reliably converge with feedback.

Correctness improving from 4%, 0%, and 0% to 88%, 82%, and 16% demonstrates the potential of that direction. Turning it into a dependable production workflow will still require harder tasks, stronger verifiers, cross-hardware benchmarks, and strict reproduction of performance results.

---

Source: ArXiv Machine Learning (cs.LG) (2026-11-13)

— END —

Tags: #CUDA Kernels Follow #Code Generation Follow #Agentic Training Follow #Supervised Fine-Tuning Follow #Open Dataset Follow