広告

ドメイン固有言語とスピード・オブ・ライト指針を用いたGPUカーネル最適化エージェントの効率改善

arXiv cs.AI / 2026/4/1

💬 オピニオンDeveloper Stack & InfrastructureIdeas & Deep AnalysisModels & Research

要点

  • 本論文は、LLMエージェントによるGPUカーネル最適化が高コストであるのは、各試行で大きな設計空間にわたって候補を生成し、コンパイルし、検証し、プロファイルする必要があるためだと主張する。
  • そこで効率化のための2つの原則として、重要な最適化レバーを維持しつつ、低影響な推論を減らすためにコンパクトな高水準のドメイン固有言語(DSL)を用いること、そしてSpeed-of-Light(SOL)指針によって性能の余地(headroom)を見積もり、限界(逓減的な効果)が見えてきた近辺で探索を無駄にしないことを提案する。
  • 実装であるμCUTLASSは、構成(configuration)、エピローグ融合(epilogue fusion)、マルチステージ・パイプラインを対象とする、CUTLASS対応GPUカーネルのためのDSLとコンパイラを提供する。
  • KernelBenchの59問に対する実験では、低水準のコード生成からDSLベースの生成(GPT-5-mini)へ移行することで、性能が0.40×の後退からPyTorch比で1.27×の高速化へと変化し、さらにSOL指針により1.56×まで向上する。
  • SOL指針による予算管理(budgeting)により、LLMのトークン使用量を19〜43%削減しつつ、少なくともgeomeanでの高速化の95%以上を維持できる。加えて、速く見えるが意図した計算ができていない「ベンチマーク・ゲーミング」型のカーネルを検出できる可能性もある。

要旨: LLMエージェントでGPUカーネルを最適化することは、大きな設計空間に対する反復的なプロセスである。すべての候補は生成され、コンパイルされ、検証され、プロファイルされなければならないため、試行回数を減らせば実行時間とコストの両方を節約できる。私たちは2つの重要な観察を行う。第一に、エージェントが動作する抽象化レベルが重要である。抽象化レベルが低すぎると、LLMは影響の小さい詳細に対して推論を浪費する。逆に高すぎると、重要な最適化の選択肢を見落とす可能性がある。第二に、エージェントは限界(収穫逓減)の点に到達したかどうかを容易に判断できず、探索を続けて資源を浪費してしまう。
これらの観察は、効率を高めるための2つの設計原則を導く。(1) 文脈内で学習でき、重要な最適化レバーを保持しながらより高いレベルでモデルが推論できる、コンパクトなドメイン特化言語(DSL)、および(2) 第一原理に基づく性能上限を用いて探索を誘導し、探索予算を管理するSpeed-of-Light(SOL)ガイダンスである。私たちは、CUTLASSベースのGPUカーネル向けのコンパイラを備えたDSLである

CUTLASSにおいて、これらの原則を実装する。このDSLは、カーネル構成、エピローグ融合、多段パイプラインをカバーする。SOLガイダンスを用いてヘッドルームを推定し、最適化試行を導き、SOLに近い問題を優先度を下げ、ベンチマークを“稼ぐ”カーネルをフラグ付けする。
同じ反復予算での59のKernelBench問題において、低レベルコードの生成から、GPT-5-miniを用いたDSLコードへの切り替えにより、0.40xのgeomean回帰がPyTorchに対する1.27xの速度向上に変わる。SOLに基づく誘導を追加すると、1.56xまで向上する。モデルの階層を通して、
CUTLASS + SOLガイダンスは、より弱いモデルでも、トークンコストが低い状態で、より強いベースラインのエージェントを上回れることを可能にする。SOLに基づくバジェット管理は、少なくともgeomeanの速度向上を95%以上維持しつつ、トークンを19-43%節約する。最良の方針では、1.68xの効率改善に到達する。最後に、SOL分析は、意図された計算を行えないにもかかわらず速く見える、ベンチマークを“稼ぐ”ケースの検出にも役立つ。

広告