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}$ 构造边界奖励:

$$\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}$$


当 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}}$:

$$\begin{aligned} \mathcal{R}_i^{\mathcal{G}} = \begin{cases} 1 + \mathcal{T}(\mathcal{G}) \cdot \beta \cdot S_i^{\mathcal{G}}, & \text{if } \text{Corr}(m, k_i) = 1 \\ 0, & \text{otherwise} \end{cases} \end{aligned}$$
这确保了模型先稳定学会正确实现,再激发展开性能优化。两者固定一方、交替更新各 20 步,形成能力递进闭环。

效果如何 实验基于 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│
                     └───────────────────────┘

本文的核心创新点与技术贡献包括:

  1. 双模型协同进化框架(Co-Evolution Framework):构建由提议器(Proposer)与编码器(Coder)构成的博弈协同系统。Proposer 负责从 PyTorch API 集合合成计算图模块,Coder 负责将其翻译为底层的 CUDA 或 Triton 硬件算子,二者通过明确的接口交替迭代。
  2. 前沿驱动的模块生成机制(Frontier-Driven Module Generation):Proposer 依据 Coder 在前一轮中的弱点(失败用例涉及的 API 分布)自适应采样算子集合,并由基于前沿胜率的奖励函数引导,精准生成处于 Coder 当前能力前沿(Pass 率约 50%)的 Torch 模块,既避免了无意义的简单样本,又防止了过难样本导致的训练发散,实现无人工标注的自适应动态数据流。
  3. 正确性感知组相对策略优化(Correctness-Aware GRPO, CA-GRPO):提出组级别门控奖励机制。只有当候选算子组的整体正确率达到预设阈值 $\alpha$ 时,才激活基于相对加速比的性能优化信号,根本上消除了早期训练因追求极致加速而破坏算子功能等价性的退化现象。
  4. 编译指导的冷启动蒸馏(Compiler-Inspired Cold-Start Distillation):系统化归纳分块(Tiling)、融合(Fusion)、流水线(Pipeline)与重排(Reordering)四项 GPU 算子设计准则,用于指导教师模型生成具显式推理链(CoT)的高质量冷启动数据。
  5. 构建高吞吐评测基座 Keck:设计了隔离运行、多进程看门狗监控、权重状态严格对齐(load_state_dict)且支持 CPU 异步编译的评测服务基础设施 Keck,保障了奖励反馈的高保真度。
  6. 超越超大规模前沿模型的评测表现:参数量仅 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。
Figure 1: KernelZero 协同进化框架概览。系统由两个具备明确接口的模型组成:Proposer 接收一组 Torch API 作为输入并生成可运行的 Torch 模块,Coder 将该模块翻译为经过优化的 CUDA 或 Triton 算子。
Figure 1: KernelZero 协同进化框架概览。系统由两个具备明确接口的模型组成:Proposer 接收一组 Torch API 作为输入并生成可运行的 Torch 模块,Coder 将该模块翻译为经过优化的 CUDA 或 Triton 算子。

2. 背景知识、关键观察与设计原则

2.1 现有算子生成范式的局限与观察

2.2 四大 GPU 算子设计核心原则

为了在冷启动阶段消除高层抽象模型与底层硬件微架构之间的认知鸿沟,作者将现代 GPU 编译器与资深系统程序员的优化经验归纳为互补的四大体系结构优化原则:


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$ 则学习如何将其高效地在物理硬件底层落地。

Figure 2: KernelZero 协同进化框架整体设计。包含 API List Sampling、Module Proposal、Validation Check、双模型交替强化学习循环(RL Training Loop)及基于奖励反馈演进的闭环架构。
Figure 2: KernelZero 协同进化框架整体设计。包含 API List Sampling、Module Proposal、Validation Check、双模型交替强化学习循环(RL Training Loop)及基于奖励反馈演进的闭环架构。

