AI CUDA Engineer：エージェントによるCUDAカーネルの発見、最適化、生成



Sakana AIでは、より強力なAIシステムを目指すうえで取るべき道は、AIを使ってAIの開発を自動化することだと考えています。

昨年、私たちは、AI基盤モデルを従来より劇的に小さいコストで自動的に作成するAIシステムを開発し、LLMがLLMをトレーニングするためのより効率的な方法を作成できることを示しました。また、私たちはAI研究プロセス全体を完全に自動化する包括的なエージェントフレームワークである「AIサイエンティスト」を提案しました。そして、次にこの疑問が生じました。もしAIを使ってAIの研究ができるのであれば、AIを使ってAIをより速く動かすこともできるのでは？


The AI CUDA Engineer

現代のAIシステムはGPUなどのハードウェアアクセラレータによる並列処理に大きく依存しています。その点では人間の脳と同様です。しかし、エネルギーの制約下で効率的に動作するように（生物学的および文化的に）進化してきた人間の脳とは異なり、最近のAI基盤モデルは、大規模化の方向に進んできました。推論時間およびエネルギー需要が増大し、AI技術の開発と展開のコストは指数関数的に増加しています。

私たちは、現代のAIシステムは人間の脳と同じくらい効率的であるべきであり、その効率を達成するための最善の方法は、AIを使ってAIをより効率的にすることであると信じています。今回発表する AI CUDA Engineer はその一歩となるものです。これは以前のAIサイエンティストの研究に触発されたものであり、CUDAカーネルの発見と最適化のための初めての完全自動の包括的なエージェントフレームワークです。

CUDAは、NVIDIAのGPUハードウェアの並立計算を制御するための基層レイヤーのソフトウェアであり、今日のGPUプログラミングにおいて最も普及したプラットフォームです。その核となるのが、GPU上で並列実行される「CUDAカーネル」と呼ばれる関数です。CUDAカーネルのレベルで直接命令を記述することで、AIアルゴリズムのパフォーマンスを大幅に向上させることができます。

AI CUDA Engineerは、先端的なLLMを活用して、標準的なPyTorchコードを高度に最適化されたCUDAカーネルへ自動的に変換するためのエージェントフレームワークです。進化的な最適化やその他の進化計算の手法（それまでに見つけたカーネルのデータベースからの交叉演算など）を活用して、PyTorchモジュールを高度に最適化されたCUDAカーネルに自動的に変換します。しかも、最適化されたCUDAカーネルは多くの場合、大幅に高速化を実現します。

私たちは、この技術がLLMをはじめとする基盤モデルの開発（training）と利用（inference）の際の動作速度を加速させ、最終的にはNVIDIAのハードウェア上でのAIモデルの動きが格段に速くなると考えています。

The AI CUDA Engineerは、一般的なPyTorchによるものよりも10〜100倍高速なCUDAカーネルを生成できます。AI CUDA Engineerは高度な最適化を実現し、本番環境で一般的に使用されている既存のCUDAカーネルよりもはるかに（最大500％）高速なCUDAカーネルを生成することもできます。


AI CUDA Engineerエージェントフレームワークの概要


ステージ1と2（変換と翻訳）：AI CUDA Engineerは、まずPyTorchコードを正しく動作するCUDAカーネルに変換します。この時点では高速化は狙っていないものの、多少の実行時間の改善が観察されます。

ステージ3（進化的最適化）：生物学的進化にヒントを得た進化的最適化（「適者生存」）の手法にて、もっともよいCUDAカーネルのみが生成されるようにします。さらに、複数の最適化されたカーネルを組み合わせて新しいカーネルをつくる、カーネル交叉プロンプティングを行います。

ステージ4（イノベーションの保存）：文化的な進化が何千年もの文明を通して祖先からのノウハウで私たちの人間の知性を形作ったように、AI CUDA Engineerも過去のイノベーションと発見で得たものを活用すべく、既知の高性能CUDAカーネルをイノベーションアーカイブに保存します。これを活用してパフォーマンスの向上を実現します。


AI CUDA Engineerが発見したカーネルによる実行時間の高速化

AI CUDA Engineerは、一般的な機械学習に関して、PyTorchネイティブまたはコンパイルされたPyTorchコードよりも10〜100倍高速にするCUDAカーネルを安定的に発見しました。このアプローチはまた、機械学習アーキテクチャ全体を最適化されたCUDAカーネルに変換することもできます。下表は、完全に自律的に行われた重要な高速化の発見の例です。

AI CUDA Engineerは、行列積算のような基本的な演算から、一般的なディープラーニングの演算まで、より効率的なCUDAカーネルを発見します。そして、本記事の執筆時点では、私たちが発見したCUDAカーネルのパフォーマンスはKernelBenchで最高水準の成果を達成しました。

