Sakana AI、CUDAカーネルの自動最適化フレームワークと新ベンチマークを発表

PyTorch で機械学習モデルを組むのは手軽ですが、GPU の性能を限界まで引き出すには低レイヤーの CUDA カーネルを手作業で書く必要があります。

この作業には高度なハードウェア知識と膨大な試行錯誤が求められます。Sakana AI はこの課題を解決するため、CUDA カーネルの自動探索、検証、最適化を行うエージェントフレームワークを発表しました。

LLM を低レイヤーの GPU 最適化へ投入

最近の大規模言語モデルは思考時間を増やすことでソフトウェア開発タスクをこなせるようになっていますが、その大半は高水準なコードに限られていました。

さらに、既存のカーネル生成ベンチマークには抜け穴が多く、多様な条件下での汎化性能を正確に測れていませんでした。

そこで Sakana AI は、様々なシナリオでカーネルの性能と正しさを厳格に評価する新ベンチマーク robust-kbench を構築しました。

この評価基盤の上で動くエージェントは、まず最先端の LLM を使って PyTorch コードを等価な CUDA カーネルへと自動変換します。

進化的メタ生成と検証器でネイティブ実装を凌駕

変換されたコードは、CUDA エコシステム向けに設計された独自の進化的メタ生成プロセスによって、実行速度を繰り返し最適化されます。

高速だが誤った結果を返すカーネルを排除するため、システム内には LLM ベースの検証器が組み込まれ、正しさを判定して効率的にフィルタリングします。

これにより、実際のハードウェア上で不完全なコードを動かす無駄な検証コストを大幅に削減できます。

robust-kbench での評価において、この手法が生成した CUDA カーネルは順伝播と逆伝播の両方を含む実践的なタスクで PyTorch のネイティブ実装を上回る実行速度を達成しました。

さらに、複数の演算をまとめる演算子融合や多様なランタイム最適化戦略も自律的に適用できることが確認されています。

なぜ重要か

モデル開発者やパフォーマンスエンジニアにとって、CUDA のチューニングは最も専門性が高く時間のかかる作業の一つです。PyTorch から直接、公式実装を超える高速なカーネルが生成できるようになれば、開発から本番投入までのサイクルが劇的に短縮されます。

業界全体で見ても、ハードウェアの性能を極限まで引き出せるシステムエンジニアは世界的に不足しています。Sakana AI のアプローチは、LLM の探索能力と進化的アルゴリズムを組み合わせることで、極めて高度な職人芸を自動化できる可能性を示しました。

ハードウェア上で直接テストする前に LLM 検証器でバグを弾くアプローチは、リソースの無駄を省く上で非常に賢明だと考えられます。一方で、提示された成果が独自のベンチマーク環境外にある極端に複雑な処理や特殊なアーキテクチャに対しても同様に機能するかどうかは、今後の検証が必要になりそうです。