特化型GPUカーネル生成による極限の効率性の実現

従来、本番環境の推論システムは、多様なモデルやワークロードを処理するために汎用的なカーネルに依存しています。しかし、GPUの演算シェイプは静的なモデルパラメータとリクエスト時の動的な要因の組み合わせによって決定されるため、これは最適とは言えません。例えば、行列乗算の次元の一方はモデルによって定義されますが、もう一方は各リクエストの具体的なトークン数に基づいて変動します。現在、エージェントによるGPUカーネル生成への関心が高まっており、最近の取り組みでも有望な成果が示されています。私たちは、ある本質的な疑問を掘り下げました。もしカーネル生成を自動化できるのであれば、なぜ（10億から1兆パラメータという）まったく異なる規模のモデルが同じカーネルに依存する必要があるのでしょうか？実行時に発生する特定のシェイプにカーネルを特化させることで、極めて高い効率性を実現できます。

このブログでは、エージェントを使用してGPUカーネルを生成したことによる成功体験とインサイトを共有します。私たちは、極限の特化を実現するために設計されたシステムであるProteusを構築しました。これには、厳格な最適化、検証、およびコンテキスト管理に合わせてカスタマイズされたハーネスが必要です。

従来型のコーディングハーネスは、ここで失敗することがよくあります。なぜなら、エージェントは「報酬ハック（reward-hack）」、つまり本質ではなくルールを文字通りに解釈して従う傾向があるからです。エージェントにベンチマークを与えると、意図した演算ではなく、そのベンチマークだけを最適化してしまう可能性があります。

これに対処するため、Proteusはカーネルを提案し、管理されたリファレンス実装と比較して検証し、成功したものの実行時間を測定して、最良の結果を反復的に改善します。プロセス自体はシンプルですが、その成功は2つの根本的な課題を解決できるかどうかに完全にかかっています。図1は、私たちの設計の簡略化されたアーキテクチャを示しています。Proteusハーネスを使用することで、vLLMで利用可能な最速のカーネルよりも1.8〜5.2倍高速なQwen 3.5 122Bカーネルを生成しました。

検証

私たちは当初、カーネル探索こそが難しい部分だと考えていました。つまり、改善が見られない停滞状態（プラトー）に陥ることなく、どのようにして膨大なプログラム空間を探索するか、ということです。しかし実際には、最初の疑問はもっと基本的なものでした。それは、「私たちは本当に測定したいものを測定できているのか？」という点です。

モデルは、与えられたスコアを最適化します。特別な裏技を使う必要はありません。単に評価プロセスが何らかの前提を置いているだけという場合もあります。一例として、アテンションレイヤーで一般的なステップである回転位置埋め込み（RoPE）のカーネルが挙げられます。ある候補は、以前の試行から残されたコンパイル済みコードを再利用することで、ゼロから公正に再構築するよりもコストが低く見える（高速である）ように見せかけることができました。また別の候補は、GPUの起動バッチをグラフ（CUDAグラフなど）に記録して1つのユニットとして再生する一方で、比較対象のベースラインは依然として各パーツを個別に起動していたため、両者が同じ処理を行っていないというケースもありました。さらに、公開されているテストセットに含まれる入力サイズに対しては強力であるものの、提示されていないサイズに対しては脆弱であるというケースもありました。

そのため、私たちは初期の設計作業をプロンプトではなく、チェッカー（検証器）に費やしました。クロスチェックが必要な場合は、複数のタイマー（CUDAイベントタイマー、実時間、CUPTIタイマーなど）を使用し、両者を同じ方法で測定します。残すべきではないコンパイル済みの状態をクリアし、セットアップとクリーンアップの順序を一貫させることで、一方が他方の負担する処理をスキップできないようにします。次のラウンドの開始点として使用する前に、勝ち残った候補の実行時間を再測定します。候補が見ることのできないテストを一部残しておくことで、テスト対策だけに特化してしまうのを防ぎます。パフォーマンスが人為的に誇張される評価の「不正行為」を防ぐため、物理的なGPUの帯域幅や計算限界を超える、理論上不可能な高速化（例：>100x）にフラグを立てる自動整合性チェックを実装しています。これにより、エージェントが真のパフォーマンス向上ではなくハーネスのメトリクスに合わせて最適化してしまうという、過去の業界事例で見られた報酬ハックの落とし穴を防ぐことができます。これらの制約がなければ、カーネルをさらに生成しても、ほとんどノイズが増えるだけでした。

