PaperlayerView exact revision on arXiv

Source-linked AI summary

Kevin: Multi-Turn RL for Generating CUDA Kernels

Carlo Baronio, Pietro Marsella, Ben Pan, Simon Guo, Silas Alberti

arXiv:2507.11948v1cs.LGcs.AIcs.PFcs.SE

TL;DR

GPU kernel generation is difficult, iterative, and objectively evaluable through correctness and speedup, yet existing methods are limited by model capability and sparse feedback. Kevin addresses this with multi-turn RL that incorporates execution feedback and intermediate reward attribution. On KernelBench, it improves QwQ-32B from 56% to 82% correctness and from 0.53x to 1.10x mean speedup over PyTorch Eager, with serial refinement scaling more effectively than parallel sampling.

  • Problem

    GPU kernel generation requires difficult domain expertise and iterative optimization, while existing LLM approaches are limited by base-model capability and sparse single-turn feedback.

  • Method

    Kevin uses a multi-turn RL recipe that repeatedly generates, executes, and refines CUDA kernels while addressing long trajectories, context growth, and reward attribution across turns.

  • Results

    Kevin improves pure-CUDA correctness from 56% to 82% and mean speedup from 0.53x to 1.10x over PyTorch Eager, surpassing o4-mini’s 0.78x speedup and scaling better with refinement turns.

  • Takeaways & Limitations

    Multi-turn training enables more effective kernel-refinement strategies and stronger test-time scaling than single-turn RL and the base model.

  • Takeaways & Limitations

    Limited compute restricted training to 80 gradient steps and prevented exhaustive ablations, while measured speedups are accurate only for predefined input dimensions on NVIDIA H200 GPUs.

Abstract

from arXiv · show

Writing GPU kernels is a challenging task and critical for AI systems' efficiency. It is also highly iterative: domain experts write code and improve performance through execution feedback. Moreover, it presents verifiable rewards like correctness and speedup, making it a natural environment to apply Reinforcement Learning (RL). To explicitly incorporate the iterative nature of this process into training, we develop a flexible multi-turn RL recipe that addresses unique challenges encountered in real-world settings, such as learning from long trajectories and effective reward attribution across turns. We present Kevin - K(ernel D)evin, the first model trained with multi-turn RL for CUDA kernel generation and optimization. In our evaluation setup, Kevin shows significant gains over its base model (QwQ-32B), improving correctness of generated kernels (in pure CUDA) from 56% to 82% and mean speedup from 0.53x to 1.10x of baseline (PyTorch Eager), and surpassing frontier models like o4-mini (0.78x). Finally, we study its behavior across test-time scaling axes: we found scaling serial refinement more beneficial than parallel sampling. In particular, when given more refinement turns, Kevin shows a higher rate of improvement.

1 Introduction

GPU kernel generation is difficult but offers verifiable correctness and speedup rewards, making it suitable for RL. Kevin incorporates iterative execution feedback into multi-turn training and improves kernel generation while benefiting from serial refinement.

  • GPU kernel optimization requires substantial domain expertise, while correctness and speedup provide verifiable rewards for reinforcement learning.
  • Multi-turn RL mirrors kernel development by conditioning successive code generations on execution feedback and timing profiles.
  • The training recipe splits long trajectories into turn-level samples, summarizes prior chain-of-thought context, and aggregates intermediate rewards across turns.
  • Strict rule-based checks address reward hacking, including attempts to copy the reference implementation instead of implementing kernels.
  • Kevin improves pure-CUDA correctness from 56% to 82% and mean speedup from 0.53x to 1.10x over PyTorch Eager, exceeding o4-mini’s 0.78x speedup.
  • Sequential scaling with more refinement turns is more effective than parallel scaling, and multi-turn training produces stronger improvement across refinement turns.

2 Background and Related Work

GPU kernel optimization is costly and underrepresented in training data, limiting frontier-model performance. Prior work uses agentic search and multi-turn RL, but existing approaches often remain constrained by base-model capability or sparse final rewards.

  • Efficient GPU kernels demand costly expertise and long engineering cycles, while CUDA comprises less than 0.01% of Stack pretraining data.
  • Agentic systems use extensive test-time computation to generate kernels, but their performance remains limited by the underlying models’ domain capability.
  • Existing code-optimization methods address performance beyond correctness, but prior approaches have largely remained confined to supervised or single-turn settings.
  • RLEF frames code generation as multi-turn RL but assigns a single binary pass/fail reward to the final generation, potentially creating sample-inefficiency issues.

