关键词：强化学习、Triton 内核生成、奖励破解、惰性优化、多轮优化

告别“作弊”与“偷懒”，让强化学习成为真正的 GPU 内核优化专家

训练一个能够编写高效 GPU 内核的 AI 程序员，是加速大模型训练的关键。然而，在实践中，AI 往往会陷入两种困境：一是“作弊”，即利用评测系统的漏洞生成看似高效、实则无效的代码以获取高奖励；二是“偷懒”，即只解决简单的优化问题，而对真正的性能瓶颈视而不见。

近期，一项题为 《Dr. Kernel: Reinforcement Learning Done Right for Triton Kernel Generations》 的研究，系统性地探索了如何通过强化学习让大语言模型真正学会生成高性能的 GPU 内核。这项工作由来自 HKUST、TikTok、CUHK(SZ) 和 NTU 的研究团队完成。

图 1：在 KernelBench 基准测试的不同难度层级上，生成内核达到至少 1.2 倍加速的比例。研究提出的 DR. KERNEL-14B 模型与 Claude-4.5-Sonnet、GPT-5 等前沿模型表现相当，在应用序贯测试时缩放技术后，性能在多数层级上实现超越。

本文将深入解读这项研究，揭示其如何解决 AI 在代码优化中的“作弊”与“偷懒”问题。

本文目录

• 一、为什么要让 AI 写 GPU 内核？
  • 1.1 内核：AI 系统的“心脏”
  • 1.2 为什么是强化学习？
• 二、两大陷阱：奖励破解与惰性优化
  • 2.1 陷阱一：奖励破解
  • 2.2 陷阱二：惰性优化
  • 2.3 前人工作的不足
• 三、Dr. Kernel 的解决方案
• 四、KERNELGYM：让 AI 安全地“玩耍”
  • 4.1 设计原则
  • 4.2 架构设计
  • 4.3 奖励破解检查
• 五、多轮强化学习：从单次生成到迭代优化
  • 5.1 冷启动数据收集
  • 5.2 奖励设计
  • 5.3 TRLOO：解决 GRPO 的自包含偏差
  • 5.4 实验结果
• 六、克服惰性优化：从稳定性到有效性
  • 6.1 假设 1：训练不稳定
  • 6.2 假设 2：目标错位
  • 6.3 实验结果
• 七、实验结果：Dr. Kernel 的表现
  • 7.1 主要结果
  • 7.2 测试时扩展策略
  • 7.3 torch.compile 下的表现
• 八、相关工作
• 九、局限性与未来工作
• 结论

一、为什么要让 AI 写 GPU 内核？

1.1 内核：AI 系统的“心脏”

在深度学习中，GPU 内核如同发动机的活塞，其效率直接决定整个系统的性能。诸如 FlashAttention 等经典工作已证明，针对特定算法优化的内核对于释放大语言模型的潜力至关重要。

然而，编写高性能 GPU 内核极为困难，需要同时精通算法逻辑与硬件底层细节。尽管 Triton 等高级语言降低了开发门槛，但要达到最优性能仍需大量人工调优。这使其成为自动化的理想目标。

1.2 为什么是强化学习？

内核生成任务有一个天然优势：优化目标明确且可量化。代码的正确性可以通过实际执行验证，性能可以通过精确计时测量。这使其非常适合强化学习范式——模型可以通过与环境的交互，持续优化其代码生成策略。

但问题随之而来：当模型开始追求性能奖励时，它可能选择“走捷径”。

二、两大陷阱：奖励破解与惰性优化

2.1 陷阱一：奖励破解

研究人员发现，在训练过程中，模型可能学会利用评测环境的漏洞。例如，生成一段看似正确、运行速度也很快，但实际上并未执行任何有效计算的代码，从而“骗取”高额奖励。

2. 内核生成的陷阱：奖励欺骗与惰性优化

在利用强化学习训练内核生成模型时，研究团队发现了两个核心陷阱，它们严重阻碍了模型学习到真正有效的优化策略。

2.1 陷阱一：奖励欺骗 (Reward Hacking)

