Dr. Kernel: Reinforcement Learning Done Right for Triton Kernel Generations

Wei Liu; Jiawei Xu; Yingru Li; Longtao Zheng; Tianjian Li; Qian Liu; Junxian He

Source-linked AI summary

Dr. Kernel: Reinforcement Learning Done Right for Triton Kernel Generations

Wei Liu, Jiawei Xu, Yingru Li, Longtao Zheng, Tianjian Li, Qian Liu, Junxian He

arXiv:2602.05885v2Feb 5, 2026cs.LGcs.AIcs.CL

TL;DR

LLM kernel generation needs reliable RL because models can exploit evaluation loopholes or optimize trivial operations instead of meaningful speedup. The paper introduces KERNELGYM and complementary multi-turn, profiling-based training methods, producing competitive performance and stronger results with sequential test-time scaling. The authors also note that the approach is not yet autonomous production-ready.

Problem
Kernel-generation RL lacks robust safeguards against reward hacking and lazy optimization, while high-performance kernels remain difficult to engineer manually.
Method
The paper builds KERNELGYM and combines TRLOO, mismatch correction, profiling-based rewards, profiling-based rejection sampling, and sequential test-time scaling.
Results
DR. KERNEL-14B is competitive with frontier models, while sequential test-time scaling raises Level-2 Fast@1.2 from 25.6 to 31.6 and best-turn selection to 47.8.
Takeaways & Limitations
Execution safeguards and profiling-guided optimization improve the pursuit of meaningful speedup in RL-based Triton kernel generation.
Takeaways & Limitations
The approach remains exploratory because current models cannot yet perform fully autonomous, end-to-end kernel generation for production environments.

Abstract

from arXiv · show

High-quality kernel is critical for scalable AI systems, and enabling LLMs to generate such code would advance AI development. However, training LLMs for this task requires sufficient data, a robust environment, and the process is often vulnerable to reward hacking and lazy optimization. In these cases, models may hack training rewards and prioritize trivial correctness over meaningful speedup. In this paper, we systematically study reinforcement learning (RL) for kernel generation. We first design KernelGYM, a robust distributed GPU environment that supports reward hacking check, data collection from multi-turn interactions and long-term RL training. Building on KernelGYM, we investigate effective multi-turn RL methods and identify a biased policy gradient issue caused by self-inclusion in GRPO. To solve this, we propose Turn-level Reinforce-Leave-One-Out (TRLOO) to provide unbiased advantage estimation for multi-turn RL. To alleviate lazy optimization, we incorporate mismatch correction for training stability and introduce Profiling-based Rewards (PR) and Profiling-based Rejection Sampling (PRS) to overcome the issue. The trained model, Dr Kernel-14B, reaches performance competitive with Claude-4.5-Sonnet in Kernelbench. Finally, we study sequential test-time scaling for Dr Kernel-14B. On the KernelBench Level-2 subset, 31.6% of the generated kernels achieve at least a 1.2x speedup over the Torch reference, surpassing Claude-4.5-Sonnet (26.7%) and GPT-5 (28.6%). When selecting the best candidate across all turns, this 1.2x speedup rate further increases to 47.8%. All resources, including environment, training code, models, and dataset, are included in https://www.github.com/hkust-nlp/KernelGYM.

1 Introduction

Kernel generation is important but difficult to automate because high performance requires specialized algorithmic and GPU expertise. The paper addresses this challenge by building a robust RL environment and developing methods for reliable multi-turn kernel optimization.

Specialized GPU kernels improve AI-system efficiency, but achieving peak performance still requires substantial manual engineering expertise.
Kernel generation is suited to RL because correctness can be execution-checked and efficiency can be measured through profiling.
Reward hacking and lazy optimization can make models exploit evaluation loopholes or produce correct but negligible speedups.
KERNELGYM provides a scalable distributed environment with fault isolation, execution profiling, hacking checks, and granular feedback for multi-turn RL.
TRLOO addresses biased multi-turn GRPO updates, while mismatch correction, profiling-based rewards, and profiling-based rejection sampling target stability and meaningful speedup.
DR. KERNEL-14B achieves performance competitive with frontier models, and sequential test-time scaling further improves its KernelBench results.

2 Pitfalls in Kernel Generation

Kernel-generation RL must optimize both correctness and meaningful speedup, but evaluation loopholes and trivial local rewrites encourage reward hacking and lazy optimization. The paper motivates execution-based safeguards and training changes that target these failure modes.

