AutoKernel

针对 GPU 内核的自动研究工具。 只需提供任意 PyTorch 模型，然后安心入睡，醒来时即可获得优化后的 Triton 或 CUDA C++ 内核。

灵感来源于 @karpathy/autoresearch —— 该工具展示了用于大语言模型训练研究的自主 AI 代理。AutoKernel 将同样的理念应用于 GPU 内核优化：代理会修改单个文件，运行固定的评估流程，决定保留或回滚更改，并无限循环重复这一过程。

工作原理

向 AutoKernel 提供任意 PyTorch 模型，它将执行以下步骤：

1. 性能剖析：识别哪些 GPU 内核是性能瓶颈。
2. 提取瓶颈内核：将每个瓶颈内核单独提取为独立的 Triton 或 CUDA C++ 内核。
3. 自主优化：对每个内核进行持续优化（编辑、基准测试、保留或回滚——直至满意）。
4. 端到端验证：确保整体正确性，并报告总加速效果。

代理会读取 program.md 文件——即“研究组织代码”——其中包含详细的自主操作指南。它会逐个内核编辑 kernel.py 文件，运行 bench.py 脚本（包含五阶段正确性检查和 Roofline 分析的固定基准测试），并根据结果决定是否保留或回滚更改。调度器会利用阿姆达尔定律来判断何时切换到下一个内核。

每次实验大约耗时 90 秒，因此每小时可完成约 40 次实验，整夜下来则能完成约 320 次，覆盖所有内核。

快速入门

要求：NVIDIA GPU（已测试 H100/A100/RTX 4090）、Python 3.10+、uv。

# 安装 uv（如果尚未安装）
curl -LsSf https://astral.sh/uv/install.sh | sh

# 克隆并设置环境
git clone https://github.com/RightNow-AI/autokernel.git
cd autokernel
uv sync

# 一次性设置：测试数据与基线
uv run prepare.py

# 对模型进行性能剖析（内置 GPT-2、LLaMA、BERT 模型，无需 transformers 库）
uv run profile.py --model models/llama_7b.py --class-name LlamaModel \
 --input-shape 1,512 --dtype float16

# 提取前 5 个瓶颈内核
uv run extract.py --top 5

# 验证基准测试是否正常运行
uv run bench.py

启动代理

在此目录下启动 Claude、Codex 或任何编码代理：

请阅读 program.md，我们开始新的实验吧。先从设置开始。

代理将执行以下操作：

1. 对您的模型进行性能剖析，并给出优化计划。
2. 创建一个分支（例如 autokernel/mar10-llama7b）。
3. 按优先级顺序优化每个瓶颈内核。
4. 验证端到端正确性，并报告总加速效果。

program.md 文件内容详尽，旨在让代理能够连续运行 10 小时以上而不会卡住。它包含了六层优化策略、决策框架、异常处理机制以及基于阿姆达尔定律的推理逻辑。

流程概述

                 profile.py              extract.py           bench.py (循环)         verify.py
任意 PyTorch  ──>  按 GPU 时间排序内核  ──>  生成基线  ──>  优化每个  ──>  端到端
   模型          内核        Triton/CUDA 内核     内核（代理）       验证

工具	功能
profile.py	使用 torch.profiler 对任意 PyTorch 模型进行性能剖析，按 GPU 时间对内核排序，并分类为计算密集型或内存密集型。
extract.py	将前 N 个瓶颈内核提取为独立的 Triton 或 CUDA C++ 内核文件（通过 --backend triton|cuda 指定后端）。
orchestrate.py	多内核调度器：使用阿姆达尔定律决定下一个优化的内核，并跟踪整体进度。
bench.py	固定基准测试：包括五阶段正确性检查（烟雾测试、形状扫描、数值稳定性、确定性、边界情况）以及性能和 Roofline 分析。
verify.py	将优化后的内核重新集成到模型中，检查端到端正确性，并报告总加速效果。

支持的内核类型

涵盖现代深度学习核心操作的 9 种内核类型：

