Kernel Forge：LLMによるCUDAカーネルの生成と最適化のためのエージェントハーネス
Kernel Forgeは、未修正のPyTorchモデルをそのまま受け入れ、モンテカルロ木探索（MCTS）を用いて複数の最適化経路を探索し、GUI監視機能を備えたオープンソースのエンドツーエンドエージェントフレームワークです。NVIDIA DGX Spark上で4つのPyTorchモデルを評価し、カーネルあたりわずか50回の最適化反復で14のカーネルがPyTorch eagerモードを上回り、最大2.83倍の高速化を達成しました。
機械学習モデルは日常的なソフトウェアにますます組み込まれており、その実行時間の大半は行列乗算、畳み込み、正規化などの少数の計算カーネルに費やされています。これらのカーネルを最適化することはレイテンシとコストを削減する最も直接的な方法の一つですが、従来は専門技術者による手書きの低レベルGPUコードが必要でした。大規模言語モデル（LLM）に基づくエージェントシステムは、はるかに少ない人間の労力でカーネルを生成・最適化できるようになりましたが、既存のツールは主にランダム生成テンソルと孤立カーネルで評価され、開発者が手動で再統合する必要があるスタンドアロンCUDAコードを出力し、主にLLM PyTorchモデルのみを対象とし、結果の検査とデバッグのサポートが限られています。
これらの課題に対処するため、研究チームはKernel Forgeを提案しました。これは、未修正のPyTorchモデルをそのまま受け入れるオープンソースのエンドツーエンドエージェントフレームワークです。Kernel Forgeはビジョン、拡散、LLMワークロードをサポートし、単一の線形リファインメントチェーンではなくモンテカルロ木探索（MCTS）を使用して複数の最適化経路を探索し、進捗監視、候補カーネルの検査、障害デバッグのためのグラフィカルユーザーインターフェースを備えています。
研究チームは、NVIDIA DGX Spark（GB10 GPU搭載）上で、ビジョン、拡散、LLMワークロードをカバーする4つのPyTorchモデルでKernel Forgeを評価しました。カーネルあたりわずか50回の最適化反復で、14のカーネルがPyTorch eagerモードを上回る性能を達成しました。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倍の高速化を記録しました。これらの結果は、Kernel Forgeがカーネル最適化のハードルを下げ、AIデプロイメントに顕著な性能向上をもたらすことを示しています。コードはオープンソースとして公開されています。