arXiv (AI)AI
TITLE_JA: Kernel Forge:LLMベースのCUDAカーネル生成・最適化のためのエージェントハーネス
Kernel Forge: An Agent Harness for LLM-based Generation and Optimization of CUDA Kernels
この記事についてAIに質問する →
日本語要約青い用語にマウスを合わせると解説が表示されます
機械学習モデルは現代のソフトウェアに広く組み込まれており、その実行時間の大部分は行列乗算、畳み込み、正規化といった少数の計算カーネルで消費されている。これらのカーネルの最適化は遅延とコストを削減する最も直接的な手段だが、従来はGPUの低レベルコードを手書きする専門家エンジニアが必要とされてきた。大規模言語モデル(LLM)に基づくエージェントシステムは現在、人手をはるかに減らしてカーネルを生成・最適化できるようになっているが、既存ツールのほとんどはランダムに生成されたテンソルと孤立したカーネルで評価され、開発者が手動で再統合する必要があるスタンドアロンのCUDAコードを出力し、主にLLM PyTorchモデルのみを対象とし、結果の検査とデバッグのサポートが限定的である。
本研究では、Kernel Forgeというオープンソースのエンドツーエンドエージェントハーネスを提案する。任意の未修正PyTorchモデルをそのまま受け入れ、ビジョン、拡散、LLMワークロードをサポートし、単一の線形改善チェーンではなくモンテカルロ木探索(MCTS)を使用して複数の最適化パスを探索する。さらに進行状況の監視、候補カーネルの検査、障害のデバッグのためのグラフィカルユーザーインターフェースを備えている。
NVIDIA DGX SparkのGB10 GPUを使用した4つのPyTorchモデル(ビジョン、拡散、LLMワークロード)での評価において、カーネルあたり50回の最適化イテレーションのみで、PyTorchイーガーモードを上回る性能に最適化された14個のカーネルを達成した。ResNet-50のadaptive_avgpool2dで1.52倍、Stable Diffusion 3.5 Mediumのgroup_normで1.70倍、Gemma 4 E2BのSoftmaxで2.83倍、Qwen 3.5 35B-A3BのSoftmaxで1.54倍の高速化を実現した。