🤖 AI Summary
Manually crafting high-performance GPU kernels is inefficient and struggles to meet the computational demands of modern machine learning. This work proposes the first end-to-end verifiable autonomous agent framework that automatically analyzes bottlenecks in PyTorch models—guided by Amdahl’s Law—and iteratively generates and optimizes Triton or CUDA C++ kernels without human intervention, employing a five-stage correctness verification mechanism. The approach comprehensively covers core operators in mainstream Transformer architectures and achieves substantial speedups over both PyTorch eager execution and torch.compile (max-autotune) on H100 GPUs—for instance, accelerating RMSNorm by 5.29× and softmax by 2.82×—while attaining top performance on the B200 vectorsum_v2 benchmark.
📝 Abstract
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.