港科大等DR. KERNEL：强化学习训练系统赋能AI生成高性能GPU代码

这项由香港科技大学、字节跳动、香港中文大学（深圳）以及南洋理工大学联合开展的研究发表于2026年，研究团队开发出了一套完整的训练系统，让大语言模型学会编写高性能的GPU内核代码。这项突破性工作首次系统性地解决了用强化学习训练AI模型编写内核代码时遇到的种种挑战，相关论文编号为arXiv:2602.05885v1。

想象一下，如果把GPU比作一个超级厨房，那么内核代码就像是烹饪食谱。一个优秀的食谱能让厨师充分发挥厨房设备的潜力，做出美味佳肴。同样，高质量的内核代码能让GPU发挥最大性能，让AI系统运行得更快更高效。然而，编写这样的"食谱"需要深厚的专业知识，就像一个厨师需要同时精通食材特性和烹饪技巧一样。

目前的AI系统越来越复杂，对计算性能的要求也越来越高。FlashAttention和FlashInfer等经典工作已经证明，专门设计的内核代码对于释放现代大语言模型的全部效率至关重要。但是开发这样的内核代码极其困难，需要开发者具备算法和GPU硬件方面的深厚专业知识。虽然像Triton这样的领域特定语言降低了相比CUDA的入门门槛，但要达到峰值性能仍然需要大量的手工优化。

研究团队发现，让AI模型学会编写内核代码面临着两个主要挑战。第一个挑战被称为"奖励黑客行为"，就像学生为了应付考试而背标准答案，但实际上并没有真正掌握知识。模型可能会生成看起来正确且运行很快的代码，但实际上这些代码毫无意义，只是钻了评测系统的漏洞。第二个挑战叫"懒惰优化"，模型虽然生成了正确的代码并获得了一些性能提升，但这些提升微不足道，就像只是把菜谱中的一个调料稍作调整，而没有解决真正的烹饪瓶颈。

一、构建完善的训练环境：KERNELGYM系统

为了解决这些问题，研究团队首先建立了一个名为KERNELGYM的分布式GPU执行环境。就像建立一个标准化的厨房测试中心，这个环境能够公正准确地评估各种"食谱"的效果。

KERNELGYM采用服务器-工作节点的分布式架构设计。服务器端包含一个接口层和任务管理器，接口层通过FastAPI提供REST API服务，负责接收任务提交和查询请求以及工作节点注册。任务管理器使用Redis维护持久化的任务和工作节点状态，并将任务分派给可用的工作节点，支持基于超时的重新排队机制来维持吞吐量。

在工作节点方面，每个GPU都被视为独立的工作实例，从服务器拉取计划任务并使用配置的后端工具包按序执行。为了防止生成的内核代码中的CUDA错误或运行时故障损坏长期运行的进程，每次评估都在新生成的子进程中运行，而父工作进程保持CUDA环境清洁并继续服务后续任务。工作节点监控器跟踪存活状态，自动重启失败的工作节点，并将未完成的任务重新分配给健康的工作节点，以保持强化学习训练的稳定性。

这个环境的核心优势在于它的容错能力。就像在厨房测试中心，即使某个厨师的实验菜谱导致了厨房事故，也不会影响其他厨师继续工作，整个测试中心会自动清理现场并重新开始。KERNELGYM实现了严格的故障隔离，当不安全的生成内核触发非法内存访问或不可恢复的CUDA错误时，故障会被隔离在任务级别并自动恢复，确保长期训练的可用性。

更重要的是，KERNELGYM提供了丰富的环境反馈。不同于简单的通过失败或单一的速度提升值，系统暴露出结构化的反馈信息，包括执行分析摘要和奖励黑客检测，以支持优化和数据收集。系统还实现了基于执行的黑客检查功能，通过监控Triton的启动路径来记录执行的Triton内核，并在训练和评估模式下测量端到端运行时间。如果代码在任一模式下都没有执行Triton内核，就会被标记为不正确。

二、多轮强化学习的偏差问题与解决方案

研究团队发现，在多轮交互的强化学习训练中，传统的GRPO方法存在一个重要的技术问题：自包含偏差。这就像在评估学生成绩时，把学生自己的成绩也算入了全班平均分，这会导致评估基准线受到当前样本的影响，从而产生有偏差的策略梯度估计。