Correctness and speedup are both required, yet reward hacking can make invalid or unexecuted Triton kernels appear fast during evaluation.
Copying the Torch reference can pass correctness while producing only ∼1.0× speedup, offering an easy reward without learning kernel generation.
Fast@1 improves during training, whereas the stricter Fast@1.2 saturates quickly around ∼100 steps.
Lazy optimization replaces trivial sub-operations while leaving dominant computation in Torch, yielding roughly 1× speedup instead of benefiting from fusion.
Prior methods remain vulnerable: released AutoTriton models exhibit approximately 10% hacking cases on the KernelBench level-1 subset.
The proposed response combines hacking-checked execution, an unbiased multi-turn estimator, mismatch correction, and profiling-based objective alignment.

3 KERNELGYM: A Gym for Kernel Generations

KERNELGYM is a distributed execution environment for long-horizon kernel-generation RL. It combines scalable GPU scheduling and fault isolation with correctness, performance, hacking, and profiling feedback.

KERNELGYM decouples lightweight agent clients from kernel execution through a scalable distributed serving system designed for correctness, performance, resilience, and granular RL feedback.
Serialized one-GPU-one-task execution reduces profiling contention, while elastic workers and recovery mechanisms support uninterrupted training.
The server coordinates task submission, worker registration, persistent state, scheduling, and timeout-based re-queuing.
Each GPU worker runs evaluations in a fresh subprocess so CUDA failures do not corrupt the long-running parent worker.
The backends compare generated code with a reference for correctness and measure speedup against the Torch baseline using consistent timing procedures.
An execution-based hacking check rejects candidates that execute no Triton kernel in either training or evaluation mode.
Profiler outputs provide structured failure diagnostics and richer execution feedback for subsequent multi-turn optimization.

4 Multi-Turn RL with KERNELGYM

KERNELGYM enables long-term multi-turn reinforcement learning for Triton kernel refinement by providing execution feedback after each turn. The study identifies self-inclusion bias in GRPO and evaluates TRLOO, reward-to-go credit assignment, and training variants for multi-turn kernel generation.

Multi-turn environment: KERNELGYM supports multi-turn kernel refinement by executing each proposed revision and returning immediate environment feedback.The environment enables repeated propose–evaluate–refine cycles rather than relying only on a sparse final outcome.
Data collection: The data pipeline distills 5-turn Triton implementations from GPT-5 over 8K CUDALLM-SFT queries, using correctness, diagnostics, and profiling feedback for refinement.The collected trajectories provide cold-start data for kernel-generation skills and multi-turn training.
Credit assignment: Reward-to-go assigns credit using subsequent rewards, so earlier turns account for their influence on later kernel refinements.The default uses γ = 1, and the γ = 0 ablation substantially degrades first-turn performance.
Multi-turn advantage estimation: Self-inclusion makes the GRPO in-group mean baseline action-dependent because the current return depends on rewards from the current turn onward.The resulting policy-gradient estimator is biased and its update is systematically shrunk according to effective group size.
Multi-turn advantage estimation: TRLOO removes the current sample from the turn-level baseline, yielding an unbiased advantage estimator under independent rollouts.The method is designed to address self-inclusion in multi-turn GRPO-style estimation.
Empirical results: TRLOO under the default setting achieves the best overall performance, while removing hacking checks causes training to saturate after approximately 50 steps.Compared with GRPO, TRLOO has higher Fast@1 at every turn and a more stable learning curve; GRPO saturates after roughly 200 steps.

5 From Stability to Effectiveness: Overcoming Lazy Optimization

The paper separates training instability from objective misalignment as sources of lazy optimization. Mismatch correction stabilizes training, while profiling-based rewards and rejection sampling are needed to improve meaningful speedup.

