用强化学习生成高效GPU内核，解决奖励作弊和懒惰优化问题。
Dr. Kernel: Reinforcement Learning Done Right for Triton Kernel Generations
- 构建分布式训练环境KernelGYM，支持多轮交互与奖励校验。
- 提出TRLOO方法，消除多轮强化学习中的偏差，提升策略稳定性。
- 在内核加速测试中表现优于Claude-4.5-Sonnet，适合高性能计算场景。
高质量内核对可扩展的AI系统至关重要，让大模型生成此类代码将推动AI发展。然而，训练大模型进行内核生成需充足数据、稳健环境，且易受奖励作弊和懒惰优化影响，导致模型仅关注表面正确性而非实际性能提升。本文系统研究内核生成中的强化学习方法，首先设计了KernelGYM——一个支持奖励作弊检测、多轮交互数据收集与长期训练的分布式GPU环境。基于此，我们研究有效的多轮强化学习方法，发现自包含机制引发的策略梯度偏差，并提出逐轮强化留一法（TRLOO）以实现无偏优势估计。为缓解懒惰优化，引入匹配纠正机制保障训练稳定，提出基于性能分析的奖励（PR）与拒绝采样（PRS）。训练得到的Dr Kernel-14B模型在Kernelbench上表现媲美Claude-4.5-Sonnet。进一步研究测试时序列缩放，在KernelBench Level-2子集上，31.6%的生成内核实现至少1.2倍加速，超过Claude-4.5-Sonnet（26.7%）与GPT-5（28.6%）；若选取所有轮次的最佳候选，该比例升至47.8%。所有资源（环境、代码、模型、数据集）均开源于https://www.github.com/hkust-nlp/KernelGYM。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。