LLMによるCUDAカーネル自動最適化システム Kernel Forge──PyTorchモデルで既存手法比最大2.83倍の高速化を実証

論文概要

機械学習モデルのランタイムの大半は、行列積や畳み込み、正規化などの少数の計算カーネルで消費されている。これら低レベルのGPUコードの手動最適化には専門エンジニアの多大な労力が必要であり、既存のLLMベースのエージェントツールは主に単体のランダムテンソルを対象にし、手動での再統合が必要となる課題を抱えていた。本論文では、未修正のPyTorchモデルを直接入力として受け付け、ビジョン・拡散モデル・LLMの各ワークロードを包括的にサポートするオープンソースのエンドツーエンドエージェント基盤「Kernel Forge」を提案する。本手法はMonte Carlo Tree Search（MCTS）を用いて複数の最適化パスを探索し、NVIDIA DGX Spark（GB10 GPU）を用いた実験において、わずか50回の最適化イテレーションでPyTorch eager modeを上回る高速化を達成した。

関連研究

従来、LLMを活用したCUDAコード生成や自動チューニングに関する研究はいくつか存在していた。しかし、それらの多くは単体の孤立したカーネルやランダムに生成されたテンソルを対象として評価が行われており、実際の複雑なディープラーニングモデル全体への組み込みを考慮していなかった。また、生成されたコードの統合を手動で行う必要性や、LLMベースのモデルに特化しすぎてビジョンや拡散モデルへの汎用性が欠けている点がボトルネックとなっていた。Kernel Forgeは、エンドツーエンドでのモデル統合と、MCTSによる多角的な探索パスの確保により、これらの制約を克服している。

新規性と貢献

Kernel Forgeの主要な貢献は以下の3点にある。

1. 未修正のPyTorchモデル（ビジョン、拡散モデル、LLM）をそのまま入力として受け付ける、エンドツーエンドのエージェント基盤の構築。
2. 単一の線形リファインメントチェーンではなく、MCTSを採用することで多様な最適化パスを網羅的に探索する仕組みの導入。
3. 進捗監視、候補カーネルの検査、障害デバッグを統合したグラフィカルユーザーインターフェース（GUI）の提供。

提案手法の詳細

Kernel Forgeは、モデルの取り込みフェーズから演算子カードの生成、そしてMCTSコントローラーによるCUDA候補の反復的な最適化に至る一連のパイプラインで構成されている。なぜMCTSを採用したかといえば、単一のコード修正の連鎖では局所最適解に陥りやすく、多様なカーネル構造のバリエーションを同時に評価・探索することが困難であるためである。MCTSによって複数方向の最適化アプローチを並行してシミュレートし、最も高いパフォーマンスを示すCUDAコード片を効率的に選択・洗練させることが可能となる。さらに、開発者がリアルタイムで挙動を把握しデバッグできるように、専用のGUIが統合されている。

評価・考察

NVIDIA DGX SparkのGB10 GPU環境において、ビジョン、拡散モデル、LLMにわたる4つのPyTorchモデルで検証を実施した。各カーネルあたりわずか50回の最適化イテレーションという制限下において、14のカーネルがPyTorch eager modeのベースラインを超える性能を記録した。具体的には、ResNet-50の adaptive_avgpool2d で $1.52 imes$、Stable Diffusion 3.5 Mediumの group_norm で $1.70 imes$、Gemma 4 E2Bの softmax で $2.83 imes$、Qwen 3.5 35B-A3Bの softmax で $1.54 imes$ の高速化を実証している。

応用例と今後の展望

本手法は、LLM推論サービスや画像生成インフラを展開するクラウド事業者およびAI開発企業において、推論レイテンシの削減とGPUコストの最適化に直接寄与する。特に数千億パラメータ規模のLLMや複雑なマルチモーダルモデルを運用する国内のAIインフラ事業者にとって、手動でのカーネルチューニングコストを削減する有力な選択肢となる。今後の課題としては、より多様なハードウェアアーキテクチャへの対応範囲の拡大や、最適化イテレーション数のさらなる削減によるオーバーヘッドの低減が挙げられる。

結論

Kernel Forgeは、LLMを活用したCUDAカーネルの生成と最適化において、未修正のPyTorchモデルを直接処理可能なエンドツーエンドのMCTSベースエージェント基盤を提示した。主要なモデル群における顕著な高速化実績は、低レベルGPUコードの自動最適化の実用性を大きく前進させるものである。

注釈

• CUDAカーネル: NVIDIA製GPU上で並列実行されるプログラムの基本単位。
• MCTS (Monte Carlo Tree Search): 状態空間木をランダムサンプリングに基づいて探索し、最善の決定を行うアルゴリズム。