KernelZero: Co-Evolving Proposer and Coder for Continuously Improved GPU Kernel Generation

• LLM-based GPU kernel generation remains challenging: high-quality data aligned with current model capabilities is scarce, and correctness competes with performance.Developing high-performance kernels also requires domain-specific language and computer-architecture expertise.Source: §1 · pp. 1–2
• KernelZero co-evolves a Proposer that generates Torch modules from API sets and a Coder that translates them into CUDA or Triton kernels.Frontier-driven generation targets the Coder’s current capability, while CA-GRPO gates performance optimization on sufficiently reliable correctness.Source: §1 · p. 2, Figure 1 · p. 2
• KernelZero-7B surpasses Claude-4.5-Sonnet on CUDA and DeepSeek-V4-Pro on Triton, reaching 75.8%/69.6% CUDA pass@1 and 77.2%/72.5% Triton pass@1 on Levels 1/2.Source: §1 · p. 2

2.1. KernelZero Framework

• The Proposer maps sampled Torch API sets to runnable modules, and the Coder maps those modules to executable CUDA or Triton kernels; alternating updates use each model’s reward while the other is fixed.The Proposer’s module distribution moves toward the Coder’s capability frontier, while the Coder adapts to that distribution.Source: §2.1 · pp. 2–3

2.2. Proposer: Proposing Frontier Torch Modules

• API co-occurrence sampling encourages realistic, diverse operator combinations, while validation checks module semantics, execution, and numerical stability.Validation includes AST consistency, runnable execution on CUDA, and filtering outputs containing NaN or Inf.Source: §2.2 · pp. 3–4
• The Proposer’s Frontier Reward is highest when approximately half of sampled Coder kernels solve a proposed module, targeting the Coder’s capability frontier.Invalid modules are penalized; modules should be neither trivial nor impossible.Source: §2.2 · p. 4

2.3. Coder: Generating High-Performance Kernels

• Cold-start distillation teaches four GPU optimization principles: tiling, fusion, pipelining, and reordering.These principles address data reuse, kernel-launch and memory costs, latency hiding, and access-pattern efficiency.Source: §2.3 · pp. 4–5
• The four principles form a taxonomy covering data reuse, latency hiding, execution overlap, and access-pattern optimization across GPU bottlenecks.Source: §2.3 · p. 5
• CA-GRPO separates correctness and performance signals, activating performance optimization only after group correctness reaches a reliability threshold.For correct kernels, the reward incorporates normalized speedup; β controls the strength of that performance signal.Source: §2.3 · pp. 5–6

3.1.1. Training Details

• Both models start from Qwen2.5-Coder-7B; the Coder is cold-start distilled on correctness-filtered CUDA or Triton traces before alternating 20-step GRPO updates for each model.The cold-start traces are generated from cudaLLM instances using gpt-oss-120B as teacher; the alternating training runs for four rounds.Source: §3.1.1 · p. 6

3.1.2. Evaluation Benchmark and Metrics

• The evaluation uses one-shot CUDA or Triton generation from PyTorch references on KernelBench, which covers 250 tasks across three difficulty levels.It reports pass@k execution accuracy and fastp@k, which counts correct kernels exceeding a specified speedup over PyTorch; average output tokens reflect inference efficiency.Source: §3.1.2 · p. 7

3.1.3. Baselines

• The baselines include general foundation models and GPU-kernel specialists, including cudaLLM-8B and Kevin-32B for CUDA and Dr.Kernel variants for Triton.Foundation-model comparisons include Claude and Gemini series for CUDA, and DeepSeek-V4-Pro and GLM-5.2 for Triton.Source: §3.1.3 · p. 7

3.2. Main Results

