テクニカル・サマリー：GPUフォキャスター：カーネルの実行時最適化のための選択的サロゲートとしての言語モデル

問題提起

GPUカーネルは現代のディープラーニングにおける基本的なワークホースであり、そこでの微細な効率向上は、トレーニングおよび推論コストの大幅な削減につながります。これらのカーネルを最適化するには、通常、大規模言語モデル（LLM）が候補となるカーネルを提案し、それらをコンパイルして物理ハードウェア上で実行して性能を測定するという「生成・評価サイクル」を伴います。しかし、この評価プロセスは深刻なボトルネックとなります。各評価にはコンパイル（CUDAの場合60〜120秒）とデバイス上での実行が必要であり、そのため評価のスループットは生成のスループットよりも数桁遅くなります。探索予算が増大するにつれ、繰り返される物理的な測定のコストが、潜在的なカーネル実装の広大な空間を探索する能力を制限します。

先行研究はLLMを用いてより優れたカーネルを「生成」することに焦点を当ててきましたが、本論文は、評価コストを削減するという補完的な問題に対処しています。著者らは、LLMを単なる生成器としてではなく、物理GPUに対する**選択的サロゲート（代理モデル）**として使用することを提案しています。目標は、候補となるカーネルの相対的な性能をリファレンス（参照）実装と比較して予測し、モデルが不確実性を検知した際に、高価な物理GPUへの問い合わせを保留することです。

手法

タスク定式化：選択的予測

著者らは、ハードウェア仕様が与えられた条件下で、候補となるカーネルの（リファレンス・カーネルに対する）相対的なスピードアップ（$S_i = T_{ref} / T_i$）を予測する問題として定式化しています。コンパイラの挙動やマイクロアーキテクチャの詳細をモデリングすることは困難であるため、実行時間を秒単位の絶対値として予測することは困難であると判断されました。代わりに、ターゲットは8つの順序的なスピードアップ・ビン（「大幅な低速化」から「極端な高速化」まで）に離散化されます。

サロゲート（LLM）は、リファレンス・カーネル、候補カーネル、およびハードウェア仕様を入力として受け取ります。モデルは推論の連鎖（Chain of Reasoning）を生成し、8つのビンに対する確率分布を出力します。モデルが表明する信頼度は、任意の単一のビンに割り当てられた最大確率として定義されます。

選択的予測戦略

システムは選択的予測フレームワークの下で動作します。サロゲートは候補のランキングに使用されます。サロゲートの信頼度が特定の閾値を超えた場合、その予測は受理され、候補はGPU上で測定されません。信頼度が低い場合（不確実性を示唆する場合）、その候補は真値の測定のために物理GPUへと転送されます。これにより、固定されたGPU測定予算内で、より多くの候補を探索プロセスで検討することが可能になります。

キャリブレーションのための強化学習（RL）

サロゲートの精度、そして極めて重要な点として、そのキャリブレーション（表明された信頼度と実際の精度の間の整合性）を向上させるために、著者らはオープンウェイトモデル（GPT-OSS-20B）に対して強化学習を適用しています。彼らは以下の要素に基づく報酬を定義しています：

1. 正確性（Correctness）： 予測されたビンが測定されたビンと一致しているかどうか。
2. カテゴリカル・キャリブレーション（Brier Score）： 真のワンホット・ターゲットからの距離に基づいて分布にペナルティを課す。
3. 距離を考慮したキャリブレーション（CRPS）： スピードアップ・ビンの順序性を考慮して累積分布にペナルティを課す（例：真の値が1.5倍であるときに4倍と予測することは、1.5倍と予測することよりも深刻なエラーであるとする）。

訓練データは、カーネル探索自体の副産物として、過去の探索ランにおける成功した候補を使用して構築されています。

主な貢献

