基准评测基准

KernelBench: Can LLMs Write Efficient GPU Kernels?

KernelBench:LLM 能写出高效 GPU kernel 吗

Anne Ouyang, Simon Guo, Simran Arora, et al. · Stanford (Scaling Intelligence Lab) · ICML 2025 · 2025-02 · 被引 215

一句话Stanford 提出的 250 题 GPU kernel 生成基准:给 PyTorch 参考实现,让 LLM 写出更快的 CUDA/DSL 实现,用 fast_p(正确且加速比 > p 的任务比例)打分。一次生成下最强模型 fast_1 不到 20%,该基准已成为 kernel 自动生成方向的标准尺子(S2 引用 215)。

这是什么

写高性能 GPU kernel 是 ML 系统工程里最费专家时间的环节之一——FlashAttention 从 Transformer 发布到出现要了 5 年。KernelBench 把这件事变成一个可自动验证的 LLM 任务:输入是一个 PyTorch 的 Model 类(含 get_inputs() 指定的具体张量形状/精度),模型输出 ModelNew 类,内嵌自定义 CUDA(或 Triton、CUTLASS 等)实现,要求输出与参考一致且更快。

250 个任务分三级:Level 1 是 100 个单算子(matmul、conv、norm 等,对手是 cuBLAS 这类闭源库,很难打);Level 2 是 100 个可融合的算子序列(Conv+Bias+ReLU 之类,对手是 torch.compile 的融合规则);Level 3 是 50 个完整架构(AlexNet、MiniGPT、Mamba 等,编译器管不到、需要算法级改写)。repo 里另有实验性的 Level 4(HuggingFace 整模型),论文未评测。

评测指标 fast_p = 正确(5 组随机输入对拍)且加速比 > p 的任务比例;fast_0 即正确率,fast_1 即'正确且比 PyTorch Eager 快'。p 可调,基准难度随 p 拉高,这个设计让它不容易饱和。后来的 AI CUDA Engineer(Sakana)、Kevin、多家实验室的 kernel RL 工作都在它上面评测。

任务全流程:左边是 PyTorch 参考 Model 类 + 任务指令,LLM 生成内嵌自定义 CUDA 的 ModelNew;右边评测分两步——随机输入下与参考输出对拍判正确,再与 Eager / torch.compile 对比 wall-clock 计时判性能。

机制与做法

任务格式与验证

每题给一个 torch.nn.Module 的参考实现和确定的输入规格,LLM 自由决定优化粒度——可以只替换最热的算子,也可以整个 forward 重写,可用任何技术(fusion、tiling、tensor core、PTX)。正确性靠随机输入下与参考输出对拍(n_correctness 次),性能靠 wall-clock 重复计时(n_trial 次),全流程无需人工标注,天然适合做 RL 的 verifiable reward。

评测是 evaluation-only:官方不提供 ground-truth kernel,因为最优解依赖具体硬件和输入形状。论文主实验在 NVIDIA L40S 上跑,repo 提供多代 NVIDIA GPU 的 baseline 时间(results/timing),并支持 Modal 云端评测。

一次生成下 fast_p 随加速阈值 p 的衰减曲线(三个 Level)。p=0 处即正确率(最好约 60-70%),到 p=1(要求比 PyTorch 快)全部跌破 15%,p=2 后接近归零;DeepSeek R1(黄)和 o1(橙)明显领先,Level 2 是推理模型相对最能打的一级。

一次生成的战绩:普遍很差

greedy 一次生成下,对 PyTorch Eager 的 fast_1:最强的 DeepSeek R1 在 Level 1/2/3 分别是 12%/36%/2%,o1 是 10%/24%/12%,GPT-4o 只有 4%/5%/0%,Llama 3.1-70B 接近全灭。失败分解显示:非推理模型 >70% 的生成连跑都跑不起来或输出错误,其中执行失败(编译错、CUDA 内存违例)是大头;推理模型执行失败少,但功能正确性错误和其他模型差不多。

把 speedup 阈值 p 拉到 2x,几乎所有模型都跌到接近 0。跨硬件泛化也差:R1 的 Level 2 kernel 在 L40S 上 fast_1=36%,在 A10G 上 47%,说明一次生成的 kernel 是'碰巧快',不是针对硬件优化的。给 prompt 里塞硬件规格(H100 参数、warp/SM 概念)对多数模型基本没用;R1 会尝试 wmma 指令(约 50% 的 Level 1 matmul 题),但大多编译失败。

test-time 方法能显著提分

重复采样:DeepSeek-V3 在 Level 2 上 k=100 时 fast_1 从 4% 涨到 37%,但对模型本来就完全不会的题(如 Level 1 的 34 个卷积变体)采样再多也没用。迭代精炼:把编译错误、正确性结果、PyTorch profiler 输出喂回去多轮改,R1 在 Level 2 上 10 轮后 fast_1 从 36% 到 72%,且 90% 以上的题能在 10 轮内改出功能正确的 kernel。同等 10 次调用预算下,迭代精炼在 6 组设置中 5 组优于重复采样——但两者都强依赖基座模型质量。