• KernelZero reaches 75.8 pass@1 on CUDA Level 1 and 69.6 on Level 2, and 77.2 and 72.5 on Triton Levels 1 and 2.Table 1 covers execution accuracy and token cost; KernelZero is competitive with much larger foundation models and has the highest Triton pass@1 on Level 1.Source: §3.2 · p. 7, Table 1 · p. 8
• On Level 2, KernelZero achieves 2.4 fast1@1 and 1.2 fast2@1, exceeding both specialized baselines under the stricter fast2 criterion.On Level 1, it achieves 17.6 fast1@1 and 7.2 fast2@1; Triton outperforms CUDA on fast1@1, especially on Level 2, where CUDA kernels lack external high-performance library calls.Source: §3.2 · p. 8

3.3. Ablation Study

• Updating the Proposer raises pass@1 from 64.4% to 69.6% on CUDA Level 2 and from 68.3% to 72.5% on Triton Level 2.Level 1 pass@1 also rises from 74.9% to 75.8% on CUDA and from 74.9% to 77.2% on Triton; fast1@1 improves on both levels and backends.Source: §3.3 · pp. 8–9
• At Step 80, α=1.1 gives the highest pass@1, 74.85%, but only 9.75% fast1.5@1, while α=0.9 gives the highest fast1.5@1, 11.00%.The study tests α values on 200 combined KernelBench Level 1 and 2 tasks; α=0.5 is reported as the best overall configuration with per-operator results considered.Source: §3.3 · p. 9
• With α=0.5, the fastest correct kernel among ten generations is best across six representative operators, including a Tensor-MM improvement from 0.12× to 4.42× over SFT.The reported kernels were manually checked for complete output coverage and valid boundary handling.Source: Figure 4 · p. 7, §3.3 · p. 9

Related work

• Kernel-generation studies evaluate correctness and runtime performance on benchmarks, revealing gaps that motivate further training research [15] [20].The area has attracted growing attention for generating kernels from high-level specifications such as PyTorch modules [20] [25].Source: §4 · p. 9
• Supervised fine-tuning on curated kernel corpora shows promise but requires substantial compute, and trained models cannot be continuously improved without repeating the full pipeline.Reinforcement-learning approaches use execution feedback to optimize correctness or speedup, but typically combine them in one undifferentiated reward [6] [16] [17].Source: §4 · p. 10
• Search-based and agent-based methods improve kernel quality through iterative inference-time refinement, but significantly increase inference costs [5] [9] [24].Source: §4 · p. 10

5. Conclusion

• The Proposer generates frontier Torch modules according to the Coder’s current capability, while the Coder produces CUDA or Triton kernels with CA-GRPO.Source: §5 · p. 10
• Alternating training forms an automatic curriculum without requiring manually constructed tasks at each stage.Source: §5 · p. 10
• KernelZero-7B achieves 75.8% and 69.6% pass@1 on CUDA Levels 1 and 2, and 77.2% and 72.5% on Triton.The paper concludes that capability-aligned data generation and correctness-aware optimization can improve execution accuracy and kernel performance across programming backends.Source: §5 · p. 10

A.1. Reward Server Infrastructure

• Keck isolates candidate execution and measures performance only after the generated kernel passes numerical correctness checks.It uses separate execution paths for CUDA compilation and Triton JIT compilation, and compares each implementation with the PyTorch reference.Source: Figure 6 · p. 15, §A.1 · pp. 15–16
• Fixed random seeds can fail to produce matching parameters when reference and generated implementations invoke random operations in different orders.Keck addresses parameter consistency by transferring the reference state when the models have compatible parameters.Source: §A.1 · p. 15

B. Preliminaries: GRPO

• GRPO removes the value model and estimates advantages relative to a group of outputs, with a KL penalty to stabilize policy updates.The advantage is broadcast to the output tokens; the KL regularization strength is controlled by γ.Source: B. Preliminaries: GRPO · p. 16

C. Experiment Details

