AI News
商业产业 Hacker News LLM 2026-10-02T01:13 正文 ok 1714 words

KernelBench：LLM 能编写高效 GPU Kernel 吗？[ICML '25]

英文标题：🔶 KernelBench: Can LLMs Write GPU Kernels? – Benchmark and Toolkit, Torch –> CUDA

KernelBench 是一套用于评估 LLM 能否将 PyTorch 程序转译为正确且高效的 CUDA 或其他 GPU DSL Kernel 的基准与工具包。基准包含四个层级：100 个单 Kernel 算子、100 个简单融合模式、50 个完整模型架构任务，以及面向 Hugging Face 整体模型优化的 Level 4。项目以 `fast_p` 衡量结果，即同时满足正确性且相较 PyTorch 基线加速超过阈值 p 的任务比例。该工具支持 CUDA、Triton、HIP、CuTe、TileLang 和 ThunderKittens，并提供本地或 Modal 云端评测、迭代优化、RLVR 和演化搜索等扩展方向。
打开原文相关 3 · 机会 0
media124.4 KB media657.2 KB media82.1 KB

中文

KernelBench：LLM 能编写高效 GPU Kernel 吗？[ICML '25]

一个用于评估 LLM 生成高效 GPU Kernel 能力的基准与环境。

具体而言，我们要求 LLM 为目标 GPU 上的 PyTorch 程序生成正确且高效的 CUDA / DSL Kernel。

arXiv | 博客文章 | HuggingFace 数据集

最新稳定版本位于 main 分支。我们会持续更新和改进该仓库。

Huggingface 数据集已更新至 v0.1。

本仓库提供 KernelBench 的核心功能，以及一套易于使用的评测脚本。其目的并非提供能够解决该任务的复杂 Agent 脚手架；我们建议你克隆并修改本仓库以开展实验，或将其作为 git 子模块使用。

我们将问题设计为：让 LLM 以其选择的任意粒度，将 PyTorch 所描述的算子转译为 CUDA Kernel。KernelBenchMascot

我们将 KernelBench 构建为 4 个层级的类别：

  • Level 1 🧱：单 Kernel 算子（100 个问题）。神经网络的基础构件，包括卷积、矩阵乘法和 Layer Normalization。
  • Level 2 🔗：简单融合模式（100 个问题）。融合后的 Kernel 应比拆分的多个 Kernel 更快，例如 Conv + Bias + ReLU、Matmul + Scale + Sigmoid。
  • Level 3 ⚛️：完整模型架构（50 个问题）。端到端优化完整模型架构，例如 MobileNet、VGG、MiniGPT、Mamba。
  • Level 4 🤗：Hugging Face 层级。优化来自 Hugging Face 的完整模型架构。

我们也在积极将 KernelBench 扩展至 cuda 之外的其他 DSL（见下文），并增加对 AMD GPU 的支持。

要评估模型生成的 Kernel，我们需要检查它们是否：

  • 正确 ✅：在随机输入上，与参考 torch 算子进行 n_correctness 次对比检查。
  • 具备性能优势 ⏱️：与参考 torch 算子进行 n_trial 次比较，以测量运行时间之间的加速比。

请查看 src/eval.py，了解我们如何实现正确性检查和计时；也可查看 EVAL.md 获取评估及基准测试指南说明 [WIP]。

我们提供便捷脚本 scripts/run_and_check.py，用于针对参考源代码评估单个样本源代码、检查正确性并计算加速比。你可以通过设置 eval_mode=local 或 eval_mode=modal，在本地或远程评估 Kernel。

由于我们需要同时衡量正确性和性能，因此定义了指标 fast_p：同时正确且加速比大于阈值 p 的任务占比；加速比计算为 PyTorch 参考实现的墙钟时间与生成 Kernel 时间的比值。

以下示例说明了这个根据加速比进行筛选的指标：

  • fast_1：LM 生成的 Kernel 同时正确且快于 PyTorch 基线的任务占比。
  • fast_2：LM 生成的 Kernel 同时正确且比 PyTorch 基线至少快 2 倍的任务占比。
  • fast_0：LM 生成的 Kernel 正确的任务占比（与正确率相同）。

你可以提高加速比阈值 p，使任务更具挑战性。

计算整体基准性能

我们提供脚本 scripts/greedy_analysis.py 来计算整体基准性能。由于我们需要同时衡量正确性和性能，因此使用指标 fast_p：同时正确且加速比大于阈值 p 的任务占比；加速比计算为 PyTorch 参考实现的墙钟时间与生成 Kernel 时间的比值。

我们将仓库组织为以下结构：

KernelBench/
├── assets/
├── KernelBench/ # 基准数据集文件
├── src/kernelbench/ # KernelBench 逻辑代码
│   ├── unit_tests/  
│   ├── prompts/
│   ├── ....
├── scripts/ # 用于运行基准测试的辅助脚本
├── results/ # 跨硬件的基线时间 
├── runs/ # 存储运行结果的位置
├── notebooks/ # 用于分析的示例 notebook
├── pyproject.toml # 项目配置与依赖项

我们已迁移至使用 pyproject.toml 和 uv 管理依赖。如果你尚未安装，请安装 uv。

# 安装基础依赖（无需本地 GPU 即可运行）
uv sync

# 使用 AMD ROCm 后端安装（需要 ROCm>=7.1）
uv add torch --index pytorch=https://download.pytorch.org/whl/rocm7.1

# 安装 GPU 依赖（用于本地 GPU 评估）
uv sync --extra gpu

# 使用 uv 运行命令（会调用正确的环境）
uv run python scripts/<script_name>.py ...

对于 AMD GPU，即 ROCm 后端（ROCm>=7.1），请添加 uv remove torch && uv add torch --index pytorch=https://download.pytorch.org/whl/rocm7.1，以使用兼容 ROCm 的 PyTorch 配置依赖。由于 ROCm 配置较为复杂，建议在 Docker 镜像中运行。

你仍可使用 conda (python=3.10) 创建环境，并通过 requirements.txt 安装依赖。

我们使用 litellm 进行 API 调用。请按照 .env.example 创建 .env 文件并设置你的密钥。

运行及性能分析 Kernel 需要 GPU。如果你本地没有可用 GPU，可以配置 Modal 以进行云端无服务器 GPU 评估。创建账户后，运行 modal token new 设置 Modal token。随后使用 generate_and_eval_single_sample_modal.py 脚本。

你也可以试用我们的教程 notebook（也位于 notebooks/tutorial.ipynb）以及 Google Colab。

从单个问题开始会更容易。以下命令将获取问题、生成一个样本并评估该样本。

# 例如：从 Hugging Face 运行 Level 2 的第 40 个问题，并使用 Google Gemini 2.5 Flash 生成

uv run python scripts/generate_and_eval_single_sample.py dataset_src=huggingface level=2 problem_id=40 server_type=google model_name=gemini/gemini-2.5-flash

# dataset_src 可以为 "local" 或 "huggingface"
# 添加 .verbose_logging 可获得更详细的日志

你可能需要修改的内容

  • gpu_arch - 视你的 GPU 而定，你可能需要调整 gpu_arch 参数以反映硬件配置。
  • precision - 可通过 precision=fp32 指定张量精度。当前我们报告的所有结果均为 fp32，但已增加对 fp16 和 bf16 的支持。
  • backend - 我们也支持 cuda 之外的其他 GPU 编程语言。例如，只需指定 backend=triton 或 backend=hip。目前，我们支持 NVIDIA GPU 及其编程框架和 DSL：cuda、triton、cute、tilelang、thunderkittens。

AMD GPU 注意事项：请使用 hip 后端；当前支持的 gpu_arch 为 gfx942、gfx950。

本地配置 ThunderKittens（TK）的注意事项：如需使用 backend=thunderkittens，你需要 git clone ThunderKittens 仓库，并设置以下环境变量，使其指向本地 ThunderKittens 目录：export THUNDERKITTENS_ROOT=<PATH to ThunderKittens folder>。此外，如示例所示，所有 ThunderKitten 程序都应包含 tk_root = os.environ.get("THUNDERKITTENS_ROOT", "/root/ThunderKittens")，从而让 Kernel 可以包含正确的 TK 原语。另请注意，目前我们仅支持 TK 的 BF16。

请查看配置字段以了解完整选项集。注意，默认情况下，我们会向模型提供一个 one-shot 示例及最小必要信息；你可以在 src/prompt_constructor_toml.py 中查看其他提示词设置，或自行构建设置。

# 1. 生成响应，并将 Kernel 本地存储至 runs/{run_name} 目录
uv run python scripts/generate_samples.py run_name=test_hf_level_1 dataset_src=huggingface level=1 num_workers=50 server_type=deepseek model_name=deepseek-chat temperature=0

# 2. 评估 runs/{run_name} 目录中的所有已生成 Kernel
uv run python scripts/eval_from_generations.py run_name=test_hf_level_1 dataset_src=local level=1 num_gpu_devices=8 timeout=300

# 如果希望加快评估，可以在进入 GPU 评估前先通过 CPU 并行执行编译
# 在命令中添加 build_cache=True 和 num_cpu_workers=<num_cpu_workers>

分析评估结果以计算基准性能

我们提供 scripts/benchmark_eval_analysis.py，用于分析评估结果，计算成功率、计时指标以及整体基准性能 fast_p。

uv run python scripts/benchmark_eval_analysis.py run_name=test_hf_level_1 level=1 hardware=L40S_matx3 baseline=baseline_time_torch

如果使用不同硬件，可以通过 scripts/generate_baseline_time.py 脚本生成基线时间。我们在 results/timing 中提供了跨多代 NVIDIA GPU 的一些参考基线时间，但建议你自行生成基线时间以获得更准确的结果，因为集群功耗、软件版本等因素都会影响计时结果。详情请参见 results/timing/README.md。

请查看我们的路线图，了解计划增加的功能。我们欢迎社区在这些方向上贡献和讨论。

你也可以将 KernelBench 作为项目中的库使用，例如：from kernelbench import timing、from kernelbench import eval as kb_eval，或 from kernelbench.utils import set_gpu_arch。

  • 与 Harbor 的适配器 —— 我们正与 Harbor 集成，以实现更高吞吐量的评估，并提供超越模型 pass@1/k 的、更丰富的 Agent 性能评估。(进行中)

  • 多轮 / 测试时扩展 —— Caesar 是我们面向吞吐量优化的多轮推理引擎（ICML '25），用于论文中的迭代优化实验。它会批量运行生成轨迹，并在各轮之间反馈正确性、运行时间和性能分析信号，以实现序列式测试时扩展。

  • 强化学习（RLVR） —— kernelbench-tinker 是与 Thinking Machines Lab 的 Tinker RL library 的端到端集成。该流水线让策略模型生成 Kernel，通过 Modal 在云端 GPU 上评估，并将结果转化为 RL 奖励——这是一个用于实验 GPU Kernel 优化 RLVR 的轻量级试验场。

  • 演化搜索 —— AlphaEvolve 等演化搜索已展现出为优化问题发现创新解法的潜力。我们正在开发与 OpenEvolve 的集成，即将发布。

  • Roofline / 最大加速比分析 —— （实验性，WIP） simple-torchroofline 为 PyTorch 程序提供分析式 Roofline 分析，可估算目标 GPU 的理论极限（SoL）计算与内存边界——无需硬件。结合基于硬件计数器的性能分析进行经验 Roofline 分析，该工具有助于合理性检查报告的加速比在物理上是否现实。

自发布以来，我们收到了众多研究人员、研究实验室和企业的关注，他们使用 KernelBench 探索这一方向。我们已记录 KernelBench 及相关自动化 Kernel 生成工作已知的使用案例。如果你正在使用 KernelBench，我们很希望了解更多信息！

免责声明：KernelBench 被设计为一个开源评估框架和工具包。KernelBench 团队不会审查、验证或认可单个 Kernel 或已报告的结果。用户有责任独立验证使用该框架获得的任何结果。有关基准测试与 Kernel 评估的更多指南，请查看 EVAL.md。

MIT。更多详情请查看 LICENSE.md。

@misc{ouyang2025kernelbenchllmswriteefficient,
      title={KernelBench: Can LLMs Write Efficient GPU Kernels?}, 
      author={Anne Ouyang and Simon Guo and Simran Arora and Alex L. Zhang and William Hu and Christopher Ré and Azalia Mirhoseini},
      year={2025},
      eprint={2502.10517},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2502.10517}, 
}

我们感谢 GPU Mode、PyTorch、Modal Labs 以及更广泛开源社区所提供的支持，这些支持使本项目成为可能。

English / 原文

KernelBench: Can LLMs Write Efficient GPU Kernels? [ICML '25]

A benchmark and environment for evaluating LLMs' ability to generate efficient GPU kernels

Specifically we task LLM to generate correct and efficient CUDA / DSL kernels for PyTorch programs on a target GPU.

arXiv | blog post | HuggingFace Dataset

The latest stable version will be on main branch. We continue to update and improve the repo.

The Huggingface dataset is updated to v0.1.

This repo provides core functionality for KernelBench and an easy-to-use set of scripts for evaluation. It is not intended to provide complex agentic scaffolds that solve this task; we recommend cloning and modifying this repo for your experiment, or using it as a git submodule.

We structure the problem for LLMs to transpile operators described in PyTorch to CUDA kernels, at whatever level of granularity they desire. KernelBenchMascot

We construct KernelBench to have 4 Levels of categories:

  • Level 1 🧱 : Single-kernel operators (100 Problems) The foundational building blocks of neural nets (Convolutions, Matrix multiplies, Layer normalization)
  • Level 2 🔗 : Simple fusion patterns (100 Problems) A fused kernel would be faster than separated kernels (Conv + Bias + ReLU, Matmul + Scale + Sigmoid)
  • Level 3 ⚛️ : Full model architectures (50 Problems) Optimize entire model architectures end-to-end (MobileNet, VGG, MiniGPT, Mamba)
  • Level 4 🤗 : Level Hugging Face Optimize whole model architectures from HuggingFace

We are actively extending KernelBench to other DSLs beyond cuda as well (see below), as well as AMD GPU support.

To evaluate model-generated kernels, we need to check if they:

  • are correct ✅ : check against reference torch operators n_correctness times on randomized inputs.
  • are performant ⏱️ : compare against reference torch operators n_trial times to measure speedup between runtimes.

Check out src/eval.py for details on how we implement correctness check and timing and EVAL.md for notes on evaluation and benchmarking guidelines [WIP].

We provide a convenient script scripts/run_and_check.py to evaluate one single sample source code against a reference source code, check correctness and compute speedup. You can use this to evaluate a kernel either locally or remotely by setting eval_mode=local or eval_mode=modal.

Since we need to capture both correctness and performance, we define a metric fast_p: fraction of tasks that are both correct and have a speedup greater than threshold p; speedup is computed as the ratio of PyTorch reference wall-clock time to generated kernel time.

Some examples to illustrate this metric that filters based on speedups:

  • fast_1 is the fraction of tasks that LM-generated kernels are both correct and faster than PyTorch baseline
  • fast_2 is the fraction of tasks that LM-generated kernels are both correct and at least 2x faster than PyTorch baseline
  • fast_0 is the fraction of tasks that LM-generated kernels are correct. (same as correctness rate)

You can increase speedup threshold p to make the task more challenging.

Compute Overall Benchmark Performance

We provide a script scripts/greedy_analysis.py to compute the overall benchmark performance. Since we need to capture both correctness and performance, we use a metric fast_p: fraction of tasks that are both correct and have a speedup greater than threshold p; speedup is computed as the ratio of PyTorch reference wall-clock time to generated kernel time.

We organize the repo into the following structure:

KernelBench/
├── assets/
├── KernelBench/ # Benchmark dataset files
├── src/kernelbench/ # KernelBench logic code
│   ├── unit_tests/  
│   ├── prompts/
│   ├── ....
├── scripts/ # helpful scripts to run the benchmark
├── results/ # baseline times across hardware 
├── runs/ # where your runs will be stored
├── notebooks/ # example notebooks for analysis
├── pyproject.toml # Project configuration and dependencies

We have transitioned to using pyproject.toml and uv for dependency management. Install uv if you haven't already

# Install base dependencies (works without a local GPU)
uv sync

# Install with AMD ROCm backend (ROCm>=7.1 is required)
uv add torch --index pytorch=https://download.pytorch.org/whl/rocm7.1

# Install with GPU dependencies (for local GPU evaluation)
uv sync --extra gpu

# Run commands with uv (which invoke the right env)
uv run python scripts/<script_name>.py ...

For AMD GPU aka ROCm backend (ROCm>=7.1), please add uv remove torch && uv add torch --index pytorch=https://download.pytorch.org/whl/rocm7.1 for RoCm compatible PyTorch to configure your dependencies. Running in a docker image is recommended for this due to complexity of ROCm setup.

You can still use conda (python=3.10) to create your environment and install dependencies with requirements.txt.

We use litellm for API calls. Please set your keys by creating a .env following our .env.example.

Running and profiling kernels require a GPU. If you don't have a GPU available locally, you can set up Modal for cloud serverless GPU evaluation. Set up your modal token after creating an account by running modal token new. Then, use the generate_and_eval_single_sample_modal.py script.

You can also try out our tutorial notebook (also in notebooks/tutorial.ipynb) with Google Colab.

It is easier to get started with a single problem. This will fetch the problem, generate a sample, and evaluate the sample.

# for example, run level 2 problem 40 from huggingface and use google gemini 2.5 flash for generation

uv run python scripts/generate_and_eval_single_sample.py dataset_src=huggingface level=2 problem_id=40 server_type=google model_name=gemini/gemini-2.5-flash

# dataset_src could be "local" or "huggingface"
# add .verbose_logging for more visibility

What you might need to modify

  • gpu_arch - Depending on your GPU, you might need to adjust the gpu_arch argument to reflect your hardware.
  • precision - You can specify the precision of tensor by precision=fp32. Currently all of our reported results are fp32 but we added support for fp16 & bf16.
  • backend - We are also supporting other GPU programming languages beyond cuda. For example, simply specify backend=triton or backend=hip. For now we support NVIDIA GPUs with programming frameworks and DSLs: cuda, triton, cute, tilelang, thunderkittens.

Note for AMD GPUs: Use hip backend, gpu_arch currently supported: gfx942, gfx950.

Note on setting up ThunderKittens (TK) locally: to use backend=thunderkittens, you need to git clone the ThunderKittens repo and set the following environment variable to point to your local ThunderKittens directory, export THUNDERKITTENS_ROOT=<PATH to ThunderKittens folder>, and all ThunderKitten programs as shown in the example, should contain tk_root = os.environ.get("THUNDERKITTENS_ROOT", "/root/ThunderKittens"), which enable the kernel to include the right TK primitives. In addition, we only support BF16 for TK right now.

Check the config fields for a comprehensive set of options. Note we provide the model with a one-shot example by default along with the minimum set of info; you can check out other prompt settings or construct your own in src/prompt_constructor_toml.py.

# 1. Generate responses and store kernels locally to runs/{run_name} directory
uv run python scripts/generate_samples.py run_name=test_hf_level_1 dataset_src=huggingface level=1 num_workers=50 server_type=deepseek model_name=deepseek-chat temperature=0

# 2. Evaluate on all generated kernels in runs/{run_name} directory
uv run python scripts/eval_from_generations.py run_name=test_hf_level_1 dataset_src=local level=1 num_gpu_devices=8 timeout=300

# If you like to speedup evaluation, you can use parallelize compilation on CPUs before getting to evaluation on GPUs
# add build_cache=True and num_cpu_workers=<num_cpu_workers> to the command

Analyze the eval results to compute Benchmark Performance

We provide scripts/benchmark_eval_analysis.py to analyze the eval results to compute success rate, timing metric, and overall benchmark performance fast_p.

uv run python scripts/benchmark_eval_analysis.py run_name=test_hf_level_1 level=1 hardware=L40S_matx3 baseline=baseline_time_torch

If you are using a different hardware, you can generate the baseline time with scripts/generate_baseline_time.py script. We provide some reference baseline times for a variety of NVIDIA GPUs across generations in results/timing, but we recommend you generate your own baseline time for more accurate results (cluster power, software version, all affect timing results). See results/timing/README.md for more details.

Check out our roadmap for what we plan to add as features. We welcome community contributions and discussions in these directions.

You can also use KernelBench as a library for your projects, for example: from kernelbench import timing, from kernelbench import eval as kb_eval, or from kernelbench.utils import set_gpu_arch.

  • Adapter with Harbor — We are integrating with Harbor to enable higher-throughput eval and richer evaluation of agentic performance beyond model pass@1/k. (Ongoing)

  • Multi-Turn / Test-Time Scaling — Caesar is our throughput-oriented multi-turn inference engine (ICML '25), used for the iterative refinement experiments in the paper. It runs generation trajectories in batch, feeding back correctness, runtime, and profiling signals across turns for sequential test-time scaling.

  • Reinforcement Learning (RLVR) — kernelbench-tinker is an end-to-end integration with Thinking Machines Lab's Tinker RL library. The pipeline has the policy model generate kernels, evaluates them on cloud GPUs via Modal, and converts results into RL rewards — a minimal playground for experimenting with RLVR on GPU kernel optimization.

  • Evolutionary Search — Evolutionary search like AlphaEvolve has shown promise for discovering innovative solutions for optimization problems. We are working on an integration for OpenEvolve. Releasing soon.

  • Roofline / Max Speedup Analysis — (Experimental, WIP) simple-torchroofline provides analytical roofline analysis for PyTorch programs, estimating the speed-of-light (SoL) compute and memory bounds for a target GPU — no hardware required. Combined with hardware counter-based profiling for empirical roofline analysis, this helps sanity-check whether a reported speedup is physically realistic.

Since release, we have gotten a lot of interest from researchers, research labs, and companies that use KernelBench to explore this direction. We have documented known usage of KernelBench and related efforts towards automated kernel generations. If you are using KernelBench, we love to hear more about it!

Disclaimer: KernelBench is designed as an open-source evaluation framework and toolkit. The KernelBench team does not review, validate, or endorse individual kernels or reported results. Users are responsible for independently verifying any results obtained using the framework. Please check out EVAL.md for more guidance on benchmarking and evaluating kernels.

MIT. Check LICENSE.md for more details.

@misc{ouyang2025kernelbenchllmswriteefficient,
      title={KernelBench: Can LLMs Write Efficient GPU Kernels?}, 
      author={Anne Ouyang and Simon Guo and Simran Arora and Alex L. Zhang and William Hu and Christopher Ré and Azalia Mirhoseini},
      year={2025},
      eprint={2502.10517},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2502.10517}, 
}

We are grateful for the support from GPU Mode, PyTorch, Modal Labs and the broader open-source community that made this project possible.