AIによる驚異的な高速カーネル生成 — サンフォード大学CRFMの最新研究成果 サンフォード大学の計算研究機械学習チーム（CRFM）は、PyTorchなどの専門家最適化済みのカーネルを上回る高速なAI生成カーネルを無意図に開発したと報告しています。これらは、CUDA-Cで完全に書かれており、CUTLASSやTritonなどのライブラリやDSLを使わずに実現しました。主要成果は以下の通りです： 行列乗算 (Matmul): 4096x4096の正方行列では参考コードの101.3%の性能 2次元畳み込み (Conv2D): (100, 3, 224, 224)の入力テンソルでは参考コードの179.9%の性能 ソフトマックス (Softmax): (4096, 65536)の入力テンソルでは参考コードの111.8%の性能 レイヤー正規化 (LayerNorm): (16, 64, 256, 256)の入力テンソルでは参考コードの484.4%の性能 2次元畳み込み + ReLU + MaxPool: (100, 3, 224, 224)の入力テンソルでは参考コードの290.1%、torch.compile()参考コードの189.0%の性能 これらの結果はNvidia L40S GPU上でベンチマークを実施し、性能の％は参考時間／生成カーネル時間で定義されました。 生成手法 KernelBenchという、2024年12月にリリースされたAIベースのカーネル生成向けベンチマークを使用しました。与えられたPyTorchコードに対して、LLM（大規模言語モデル）がカスタムカーネルを書き、その性能向上を目指します。デフォルト設定はFP32であり、精度の閾値（1e-02）を超える解も有効でした。カーネルの生成と性能評価は複数の無作為な入力に対して反復的に実施され、出力の数値的一致性によって正しさが確認されます。 変更点 自然言語での最適化アイデアの生成: 各ステップで、事前に試したアイデアに基づいて自然言語で新しい最適化アイデアを生成し、それを新しいコード版に実装します。 各最適化ステップでのブランチング: 単一の候補を修正するのではなく、それぞれのアイデアから複数の実装を派生させ、最高性能のカーネルを次のラウンドに使用します。これにより、並列的な探索が可能になり、異なる方向への探査が可能となります。 例：Conv2Dカーネルの最適化軌跡 Round 0: 7.02 ms, 参考の20.1% Round 1: 7.54 ms, 参考の18.8% Round 2: 3.46 ms, 参考の41.0% Round 3: 3.67 ms, 参考の38.7% Round 4: 3.46 ms, 参考の41.0% Round 5: 1.91 ms, 参考の74.9% Round 6: 1.37 ms, 参考の103.6% Round 7: 1.38 ms, 参考の102.9% Round 8: 1.37 ms, 参考の103.6% Round 9: 1.36 ms, 参考の105.1% Round 10: 1.07 ms, 参考の133.6% Round 11: 1.21 ms, 参考の117.4% Round 12: 1.01 ms, 参考の141.2% Round 13: 0.795 ms, 参考の179.9% 得られた教訓 この手法は、AI研究において「強い推論と多数の仮説の並列探索を組み合わせる」ことが改善につながるというトレンドを反映しています。一部の最近的工作では、大規模な再学習が必要ない場合もあり、巧妙な検索と分岐戦略が科学的イノベーションや複雑な問題解決につながることが示されています。しかし、さらなるトレーニングが必要であることに変りはありません。特に、さまざまな問題セットに対するより多くの合成データを生成することで、将来のモデル開発を改善することが期待されます。 未来の展望 今回の成果はまだ初期段階ですが、非常に有望であると感じています。特に、Flash Attentionでの性能向上が著しく、以前は機能的なバージョンすら生成できませんでした。これは、FP32における現代的なMLワークロードが少なく、最適化が難しいことも要因の一つかもしれません。それでも、研究チームは今後の進展に自信を持っています。 業界関係者のコメント Christopher Rinard、Saman Amarasinghe、Allen Nieは本研究の有用な議論に貢献しました。Standard Kernel Co.とPrime Intellectもこの研究をサポートしています。これらの成果はAIによる自己改善システムの開発に向けて重要な一歩となる可能性が高いとの見方が出ています。