摘要：本周三，知名 AI 创业公司，曾发布「全球首个 AI 软件工程师」的 Cognition AI 开源了一款使用强化学习，用于编写 CUDA 内核的大模型Kevin-32B。

机器之心报道

编辑：蛋酱、泽南

本周三，知名 AI 创业公司，曾发布「全球首个 AI 软件工程师」的 Cognition AI 开源了一款使用强化学习，用于编写 CUDA 内核的大模型 Kevin-32B。

Kevin-32B 基于 QwQ-32B 在 KernelBench 数据集上使用 GRPO 进行了多轮强化学习训练，实现了超越 o3 和 o4-mini 的顶级推理表现。

对此，机器学习社区表现出了极大的兴趣。有人表示期待 DeepSeek R1 风格的训练方法用来提升代码效率已久，这回终于有人站出来了。

在一篇博客中，Cognition AI 详细介绍了新模型强化学习训练的机制。

代码是一个不断迭代的过程 —— 需要我们编写、执行程序，评估结果，并根据反馈优化代码。大语言模型（LLM）在代码生成方面的最新进展尝试将此过程融入推理阶段，并使用并行采样等方法。虽然这些方法是有效的，但它们依赖于搜索而非实际学习 —— 在这其中模型权重被冻结。

Cognition AI 探索了多轮强化学习，使用来自环境的中间反馈，并屏蔽模型思维以避免在多轮训练中上下文爆炸。

他们提出的模型 Kevin-32B（即 Kernel Devin）在内核生成方面优于前沿推理模型。此外他们通过实验结果表明，与单轮训练相比，多轮训练使模型在自我优化方面更有效。

多轮训练方法

在新模型上，作者使用了 KernelBench，这是一个包含 250 个基于 PyTorch 的经典深度学习任务的数据集。它衡量优化 CUDA 内核替换 PyTorch 算子的能力。训练专注于前两个级别，每个级别包含 100 个任务。级别 1 包含矩阵乘法、卷积和损失函数等基础任务，级别 2 则包含融合算子，再使用这两个级别的 180 个任务进行训练，并留出 20 个任务作为保留集。

在训练过程中，模型会经历一个迭代反馈循环：从生成的内核中提取反馈，并再对其进行优化。如果内核编译失败，错误轨迹会传递给模型并要求其修复。如果编译正确，系统则会测量运行时间并要求模型进一步改进。

初始方法的机制如下。从初始提示开始，先在每个优化步骤后附加思路链、内核和评估信息，再为整个轨迹分配一个奖励（定义为任何内核获得的最高分数），并使用该序列进行训练。

不过，这种方法存在一些挑战：

为了解决上下文长度爆炸的问题，Kevin-32B 丢弃了轨迹中最长的部分 —— 思维链。现在每个 token 将只包含先前生成的内核和评估结果。为了仍然保留上一步思考过程的信息，我们特意要求模型生成其自身思考过程的简要摘要，然后将其传递给后续上下文。

移除推理的思路链。

为了解决样本效率低下的问题，Kevin-32B 选择了一个更具表现力的奖励函数，将内核的细化建模为马尔可夫决策过程，将给定响应的奖励设置为当前内核与所有后续内核得分的折扣总和。至此，每个细化步骤都变成了一个训练样本。

将奖励设为分数的折扣总和。

结果

对于每项任务，作者并行采样了 16 条轨迹，并进行 8 个连续的细化步骤。对于每项任务的正确性或性能，他们将 best@16 定义为所有轨迹的最大值，avg@16 定义为所有轨迹的平均值。

经过 8 个优化步骤，Kevin-32B 在整个数据集上的平均正确率为 65%，显著超越了 QwQ-32B 和前沿模型。它解决了 89% 的数据集，而 o4-mini 和 o3 分别只解决了 53% 和 51%。在整个数据集中，Kevin-32B 实现了 1.41 倍的 best@16 加速比，优于前沿模型。

Kevin-32B 在二级任务上尤其有效，平均正确率为 48%（o4-mini 为 9.6%，o3 为 9.3%）。这表明多轮训练可以提高模型解决更具挑战性、更长期任务的能力。同样，我们能注意到模型实现了 1.74 倍的 best@16 加速比（是 o4-mini 和 o3 为 1.2 倍）。

多轮训练 vs 单轮训练

Kevin-32B 也展现出比 QwQ-32B 和单轮训练模型显著的提升。在 4 个优化步骤下，Kevin-32B 的表现略优于单轮模型，但随着优化步骤增加到 8 个，两者之间的差距进一步扩大。这表明，通过鼓励更积极的优化，多轮训练在串行轴上具有更好的扩展性。