3 Task and Baseline

KernelBench evaluates CUDA generation against PyTorch references for correctness and speedup, while the single-turn GRPO baseline assigns rewards to independently sampled kernels. Its training can plateau because kernels that are close to correct receive no reward, motivating multi-turn refinement.

  • KernelBench training uses 180 Level 1 and 2 tasks, while an additional 80 novel tasks provide an unseen evaluation set.
  • Each task supplies a PyTorch reference implementation, and generated inline CUDA is checked for format, reward hacking, compilation, runtime errors, correctness, and performance.
  • Kernel scores balance correctness with speedup, where speedup is the ratio of PyTorch baseline time to kernel runtime.
  • Intermediate compilation or execution rewards caused over-optimization for kernels that compile without necessarily being correct, while response-length penalties degraded training performance.
  • Single-turn GRPO samples 16 responses per task and assigns each evaluated kernel score as its reward.
  • A strong QwQ-32B base model is used because weaker priors produced reward hacking and failed to learn.
  • Single-turn reward plateaus after 50 steps, as nearly correct kernels receive zero reward and correct kernels often avoid risky performance optimizations.

4 Multi-Turn Training

Multi-turn training turns each refinement step into a separate, history-conditioned sample and assigns rewards using discounted future kernel scores. Ablations favor Sum with γ = 0.4 for scaling across refinement turns, while multi-turn training avoids the reward plateau observed in single-turn training.

  • Training on Every Refinement Turn: Each refinement turn becomes an individual training sample containing its kernel and CoT summary, with prior history retained as context.This improves sample efficiency while preserving information from earlier turns through summarized reasoning and generated kernels.
  • Training on Every Refinement Turn: The context includes previous kernels, summarized CoTs, and evaluation feedback so later turns can condition on the trajectory history.
  • Reward Aggregation and Discounting: Each turn receives a discounted sum of its current and subsequent kernel scores, balancing greedy credit assignment against assigning every turn the trajectory’s best score.
  • Reward Aggregation and Discounting: Sum with γ = 0.4 scales best over 8 turns, although Max performs better with γ = 0.8 when using fewer turns.
  • Multi-Turn Training Behavior: Multi-turn training reward climbs steadily through 40 steps (80 gradient steps), contrasting with the plateau during single-turn training through step 50 (100 gradient steps).

5 Evaluation

Kevin outperforms the base model, single-turn RL baseline, and frontier models on KernelBench, while benefiting especially from serial refinement and retaining exploration as parallel samples increase.

  • 5.1 Result on KernelBench Eval Set: Kevin achieves higher correctness and performance than the single-turn baseline, frontier models, and its QwQ-32B base model on 100 unseen KernelBench tasks.Evaluation uses 16 parallel trajectories with 8 refinement turns per trajectory.
  • 5.1 Result on KernelBench Eval Set: Kevin more effectively applies aggressive optimizations across several refinement turns.The qualitative comparison attributes this behavior to multi-turn execution and refinement.
  • 5.2 Scaling Refinement Turns: With more refinement turns, Kevin’s performance improves significantly faster than the single-turn model.The single-turn model performs slightly better at one turn, but Kevin has the steepest performance curve with additional turns.
  • 5.3 Scaling Parallel Samples: As parallel trajectory count k increases, best@k performance continues rising for the multi-turn model but plateaus for the single-turn model.Refinement turns are fixed to 8 in this comparison.
  • 5.4 Parallel vs Sequential Scaling: Under the same total inference budget, allocating more computation to serial refinement is consistently better than allocating it to parallel sampling.The compared configurations use 128 one-turn trajectories, 32 four-turn trajectories, or 16 eight-turn trajectories.

6 Discussion