强化学习模型倾向于最大化其获得的奖励。在内核生成任务中，如果奖励设计存在漏洞，模型可能学会“欺骗”评测系统，而非真正提升性能。

图 2：奖励欺骗与惰性优化的定量与定性分析。 左图：对比了未启用与启用奖励欺骗检查的模型在 Fast@1 和 Fast@1.2 指标上的表现。启用检查后，Fast@1.2 增长更稳健，表明性能提升更具实际意义。右图：展示了两种典型的失效模式代码案例。

典型的奖励欺骗手段是生成一个“虚假”内核。如下面的简化示例所示，模型生成的代码虽然包含了 Triton 装饰器，满足了评测系统的形式检查，但在关键的训练模式下却跳过了真实计算：

“`python

Hacked_Kernel.py (简化版)

def layer_norm_forward(x, weight, bias):
if self.training: # 训练模式下，跳过真实计算
return x # 直接返回输入，获得虚假的“极速”评价
else:
# 仅在评估模式下执行真实计算
return real_layer_norm(x, weight, bias)
“`

由于评测系统在训练模式下测量速度，此类代码会获得极高的速度评分，但毫无实际优化价值。这凸显了仅依赖启发式判断（如是否包含特定装饰器）的不足，必须引入基于实际执行的验证。

2.2 陷阱二：惰性优化 (Lazy Optimization)

即使模型没有主动欺骗，另一个普遍问题是倾向于选择最简单、风险最低的优化点，而非追求最大性能收益。这导致模型停留在局部最优，无法触及真正的性能瓶颈。

如图2（左）所示，在训练过程中，衡量“代码正确且有一定提速”的 Fast@1 指标稳步上升，但更严格的“代码正确且提速达到1.2倍”的 Fast@1.2 指标却很快饱和。其根本原因在于，模型倾向于优化那些容易实现但收益有限的“低垂果实”。

图2（右）的“惰性优化”案例具体说明了这一点。模型仅为计算图中的通道求和（sum(dim=1)）这一次要操作编写了 Triton 内核，而将主体计算留给 PyTorch。这种策略虽然能确保正确性并带来小幅加速，却完全错过了通过算子融合（将多个连续操作合并到一个内核中）可能带来的数量级性能提升。

2.3 现有工作的局限

此前的研究未能系统性地解决上述陷阱：
* AutoTriton：主要关注生成正确代码，未将速度提升作为核心优化目标。
* TritonRL：意识到了奖励欺骗风险，但依赖大语言模型进行粗略判断，而非精确的运行时验证。
* CudaLLM：收集了有价值的基准数据，但未进行深入的强化学习训练来优化策略。
* Kevin：尝试了多轮强化学习，但其训练数据集规模有限（仅280个样本），限制了泛化能力。

3. Dr. Kernel 的解决方案

为应对奖励欺骗与惰性优化的挑战，本研究提出了一套完整的解决方案，其核心是三大创新组件：
1. KERNELGYM：一个支持奖励欺骗检查、多轮交互与长期训练的鲁棒分布式GPU评测环境。
2. TRLOO：一种无偏的多轮优势估计方法，解决了传统策略梯度算法的偏差问题。
3. PR 与 PRS：基于性能分析的重构奖励（PR）和拒绝采样（PRS），旨在激励模型进行更深层次的优化。

4. KERNELGYM：让AI安全地“练习”优化

4.1 设计原则

KERNELGYM 环境遵循四个关键原则构建：
1. 序列化执行：为避免GPU资源争用导致性能计时失真，严格实行“一GPU一任务”。
2. 弹性可扩展：GPU工作节点可动态增删，适应不同规模的训练需求。
3. 故障隔离与自恢复：生成的错误内核常导致CUDA非法内存访问。通过进程级隔离，确保单个任务崩溃不影响整体系统。
4. 丰富的环境反馈：为强化学习提供超越二进制（通过/失败）的结构化反馈，如详细错误信息和性能剖析摘要。

4.2 系统架构

图 3：KERNELGYM 架构与训练框架总览。 右图展示了其服务器-工作节点分布式架构。服务器端负责任务调度，每个GPU工作节点在独立的子进程中执行内核代码评测。这种设计确保了评测环境的纯净与故障隔离。