1. GPUサロゲートとしてのLLM： 本論文は、LLMが物理GPUの応答を予測する上で非自明な精度を達成できることを示しており、測定に基づくのではなく予測に基づくランキングによって、利用可能なスピードアップの大部分を回収できることを実証しています。
2. 選択的予測フレームワーク： 著者らはGPU予測を選択的予測問題として捉え、サロゲートがキャリブレーションされた不確実性を報告できることを示しました。これにより、高信頼度の予測を信頼し、不確実なケースに対してのみ物理的な測定を予約するという探索戦略が可能になります。
3. サロゲートの挙動へのRLの影響： 実験を通じて、本論文はRL訓練がサロゲートの特性をどのように変化させるかを特徴付けています。具体的には、キャリブレーション報酬（BrierまたはCRPS）を追加することで、信頼度と精度の相関関係が改善される一方で、正確性のみを学習させると、信頼度と精度の関係が劣化（誤った予測に対して過度に自信満々になる）することが示されました。
4. エンドツーエンドの探索改善： サロゲートをカーネル探索ループ（具体的にはPUCTベースの探索）に統合することで、GPU測定あたりの候補探索数が大幅に増加することを示しています。ほとんどのタスクにおいて、これは同一のGPU予算下で全ての提案を測定するベースラインよりも優れたカーネルを見つけ出す結果となりました。
5. データセットの公開： 著者らは、118の問題、複数のGPUタイプ（A100, L40S, H100）、および様々な生成手法にわたる、測定された実行時間を伴う12,388個のLLM生成GPUカーネルの再利用可能なコレクションを公開しています。このデータセットは、GPU性能予測およびカーネル探索の幅広い研究をサポートします。

結果

• 予測精度： オフザシェルフのLLM（例：GPT-OSS-20B）は、上位1%の候補のみをGPUで測定する場合でも、利用可能なベスト・スピードアップの約83%を回収します。この性能は、上位50%を測定する場合に94%まで向上します。
• キャリブレーションとRL： RL訓練はサロゲートの挙動を大きく変えます。
  • 正確性のみの訓練は、正確なビンの一致率を向上させますが、信頼度と誤差の相関を悪化させます（モデルが誤った予測に対して過信状態になります）。
  • **キャリブレーション訓練（Brier/CRPS）**は、予測誤差と信頼度の単調性を改善します。特筆すべきは、Brierでキャリブレーションされた訓練が、ベースラインモデルよりも「回収されたスピードアップ」という指標を約2.1パーセントポイント向上させたことであり、これは、より適切にキャリブレーションされた分布が、探索アルゴリズムに対してより優れたランキング信号を提供することを示唆しています。
• 探索効率： エンドツーエンドのカーネル探索において、サロゲートを有効にした探索は、6つのタスクのうち4つでベースラインよりも高速なカーネルを見つけ出し、多くの場合、ベースラインのベスト・スピードアップに到達する際により少ないGPU測定数で済みました。ベースラインが勝利したタスクでは、探索が早期に飽和しているように見受けられ、これはサロゲートが探索空間が広く未探索である場合に最も有益であることを示しています。
• 発見の瞬間（Discovery Moments）： サロゲートは、「子カーネルが親カーネルよりも大幅に高速である変異」をランダムな試行よりも高い精度で特定できますが、精度には依然として課題があり、重大な決定にはGPUによる確認が必要となります。

意義と主張

本論文は、LLMが単なるコード生成を超えて、より広範な役割を果しくできることを主張しています。GPUの仮想モデルとして機能することで、物理的なハードウェアリソースの希少性を管理する、高スループットかつ低コストなフィルターとして機能できます。

著者らは、このアプローチの主な価値は速度だけでなく、あらゆるステップで元のシステムに問い合わせることなく、モデルベースの最適化を実行できることにあると強調しています。サロゲートは、物理GPUよりも高いスループットと低いクエリコストを提供し、その代わりに、キャリブレーションされた不確実性推定を通じて管理可能な「予測ノイズ」を受け入れます。この研究は、カーネル最適化を複雑な物理システムのニューラルシミュレータ構築という広範な分野へと結びつけており、ハードウェア測定を「選択的に使用すべき希少なリソース」として扱うオプティマイザーへの道筋を示唆しています。