内核	描述	关键指标
matmul	密集矩阵乘法 (M x K) @ (K x N)	TFLOPS
softmax	行并行数值稳定的 softmax	GB/s
layernorm	带仿射变换的层归一化	GB/s
rmsnorm	RMS 归一化（LLaMA 风格）	GB/s
flash_attention	缩放点积注意力，带因果掩码	TFLOPS
fused_mlp	SwiGLU 风格的融合 MLP（门控 + 上升 + 下降）	TFLOPS
cross_entropy	融合交叉熵损失	GB/s
rotary_embedding	旋转位置嵌入（RoPE）	GB/s
reduce	并行规约（求和）	GB/s

每种内核在 reference.py 中都有 PyTorch 参考实现，在 kernels/ 目录中有 Triton 初始内核，在 kernels/cuda/ 目录中有 CUDA C++ 初始内核。

示例模型

AutoKernel 自带了自包含的模型定义（无需 transformers 库）：

模型	文件	参数量	使用方法
GPT-2 Small	models/gpt2.py	1.24 亿	--class-name GPT2 --input-shape 1,1024
LLaMA（精简版）	models/llama_7b.py	1.6 亿	--class-name LlamaModel --input-shape 1,512
LLaMA 7B	models/llama_7b.py	70 亿	--class-name LlamaModel7B --input-shape 1,2048
BERT-base	models/bert_base.py	1.1 亿	--class-name BertModel --input-shape 8,512
自定义	models/custom.py	--	您可以在此基础上构建自己的模型模板。

对于 HuggingFace 模型（需运行 uv sync --extra models）：

uv run profile.py --module transformers --class-name AutoModelForCausalLM \
 --pretrained meta-llama/Llama-2-7b-hf --input-shape 1,2048 --dtype float16

KernelBench 集成

AutoKernel 与 KernelBench 集成，后者是评估 AI 生成 GPU 内核的标准基准测试平台（涵盖 4 个难度级别的 250 多个问题）。虽然大多数 KernelBench 评估采用一次性的 LLM 生成方式，但 AutoKernel 每个问题会进行 50–300+ 次迭代优化实验——系统性地探索优化空间，而非随机猜测。

# 安装 KernelBench 的依赖
uv sync --extra kernelbench

# 从 HuggingFace 获取 Level 1 问题
uv run kernelbench/bridge.py fetch --source hf --level 1

# 设置特定问题以进行优化
uv run kernelbench/bridge.py setup --level 1 --problem 1 --source hf

# 评估（正确性及相对于 PyTorch 参考的加速效果）
uv run kernelbench/bench_kb.py

# 批量评分整个级别（计算 fast_p 指标）
uv run kernelbench/scorer.py --level 1

代理会读取 kernelbench/program_kb.md 文件，以获取 KernelBench 特定的优化指导： 如何编写 ModelNew 类、何时使用 CUDA C++ 或 Triton、针对不同问题级别的融合策略， 以及为 KernelBench 的 fast_p 指标量身定制的编辑-保留/回滚循环。

工具	功能
kernelbench/bridge.py	从 HuggingFace 或本地仓库加载问题，缓存并生成初始 kernel.py 文件
kernelbench/bench_kb.py	对比评估 ModelNew 和 Model：进行 5 次正确性测试 + CUDA 事件计时 + 稳定性 + 确定性
kernelbench/scorer.py	在一个级别内批量评估，并在不同阈值下（1.0x、1.5x、2.0x、3.0x、5.0x）计算 fast_p
kernelbench/program_kb.md	针对 KernelBench 模式的代理操作指南

HuggingFace 内核导出

将优化后的内核导出到 HuggingFace Hub，便于分发。用户只需一行代码即可加载您的内核：

from kernels import get_kernel
module = get_kernel("your-username/kernel-name")

# 导出一个优化后的 CUDA 内核
uv run export_hf.py --name my_matmul

# 上传至 Hub（需先安装 `pip install kernels` 并登录 HuggingFace）
cd workspace/hf_export/my_matmul
kernels upload . --repo_id your-username/my_matmul

项目结构