双模型交替优化机制:整个训练过程采用交替优化策略。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 采样的候选模块必须串行通过三道校验关卡:

  1. 抽象语法树语义检查(AST Check):利用静态语法解析器遍历模块 AST,验证代码语义的合规性,严格检查生成的计算图中是否无遗漏地包含了所分配的采样 API 集合 $a$。
  2. 硬件可执行性检查(Runnable Check):依据标准评测框架规范,生成的模块必须实现用于张量初始化的 get_init_inputs 方法和用于生成随机测试输入的 get_inputs 方法。系统在物理 CUDA 显卡上实例化该模块并注入随机张量运行前向传播;若发生编译错误、形状不兼容或运行时异常抛出,该模块被立即丢弃。
  3. 数值稳定性检查(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.2 软件与算法配置

超参数名称 提议器 (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 评测基准与度量衡


5. 实验结果与深度分析

5.1 主要评测对比

实验对基础通用大模型(Claude 系列、Gemini 系列、GPT-OSS、GLM 系列、Qwen 系列)和领域专用内核生成模型(cudaLLM-8B、Kevin-32B、Dr.Kernel-8B/14B)进行了全方位对比。

Table 1: KernelBench 基准上 CUDA 和 Triton 算子生成的执行正确率与 Token 成本对比。加粗表示最优,下划线表示次优。
Table 1: KernelBench 基准上 CUDA 和 Triton 算子生成的执行正确率与 Token 成本对比。加粗表示最优,下划线表示次优。

CUDA 算子生成性能剖析:

Triton 算子生成性能剖析:

生成效率(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% 的显存和计算资源。

Figure 5: KernelBench 基准上的 Fast@k 性能分布。左列为 Triton 生成,右列为 CUDA 生成。Fast_p@k 表示 k 次采样中至少有 1 个算子兼具正确性与 p 倍以上加速比的百分比。
Figure 5: KernelBench 基准上的 Fast@k 性能分布。左列为 Triton 生成,右列为 CUDA 生成。Fast_p@k 表示 k 次采样中至少有 1 个算子兼具正确性与 p 倍以上加速比的百分比。

内核加速比实测表现与后端鸿沟分析:

5.2 消融实验与收敛动态

Proposer 持续进化的必要性:
为验证 Proposer 动态前沿采样的有效性,作者设定对照组:在 Step 40 处强行冻结 Proposer,使 Coder 仅在固化分布上继续迭代。

Figure 3: KernelZero 的训练演化动态。左图展示 Proposer 训练过程中的有效率、Coder 准确率及前沿奖励变化趋势。右图展示 CUDA 和 Triton 算子训练中 pass@1 和 fast_1@1 的收敛曲线,虚线代表 Step 40 处冻结 Proposer 的消融对照组。
Figure 3: KernelZero 的训练演化动态。左图展示 Proposer 训练过程中的有效率、Coder 准确率及前沿奖励变化趋势。右图展示 CUDA 和 Triton 算子训练中 pass@1 和 fast_1@1 的收敛曲线,虚线代表 Step 40 处冻结 Proposer 的消融对照组。

CA-GRPO 正确性阈值 $\alpha$ 的调优消融:
作者在 200 个 KernelBench 任务全集上针对正确性门控阈值 $\alpha \in \{0, 0.5, 0.9, 1.1\}$ 进行了详细消融对比(其中 $\alpha=0$ 意味着无门控直接注入加速信号,$\alpha=1.1$ 意味着完全关闭加速比奖励,仅训练正确性)。

Figure 4: CA-GRPO 正确性阈值 alpha 的消融实验。各柱状图报告 10 次生成中最快正确内核的加速比。对比了冷启动基线 KernelZero-SFT、KernelZero-RL (alpha=0.5) 以及 alpha=0, 0.9, 1.1 变体,并引入 DeepSeek-V4-Pro 与 GLM-5.2 作为外部参考。
Figure 4: CA-GRPO 正确性阈值 alpha 的消融实验。各柱状图报告 10 次生成中最快正确内核的加速比。对比了冷启动基线 KernelZero-SFT、KernelZero-RL (alpha=0.5) 以及 alpha=0, 0.9, 1.1 变体,并引入 DeepSeek-V4-Pro 与 GLM-5.2 作为外部参考。

6. 基础设施细节:Keck 评测与验证服务端设计

针对原始 KernelBench 评测框架中由于并发运行导致的 GPU 显存资源争抢、延迟抖动,以及算子内部 __init__ 中随机初始化顺序不一致导致的数值误差问题,作者基于 cudaLLM 开发了高度可控的算子判题服务系统 Keck。

Figure 6: Keck 算子执行与验证服务器架构图。CUDA 扩展在 CPU 工作节点并发编译,随后分发至隔离的 GPU 工作节点;Triton 算子在目标 GPU 上 JIT 即时编译。两条链路均将生成算子与 PyTorch 参考基准严格对齐,并在通过正确性验证后测量延迟与加速比。
Figure 6: Keck 算子执行与验证服务器架构图。CUDA 扩展在 CPU 工作节点并发编译,随后分发至隔离的 GPU 工作节点;Triton 算子在目标 GPU 上 JIT 即时编译。两条链路均将生成算子与 PyTorch 参考基准严格对齐,并在通过正确性验证后测量延迟与加速比。

Keck 具备六大核心技术特性:

  1. 后端感知的解耦编译与执行:针对 CUDA 扩展,利用 CPU Worker 资源池执行多进程并发 load_inline C++ 编译,编译成功后立即分发至空闲的 GPU 执行器;Triton 算子无单独 C++ 编译开销,直接在目标 GPU 上 JIT 编译并利用预热运行完成缓存。
  2. 单样本单卡物理隔离(Single-sample GPU execution):严密约束每个物理 GPU Worker 在同一时刻仅独占运行一个算子候选进程,并在超时或执行崩溃后由主控安全销毁子进程,杜绝多任务并发争抢片上缓存与高带宽内存带来的性能抖动。
  3. 模块参数严格等价迁移(load_state_dict):在通过统一随机种子初始化参考模型与生成模型后,显式将参考模型的 state_dict 复制并强制载入生成模型实例,消除算子构造内部调用伪随机数生成器顺序差异对数值等价性判定的干扰。
  4. 统一双容差数值校验:向参考与生成模型分发完全独立的输入张量拷贝(防止原位 In-place 修改污染),设定绝对容差(Absolute Tolerance)与相对容差(Relative Tolerance)均为 $10^{-2}$。
  5. 严苛的纯硬件算子实现约束:静态与动态监测生成代码,CUDA 必须原生手写全局内核(__global__),严禁委托给 torch.nn.functional;Triton 必须包含 @triton.jit,严禁通过 Python 逻辑伪造算子实现。
  6. 正确性门控的精确时钟测速:唯有通过数值校验的算子方可激活性能剖析。采用物理显卡上的 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 提取真实计算图,模型逐渐从初期的参数与维度不匹配、机械单链串联,自主演化出多分支合并、特征分层与显式目标节点:

Figure 7: PyTorch API 共现统计矩阵与包含全部 442 个候选 API 的采样环形图。
Figure 7: PyTorch API 共现统计矩阵与包含全部 442 个候选 API 的采样环形图。

案例 1:批矩阵乘法融合(Dropout + torch.bmm)

Base 模块计算图:直接双分支输入经过 Dropout 后进入 BMM。
Step 20 模块计算图:引入 Linear 权重并通过 unsqueeze 进行 BMM。
Step 60 模块计算图:单输入通过 Linear、Dropout 后通过双向 unsqueeze 自身实现自注意力形态的 BMM 交互。

案例 2:多分支时间卷积(Conv1d + torch.cat)

Step 20 计算图:单路卷积后接自身拼接与线性变换。
Step 60 计算图:单输入分支分发至三个独立 Conv1d 并行节点,随后在通道维进行三路 cat 汇聚。
Step 60 对应生成的标准 PyTorch 代码实现。

案例 3:正负互补激活分支(ReLU + torch.cat)

Step 20 计算图:ReLU 激活与负值 ReLU 串联后拼接。
Step 40 计算图:持续微调负值响应特征。
Step 60 计算图:分化出明确的正负激活双支路并在最终融合。
Step 60 正负激活分支完整 Python 源码。

案例 4:具显式目标的嵌入分类网络(Embedding + CrossEntropyLoss + torch.cat)

Base 模块计算图:简单的 Embedding 输出后拼接接 Loss。
Step 60 计算图:完整的特征处理流,结合双层全连接后与目标汇入 CrossEntropyLoss。
Step 60 嵌入分类网络完整 Python 源码。

案例 5:并行 3D 池化融合(Dropout + MaxPool3d + torch.cat)

Step 40 计算图:3D 卷积及 Dropout 后接双路最大池化汇集。
Step 60 计算图:结构清晰的双路 3D 空间池化并行网络。
Step 60 并行 3D 池化融合完整 Python 源码。

案例 6:多输出头共享主干(torch.sqrt + torch.mean)

Step 40 计算图:单路串行 sqrt -> mean -> linear。
Step 60 计算图:复杂主干并在末端解耦输出均值与开方两个头。


7.3 生成算子底层实现解析(Case Study on Generated Kernels)

1. 对角矩阵乘法(Diagonal Matrix Multiplication, 1-12)

__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)

3. 三元组边际损失函数(Triplet-Margin Loss, 1-99)

4. KL 散度损失函数(KL-Divergence Loss, 1-98)

5. GEMM、除法、求和与缩放融合算子(2-14)

6. 转置卷积与全局平均池化融合(2-44)


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 参数规模的专用模型完全能够在低级底层代码生成任务上超越参数量达数千亿的前沿通用基座模型,证明了能力对齐合成数据流与正确性优先强化学习范式在底层编译系统领域的巨大潜力。