AutoKernel：反復的なエージェント駆動型探索による自律的なGPUカーネル最適化
AIによるGPUカーネル自律機械学習システムにおいて、高性能なGPUカーネルの作成は非常に労力のかかる課題です。
本研究では、この課題を解決するオープンソースフレームワーク「AutoKernel」を提案します。
AutoKernelは、自律的なエージェントループを採用し、モデルの計算ボトルネックを特定した後、人間の介入なしにTritonやCUDA C++カーネルを反復的に洗練させます。
候補となるカーネルは、数値安定性や決定性など、5段階の厳格な検証ハネスを通過する必要があります。
NVIDIA H100上での検証では、AutoKernelが従来の手法やPyTorch eagerを大幅に凌駕する高速化を実現しました。
これにより、複雑なGPUカーネルの最適化プロセスを自動化し、効率的な高性能AIシステムの構築を支援します。
原文の冒頭を表示（英語・3段落のみ）
View PDF
HTML (experimental)
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 this http URL (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 this http URL 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 this https URL.
※ 著作権に配慮し、引用は冒頭3段落までです。続きは元記事をご覧ください。