autokernel/
  kernel.py             代理每次修改的文件（一次只改一个内核）
  program.md            代理的操作说明——“研究组织代码”

  bench.py              固定的基准测试 + 5 阶段正确性测试框架
  reference.py          PyTorch 参考实现（真实结果）
  prepare.py            一次性设置：测试数据、基线

  profile.py            分析任意 PyTorch 模型，按 GPU 时间对内核排序
  extract.py            将瓶颈内核提取到 workspace/ 目录
  orchestrate.py        多内核调度器（阿姆达尔定律）
  verify.py             端到端模型验证 + 加速效果报告
  export_hf.py          将优化后的内核导出为 HuggingFace Kernels 格式
  analysis.py           实验可视化（生成 https://oss.gittoolsai.com/images/RightNow-AI_autokernel_readme_18212f72a48c.png）

  kernels/              初始 Triton 内核（9 种类型）
  kernels/cuda/         初始 CUDA C++ 内核（9 种类型，支持 Tensor Core 加速）
  kernelbench/          KernelBench 集成（桥接、评估框架、评分器）
  models/               自包含的模型定义（GPT-2、LLaMA、BERT）
  workspace/            运行时生成的中间文件（已忽略在 Git 中）

设计选择

双后端：Triton + CUDA C++。 Triton 适合快速迭代（类似 Python 的语法，编译只需几秒）。CUDA C++ 则用于追求极致性能（通过 wmma、PTX 内建指令和共享内存无冲突布局直接访问 Tensor Core）。Triton 的性能通常能达到 cuBLAS 的 80%-95%；而 CUDA C++ 则可以持平甚至超越 cuBLAS。两种后端共享相同的 kernel_fn() 接口——bench.py 在两者上都能以相同方式运行。

正确性第一。 基准测试会在测量性能之前先对比内核输出与 PyTorch 结果。如果内核速度很快但结果错误，则会立即回滚。这样可以防止代理通过生成无效结果来“优化”。

阿姆达尔定律调度。 调度器会根据影响优先级排序。例如，在占总耗时 60% 的内核上实现 1.5 倍加速（整体加速 1.25 倍），远胜于在仅占 5% 的内核上实现 3 倍加速（整体加速仅 1.03 倍）。当收益递减时，调度器会停止进一步优化。

单一可修改文件。 代理仅需修改 kernel.py 文件。这样范围可控，差异易于审查，回滚也更干净。

TSV 日志记录。 结果会保存到纯文本的 results.tsv 文件中。格式简单易读、适合 Git 管理且易于解析，无需额外基础设施。

结果格式

每次实验都会记录到 results.tsv 文件中（制表符分隔）：

列名	描述
experiment	实验序号（0 表示基线）
tag	简短标识符
kernel_type	内核类型（如 matmul）
throughput_tflops	测得的吞吐量（越高越好）
latency_us	执行时间，单位为微秒
pct_peak	占 GPU 理论峰值的比例
speedup_vs_pytorch	相对于 PyTorch/cuBLAS 的加速倍数
correctness	PASS、FAIL、TIMEOUT 或 CRASH
peak_vram_mb	GPU 内存峰值用量
description	尝试的内容

致谢

本项目是 面向 GPU 内核的自动研究 —— 直接受到 Andrej Karpathy 的 autoresearch 启发，该实验最初旨在探索由 AI 代理自主进行 LLM 训练的研究。Karpathy 证明了 AI 代理可以在一夜之间完成数百次实验，系统地探索搜索空间并记录每一步结果。AutoKernel 将这一流程应用到 GPU 内核优化领域，结合 Triton 和原生 CUDA C++，采用迭代优化的方式（每个问题进行 300 多次实验），而非一次性生成。

KernelBench 集成基于斯坦福大学 Scaling Intelligence Lab 的 Simon Guo、Sean Resta 等人的工作。他们的论文《KernelBench：LLM 能否编写 GPU 内核？》（2025 年）确立了评估 AI 生成 GPU 内核的标准基准。AutoKernel 在此基础上进一步扩展，通过迭代优化（每个问题进行 300 多次实验）取代了一次性生成的方式。KernelBench 数据集及评估协议详见：ScalingIntelligence/KernelBench。

由 RightNow AI 构建。如需企业级 GPU 优化服务，请访问 RightNow Enterprise。

更改日志

v1.3.0