チェッカーを重視することは、エージェントによるカーネル生成のボトルネックをも変化させます。単なるプログラム探索（システムがバリアントを繰り返し生成してプログラムを探索する反復的な最適化）の作業では、優れた候補は稀であるため、それらを作成することにコストの大部分が費やされます。多くのドラフトを並行して作成することはできますが、検証をスキップすることはできません。検証を慎重に設計する必要があり、チェックは実際のGPU上で、隔離された環境で、複数回実行されなければなりません。システムが進化するスピードは、カーネルをいかに速く書けるかではなく、いかに速くそのカーネルを信頼できるかによって決まります。

コンテキスト管理

もう1つの課題は、カーネル生成モデルに何を提示するかを決定することです。これはトレードオフの関係にあります。モデルに大きなプロンプトを与えれば、現在の最適なカーネル、最近の失敗、プロファイラのヒント、以前の実行時のメモなど、より多くの情報が得られます。これは役立つ場合があります。しかし、モデルが読み取るすべてのトークンに対してコストが発生するため、費用も高くなります。また、プロンプトが大きくなるにつれて、次の試行が本来の方向から逸脱しやすくなります。有用なシグナルが、古いアドバイス、矛盾するヒント、異なる入力サイズや異なる演算に適用される詳細情報などと混ざり合ってしまうからです。モデルは常にどの文を信頼すべきかを判断できるわけではないため、最も目立つ指示に従ったり、すべての指示を少しずつ取り入れたりしてしまいます。

逆に、与える情報が少なすぎると、その反対のことが起こります。すべての試行がゼロからスタートすることになります。同じ行き止まりに何度も突き当たります。前回の実行や関連する演算から何も引き継がれず、ループが前進しません。

私たちは、これを支援するためのナレッジレイヤーを必要としていました。つまり、何がうまくいったかを記憶し、後で再利用し、それを人間を介さずに行うというものです。このレイヤーには、保存される教訓の具体性と、それがどれだけ広く適用できるかという、2つ目のトレードオフが存在します。

非常に具体的なメモ（「この入力サイズを持つこのカーネルでは、このループを展開する」など）は、次の試行にまさに必要なものである場合があります。しかし、次の演算、次のGPU、あるいは異なる入力サイズに対して誤って適用されやすいという側面もあります。非常に一般的なメモ（「オンチップメモリをより有効に活用する」など）は、ほぼどこにでも当てはまりますが、モデルに具体的なアクションをほとんど指示しません。私たちは、これら両方の失敗パターンを目にしました。教訓が一般的すぎると、具体的なアクションを示すことなく失敗を再確認するだけになってしまいます。より詳細な情報を保存すると、それらは特定の実行に密接に結びつきすぎて、次の実行の役には立たないことがよくありました。ある長期の実行では、モデルが読み書きした内容の大部分は、カーネルの作成ではなく、そのメモリの取得とルーティングに費やされていました。メモリレイヤーは多くの処理を行っていましたが、次の候補を改善することにはつながっていませんでした。図2は、そのようなシステムのトークンコストの内訳を示しています。コストの大部分をナレッジレイヤーが占めています。

保持する価値のあるナレッジは、よりコンパクトであり、一般性と具体性のバランスが取れたものです。モデルがカーネルを書き込もうとするとき、そのプロンプトには信頼性の高いコンテキストのみを含めるべきです。具体的には、特定の状況とアクションを組み合わせた実行可能なテイクアウェイ（過去の「変更と影響のマッピング」から抽出されたもの）や、密接に関連する親実行からの簡潔な失敗メモなどです。階層的なタグフィルタリングとハイブリッド（キーワード＋セマンティック）検索を組み合わせて取得される教訓は、実行に移せるほど具体的であり、かつ適用されないケースが明確になるよう範囲が限定されている必要があります。教訓ストアの再編成やさらなる抽出といったより深い処理は、毎回の試行における過去の実行に対する同期的なマルチホップトラバーサルではなく、バックグラウンドジョブで行うべきです。テイクアウェイが状況とアクションを特定できないのであれば、プロンプトに含める価値はありません。図3は、ナレッジレイヤーを修正した後のトークンコストの内訳を示しており、この修正後はトークンの大部分が候補の生成に費やされています。

ケーススタディ：Gated DeltaNetのパックドデコード

