RightNow-AI/autokernel

它解决了什么问题

编写快速的 GPU 内核（运行神经网络的底层代码）既缓慢又需要深厚的专业知识。AutoKernel 通过让编码 AI 智能体处理你的模型来实现自动化：它会找出最慢的部分，将其提取为独立的内核，并不断尝试优化，直到它们运行得更快。你无需手动调整每个操作，只需提供模型，即可获得经过优化且具有测量加速比的内核。

工作原理

你提供任何 PyTorch 模型，AutoKernel 会运行以下流水线：

1. Profile（分析）：使用 PyTorch 的分析器对耗时最长的 GPU 内核进行排名。
2. Extract（提取）：将每个瓶颈提取为独立的 Triton 或 CUDA C++ 内核文件。
3. Optimize（优化）：AI 编码智能体（如 Claude 或 Codex）读取一组指令 (program.md)，编辑一个内核文件，运行固定的基准测试，并且仅在通过正确性检查且速度提升时保留更改。它会不断重复此循环，每次实验大约耗时 90 秒。
4. Verify（验证）：在优化所有瓶颈后，它会将新内核插回模型中，并报告总体的端到端加速比。

基准测试在测量性能之前包含 5 阶段正确性检查（冒烟测试、形状扫描、数值稳定性、确定性、边缘情况）。编排器使用 Amdahl 定律来决定接下来处理哪个内核，优先处理对整个模型影响最大的内核。

适用人群

需要为 PyTorch 模型提供更快的 GPU 推理或训练，但不想手写 CUDA 或 Triton 内核的开发者和研究人员。它需要 NVIDIA GPU（已在 H100/A100/RTX 4090 上测试）和 Python 3.10+。它附带了可直接运行的模型定义（GPT-2, LLaMA, BERT），因此你无需外部模型库即可尝试。

亮点

• 自动优化：无需监督，在夜间运行数百次实验，编辑文件并仅保留成功的更改。
• 双后端：生成 Triton（快速迭代）和原生 CUDA C++（最大性能）内核，共享相同的接口。
• 9 种内核类型：涵盖常见的深度学习操作，如 matmul, softmax, layernorm, flash attention, fused MLP 等。
• KernelBench 集成：与 AI 生成的 GPU 内核标准基准测试配合使用，进行迭代细化而非一次性猜测。
• 正确性优先：任何未通过 5 阶段正确性测试的内核都会立即还原，因此加速始终是安全的。
• 简单输出：每次实验都会记录到简单的 results.tsv 文件中，以便于跟踪。