Kernel Forge: An Agent Harness for LLM-based Generation and Optimization of CUDA Kernels
Kernel Forgeは、大規模言語モデルとモンテカルロ木探索を活用して、多様なPyTorchワークロードに対してCUDAカーネルを自動生成および最適化し、PyTorchのイーガーモードに対して大幅な高速化を実現すると同時に、検査とデバッグのためのグラフィカルなインターフェースを提供する、オープンソースのエンドツーエンドのエージェンティック・ハーネスです。
原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。 免責事項の全文を読む
コンピューターでビデオゲームを動かそうとしている場面を想像してみてください。時にはゲームがスムーズに動きますが、他の時にはカクついたり、ラグが発生したりします。これは、コンピューターが絵を描画したり物理演算を処理したりするために、毎秒何百万もの微細な数学的計算を行う必要があるからです。人工知能（AI）の世界では、これらの計算はさらに激しいものになります。AIに「思考」させるためには、カーネルと呼ばれる特別な高速命令に依存しています。カーネルを、シェフが野菜を刻む時に使う特定の、非常に効率的なレシピだと考えてください。もしシェフが切れ味の悪いナイフと遅い方法を使えば、ディナーは遅れてしまいます。もし鋭く専門的なナイフを使えば、料理は瞬時に完成します。
長年、これらの「レシピ」を書くことは、ほとんど理解できていない言語で小説を書こうとするようなものでした。コンピューターチップ（GPU）をより速く動作させるために、高度なスキルを持つ専門家が複雑な指示をハンドコーディングする必要がありました。しかし今、新しい種類の助っ人が現れました。それが**大規模言語モデル（LLM）**です。物語を書いたり数学の問題を解いたりできるスマートなチャットボットとして、皆さんも知っているかもしれません。科学者たちは今、これらのチャットボットに「レシピ」（カーネル）を書くことを教え、人間よりも速く、より良くこれらを作成できることを期待しています。しかし、ここには落とし穴があります。チャットボットが紙の上で良さそうなレシピを書いたとしても、それが実際のキッチンで実際に食事を早く調理できるとは限らないのです。
ここで、Kernel Forgeと呼ばれる新しいプロジェクトが登場します。ミシガン大学の研究者たちは、単にチャットボットにコードを書かせるだけでなく、フルタイムのコーチ、品質検査官、そしてプロジェクトマネージャーのすべてを兼ね備えたスマートなシステムを構築しました。彼らは、AIが単なる架空の、作られたテストケースに対してではなく、現実世界のAIモデルに対して、これらのレシピを最適化する仕事を実際に引き継げるかどうかを確認したいと考えました。
問題点：「孤立」という罠
ロボットに走る訓練をしている場面を想像してください。もし、完璧に平坦で空っぽのトラックでしかテストしなければ、そのロボットは世界クラスのスプリンターに見えるかもしれません。しかし、いざデコボコした混雑した街路に出されると、つまずいて転んでしまうかもしれません。これまでのGPUカーネル作成における試みで起きていたのは、まさにこのことでした。多くのツールは、ランダムで架空の数学問題という「空っぽのトラック」をAIに与えてテストしていました。彼らはコードを生成し、それを孤立した状態でテストして、「見て、こんなに速いよ！」と言っていたのです。
しかし、現実世界において、AIモデルは忙しい都市のようなものです。「レシピ」（カーネル）は、特定の形状のデータ、特定のメモリ量と連携しなければならず、また、すぐ隣にある他のレシピとも上手く調和しなければなりません。ランダムなテストでは速かったレシピが、画像生成やチャットを行うような実際のAIモデルの中で使用されると、遅くなったり、あるいは壊れてしまったりすることがあります。従来のツールは、これらの新しいAI作成レシピを複雑な機械の中にどうやって適合させるかを、人間の開発者に判断させるままにしていました。これは非常に手間がかかり、ミスが起きやすい作業でした。
解決策：Kernel Forge
研究者たちは、AIの書き手と現実世界との間の架け橋として機能するオープンソースツール、Kernel Forgeを作成しました。Kernel Forgeは、AIに推測させるのではなく、まずコンピューター上で動いている実際のAIモデル（顔認識や物語作成を行うものなど）を観察します。そして、どの「レシピ」が使われているのか、データの大きさはどれくらいか、実行にどれくらいの時間がかかるのかといった詳細なメモを取ります。
この現実世界のデータを入手した後、システムはAIエージェント（大規模言語モデルによって駆動されるスマートなボット）に仕事を渡します。しかし、これは一度きりのやり取りではありません。このシステムは、**モンテカルロ木探索（Monte Carlo Tree Search）**と呼ばれる巧妙な戦略を使用しています。探偵が謎を解く場面を想像してください。一つの直感に従うだけでなく、探偵は多くの異なる経路を試します。もし一つの経路が行き止まりであれば、諦めるのではなく、戻って別の角度から試します。同様に、Kernel Forgeは単に一つのバージョンのレシピを書いて期待するだけではありません。多くのバージョンを生成し、テストし、もし一つが遅ければ、それを修正するか、あるいは全く異なるアプローチを試します。システムはすべての試みの「家系図」を保持し、最初は完璧ではなくても、有望であったものを記憶しておきます。
結果：スピードアップの勝利、しかしどこでもではない
チームは、4種類の異なるAIモデル（画像認識用のResNet-50、アート生成用のStable Diffusion 3.5 Medium、そしてチャットと推論用のGemma 4とQwen 3.5）を用いてKernel Forgeをテストしました。彼らはNVIDIA GB10 GPUを搭載した強力なコンピューターでこれらのテストを実行しました。
判明したことは以下の通りです：
- 機能するが、選り好みする： システムは、AIの多くの部分に対して、新しく、より高速なレシピの生成に成功しました。いくつかのケースでは、AIが書いたコードは、コンピューターが通常使用する標準的なコードよりも大幅に高速でした。例えば、画像生成器において、「グループ正規化（group normalization）」のための新しいコードは1.70倍高速でした。チャットボットのGemma 4では、「ソフトマックス（softmax）」（意思決定に使用される数学的ステップ）の新しいコードは、驚異的な2.83倍の速さでした。
- 「守られた」セーフティネット： システムは非常に慎重です。新しいレシピをチェックする「ガード（監視役）」を備えています。もし新しいAI作成レシピが遅かったり、間違いを犯したりした場合、システムは即座に元の信頼できるコードへと切り替えます。ユーザーに悪いレシピを強制することはありません。つまり、AIがコードの最適化を試みて失敗したとしても、コンピューターがクラッシュしたり遅くなったりすることはありません。単に、以前の安全な方法を使用するだけです。
- 「大物」を倒すのは難しい： 研究者たちは興味深いことに気づきました。AIの中で最も時間を占める部分（「ビッグプレイヤー」）は、多くの場合、NVIDIAのような企業による非常に成熟した、専門家が書いたコードによって処理されています。AIはこれらの専門家を打ち負かすのに苦戦しました。例えば、画像生成器において、「線形（linear）」演算（時間の50%以上を占める）は、AIが書き換えを試みた際に実際には遅くなっていました。システムは賢明にも、その部分については元のコードを使い続けるという判断を下しました。
- 小さな勝利の積み重ね： 最大のスピードアップは、コードのより小さく、それほど重要ではない部分で見られました。AIは最大のタスクにおいて専門家に勝つことはできませんでしたが、小さなタスクを1.5倍から2.8倍に高速化する巧妙な方法を見つけ出しました。
好奇心のコスト
研究者たちは、この「思考」にどれほどのコストがかかるかも調査しました。彼らは強力なAIモデルを使用してコードを生成しましたが、AIがコードを一行書いたり結果をチェックしたりするたびに、（APIの使用量に基づいた）わずかな費用が発生しました。彼らは、AIがより多くの修正と試行（最大50ラウンドの試行錯誤）を行うにつれて、コストが増加することを発見しました。しかし、得られた追加のスピードが、費やされた追加の金額と必ずしも一致しないことも分かりました。これは、AIが素晴らしい近道を見つけることはできるものの、特に修正しようとしているコードの部分が全体の速度にとってあまり重要ではない場合、いつ探索を止めるべきかについて、私たちは賢明である必要があることを示唆しています。
結論
Kernel Forgeは、AIが他のAIをより速く動かすための低レベルコードを書くのを助ける、新しい時代に入っていることを示しています。それは、あらゆる問題を即座に解決する魔法の杖ではありません。現時点では、最大のタスクにおいて世界のベストな人間エキスパートを簡単に打ち負かすことはできません。しかし、これは、システムの安全性と安定性を保ちながら、小さな詳細の中に隠れたスピードアップを見つけ出すことができる強力なツールです。
研究者たちは、このツールをオープンソースとして公開しました。つまり、誰でも自分のAIモデルを高速化するためにこれを利用できるということです。彼らは、AIの創造性と現実世界のテストの安全性を組み合わせることで、コンピューターサイエンスの博士号を必要とすることなく、デジタル世界のエンジンを最適化できることを証明しました。これは、私たちのコンピューターがより賢くなるだけでなく、より効率的になる未来への一歩なのです。
自分の分野の論文に埋もれていませんか？
研究キーワードに一致する最新の論文のダイジェストを毎日受け取りましょう——技術要約付き、あなたの言語で。