• 4,000 API combinations form the initial pool, from which 400 unique prompts are sampled for Proposer training; the Coder uses 1,224 validated modules.Figure 7 describes sampled combinations covering all 442 candidate APIs.Source: C. Experiment Details · pp. 16–17, Figure 7 · p. 17
• A 20-step Proposer update takes about 3.1 hours on 12 GPUs, while a 20-step Coder update takes about 4.5 hours on eight GPUs.The optimization jobs use NVIDIA A100-80GB GPUs, and the reported runtimes include generation, inference, validation, and policy updates.Source: C. Experiment Details · pp. 17–18

D.1. Meta Prompt

• The CUDA prompt targets an A100-80GB and asks for native kernels using strategies such as tiling, fusion, pipelining, and reordering.It prohibits PyTorch computation APIs and cuBLAS or cuDNN calls as substitutes for the kernel computation.Source: §D.1 · pp. 18–19

D.2. Prompt of Proposer

• The Proposer prompt requests a runnable PyTorch module that uses all supplied operators and defines inputs and initialization arguments.It specifies the Model class format and requests code only.Source: §D.2 · p. 19

E.1. Qualitative Evolution of Proposer Outputs

• The compared computation graphs come directly from generated modules via torch.fx, rather than manual reconstruction.The comparison covers the Base Proposer and three training checkpoints for identical API combinations.Source: §E.1 · p. 19
• Training repairs operator mismatches and shape inconsistencies in generated modules.These repairs are reported across the compared examples.Source: §E.1 · p. 19
• Training replaces mechanical operator chains with coherent branches, feature hierarchies, and task objectives.These structural changes are reported across the compared examples.Source: §E.1 · p. 19

E.1.1. Batched multiplicative fusion

• The Base module transforms two external inputs independently before BMM.Source: §E.1.1 · p. 20
• The final module derives two differently shaped branches from one feature tensor.Source: §E.1.1 · p. 20
• The final module merges those branches with BMM.Source: §E.1.1 · p. 20

Base

• The Base examples cover dropout with batch matrix multiplication, as well as linear transformations with ReLU or dropout and a matrix-product output.Source: Base · pp. 20–23

E.1.2. Multi-branch temporal convolution

• Three independently parameterized convolution branches are concatenated, replacing the Base module’s duplicated convolution result with a genuine multi-branch graph.Source: §E.1.2 · p. 23, Proposer Step 60 · p. 26

Proposer Step 20

• Proposer Step 20 applies a 1D convolution, duplicates its output by concatenation, then passes it through a linear layer.Source: Proposer Step 20 · pp. 24–25

Proposer Step 40

• Proposer Step 40 applies a 1D convolution and concatenates its output with itself along the channel dimension.Source: Proposer Step 40 · p. 25

Proposer Step 60

• Proposer Step 60 creates three independent 1D convolution outputs and concatenates them along the channel dimension.Source: Proposer Step 60 · p. 26

E.1.3. Complementary activation branches

• The module applies convolution, then concatenates positive and negative ReLU responses, replacing repeated copies of one activation with complementary branches.Source: §E.1.3 · pp. 26, 28

E.1.4. Embedding classifier with an explicit objective

• The embedding-classifier module merges an embedded feature branch and two-layer classifier with a separate external target input at the loss node.Source: §E.1.4 · p. 30
• Embedding-classifier variants look up embeddings, concatenate representations, and use linear layers with CrossEntropyLoss for classification.One implementation concatenates an embedding with itself before its linear layers and applies the loss to the resulting logits.Source: §E.1.4 · p. 34

E.1.5. Parallel volumetric pooling fusion

• The Base module omits the requested concatenation, while the final module pools one shared convolutional feature tensor through two distinct branches and concatenates their outputs.The two branches apply separate MaxPool3d operations to the same feature tensor.Source: §E.1.5 · p. 34, Base · p. 37

Base

• The 2D example uses convolution and dropout before concatenating duplicated features, whereas the 3D example applies dropout, pools along two branches, and concatenates those branch outputs.Source: Base · pp. 35–37

E.1.6. Shared trunk with two output heads