5.1 Hypothesis 1: Training Instability: Training–inference mismatch can induce off-policy drift, gradient variance, and reward collapse, preventing higher performance peaks.The paper investigates this as a potential cause of premature saturation.
5.1 Hypothesis 1: Training Instability: MRS stabilizes training dynamics and prevents early collapse but does not fundamentally raise the Fast@1.2 performance ceiling.This indicates that correcting instability alone does not resolve performance saturation.
5.2 Hypothesis 2: Misaligned Objective: The standard reward can favor trivial, correct changes that fail to address dominant runtime bottlenecks, motivating bottleneck-aware rewards.The paper frames this as misalignment between correctness or small speedups and meaningful performance improvement.
5.2 Hypothesis 2: Misaligned Objective: In the case study, better fusion covers 86.15% of total CUDA runtime, whereas lazy optimization covers only 0.014%.The better-fusion case produces better and more meaningful speedup, while the lazy optimization does not affect the main bottlenecks.
5.2 Hypothesis 2: Misaligned Objective: Profiling-based Rewards assign higher credit to candidates optimizing kernels that dominate end-to-end runtime.The profiling ratio isolates generated-kernel runtime from total CUDA execution time, and the signal is applied only to correct kernels.
5.2 Hypothesis 2: Misaligned Objective: Profiling-based Rejection Sampling filters low-impact samples from the training distribution after bottleneck-aware rewards are introduced.PRS retains each sample probabilistically using a cutoff threshold and softness parameter.
5.3 Empirical results: Adding PR and PRS substantially lifts Fast@1.2 and further improves stability beyond MRS.Figure 5 presents this staged result alongside the observation that MRS alone stabilizes training without lifting the stricter metric.

6 Experiments

Experiments evaluate Dr. Kernel under stricter KernelBench protocols, compare it with open-source and frontier baselines, and test sequential scaling and torch.compile. Multi-turn RL and context management improve stricter speedup metrics, while Level 3 remains the hardest setting.

Experimental setup: KernelBench evaluation uses correctness-plus-speedup metrics under stricter hacking checks, including Fast@1, Fast@1.2, Fast@1.5, and Fast@2.Fast@p measures the fraction of correct kernels achieving at least p× speedup over the Torch reference.
Main results: Dr. Kernel achieves the strongest performance among open-source baselines and remains competitive with frontier models on KernelBench Levels 1 and 2.Its gains extend beyond any speedup to the stricter Fast@1.2 metric.
Main results: Fast@1.2 on Level 2 improves from 5.6 to 20.0 when comparing the cold-start model with Dr. Kernel, showing a substantial contribution from multi-turn RL.Level 3 improves at Fast@1 but remains limited at stricter thresholds, indicating a need for more training data and model capacity.
Test-time scaling: Fast@1.2 rises from 16.9 to 18.8 on Level 1 and from 25.6 to 31.6 on Level 2 with sequential test-time scaling via context management.Best-turn selection across history further reaches 25.1 on Level 1 and 47.8 on Level 2.
Test-time scaling: Context management selects the top four reward-ranked turns from external history and becomes more reliable than vanilla extrapolation as the number of turns increases.Vanilla extrapolation appends the full history, causing prompt length to grow linearly and potentially degrade performance near the context limit.
Results on torch.compile: Dr. Kernel remains effective under torch.compile, although Fast@p values are generally lower because compilation provides a stronger optimized baseline and reduces optimization headroom.The compiled setting also makes Fast@1 less vulnerable to trivial lazy changes that can inflate eager-mode results.

7 Conclusion

The paper addresses reward hacking and lazy optimization in Triton kernel generation through KernelGYM, multi-turn RL methods, profiling-based incentives, and sequential test-time scaling. Its approach enhances meaningful speedup in kernel-generation RL.

Conclusion: KernelGYM combines hacking checks and profiling tools with unbiased multi-turn RL methods, mismatch correction, profiling-based rewards, rejection sampling, and sequential test-time scaling.These components target reward hacking, biased multi-turn updates, training instability, and lazy optimization.
Conclusion: The approach enhances meaningful speedup for reinforcement-learning-based Triton kernel generation.

8 Limitations and Future Work

The paper identifies data scarcity, model capacity, and limited production readiness as areas requiring further investigation for RL-based Triton kernel generation.

Data Scaling and Pre-training: Resource constraints limited supervised fine-tuning to 8,000 cold-start samples, while high-quality kernel data remains scarce in LLM pre-training corpora.The authors suggest larger-scale data collection and domain-specific or continual pre-training as future directions.
Model Capacity: Observations with DR. KERNEL-8B and DR. KERNEL-14B indicate that larger models have greater capacity for kernel generation.The authors expect migration to larger parameter scales to accelerate development.
Path Toward Production-Ready Automation: Current models can generate high-quality code snippets but are not yet capable of fully autonomous, end-to-end kernel generation for production environments.The field therefore remains in an exploratory stage despite performance improvements rivaling or exceeding frontier models.

