技術要約：ML GPUカーネルの等価性検証

問題提起

ディープラーニングと大規模言語モデル（LLM）の急速な進歩により、GPUカーネルの実行に膨大なコストが費やされるようになり、これらは積極的な最適化の主要な対象となっています。これらの最適化は、手動、コンパイラ、そしてますます増加しているLLMによって行われます。経験的なテストやプログラム解析がカーネルの検証に使用されていますが、それらには形式的な保証が欠けています。これは、GPUプログラミング特有の課題を考慮すると特に重要です。

1. 並行性のエラー： 数千のスレッド間で行われるきめ細かな同期は、テストだけでは検出が困難な微細なデータレースやデッドロックを導入します。これらのエラーは、稀な実行スケジュールにおいてのみ発生したり、ハードウェアバージョン間で保証されていない暗黙的な同期に依存したりする可能性があります。
2. 浮動小数点演算のセマンティクス： 最適化はしばしば算術演算の順序を入れ替えるため、リファレンス実装と最適化された実装との間に小さな数値的な差異が生じます。テストではこれらを許容することが多いですが、これによって真のセマンティックなバグ（例：クリッピング操作の欠落）が見逃される可能性があります。
3. 形式的手法の欠如： 既存の等価性検証技術は、シングルスレッドの整数プログラムや高レベルのテンソル演算をサポートしていますが、GPUカーネル固有の並列性、同期、および非決定的なスケジューリングを扱うことができません。GPUプログラムのための等価性チェッカーはこれまで存在しませんでした。

手法

著者らは、実用的なクラスのML GPUカーネル向けに特別に設計された初の等価性チェッカーであるVoltaを提案します。Voltaはブラックボックス・ベリファイアとして動作し、最適化プロセスに関する知識を必要とせずに、PTXアセンブリ（NVIDIA GPUスタックにおける最も文書化された低レベル層）を解析します。

1. 対象クラス：Structured-CTAs

VoltaはStructured-CTAs（Cooperative Thread Arrays）を対象としています。このクラスは以下を前提としています：

• スレッド数、テンソルサイズ、およびポインタターゲットが静的に既知であること。
• ブランチターゲットおよびメモリアクセスアドレスが、スレッドIDおよびループインデックスに基づいて静的に解決可能であること。
• 制御フローが実行時のデータ値に依存しないこと（効率的なMLカーネルはデータ依存の実行時決定を行いません）。
• これらの仮定は多くの高性能MLカーネル（JAXやXLAにおけるものなど）で成立しており、Voltaによってチェックされます。もし違反がある場合は例外が発行されます。

2. シンボリック実行と合流性（Confluence）

Voltaの核となるのは、テンソル値を実数として扱うシンボリック実行エンジンです（これは--ffast-mathのような最適化手法で使用される慣習をモデル化したものです）。

• シンボリック状態： エンジンはリファレンスカーネルと最適化されたカーネルをシンボリックに実行し、入力テンソルを関数とする出力テンソル要素の式を導出します。
• スケジューリング： ラウンドロビン・スケジューリング戦略を採用しています。スレッドは障壁（例：syncthreads, syncwarp）に到達するまで実行されます。
• 合流性： 主要な技術的成果は、合流性の証明です。シンボリック評価器は、すべての可能な実行スケジュールが同じシンボリック式を導く（またはレース/デッドロックを報告する）ことを保証します。この特性はAgda証明助手内で形式的に検証されており、チェッカーの結果が分析中に選択された特定のスケジュールに依存しないことを保証します。

3. レースおよびデッドロックの検出

Voltaは、障壁同期を利用してデータレースとデッドロックを検出するために、シンボリック実行を拡張しています。

• コンテキスト追跡： 各スレッドのメモリイベント（読み取り/書き込み）および同期ステータスを追跡するコンテキスト X を保持します。
• レース検出： スレッドが直前の書き込み手または以前の読み取り手と同期することなく、メモリロケーションを読み書きした場合にレースが検出されます。
• デッドロック検出： スレッドが、満たされることのない障壁を待ってブロックされた場合にデッドロックが報告されます。
• 保証： システムは、健全性（レースが報告された場合、具体的なレースが存在する）を満たし、「真陽性」の特性を満たすことが証明されています。

4. 決定手続き

