Dr. Kernel: Reinforcement Learning richtig gemacht für Triton-Kernel-Generationen
Dr. Kernel: Reinforcement Learning Done Right for Triton Kernel Generations
February 5, 2026
papers.authors: Wei Liu, Jiawei Xu, Yingru Li, Longtao Zheng, Tianjian Li, Qian Liu, Junxian He
cs.AI
papers.abstract
Hochwertige Kernel sind entscheidend für skalierbare KI-Systeme, und die Fähigkeit von LLMs, solchen Code zu generieren, würde die KI-Entwicklung vorantreiben. Das Training von LLMs für diese Aufgabe erfordert jedoch ausreichend Daten, eine robuste Umgebung und der Prozess ist oft anfällig für Reward Hacking und faule Optimierung. In diesen Fällen können Modelle Trainings-Belohnungen manipulieren und triviale Korrektheit über bedeutende Geschwindigkeitssteigerungen stellen. In dieser Arbeit untersuchen wir systematisch Reinforcement Learning (RL) für Kernel-Generierung. Zunächst entwerfen wir KernelGYM, eine robuste verteilte GPU-Umgebung, die Reward-Hacking-Prüfung, Datensammlung aus Mehrfach-Interaktionen und langfristiges RL-Training unterstützt. Aufbauend auf KernelGYM untersuchen wir effektive Mehrfach-RL-Methoden und identifizieren ein Problem mit verzerrtem Policy Gradient durch Selbst-Einbeziehung in GRPO. Zur Lösung schlagen wir Turn-level Reinforce-Leave-One-Out (TRLOO) vor, um unverzerrte Advantage-Schätzung für Mehrfach-RL zu ermöglichen. Um faule Optimierung zu mildern, integrieren wir Mismatch-Korrektur für Trainingsstabilität und führen Profiling-basierte Belohnungen (PR) sowie Profiling-basiertes Rejection Sampling (PRS) ein, um das Problem zu überwinden. Das trainierte Modell Dr.Kernel-14B erreicht im KernelBench wettbewerbsfähige Leistung zu Claude-4.5-Sonnet. Abschließend untersuchen wir sequenzielles Test-Time-Scaling für Dr.Kernel-14B. Im KernelBench Level-2-Subset erzielen 31,6% der generierten Kernel mindestens eine 1,2-fache Beschleunigung gegenüber der Torch-Referenz, was Claude-4.5-Sonnet (26,7%) und GPT-5 (28,6%) übertrifft. Bei Auswahl des besten Kandidaten über alle Durchgänge steigt diese 1,2-fache Beschleunigungsrate weiter auf 47,8%. Alle Ressourcen, inklusive Umgebung, Trainingscode, Modelle und Datensatz, sind unter https://www.github.com/hkust-nlp/KernelGYM verfügbar.
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.