• A shared volumetric feature trunk branches into global-mean and elementwise-square-root heads, and returns both outputs explicitly.This describes the final module in the shared-trunk subsection.Source: §E.1.6 · p. 38

Proposer Step 20

• Proposer Step 20 applies a 3D convolution and channel-wise softmax, then takes a global mean, square root, and returns that result.The operations are applied in that order in the shown forward method.Source: Proposer Step 20 · p. 40

Proposer Step 40

• Proposer Step 40 defines a linear layer and shows the forward path applying square root followed by a mean over dimension 1.The excerpt specifies an output shape of (batch_size, out_features).Source: Proposer Step 40 · p. 40

Proposer Step 60

• The Proposer Step 60 excerpt reports a Conv3d intermediate with shape (128,1...30,30).Source: Proposer Step 60 · p. 41
• Proposer Step 60 composes 3D convolution, channel-wise softmax, two max-pooling operations, and returns the output mean and square root.Source: Proposer Step 60 · p. 41

F. Case Study on Generated Kernels

• The examples are selected from verified CUDA evaluation runs by reported speedup and implementation complexity, with speedups relative to PyTorch Eager on an NVIDIA A100-80GB GPU.Source: F. Case Study on Generated Kernels · p. 42

F.1.1. Diagonal Matrix Multiplication

• The diagonal-multiplication CUDA example computes C[i, j] = A[i] * B[i, j] using 32×32 output tiles and shared-memory caching of the corresponding A values.Source: §F.1.1 · pp. 42–43

F.1.2. Cross-Entropy Loss

• The cross-entropy CUDA example uses a numerically stable log-softmax and negative log-likelihood with mean reduction, assigning one block to each sample.Source: §F.1.2 · pp. 45–46
• The cross-entropy model calls a custom CUDA extension that returns a scalar mean loss, with synchronization before reading the result.The implementation divides the total by N.Source: §F.1.2 · p. 47

F.1.3. Triplet-Margin Loss

• Triplet-margin loss assigns one CUDA block to each sample, reduces squared distances, applies the margin hinge, and divides the accumulated loss by the batch size.The kernel uses p=2, eps=1e-6, no swap, and mean reduction.Source: §F.1.3 · pp. 48–49

F.1.4. KL-Divergence Loss

• The KL-divergence kernel accumulates target-weighted log differences across elements, reduces partial sums, and applies batchmean reduction.The model expects predictions and targets to be already-softmaxed probabilities.Source: §F.1.4 · pp. 51–53

F.1.5. Fused GEMM, Divide, Sum, and Scaling

• The model computes xW^T, divides by 2, sums across hidden units, and scales the resulting value for each batch row.The output has shape (batch, 1).Source: §F.1.5 · p. 53
• CUDA first computes weight-column sums, then reduces each batch row’s input-weighted sum and applies the 0.5 and scaling factors.The implementation uses a separate column-sum kernel and a per-batch fused reduction kernel.Source: §F.1.5 · p. 54

F.1.6. Fused Transposed Convolution and Global Pooling

• The module applies ConvTranspose2d, multiplies its output by a scalar, and then performs two spatial means with dimensions retained.The model constructs a ConvTranspose2d submodule and returns the resulting tensor after both mean operations.Source: §F.1.6 · pp. 55–56
• The configured example uses batch size 16, 64 input channels, 128 output channels, 128×128 inputs, a 3×3 kernel, and multiplier 0.5.Its stride is 2, padding is 1, and output_padding is 1.Source: §F.1.6 · p. 56
• One CUDA block handles each batch/output-channel pair, while threads process input spatial elements and the kernel slice for each input channel.The kernel sums the input plane and corresponding weight slice, then reduces thread contributions across the block.Source: §F.1.6 · pp. 56–57
• The kernel scales the accumulated value by multiplier/(H_out*W_out) and adds the bias contribution multiplied by the output spatial area.H_out and W_out are computed from the transposed-convolution input size, stride, padding, kernel size, and output_padding.Source: §F.1.6 · pp. 56–57