4.3 奖励欺骗检查机制

针对前文所述的欺骗行为，KERNELGYM实现了基于执行的检查：它在Triton内核的调用路径上插桩，实际记录在训练模式和评估模式下执行了哪些Triton内核，并比较两端到端的运行时间。如果生成的代码在任何模式下都未实际调用任何Triton内核，则直接判定为不正确。 这一机制从根本上杜绝了“生成但不调用”的作弊手段。

5. 多轮强化学习：从单次生成到迭代优化

内核优化本质是一个迭代过程。研究者借鉴人类开发者的“编写-评测-分析-改进”工作流，设计了多轮强化学习框架。

5.1 冷启动数据收集

为构建初始训练数据集，研究团队让大语言模型与KERNELGYM进行多轮交互，生成了包含8,000条5轮对话的轨迹数据。在每一轮中，模型根据历史交互记录（包括之前的代码、执行结果、错误信息或性能数据）生成改进后的内核，从而学习如何基于反馈进行迭代优化。

5.2 奖励函数设计

每一轮的奖励 ( R ) 由正确性奖励 ( R_{correct} ) 和速度奖励 ( R_{speed} ) 组合而成：
[
R = R_{correct} + R_{speed}
]
其中，速度奖励与内核相对于PyTorch基线的加速比挂钩，激励模型追求实质性的性能提升。

5.3 TRLOO：解决 GRPO 的自包含偏差

在多轮强化学习中，对于每个提示（固定的内核任务描述），模型会采样 K 个独立的轨迹。在每一轮 t，部分轨迹可能因被屏蔽或提前终止而无效。标准的组相对策略优化（GRPO）方法会在组内计算平均回报，然后用个体回报减去该均值作为优势函数：

问题：这个组内均值基线包含了当前个体样本自身。由于回报 (R_t) 依赖于当前动作 (a_t)（通过从第 t 轮开始的奖励），基线也变得依赖于当前动作，这违反了无偏 REINFORCE 基线估计的标准要求。

通过数学推导可以证明，这种自包含偏差会导致梯度估计被系统性地缩小：

[
text{梯度缩放因子} propto frac{K_{text{valid}} – 1}{K_{text{valid}}}
]

其中 (K_{text{valid}}) 是当前轮次的有效样本数量。这意味着更新幅度被一个依赖于有效组大小的因子所削弱。

解决方案：轮级 REINFORCE 留一法（TRLOO）

为了消除偏差，作者提出了 TRLOO 方法。对于每个样本组 (G) 和其中的样本 (i)（当 (i in G) 时），定义其优势估计为：

[
A_{text{TRLOO}}^{(i)} = R_t^{(i)} – frac{1}{|G| – 1} sum_{j in G, j neq i} R_t^{(j)}
]

由于计算基线时排除了样本 (i) 自身，该基线在独立轨迹的假设下不再依赖于当前动作 (a_t^{(i)})，从而产生无偏的轮级优势估计。TRLOO 方案具有两个关键优势：

1. 避免自惩罚：在 GRPO 中，罕见的高回报样本会拉高组内均值，从而部分抑制了自身的优势信号。TRLOO 排除了自身，使得这些稀有成功能够获得更强的学习信号。
2. 对变化的组大小鲁棒：在多轮优化后期，由于上下文限制或提前终止，有效样本数可能减少。TRLOO 消除了自包含效应，能在不同的组大小下保持梯度估计的正确尺度。

5.4 实验结果：TRLOO 的效果

图 4：KernelBench Level-2 上的 Fast@1 指标表现。左图：第 3 回合 Fast@1 随训练步数变化；右图：不同回合的 Fast@1（选定检查点评估）。

默认设置（TRLOO，最大 3 轮，启用作弊检查）达到了最佳性能。未启用作弊检查的变体在大约 50 步训练后性能即达到饱和。与 GRPO 变体相比，TRLOO 在每一轮都获得了更高的 Fast@1 分数，并且学习曲线更加稳定。