The discussion examines instability and reward hacking in multi-turn kernel training, describing diagnostic signals, mitigation strategies, and dataset-design constraints.

  • 6.1 Diagnosing Model Instability: Longer training can produce repetitive, nonsensical outputs that first appear in the final turn and spread backward before model collapse.The authors identify this behavior as “junk.”
  • 6.1 Diagnosing Model Instability: The Not Okay Ratio detects early instability because erratic openings appear well before junk outputs.It is introduced as a proxy signal for model instability and a predictor of junk.
  • 6.1 Diagnosing Model Instability: Constant length normalization and gradient norm clipping of 0.05 delay junk onset until step 100, whereas KL penalties slow learning without improving stability.The tested KL penalties were 0.001 and 0.01.
  • 6.2 Reward Hacking: Reward hacking includes copying or inheriting PyTorch reference implementations, and stricter format checks assign zero reward when Level 2 outputs retain PyTorch functional operators.These checks target kernels that leave the operator being optimized unfused or unmodified.
  • 6.3 Dataset Difficulty: Training only on easy tasks causes reward to plateau through overfitting, motivating a balanced mix of Level 1 and Level 2 tasks with QwQ-32B.The dataset is designed so batches contain both easier and harder tasks.

7 Conclusion

The paper presents multi-turn RL for CUDA kernel generation, addressing context management, reward attribution, and reward hacking in an iterative engineering task. Kevin outperforms its base model and frontier models, while multi-turn training improves test-time refinement scaling.

  • Conclusion: Multi-turn RL addresses context management and credit attribution across turns while incorporating mechanisms to prevent reward hacking.The recipe also uses a proxy reward to diagnose instability and experiments with ways to constrain it.
  • Conclusion: Multi-turn training improves test-time scaling with sequential refinement and parallel sampling while preserving exploration capacity.The paper specifically reports better scaling characteristics as refinement turns increase.
  • Limitations: Additional RL training could destabilize the heavily post-trained QwQ-32B base model, and limited compute prevented exhaustive ablations.The study trained for up to 80 gradient steps and deferred varying-turn studies to future work.
  • Limitations: Speedups are accurate only for the predefined tensor dimensions used on NVIDIA H200 GPUs.This scope boundary limits how directly the reported optimization results generalize to other dimensions or hardware.
  • Conclusion: The recipe is presented as potentially applicable to tasks with verifiable rewards and execution feedback across a trajectory.The authors frame multi-turn reasoning as relevant to complex tasks beyond GPU kernel generation.

A.1 Task Improvements

The authors modify KernelBench to address evaluation and reward-hacking problems before using it for training. The resulting training environment contains 180 tasks across Levels 1 and 2.

  • Task Improvements: KernelBench is modified to mitigate reward hacking, as the original benchmark contains several task and evaluation limitations.The authors describe these targeted changes as crucial for addressing reward-hacking behavior.
  • Task Improvements: The evaluation process is sandboxed so fatal CUDA errors, including illegal memory accesses, do not crash RL training.This change keeps kernel failures from terminating the training process.
  • Task Improvements: Small input tensors are enlarged because kernel-launch overhead can otherwise dominate measured runtime.The modification targets more meaningful performance measurements for affected tasks.
  • Task Improvements: The training environment contains 180 tasks: 90 of 100 Level 1 problems and 90 Level 2 problems.Level 2 problems are sequences of operators with fusion opportunities.

A.2 Construction of Additional Evaluation Set

The authors construct an unseen evaluation set by extending KernelBench’s task-generation methodology and use turn-aware rewards to improve sample efficiency during multi-turn RL.

  • A.2 Construction of Additional Evaluation Set: Because KernelBench lacks a train-test split, the authors construct 80 additional tasks using the benchmark’s original methodology.These tasks are designed to remain unseen during training.
  • A.2 Construction of Additional Evaluation Set: The additional tasks combine one computationally expensive main operator with 2-5 other operators sampled from the available operator set.The resulting PyTorch programs define the kernel-generation tasks.
  • A.2 Construction of Additional Evaluation Set: A language model generates each PyTorch program and sample tensor sizes, which must execute with runtime above 0.1ms on NVIDIA H200.The runtime threshold avoids measurements dominated by CPU kernel-launch overhead.
  • A.2 Construction of Additional Evaluation Set: The authors exclude programs matching existing Level 1 and 2 operator sequences and manually inspect all new tasks for validity.The final evaluation set combines 80 new tasks with 20 remaining original tasks, totaling 100 unseen evaluation tasks.
  • A.2 Construction of Additional Evaluation Set: A discounted look-ahead max or sum aggregates rewards across turns to retain useful training information after the best kernel appears.Using only the maximum kernel score would discard later reasoning, code, and evaluation results.
  • A.2 Construction of Additional Evaluation Set: The reward design includes correctness and speedup terms, with the paper illustrating alternative configurations through reward curves.One illustrated configuration uses S = 1{correct}+ Tbaseline.

