KernelArc: A Multi-Agent Framework for GPU Kernel Optimization
KernelArcは、共有メモリと決定論的なガードを介して連携する、戦略に特化した並列エージェントを通じてGPUカーネルを自律的に最適化するマルチエージェントフレームワークであり、NVIDIA H100およびB200 GPUにおける多様なワークロードにおいてSOL-ExecBenchリーダーボードでトップのランキングを獲得しています。
原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。 免責事項の全文を読む
現代のコンピュータは、人工知能に必要な膨大な計算を処理するために、GPU（グラフィックス・プロセッシング・ユニット）と呼ばれる専用のチップに依存しています。これらのチップは非常に強力ですが、同時に多くの異なる部品が完璧に同期して動作しなければならない複雑な機械でもあります。その性能を最大限に引き出すために、エンジニアは「カーネル」と呼ばれる、チップに対してデータの移動方法や数学的演算を正確に指示する、極めて限定的で高度なプログラムを書かなければなりません。長年、これは人間の専門家による仕事であり、彼らはコードを微調整し、メモリ使用量を調整し、タイミングを計ることで、速度を極限まで引き出すために膨大な時間を費やしてきました。チップがより洗練されるにつれ、この手作業はより困難になり、ハードウェアができることとソフトウェアが達成していることの間の溝は広がっています。
研究者たちは、このコードの作成や改善を支援するために、膨大な量のテキストで学習されたコンピュータプログラムである大規模言語モデルを使用し始めました。これらのモデルは、変更を提案したり、テストを行い、その結果から学習したりすることができ、自動化されたエンジニアのように機能します。しかし、単一の自動エージェントでは行き詰まってしまうことがよくあります。あるエージェントは優れた解決策を見つけると、それを磨き上げ続けることがありますが、別の方向にあるより優れたアプローチを見逃してしまうことがあります。それは、山の上にある素敵な道を見つけ、その道を登り続けるあまり、近くにある別の道がもっと高い頂上へと続いていることに気づかないハイカーのようなものです。これを解決するために、IMECの研究チームは、複数の自動エージェントが協力して一度に多くの経路を探索する「KernelArc」と呼ばれる新しいシステムを開発しました。
KernelArcシステムは、複数のエージェントに対して並行して異なる戦略を割り当てることで動作します。一つのエージェントが単独で問題を解決しようとするのではなく、グループ内の各エージェントは、データの保存方法の変更、使用される数学的精度の変更、あるいは複数のステップの融合といった、異なる攻撃の角度に焦点を当てます。これらのエージェントは、プロセス全体やすべてのステップを共有することはありません。それを行うと、あまりにも煩雑で混乱を招くからです。代わりに、彼らは最終的な結論、つまり何が機能し、何が失敗したかのみを共有します。彼らはこれらの結果を、掲示板のような役割を果たす共有メモリ空間に書き込みます。もし一つのエージェントが計算を高速化するコツを発見すれば、その結果を投稿します。他のエージェントはその投稿を読み、その洞察を利用して自身の作業を導き、行き止まりを回避し、成功の上に積み上げていきます。
エージェントが壊れたコードで時間を無駄にしないように、システムには厳格な自動レフェリーが含まれています。このレフェリーは、すべての新しいコードの提案を一連のテストに通し、それが正しいかどうか、そしてどれほど速いかをチェックします。もし提案がテストに失敗すれば、即座に破棄されます。もし正しくても、現在のベストバージョンよりも速くなければ、バックアップとして保持されますが、リーダー（最良版）を置き換えることはありません。新しいバージョンが正しく、かつ高速である場合にのみ、それが新しい標準となります。このプロセスにより、エージェントのチームは個々の試行の細部に迷い込むことなく、幅広い可能性を探索することができます。また、システムには、エージェントが行き詰まり（さらなる改善が見出せないプラトー状態）に陥った場合に作動するセーフティメカニズムも備わっており、この場合、システムはエージェントに全く異なるアプローチを試みるよう強制し、探索が前進し続けることを保証します。
研究者たちは、利用可能な最も先進的なGPUであるNVIDIA H100とB200を用いてこのシステムをテストしました。彼らは、行列乗算、言語モデルで使用されるアテンション・メカニズム、および様々な形式のデータ融合を含む、パフォーマンスを測定するための標準的なタスクに焦点を当てました。特定のタスクにおいて、詳細なガイドと共に単一のエージェントが作業した特定のテストでは、システムは766テラフロップスの速度に達しました。これは、その特定のタスクにおける既存の最高のライブラリよりも約3.2パーセント高速です。これは、明確な経路が与えられれば、単一のエージェントが特定の課題に対して非常に深く掘り下げることができることを示しました。しかし、研究者がより広範なタスクに取り組むためにマルチエージェントのKernelArcシステムに切り替えると、結果はさらに印象的なものとなりました。システムは、特殊なアテンション・メカニズムや大規模言語モデル用の融合レイヤーを含む、さまざまな複雑な操作のカスタム実装を生成しました。
データの形状やサイズに応じて操作の速度をランク付けする公開リーダーボードにおいて、KernelArcの提出物はいくつかのカテゴリーで第1位を獲得しました。アテンションと残差加算を含む一つのタスクでは、単一のエージェントは0.441のパフォーマンススコアで停滞しました。しかし、マルチエージェントシステムは、洞察を共有し異なる方向を探索することで、その障壁を打ち破り、0.481のスコアに到達しました。別の複雑なアテンション・タスクを含むテストでは、マルチエージェントシステムは標準的なリファレンス実装と比較して291倍近く、高度に最適化されたベースラインと比較して143倍以上の高速化を達成しました。これらの結果は、複数のエージェントがそれぞれの知見を共有することで、特に問題の領域が広く多様である場合、単一のエージェントが単独で作業する場合よりも優れた解決策をより速く見つけられることを示唆しています。
この研究では、エージェントの数と共有メモリの量が変化した場合にシステムがどのように振る舞うかについても調査しました。彼らは、エージェントが成功と失敗を投稿できる共有メモリを与えることが、一定の試行回数内でより強力な結果に到達するのに役立つことを発見しました。メモリが無制限であった場合、システムは最も高い性能を発揮し、開始時点から290倍以上の高速化を達成しました。これは、集団の経験から学ぶ能力が成功の鍵であることを示しています。研究者たちは、共有メモリや戦略の特化といった各機能の価値は、特定のタスクや探索の段階に依存することを指摘しました。時にはエージェントは広く探索する必要があり、またある時には特定の解決策を洗練させる必要があります。
この研究は、コンピューティングの最適化におけるすべての問題を解決したと主張するものではなく、また人間のエンジニアがもはや不要であることを示唆するものでもありません。結果はテストされたタスクとハードウェアに特有のものであり、システムは依然としてエージェントを導くための人間が設計したフレームワークに依存しています。しかし、得られた知見は、調整されたエージェントのグループが、単一のエージェントよりも広い範囲の解決策を探索できることを実証しています。最も価値のある結論のみを共有することで、これらのエージェントは間違いを繰り返すことを避け、互いの発見の上に積み上げていくことができます。このアプローチは、ハードウェアの潜在能力とソフトウェアのパフォーマンスの間の溝が広がり続けている現代のコンピューティングの複雑さに対処するための、有望な方法を提供します。このシステムは、適切な調整があれば、自動化されたツールが次世代の人工知能を駆動するチップの全力を引き出す助けとなることを示しています。
自分の分野の論文に埋もれていませんか？
研究キーワードに一致する最新の論文のダイジェストを毎日受け取りましょう——技術要約付き、あなたの言語で。