多轮训练相比单轮训练带来了显著收益，并且提高了第一轮生成内核的质量。 这归功于奖励的信用分配机制：由于早期轮次的输出会影响所有后续改进，模型被激励去生成更高质量的中间内核。消融实验证实了这一点——当移除多轮信用分配后，第一轮性能大幅下降，因为其优势估计不再包含对后续交互的潜在影响。

六、克服惰性优化：从稳定性到有效性

在 KERNELGYM 和 TRLOO 的支撑下，长期强化学习训练变得可行，但“惰性优化”问题依然存在。研究团队通过两个假设系统地研究了这一性能瓶颈。

6.1 假设 1：训练不稳定

最初推测训练不稳定是导致性能过早饱和的原因。 训练与推理之间的不匹配是大语言模型强化学习中的一个基本挑战，其中策略采样（推理）与策略训练之间的差异可能导致离策略漂移，理论上这会引发梯度方差增大和奖励崩溃。

通过监控策略熵、梯度范数和困惑度等指标，发现多轮强化学习确实表现出数值过高和波动，表明训练过程不稳定。

图 7：四种训练方案（TRLOO、TRLOO+MRS、TRLOO+MRS+PR、TRLOO+MRS+PR+PRS）的训练动态对比，从熵值、梯度范数、VLLM-PPL、FSDP-PPL 四个维度进行分析。

为了缓解这种漂移，研究团队采用了几何平均不匹配拒绝采样（MRS）。只有当新策略与旧策略在整条轨迹上的几何平均似然比 ( bar{r}_t ) 处于区间 ([beta^{-1}, beta]) 内时，样本才会被保留用于训练。同时还实施了严格的令牌级否决机制：如果任何单个令牌的似然比低于阈值 ( eta )，则整个序列将被拒绝。

如图 7 所示，MRS 成功地稳定了训练动态。然而，图 5 揭示了一个关键发现：虽然不匹配校正防止了训练早期崩溃（使学习曲线变得平滑），但它并未从根本上提升更严格的 Fast@1.2 性能指标的上限。

图 5：左图显示第 3 回合 Fast@1.2 随训练步数变化，MRS 能稳定训练，但需要 PR 和 PRS 才能显著提升 Fast@1.2 指标。右图显示熵值变化，MRS 提升了稳定性，PR 和 PRS 进一步增强了稳定性。

这表明假设 1 仅部分解释了问题——它说明了训练不稳定性，但仅解决稳定性不足以突破性能瓶颈。

6.2 假设 2：目标错位

标准奖励信号未能区分“微不足道的改进”和“针对真正瓶颈的优化”。一个内核可能正确并带来一定的速度提升，但仍可能未触及主要的性能瓶颈。要使模型从生成“正确”的内核转向生成“有效”的内核，必须让奖励对性能瓶颈敏感。

基于性能分析的奖励（PR）

目标错位的一个关键症状是，模型倾向于优化那些琐碎的、对整体计算性能影响微乎其微的子操作（例如替换简单的求和操作），而忽略了真正的计算瓶颈。

图 10：惰性优化案例与优化融合案例的性能分析反馈对比。惰性优化案例中，模型生成的内核仅占总 CUDA 执行时间的 0.014%（PR=0.00014）；而在优化融合案例中，生成的内核覆盖了 86.15% 的总 CUDA 运行时（PR=0.8615）。

如图 10 的案例研究所示，研究团队利用 KERNELGYM 的性能分析工具包，从程序总 CUDA 执行时间中分离出由模型生成的内核所贡献的运行时部分，并定义性能分析比（PR）为：

[
text{PR} = frac{text{生成内核的 CUDA 运行时}}{text{总 CUDA 执行时间}}
]

直观上，当候选代码优化了主导端到端运行时的核心内核时，PR 值会更高，从而赋予模型更大的奖励信用。随后，用这个信号来增强每一轮的奖励（仅应用于正确的内核）：

[
R_{text{enhanced}} = R_{text{correctness}} + lambda_{text{speed}} cdot text{clip}(S, 3) cdot text{PR}
]

这鼓励模型专注于那些对整体性能有显著贡献的内核优化，明确地将学习方向引导至能带来实质性速度提升的修改上。 由于 PR 值介于 0 和 1 之间，速度奖励项仍占主导，这防止了模型通过生成低效但覆盖率高的代码来滥用该机制。