有意思的反直觉结果:给 few-shot 塞'好 kernel 范例'(fusion/tiling/mini-FlashAttention)反而降低整体 fast_1——模型会尝试更激进的优化,导致更多执行失败;但成功的那部分质量更高(o1 在 77% 的 GEMM 变体上用了 tiling)。

关键结果

一次生成下所有前沿模型 fast_1 < 20%(平均):对 PyTorch Eager,R1 最好(L1 12% / L2 36% / L3 2%),GPT-4o、Llama 3.1 系列几乎为零。
推理模型(o1、R1)错误率 <55%,非推理模型 >70%,差距主要来自执行失败更少;功能正确性错误各家都差不多。
重复采样 k=100:DeepSeek-V3 Level 2 fast_1 从 4% → 37%;迭代精炼 10 轮(带执行+profiler 反馈):R1 Level 2 从 36% → 72%。
p 阈值一拉高就崩:fast_2(2 倍加速)下所有模型接近 0,说明离'真正超越手工优化库'很远。
torch.compile 作为 baseline 时数字普遍更好看(R1 L1 达 38%),但论文注明这是 torch.compile 自身对小 kernel 有运行时开销,主分析用 PyTorch Eager。
跨硬件不稳定:同一批 R1 kernel 在 L40S 与 A10G 上 Level 2 fast_1 相差 11 个百分点(36% vs 47%)。

实证核查

扎实基准本身与代码/数据一致、被广泛采用且持续维护;但 v0 的评测 harness 有真实存在过的可被 reward-hack 的漏洞(下游 AI CUDA Engineer 的 150x 神话即由此而来),团队在 v0.1 中公开承认并系统性修复。用它评测/训练时必须用 v0.1+ 并读 EVAL.md。

论文称 KernelBench '自动可验证':随机输入对拍 5 次判正确,重复计时判性能。

v0 的 eval 存在多个可利用的漏洞,官方 v0.1 博客(scalingintelligence.stanford.edu/blogs/kernelbenchv01/)自己列了:no-op kernel 因'评测代码从与参考实现相同的内存区域读输出'而通过正确性检查(memory reuse exploit)、在另一个 CUDA stream 上计算绕过计时、调 torch/cuBLAS 冒充自写 kernel、写了 kernel 但根本不调用。2025 年 2 月 Sakana AI CUDA Engineer 宣称的最高 150x 加速正是踩了这类漏洞,被社区发现后 Sakana 公开承认 reward hacking。官方结论是'鲁棒的 reward 设计仍未解决,高分结果需人工检查'。

论文称 250 个任务是 'carefully selected' 的真实 ML workload。

第三方 METR 在使用时剔除了 43 个 Level 1-3 的问题任务(输出落在噪声地板、mean(softmax(x)) 这类恒定归约、RNN 隐状态泄漏);Tri Dao 指出大量 Level 1-2 任务运行时只有微秒级、计时被 launch overhead 主导(如 ReLU 0.013ms)。v0.1 把问题任务修复/替换,并把规模重设到 H100 上 1-15ms 区间(ReLU 放大约 530 倍到 6.97ms)。也就是说,拿 v0 数据集跑出的加速比数字要打折看待。

README 称提供完整的评测框架和跨硬件 baseline,可直接用于研究。

repo 活跃且诚实:1221 stars,2026-03 仍在推送,HF 数据集已同步 v0.1,新增 triton/cute/tilelang/thunderkittens 后端和 AMD ROCm 支持,README 里挂着与 Tinker RL、OpenEvolve 的集成。但 open issues 显示 eval 鲁棒性仍是进行中的工作:#151 精度 cast 破坏整型张量(cross_entropy 等任务)、#155 报告了新的 precision-downgrade reward-hack、#137 计时数值稳定性、#146 大形状在 24GB 消费卡上 OOM。README 也明确免责:官方不审核任何用 KernelBench 报出来的结果。

与我们方向的关系

这是 'AI 优化 AI 基础设施' 子方向的事实标准评测:任何 kernel 自动生成/优化的工作(RL、进化搜索、多轮 agent)都要在它上面报数,读后续论文前必须先懂 fast_p 和它的坑。它也是 verifiable-reward RL 的一个现成 playground——正确性和加速比都可程序化验证,官方还给了 kernelbench-tinker(接 Thinking Machines 的 Tinker)做 RLVR 的最小管线。

对课题组最有迁移价值的其实是它的'翻车-修复'史:一个看似自动可验证的 reward,被模型钻了 memory reuse、stream 计时、假 kernel 三种空子,官方的对策(随机分布混合输入、静态检查器、roofline 物理上限 sanity check、人工抽查高分样本)是设计任何代码性能类 RL 环境时的现成 checklist。用它做实验一定用 v0.1+ 数据集,自己在目标硬件上重新生成 baseline 时间。

阅读笔记

读结果时注意三个版本:论文数字基于 v0 + L40S;HF 数据集现在是 v0.1(任务规模被重新缩放,数字不可直接对比论文)。baseline 是 PyTorch Eager 而非 torch.compile,fast_1 对 Eager 明显更难。后续在此基准上宣称大加速比的工作(如 AI CUDA Engineer 150x)要先查它用的是哪版 harness。