• 支持 AMD ROCm GPU：检测并记录 MI300X、MI325X、MI350X、MI355X 的规格信息（感谢 @andyluo7）
• 修复了 verify.py 在 Python 3.13+ 上的 SyntaxError
• 修复了 CUDA flash_attention 忽视 sm_scale 参数的问题
• 修复了 CUDA cross_entropy 返回错误数据类型的问题
• 修复了 Triton rotary_embedding 广播截断的问题
• 修复了 Triton reduce 在非最后一维归约时输出形状不正确的问题

v1.2.0

• 增强了性能分析工具：新增 --export-trace、--memory-snapshot、--torch-compile-log 标志
• 新增通过 export_hf.py 导出至 HuggingFace Kernels 的功能

v1.1.0

• 增加了原生 CUDA C++ 后端，包含 9 个初始内核（支持 Tensor Core、Warp 指令和共享内存分块）
• 集成了 KernelBench（250 多个标准化的 GPU 内核问题）
• 为 extract.py 添加了 --backend triton|cuda 标志

v1.0.0

• 初始版本：基于 Triton 的内核优化流水线，配备 5 阶段正确性测试框架

完整详情请参阅 CHANGELOG.md。

许可证

MIT

AutoKernel 快速上手指南

AutoKernel 是一个用于 GPU 内核自动优化的开源工具。它借鉴了自主 AI 研究代理的理念，能够自动分析 PyTorch 模型，识别性能瓶颈，并通过迭代实验自动生成优化后的 Triton 或 CUDA C++ 内核。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 硬件：NVIDIA GPU（已在 H100/A100/RTX 4090 上测试）。
  • 注：最新版本已支持 AMD ROCm GPU (MI300X 等)。
• 系统：Linux 或 macOS。
• Python：版本 3.10 或更高。
• 包管理器：uv（推荐，用于快速管理依赖和虚拟环境）。

安装步骤

使用 uv 进行安装是最快捷的方式。以下是具体命令：

# 1. 安装 uv (如果尚未安装)
curl -LsSf https://astral.sh/uv/install.sh | sh

# 2. 克隆仓库并进入目录
git clone https://github.com/RightNow-AI/autokernel.git
cd autokernel

# 3. 同步依赖并创建虚拟环境
uv sync

# 4. 一次性初始化设置（生成测试数据和基线）
uv run prepare.py

基本使用

以下是针对一个 PyTorch 模型进行自动内核优化的标准流程。本示例以内置的 LLaMA 模型为例。

1. 性能剖析 (Profile)

首先对模型进行剖析，找出占用 GPU 时间最多的内核瓶颈。

uv run profile.py --model models/llama_7b.py --class-name LlamaModel \
 --input-shape 1,512 --dtype float16

支持内置模型：GPT-2, LLaMA, BERT。无需安装庞大的 transformers 库即可运行内置模型。

2. 提取瓶颈内核 (Extract)

根据剖析结果，提取排名靠前的瓶颈内核，生成独立的 Triton 或 CUDA C++ 文件。

# 提取前 5 个最耗时的内核
uv run extract.py --top 5

3. 验证基准测试 (Verify Benchmark)

在启动 AI 代理之前，先确保基准测试脚本能正常运行。

uv run bench.py

4. 运行 AI 代理 (Run Agent)

启动您的编码代理（如 Claude、Codex 或其他支持代码生成的 LLM），并在项目根目录下输入以下指令：

Read program.md and let's kick off a new experiment. Start with setup.

代理将自动执行以下循环：

1. 读取 program.md 中的优化策略。
2. 创建新的 Git 分支。
3. 按优先级逐个修改 kernel.py 文件。
4. 运行 bench.py 进行 correctness 检查（5 阶段验证）和性能测试。
5. 根据结果保留或回退更改，并持续迭代优化。
6. 最终通过 verify.py 报告端到端的加速比。

提示：每个实验约需 90 秒。代理可整夜运行，自动完成数百次迭代实验。

可选：导出优化后的内核

优化完成后，您可以将内核导出到 HuggingFace Hub 以便分享或使用：

# 导出 CUDA 内核
uv run export_hf.py --name my_matmul

# 上传到 Hub (需先登录: huggingface-cli login)
cd workspace/hf_export/my_matmul
kernels upload . --repo_id your-username/my_matmul