Dr. Kernel : Un Renforcement de l'Apprentissage Bien Conçu pour la Génération de Noyaux Triton

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

La qualité des noyaux de calcul est cruciale pour les systèmes d'IA évolutifs, et permettre aux LLMs de générer un tel code ferait progresser le développement de l'IA. Cependant, l'entraînement des LLMs pour cette tâche nécessite des données suffisantes, un environnement robuste, et le processus est souvent vulnérable au détournement des récompenses (reward hacking) et à l'optimisation paresseuse (lazy optimization). Dans ces cas, les modèles peuvent détourner les récompenses d'entraînement et privilégier une exactitude triviale plutôt qu'une accélération significative. Dans cet article, nous étudions systématiquement l'apprentissage par renforcement (RL) pour la génération de noyaux. Nous concevons d'abord KernelGYM, un environnement distribué robuste sur GPU qui prend en charge la détection du détournement des récompenses, la collecte de données à partir d'interactions multi-tours et l'entraînement RL à long terme. En nous appuyant sur KernelGYM, nous étudions des méthodes RL multi-tours efficaces et identifions un problème de biais dans le gradient de politique causé par l'auto-inclusion dans GRPO. Pour résoudre ce problème, nous proposons Turn-level Reinforce-Leave-One-Out (TRLOO) pour fournir une estimation non biaisée de l'avantage pour le RL multi-tours. Pour atténuer l'optimisation paresseuse, nous intégrons une correction de décalage (mismatch correction) pour la stabilité de l'entraînement et introduisons les Récompenses Basées sur le Profilage (PR) et l'Échantillonnage par Rejet Basé sur le Profilage (PRS) pour surmonter le problème. Le modèle entraîné, Dr.Kernel-14B, atteint des performances compétitives avec Claude-4.5-Sonnet sur KernelBench. Enfin, nous étudions la mise à l'échelle séquentielle en phase de test (sequential test-time scaling) pour Dr.Kernel-14B. Sur le sous-ensemble KernelBench Level-2, 31,6 % des noyaux générés atteignent au moins une accélération de 1,2x par rapport à la référence Torch, surpassant Claude-4.5-Sonnet (26,7 %) et GPT-5 (28,6 %). En sélectionnant le meilleur candidat parmi tous les tours, ce taux d'accélération de 1,2x augmente encore pour atteindre 47,8 %. Toutes les ressources, y compris l'environnement, le code d'entraînement, les modèles et le jeu de données, sont disponibles sur 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.