autokernel

AutoKernel：让AI通宵为你的PyTorch模型「私人定制」GPU内核

一套脚本扔进去，醒来发现模型快了两倍——这是不是魔法？

不是魔法，是AutoKernel。

想象这样一个场景：你训练了一个LLaMA 7B模型，推理速度慢得像老牛拉车。常规优化思路是手动写Triton内核，翻文档、调参数、改block size，一搞就是好几天。AutoKernel的思路完全不同——你只需要跑三行命令，然后去睡觉。等你醒来，它已经自动分析了GPU瓶颈、生成了优化内核、跑完了benchmark，告诉你「这套内核让你的模型整体提速1.8倍」。

这是RightNow-AI开源的AutoKernel项目，一套自主AI Agent驱动的GPU内核自动优化框架。它的灵感直接来自Andrej Karpathy的autoresearch——用AI Agent做AI训练研究的自动化。AutoKernel把同样思路搬到了GPU内核优化领域：Agent修改一个文件，跑benchmark测试，决定保留还是回滚，然后继续循环。整个过程无人值守、持续运行，一个晚上能做300+次实验。

AutoKernel 项目 Logo

项目背景：从手工作坊到智能工厂

GPU内核优化长期以来是深度学习工程中「又脏又累」的活儿。一个LLaMA模型的推理瓶颈，往往不在模型架构本身，而在于几个核心算子（矩阵乘法、Softmax、LayerNorm）的GPU实现是否高效。手写一个Triton内核，需要深入理解CUDA/Triton编程模型、硬件内存层次结构、以及Roofline性能模型。业界真正掌握这项技能的人少之又少。

RightNow-AI团队认为，既然LLM可以写代码，AI Agent也应该能够优化GPU内核。他们借鉴了Karpathy在LLM训练自动化上的成功经验——Karpathy的autoresearch证明，给AI Agent足够的背景知识和一个固定的评测环境，它能够持续做出有意义的改进。AutoKernel把这一哲学应用到了内核优化：Agent读取一份详尽的program.md（项目把它叫做「研究组织代码」），参照内置的6层优化 playbook，在固定评测环境中对内核进行「编辑—测试—保留/回滚」的循环迭代。

这套方法论的有效性已经被验证。AutoKernel目前支持9种核心深度学习内核类型：矩阵乘法（matmul）、Softmax、LayerNorm、RMSNorm、Flash Attention、SwiGLU风格的Fused MLP、Cross Entropy、RoPE位置编码、并行归约，覆盖了现代Transformer模型的全部核心算子。每个内核都有PyTorch参考实现、Starter Triton内核、以及可选的CUDA C++内核版本。

核心技术架构：五阶段自动化流水线

AutoKernel的核心是一个五阶段自动化流水线：profile.py → extract.py → bench.py（循环） → verify.py。每一阶段职责明确，通过文件接口串联，构成了一套完整的内核优化闭环。

第一阶段：Profile（性能剖析）。输入任意PyTorch模型（自带GPT-2、LLaMA、BERT定义文件，无需安装transformers），profile.py通过torch.profiler对模型进行GPU时间分布分析，按内核类型（计算密集型/内存密集型）排序，找出真正的性能瓶颈。输出瓶颈内核的优先级列表。

第二阶段：Extract（内核提取）。根据优先级列表，extract.py将每个瓶颈内核提取为独立的Triton或CUDA C++内核文件（--backend triton|cuda）。提取后的内核与原模型解耦，可以在隔离环境中独立测试和优化。

第三阶段：Benchmark循环（核心优化阶段）。这是AutoKernel的灵魂：bench.py为每个内核定义了一套固定的评测协议，包含5个正确性检查阶段（smoke test、shape sweep、数值稳定性、确定性、边界情况）加上性能测试和Roofline分析。Agent每次修改kernel.py中的一个内核（block size、tile策略、split-K参数等），运行bench.py，根据评测结果决定保留或回滚。这个循环以~90秒/次的频率持续运行，约每晚300+次实验。Agent的行为规则全部写在program.md里——这份文档刻意写得非常详尽，确保Agent能在无人监督下运行10小时以上而不陷入死循环。

第四阶段：Verify（端到端验证）。优化后的内核通过verify.py重新插入原模型，执行端到端正确性检查，确保内核优化没有引入数值偏差，并报告整体提速比例。

第五阶段：Orchestrate（多内核调度）。对于同时优化多个内核的场景，orchestrate.py负责调度决策。它基于Amdahl定律计算「哪个内核的优化对整体性能影响最大」，决定下一个优化哪个内核，避免做边际收益低的工作。

