AI生成GPUカーネルの形式検証
AIエージェントが高性能なGPUカーネルを生成するようになるにつれ、その出力への信頼構築が新たなボトルネックとなっています。Gimlet Labsは、数値テストでは見逃されるバグを形式検証で補完する初期研究システムを構築しました。スケーリングドット積注意機構における欠落したクランプ操作の例を通じて、その有効性を示します。
2026年7月8日、Gimlet Labsの研究チームは、AI生成GPUカーネルの形式検証に関する初期研究を発表しました。AIエージェントが高性能カーネルを生成する能力が向上するにつれて、その出力への信頼が新たな課題となっています。従来の数値テストは限られた入力空間しかカバーできず、AIエージェントはテストスイートに特化した最適化を行う可能性があります。
具体例として、スケーリングドット積注意（SDPA）を考えます。リファレンス実装では、中間値が[-5,5]にクランプされますが、AI生成の最適化カーネルは融合操作のためにこのクランプを削除しました。ランダムテストでは両者の出力は類似していましたが、本番環境では入力の境界によって差異が生じる可能性があります。このような欠陥は、テストケースを増やすだけでは根本的に解決できません。
Gimletの検証システムは、リファレンスモデル（PyTorch）と候補カーネル（Triton）をそれぞれテンソル代数DAGに変換し、基本演算に分解した後、単一の出力要素の式にスカラー化します。その後、Z3ソルバーを用いて二つの式を比較し、等価性の証明（UNSAT）または反例（SAT）を返します。この手法は、コンパイラ分野での形式検証（LLVMのAlive2やCompCertなど）に類似しています。
現在のシステムは研究プロトタイプですが、実際のワークロードで数値テストを通過したエラーを捕捉しています。チームはARRAY 2026会議で完全な成果を発表しました。将来的には、形式検証がAI生成カーネルの本番展開における重要な保証手段になると期待されています。
さらに、Gimlet Labsはこの手法をCUDAカーネルなどより低レベルの実装に拡張する方法も模索しています。CUDAはTTIRのような検査可能な中間表現を提供しないため、計算グラフを回復するための特別な抽出パイプラインが必要です。研究チームは、検証器の性能をさらに最適化し、より大規模なカーネルを処理できるようにし、最終的にはCI/CDパイプラインに統合して、AI生成カーネルのデプロイメントに自動化された正当性保証を提供することを計画しています。