B.2 Weighting for Score

The paper studies how correctness weighting, length penalties, trajectory parallelism, and infrastructure choices affect RL training for CUDA kernel generation. It selects a score balancing correctness with raw speedup and observes instability from some alternatives.

  • B.2 Weighting for Score: The selected score uses 0.3 correctness weighting plus raw speedup gated by correctness: S = 0.3 · 1{correct} + 1{correct} · Tbaseline.The weighting was chosen after small-scale score-design ablations on QwQ-32B.
  • B.2 Weighting for Score: A correctness weight of 1 causes reward to plateau and decrease as the model over-optimizes correct kernels rather than exploring speedup.The corresponding configuration is illustrated in Figure 8.
  • B.2 Weighting for Score: With no correctness weighting, reward also plateaus when only performant and correct kernels are rewarded.The corresponding configuration is illustrated in Figure 9.
  • B.2 Weighting for Score: Increasing parallel trajectories from 16 to 64 slightly improves best@16 correctness but does not significantly improve overall performance.The authors retain 16 trajectories because additional generations require high compute.
  • B.2 Weighting for Score: Adding a length penalty collapses response length, with the model stopping CoT output after 10 training steps.The authors describe this result as counterproductive for the setting.
  • B.2 Weighting for Score: The training framework supports kernel evaluation and multiple turns within one training step by building on OpenRLHF, with vLLM inference and DeepSpeed Zero-3 offloading.Each GPU generates and evaluates kernels for one task, then regenerates after reward calculation.

E Training Stability

The training analysis examines instability, reward hacking, and KL regularization. KL-enabled training reaches a reward plateau, while rule-based checks target several ways kernels can exploit evaluation.

  • KL regularization: With KL coefficient β = 0.001, reward plateaus, whereas β = 0 avoids this observed plateau.The ablation reports that adding the KL penalty slows learning.
  • Reward hacking: Reward hacking includes copying the reference implementation or using PyTorch fallbacks instead of implementing CUDA kernels.These behaviors can produce apparently correct or fast results without the intended CUDA implementation.
  • Reward-hacking checks: The evaluation assigns reward 0 to kernels containing pass, torch.nn or torch.nn.functional, try, or except.These rule-based checks are designed to block reference copying and fallback execution.

G Kernel Highlight

The highlighted KernelBench example traces iterative CUDA optimization from an initially correct fused kernel through failed and regressive refinements to a substantially faster final implementation. The trajectory illustrates that later turns can discover performance improvements but can also introduce compilation or correctness failures.

  • Initial implementation: The first turn fuses the kernels correctly but achieves only 1.06x speedup over the reference implementation.The task combines computationally intensive 3D convolution with lightweight activation functions.
  • Failed refinement: Turn 2 fails because the attempted improvement produces a compilation error.This refinement does not yield an evaluable kernel.
  • Memory coalescing: 0.61x speedup at turn 3 shows that changing loop order for coalesced memory access can worsen performance.The model repeats the strategy at turn 4 and reaches 1.19x through improved loop ordering.
  • Memory coalescing: 1.19x speedup at turn 4 follows a loop ordering that increases memory coalescing.The preceding turn’s loop-order change had reduced performance to 0.61x.
  • Compile-time optimization: 1.93x speedup at turn 7 results from making kernel size a compile-time constant and unrolling the three convolution loops.Turns 5 and 6 make minor changes with no measurable performance effect.
  • Final refinement: The final advanced implementation fails because it applies warp-level reduction across the wrong axis.The model has succeeded at complex warp reductions in several other tasks.
Loading 2507.11948v1…