KernelBench：AI 能否写出超越人类的 GPU 内核？

想象一个场景：你是一家 AI 公司的 GPU 性能优化工程师，面对一个紧迫的产品发布 deadline，需要为一个新算子编写高度优化的 CUDA 内核。传统方式需要数周专业知识积累——查阅 NVIDIA 文档、反复用 nsight-compute 调优、对比 cuDNN 基准。而现在，有人告诉你：把这个 PyTorch 算子描述丢给大模型，它能在几分钟内生成一个正确且可能更快的 CUDA 内核。

这就是 KernelBench 试图回答的核心问题。

图1：KernelBench 评测流程：LLM 生成 CUDA 内核，与 PyTorch 参考实现对比正确性和性能

背景：从 AI 写代码到 AI 写 GPU 内核

GPU 内核（Kernel）编写长期以来被视为最高难度的编程任务之一。它要求开发者同时具备并行计算理论、硬件架构知识（CUDA/ROCm）、以及大量工程调优经验。一个优秀的 CUDA 工程师需要 years of experience 才能达到生产级水平。

KernelBench 由斯坦福大学 Scaling Intelligence 实验室于 2024 年 10 月发布，2025 年初在 ICML 会议上正式发表，旨在系统性地量化评估当前大语言模型（LLM）在 GPU 内核编写领域的真实能力边界。核心动机是探索「AI 能否理解底层硬件约束并生成高效并行代码」。

这个问题的答案意义重大：若 AI 能可靠地写出高质量 GPU 内核，将彻底改变 AI 编译器和硬件加速器的工作方式——从手工调优时代进入「AI 生成 + AI 评测」的自动化循环。

核心功能：四级难度阶梯，覆盖从算子到模型全链路

KernelBench 构建了一套结构化的评测体系，将任务按难度分为四个递进级别：

Level 1（基础）：单算子内核（100 题） 考察模型将独立 PyTorch 算子翻译为 CUDA 的能力。题目覆盖神经网络中最基础的构建块：矩阵乘法（MatMul）、卷积（Conv2d）、LayerNorm、Softmax 等。这些问题的共同特点是：参考实现清晰、正确性验证相对直接、性能对比有明确基准。

Level 2（进阶）：融合算子模式（100 题） 将多个算子融合为单一内核，典型场景如 Conv + Bias + ReLU 或 MatMul + Scale + Sigmoid。融合内核的难点在于：需要模型理解算子间的数据流依赖、寄存器复用策略、以及共享内存的合理使用。融合优化的性能收益可能达到数倍，但正确性验证也更复杂。

Level 3（挑战）：完整模型架构（50 题） 要求模型优化整张神经网络：MobileNet、VGG、MiniGPT、Mamba 等。题目需要模型理解端到端的数据流、内存布局、以及跨层优化机会。评测维度从单内核扩展到了多内核协同和内存带宽利用。

Level 4（前沿）：HuggingFace 集成 直接从 HuggingFace 模型库选取架构进行评测，跟踪业界最新模型的 GPU 优化潜力。

评测方法论：不止于「能跑通」

KernelBench 的评测不止验证内核是否正确执行，更重要的是衡量性能提升幅度。评测系统执行以下两步：

正确性验证：在随机化输入上重复执行 n 次（n_correctness），对比 LLM 生成内核与 PyTorch 参考实现的输出是否一致。容差标准参考 PyTorch Benchmarking Suite：FP32 采用 1e-4，FP16/BF16 采用 1e-2。

性能对比：在 n_trial 次测量中记录 wall-clock 时间，计算加速比。核心指标为 fast_p：正确且加速比超过阈值 p 的任务比例。例如 fast_1 代表「正确且比 PyTorch 快」的比例，fast_2 代表「正确且快 2 倍以上」的比例。

项目特别强调防止 Reward Hacking：模型可能在 RL 训练或进化搜索中寻找评测漏洞。项目维护了一份 Hack 和 Defense 的公开博客记录，鼓励研究者提交检测和缓解工具。

技术架构：模块化设计，支持本地+云端双模式

代码组织清晰，核心模块位于 src/kernelbench/：

eval.py（34873 字）：评测逻辑核心，含正确性验证（n 次随机输入对比）、性能计时（wall-clock 测量）、容差管理（按精度类型自动选择阈值）。支持 fetch_ref_arch_from_problem_id 从数据集获取参考架构。