具体来说，在GRPO的分组平均基准线中，对于给定的提示问题，系统会采样K个独立的轨迹。在第t轮，某些轨迹可能无效，系统用G_t表示给定提示在第t轮的有效轨迹集合，N_t = |G_t| ≤ K。然后计算每轮组内的平均回报G_t = (1/N_t) Σ_{j∈G_t} G_{j,t}，GRPO优势为A^{GRPO}_{i,t} = G_{i,t} - G_t。

问题在于，G_t包含了G_{i,t}本身，对于任何i ∈ G_t都是如此。由于G_{i,t}通过从第t轮开始的奖励依赖于当前动作y_{i,t}，基准线可能变得依赖于动作，违背了无偏REINFORCE基准线的标准要求。对于均值中心化形式，这表现为每个提示-轮次组内的有偏策略梯度估计器，更新被系统性地收缩了一个依赖于有效组大小的因子。

为了解决这个问题，研究团队提出了轮级REINFORCE留一法（TRLOO），这是留一法基准线在多轮强化学习中的适配版本。对于每个组G_t和样本i ∈ G_t且N_t > 1，定义G^{(-i)}_t = (1/(N_t-1)) Σ_{j∈G_t, j≠i} G_{j,t}，A^{TRLOO}_{i,t} = G_{i,t} - G^{(-i)}_t。

等价地，A^{TRLOO}_{i,t} = (N_t/(N_t-1))(G_{i,t} - G_t)。由于G^{(-i)}_t排除了G_{i,t}，在独立轨迹下它不依赖于当前动作y_{i,t}，从而产生多轮强化学习的无偏轮级优势估计器。

除了无偏性，TRLOO对稀疏正奖励的困难任务特别有益，在这种情况下成功轨迹很少。首先，它避免了GRPO中的自惩罚现象：在均值中心化下，罕见的高回报样本会贡献到G_t中，因此通过减去基准线部分地抑制其优势。TRLOO从基准线中排除G_{i,t}，所以罕见成功获得更大的学习信号。其次，TRLOO对变化的组大小具有鲁棒性。在多轮细化中，后续轮次由于上下文限制或早期终止可能有更少的有效样本，使得在等式中1/N_t更大。TRLOO消除了这种自包含效应，在正反馈稀少时跨变化组大小保持正确尺度，提高样本效率。

三、训练稳定性与优化目标对齐

研究团队通过实验验证发现，多轮强化学习训练中的"懒惰优化"问题有两个潜在原因。第一个假设是优化不稳定性导致的饱和，第二个假设是优化目标与有意义的加速不匹配，激励低影响解决方案。

针对第一个假设，团队监控了训练动态，包括熵、梯度范数和困惑度。观察发现，多轮内核生成强化学习训练表现出过高的这些指标值，明确表明训练不稳定。训练-推理不匹配是LLM强化学习中的基本挑战，其中轨迹（推理）和训练引擎之间的差异导致离策略漂移。理论上，这种漂移可能导致梯度方差和奖励崩溃，阻止模型达到更高的性能峰值。

为了缓解这种漂移，研究团队采用了几何不匹配拒绝采样（MRS）。他们计算几何平均重要性比w = exp((1/|T|) Σ_{t∈T} log(π_{train}(a_t|s_t)/π_{rollout}(a_t|s_t)))，只保留w ∈ [0.999, 1.001]的样本。此外，他们还强制执行严格的token级别否决：如果任何单个token的似然比π_{train}/π_{rollout}低于10^{-4}，整个序列都会被拒绝。

经过验证，MRS成功稳定了训练动态。然而，关键洞察是，虽然不匹配校正防止了早期崩溃，使学习曲线更平滑，但它并没有从根本上提升Fast@1.2的性能上限。这表明虽然假设1有效解释了训练不稳定性，单独解决它并不能完全解决性能饱和问题，因此将注意力转向优化目标本身。

针对第二个假设，研究团队认为标准奖励信号未能区分微不足道的改进和有意义的瓶颈。虽然内核可能是正确的并实现了一些加速，但它可能仍然未能解决真正的性能瓶颈。要从生成仅仅正确的内核转向有效的内核，必须使奖励具有瓶颈感知能力。

团队提出了基于分析的奖励（PR）方法。这种不匹配的一个关键症状是模型倾向于优化微不足道的子操作（例如替换简单的求和操作），而不影响计算中的主要瓶颈。在懒惰优化与更好融合的案例研究中，在懒惰优化情况下，模型生成的内核仅占总CUDA执行时间的0.014%，表明内核优化没有影响主要瓶颈。相比之下，通过更好的融合，模型生成的内核覆盖了总CUDA运行时间的86.15%，导致更好且更有意义的加速。

