KernelZero: Co-Evolving Proposer and Coder for Continuously Improved GPU Kernel Generation
KernelZero: Co-Evolving Proposer and Coder for Continuously Improved GPU Kernel Generation
发表时间: 2026-09 · arXiv:2609.33074 (ICT, CAS)
原文: https://arxiv.org/abs/2609.33074
作者/机构:Changxin Ke$^{1,2}$, Rui Zhang$^{1,2}$, Zixiang Fang$^{1,2}$, Zhenghong Li$^{1,2}$, Yuanbo Wen$^{1,2}$, Jiashuo Shen$^{1,2}$, Shuo Wang$^{1,2}$, Jiaming Guo$^{1,2}$, Ling Li$^{2,3}$, Qi Guo$^{1,2}$, Yunji Chen$^{4,1,2*}$
$^{1}$中科院计算所 处理器芯片全国重点实验室 (State Key Lab of Processors, Institute of Computing Technology, CAS)
$^{2}$中国科学院大学 (University of Chinese Academy of Sciences)
$^{3}$中科院软件所 智能软件研究中心 (Intelligent Software Research Center, Institute of Software, CAS)
$^{4}$中国科学院 产业人工智能研究所 (Institute of AI for Industries, Chinese Academy of Sciences)
速读
一句话结论 KernelZero 提出了由生成算子组合的 Proposer 和编写底层算子的 Coder 构成的协同进化框架,利用能力边界驱动的数据生成与正确性门控强化学习,使 7B 规模模型在 GPU 算子生成任务上超越了千亿级通用与专用模型。
要解决什么问题 利用大语言模型自动生成兼具正确性与硬件高效性的 GPU 算子面临两个核心卡点。首先是缺乏与模型当前能力动态对齐的高质量训练数据:静态数据容易导致训练无法收敛或成本极高,算子过易则无效、过难则学不会。其次是算子正确性与运行性能之间存在天然权衡:若仅优化正确性,模型倾向于生成极度保守且无加速比的代码;若直接引入激进的性能优化奖励,则极易破坏功能等价性,引发训练早期不稳定、输出错误算子。
怎么做的 KernelZero 将算子生成建模为两个模型交替优化的自动化课程体系。核心思路是用 Proposer 动态探测 Coder 的弱点并生成处于其能力边界的 PyTorch 模块,再由 Coder 将其翻译为高性能 CUDA 或 Triton 算子,绕开了对人工标注数据的依赖与性能正确性失衡的陷阱。该系统包含三大关键部件:第一,能力边界驱动的 Proposer。从真实代码的 API 共现分布中采样 API 集合,生成计算图模块后通过抽象语法树(AST)静态检查、动态可运行检查和数值稳定性过滤(剔除 NaN/Inf)。随后采用边界导向强化学习(FO-GRPO),以 Coder 的平均正确率 $\bar{l}$ 构造边界奖励:
当 Coder 采样生成的算子成功率接近 50% 时奖励最高,强迫 Proposer 持续生成既不过难也不过易的进阶任务。第二,融入编译知识的 Coder 冷启动蒸馏。在强化学习前,让大模型基于平铺(Tiling)、算子融合(Fusion)、流水线重叠(Pipeline)和循环访存重排(Reordering)四项专家启发式原则进行链式思考推理,生成冷启动微调数据集。第三,正确性感知强化学习(CA-GRPO)。引入组级门控机制,只有当采样组 $\mathcal{G}$ 的整体正确率 $\text{Acc}_{\mathcal{G}} \ge \alpha$ 时(阈值 $\alpha=0.5$),门控函数 $\mathcal{T}(\mathcal{G})$ 才置为 1,向正确算子赋予组内归一化的加速比得分 $S_i^{\mathcal{G}}$:
效果如何 实验基于 Qwen2.5-Coder-7B 进行初始化,在 4 块或 8 块 A100-80GB GPU 上迭代 4 轮共 80 步,在包含 250 个任务的 KernelBench(涵盖单算子 Level 1 与融合算子 Level 2)上进行评估。对比基线涵盖通用大模型(Claude-4.5-Sonnet、GPT-oss-120B、DeepSeek-V4-Pro-1.6T)以及专用算子模型(由 128 块 H100 训练的 cudaLLM-8B、多轮强化学习模型 Kevin-32B、Dr.Kernel-8B/14B)。在 CUDA 生成上,KernelZero-7B 在 Level 1 和 Level 2 取得 75.8% 和 69.6% 的 pass@1,pass@10 达到 100% 和 97%,整体 pass@1 逼近甚至超越 Claude-4.5-Sonnet(76.4%/66.5%),并在 2 倍加速比指标($\text{fast}_2@1$)上超过 cudaLLM-8B 与 Kevin-32B;在 Triton 生成上,Level 1 和 Level 2 的 pass@1 分别为 77.2% 和 72.5%,大幅领先 DeepSeek-V4-Pro(49.7%/63.4%)。局限在于,由于生成的 CUDA 算子均为手写原生代码、禁止调用外部闭源库,面对 Level 2 中由 cuBLAS/cuDNN 强力优化的 GEMM 和卷积等算子时,难以取得显著超越 PyTorch 底层原生库的加速比(Level 2 下 CUDA 的 $\text{fast}_2@1$ 仅为 1.2%)。
1. 主要贡献
高性能 GPU 算子(Kernel)是现代机器学习系统的基石,但自动化生成兼具正确性与高效执行性能的 GPU 算子仍具极大挑战。针对大语言模型(LLM)在底层算子生成领域所面临的与模型当前能力相匹配的高质量训练数据极度匮乏以及算子正确性与加速性能之间的天然权衡(Trade-off)两大核心瓶颈,本文提出了 KernelZero,一个通过两个专有模型协同进化(Co-evolution)以实现 GPU 算子生成能力持续提升的强化学习框架。
┌─────────────────────────┐
│ sampled API set │
└────────────┬────────────┘
│
▼
┌───────────────────────┐
│ Proposer P │◄───────────┐
└───────────┬───────────┘ │
│ │ FO-GRPO
▼ module m │ (Frontier Reward Rp)
┌───────────────────────┐ │
│ Validation Pipeline │ │
│ (AST, Run, Stability)│ │
└───────────┬───────────┘ │
│ valid m │
▼ │
┌───────────────────────┐ │
┌───────────────────►│ Coder C │────────────┘
│ └───────────┬───────────┘
│ CA-GRPO │ rollout kernels {k_1, ..., k_n}
│ (Correctness-Gated ▼
│ Reward R_c) ┌───────────────────────┐
└────────────────────┤ Keck Validation Server│
│(Isolated Exec & Profil│
└───────────────────────┘
本文的核心创新点与技术贡献包括:
- 双模型协同进化框架(Co-Evolution Framework):构建由提议器(Proposer)与编码器(Coder)构成的博弈协同系统。Proposer 负责从 PyTorch API 集合合成计算图模块,Coder 负责将其翻译为底层的 CUDA 或 Triton 硬件算子,二者通过明确的接口交替迭代。
- 前沿驱动的模块生成机制(Frontier-Driven Module Generation):Proposer 依据 Coder 在前一轮中的弱点(失败用例涉及的 API 分布)自适应采样算子集合,并由基于前沿胜率的奖励函数引导,精准生成处于 Coder 当前能力前沿(Pass 率约 50%)的 Torch 模块,既避免了无意义的简单样本,又防止了过难样本导致的训练发散,实现无人工标注的自适应动态数据流。
- 正确性感知组相对策略优化(Correctness-Aware GRPO, CA-GRPO):提出组级别门控奖励机制。只有当候选算子组的整体正确率达到预设阈值 $\alpha$ 时,才激活基于相对加速比的性能优化信号,根本上消除了早期训练因追求极致加速而破坏算子功能等价性的退化现象。
- 编译指导的冷启动蒸馏(Compiler-Inspired Cold-Start Distillation):系统化归纳分块(Tiling)、融合(Fusion)、流水线(Pipeline)与重排(Reordering)四项 GPU 算子设计准则,用于指导教师模型生成具显式推理链(CoT)的高质量冷启动数据。
- 构建高吞吐评测基座 Keck:设计了隔离运行、多进程看门狗监控、权重状态严格对齐(
load_state_dict)且支持 CPU 异步编译的评测服务基础设施 Keck,保障了奖励反馈的高保真度。 - 超越超大规模前沿模型的评测表现:参数量仅 7B 的 KernelZero-7B 在 KernelBench 基准上超越了 Claude-4.5-Sonnet(CUDA 后端)和 DeepSeek-V4-Pro(Triton 后端)。在 CUDA Level 1 和 Level 2 任务上分别取得 $75.8\%$ 和 $69.6\%$ 的 pass@1,pass@10 分别达 $100\%$ 与 $97\%$;在 Triton Level 1 和 Level 2 上分别取得 $77.2\%$ 和 $72.5\%$ 的 pass@1。
2. 背景知识、关键观察与设计原则
2.1 现有算子生成范式的局限与观察
- 高质量能力对齐数据极度匮乏:高质量底层硬件代码天然稀缺,人工标注代价极高。随机或机械合成的高层 PyTorch 代码往往脱离实际业务计算图模式,导致 SFT 或经典 RL 要么因缺乏难度阶梯而无法收敛,要么耗费数以百计的高端算力卡(如 cudaLLM 使用 128 张 H100 训练,在 KernelBench Level 1 上 pass@10 仅为 81%)。
- 正确性与执行性能的天然矛盾:生成低级硬件代码时,过度关注正确性会导致模型倾向于生成保守、未充分利用硬件并行能力的浅层实现;反之,若在强化学习早期强行引入运行时加速比奖励,模型极易通过省略边界检查、错误削减同步屏障或直接破坏逻辑以减少时钟周期,导致严重的算子功能错误。
2.2 四大 GPU 算子设计核心原则
为了在冷启动阶段消除高层抽象模型与底层硬件微架构之间的认知鸿沟,作者将现代 GPU 编译器与资深系统程序员的优化经验归纳为互补的四大体系结构优化原则:
- 分块(Tiling):将输入多维张量切分为能够驻留在片上存储(如寄存器堆、共享内存 Shared Memory)的 Block 级与 Warp 级分块。该设计最大化了片上数据局部性复用,成倍减少高延迟且昂贵的全局内存(Global Memory/DRAM)访问,同时便于利用 Warp 内洗牌指令(Shuffle)完成 Block/Warp 级的高效归约。
- 算子融合(Fusion):将存在生产者-消费者依赖关系的多个连续计算算子串联并合并进单一内核中。算子产生的中间特征直接保存在寄存器或共享内存中,彻底消除了写回 DRAM 再重新加载的往返开销,大幅降低内核启动延迟(Kernel Launch Overhead)与设备端全局同步开销。
- 流水线编排(Pipeline):将内存访问与算术计算组织成多阶段流水重叠执行。利用硬件异步拷贝(如 Ampere 架构异步全局内存至共享内存数据搬运)使数据预取与算术运算重叠,隐藏访存长延迟,显著拉高流式多处理器(SM)的持续利用率。
- 访存重排(Reordering):重新调整循环迭代次序、执行顺序及线程-数据之间的映射拓扑,确保全局内存访问达成对齐的合并访存(Coalesced Memory Access),并在共享内存访问中规避存储体冲突(Bank Conflicts)与分支分化(Warp Divergence)。
3. 方法细节
3.1 KernelZero 整体架构
算子生成的端到端协同建模:KernelZero 将底层代码生成抽象为提议器 $P$ 与编码器 $C$ 之间的协同进化框架。提议器输入一组 Torch API 集合 $a \in \mathcal{A}$,输出一个结构化的高层计算图模块 $m \sim \pi_P(\cdot \mid a; \theta_P), m \in \mathcal{M}$,其中 $\mathcal{M}$ 为语法合法且能实际执行的 PyTorch 模块空间。随后,编码器 $C$ 接收模块 $m$ 作为高级规格说明,将其生成为底层 CUDA 或 Triton 硬件内核代码 $k \sim \pi_C(\cdot \mid m; \theta_C), k \in \mathcal{K}$,其中 $\mathcal{K}$ 表示可执行硬件算子空间。由此建立起从高层抽象到极致性能的两阶段动态映射:
$$ a \xrightarrow{P} m \xrightarrow{C} k $$其中 $P$ 负责在广泛的算子组合空间中开发展示多样性的前沿挑战,而 $C$ 则学习如何将其高效地在物理硬件底层落地。
双模型交替优化机制:整个训练过程采用交替优化策略。Proposer 与 Coder 依次轮流更新:在优化某一模型时保持另一模型权重冻结,各自利用其专有奖励(Proposer 使用前沿奖励 $r_P$,Coder 使用感知正确性的奖励 $r_C$)通过组相对策略优化(GRPO)进行更新。这种机制构成了自动递进的自适应课程学习(Automatic Curriculum):Proposer 伴随 Coder 能力边界的外推持续生成处于挑战边缘的新模块分布,Coder 则在该演化的数据分布驱动下稳步提升内核正确性与硬件加速比。
3.2 提议器(Proposer):提出前沿 Torch 模块
真实 API 列表共现采样(API List Sampling):为确保生成的模块既符合真实深度学习计算图的拓扑模式,又具备足够的组合多样性以增强 Coder 的泛化能力,系统从大规模真实 PyTorch 开源模块中统计并构建了 API 共现概率分布。在生成开始前,系统从该共现分布中采样 API 子集 $a \subset \mathcal{A}$ 提供给 Proposer,约束 Proposer 必须将采样到的 API 完整合成为一个完整的复合计算图 $m$。
严格的三阶段验证流水线(Validation Check):为确保流入 Coder 训练池的模块具有绝对的可靠性,避免无效或病态计算图干扰 Coder 策略,所有由 Proposer 采样的候选模块必须串行通过三道校验关卡:
- 抽象语法树语义检查(AST Check):利用静态语法解析器遍历模块 AST,验证代码语义的合规性,严格检查生成的计算图中是否无遗漏地包含了所分配的采样 API 集合 $a$。
- 硬件可执行性检查(Runnable Check):依据标准评测框架规范,生成的模块必须实现用于张量初始化的
get_init_inputs方法和用于生成随机测试输入的get_inputs方法。系统在物理 CUDA 显卡上实例化该模块并注入随机张量运行前向传播;若发生编译错误、形状不兼容或运行时异常抛出,该模块被立即丢弃。 - 数值稳定性检查(Numerical Stability):向模块连续注入多组随机张量,严密监控前向计算结果;一旦产生
NaN、Inf等非正常数值,即被判定为不稳定样本予以剔除。只有全部通过上述三道工序的样本才被接纳为合法模块。
前沿驱动的组相对策略优化(Frontier-Oriented GRPO, FO-GRPO):为使训练具有最大信息增益,Proposer 不应生成 Coder 能轻易解决的简单算子,亦不能生成超出其理解极限的不可行算子。给定合法模块 $m$,系统调用当前固定的 Coder 采样展开一组算子实现 $\mathcal{G} = {k_1, \ldots, k_K}$。对其中每个硬件算子 $k_i$,结合 Keck 验证服务评估其功能等价性 $\text{Corr}(m, k_i) \in {0, 1}$,并计算组平均正确率 $\bar{l}$:
$$ \bar{l} = \frac{1}{K} \sum_{i=1}^K \mathbf{1}[\mathrm{Corr}(m, k_i) = 1] $$基于此,作者设计了针对 Proposer 的前沿奖励函数 $R_p(m)$:
$$\begin{aligned} R_p(m) = \begin{cases} 1 - 2 |\bar{l} - 0.5|, & m \text{ is valid}, \\ \rho, & m \text{ is invalid}, \end{cases} \end{aligned}$$其中 $\rho < 0$ 为对不合法模块的恒定惩罚系数(实验设定为 $\rho = -1$)。当 Coder 采样算子的正确率恰好处于 50% 附近时(即该模块处于 Coder 当前能力的模糊临界前沿,既不平凡也非绝对无法突破),$R_p(m)$ 取得极大值 $1.0$;一旦模块过难($\bar{l} \to 0$)或过易($\bar{l} \to 1$),奖励均线性衰减至 0。该机制驱使 Proposer 始终紧咬 Coder 的能力边界进行针对性出题。
3.3 编码器(Coder):生成高性能底层算子
基于体系结构准则的冷启动蒸馏(Cold Start Distillation):在强化学习探索前,必须使基础模型建立底层的并行编程与硬件映射逻辑。作者以大型开源算子集 cudaLLM 中的 71,996 个算子实例为源头,使用参数规模达 120B 的开源大模型 gpt-oss-120B 作为教师模型进行重采样。在 Prompt 中显式注入四大编译优化原则(分块、融合、流水线、重排),驱动教师模型在生成代码前输出密集的思维链(Chain-of-Thought, CoT)推理过程,细致规划硬件资源分配(如共享内存字节数、线程块布局)。经过严格的端到端数值正确性校验,最终筛选出 42,454 条 CUDA 优质轨迹和 59,998 条 Triton 优质轨迹,用于对基础模型执行监督微调(SFT),固化底层高性能算子编写骨干能力。
正确性感知组相对策略优化(CA-GRPO):针对强化学习阶段正确性与加速比之间的优化失衡问题,作者设计了组级别动态门控机制。给定输入模块 $m$,Coder 采样生成包含 $n$ 个候选硬件内核的展开组 $\mathcal{G} = \{k_1, \ldots, k_n\}$。系统首先统计该组样本的全局正确率:
$$ \operatorname{Acc}_{\mathcal{G}} = \frac{1}{|\mathcal{G}|} \sum_{i=1}^n \mathbf{1}[\operatorname{Corr}(m, k_i) = 1] $$并引入基于阈值 $\alpha$ 的二值门控函数 $\mathcal{T}(\mathcal{G})$:
$$\begin{aligned} \mathcal{T}(\mathcal{G}) = \begin{cases} 1, & \text{if } \operatorname{Acc}_{\mathcal{G}} \ge \alpha, \\ 0, & \text{otherwise}. \end{cases} \end{aligned}$$只有当整组样本的成功率不低于 $\alpha$ 时,门控才允许性能加速信号介入。设 $\mathcal{G}_c \subset \mathcal{G}$ 为组内所有功能正确的内核子集,对于每个正确内核 $k_i \in \mathcal{G}_c$,获取物理显卡上实测的相对于 PyTorch Eager 的加速比 $\mathbf{S}(k_i)$,并计算组内极值归一化加速得分 $S_i^{\mathcal{G}}$:
$$\begin{aligned} S_i^{\mathcal{G}} = \begin{cases} \frac{\mathbf{S}(k_i) - \min(\mathbf{S}(\mathcal{G}_c))}{\max(\mathbf{S}(\mathcal{G}_c)) - \min(\mathbf{S}(\mathcal{G}_c))}, & \max(\mathbf{S}(\mathcal{G}_c)) > \min(\mathbf{S}(\mathcal{G}_c)), \\ 0, & \text{otherwise}. \end{cases} \end{aligned}$$若组内正确内核加速比完全相同,归一化得分置为 0。最终赋予每个样本 $k_i \in \mathcal{G}$ 的复合奖励定义为:
$$\begin{aligned} \mathcal{R}_i^{\mathcal{G}} = \begin{cases} 1 + \mathcal{T}(\mathcal{G}) \cdot \beta \cdot S_i^{\mathcal{G}}, & \text{if } \operatorname{Corr}(m, k_i) = 1, \\ 0, & \text{otherwise}, \end{cases} \end{aligned}$$其中 $\beta$ 为性能信号缩放因子(实验设定为 $0.1$)。在策略优化阶段,未通过数值检验的代码奖励恒为 0;在组正确率较低的初期阶段,门控强行切断性能偏置,使得模型专注于攻克功能正确性;唯有算子正确性得到稳固保证时,组内加速比差异才被作为优势函数更新策略梯度,引导模型在正确性约束内部探索更深层的硬件级加速手法。
4. 实验环境与实施细节
4.1 硬件环境
- 计算与训练平台:所有训练与离线优化均在单节点配备 4 张 NVIDIA A100-SXM4-80GB GPU 的服务器上进行。
- 协同推断拓扑架构:
- Proposer 轮次更新:共占用 12 张 A100-80GB GPU,其中 4 张负责 Proposer 策略的反向梯度更新,4 张作为固定参数的 Coder 部署独立推断服务,剩余 4 张搭建为独立的 Keck 执行与数值验证集群。
- Coder 轮次更新:共占用 8 张 A100-80GB GPU,其中 4 张用于 Coder 自身的梯度更新,另外 4 张独立承载 Keck 编译与真实执行测速服务。
4.2 软件与算法配置
- 基座模型初始化:Proposer 和 Coder 统一采用开源预训练模型
Qwen2.5-Coder-7B。 - 数据流与超参数设定(详见下表):
| 超参数名称 | 提议器 (Proposer) | 编码器 (Coder) |
|---|---|---|
| 优化算法 | GRPO | GRPO |
| 初始化策略 | Qwen2.5-Coder-7B 原生权重 | KernelZero-SFT 蒸馏微调权重 |
| 强化学习训练实例总数 | 400 个 API 组合 Prompt | 1,224 个通过验证的合成模块 |
| 策略展开组大小(Rollout Group Size) | 4 | 5 |
| 每个模块嵌套 Coder 采样数 | 5 | - |
| 每轮交替优化步数(Steps per round) | 20 steps | 20 steps |
| 交替轮次总数(Training rounds) | 4 轮 | 4 轮 |
| 累计训练迭代步数 | 80 steps | 80 steps |
| 全局 / Rollout Batch Size | 8 | 64 |
| 优化器与学习率设置 | AdamW, $\text{LR}=1.0 \times 10^{-6}$ | AdamW, $\text{LR}=1.0 \times 10^{-6}$ |
| 权重衰减 / 梯度裁剪 | $1.0 \times 10^{-2}$ / 1.0 | $1.0 \times 10^{-2}$ / 1.0 |
| KL 惩罚类型与正则系数 $\gamma$ | low_var_kl, $1.0 \times 10^{-2}$ | low_var_kl, $1.0 \times 10^{-2}$ |
| 采样温度 / Top-P / Top-K | 0.6 / 0.95 / 20 | 0.6 / 0.95 / 20 |
| 最大输入 / 输出序列长度 | 8,192 / 6,144 tokens | 8,192 / 16,384 tokens |
| 张量并行度(Tensor Parallel Size) | 4 | 4 |
| 数值精度 | BF16 | BF16 |
| 单轮 20 步执行总耗时 | 约 3.1 小时 | 约 4.5 小时 |
| FO-GRPO 前沿目标值 / 非法惩罚 $\rho$ | 0.5 / -1.0 | - |
| CA-GRPO 正确性阈值 $\alpha$ / 性能权重 $\beta$ | - | 0.5 / 0.1 |
4.3 评测基准与度量衡
-
基准数据集:采用权威的 KernelBench 基准。涵盖两类评测集:
- Level 1:包含 100 个独立的单算子/基础组件级任务(如 GEMM、标准 1D/2D 卷积、Pooling、各种 Loss 函数等)。
- Level 2:包含 100 个融合模式任务(如
Matmul + Sigmoid、Conv + BatchNorm + Swish等跨算子复合计算)。
-
度量指标:
-
pass@k:借鉴标准代码评估 HumanEval 公式计算 $k$ 次采样中至少有 1 个算子功能正确的期望胜率:
$$\mathrm{pass}@k = 1 - \frac{\binom{n-c}{k}}{\binom{n}{k}}$$
-
$\mathrm{fast}_p@k$:在 pass@k 的基础上增加严苛的性能门槛,只有当生成的内核功能正确且实测执行延迟相比 PyTorch Eager 参考实现达到 $p\times$ 以上加速比时,才被认定为有效样本。
- Token 消耗:统计包含思考 CoT 与内核完整源码的平均输出 Token 长度,反映推断计算成本。
-
5. 实验结果与深度分析
5.1 主要评测对比
实验对基础通用大模型(Claude 系列、Gemini 系列、GPT-OSS、GLM 系列、Qwen 系列)和领域专用内核生成模型(cudaLLM-8B、Kevin-32B、Dr.Kernel-8B/14B)进行了全方位对比。
CUDA 算子生成性能剖析:
- 在 Level 1 上,KernelZero-7B 取得了 $75.8\%$ 的 pass@1,大幅领先同量级的 Qwen2.5-Coder-7B($6.5\%$)与 Qwen3-8B($27.4\%$),以 7B 体量逼近超大型专有模型 Claude-4.5-Sonnet($76.4\%$,差距仅 0.6%),并在多样本指标上全面反超(pass@5 达 $98.87\%$ vs $95.35\%$,pass@10 达成全集通过 $100\%$ vs $97\%$)。
- 在 Level 2 复合融合任务上,KernelZero-7B pass@1 达到 $69.6\%$,高居榜首,超越 Claude-4.5-Sonnet($66.5\%$)与拥有强化学习多轮优化的专用模型 Kevin-32B($19.0\%$)及 cudaLLM-8B($67.4\%$)。在 pass@10 上达到 $97\%$,显著超出 Claude-4.5-Sonnet 的 $93\%$。
Triton 算子生成性能剖析:
- 在 Level 1 上,KernelZero-7B 的 pass@1 达到 $77.2\%$,超越包含 1.6T 参数的 DeepSeek-V4-Pro($49.7\%$)、GLM-5.2($66.7\%$)以及三轮迭代推断的专用算子模型 Dr.Kernel-8B($34.4\%$)。
- 在 Level 2 上,KernelZero-7B pass@1 取得 $72.5\%$,比肩 Dr.Kernel-8B 在三次多轮循环后的 $73.0\%$,并超越 DeepSeek-V4-Pro 的 $63.4\%$。在多样本采样方面,pass@5 达到 $94.96\%$,pass@10 达到 $98\%$。
生成效率(Token 成本)优势:
KernelZero 保持极高的推断紧凑度。生成单个 CUDA 样本在 Level 1 和 Level 2 的 Token 消耗仅约为 2.8k 与 4.5k,远低于 Gemini-2.5-Flash(14.8k/21.2k)与 Kevin-32B(9.2k/13.7k);生成 Triton 算子仅消耗 3.8k/3.7k tokens,相比 GLM-5.2 的 19.6k/15.8k 节省了超 75% 的显存和计算资源。
内核加速比实测表现与后端鸿沟分析:
- Triton vs CUDA 架构级差异:如 Figure 5 所示,KernelZero 生成 Triton 算子在 Level 1 上展现出极高的加速通过率($\mathrm{fast}_1@1$ 为 $43.9\%$,$\mathrm{fast}_1@10$ 高达 $87\%$),Level 2 的 $\mathrm{fast}_1@1$ 达到 $65.4\%$,$\mathrm{fast}_1@10$ 达 $96\%$。而 CUDA 后端在 Level 2 上 $\mathrm{fast}_1@1$ 仅为 $2.4\%$,$\mathrm{fast}_1@10$ 为 $12\%$。
- 原因剖析:Triton 提供了天然以分块为中心的编程原语与由编译器深度管辖的硬件映射机制,能够自动优化寄存器溢出与片上调度;相比之下,本项工作对 CUDA 代码提出了极其严苛的工程限制——绝对禁止调用 cuBLAS 与 cuDNN 等底层闭源商业加速库。Level 2 任务包含大量卷积与高维 GEMM,其 PyTorch 参考实现直接分发到底层高度调优的手动汇编级供应商库。纯手写且未引入汇编级微优化的单内核 CUDA 代码,在面对高度工程化的 cuDNN/cuBLAS 基线时,极难在无外部库绑定的纯裸核下实现超高倍率加速。
5.2 消融实验与收敛动态
Proposer 持续进化的必要性:
为验证 Proposer 动态前沿采样的有效性,作者设定对照组:在 Step 40 处强行冻结 Proposer,使 Coder 仅在固化分布上继续迭代。
- 如 Figure 3 所示,持续更新 Proposer 时,生成的模块有效率持续上扬至接近 90%,Coder 准确率稳定围绕 50% 的前沿中枢波动,前沿奖励显著增加。
- 在 Coder 最终指标上,持续演化 Proposer 相比冻结版本,在 CUDA Level 1 和 Level 2 上分别获得 0.9% 和 5.2% 的 pass@1 提升;在 Triton 上分别取得 2.3% 和 4.2% 的 pass@1 增益。同时在 $\mathrm{fast}_1@1$ 上亦取得稳定增长。Level 2 上更大的提升幅度印证了当 Coder 能力不断外延时,一个协同演进的 Proposer 能够持续提供更具价值与结构丰富度的深层挑战。
CA-GRPO 正确性阈值 $\alpha$ 的调优消融:
作者在 200 个 KernelBench 任务全集上针对正确性门控阈值 $\alpha \in \{0, 0.5, 0.9, 1.1\}$ 进行了详细消融对比(其中 $\alpha=0$ 意味着无门控直接注入加速信号,$\alpha=1.1$ 意味着完全关闭加速比奖励,仅训练正确性)。
- 在 Step 80 评估时,纯正确性导向的 $\alpha=1.1$ 取得了最高的 pass@1($74.85\%$),但其 $\mathrm{fast}_{1.5}@1$ 仅有 $9.75\%$;过于激进的 $\alpha=0.9$ 使得 $\mathrm{fast}_{1.5}@1$ 升至 $11.00\%$,但 pass@1 下滑至 $73.15\%$。
- 如 Figure 4 所示,在具有代表性的 6 个算子详细测试中,设定 $\alpha=0.5$ 在正确性与极端性能之间取得了最佳平衡,在 Tensor-MM 任务上将 SFT 的 $0.12\times$ 加速比激增至 $4.42\times$,并在 GroupNorm($1.25\times$)、Conv-BN-Scale($1.16\times$)与 GEMM-BN-Swish($1.09\times$)上均拔得头筹。
6. 基础设施细节:Keck 评测与验证服务端设计
针对原始 KernelBench 评测框架中由于并发运行导致的 GPU 显存资源争抢、延迟抖动,以及算子内部 __init__ 中随机初始化顺序不一致导致的数值误差问题,作者基于 cudaLLM 开发了高度可控的算子判题服务系统 Keck。
Keck 具备六大核心技术特性:
- 后端感知的解耦编译与执行:针对 CUDA 扩展,利用 CPU Worker 资源池执行多进程并发
load_inlineC++ 编译,编译成功后立即分发至空闲的 GPU 执行器;Triton 算子无单独 C++ 编译开销,直接在目标 GPU 上 JIT 编译并利用预热运行完成缓存。 - 单样本单卡物理隔离(Single-sample GPU execution):严密约束每个物理 GPU Worker 在同一时刻仅独占运行一个算子候选进程,并在超时或执行崩溃后由主控安全销毁子进程,杜绝多任务并发争抢片上缓存与高带宽内存带来的性能抖动。
- 模块参数严格等价迁移(load_state_dict):在通过统一随机种子初始化参考模型与生成模型后,显式将参考模型的
state_dict复制并强制载入生成模型实例,消除算子构造内部调用伪随机数生成器顺序差异对数值等价性判定的干扰。 - 统一双容差数值校验:向参考与生成模型分发完全独立的输入张量拷贝(防止原位 In-place 修改污染),设定绝对容差(Absolute Tolerance)与相对容差(Relative Tolerance)均为 $10^{-2}$。
- 严苛的纯硬件算子实现约束:静态与动态监测生成代码,CUDA 必须原生手写全局内核(
__global__),严禁委托给torch.nn.functional;Triton 必须包含@triton.jit,严禁通过 Python 逻辑伪造算子实现。 - 正确性门控的精确时钟测速:唯有通过数值校验的算子方可激活性能剖析。采用物理显卡上的 CUDA Events 进行精准计时,统一经历 2 次预热(Warmup)与 3 次正式评测计算均值,计算相对 PyTorch Eager 的净加速比。
7. 附录精要与案例研究
7.1 GRPO 理论基础回顾
组相对策略优化(Group Relative Policy Optimization, GRPO)省去了传统 PPO 算法中庞大的价值网络(Critic Model),以降低在线推理开销。针对包含奖励集合 $\{\mathcal{R}_i^{\mathcal{G}}\}_{i \in \mathcal{G}}$ 的输出组 $\mathcal{G}$,其组归一化优势函数计算为:
$$ \mathcal{A}_i^{\mathcal{G}} = \frac{\mathcal{R}_i^{\mathcal{G}} - \operatorname{mean}(\mathcal{R}_j^{\mathcal{G}})}{\operatorname{std}(\mathcal{R}_j^{\mathcal{G}})} $$优化目标由重要性采样比率 $r_{i,t}(\theta) = \frac{\pi_\theta(o_{i,t} \mid x, o_{i,<t})}{\pi_{\theta_{\text{old}}}(o_{i,t} \mid x, o_{i,<t})}$ 与带权裁剪项构成,并辅以参考模型 $\pi_{\text{ref}}$ 的无偏低方差 KL 散度约束项:</p>
$$ \mathcal{T}(\theta) = \mathbb{E} \left[ \frac{1}{|\mathcal{G}|} \sum_{i \in \mathcal{G}} \frac{1}{|o_i|} \sum_{t=1}^{|o_i|} \min\left(r_{i,t}(\theta)\mathcal{A}_i^\theta, \operatorname{clip}(r_{i,t}(\theta), 1-\epsilon, 1+\epsilon)\mathcal{A}_i^\theta\right) - \gamma \mathrm{KL}(\pi_\theta \parallel \pi_{\mathrm{ref}}) \right] $$7.2 Proposer 输出的定性演化分析
附录展示了在给定相同 API 组合输入下,Proposer 从基座(Base)到 Step 60 Checkpoint 的进化过程。利用 torch.fx 提取真实计算图,模型逐渐从初期的参数与维度不匹配、机械单链串联,自主演化出多分支合并、特征分层与显式目标节点:
案例 1:批矩阵乘法融合(Dropout + torch.bmm)
- 初始状态(Base):机械地将两个外部输入张量分别经过 Dropout,然后调用 BMM,拓扑平庸。
- 演化到 Step 60:模型学会从单一特征张量 $x$ 出发,经过 Linear 层与 Dropout 提取表征后,通过维度重塑(
unsqueeze(1)与unsqueeze(2))构建出双分支,再利用 BMM 计算内积,最终完成降维输出,生成了复杂的结构依赖。
案例 2:多分支时间卷积(Conv1d + torch.cat)
- 初始状态(Base):仅执行单个 Conv1d,随后简单地通过
torch.cat([x, x], dim=1)复制自身拼接。 - 演化到 Step 60:自主构建出具有多尺度或独立参数的三条独立卷积分支(
conv1,conv2,conv3),前向运算并行提取特征后,在通道维度进行真实的三路融合拼接,极大地提升了图拓扑丰富度。
案例 3:正负互补激活分支(ReLU + torch.cat)
- 初始状态(Base):通过多次重复调用同一个变量并进行拼接来凑数。
- 演化到 Step 60:自动设计出双路径特征提取网络——一条分支执行
relu1(x),另一条分支执行求负后激活relu2(-x),随后执行跨路径特征聚合,形成了类似 CReLU(Concatenated ReLU)的自适应神经网络结构。
案例 4:具显式目标的嵌入分类网络(Embedding + CrossEntropyLoss + torch.cat)
- 初始状态(Base):直接对双输入查表并拼接,直接输入损失函数。
- 演化到 Step 60:网络分化为主干特征分支与监督分支。输入分别提供类别标记与索引,特征支路历经查表、自拼接增维、双层 MLP 映射(
fc1 -> relu -> fc2),最终汇聚至CrossEntropyLoss节点,构建了规范完整的监督学习流。
案例 5:并行 3D 池化融合(Dropout + MaxPool3d + torch.cat)
- 演化至 Step 60:模型演化出高维空间处理结构,在 3D 卷积与随机失活之后,分化出两个独立的 3D 最大池化核(
pool1与pool2),最终通过torch.cat实行通道聚合。
案例 6:多输出头共享主干(torch.sqrt + torch.mean)
- 演化至 Step 60:自发生成包含 3D 卷积、Softmax 激活及两级级联 3D 最大池化的共享主干,并在尾端分化为全局均值归约头与逐元素平方根头,同时显式返回两个并行的张量输出。
7.3 生成算子底层实现解析(Case Study on Generated Kernels)
1. 对角矩阵乘法(Diagonal Matrix Multiplication, 1-12)
- 测定加速比:$13.59\times$(相对 PyTorch Eager)。
- 底层优化手法:PyTorch 原生实现计算
torch.diag(A) @ B会产生巨大的稀疏临时矩阵,导致显存带宽浪费。KernelZero 生成的代码将其数学等价为 $C[i,j] = A[i] \cdot B[i,j]$。在 CUDA 内核中利用 $32 \times 32$ 的二维线程块(dim3 blockDim(32, 32))划分输出区域,每个 Block 分配共享内存数组__shared__ float As[TILE_N],仅由第一列线程(threadIdx.y == 0)协作将向量 $A$ 对应行的一次性加载到共享内存,后续所有列计算完全复用片上缓存,规避了全局内存的反复读取。
__global__ void diag_mul_kernel(const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C, int N, int M) {
int row = blockIdx.x * TILE_N + threadIdx.x;
int col = blockIdx.y * TILE_M + threadIdx.y;
__shared__ float As[TILE_N];
if (threadIdx.y == 0 && row < N) {
As[threadIdx.x] = A[row];
}
__syncthreads();
if (row < N && col < M) {
float a = As[threadIdx.x];
float b = B[row * M + col];
C[row * M + col] = a * b;
}
}
2. 多分类交叉熵损失函数(Cross-Entropy Loss, 1-95)
- 测定加速比:$4.79\times$。
- 底层优化手法:设计了单线程块处理单样本的高效归约结构。内核通过四阶段流水完成:利用 Warp 原语
__shfl_xor_sync执行块内两级最大值规约获取稳定性偏置max_val;跨步长循环累加指数和并通过共享内存二次归约得到sum_exp;计算样本级对数似然损失;最终通过原子加指令(atomicAdd)聚合批次总损失,避免了 PyTorch 产生多次中间张量所带来的显存往返开销。
3. 三元组边际损失函数(Triplet-Margin Loss, 1-99)
- 测定加速比:$3.88\times$。
- 底层优化手法:将高维特征($D=8192$)划分给线程块内的 256 个线程通过跨步(Grid-stride loop)并行累加差值平方和。利用 Warp 级洗牌规约指令
__shfl_down_sync与共享内存分层合并,单 Block 内部完成锚点-正样本与锚点-负样本的二范数距离开方运算,最后由辅内核divide_by_batch_kernel执行批归约除法,访存连续性与寄存器复用极高。
4. KL 散度损失函数(KL-Divergence Loss, 1-98)
- 测定加速比:$2.99\times$。
- 底层优化手法:将
torch.log预处理与逐元素贡献求和深度融合进单个内核。每个线程负责跨步处理大张量切片,在寄存器局部累积thread_sum += t * (logt - logp),之后在共享内存中以二叉树折叠方式归约,利用atomicAdd将各 Block 的局部和累加至全局标量,最后通过独立微内核执行batchmean归一化除法。
5. GEMM、除法、求和与缩放融合算子(2-14)
- 测定加速比:$20.11\times$。
- 底层优化手法:针对连续串行的矩阵乘法、常数除法、按行求和与标量缩放计算,KernelZero 识别并执行了代数重排优化。将原本昂贵的完整矩阵乘法和随后的逐元素求和,数学等价重构为两个级联小内核:首先由
col_sum_kernel计算权重矩阵在列维度上的投影累加和向量,随后fused_kernel直接将输入向量与该列和向量计算点积,并在片上直接执行除以 2 与常数乘法缩放,彻底消除了计算大尺寸中间激活矩阵的巨大开销。
6. 转置卷积与全局平均池化融合(2-44)
- 测定加速比:$2.50\times$。
- 底层优化手法:原本 PyTorch 需先通过
ConvTranspose2d展开完整的空间张量($128 \times 128 \times 254 \times 254$),再调用全局平均池化。KernelZero 生成的transposed_conv_avg_kernel利用代数等价性,直接将输入特征图的空间总和与卷积核切片的空间总和在 Block 内部利用 Warp 原语并行累加求积,完全跳过了生成庞大转置卷积空间中间特征图的过程,将时间与空间复杂度降至最低。
8. 方法细节引用汇总
| 引用标号 | 对应完整参考文献条目 | 论文中出现的段落与上下文描述 |
|---|---|---|
| [1] | Mohammad Ashraful Anam, Paul N Whatmough, and Yiannis Andreopoulos. Precision-energy-throughput scaling of generic matrix multiplication and discrete convolution kernels via linear projections. IEEE SETS 2013. | 2.3 节 Cold Start Distillation - Reordering:"Reordering rearranges loop/order-of-operations and thread/data mapping to improve coalesced access and temporal reuse, which strengthens cache/shared-memory locality and alleviates bank conflicts and divergence [1]." |
| [2] | Carlo Baronio, Pietro Marsella, Ben Pan, Simon Guo, and Silas Alberti. Kevin: Multi-turn rl for generating cuda kernels. arXiv:2507.11948, 2025. | 1 节 Introduction:"recent studies including Kevin [2] and Dr.Kernel [17] have demonstrated the capability of LLMs..." 及 3.1.3 节 Specialized Models 作为多轮 RL 基线对比。 |
| [3] | ByteDance-Seed. cudaLLM: Training language models to generate high-performance cuda kernels. GitHub 2025. | 1 节、3.1.1 节 及 A.1 节:作为 SFT 数据集来源及 Keck 基础架构支撑文献。 |
| [4] | Cambricon. Cambricon BANG C Developer Guide. Cambricon 2023. | 1 节 Introduction:列举领域特定编程语言时引用。 |
| [7] | William J Dally, Stephen W Keckler, and David B Kirk. Evolution of the graphics processing unit (gpu). IEEE Micro 2021. | 2.3 节 Cold Start Distillation - Tiling:"Tiling partitions input tensors into block-level and warp-level tiles... enables efficient block/warp reductions [7]." |
| [8] | DeepSeek-AI et al. Deepseek-v4: Towards highly efficient million-token context intelligence. arXiv:2606.19348, 2026. | 3.1.3 节 Foundation Models:作为 Triton 评测的高性能超大规模基线模型。 |
| [11] | GLM-5-Team et al. GLM-5: from vibe coding to agentic engineering. arXiv:2602.15763, 2026. | 3.1.3 节 Foundation Models:作为对比的基础模型基线。 |
| [12] | Chengsong Huang et al. R-zero: Self-evolving reasoning llm from zero data. arXiv:2508.05004, 2026. | 2.2 节 Frontier Reward:"This formulation encourages the Proposer to generate modules that are neither trivial nor impossible [12]." |
| [13] | Binyuan Hui et al. Qwen2.5-coder technical report. arXiv:2409.12186, 2024. | 3.1.1 节 Training Details:作为 Proposer 与 Coder 权重的初始基座模型。 |
| [17] | Wei Liu, Jiawei Xu et al. Dr. kernel: Reinforcement learning done right for triton kernel generations. arXiv:2602.05885, 2026. | 1 节 与 3.1.3 节:作为 Triton 专业算子生成模型的重点对比基线。 |
| [18] | John Nickolls et al. Scalable parallel programming with cuda. ACM SIGGRAPH 2008. | 1 节 Introduction:作为 GPU 并行编程领域的奠基性框架引用。 |
| [19] | OpenAI. gpt-oss-120b & gpt-oss-20b model card. arXiv:2508.10925, 2025. | 3.1.1 节 Coder Cold Start:作为冷启动思维链蒸馏过程中的教师模型。 |
| [20] | Anne Ouyang et al. Kernelbench: Can llms write efficient gpu kernels? arXiv:2502.10517, 2025. | 1 节、2.2 节、3.1.2 节:作为全篇的核心评测基准套件及接口规范引用。 |
| [21] | Daniel Snider and Ruofan Liang. Operator fusion in xla: Analysis and evaluation. arXiv:2301.13062, 2023. | 2.3 节 Cold Start Distillation - Fusion:"Fusion merges producer–consumer operator chains into a single kernel... reducing synchronization cost [21]." |
| [22] | Guangming Tan et al. Fast implementation of dgemm on fermi gpu. SC 2011. | 2.3 节 Cold Start Distillation - Pipeline:"Pipeline organizes computation into staged execution that overlaps memory movement... thereby hiding memory latency [22]." |
| [23] | Philippe Tillet, H. T. Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. MAPL 2019. | 1 节 Introduction:作为分块神经网络高级编程语言引用。 |
| [25] | Shanli Xing et al. Flashinfer-bench: Building the virtuous cycle for ai-driven llm systems. arXiv:2601.00227, 2026. | 1 节 Introduction:作为系统性算子评测套件引用。 |
| [26] | An Yang et al. Qwen3 technical report. arXiv:2505.09388, 2025. | 3.1.3 节 Foundation Models:作为基线大模型引用。 |
| [27] | Qirui Zhou, Yuanbo Wen et al. Qimeng-gemm: Automatically generating high-performance matrix multiplication code by exploiting large language models. 2025. | 2.3 节 Cold Start Distillation:作为 GPU 算子优化分类学(Taxonomy)与设计启发式规则的核心依据。 |
9. 结论
KernelZero 构建了一套无需人工构建复杂训练用例、可自主演进的代码生成框架。通过 Proposer 依托 API 共现与胜率前沿自适应出题,以及 Coder 在冷启动编译原则指引与 CA-GRPO 门控奖励下的双重规范,模型实现了在算子功能正确性与极端硬件执行效率之间的协同跃升。在 KernelBench 基准上的系统评测证明,仅 7B 参数规模的专用模型完全能够在低级底层代码生成任务上超越参数量达数千亿的前沿通用基座模型,证明了能力对齐合成数据流与正确性优先强化学习范式在底层编译系统领域的巨大潜力。
💬 评论讨论
欢迎在这里分享您的想法和见解!