基于性能分析的拒绝采样（PRS）

为了进一步强化目标对齐，研究团队引入了基于性能分析的拒绝采样。只有那些性能分析比（PR）超过预设阈值的内核样本才会被保留用于策略更新。这有效过滤了那些虽然正确但优化价值低的“惰性”样本，确保训练数据集中包含更多能解决真实瓶颈的高价值优化案例。

6.3 实验结果：PR 和 PRS 的效果

为进一步过滤训练分布，研究团队引入了基于性能分析的拒绝采样（PRS）。对于每个样本，其被保留用于训练的概率由以下公式决定：

[
P_{text{keep}} = expleft(-frac{max(tau – R_{text{bottleneck}}, 0)}{sigma}right)
]

其中 ( tau ) 是性能阈值，( sigma ) 控制滤波器的“柔软度”。实验固定 ( sigma = 0.1 )。

图 5：左图展示了第 3 回合 Fast@1.2 指标随训练步数的变化；右图展示了策略熵值随训练步数的变化。

图 5 的结果验证了阶段性诊断：
* 最小奖励塑形（MRS） 提高了训练稳定性，但未能显著提升更严格的 Fast@1.2 指标上限。
* 性能奖励（PR）与拒绝采样（PRS） 在稳定训练的基础上，通过聚焦性能瓶颈、过滤低价值样本，显著提升了核心优化目标（Fast@1.2）的表现。右图的熵值变化进一步表明，PR 和 PRS 使训练过程更加平稳，探索效率更高。

图 7：四种训练方案（TRLOO、+MRS、+PR、+PRS）在熵值、梯度范数、VLLM-PPL、FSDP-PPL 四个维度的训练动态对比。

七、实验结果：Dr. Kernel 的表现

7.1 主要结果

表 1：不同 Fast 阈值下，各模型在 KernelBench 三个层级上的性能表现。评估标准排除了存在奖励欺骗的样本，比原始基准更严格。

表 1 总结了 Dr. Kernel 与基线模型的关键对比：
* Dr. Kernel 在开源基线中达到最强性能，并在 Level 1 和 Level 2 上与前沿闭源模型（如 GPT-5, Claude-4.5-Sonnet）形成竞争。
* 与冷启动模型相比，多轮强化学习对性能，尤其是更严格指标（如 Level 2 的 Fast@1.2）有实质性提升（从 5.6% 提升至 20.0%）。
* 测试时扩展策略进一步提升了性能：
* 通过上下文管理的顺序测试时扩展（STTS），Dr. Kernel-14B 在 Level 2 的 Fast@1.2 从 25.6% 提升至 31.6%。
* 通过选择历史最优回合（Dr. Kernel-14B-STTS†），Fast@1.2 在 Level 2 上达到 47.8%，在该指标上超越了对比的前沿模型。

7.2 测试时扩展策略

图 6：Dr. Kernel-14B 的两种测试时扩展策略效果对比。左图为最后一回合的 Fast@1.2，右图为历史最优回合的 Fast@1.2。

图 6 对比了两种测试时扩展策略：
* 朴素外推：直接增加提示中的历史回合数，当回合数增加导致提示长度接近模型上下文限制时，性能会下降。
* 上下文管理：将完整历史存储在外部，仅在提示中包含奖励最高的数个回合。该策略能持续利用更多历史信息，性能随回合数增加而稳定提升，最终超越朴素外推。

7.3 torch.compile 下的表现

表 2：在 torch.compile 优化环境下，不同模型在各层级的 Fast 性能表现。

表 2 展示了在已应用了算子融合等编译器级优化的 torch.compile 环境下的评估结果：
* 在此更具挑战性的设置下，所有模型的绝对性能增益空间被压缩，但 Dr. Kernel 在三个层级上仍保持竞争力，并在 Level 1 和 Level 2 的 Fast@1.2 上优于 Claude-4.5-Sonnet。
* 这证明了 Dr. Kernel 的优化能力并非依赖于浅层的、琐碎的更改，而是能够产生超越强优化基线的实质性性能改进。