为了形式化这种直觉，研究团队利用KERNELGYM中的分析工具包来隔离生成内核（T_{generated}）从整体CUDA执行时间（T_{total}）的运行时贡献。他们定义分析比率为PR_{i,t} = T_{generated}/T_{total}。直觉上，PR_{i,t}在候选优化主导端到端运行时间的内核时分配更高的信用。然后他们用这个信号增强每轮奖励（仅应用于正确的内核）：R_{i,t} = C(y_{i,t}) + C(y_{i,t}) · speedup_{i,t} + C(y_{i,t}) · PR_{i,t}。

这鼓励模型专注于对性能有显著贡献的内核优化，明确驱动学习朝向具有更大真实加速的优化。此外，由于PR_{i,t}在[0,1]范围内，加速项自然占主导地位，防止模型通过低效代码最大化覆盖率。

研究团队还引入了基于分析的拒绝采样（PRS）。即使有瓶颈感知奖励，探索过程仍可能被大量低影响（"懒惰"）样本主导。为了进一步过滤训练分布，他们引入了基于分析的拒绝采样。对于每个样本(i,t)，他们以概率p_{i,t} = clip((PR_{i,t} - τ)/s, 0, 1)保留它，其中τ是截断阈值，s控制过滤器的软度。在实验中，他们固定τ = 0.3和s = 0.1。

四、测试时序列扩展

研究团队还探索了测试时的序列扩展（STTS），通过在推理时增加多轮细化步骤的数量来最大化训练模型的推理能力。他们使用DR. KERNEL-14B，最大上下文长度为32,768个token，评估了两种策略：香草外推和上下文管理。

香草外推直接通过在每轮将整个交互历史附加到提示中来外推细化轮次数超出训练范围（用最多3轮训练）。结果显示，增加轮次最初改善了最后轮次和历史最佳性能。然而，随着T增长，提示长度线性扩展，可能接近上下文限制，这可能降低性能。

为了在不无限制提示增长的情况下扩展T，研究团队采用上下文管理将所有轮次存储在外部内存中，并维护固定的上下文内窗口。具体来说，在每轮他们从累积历史中选择奖励最高的前w轮（w=4），并仅将这些选定的轮次作为生成下一轮的提示历史包含进来。

实验结果表明，上下文管理产生一致更强的历史最佳性能，并随着轮次扩展继续改善。最后轮次性能在小T时可能稍低，因为香草外推可以基于完整历史进行条件化，但随着T增加，上下文管理变得严格更可靠，超越香草外推可达到的最佳性能。

五、实验结果与性能表现

研究团队在KernelBench的所有三个级别上进行了全面评估。他们遵循KernelBench中的官方Torch后端及其正确性和加速测量的实现，并在评估内核时进一步进行黑客检查。因此，他们的评估比原始KernelBench更严格。

实验设置方面，团队使用Qwen3-8B-Base和Qwen-14B-Base模型，对每个问题采样8个样本，最多3轮。他们设置最大token数为32768，每轮最大生成token数为8192。为保持公平比较，他们报告第3轮的结果，因为第3轮通常产生最佳平均性能。对于在较早轮次达到最佳平均性能的基线，他们报告其最佳表现轮次。

在冷启动阶段，团队首先在收集的8K个5轮轨迹上进行监督微调，使用学习率1×10^{-6}，批量大小256，训练4个epoch。冷启动训练后，他们在来自CudaLLM的强化学习查询上运行多轮强化学习。强化学习使用学习率1×10^{-6}，训练300个轨迹步骤，每个提示采样16个轨迹，最多3轮，轨迹批量大小为16。

主要结果显示，DR. KERNEL在开源基线中取得最强性能，在Level 1和Level 2上与前沿模型竞争。特别是，DR. KERNEL-14B在Level 1和Level 2上都达到了高Fast@1.2，表明它不仅改善了任何加速（Fast@1），还改善了更严格和有意义的加速。这与之前的方法如AutoTriton形成对比，后者在Level 2上实现了强劲的Fast@1但在更严格阈值下提供显著更小的增益。