私たちは、これは始まりにすぎず、今後AIの大いなる最適化が進んでいくと確信しています。ぜひ私たちの論文「The AI CUDA Engineer: Agentic CUDA Kernel Discovery and Optimization AI CUDA Engineer：エージェントによるCUDAカーネルの発見と最適化）」をご覧ください。

論文の内容：

AI CUDA Engineerが発見した注目のカーネル

この新しいLLM駆動による進化論的カーネル最適化手順を活用することで、多様な用途で安定して高速化を実現できます。具体的には、検討した229個のタスクのうち81%でPyTorchネイティブの実行時間を短縮しています。さらに、発見されたすべてのCUDAカーネルの20%は、PyTorchの実装よりも少なくとも2倍高速です。


AI CUDA Engineerは、PyTorchの実装を上回るCUDAカーネルを安定して発見。


以下に、カーネルの一部を示します。AI CUDA Engineerはこのように幅広い用途に用いることができます。正規化手法、損失関数、特殊な行列積算、さらにはニューラルネットワークアーキテクチャ全体の演算も含まれます。

AI CUDA Engineerによって生成された高度に最適化されたCUDAカーネルの例。 それぞれのCUDAカーネルのサムネイルをクリックすると、実行時間の高速化などの詳細・分析がインタラクティブなウェブサイトでご覧になれます。.


AI CUDA Engineerアーカイブ：検証済みCUDAカーネルのデータセット

AI CUDA Engineerアーカイブのテキスト埋め込みを可視化したもの。発見されたカーネルがタスク（例：MatMul、プーリング、畳み込み）と実装戦略（アンローリング、融合、ベクトル化）にグループ化されています。アーカイブは公開されており、LLMのファインチューニングに使用できます。


この論文とともに、AI CUDA Engineerアーカイブを公開します。これは、AI CUDA Engineerによって生成された30,000個以上のCUDAカーネルで構成されるデータセットです。CC-By-4.0ライセンスで公開され、HuggingFace経由でアクセスできます。データセットには、torchのリファレンス実装、torch、NCU、Clang-tidyのプロファイリングデータ、タスクごとの複数のカーネル、エラーメッセージ、およびtorch Nativeとコンパイル済みランタイムに対する高速化スコアが含まれています。


約30,000個のカーネルと17,000個以上の正しい検証済み実装で構成されるAI CUDA Engineerアーカイブの概要。約50%のカーネルがtorch Nativeの実行時間を改善しています。


このデータセットにより、オープンソースモデルの事後トレーニングにおいて、より高性能なCUDAモジュールを実行できるでしょう。用途には、オフライン強化学習、 Preference Optimization、標準的な教師ありファインチューニングが含まれます。


AI CUDA Engineerアーカイブを探索する

さらに、17,000個の検証済みカーネルとそのプロファイル（torch、NCU、Clang-Tidyデータを含む）をインタラクティブに探索できるWebサイトも公開しました。Webページにはこちらからアクセスできます。

このWebサイトでは、230個のタスクにわたるさまざまな高性能カーネルを探索できます。個別の実験やLLMをまたいでカーネルを比較できるカスタムリーダーボードが付属しています。


AI CUDA Engineerによって発見されたカーネルのインタラクティブリーダーボード


さらに、カーネルを可視化したり、関連するカーネルを取得したり、コードをダウンロードして実装と高速化を検証したり、取得したプロファイリングデータを表示することができます。最適化実験の中身を詳細に見ることもできます。


プロファイリングデータ、評価スクリプトのダウンロード、関連カーネル、発見実験の詳細など、カーネルの詳細。


AI CUDA Engineerの今後の展望

AI革命は始まったばかりです。今日のLLMは、コンピュータでいえば20世紀半ばの「メインフレームコンピュータ」の段階にあると私たちは見ています。AIの発展はまだまだこれからであり、市場競争の激しさとグローバルに行われる（とりわけリソースの制約下でなされている開発による）イノベーションによって、AIが今後10万倍、100万倍と効率化されることは間違いありません。

現在、私たちのAIシステムは膨大なリソースを消費しています。効率や消費エネルギーを考慮しない技術の拡大は持続可能ではありません。AIシステムが人間の知能と同じくらい（あるいはそれ以上に）効率的になれない理由はありません。それを実現する最善の方法は、AIを使ってAIをより効率的にすることだと、私たちは信じています。

これがSakana AIが目指す方向性であり、本プロジェクトは100万倍高速化したAIへの重要なステップです。初期の巨大で不格好なメインフレームコンピュータから現代のコンピューティングへと進化したように、数年後に私たちのAIの使い方は様変わりし、今日の不格好で非効率なLLMとは大きく異なるものになるでしょう。


Sakana AI

日本でのAIの未来を、Sakana AIと一緒に切り拓いてくださる方を募集しています。当社の募集要項をご覧ください。