Dr. Kernel: Aprendizado por Reforço Aplicado Corretamente para Gerações de Kernels Triton
Dr. Kernel: Reinforcement Learning Done Right for Triton Kernel Generations
February 5, 2026
Autores: Wei Liu, Jiawei Xu, Yingru Li, Longtao Zheng, Tianjian Li, Qian Liu, Junxian He
cs.AI
Resumo
Kernels de alta qualidade são críticos para sistemas de IA escaláveis, e permitir que LLMs gerem esse código avançaria o desenvolvimento de IA. No entanto, treinar LLMs para essa tarefa requer dados suficientes, um ambiente robusto, e o processo é frequentemente vulnerável a "hacking de recompensa" e otimização preguiçosa. Nestes casos, os modelos podem burlar as recompensas de treinamento e priorizar correção trivial em detrimento de aceleração significativa. Neste artigo, estudamos sistematicamente o aprendizado por reforço (RL) para geração de kernels. Primeiro, projetamos o KernelGYM, um ambiente distribuído robusto em GPU que suporta verificação de hacking de recompensa, coleta de dados de interações multi-turno e treinamento de RL de longo prazo. Com base no KernelGYM, investigamos métodos eficazes de RL multi-turno e identificamos um problema de gradiente de política tendencioso causado pela auto-inclusão no GRPO. Para resolver isso, propomos o Turn-level Reinforce-Leave-One-Out (TRLOO) para fornecer uma estimativa de vantagem não tendenciosa para RL multi-turno. Para aliviar a otimização preguiçosa, incorporamos correção de incompatibilidade para estabilidade do treinamento e introduzimos Recompensas Baseadas em Perfilamento (PR) e Amostragem de Rejeição Baseada em Perfilamento (PRS) para superar o problema. O modelo treinado, Dr.Kernel-14B, atinge desempenho competitivo com o Claude-4.5-Sonnet no Kernelbench. Finalmente, estudamos o escalonamento sequencial em tempo de teste para o Dr.Kernel-14B. No subconjunto KernelBench Level-2, 31,6% dos kernels gerados alcançam pelo menos uma aceleração de 1,2x em relação à referência do Torch, superando Claude-4.5-Sonnet (26,7%) e GPT-5 (28,6%). Ao selecionar o melhor candidato em todos os turnos, essa taxa de aceleração de 1,2x aumenta ainda mais para 47,8%. Todos os recursos, incluindo ambiente, código de treinamento, modelos e conjunto de dados, estão incluídos em 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.