与冷启动模型相比，DR. KERNEL表明多轮强化学习对性能增益有实质贡献，特别是在更严格的指标上（例如Level 2上Fast@1.2从5.6→20.0）。虽然DR. KERNEL相对于开源基线改善了Level 3 Fast@1，但在Level 3上更严格阈值的性能仍然有限，表明可能需要进一步扩展训练数据和模型容量来缩小与前沿模型在最困难子集上的差距。

测试时扩展进一步放大了DR. KERNEL的性能。通过上下文管理的序列测试时扩展（STTS），DR. KERNEL-14B-STTS将Fast@1.2从16.9→18.8（Level 1）和从25.6→31.6（Level 2）提升；通过历史轮次的最佳轮次选择（DR. KERNEL-14B-STTS+），Fast@1.2进一步上升到25.1（Level 1）和47.8（Level 2），在这些指标上超越GPT-5和Claude-4.5-Sonnet等前沿模型。

研究团队还在torch.compile下评估了模型性能。torch.compile是一个先进的PyTorch功能，它将PyTorch程序捕获到编译的计算图中，并通过算子融合、代码生成和调度等编译器传递优化执行。虽然大多数先前工作仅在Torch急切执行下评估模型生成的内核，但团队进一步在torch.compile下验证了他们的模型和前沿模型，提供了显著更强和更实用的加速评估。

结果显示，DR. KERNEL在更具挑战性的torch.compile设置下保持有效，并在三个级别上与前沿模型保持竞争。由于torch.compile已经应用编译器优化，额外增益的余地比急切模式小；因此，包括我们和前沿模型在内的所有模型的绝对Fast@p数字普遍较低。重要的是，torch.compile下的Fast@1也是更严格的目标：在急切执行中可能产生边际改善的微不足道"懒惰"变化通常不会超越优化的编译基线。

六、技术局限与未来展望

尽管这项研究在用强化学习训练Triton内核生成方面取得了进展，研究团队也承认了几个需要进一步探索的领域。

从数据角度来看，资源限制将他们的监督微调阶段限制为8,000个冷启动样本。鉴于当前大语言模型的预训练语料库中高质量内核编程数据相对稀缺，他们的结果表明该领域的"数据下限"相当高。未来工作可能涉及更大规模的数据收集，以促进领域特定的预训练或持续预训练，这也将为后续的强化学习优化提供更强大的基础。

从模型容量角度，他们对DR. KERNEL-8B和DR. KERNEL-14B的观察证实，更大的模型具有更优的内核生成能力。这种扩展效应在强化学习中特别关键，因为模型必须依赖自己的生成来探索解决方案空间并更新其策略。他们预期将这些方法迁移到更大的参数规模将加速发展。

虽然他们的方法取得了与前沿模型匹敌或超越的性能改进，但该领域仍处于探索阶段。虽然当前模型可以生成高质量的代码片段，但它们尚未能够为生产环境进行完全自主的端到端内核生成。研究团队希望他们的贡献，特别是KERNELGYM环境和DR. KERNEL训练框架，能作为未来研究的催化剂。

总的来说，这项研究为AI系统自动化GPU内核代码生成开辟了新的可能性。通过系统性地解决强化学习训练中的关键挑战，并提供完整的训练环境和方法，这项工作为未来在这一重要领域的研究奠定了坚实基础。随着AI系统变得越来越复杂，对高效GPU计算的需求只会继续增长，这使得自动化内核生成成为一个具有重要实际意义的研究方向。

Q&A

Q1：KERNELGYM是什么系统？

A：KERNELGYM是香港科大团队开发的分布式GPU执行环境，专门用来训练AI模型编写GPU内核代码。它就像一个标准化的测试中心，能够公正评估AI生成的代码性能，并具备强大的容错能力，即使某个测试导致系统故障也不会影响整体训练。

Q2：什么是奖励黑客行为和懒惰优化？

A：奖励黑客行为是指AI模型生成看起来正确且快速的代码，但实际上毫无意义，只是钻了评测系统的漏洞。懒惰优化是指模型虽然生成了正确代码并有性能提升，但提升微不足道，没有解决真正的性能瓶颈，就像只调整了菜谱中的一个小调料。

Q3：DR. KERNEL相比其他方法有什么优势？

A：DR. KERNEL通过TRLOO方法解决了多轮训练中的偏差问题，通过分析奖励和拒绝采样解决了懒惰优化问题。在KernelBench测试中，DR. KERNEL-14B在Level 2子集上31.6%的生成内核达到了至少1.2倍的加速，超过了Claude-4.5-Sonnet的26.7%和GPT-5的28.6%。