技术栈：Triton为主，CUDA为辅，uv管理依赖

从技术实现来看，AutoKernel是围绕Triton构建的。Triton是OpenAI开源的GPU编程语言，可以用Pythonic的语法写高性能GPU内核，无需手写CUDA C++。PyTorch从2.0版本开始原生支持Triton，这使得AutoKernel能够直接与PyTorch的张量无缝对接。

依赖管理上，AutoKernel选择uv而非传统的pip。uv是Astral公司出品的超快速Python包管理器（Rust实现），通过pyproject.toml声明依赖，uv sync即可完成环境搭建。PyTorch通过自定义索引pytorch-cu128拉取CUDA 12.8版本。项目支持多套可选依赖：transformers（模型摄入）、ninja（加速CUDA编译）、datasets（KernelBench基准集成）、HolisticTraceAnalysis（高级性能分析）、huggingface-hub（导出优化内核到Hub）。

当前优化阶段Agent修改的核心文件是kernel.py。这是一个@triton.jit装饰的Triton内核模版，Agent可以在其中调整block大小、tile策略、split-K persistent kernel配置、以及autotune参数。评测框架通过reference.py（PyTorch参考实现）和kernel.py（Triton实现）的输出一致性来判断正确性，通过TFLOPS或GB/s指标判断性能。

部署体验：uv一键安装，但GPU是硬性门槛

从部署角度，AutoKernel的核心约束是必须有NVIDIA GPU——这是最直接的门槛。项目明确测试过H100、A100、RTX 4090三款显卡，其他型号需要自己验证。

安装过程非常顺滑：

git clone https://github.com/RightNow-AI/autokernel.git
cd autokernel
uv sync           # 自动搭建环境，拉取 CUDA 版 PyTorch
uv run prepare.py  # 一次性准备测试数据和基准

之后三行命令完成一次优化流程：profile.py找瓶颈 → extract.py提取内核 → Agent在kernel.py上迭代优化。官方推荐把Claude Code、Codex等编码Agent召唤进来，读program.md后自动运行。program.md包含了完整的优化playbook（6层优化策略：内存合并→tile优化→register tiling→autotune→split-K→persistent kernel），Agent参照执行即可。

没有Dockerfile，意味着无法容器化部署；没有Web UI，是纯CLI工具。这对部署体验有一定影响——用户必须熟悉命令行和GPU环境配置。但考虑到GPU内核优化本身就是在本地高性能环境进行的，CLI模式反而是最自然的选择。

局限与挑战：GPU碎片化、Agent可靠性、适用边界

AutoKernel并非万能解药。首先，GPU硬件环境是硬性依赖，且不同GPU架构（Hopper、Ada Lovelace、Ampere）的优化策略差异很大。在非测试过的显卡上，Agent给出的优化方向可能无效。其次，Agent的行为可靠性完全依赖program.md的质量——如果Agent陷入了局部最优或跑偏了方向，人需要介入。第三，适用范围有边界：AutoKernel擅长的是Triton可表达的算子融合优化，对于需要手写CUDA寄存器级别的深度优化，自动化方法的效果有限。

从代码质量看，项目结构清晰，单文件kernel.py设计降低了Agent的操作复杂度；但测试覆盖和文档质量（除program.md外）相对薄弱——这是一个原型导向的项目，代码本身不是重点，持续的Agent实验才是。

行业意义：内核优化的「自动驾驶」时代来临

AutoKernel代表了GPU性能优化领域的一个趋势转变：从专家手工优化，到AI Agent持续自动化搜索。传统上，内核优化是少数CUDA高手的专利；AutoKernel通过把优化过程形式化为「评测—修改—决策」循环，让AI Agent成为内核工程师的「数字实习生」——做一些试错性的baseline优化，而专家去做更复杂的架构决策。

这个方向的天花板取决于Agent的推理能力。当前Agent在program.md的引导下能够做参数调优（block size、tile factor），但真正创新的优化策略（全新的算法替换）仍然需要人类介入。随着LLM推理能力提升，AutoKernel式的方法论有望渗透到更多硬件后端（AMD ROCm、Apple Metal等），成为硬件无关的性能优化基础设施。

如果你有NVIDIA GPU、有PyTorch模型、有耐心等一晚上，AutoKernel值得一试。毕竟——「去睡觉，醒来模型变快」，这本身就是一个足够有吸引力的价值主张。