ハクソク

世界を動かす技術を、日本語で。

予想外に速いAI生成カーネル、まだ公開するつもりではなかった

360日前原文(crfm.stanford.edu)

概要

  • AI生成による CUDA-Cカーネル がPyTorchの最適化カーネルを上回る性能を実現
  • CUTLASSやTriton などのライブラリやDSLを使用せず、純粋なCUDA-Cで開発
  • KernelBench ベンチマークで多様なMLオペレータに対して高い最適化効果を確認
  • 並列的なアイデア探索と自然言語推論による新しいカーネル最適化手法の紹介
  • 今後の課題と展望として、FP16やFlash Attentionなど更なる性能向上を目指す

驚異的な高速AI生成カーネルの早期公開

  • PyTorch標準カーネル と同等またはそれ以上の性能を示すAI生成CUDA-Cカーネルの実現
  • 代表的な結果
    • Matmul (FP32): torch.matmul比101.3% （4096x4096行列）
    • Conv2D: torch.nn.Conv2D比179.9% （(100, 3, 224, 224)入力）
    • Softmax: torch.softmax比111.8% （(4096, 65536)入力）
    • LayerNorm: torch.nn.LayerNorm比484.4% （(16, 64, 256, 256)入力）
    • Conv2D + ReLU + MaxPool: torch参照比290.1%
  • Nvidia L40S GPU 上でベンチマークを実施、%性能＝参照時間÷生成カーネル時間で算出

研究の背景と驚き

  • 合成データ生成 を通じたカーネル生成モデルの訓練が当初の目的
  • 予想外にテスト時の合成データ生成で 高性能なカーネル が生成
  • 高度な最適化 やハードウェア機能を活用し、人間専門家の手書きカーネルに匹敵
  • 本記事は新規手法の提案ではなく、 シンプルな設計 による驚きの成果の共有
  • 5つの最適化カーネル と最適化過程、今後の展望を紹介

手法概要

  • KernelBench （2024年12月リリース）のタスクセットアップを利用
    • torchコードをもとにLLMがカスタムカーネルを生成し高速化を目指す
    • 参照はデフォルトで FP32、許容誤差1e-02で低精度も許容
    • 問題ごとにサイズを固定し、 サイズ特化最適化 を評価
    • 参照コードと生成コードをランダム入力で 数値一致検証
  • 一般的な最適化ループ
    • モデルがカーネルを逐次修正し、 正しさと性能を検証 しながら改良
    • アイデアの多様性不足や 局所最適化 への陥りやすさが課題

新しい最適化アプローチ

  • 自然言語による最適化アイデア生成
    • 各ステップで直接カーネルを生成せず、自然言語で 最適化アイデア を出力
    • そのアイデアから新たなコードバリアントを実装
  • 各ステップでの分岐探索
    • 1ステップごとに複数の実装を生成し、 最良カーネルを次ラウンドのシード に
    • 大規模並列探索 により多様な方向性を同時に模索
  • 結果： 構造化された探索的最適化ループ を実現し、従来よりも新規性の高い最適化案を発見

最適化戦略の主なカテゴリ

  • メモリアクセス最適化 ：データ移動効率向上、帯域幅最大化
  • 非同期処理・レイテンシ隠蔽 ：計算とメモリアクセスの重畳
  • データ型・精度最適化 ：FP16/BF16利用による帯域幅削減
  • 計算・命令最適化 ：算術計算効率化、特殊命令活用
  • 並列性・オキュパンシー向上 ：SM上のアクティブワープ数最大化
  • 制御フロー・ループ最適化 ：ループや分岐のオーバーヘッド削減

