Hugging Face、ブラウザで動くAI向けWebGPUカーネル207本を公開

WebGPUカーネル207種を公開

Hugging Faceは、WebGPUを利用してブラウザ上でAIモデルの推論処理を高速化するためのカーネル群を公開しました。ここで言うカーネルとは、ニューラルネットワークに現れる行列積や畳み込みなどの基本演算をGPU上で実行するためのプログラムです。公開されたのは207個のカーネルで、Hugging Face Hubの「webgpu-kernels」グループで個別のリポジトリとして管理されます。ライセンスはApache-2.0です。

あわせて、カーネルを読み込むためのJavaScriptライブラリ「@huggingface/kernels」と、ブラウザ上でカーネルの正しさや性能をテストするスイート「Fleet」も発表されました。これらは、ブラウザベースのAIエコシステムを支える低レイヤーの土台を目指しています。

単なるシェーダーではない設計

今回の特徴は、各カーネルが単体のシェーダーコードではなく、契約情報を含む完成されたソフトウェア成果物として公開されている点です。各リポジトリには、入力・出力・属性などを定義する「manifest.json」、カーネル識別子と由来を記録する「metadata.json」、正しさを検証する「test.json」、ベンチマーク用の「bench.json」、パラメータ付きのWGSLシェーダーテンプレートが含まれます。これにより、WGSLを直接読まなくてもインターフェースを把握でき、正しさの検証と性能測定を実装と一緒に共有できます。

実行方法もシンプルです。npmから「@huggingface/kernels」をインストールし、getKernel関数にリポジトリIDとバージョンを渡せば、配列とテンソル形状を渡してカーネルを呼び出せます。たとえば「ai.onnx.Add」は、加算とブロードキャストの出力形状を自動で計算し、結果のバッファを割り当てます。バージョンはアプリケーションが依存できる安定した契約として管理される設計です。

ORT WebGPUとの性能比較

同社はApple M4 GPU上のONNX Runtime Web（ORT WebGPU）と比較した結果を公表しています。1,756件のテストケースのうち、両方の実装で出力が一致し、時間が計測できた809件では、Hugging Faceカーネルは幾何平均で2.57倍、中央値で1.90倍高速でした。勝敗は629勝176敗4引き分けです。

個別のケースでは、極端な差も見られます。特定条件のEinsum（テンソルの縮約演算）では、Hugging Faceカーネルが0.136ミリ秒だったのに対し、ORT WebGPUは1,396ミリ秒と、1万倍以上の差が出ました。CumSum（累積和）でも約301倍の差がつきました。これは汎用実装が低速な経路に入った場合の改善例であり、あらゆる場面でこの速度差が期待できるわけではありません。

なお、計測はGPU上での処理時間のみで、シェーダーのコンパイルやデータ転送は含まれません。また単一の演算レベルの比較であり、モデル全体の推論時間を保証するものではない点に注意が必要です。性能はGPUやブラウザによって異なるため、Fleetによる実機データの収集が重要だとしています。

Fleetと今後の展開

WebGPUの性能は、GPU、ブラウザ、ドライバの組み合わせで大きく変わります。Fleetはブラウザ上でカーネルのテストとベンチマークを実行し、ユーザー自身のハードウェアでの結果を確認できるようにするツールです。ユーザーが同意すれば、実行結果は送信され、デバイス固有の障害や低速ケースの発見、カーネルの改良に使われます。同社はこれを、従来のテストラボではカバーできない幅広い実機を利用する仕組みと位置づけています。

今回のリリースは出発点であり、今後はカバレッジを増やしつつ、上位のモデルツールとの連携も予定されています。また、ONNX Runtimeチームとも協力し、改善をONNX Runtime Webエコシステムに還元する方針も示されました。