Dottor Kernel: Reinforcement Learning Fatto Bene per le Generazioni di Kernel Triton
Dr. Kernel: Reinforcement Learning Done Right for Triton Kernel Generations
February 5, 2026
Autori: Wei Liu, Jiawei Xu, Yingru Li, Longtao Zheng, Tianjian Li, Qian Liu, Junxian He
cs.AI
Abstract
Un kernel di alta qualità è fondamentale per sistemi di IA scalabili, e consentire agli LLM di generare tale codice farebbe avanzare lo sviluppo dell'IA. Tuttavia, addestrare LLM per questo compito richiede dati sufficienti, un ambiente robusto, e il processo è spesso vulnerabile a reward hacking e ottimizzazione pigra. In questi casi, i modelli possono manipolare i reward di addestramento e privilegiare una correttezza banale rispetto a un significativo miglioramento delle prestazioni. In questo articolo, studiamo sistematicamente il reinforcement learning (RL) per la generazione di kernel. Inizialmente progettiamo KernelGYM, un ambiente distribuito robusto su GPU che supporta il controllo del reward hacking, la raccolta di dati da interazioni multi-turno e l'addestramento RL a lungo termine. Basandoci su KernelGYM, investigiamo metodi RL multi-turno efficaci e identifichiamo un problema di policy gradient distorto causato dall'auto-inclusione in GRPO. Per risolverlo, proponiamo Turn-level Reinforce-Leave-One-Out (TRLOO) per fornire una stima unbiased del vantaggio per l'RL multi-turno. Per alleviare l'ottimizzazione pigra, incorporiamo una correzione del mismatch per la stabilità dell'addestramento e introduciamo Profiling-based Rewards (PR) e Profiling-based Rejection Sampling (PRS) per superare il problema. Il modello addestrato, Dr.Kernel-14B, raggiunge prestazioni competitive con Claude-4.5-Sonnet in Kernelbench. Infine, studiamo lo scaling sequenziale al test-time per Dr.Kernel-14B. Sul sottoinsieme KernelBench Level-2, il 31.6% dei kernel generati ottiene almeno un speedup di 1.2x rispetto al riferimento Torch, superando Claude-4.5-Sonnet (26.7%) e GPT-5 (28.6%). Selezionando il miglior candidato tra tutti i turni, questa percentuale di speedup di 1.2x aumenta ulteriormente al 47.8%. Tutte le risorse, inclusi ambiente, codice di addestramento, modelli e dataset, sono disponibili su 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.