Conv2Dカーネル最適化の例

  • torch参照基準：1.41ms
  • ラウンドごとの最適化案と性能推移
    • CUDAカーネルへの単純置換、 7.02ms（20.1%）
    • Read-onlyキャッシュ利用、 7.54ms（18.8%）
    • FP16 Tensor-Core GEMM変換、 3.46ms（41.0%）
    • cp.asyncによるダブルバッファリング、 3.67ms（38.7%）
    • 既存の良好なGEMMカーネルをシードに、 3.46ms（41.0%）
    • k_idx分解インデックスの共有メモリ事前計算、 1.91ms（74.9%）
    • N次元GEMMインデックスの共有メモリキャッシュ、 1.37ms（103.6%）
    • ワープ単位の共有メモリバッファによる出力並列化、 1.38ms（102.9%）
    • 入力座標の共有メモリキャッシュ、 1.37ms（103.6%）
    • Bフラグメントのソフトウェアパイプライン、 1.36ms（105.1%）
    • N次元GEMM分解の出力アドレス計算再利用、 1.07ms（133.6%）
    • half WMMAのhi/lo分解省略、 1.21ms（117.4%）
    • Kループのダブルバッファリング、 1.01ms（141.2%）
    • half2によるベクトル化共有メモリ書き込み、 0.795ms（179.9%）
  • 最終コードは高度なCUDA技術を駆使し、手書きが困難なレベル

考察と今後の展望

  • 強力な 推論力 と 並列探索 の組み合わせが最適化を大きく前進
  • AlphaEvolve や Gemini 2.5 Pro Deep Think の示す通り、賢い探索と分岐戦略で大規模再学習不要な場合も
  • 本手法は 合成データ生成 の質向上にも寄与し、今後のモデル訓練にも有用
  • 依然として課題も存在
    • FP16 Matmul: torch.matmul比52%
    • FP16 Flash Attention: torch.nn.functional.scaled_dot_product_attention比9%
  • FP32 は近年のハードウェアで最適化が遅れがちで、FP16/BF16対応の最適化が今後の焦点
  • 探索予算が限られる中でも着実に性能向上 を実現し、AIによるカーネル最適化の将来性に期待

まとめ

  • AI生成カーネルが既存の専門家最適化カーネルに匹敵、もしくは凌駕
  • 自然言語推論＋並列分岐探索 による新しい最適化アプローチ
  • 今後はFP16/Flash Attentionなどさらに難易度の高いカーネル最適化に挑戦予定
  • AIによる自動カーネル生成と最適化の実用化 に向けた大きな一歩

Hackerたちの意見

すごくクールだね。o3とGemini 2.5 Proを使ったみたいだけど、どっちがより良いカーネルを生み出したかは残念ながら言及されてないね。

とても興味深い結果だね。彼らはこのブログ記事を書いたのは、自分たちの発見を共有したいという純粋な興奮からだと思うし、もしかしたら誰かに冷や水をかけてもらうためかも、ハハ。これが実際の「自己改善」の伝説の道かどうかは分からないけど、こういう結果が出るのがその道で期待されることだよね。

└

これが実際の「自己改善」の伝説の道かどうかは分からない それは疑わしいね。これは非常に明確に定義された評価関数でしか機能しないから。

僕の考えは、この記事やGoogleのAlphaEvolve [1]、最近のo3がLinuxカーネルでゼロデイを見つけたという発表 [2] から得たもので、特にGemini Pro 2.5とo3は、新しい能力のレベルに達していて、他のモデルでうまくいかなかったアイデアが突然うまくいくようになったってことだね。 [1] https://deepmind.google/discover/blog/alphaevolve-a-gemini-p... [2] https://sean.heelan.io/2025/05/22/how-i-used-o3-to-find-cve-...

└

Gemini Pro 2.5は、人間の言語翻訳以外で生産的に使える初めてのAIなんだけど、ギリギリその境界を越えた感じ。成功率が20%以下になることもあるし。3.0が出たら、ちょっと怖くなりそうだね。

└

え、何言ってるの？これ、Linuxカーネルとは全く関係ないよ。GPUプログラミングの意味での「カーネル」だから。もしかして、このコメント全部妄想だったの？

└

私の意見としては、突然動き出したってわけじゃないと思う。むしろ、彼らが人間よりもずっと早く反復してテストできるようになったって感じかな。さらに、すぐに使える情報をたくさん引き出せるから、その結果、情報の組み合わせや進歩、賢く適用された力技が特定のアプリケーションで成功してるみたい。

