GPUカーネル自動生成を促すHawkeye、専門家に頼らず高性能を実現

発表内容と開発体制

米ハーバード大学、スタンフォード大学、Together AI、カリフォルニア工科大学の研究者らは、AIエージェントがGPUカーネルを効率的に記述できるようにするためのオープンソースフレームワーク「Hawkeye」を発表した。特定のGPUハードウェア向けに最適化されたカーネルを、最小限の人間の監督で生成することを目指している。

「最小限かつ包括的な分類体系」に基づき、自律的なカーネル生成を実現する点が特徴だ。研究チームは、このフレームワークにより、新興のハードウェアアクセラレータを支援する際のオーバーヘッドを削減できるとしている。

仕組み: 単位テストの分類体系

Hawkeyeの核となるのは、ハードウェアプラットフォームと最適化戦略を単位テストとしてパッケージ化した分類体系だ。各単位テストは、人間が作成した解決カーネルと、最適化を検証するプロファイリング指標を最小限の組み合わせとして提供する。

エージェントは、このカーネルを構文例として読み、直接呼び出し、または断片を組み合わせてより大きなカーネルを構成することができる。これにより、アーキテクチャ固有の機能を最大限に活用したコード生成が可能になる。

性能評価の結果

研究チームは、PyTorchワークロードをNVIDIA Ampere、Hopper、Blackwell、AMD MI350向けの高性能カーネルに移植して評価した。精度はBF16、FP8、NVFP4、MXFP4の各形式で検証している。

確立済みのワークロードでは、torch.compileがcuBLASやcuDNNなどの専門家向けライブラリにディスパッチする場合と比較して、Hawkeyeは同等以上の性能を示した。新興のアテンション変種では、Flash Linear Attentionライブラリの専門家作成Tritonカーネルに対して相乗平均で18.9倍の高速化を達成し、Blackwellでは1.22倍、MI350では1.00倍という結果だった。

意義と注意点

Hawkeyeの成果は、人間が厳選した知識を少し与えるだけで、AIシステムが高度に最適化された人間の作業に追いつき、上回ることができることを示している。GPUカーネルのような複雑なチューニング作業が自動化される可能性がある。

ただし、この手法は単位テストの分類体系を事前に人間が整備することが前提だ。また、カーネル生成の性能は、ベースとなるモデルの能力に依存すると考えられる。現時点では、評価対象外のGPUや新しいアーキテクチャへの汎用性は不明である。