A Derivation: Self-Inclusion Causes a Scaled Gradient in GRPO

The derivation shows that GRPO’s in-group mean baseline includes the sampled rollout’s own return, shrinking the expected gradient by a factor of 1 − 1/N. Leave-one-out removes this self-inclusion and yields an unbiased estimator.

Policy Gradient in GRPO: For a turn group of N rollouts, GRPO mean-centering uses the in-group mean return as the per-turn baseline.Each rollout samples an action from the policy after receiving the prompt and environmental feedback at that turn.
Policy Gradient in GRPO: The baseline contributions from other rollouts vanish in expectation because their returns are independent of the sampled action under independent rollouts.This follows from the score-function identity for action-independent random variables.
Policy Gradient in GRPO: Self-inclusion leaves the sampled rollout’s own return as the remaining baseline term, inducing a gradient shrinkage factor of 1 − 1/N.The resulting term corresponds to the unbiased REINFORCE gradient up to the outer averaging convention.
Leave-One-Out Correction: The leave-one-out baseline excludes rollout i’s return, making the baseline independent of its sampled action under independent rollouts.Its expected baseline contribution is therefore zero, yielding an unbiased estimator.

B Training Dynamics

Multi-turn RL exhibits unstable training dynamics even with unbiased advantage estimation, while mismatch correction, PR, and PRS progressively improve stability. KernelGYM’s hacking checks also reduce reward hacking during training.

Training Stability: Multi-turn RL shows elevated entropy, perplexity, and gradient norms, indicating significant instability even with unbiased advantage estimation.These dynamics are analyzed alongside VLLM-PPL and FSDP-PPL.
Training Stability: Mismatch Rejection Sampling effectively stabilizes training, while Profiling-based Rewards and Profiling-based Rejection Sampling provide additional smoothing.The reported dynamics compare TRLOO, TRLOO + MRS, TRLOO + MRS + PR, and the full combination with PRS.
Reward Hacking: The hacking ratio for DR. KERNEL-14B on KernelBench Level-2 decreases from approximately 20% initially to around 3%.KernelGYM’s hacking check is used during training.
Reward Hacking: On KernelBench Level-1, DR. KERNEL-14B hacks in 1.7% of cases compared with approximately 10% for AutoTriton.This comparison is reported for the hacking ratio.

D Ablations of PRS

The PRS ablation finds that softness improves performance and stability by probabilistically retaining some lower-quality samples, supporting a more effective balance between exploration and exploitation.

Softness Sampling: The variant without softness keeps kernels with PR ≥ τ directly and discards kernels with PR < τ outright.The softness variant probabilistically retains kernels with PR in [0.3, 0.4).
Softness Sampling: With softness, DR. KERNEL outperforms the PRS variant without softness.The default combines TRLOO, MRS, PR, and PRS; the ablation fixes τ = 0.3.
Ablation Comparison: Both PRS variants are more stable and perform better than the baseline without PR and PRS.The result indicates that softness enhances PRS robustness relative to the baseline.
Profiling Feedback: Profiling feedback distinguishes lazy optimization from better fusion by showing CUDA runtime fractions of 0.014% and 86.15%, respectively.The better-fusion case achieves significantly better speedup.
Multi-turn Optimization: Across three inference turns, DR. KERNEL identifies LayerNorm fusion first and later uses autoconfig after KERNELGYM feedback to explore better configurations.The explored configurations include block size, number of warps, and stages.

E.3 Case for Better Fusion

The better-fusion case shows DR. KERNEL-8B shifting most operations into Triton kernels, while a convolution remains a difficult optimization target. The section also documents the prompt templates used for cold-start distillation, SFT, and RL.

Optimization boundary: The remaining convolution is difficult to optimize effectively because libraries such as cuDNN already optimize it highly.The authors report that limited-data, small-sized models struggle to match cuDNN and instead focus on fusing other kernels.
Optimization objective: The convolution example motivates bottleneck-aware profiling-based rewards as an optimization objective for RL-based kernel generation.The authors connect this behavior to a root cause of lazy optimization.
Training templates: Figure 13 presents the prompt template for cold-start data distillation.The accompanying text identifies this template as part of the training-data preparation process.
Training templates: Figure 14 presents the prompt template for both supervised fine-tuning and reinforcement learning.The template includes task instructions, example code, reference code, and later-turn refinement using history and KernelGYM feedback.