KernelZero: Co-Evolving Proposer and Coder for Continuously Improved GPU Kernel Generation
Changxin Ke, Rui Zhang, Zixiang Fang, Zhenghong Li, Yuanbo Wen, Jiashuo Shen, Shuo Wang, Jiaming Guo, Ling Li, Qi Guo, Yunji Chen
cs.LG, cs.SE
2026-09-27
出题模型按写核模型当前弱点持续造训练题，共进化得到的7B模型在CUDA一级与二级上的一次通过率分别达到75.8%和69.6%。
用大模型写 CUDA、Triton kernel 已经能在 KernelBench 上交卷，训练成本和目标冲突还没解开。cudaLLM 用 128 张 H100 做监督训练，8B 模型在 Level 1 的 pass@10 仍只有 81%。现成语料要么对当前模型太简单，RL 刷不到边界；要么太难，梯度几乎是噪声。算力砸下去，能力却不跟着长。
奖励设计把问题再拧紧一圈。只奖励正确，模型会写出几乎不加速的保守实现；过早按 speedup 给分，功能等价又先崩掉。KernelZero 要同时处理这两件事：持续造出刚好卡在当前能力边界上的训练模块，以及把「先写对、再写快」写进强化学习的门控里。
系统是两个专门化的 7B，都从 Qwen2.5-Coder-7B 初始化。Proposer 是出题模型，吃一组 Torch API，生成可运行的 nn.Module；Coder 是写核模型，把模块翻译成 CUDA 或 Triton kernel。两边用各自奖励轮流更新、对方冻结，4 轮、每轮各 20 步，总共 80 步。两边不同时更新，是为了让出题分布在一段窗口内保持稳定，Coder 才有可学的目标；出题端则拿固定 Coder 的失败样本当边界信号。
出题侧先按真实 Torch 模块的 API 共现分布采样。题库 4000 组，两千组两算子、两千组三算子，覆盖 442 个候选 API。生成结果过三道闸：AST 核对是否用上了指定 API；调用 getinitinputs / getinputs 在 GPU 上真跑；出现 NaN 或 Inf 直接丢。过关才算有效模块。
Frontier Reward 只看 Coder 对这道题的组内平均正确率，对应算法叫 FO-GRPO。成功率靠近 0.5 时奖励最高，太简单或全军覆没都往下掉；无效模块固定罚 −1。出题模型被推去找「一半对、一半错」的题，简单组合拿不到高分。
写核侧先做编译器启发式的冷启动。从 cudaLLM 公开的 71996 条实例出发，用 gpt-oss-120B 当老师，按 Tiling、Fusion、Pipeline、Reordering 四条原则写出带推理过程的 kernel，只留功能正确的样本，得到 42454 条 CUDA 和 59998 条 Triton，再做 SFT。这四条分别对应切块复用、算子融合、访存与计算重叠、访问模式重排，冷启动就把常见优化套路灌进去。
CA-GRPO 把正确性和速度拆成两段信号。一组 rollout 的正确率不到阈值 α（取 0.5）时，只给 0/1 对错分；过线之后，正确样本才按组内 min-max 归一化的 speedup 加权，系数 β=0.1。全组速度相同则性能项为 0。错的样本始终是 0，不会因为很快但算错而拿到正奖励。
训练在 A100-80GB 上。更新 Proposer 时要挂着冻结的 Coder 和执行服务 Keck，一共 12 卡，20 步大约 3.1 小时；更新 Coder 用 8 卡，大约 4.5 小时。Proposer 实际用 400 条 API 列表 prompt，组大小 4，每道题再嵌套抽 5 个 Coder 样本；Coder 用 1224 个过检模块，组大小 5。
评测覆盖 KernelBench Level 1 的 100 个单算子和 Level 2 的 100 个融合模式，one-shot。执行器用自研 Keck：一卡一次只跑一个候选，数值容差绝对、相对都是 10⁻²，CUDA 禁止调 cuBLAS/cuDNN，Triton 必须落到 @triton.jit。
| 模型 | CUDA L1 pass@1 | CUDA L2 pass@1 | Triton L1 pass@1 | Triton L2 pass@1 |
| Claude-4.5-Sonnet | 76.4 | 66.5 | - | - |
| cudaLLM-8B | 72.5 | 67.4 | - | - |
| DeepSeek-V4-Pro | - | - | 49.7 | 63.4 |
| Dr.Kernel-8B（三轮） | - | - | 34.4 | 73.0 |
| KernelZero-SFT-7B | 70.6 | 59 | 69.4 | 64.9 |
| KernelZero-7B | 75.8 | 69.6 | 77.2 | 72.5 |
CUDA Level 1 上，KernelZero 的 pass@1 比 Claude-4.5 低 0.6 个百分点，pass@10 到 100%（Claude 为 97%），Level 2 的 pass@1 高 3.1 点，pass@10 为 97% 对 93%。Triton Level 1 一次通过 77.2%，DeepSeek-V4-Pro 是 49.7%，GLM-5.2 是 66.7%；Level 2 的 72.5% 略低于 Dr.Kernel-8B 三轮推理的 73.0%，但对方没报 pass@5/10。CUDA 输出大约 2.8k/4.5k token，短于 Kevin-32B 的 9.2k/13.7k。
加速数字在两个后端上差得很远。fast1@1 统计的是一次采样既正确、又对 PyTorch Eager 加速比大于 1 的比例。Triton Level 1 的 fast1@1 为 43.9%、fast1@10 为 87%；同设置 CUDA 只有 17.6% 和 29%。Level 2 更明显：Triton 65.4% / 96%，CUDA 2.4% / 12%。生成的 CUDA 不调用厂商库，PyTorch 参考实现却会把卷积和 GEMM 丢给 cuDNN/cuBLAS，这条基线本身就很难超。
从第 40 步冻结 Proposer、只继续训 Coder，CUDA Level 2 的 pass@1 停在 64.4%，完整共进化到 69.6%；Triton Level 2 从 68.3% 到 72.5%。Level 1 增益更薄，CUDA 只有 0.9 点。α 消融里，永远不给速度奖励（α=1.1）的 pass@1 最高，达到 74.85%，fast1.5@1 只有 9.75%；α=0.9 把后者拉到 11.00%。α=0.5 在六个代表算子的 best-of-10 速度上都最好，Tensor-MM 从 SFT 的 0.12× 到 4.42×，GroupNorm 1.25×，Conv-BN-Scale 1.16×。
在 KernelBench 这种算子和短融合题上刷正确率，7B 加针对性 RL 已经能摸到 Claude-4.5 和专用 8B、32B 的区间，推理比 Kevin 短一截。冷启动仍然吃 cudaLLM 公开语料，SFT 基线 CUDA Level 1 就有 70.6%，RL 再加 5.2 点。共进化的增量主要落在更难的 Level 2，不是从零长出写核能力。
想替换生产里的卷积和 GEMM，CUDA 这条线现在追不上 cuDNN。Triton 更接近能用，因为它的编程模型和编译器已经把 tiling、fusion 接住了。出题-写核轮流更新这个结构，可以挪到其他训练题必须贴着当前能力的代码生成任务，论文只在 KernelBench 上给出了证据。
正文没有独立的 Limitations 节，缺口得从实验设置里抠。
KernelBench 共 250 题、三档，只报了 Level 1 和 Level 2 的 200 题，第三档没有分数。评测也不是原版脚本，Keck 改了参数初始化对齐、单卡隔离和 10⁻² 容差，跟 cudaLLM、Kevin 的表不能当成同一套 harness 的严格对比。
摘要里「超过 Claude-4.5」只覆盖 CUDA 的 pass@5/10 和 Level 2 pass@1，Level 1 一次通过其实低 0.6 点。Dr.Kernel 是三轮，KernelZero 是 one-shot，Triton Level 2 的 72.5% 对 73.0% 不宜直接论输赢。CUDA Level 2 几乎加速不动，是禁止库调用撞上了厂商库基线，不能读成 Coder 在融合题上突然变差。
α=1.1 的正确率更高，选 0.5 是在用正确率换速度，论文没有把这条权衡画成完整 Pareto。训练仍要 8 到 12 张 A100 加在线编译执行，并不是零标注、零算力的自举。