新智元レポート
編集：Aeneas Hao Kun
【新智元導読】手慣らしでデータを合成しようとしたら、PyTorchの専門家カーネルをうっかり打ち負かしてしまった！スタンフォードの中国人チームが純粋なCUDA-Cで記述したAI生成カーネルは、瞬く間に業界を驚かせ、Hacker Newsのホットリストにランクインしました。チームは「この結果はもともと発表するつもりはなかった」とさえ述べています。
つい先ほど、スタンフォードHAIの中国人天才チームが再び驚異的な傑作を発表しました。
彼らが純粋なCUDA-C言語で記述した高速AI生成カーネルが、PyTorchを驚くほど凌駕しました！
このプロセスでは、CUTLASSやTritonなどのライブラリやドメイン固有言語（DSL）に頼ることなく、PyTorchに組み込まれた専門家によって最適化された標準的な生産レベルのカーネルに近い、あるいは場合によってはそれを上回るパフォーマンスを発揮することができました。
著者チームは、Anne Ouyang、Azalia Mirhoseini、Percy Liangというおなじみの名前です。興味深いことに、彼らは「この結果はもともと公表するつもりはなかった」とさえ率直に語っています。
発表されるやいなや、この発見は技術界に大きな衝撃を与え、現在Hacker Newsの総合ランキングで2位に浮上しています。
実を言うと、この発見には多くの偶然の要素が含まれています。
もともと、彼らの目標は、より優れたカーネル生成モデルを訓練するための合成データを生成することであり、合成データ生成の設計も非常に単純でした。
しかし、予期せぬ事態が発生しました。テスト目的のみで使用された合成データ生成自体が、非常に優れたカーネルを生成し始め、人間の専門家が最適化したPyTorchのベースラインさえも上回り、さらに高度な最適化とハードウェア機能を活用していました。
これまでは、これは非常に困難な課題でした。
このため、研究者たちは、この発見を共有するために、予定を繰り上げてブログ記事を執筆することを決定しました。
まとめると、研究の主要な成果は以下の通りです。
行列乗算（Matmul, FP32）：PyTorch FP32 torch.matmulの101.3%の性能を達成
2D畳み込み（Conv2D, FP32）：PyTorch FP32 torch.nn.Conv2Dの179.9%の性能を達成
Softmax（FP32）：PyTorch FP32 torch.softmaxの111.8%の性能を達成
層正規化（LayerNorm, FP32）：PyTorch FP32 torch.nn.LayerNormの484.4%の性能を達成
2D畳み込み + ReLU + Maxプーリング（Conv2D + ReLU + MaxPool, FP32）：PyTorch FP32参照実装の290.1%、PyTorch FP32 torch.compile()参照実装の189.0%の性能を達成
上記の結果は、Nvidia L40S GPUでベンチマークされ、性能パーセンテージは参照時間を発行されたカーネル時間で割ったものとして定義されます。
ネットユーザー：LLMの推論を強制するのは本当に面白い
Hacker Newsでも、ネットユーザーはこの件について活発な議論を繰り広げました。
例えば、FP32カーネルを使用する方がPyTorchよりも性能向上を達成しやすい理由については、かなり興味深いものがあります。
AIが低コストでより最適化されたカーネルを実現できるのであれば、確かにその可能性は計り知れません。
最も衝撃的なのは、最近のGoogleのAlphaEvolveも、o3がLinuxカーネルでゼロデイ脆弱性を発見したことも、すべて私たちに思い出させてくれることです—
Gemini Pro 2.5とo3は、新たな能力レベルに到達しました。以前他のモデルで試して失敗したアイデアが、今では突然機能するようになったのです。
私たちは、LLMが人間よりもはるかに速い速度で反復とテストを行うことができる段階に到達したと言えます。情報の組み合わせ、進歩、そしてインテリジェントなアプリケーションの力技が、成功しつつあるようです！
次に、スタンフォードの研究者たちのブログの具体的な内容を見ていきましょう。
ブログ全文
方法
研究者たちはKernelBenchのタスク設定を採用しました（これは彼らが2024年12月に発表したAIベースのカーネル生成ベンチマークです）。
具体的には、torchコードが与えられると、LLMは元のtorch演算子を置き換えるカスタムカーネルを記述し、高速化を目指します。
KernelBenchの初期設計によれば、参照コードはデフォルトでFP32精度を使用しますが、与えられた許容誤差閾値（1e-02）の下では、低精度ソリューションも許可されます。
さらに、特定の規模に特化した最適化手法が多数存在するため、KernelBenchの各問題には具体的な入力サイズが設定されています。
したがって、このベンチマークは、任意の規模に適用できる高速カーネルではなく、特定の規模の問題に対して最速のカーネルを見つけることを目的としています。
また、研究者たちはtorch参照コードと生成されたコードを同時に実行し、複数のランダムな入力で両者の出力値が一致するかを比較することで、その正確性を検証します。
現在、カーネルの最適化という問題において、業界でテスト時の計算リソースを拡張する最も一般的な方法は、逐次的な修正（sequential revision）です。
これは多段階反復ループであり、モデルはまずカーネルを段階的に変更し、次にその正確性と性能をチェックし、その結果に基づいて再度試行します。
つまり、問題のあるカーネルを修正するか、既存のカーネルの性能をさらに向上させるかのどちらかです。
この反復プロセスは非常に直感的で、実装も容易です。モデルは失敗したカーネルを修正し、利用可能なカーネルを微調整し、段階的に性能の向上したバージョンを最適化していきます。
この方法の主な限界は、最適化のアイデアの多様性が欠如している点にあります。
逐次ループは、同じ種類の変換を繰り返し試したり、潜在力に乏しい最適化パスを際限なく調整したりするなど、局所最適に陥りやすい傾向があります。
その結果、テスト時の計算リソースの非効率な利用につながり、モデルが根本的に革新的な最適化アイデアを生み出すことを促進するのが困難になります。
この問題を解決するため、研究者たちは2つの重要な変更を導入しました。
最適化のアイデアを自然言語で推論する
彼らは、各ステップで直接新しいカーネルを生成するのではなく、以前試したアイデアを条件として自然言語で最適化のアイデアを生成し、その後、これらのアイデアを新しいコードバリアントとして具体化します。
各最適化ステップで分岐を拡張する
彼らは各ステップで1つの候補ソリューションのみを改善するのではなく、分岐を拡張し、各アイデアから複数の実装バージョンを派生させ、その中で最も性能の良いカーネルを次ラウンドの最適化のシードとして使用します。
（研究者たちは、優れた性能を持つ既存のカーネルライブラリも保持し、シードとして提供します）。
この方法は、大規模な並列処理能力を解放し、各ラウンドで根本的に異なる最適化方向を探索することを可能にし、狭い最適化パスに閉じ込められることを回避します。
その結果、このテスト時ループは、逐次的な修正のように単にコンパイラと「対話」するだけでなく、より構造化された探索的検索に近いものになりました。
この検索は、明確な最適化仮説によって導かれ、大規模な並列評価を用いて実行されます。
研究者たちは、テストのためにKernelBenchレベル1の10の問題を実行しました。
彼らは問題の規模を調整し、カーネルの起動オーバーヘッドが問題全体の実行時間に対して無視できることを確認しました。
その後、OpenAI o3とGemini 2.5 Proモデルを使用して5ラウンドの実験が行われました。
以下の図は、最も性能の良いカーネルが最初に発見されたラウンドの分布を示しています。
ほとんどの最適な結果は、後半のラウンド（合計5ラウンド）で現れており、その大部分は4ラウンドまたは5ラウンドで出現しています。
検索範囲を拡大するにつれて、研究者たちはまた、多くの高性能カーネルの最適化戦略が非常に類似しており、いくつかの一般的なパターンに集中していることを発見しました。これは、彼らが手動でカーネルを記述した経験とも一致しています。
主要な最適化カテゴリは以下の通りです—
メモリアクセスの最適化：異なるメモリレベル（グローバルメモリ、共有メモリ、レジスタ）間のデータ転送効率を向上させ、データアクセスパターンが帯域幅を最大化し、競合を最小化するようにします。
非同期操作とレイテンシの隠蔽：時間のかかる操作（例：グローバルメモリアクセス）を計算や他のメモリ転送と重複させることで、それらによって引き起こされるレイテンシを隠蔽します。
データ型と精度の最適化：許容される範囲で、可能な限り低精度のデータ型（FP16やBF16など）を使用し、メモリ帯域幅の要件を減らし、キャッシュ効率を向上させ、専門のハードウェアアクセラレーションユニットを活用する可能性を高めます。
計算と命令の最適化：算術演算自体の効率を向上させ、命令数を削減するか、専門のハードウェア命令を利用します。
並列性と占有率の向上：ストリーミングマルチプロセッサ（SM）上のアクティブなワープの数を最大化し、レイテンシをより良く隠蔽し、全体的なスループットを向上させます。
制御フローとループの最適化：ループ、分岐、およびインデックス計算によって導入される追加のオーバーヘッドを削減します。
まとめ
今回研究者たちが採用した方法は、AI研究においてますます顕著になっているトレンドと一致しています—
強力な推論能力と複数の仮説の並列探索を組み合わせることで、性能向上がもたらされます。
最近の研究（例：AlphaEvolve、Gemini 2.5 Pro Deep Think）が強調しているように、私たちは常に大規模な再学習を必要とするわけではありません。
時には、巧妙な探索と分岐戦略だけで科学的革新を生み出し、複雑な問題を克服することができます。そして、検証器を用いた広範な探索は、さらに大きな利益をもたらす可能性があります。
しかし、これはさらなる学習が不要であることを意味するものではありません。
むしろ、研究者たちのこの方法は、将来のモデル学習を改善するためのより高品質な合成データを生成するのにも役立ちます（これにはより多くの問題インスタンスが必要です）。
したがって、これは強力なテスト時拡張方法であると同時に、より知的でデータ効率の高いモデル開発への一歩でもあります。
さらに、今回研究者たちが示したのは予備的な成果に過ぎません。これらの最適化結果の品質はかなりのものに見えますが、例えば、より優れた最適化アイデアの生成、より高品質な最終コードの生成、そしてこの方法をますます複雑なカーネルに適用するなど、広大な改善の余地がまだあります。
現在、研究者たちが積極的に改善に取り組んでいる具体的な2つの例は以下の通りです。
FP16 Matmul：torch.matmulの52%の性能を達成
FP16 Flash Attention：torch.nn.functional.scaled_dot_product_attentionの9%の性能を達成
現代の機械学習タスクでは、FP32の利用はFP16やBF16ほど一般的ではなく、新しいハードウェアではFP32向けの最適化が少ない傾向があります。
これが、FP32ベースのカーネルがPyTorchの性能を上回りやすい理由の一部を説明できるかもしれません。
著者紹介
Anne Ouyang
Anne Ouyangは現在、スタンフォード大学コンピュータサイエンス（CS）の博士課程学生で、Scaling Intelligence Labで研究を行っています。
彼女の研究関心は、スケーラブルな自己改善型機械学習システムに主に焦点を当てており、経験的機械学習（empirical ML）と性能工学（performance engineering）にも広く関心を持っています。
以前はMITで学士号と修士号を取得し、NVIDIA cuDNNチームで勤務し、GPU上の深層学習ワークロードを高速化するためのCUDAカーネルの記述を担当していました。
Azalia Mirhoseini
Azalia Mirhoseiniは、スタンフォード大学コンピュータサイエンスの助教であり、Scaling Intelligence Labの創設者でもあり、Google DeepMindのシニア研究科学者も兼任しています。
彼女のラボは、汎用人工知能の発展を推進するため、スケーラブルな自律進化型AIシステムと方法論の開発に専念しています。
スタンフォード大学に加わる前は、Google BrainやAnthropicなど、業界トップのAIラボで長年勤務していました。
彼女の過去の卓越した業績には、以下のようなものがあります。
混合エキスパート（MoE）ニューラルアーキテクチャの提案—現在、最先端のAIモデルで広く使用されています。
AlphaChipプロジェクトのリーダー—深層強化学習をレイアウト最適化に応用した画期的な研究であり、Google AIアクセラレーター（TPU）やデータセンターCPUなどの先進チップの設計に成功裏に適用されました。
テスト時の計算のスケーリングに関する詳細な研究
Percy Liang
Percy Liangはスタンフォード大学コンピュータサイエンスの准教授であり、基礎モデル研究センター（CRFM）の所長を兼任しています。また、CodaLab Worksheetsの作成者でもあり、それを通じて研究成果の再現性を強く提唱しています。
彼は現在、オープンソースと厳格なベンチマークを通じて、基礎モデル（特に大規模言語モデル）のアクセシビリティと理解可能性の向上に注力しています。
彼はこれまで、機械学習と自然言語処理の分野で幅広い研究を行っており、具体的な方向性としては、ロバスト性、解釈可能性、人間とコンピュータのインタラクション、学習理論、知識の定着、意味論、推論などが含まれます。
以前は、2004年にMITで学士号を、2011年にUCバークレーで博士号を取得しました。
カーネル最適化プロセス
最後に、Conv2Dの自動生成アイデアの最適化軌跡の例を示します。torch参照の基準時間は1.41ミリ秒でした。
ラウンド0：7.02ミリ秒、参照性能の20.1%
アイデア：与えられたPyTorchコードについて、CUDAカーネルで操作を置き換える。
ラウンド1：7.54ミリ秒、参照性能の18.8%
アイデア：__ldgを使用して不変テンソルをロードし、読み取り専用キャッシュを利用する。
ラウンド2：3.46ミリ秒、参照性能の41.0%
アイデア：畳み込みをFP16 Tensor-Core GEMMに変換する。
著者コメント：これはアルゴリズム最適化であり、畳み込みを暗黙的なGEMMに変換します。これはTensor Core上で畳み込みを効率的に実行するために不可欠です。
ラウンド3：3.67ミリ秒、参照性能の38.7%
アイデア：ダブルバッファリングされたcp.asyncパイプラインを採用し、グローバルメモリのロードとTensor-Core計算をオーバーラップさせる。
ラウンド4：3.46ミリ秒、参照性能の41.0%
アイデア：与えられたPyTorchコードについて、暗黙的な行列乗算（implicit matmul）用のCUDAカーネルで操作を置き換える。与えられたGEMMカーネルが役立つかもしれない。
著者コメント：最適化にGEMMの使用が含まれるため、このラウンドの開始時に、以前に生成された既存の優れたGEMMカーネルをシードとして使用しました。このアイデアは手動で記述されました。
ラウンド5：1.91ミリ秒、参照性能の74.9%
アイデア：各K-tileループ内で、k_idx分解されたカーネル/入力インデックスを事前計算し、共有メモリで再利用することで、冗長な算術演算を回避する。
ラウンド6：1.37ミリ秒、参照性能の103.6%
アイデア：共有メモリでN次元GEMMインデックスを事前計算しキャッシュすることで、入力データロードループ内の冗長な算術演算を削減する。
ラウンド7：1.38ミリ秒、参照性能の102.9%
アイデア：専用のワープごとの共有メモリバッファを使用することで、CUDAカーネル出力を並列化し、グローバルメモリへの書き込み時のワープレベルの直列化を排除する。
ラウンド8：1.37ミリ秒、参照性能の103.6%
アイデア：共有メモリで基本入力座標を事前計算しキャッシュすることで、入力データロードプロセスにおける冗長な算術演算を削減する。
ラウンド9：1.36ミリ秒、参照性能の105.1%
アイデア：B-フラグメントのロードをソフトウェアパイプライン化し、次のB-タイルからの共有メモリ読み出しと、現在のB-タイルのWMMA計算をオーバーラップさせる。
ラウンド10：1.07ミリ秒、参照性能の133.6%
アイデア：共有メモリで事前計算されたN次元GEMM分解を再利用して出力アドレス計算を行い、冗長で高コストな除算/剰余演算を回避する。
ラウンド11：1.21ミリ秒、参照性能の117.4%
アイデア：half WMMA演算からhi/lo分解を削除する。最終的な精度が許容できる場合は、標準のFP16累積に依存して性能を向上させる。
ラウンド12：1.01ミリ秒、参照性能の141.2%
アイデア：パイプラインのロード段階でK次元インデックスを動的に計算することで、ダブルバッファリングを使用してAsub（重み）とBsub（入力）のKループのグローバルメモリロードとMMA計算をオーバーラップさせる。
ラウンド13：0.795ミリ秒、参照性能の179.9%
アイデア：half2のようなより広いデータ型を使用することで、Asub_pipeとBsub_pipeのロードに対するベクトル化された共有メモリ書き込みを実現する。
最終コード
最終的に生成されたConv2Dカーネルコードは、高度なCUDA技術を使用しており、人間が自分で書くには非常に難しい種類のものでした！
参考文献：
https://crfm.stanford.edu/2025/05/28/fast-kernels.html
https://news.ycombinator.com/item?id=44139454