Dr. Kernel: Aprendizaje por Refuerzo Aplicado Correctamente para la Generación de Kernels en 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

Resumen

La calidad del kernel es crucial para sistemas de IA escalables, y permitir que los LLMs generen dicho código impulsaría el desarrollo de la IA. Sin embargo, entrenar LLMs para esta tarea requiere datos suficientes, un entorno robusto, y el proceso suele ser vulnerable a la manipulación de recompensas y a optimizaciones superficiales. En estos casos, los modelos pueden manipular las recompensas del entrenamiento y priorizar una corrección trivial sobre mejoras de velocidad significativas. En este artículo, estudiamos sistemáticamente el aprendizaje por refuerzo (RL) para la generación de kernels. Primero diseñamos KernelGYM, un entorno distribuido robusto en GPU que soporta verificación de manipulación de recompensas, recolección de datos de interacciones multi-turno y entrenamiento de RL a largo plazo. Basándonos en KernelGYM, investigamos métodos efectivos de RL multi-turno e identificamos un problema de gradiente de política sesgado causado por la auto-inclusión en GRPO. Para resolverlo, proponemos Refuerzo por Turno-Dejar-Uno-Fuera (TRLOO) para proporcionar una estimación de ventaja no sesgada para RL multi-turno. Para aliviar la optimización superficial, incorporamos corrección de desajuste para estabilidad del entrenamiento e introducimos Recompensas Basadas en Perfilado (PR) y Muestreo de Rechazo Basado en Perfilado (PRS) para superar el problema. El modelo entrenado, Dr.Kernel-14B, alcanza un rendimiento competitivo con Claude-4.5-Sonnet en Kernelbench. Finalmente, estudiamos el escalado secuencial en tiempo de prueba para Dr.Kernel-14B. En el subconjunto KernelBench Nivel-2, el 31.6% de los kernels generados logran al menos una aceleración de 1.2x sobre la referencia de Torch, superando a Claude-4.5-Sonnet (26.7%) y GPT-5 (28.6%). Al seleccionar el mejor candidato entre todos los turnos, esta tasa de aceleración de 1.2x aumenta aún más al 47.8%. Todos los recursos, incluidos entorno, código de entrenamiento, modelos y dataset, están disponibles en 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.