面白いのは、AI生成コードによって大規模な融合カーネルがどのように扱われるかだね。例えば、gemm + relu + gemm + 何らかのノルムを含むかもしれないけど、これは1. チューナーでスイープするのも、2. 人間が手書きするのもめちゃくちゃ面倒だよね。

└

えっと、AIの「カーネル」って何？OSのカーネルとは全然違うみたいだけど。

「FP32は現代のMLワークロードではあまり一般的ではなく、最近のハードウェアではFP16やBF16に比べて最適化されていないことが多いので、FP32カーネルでPyTorchに対する性能向上を達成するのが容易な理由の一部かもしれません。」人々は何年もFP32バージョンのカーネルを最適化することに時間を費やしていないからね。もし開発者の努力が注がれ、実際に使われているカーネルを改善できれば、もっと面白くなると思う。

└

FP16/BF16カーネルからFP32に転用できる既知の改善を使っているのかな？

└

「人々は何年もこのカーネルのfp32バージョンを最適化する時間を使っていない。」 へぇ、つまりAIが事前に解決策がない分野で新しいアルゴリズムを作ったってこと？すごいね！

時々、LLMをハイブマインドみたいに考えることがあるんだ。たくさんの人間の思考プロセスで訓練されているからね。それが、情報やコンテキストが重みの中に圧縮されているおかげで、こういうことができる理由だと思う。

└

市場自体も、なんかハイブマインドのメタファーみたいだね。考えてみる価値ありそう。

「参照コードはデフォルトのFP32で、許容誤差の閾値は1e-02」って、すごい許容範囲だね。これでfp16の操作を使って「fp32」カーネルを置き換えられるんだ。

└

うん、ここでの根本的なタスク（1ステップの推論）は、できるだけ多くのfp32操作をこのカーネルでfp16に置き換えることだったんじゃないかな。どれくらい難しいポートなのかは正確には分からないけど、直感的にはちょっとインパクトが薄いかも。もしかしたらこの直感は間違ってるかもしれないけど、もしそうならその点を明確にしてほしいな！

└

これは結果が無意味だってことだね。相対誤差を全くチェックしたのかな？float32の操作をfloat16に置き換えるのも無意味だよ。これをやると、float32の実際の精度の利点が失われちゃうから、アルゴリズムのそのバージョンを使う一番の理由がなくなっちゃうんだ。

これってOpenCLやROCmのカーネルを作るのに使えるのかな？

一番面白いのは（特に一部のケースで400%のスピードアップがあった後）その手法だね。操作のヒルクライムをする代わりに、イテレーションの間に言語的な推論ステップを挟んで、探索の多様性を促してるみたい。これがうまくいったみたいで、すごく興味深い。

└

おお、すごい！島やマップエリートの使い方を探してたのに、これを見逃してた… もっと平凡な模倣進化だと思ってた。

注意書き：これが私の得意分野だったけど、5年もこの手のことをやってなかったから、かなり鈍ってる。とはいえ、AI生成のカーネルの例をざっと見た感じ、特に新しいものは見当たらなかった。nVidiaで働いてた時に、正直言って驚くような技術をいくつか見たから、これが標準のPyTorchカーネルよりも速い理由が気になる。おそらく、cuDNNに重い処理を全部任せてるだけだと思うんだけど。私の推測だけど、特定のパラメータセットに対して生成した最速のAIカーネルと、cuDNNがそのシナリオに対して選んだカーネルを比較してるんじゃないかな。もしかしたら、cuDNNの中でどのカーネルを実行するかを選ぶサブシステムが最適でない候補を選んじゃったのかも。研究者たちはこの問題を完全に無視して、cuDNNがどんなシナリオでも常に最良のカーネルを選べると思い込んでるけど、それは現実的じゃない。何か他のことが起こってるかもしれないけど、「この重く最適化されたプロプライエタリライブラリを超えた」みたいな話は、いつもこの重要なポイントを見落としてる気がする。これを読んでるかもしれないnVidiaの内部の人たちによろしく。君たちは今まで会った中で一番頭のいい人たちだよ。