AutoKernel：反復的なエージェント駆動型探索による自律的なGPUカーネル最適化
AIによるGPUカーネル自律機械学習システムにおいて、高性能なGPUカーネルの作成は非常に労力のかかる課題です。
本研究では、この課題を解決するオープンソースフレームワーク「AutoKernel」を提案します。
AutoKernelは、自律的なエージェントループを採用し、モデルの計算ボトルネックを特定した後、人間の介入なしにTritonやCUDA C++カーネルを反復的に洗練させます。
候補となるカーネルは、数値安定性や決定性など、5段階の厳格な検証ハネスを通過する必要があります。
NVIDIA H100上での検証では、AutoKernelが従来の手法やPyTorch eagerを大幅に凌駕する高速化を実現しました。
これにより、複雑なGPUカーネルの最適化プロセスを自動化し、効率的な高性能AIシステムの構築を支援します。
機械学習の分野で、高性能なGPUカーネルを自力で開発することは非常に手間のかかる作業とされています。この課題に対し、研究チームは「AutoKernel」というオープンソースのフレームワークを発表しました。これは、自律的なエージェントループを用いて、任意のPyTorchモデルのGPUカーネルを自動で最適化するシステムです。
AIによるカーネル自動最適化の仕組み
AutoKernelは、まず与えられたモデルをプロファイリングし、計算上のボトルネックを特定します。次に、Amdahlの法則といった指標に基づき、それらのボトルネックの重要度をランク付けします。その後、人間が介入することなく、TritonやCUDA C++といったカーネル実装を数百回の実験を通じて反復的に洗練させていく仕組みです。この自動化により、開発工数を大幅に削減できると説明されています。
徹底した検証プロセスと性能向上
このシステムでは、単に速くなるだけでなく、正確性が保証されることが重要視されています。Smoke tests（スモークテスト）や数値安定性、決定論的検証、エッジケースの網羅といった5段階の厳格な検証プロセスが組み込まれています。これにより、速度向上が記録されるすべての候補カーネルが、事前に完全に検証される体制が整っています。実際にNVIDIA H100上でテストされた結果、RMSNormやsoftmaxなど複数のカーネルにおいて、既存のPyTorchのEagerモードや他の最適化手法を大きく上回る性能を示しました。
実証された高い実用性と応用範囲
AutoKernelは、現代のTransformerアーキテクチャで主流となっている9種類のカーネルタイプをカバーしています。また、コミュニティでの実運用においても、AutoKernelによって最適化されたカーネルが、B200のリーダーボードで1位を獲得した事例も報告されています。このフレームワークは、Pythonで9,000行以上のコードや、2つのバックエンドにおける18種類のスターターカーネル実装を含んでおり、高い実用性が示されています。
結論
AutoKernelは、機械学習の基盤となるGPUカーネル開発のボトルネックを、AIによる自律的な最適化で解決しようとする画期的な試みです。この技術が広く普及すれば、より高速で効率的な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段落までです。続きは元記事をご覧ください。