具体的な例の1つは、Qwen 3.5 122BのGated DeltaNetパスにおけるパックドデコード（packed decode）カーネルです。この演算は、再帰的な状態を更新し、パックされたQKV入力、ゲートパラメータ、および状態インデックスからデコード出力を書き込みます。私たちはこのタスクを使用して、Tritonバックエンドを備えたNVIDIA B200 GPU上でProteusの全ループを実行しました。具体的には、タスクコントラクトの検証、リファレンス実装の測定、エージェントへの候補カーネルの要求、静的チェックとビルドの実行、管理されたリファレンスに対する正確性の検証、検証済み候補のみのベンチマーク測定、そして最適な候補の再測定です。

図4は左から右へと読みます。ベースラインノードはベンチマークを0.025ミリ秒に固定しています。候補0000は安全なシードです。パックドデコードの構造を再現し、検証をパスしましたが、リファレンスよりも低速であったため、Proteusはこれを勝利（最適解）として扱うのではなく、測定済みの親として保持しました。そこから、Proteusはすべてのシェイプに対して1つの汎用カーネルを最適化することをやめ、探索をシェイプ固有のパスに分割しました。

Batch-1修復パスは、候補012でシェイプ固有のカーネルを生成し、シングルバッチのデコードシェイプで1.5倍の高速化を達成しました。最も強力な結果は、サービングデコード（serving-decode）パスで得られました。候補030は、測定されたカーネルレイテンシの中で最低となる0.018ミリ秒を記録し、候補036は1.6倍という最高のシェイプ高速化を達成しました。その勝利したサービングフラグメントは、Batch=4、Key=128、Value=128のレイアウトに特化し、値の次元を幅64のチャンクで処理したため、普遍的な代替カーネルではなく、その特定のシェイプに対して安全なカーネルとなっています。

タイムラインの最後の回り道は、なぜトレースが重要なのかを示しています。その後の（Tritonではなく）C++生成の試みはビルドと生成の失敗に直面し、そのブランチの試行予算を使い果たした後に長期の実行が終了しました。したがって、有用な成果物は最も高速な候補だけではありません。それは図に示されている完全なパスです。つまり、セマンティクスの失敗は拒否され、正確だが低速なカーネルが測定され、実際のパフォーマンス向上は、プロダクションカーネルへの安全な構成を可能にする形状に関連付けられたまま維持されました。

次に取り組むこと

私たちが構築した進化ループは、ライターとして厳格です。これは多くの場合、「現在の最適なカーネルを取得し、小さな編集を試し、チェックし、繰り返す」という固定されたパターンでモデルを呼び出します。これにより、エージェントに必要な自律性が奪われてしまいます。構造を簡単に変更したり、言語を切り替えたり、行き詰まった設計を放棄したりすることができなくなります。

それでもループは必要です。カーネルを作成するためではなく、エージェントに信頼できる次のヒントを与えるためです。そのヒントは2つの場所から得られる必要があります。

1つ目は、ナレッジレイヤーとのコミュニケーションです。行動を起こすのに十分なほど具体的で、適用されない状況がわかるように範囲が限定された、いくつかの要点です。これがなければ、すべての試みはゼロからのスタートになります。

2つ目は、信頼できるチェッカーからの結果です。これには、エージェント自身が測定しなかった正確性とタイミングが含まれます。これらの数値は、次の試行へのヒントになります。また、これらは私たちが信頼すべき唯一のスコアでもあります。エージェントが自身の作業を測定すると、キャッシュの残り、一致しない比較、エージェントが見ることができるテストに逆戻りしてしまいます。

したがって、私たちが望む役割分担は、「エージェント対ループ」という単純なものよりも限定的です。カーネルの記述方法についてはエージェントに自律性を与えます。ループはメモリと評価のためのチャネルとして維持します。エージェントが提案し、ループは表示が許可されている内容と、最後の提案が実際に改善されたかどうかを返します。

Proteusは、NVIDIA B200 GPU上で動作するQwen 3.5 122BのGated DeltaNetパス（線形アテンションスタイルのブロック）の一部に対して、専用のカーネルを構築しました。個々のカーネルの速度向上は1.8倍から5.2倍の範囲でした。

教訓は、生成はコストの低いステップであるということです。検証とコンテキスト管理こそが困難な部分です。そこに、慎重な設計時間とイノベーションが必要とされます。

エージェントによるGPUカーネル生成は、極限の専門化という驚くべき可能性を切り開きましたが、信頼性が高く本番環境に対応したハーネスを構築することは、依然として困難なフロンティアです。私たちはAIとシステムの交差点における最も困難な課題に取り組んでおり、効率的な推論の未来を共に形作る大胆なエンジニアを募集しています。可能性の限界に挑むことに情熱をお持ちなら、採用情報をご覧ください！

(このブログ記事はAI翻訳ツールを使用して翻訳されています) 原文記事