シンボリック式が生成された後、Voltaはリファレンスと最適化された出力が数学的に等価であるかどうかを決定する必要があります。

• 式の形式： 式は多変数多項式、加算、乗算、および指数関数（例：$p(\bar{x})e^{h(\bar{x})}$）を含みます。
• 決定可能性： 著者らは、$\sum p_i(\bar{x})e^{h_i(\bar{x})} = 0$ という形式の等式の決定可能性を証明しています。この証明は、指数 $h_i$ が互いに異なる多項式である場合、和が恒等的にゼロであるためには係数 $p_i$ がゼロでなければならないという事実に依拠しています。これにより、重厚な数論的メカニズム（Lindemann–Weierstrassなど）を回避し、任意の実係数を扱っています。
• 標準化（Canonicalization）： 実装では、有理関数や指数項を正規化するための標準化器を使用し、softmaxのリスケーリングなどのケースを処理します。

主な貢献

1. 線形時間のレース/デッドロックアルゴリズム： 障壁同期を利用するStructured-CTAにおける、レースおよびデッドロックを検出するための線形時間アルゴリズム。
2. 決定可能性の証明： 実数上の多変数多項式と指数を含む等式、特に $\sum p_i(\bar{x})e^{h_i(\bar{x})} = 0$ の形式に関する決定可能性の証明。
3. Voltaの実装： 実用的なクラスのGPUカーネルのための初の等価性チェッカー。畳み込み、行列乗算、リダクション、およびアテンションメカニズムの正当性を検証します。
4. 形式検証： 主要な合流性補題はAgdaで形式的に検証されており、定義されたプログラムクラスに対して健全性と完全性の保証を提供します。

結果

Voltaは、以下のものを含む多様なベンチマークで評価されました：

• 人間が最適化したカーネル： 標準的なチュートリアルからの行列乗算およびリダクションカーネルの等価性を検証しました。これは、非推奨のワープ同期実行パターンによるデータレースを含むリダクションカーネル（Red-5, Red-6, Red-7）を正しく拒絶しました。
• LLMが生成したカーネル： LLM（例：2D畳み込みおよび行列乗算用）によって生成されたカーネルを検証しました。あるケースでは、NVIDIAのCompute SanitizerやLLM自身のテストではフラグが立てられなかった、LLM生成カーネル内のアウトオブバウンズ共有メモリ読み取りを検出しました。
• コンパイラが生成したカーネル： Tensor Coreを使用するものを含む、TileLangコンパイラによって生成されたカーネルの等価性を検証しました。
• アテンションメカニズム： 指数関数とmax操作を含む複雑なシンボリック式が存在するにもかかわらず、標準的なsoftmaxの実装と最適化されたオンラインsoftmaxの定式化（FlashAttentionで使用されるもの）の等価性を正常に検証しました。
• パフォーマンス： Voltaは、複雑なアテンションカーネル（例：FlashAttention-2）を3分未満で検証しました。決定手続きは、指数を含むベンチマークにおいて、複数の出力要素に対して標準化コストを償却する場合、一般的なSMTソルバ（Z3など）を大幅に上回る性能を示しました。

意義と主張

本論文は、GPUカーネルのための初の形式的等価性チェッカーを提供することを主張しています。その意義は以下の通りです：

• 形式的な保証： 経験的なテストよりも強い信頼性を提供します。これは、バグがコストに直結するプロダクション環境へのデプロイメントにおいて極めて重要です。
• 並行性の処理： 従来の等価性チェッカーが扱えなかった、GPUの同期（障壁、ワープレベル同期）の特有の課題に対処します。
• 最適化への堅牢性： 値を実数としてモデル化し、ビット単位の等価性ではなく数学的な同一性をチェックすることで、最適化が浮動小数点演算の順序を変更する場合でも正当性を検証します。
• 実用的な適用可能性： 「Structured-CTA」という仮定が、現実世界の多くのMLワークロード（畳み込み、行列乗算、アテンション）をカバーしていることを示し、そのアプローチが現実的な問題サイズにスケールすることを示しています。

著者らは、本研究がNVIDIA GPU（PTX経由）に焦点を当てているものの、その手法は他のランタイムにも適用可能であると述べています。また、動的な制御フロー（例：top-k/argsort）や非同期プリミティブ（例：パイプライン、TMA）をサポートしておらず、これらは今後の課題であることも認めています。