我们可能会想，单轮训练模型是否可以通过采样并行轨迹来实现更好的加速，然而在这种环境下并非如此。在计算预算固定的情况下，即使对于单轮训练模型，多轮推理也比单轮推理更占优势。

奖励黑客攻击

最初的实验使用了规模较小的模型，例如 DeepSeek-R1-Distill-Qwen-7B，这导致了多起奖励黑客攻击事件：

奖励黑客攻击示例

为了防止奖励黑客攻击，作者对响应施加了更严格的格式检查。对于使用 PyTorch 函数或不包含 CUDA 内核的响应，奖励被设置为 0。

我们能观察到，当模型能力与数据集难度之间的差距很大时，就会发生奖励黑客攻击。模型难以解决任何任务，因此当它生成一个被黑客入侵的内核时，它是唯一具有正向优势的操作，因此会得到显著的强化。

无意义和重复的生成

多次运行后，我们可以观察到在大约 35-40 步，模型开始生成重复或无意义的响应。假设这是因为模型偏离了稳定区域，令人惊讶的是，作者偶然发现了一个可以有效预测未来垃圾响应的指标 —— 思维链不是以「OK」开头的响应比例 —— 我们称之为「Not Okay Ratio」。

模型训练越多，CoT 的开头就越不稳定：

"Okay Amigos, so I need to optimize this 3D tensor-matrix multiplication..”

"Okay Holy crap, I need to get this code optimized….”

在精炼步骤的第 8 遍时：

"Okay SHIT, this is getting frustrating. Let me see. The error is about …”

为了解决这个问题，作者尝试使用 KL 系数 0.001 和 0.01 进行运行，但发现这会减慢学习速度，同时又不能防止垃圾数据的产生。最后，作者通过使用 Dr. GRPO 提出的恒定长度损失正则化（这显著降低了梯度范数），并将梯度范数大幅削减至 0.05，将垃圾数据的出现延迟到了第 100 步。

示例内核

作为示例，以 1 级任务 40（LayerNorm）的轨迹为例，作者重点介绍了模型如何迭代地整合来自内核评估的反馈，并利用这些反馈来改进其生成过程。

步骤 1：模型生成一个正确的内核。它使用共享内存，并将均值 / 方差计算与归一化融合。这实现了 0.6 倍的加速。

__shared__ float shared_sums [512]; // 256 threads * 2 values

for (int s = block_size / 2; s > 0; s >>= 1) {

if (tid

shared_sums [2 * tid] += shared_sums [2 * (tid + s)];

shared_sums [2 * tid + 1] += shared_sums [2 * (tid + s) + 1];

}

__syncthreads ;

}

for (int d = tid; d

float x_val = x [n * D + d];

float y_val = (x_val - mean) /sqrt_var;

y_val = y_val * scale [d] + bias [d];

y [n * D + d] = y_val;

}

步骤 2：模型从评估反馈中注意到「当前的加速比是 0.6，这意味着它实际上比 PyTorch 版本慢。这不太好。让我想想瓶颈可能在哪里」。它将大部分推理 token 用于寻找更合适的块大小，并相应地修改代码。修改是正确的，加速比现在达到了 4.27 倍。

步骤 3 和 4：模型尝试进行更积极的优化。它考虑了循环展开和减少库冲突，但最终选择了 Warp 级内联函数。但最终失败了 —— 首先是因为归约逻辑存在正确性问题，其次是因为思路太长。

步骤 5：模型注意到 Warp 归约实现中存在错误的问题并进行了修复。它成功实现了两级 Warp 归约。最终加速比达到了 9.61 倍。

// Warp-level reduction using shuffle instructions

for (int delta = 1; delta

float other_sum = __shfl_xor_sync (0xFFFFFFFF, warp_sum, delta);

float other_sum_sq = __shfl_xor_sync (0xFFFFFFFF, warp_sum_sq, delta);

warp_sum += other_sum;

warp_sum_sq += other_sum_sq;

}

__shared__ float sum_warp [32];

__shared__ float sum_sq_warp [32];

__shared__ float results [2]; // [mean, inv_std]

if (warp_id == 0) {

sum_warp [warp_lane] = warp_sum;

sum_sq_warp [warp_lane] = warp_sum_sq;

}

__syncthreads ;

// Final reduction within the first warp (tid 0-31)

if (tid

float my_sum = sum_warp [tid];

float my_sum_sq = sum_sq_warp [tid];

// Reduce within the first warp (32 threads)

for (int s = 16; s >= 1; s >>= 1) {

my_sum += __shfl_xor_sync (0xFFFFFFFF, my_sum, s);

my_sum_sq += __shfl_xor_sync (0xFFFFFFFF, my_sum_sq, s);

}

...

}

