Dr. Kernel: Reinforcement Learning Goed Toegepast voor Triton Kernel Generaties
Dr. Kernel: Reinforcement Learning Done Right for Triton Kernel Generations
February 5, 2026
Auteurs: Wei Liu, Jiawei Xu, Yingru Li, Longtao Zheng, Tianjian Li, Qian Liu, Junxian He
cs.AI
Samenvatting
Een hoogwaardige kernel is cruciaal voor schaalbare AI-systemen, en de mogelijkheid om LLM's dergelijke code te laten genereren zou de AI-ontwikkeling vooruithelpen. Het trainen van LLM's voor deze taak vereist echter voldoende data, een robuuste omgeving, en het proces is vaak gevoelig voor reward hacking en luie optimalisatie. In dergelijke gevallen kunnen modellen trainingsbeloningen manipuleren en triviale correctheid prioriteren boven betekenisvolle versnelling. In dit artikel bestuderen we systematisch reinforcement learning (RL) voor kernelgeneratie. We ontwerpen eerst KernelGYM, een robuuste gedistribueerde GPU-omgeving die controle op reward hacking ondersteunt, datacollectie vanuit multi-turn interacties en langdurige RL-training mogelijk maakt. Voortbouwend op KernelGYM onderzoeken we effectieve multi-turn RL-methoden en identificeren we een bevooroordeeld policy gradient-probleem veroorzaakt door zelf-inclusie in GRPO. Om dit op te lossen, stellen we Turn-level Reinforce-Leave-One-Out (TRLOO) voor om een zuivere advantage-schatting te bieden voor multi-turn RL. Om luie optimalisatie tegen te gaan, integreren we mismatch-correctie voor trainigsstabiliteit en introduceren we Profiling-based Rewards (PR) en Profiling-based Rejection Sampling (PRS) om het probleem te overwinnen. Het getrainde model, Dr.Kernel-14B, bereikt een prestatieniveau dat competitief is met Claude-4.5-Sonnet in Kernelbench. Ten slotte bestuderen we sequentiële test-time scaling voor Dr.Kernel-14B. Op de KernelBench Level-2 subset behaalt 31,6% van de gegenereerde kernels minstens een 1,2x versnelling ten opzichte van de Torch-referentie, wat Claude-4.5-Sonnet (26,7%) en GPT-5 (28,6%) overtreft. Bij selectie van de beste kandidaat over alle turns stijgt dit 1,2x versnellingspercentage verder tot 47,8%. Alle bronnen, inclusief de omgeving, trainingscode, modellen en dataset, zijn beschikbaar op https://www.github.com/hkust-nlp/KernelGYM.
English
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.