- AIが生成したCUDA-Cカーネルが、PyTorchの専門家最適化カーネルと同等かそれ以上の性能を示した
- 単一のLLM(大規模言語モデル)が自然言語による最適化アイデア生成と多様なコード分岐を繰り返し、既存手法より最適化の多様性と並列探索で優れた性能を達成
- 代表的なベンチマーク結果では、**Matmul(101%)、Conv2D(179.9%)、Softmax(111.8%)、LayerNorm(484.4%)、Conv2D+ReLU+MaxPool(290.1%)**などでPyTorchを大きく上回った
- 従来の「逐次的なカーネル改善」の限界を超えるため、自然言語推論 + 分岐構造を適用し、高速に高性能カーネルを生成
- FP16、Flash Attentionなど最新のMLワークロードでも進展しており、将来的にはAIが自律的により高速なカーネルを探索・改善するパラダイムが主流になる可能性を示している
TL;DR 主な成果
- Stanford CRFMの研究チームは、AIが生成した高性能CUDA-Cカーネルが既存のPyTorchの専門家設計カーネルと同等か、それ以上の速度を出す事例を確認した
- 当初は合成データを通じて、より優れたカーネル自動生成モデルを学習させることが目的だったが、合成データ生成だけでも驚くほど高速なカーネルが自動生成される現象を観察した
- これを受けて、手法、性能ベンチマーク、最適化戦略、今後の方向性などを早期に公開した
- ベンチマーク: Nvidia L40S GPU基準。性能(%): PyTorch基準の実行時間 ÷ 生成カーネルの実行時間
- Matmul (FP32): PyTorch比 101.3%(4096x4096行列)
- Conv2D: 179.9%(入力: 100, 3, 224, 224; AlexNet Conv1仕様)
- Softmax: 111.8%(4096x65536テンソル)
- LayerNorm: 484.4%(16, 64, 256, 256テンソル)
- Conv2D + ReLU + MaxPool: PyTorch reference比 290.1%、torch.compile()比 189.0%
研究の動機と方法
- 当初の目的はカーネル生成モデル学習用の合成データ生成だったが、実験の過程で自動生成されたカーネル自体が予想を上回る高性能を達成した
- KernelBench(Stanford公開ベンチマーク、2024年12月公開)を活用
- 与えられたtorchコードに対して、LLMが最速カーネルを新たに作成
- 入出力結果の数値的同値性で正確性を検証
- 従来方式: カーネルを段階的に少しずつ修正しながら改善する逐次修正ループ
- 欠点: アイデアの多様性不足、同じ最適化の反復、局所最適への収束
- 改善アイデア
- 最適化アイデアを自然言語で発想・記録したうえで、そのアイデアのコード実装分岐を複数同時生成
- 各ラウンドで複数の最適化試行を並列実行 → 最高性能カーネルを次ラウンドの種(シード)に設定
- これにより、限られた探索反復の中でも多様な最適化戦略の探究と並列探索が可能になる
最適化アイデアの例
- メモリアクセス最適化: global/shared/registerメモリ階層の効率改善
- 非同期処理とレイテンシ隠蔽: 計算とデータ移動をオーバーラップ
- データ型/精度最適化: FP16/BF16活用およびハードウェア特化演算
- 計算と命令最適化: 演算量、命令数、ハードウェア固有命令の最適活用
- 並列性とオキュパンシ: SM(Streaming Multiprocessors)の活用最大化
- ループ/分岐/インデクシング最適化: ループ最小化、不要なインデックス演算の除去
カーネル最適化の実際の過程(Conv2Dの例)
- ラウンドごとの最適化アイデアと性能改善の流れ
- 初期(0ラウンド): 単純なCUDA移植(PyTorch比20%速度)
- 次ラウンド: → 読み取りキャッシュ活用 → FP16 Tensor Core演算(GEMM変換) → 二重バッファリング/パイプライン → インデックス事前計算/共有メモリ → ベクトル化 → Kループ同時バッファリング など、高度なGPUアーキテクチャ活用
- 最終(13ラウンド): PyTorch比179.9%性能を確保
- 実際のコードでは高度なCUDAプログラミング技法を多数活用
- 各ラウンドで新しいアイデアを自然言語で生成し、複数実装を並列試行して最適コードを選択
Takeaways と示唆
- AIベースのカーネル生成は、強力な探索ループと自然言語ベースの多様な最適化アイデアの組み合わせによって人間の専門家レベルを超えうる
- 一部の最新演算子(FP16 matmul、Flash Attention)は現時点ではPyTorch比でまだ性能が低い
- FP32計算は最近のハードウェアではFP16/BF16に比べて相対的に最適化が不十分であり、この領域では性能優位の可能性がある
- 限られた探索トークン(入出力合計700万トークン)の状況でも継続的な性能改善を確認
- AlphaEvolve、Gemini 2.5 Proなど最近の研究も、大規模分岐+検証ベース探索が再学習なしでも飛躍的な性能改善を可能にすることを示唆している
- 今後はこの方式で合成データと高性能カーネルを大量生成し、AIが自己改善するループ(自己改善型AI)へ発展していくとみられる
結論
- AIベースのカーネル自動生成・最適化はすでに専門家の手書きコード水準に達しており、近い将来にはモデル+分岐探索+検証の組み合わせによる自己改善型AIシステムが実現する可能性がある
- PyTorch・TensorFlowのようなフレームワークの性能限界を、AIが自律的に超えていく可能性が浮上している
付録: Conv2Dカーネル全コード
- 実際の関数とカーネル全体のソースコードを含む(詳細コードは省略)
- コード内の主な特徴:
- 共有メモリのベクトル化、K-dim階層的double-buffering、CUDA WMMA、warp-level output buffer など
- リアルタイム動的インデックス計算、bias処理、Kループ内での遅延データ同時ロード など
- 完全なサンプルコードと例示カーネルはGitHubリポジトリで確認可能
2件のコメント
一種のアンサンブル手法と言うべきでしょうか。すごいですね。
Hacker Newsの意見
この文章の著者たちの「AIエージェント」の捉え方はかなり興味深いと思う
多くの人はエージェントを人間の従業員のように考えがちだ
少数のエージェントを並列に設定し(多くの場合は1人だけ)、それぞれがループを回しながら一度に1つの仕事だけを処理させる
依然として、固定された人数で、一度に1つの仕事しかできず、タスク切り替えも遅い従業員の世界に留まっている
だがLLMはそうではない
実質的に無限のエージェントを、いつでも好きなだけ作り出せる
LLMリクエストを直列で処理しても並列で処理しても、コスト差はない
この事実を認識すると、各エージェントが必要に応じて自分を複数のサブエージェントへ分岐させるパターンが自然に思い浮かぶ
これこそが著者たちが試したやり方だ
エージェントを「タスク」や「ジョブ」と見なし、CeleryやSidekiqで学んだことを適用する視点のほうが適切だと感じる
「FP32は最新のMLワークロードではかなり使われなくなっており、最新ハードウェアではFP16やBF16に比べて最適化も不十分だ。そのため、FP32カーネルでPyTorchより性能を改善しやすかった理由の1つかもしれない」
開発者たちは何年にもわたり、fp32版カーネルの最適化にほとんど時間を投じてこなかった
本当に面白いのは、実際に集中的に開発され、人々が本当に使っているカーネル側の性能を上げられるときだと思う
NVIDIAがGPUについて十分に詳細な文書を提供していないことが、この好結果を部分的に説明していると思う
マイクロアーキテクチャが十分に文書化されたプロセッサなら、プログラマやコンパイラが決定的に最適なプログラムを書けるので、既知の解法を見つける段階を超えてML/AIを適用して大きな改善を出すのは簡単ではない
一方、NVIDIA GPUのように文書化が不十分な場合、最適なプログラムを見つけるには、過去に最適化されたプログラム例を参考にしながらランダム探索を行ったり、状況ごとの実ハードウェアの挙動をリバースエンジニアリング的に分析したりする必要があることが多い
こうした状況ではML/AIが可能性を示しうるし、よく作られたプログラム群をデータとして学習することで、人間のプログラマでも見つけにくい undocumented behavior を捉えられるかもしれない
fp16/bf16カーネルで既に知られていた改善点が、単にfp32へ移植された結果である可能性はないだろうか
「人々が何年もfp32カーネルにまったく手を入れていなかったという話?」
わあ、もしそうなら、AIが事前の解法がなかった領域で新しいアルゴリズムを生み出したことになるわけで、本当にすごいことだと思う
私の結論は、この記事と、GoogleのAlphaEvolve(こちら)、そして最近o3がLinuxカーネルでゼロデイを見つけたという話(こちら)を見る限り、
特にGemini Pro 2.5とo3が、従来のモデルでは無理だったアイデアを突然やってのける新たな能力レベルに到達したように感じた
既存ではできなかったことが突然できるようになったというより、
実質的に人間よりはるかに速い速度で反復とテストが可能になり、
大量の情報を即座に活用できるようになった結果、
膨大な情報、進歩、そして知的に適用されたブルートフォースの組み合わせが、一部の応用分野で成功を収める段階に至ったのだと見ている
あなたが挙げた事例と今回の結果には、実際かなり共通点があると思う
本文にも「テスト時の反復ループは、順番にコード修正の結果を確認する対話型チャットボットではなく、明確な最適化仮説に従って積極的に並列評価を行う構造化探索に近い」という記述がある
私の結論は、今やLLMの能力を土台に、
評価関数が明確であったり、類似パターンの解空間を大幅に絞り込めたりする方法を身につけたという点だ
あるモデルが別のモデルを追い越すとか、特定モデルにしか解けないとか……そうしたモデル間比較が論点なのではなく、
こうした形の応用が十分に現れてきているという現実のほうが、より意味があると感じる
Gemini Pro 2.5は人が使える最初のAIだと感じたが、厳密に言えば、ようやくその閾値を少し超えた程度だ
成功率が20%を下回ることもまだよくある
3.0が出たら……今度は少し怖くなるかもしれないと思う
ちょっと待って、どういう意味? これはLinuxカーネルとは関係なくて、GPUプログラミングでいう「カーネル」の話だよ
もしかして本文全体を誤解してない?
興味深くはあるが、もっと強い証拠はある? 一度きりの結果だけでは、説得力としては十分ではない
「基準コードはデフォルトFP32で、許容誤差は1e-02だ」
この程度の誤差なら、簡単にfp16演算で「fp32」カーネルを置き換えられる
もう一段考えると、今回の作業の本質は、そのカーネルで可能な限り多くのfp32演算をfp16に置き換えることだったのではないかという気もする
こうした移植作業が実際どれほど難しいかは確認が必要だが、
直感的にはそこまで印象的には感じない
もし私の考えが間違っているなら、著者たちがこの点を明確に扱ってほしい
そうなると結果は無意味になる
相対誤差をチェックしたのかもしれないが、
float32をfloat16に置き換えるのは意味がない
精度こそがfloat32を選ぶ唯一の理由と言っていいのに、
その特徴自体を失えば、バージョンごとの差別化もなくなる
このタイトルを見て、AIがOSカーネルを作り出したのかと勘違いして読んだのは、ひょっとして私だけ?
400%の高速化もすごいが、
何より興味深いのは、反復ごとに単純な演算最適化をするのではなく、言語推論の段階を強制的に入れて探索の多様性を引き出した方法論だ
このやり方が実際にうまく機能したのが非常に興味深い
ただの普通のミメティック進化だと思う
本当に興味深い結果だ
この人たちは興奮しすぎてブログで公開してしまったようにも見える
実際には出版前にフィードバックを得ようという意図もあったのかもしれない
「自己改善(self improvement)」への伝説的な道筋がこれなのかは分からないが、
こうした結果はまさにその道の途中で見られそうな事例だと感じる
こうした手法は、評価関数が極めて明確な場合にしか効果がない
経験共有: replicationを試みたところ、LayerNormカーネルは数値的に安定しておらず、検証不能だと見ている
平均0、標準偏差1だけでテストしていたため、numerical cancellationの現象がすぐには表れなかった
付け加えると、その後で数値的に安定した新バージョンが作られたことを確認した
この点は素晴らしいと思う
本当にすごい結果だ
o3とGemini 2.5 Proを使ったそうだが、
どちらがより良いカーネルを作ったのかには触れられておらず、その点は少し残念だ
AI生成コードが、fused kernelのような広い領域をどう扱うのか観察するのは興味深いポイントだ
たとえば gemm + relu + gemm + normalization のようなものがあるだろう
チューナーで総当たりしたり、人が1つずつ書いたりするには相当に骨の折れる作業だ
OSカーネルではないことだけは確かだ