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

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 or 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, reach performance competitive with Claude-4.5-Sonnet in Kernelbench. Finally, we study sequential test-time scaling for Dr. Kernel-14B, which even outperforms GPT-5 and Claude-4.5-Sonnet in the Kernelbench level-2 subset.

Modern AI systems depend on GPUs to run quickly, but getting the best performance often requires experts to hand-write small, highly tuned pieces of code. This paper studies how to train AI models to write and improve this kind of GPU code automatically. A key challenge is that models can find shortcuts: they may produce code that looks fast in a test but does not really do the work, or they may make only small, safe changes that miss the main source of slowness. We build KernelGYM, a testing environment that safely runs generated code, checks whether it is truly being used, measures its speed, and gives detailed feedback for improvement. Using this environment, we design a training method that lets the model revise its code over several rounds while learning fairly from each attempt. We also use timing information to steer the model toward changes that affect the main bottlenecks, not just easy minor operations. The resulting system, Dr. Kernel, produces correct GPU code with meaningful speedups and performs competitively with leading general-purpose AI models on KernelBench benchmarks. This suggests that carefully designed feedback and training environments can help AI systems become more useful assistants for building faster machine-learning software.