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 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.
698580004ad556f294b7ec78
Wei Liu
698580004ad556f294b7ec79
Jiawei Xu
698580004ad556f294b7ec7a
Yingru Li
claimed_verified
2026-02-06T18:50:48.867Z
67277d20eebb94a257cd6925
/avatars/c1f714b59fecb53770e28920c0c267cb.svg
Yingru Li
R1ch0rd
user
R1ch0rd
698580004ad556f294b7ec7b
Longtao Zheng
claimed_verified
2026-02-11T11:20:14.489Z
63db5dc49f2687298a1547bf
https://cdn-avatars.huggingface.co/v1/production/uploads/63db5dc49f2687298a1547bf/xVFi0kRkYud191cQgma16.jpeg
Longtao Zheng
ltzheng
user
ltzheng
698580004ad556f294b7ec7c
Tianjian Li
698580004ad556f294b7ec7d
Qian Liu
698580004ad556f294b7ec7e
Junxian He
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 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.
Reinforcement learning approach for kernel generation addresses reward hacking and optimization issues through specialized environment and unbiased policy gradient methods, achieving competitive performance with state-of-the-art models.
reinforcement learning kernel generation reward hacking policy gradient GRPO Turn-level Reinforce-Leave-One-Out multi-turn RL training stability profiling-based rewards profiling-based rejection sampling sequential test-time scaling