训练设置

这里使用的是组相对策略优化算法（GRPO），该算法由 DeepSeek 提出，是近端策略优化算法（PPO）的一种变体。

GRPO 不使用值网络来估算基线和计算优势，而是将从同一提示中采样的响应组内的奖励归一化。

作者使用 vLLM 进行推理，使用 DeepSpeed Zero-3 卸载优化器状态。每批训练 8 个任务，每个任务 16 个轨迹。我们使用 GRPO，每批 2 个梯度步骤。基础模型是 QwQ-32B。

响应生成完成后，每个 GPU 将其 vLLM 引擎卸载到 CPU 内存，并评估其生成的内核。对于每个响应，作者都会检查响应格式是否正确，并提取 CUDA 内核。然后，编译并执行代码，使用随机张量测试其正确性。如果正确，则会对内核的运行时间进行剖析。

通过正确性检查的响应将获得 0.3 的奖励，而额外的性能奖励则相当于与参考实现相比所获得的速度提升。

内核评估和基准问题

作者对评估进行了沙盒化处理，以免出现致命错误（如 CUDA 非法内存访问），导致训练过程崩溃。

由于 KernelBench 中的许多任务使用非常小的输入张量，该基准最终测量的是内核启动开销而非实际内核执行时间。为了解决这个问题，作者扩大了受影响任务的张量维度。

KernelBench 评估工具中还有一个更隐蔽的错误，导致被测内核将参考实现的输出张量作为自己的张量输出循环使用。因此，只计算（正确）部分输出张量的内核仍能通过正确性检查。为了解决这个问题，我们首先运行测试过的内核，然后再运行参考实现，从而避免了这一问题。

单轮训练设置

作者使用 max_grad_norm = 0.5、lr = 常量 2e-6（热身比为 0.03）、max_prompt_length = 8192、max_response_length = 16384，其使用 DAPO 的 Clip-High，eps_high = 0.28，并将 KL 系数设为 0，以允许模型自由偏离基本策略。

我们能观察到，单轮模型比基础模型有明显改善，但在 25 步之后，奖励开始趋于稳定。

多轮训练设置

作者根据每个轨迹的最佳内核对其正确性和性能进行评分。在最终训练运行中，每个前向路径由 16 个平行轨迹组成，每个轨迹有 4 个细化步骤，折扣系数为 0.4。与单轮训练不同的是，现在的奖励会稳步增加。

随着模型学会更高效地使用推理 token 来生成内核，响应长度最初会减少。第 25 步之后，随着模型尝试更复杂的解决方案，响应长度会增加。按照 DeepScaleR 的做法，作者在第 30 步将最大响应长度从 16K 扩展到 22K token。

更多结果 & 消融研究

推理时间缩放

作者还研究了多轮模型在并行和串行轴上的扩展。在推理时间的第一次实验中，其使用了 16 个并行轨迹和 8 个细化步骤。作者再次发现，随着细化步骤的增加，多轮模型的扩展性更好。

在第二次实验中，作者将并行轨迹的数量增加到 64 个，同时只保留 4 个细化步骤。这样，best@64 的正确率达到 89.5%，性能提高了 1.28 倍，比 8 个细化步骤的 best@16 稍差。

作者研究了沿并行轴或串行轴缩放推理的效果。此处使用 pass@k 性能来表示 k 生成的估计性能。他们使用类似于 Chen 等人的无偏估计器计算该指标，其方差低于 avg@k。

然后，作者试图找到一个合适的定律来模拟实验数据，并注意到在这个（小）数量级上，细化步骤和平行轨迹的贡献看起来都像一个幂律。此外，由于内核加速是有限的，因此性能指标应该达到饱和。因此，作者决定拟合以下定律（该定律在小数量级时呈现幂律表现，随着计算量的增加，收益会逐渐减少）：

作者发现，给定一个固定的、非微不足道的推理计算预算（例如，细化步骤 * 并行轨迹≥8），最佳计算分配会转向串行细化，而不是并行生成。

单轮模型推理

作者之前在相同的多轮推理设置中比较了多轮模型（Kevin-32B）和单轮模型。但由于单轮模型是在单轮环境下训练的，

因此，一些问题就产生了：在单轮环境下训练的模型计算量固定的情况下，我们能用单轮推理或多轮推理获得更好的推理结果吗？

在这种特定环境下，即使使用单轮训练模型，多轮推理的结果一般也优于单轮推理（平均正确率除外）。