compile.py（8048 字）：并行编译与缓存管理，使用 multiprocessing 加速批量编译，支持构建目录缓存复用。build_compile_cache 函数处理编译错误收集和日志输出。

dataset.py（19659 字）：统一数据集抽象，同时支持本地文件系统（KernelBench/ 目录下的题库文件）和 HuggingFace Dataset 两种来源，按 problem_id（1-indexed）提供统一的题目标识接口。Problem dataclass 忽略注释和空白的代码哈希用于去重。

score.py：评测指标计算，包含几何平均加速比（仅正确样本）、几何平均加速比（仅正确且加速样本）、fast_p 指标。

profile.py：性能分析工具，整合 NVIDIA nsight-compute 分析能力。

依赖层面：核心依赖包括 PyTorch 2.9+、transformers、datasets、Modal（云端运行）、Triton（GPU 编程框架）。可选 GPU 依赖（extra gpu）包含 Triton、nvidia-cutlass-dsl、Tilelang、CuPy-cuda12x、nsight-python。LLM 调用通过 litellm 统一接口，支持 OpenAI API 和代理模式。项目还支持 AMD ROCm 后端（需 ROCm >= 7.1）。

云端评测：项目推荐使用 Modal 云端函数运行评测，以保证环境一致性和可复现性。用户无需本地 GPU 即可参与评测流程，但生成和运行内核仍需 GPU。

上手体验：无 Web UI，本地部署门槛较高

KernelBench 是一个纯 Python CLI 工具，没有 Web 界面。安装依赖需要 uv 包管理器（推荐）或 conda，Python 版本固定为 3.10。

本地评测的硬件门槛相当高：需要 NVIDIA GPU（推荐 RTX 3090 / A100）、16GB+ VRAM、16GB+ RAM、10GB+ 磁盘空间，以及 CUDA 12+ 驱动。此外还需配置 LLM API Key（通过 .env 文件和 litellm 代理）。

若没有本地 GPU，项目支持通过 Modal 云端环境运行，同样需要 GPU 执行内核生成和评测。整体来看，这是一个面向 AI 系统研究者的专业工具，普通用户上手难度较大。

争议与局限：模型「作弊」与基准公平性

KernelBench 文档中明确列出了若干已知问题，值得研究者关注：

Reward Hacking：模型可能在评测过程中寻找漏洞，例如在测试输入上针对特定值优化而非泛化到所有输入。项目维护者坦言「if the model can reward hack, it will find ways」，这是 RL 训练中的经典问题。

PyTorch 参考实现的公平性：cuDNN 本身是高度优化的工业级库，若 LLM 生成的内核声称比 cuDNN 快 2 倍以上，项目文档建议「think again」（引用 NVIDIA 工程师的公开评论）。大多数「显著加速」结果需要严格审查。

评测可复现性：LLM 输出具有随机性，同一模型多次运行可能生成不同内核，性能差异可能很大。项目通过 set_seed 控制部分随机性，但 LLM 采样本身的随机性难以完全消除。

评测覆盖范围：当前基准以 NVIDIA CUDA 为主，AMD ROCm 支持仍在实验中，其他 DSL 扩展也在进行中，覆盖范围有限。

行业意义：AI 代码生成的「深水区」

KernelBench 的发布填补了 LLM 在底层系统编程领域评测的空白。在此之前，LLM 代码生成评测主要集中在 Python、JavaScript 等高级语言，或 LeetCode 算法题。GPU 内核编写代表了一个新的「深水区」——它要求模型理解硬件约束（寄存器、共享内存、warp 分支）、并行算法、以及编译器行为。

从当前结果来看，顶级模型（GPT-4、Claude 等）在 Level 1 基础算子上已有不错表现，但在 Level 3 完整模型优化上仍有显著差距。这与人类 GPU 工程师的经验判断一致：理解单一算子的 CUDA 转换相对容易，但优化整张模型的内存访问模式需要更深的硬件直觉。

随着项目持续更新（当前为 v0.2.0-dev），支持更多 DSL（除了 CUDA 还有 Triton DSL 等），以及 AMD ROCm 扩展，KernelBench 有望成为 LLM + GPU 编译器领域的标准评测基准，吸引更多研究者关注 AI 在底层系统编程的潜力与局限。