AutoKernel自动优化GPU核函数，无需人工干预即能显著提速深度学习模型。
AutoKernel: Autonomous GPU Kernel Optimization via Iterative Agent-Driven Search
- 用自主智能体循环搜索优化Triton/CUDA内核，通过大量实验迭代改进代码。
- 在H100上对RMSNorm、softmax等操作提速最高达5.29倍，超越PyTorch原生和编译器优化。
- 支持主流Transformer算子，适合需要极致性能的深度学习系统开发者。
编写高性能GPU内核是机器学习系统工程中最费力的任务之一。我们提出AutoKernel，一个开源框架，通过自主智能体循环实现对任意PyTorch模型的GPU内核自动优化。给定模型后，AutoKernel通过性能分析识别计算瓶颈，按阿姆达尔定律影响排序，并在无须人工干预的情况下，通过数百次实验迭代优化Triton或CUDA C++内核实现。五阶段正确性验证机制（包括烟雾测试、形状遍历、数值稳定性、确定性验证及边界情况覆盖）确保每个候选内核在记录加速前均经过严格验证。系统包含超过9000行Python代码，18个跨两个后端的初始内核实现，六层优化手册，以及与KernelBench基准套件的集成。AutoKernel覆盖了现代Transformer架构中九类主要运算。在NVIDIA H100上，我们的Triton内核在多数测试配置下优于PyTorch eager模式和torch.compile（max-autotune）：RMSNorm提速5.29倍，softmax提速2.82倍，交叉熵提速2.21倍；相比torch.compile分别提升2.83倍、3.44倍和2.94倍。在社区部署中，一个由AutoKernel优化的内核在vectorsum_v2 B200排行榜上获得第一名。完整系统已在https://github.com/RightNow-AI/autokernel发布。
原文摘要 · Abstract (English)
Writing high-performance GPU kernels is among the most labor-intensive tasks in machine learning systems engineering. We present AutoKernel, an open-source framework that applies an autonomous agent loop to GPU kernel optimization for arbitrary PyTorch models. Given a model, AutoKernel profiles it to identify computational bottlenecks, ranks them by Amdahl's law impact, and iteratively refines Triton or CUDA C++ kernel implementations through hundreds of experiments without human intervention. A five-stage correctness harness covering smoke tests, shape sweeps, numerical stability, determinism verification, and edge-case coverage ensures every candidate kernel is validated before any speedup is recorded. The system comprises over 9,000 lines of Python, 18 starter kernel implementations across two backends, a six-tier optimization playbook, and integration with the KernelBench benchmark suite. AutoKernel covers nine kernel types spanning the dominant operations in modern transformer architectures. On an NVIDIA H100, our Triton kernels outperform both PyTorch eager and torch.compile (max-autotune) on the majority of tested configurations: 5.29x over eager on RMSNorm, 2.82x on softmax, and 2.21x on cross-entropy, while beating torch.compile by 2.83x, 3.44x, and 2.94x respectively. In community deployment, an AutoKernel-optimized kernel achieved first place on the vectorsum_v2 B200 leaderboard. The full system is available at https://github.com/RightNow-AI/autokernel.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。