Dr. Kernel：用强化学习正确地生成 Triton 内核 Dr. Kernel: Reinforcement Learning Done Right for Triton Kernel Generations

方法的具体步骤如下：(1) 冷启动数据收集：从 CUDALLM-SFT 获取 8K 内核生成查询，使用 GPT-5 生成 5 轮 Triton 实现，每轮在 KERNELGYM 中执行并将反馈（正确性状态、错误诊断、profiling 摘要）附加到下一轮查询，构建多轮交互轨迹用于 SFT。(2) 冷启动 SFT：在收集的 8K 轨迹上进行监督微调，学习率 1e-6，batch size 256，训练 4 个 epoch，教授模型基本内核生成技能如 tiling、fusion 等。(3) 多轮 RL 训练：在 CUDALLM 的 RL 查询上训练，每 prompt 采样 16 个 rollout，最大 3 轮，rollout batch size 16，训练 300 步。每轮奖励定义为 R_{i,t} = C(y_{i,t}) + C(y_{i,t}) * speedup_{i,t}，其中 C 是二值正确性奖励，speedup 裁剪到最大 3 倍。(4) 训练稳定性优化：应用 MRS 过滤 off-policy 样本，计算几何平均重要性比率 w = exp((1/|T|) sum_{t in T} log(pi_train(a_t|s_t)/pi_rollout(a_t|s_t)))，仅保留 w 在 [0.999, 1.001] 范围内的样本。(5) 克服懒惰优化：引入 PR 将奖励扩展为 R_{i,t} = C(y_{i,t}) + C(y_{i,t}) * speedup_{i,t} + C(y_{i,t}) * PR_{i,t}，同时应用 PRS 按概率过滤低影响样本。(6) 推理时 STTS：使用上下文管理策略，存储全部历史但只保留奖励最高的 top-w 轮（w=4）在 prompt 中，避免上下文溢出同时保持高质量历史信息。

本文的技术新颖性体现在多个层面。在环境层面，KERNELGYM 是首个专为内核生成 RL 设计的分布式执行系统，其 server-worker 架构支持弹性扩展、严格故障隔离（每个评估在独立子进程中运行），并提供 hacking check 和 profiler 两种创新工具包——前者通过检测 Triton 内核是否被实际调用来防止奖励作弊，后者提供内核级执行时间摘要支持瓶颈识别。在 RL 方法层面，TRLOO 是首次将 Leave-One-Out baseline 引入多轮 RL 的工作，其数学证明了 GRPO 的系统性偏差并提供了无偏替代方案，这在理论上具有独立价值。在优化策略层面，PR 和 PRS 是首个利用 profiling 信息引导 RL 训练方向的工作，将传统的标量奖励信号扩展为包含执行时间占比的多维信号，使优化目标从'任何加速'转向'有意义的加速'。这种从稳定性（MRS）到有效性（PR/PRS）的分层诊断方法论也具有方法论意义。

本文存在几个值得探讨的弱点。第一，PR 信号的设计假设生成内核的执行时间占比越高越好，但这一假设在某些场景下可能不成立——例如当模型生成一个覆盖大部分执行时间但效率低下的内核时，高 PR 值可能误导优化方向。改进方向可以考虑结合内核效率指标（如 FLOPS 利用率）来更精确地衡量优化价值。第二，PRS 的阈值参数 tau=0.3 和 s=0.1 是固定的，缺乏自适应机制。在训练的不同阶段，合理的过滤阈值可能不同——早期需要更多探索，后期可以更激进地过滤。改进方向是设计随训练进度自适应调整的阈值策略。第三，论文的多轮 RL 最多只训练 3 轮，但 STTS 推理时可以扩展到 14 轮，这种训练-推理的轮次差距可能限制了模型的多轮优化能力。改进方向是逐步增加训练轮次或使用课程学习策略。第四，评估仅在 NVIDIA H100 上进行，不同 GPU 架构（如 A100、消费级 GPU）上的性能表现未知，这限制了方法的普适性结论。

论文作者和基于本文成果可延伸的未来研究方向包括：第一，数据扩展和预训练——进行更大规模数据收集以支持领域特定预训练或持续预训练（middle-training），为后续 RL 优化提供更坚实的基础。这可以通过自动化挖掘开源代码库中的内核实现或合成更多训练任务来实现。第二，模型容量扩展——将方法迁移到更大参数规模的模型上，作者观察到 14B 相比 8B 有显著提升，扩展到 70B+ 可能带来进一步收益。第三，多架构支持——将 KERNELGYM 扩展到支持 CUDA、TileLang 等更多内核语言，以及探索跨 GPU 架构的泛化能力。第四，生产级自动化——从当前的代码片段生成发展到端到端的内核优化流水线，包括自动识别优化机会、生成优化方案、验证和部署。第五，更先进的 RL 方法——探索如 PPO、DPO 等其他 RL 算法在内核生成中的应用，或者设计更精细的奖励函数来平衡正确性、加速和代码质量。

本文在可复现性方面做得较好。作者承诺开源所有资源，包括 KERNELGYM 环境、训练代码、模型和数据集，代码仓库位于 github.com/hkust-nlp/KernelGYM。冷启动数据从 CUDALLM-SFT 获取 8K 查询并使用 GPT-5 生成轨迹，这意味着完全复现需要 GPT-5 API 访问权限，但作者可能提供预处理好的轨迹数据。训练使用 Qwen3-8B-Base 和 Qwen-14B-Base 作为基础模型，这些都是公开可用的模型。硬件要求方面，训练在 NVIDIA H100 上进行，16 个 rollout per prompt 的 RL 训练需要显著的 GPU 资源。评估使用 KernelBench 基准的官方实现。总体而言，主要复现障碍在于算力需求和 GPT-5 API 访问，但框架设计和算法实现应该是可复现的。