Доктор Кернел: Правильный подход к обучению с подкреплением для генерации ядер Triton
Dr. Kernel: Reinforcement Learning Done Right for Triton Kernel Generations
February 5, 2026
Авторы: Wei Liu, Jiawei Xu, Yingru Li, Longtao Zheng, Tianjian Li, Qian Liu, Junxian He
cs.AI
Аннотация
Высококачественные ядра критически важны для масштабируемых систем ИИ, а возможность их генерации большими языковыми моделями (LLM) ускорила бы развитие ИИ. Однако обучение LLM для этой задачи требует достаточного объема данных, надежной среды, и процесс часто уязвим для взлома системы вознаграждений и ленивой оптимизации. В таких случаях модели могут подстраиваться под метрики обучения, отдавая приоритет тривиальной корректности вместо значительного ускорения. В данной статье мы систематически изучаем обучение с подкреплением (RL) для генерации ядер. Сначала мы разрабатываем KernelGYM — надежную распределенную среду на GPU, которая поддерживает проверку на взлом вознаграждений, сбор данных из многошаговых взаимодействий и долгосрочное RL-обучение. На основе KernelGYM мы исследуем эффективные многошаговые RL-методы и выявляем проблему смещенного градиента политики, вызванную самовключением в GRPO. Для решения этой проблемы мы предлагаем метод TRLOO (Turn-level Reinforce-Leave-One-Out) для получения несмещенной оценки преимущества в многошаговом RL. Чтобы смягчить ленивую оптимизацию, мы внедряем коррекцию рассогласования для стабильности обучения и вводим профилировочные вознаграждения (PR) и профилировочный rejection sampling (PRS) для преодоления этой проблемы. Обученная модель Dr.Kernel-14B демонстрирует результаты, сопоставимые с Claude-4.5-Sonnet в Kernelbench. Наконец, мы изучаем последовательное масштабирование на этапе тестирования для Dr.Kernel-14B. На подмножестве KernelBench Level-2 31,6% сгенерированных ядер достигают как минимум 1,2-кратного ускорения по сравнению с эталонной реализацией Torch, превосходя Claude-4.5-Sonnet (26,7%) и GPT-5 (28,6%). При выборе лучшего кандидата по всем шагам этот показатель ускорения ≥1,2x увеличивается до 47,8%. Все ресурсы, включая среду, код обучения, модели и набор данных, доступны по адресу 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.