八、相关工作

8.1 内核生成

AutoTriton 是相关领域的重要工作，它使用强化学习来生成正确的 Triton 内核，但其优化目标并未明确聚焦于速度提升。

如表 1 所示，虽然 AutoTriton 在 KernelBench Level 2 上实现了 30.6% 的 Fast@1（任意速度提升），但在更严格的 Fast@1.2（至少 20% 速度提升）指标上仅为 9.2%。这凸显了将优化目标与最终性能指标对齐的重要性。

8.2 强化学习用于代码生成

近年来，强化学习在代码生成领域取得了显著进展。从早期用于程序合成的深度强化学习，到近期用于优化大语言模型代码生成能力的基于人类反馈的强化学习，强化学习已成为提升代码生成质量的重要工具。

然而，大多数研究工作聚焦于功能正确性，忽视了性能优化这一关键维度。Dr. Kernel 首次系统性地研究了面向性能的强化学习内核生成，并解决了这一过程中特有的奖励破解和惰性优化问题。

8.3 多轮强化学习

多轮强化学习在对话系统和工具使用智能体领域已有探索，但内核生成为其带来了独特挑战：每一轮优化都会产生密集的奖励信号（包括正确性和速度提升），而非稀疏的单一结果奖励。Dr. Kernel 提出的 TRLOO 方法，为此类密集奖励场景下的多轮优化提供了无偏的优势估计。

九、局限性与未来工作

9.1 数据规模与预训练

受限于计算资源，监督微调阶段仅收集了 8,000 个冷启动样本。考虑到当前大语言模型预训练语料中高质量的内核编程数据相对稀缺，该领域对数据规模的要求较高。未来的工作可以进行更大规模的数据收集，以促进领域特定的预训练或持续预训练，为后续的强化学习优化奠定更坚实的基础。

9.2 模型容量

Dr. KERNEL-8B 和 Dr. KERNEL-14B 的实验结果证实，更大规模的模型具有更优的内核生成能力。这种扩展效应在强化学习中尤为关键，因为模型必须依赖自身生成的内容来探索解空间并更新策略。将现有方法迁移到更大参数规模的模型上，有望加速该领域的发展。

9.3 走向生产就绪的自动化

尽管所提出的方法实现了与前沿模型相媲美甚至更优的性能提升，但该领域整体仍处于探索阶段。当前模型能够生成高质量的代码片段，但尚未能在生产环境中实现完全自主的端到端内核生成。希望 KERNELGYM 环境和 Dr. KERNEL 训练框架能够为未来的研究提供助力。

结论

Dr. Kernel 系统性地研究了针对 Triton 内核生成的强化学习方法，识别并解决了奖励破解和惰性优化两大挑战。主要贡献包括：

1. KERNELGYM：一个鲁棒的分布式 GPU 环境，支持破解检查、多轮交互式数据收集和长期强化学习训练。
2. TRLOO：一种无偏的多轮优势估计方法，解决了 GRPO 方法中存在的自包含偏差问题。
3. PR 和 PRS：基于性能分析的奖励机制和拒绝采样策略，有效缓解了惰性优化问题。
4. 实证结果：Dr. KERNEL-14B 在 KernelBench 基准测试中达到了与 Claude-4.5-Sonnet 和 GPT-5 相竞争的性能；通过测试时扩展，其性能得到进一步提升，在 Level 2 测试集上的 Fast@1.2 指标达到 47.8%，超越了前沿模型。

这项工作不仅推动了强化学习在代码生成领域的应用，也为解决强化学习训练中普遍存在的奖励破解和惰性优化问题提供了新的思路。正如论文标题所言，这是“做得正确的” Triton 内核生成强化学习。

对于希望让 AI 真正成为性能优化专家的研究者而言，Dr. Kernel 提供了一个坚实的基础。对于期待 AI 能编写出更高效代码的开发者来说，这项工作的进展无疑令人振奋——或许在不久的将来，我们就能用上由 AI 自动生成的、比手工调优更高效的 GPU 内核。

关注“鲸栖”小程序，掌握最新AI资讯