为了比较这两种方法，作者使用 64 个并行轨迹和仅 1 个步骤来评估单轮模型，然后将结果与使用 16 个并行轨迹和每个轨迹 4 个细化步骤的多轮推理进行比较。作者将单轮推理的 64 个并行轨迹分成 16 组，每组 4 个内核，取每组的 best@4，然后取 16 组的平均值。这样就能将这一指标与多轮推理中的 avg@16 进行比较（因为在这种情况下，我们是在单个轨迹中取 best@4）。最后，作者将单轮推理的 best@64 与多轮推理的 best@16（4 个细化步骤）进行了比较。

奖励塑造

作者尝试了奖励塑造。在较小模型上运行时，其添加了中间奖励（成功解析、编译、执行......）来引导模型。然而，作者发现这些奖励可能会分散模型的注意力，使其无法朝着真正的目标更新 —— 生成正确且性能良好的内核。作者还按照 Kimi 的建议尝试了长度惩罚，但发现在已有的环境中，长度惩罚会降低模型的性能。

对于多轮训练，作者对不同的奖励函数进行了消融。他们尝试了不同的伽玛值（0.4 与 0.8），以及如何汇总单个轨迹的奖励 —— 求和或取最大值。

在内核生成中，作者从根本上关心的是获得具有最大轨迹的内核（而不是优化多个内核的分数折扣总和）。因此，作者认为使用最大奖励公式会带来更好的加速效果。

然而，我们可以发现，将整个 MDP 的奖励相加，gamma=0.4 的效果最好。

奖励塑造消融 — 多轮训练是在奖励总和与伽马值 = 0.4（sum_gamma_0_4）的条件下进行的。

并行轨迹

作者还尝试在训练过程中将并行轨迹的数量从 16 个增加到 64 个，即训练批次大小为 64 * 4 * 8 = 2048。理论上讲，这将使优势估计的噪声更小，但他们发现在评估的 best@16 和 avg@16 性能上没有明显差异。此外，在第 30 步（60 梯度步）左右，无用的东西开始提前出现。

数据分布

在项目初期，作者尝试在 DeepSeek-R1-Distill-Qwen-14B 上运行 level 1 任务的简单子集。实验发现，奖励会出现高原现象，模型会过拟合单一难度级别。此外，模型只能学习有限的优化技术。因此，作者认为在数据集中均衡、多样地分布不同难度的任务非常重要。

测试时搜索

在训练过程中，为了保持合理的训练时间，必须限制搜索。但在测试时，可以自由地使用更复杂的测试时间技术来进一步提高性能。

为此，Kevin-32B 使用了改进版的束搜索，其工作原理如下。作者首先在 16 条轨迹上执行 4 个串行细化步骤，就像在训练时间所做的那样。在这一过程结束时，根据生成的最快内核对轨迹进行排序，并保留最好的 4 个轨迹。然后，作者将每条轨迹复制 4 次（共 16 条轨迹），并重复这一过程。作为一种通用方法，它还显著提高了模型性能，在整个数据集上平均提速 1.56 倍。

不出所料，随着测试时间计算量的增加，得到的收益也在减少，尽管平均性能在几个小时后仍有提高。

束搜索测试时间推理。计算时间是在单个 H200 GPU 上按内核计算的。请注意，其中还包括评估时的停滞时间，有效推理时间约占总计算时间的 45%。

作者通过修改轨迹数、束宽度和细化步数，尝试了波束搜索的几种不同变体。有趣的是，他们发现中位加速度较高的设置平均加速度较低，反之亦然。

作者将这种行为归因于每种技术在探索 / 利用的帕累托前沿中的不同位置：通过更多的探索，内核可以实现一致的加速，尽管它们可能会错过非常激进的、显著提高平均速度的优化。

未来工作

我们相信，这项工作只是探索编程智能体训练方法的开始。如果有更多的时间和计算能力，我们希望尝试以下方法：

总结

在这项工作中，作者提出了一种方法，该方法适用于任何具有中间奖励的多轮环境。研究表明，这种方法比单圈 GRPO 取得了更好的结果。

端到端训练将是未来智能体的重要组成部分。虽然手工制作的多 LLM 工作流程可以带来短期收益，但它们在很大程度上依赖于人类的启发式方法，而且无法扩展。相比之下，更通用的方法可以让模型自由探索不同的轨迹（搜索），然后学习这些长期动态。通过与编程环境的交互，模型产生经验流，并通过反馈不断调整。我们希望这样的方法能成为迈向自主编程智能体的第一步。

HuggingFace: https://huggingface.co/cognition-ai/Kevin-32B

最后，让我们来复习一下《The Bitter Lesson》的道理：

我们必须吸取惨痛的教训，从长远来看，在我们的思维方式中植入知识是行不通的。这一惨痛教训基于以下历史观察：

来源：机器之心Pro一点号