Hugging Face、ブラウザーAI向けWebGPUカーネル207個を公開

ブラウザーでAIモデルを動かすとき、課題はモデルを配布できる形式に変換することだけではない。最終的な推論は、行列乗算、正規化、畳み込み、アテンション、量子化、データ配置変換など、多数のGPU演算の組み合わせになる。個々の演算がデバイスに合っていなければ、ブラウザー対応のモデルでも十分な応答速度は得られない。Hugging Faceが公開した@huggingface/kernelsは、この低レイヤーを整備する取り組みだ。

発表のポイント

• 初期コレクションは207個のWebGPUカーネル。webgpu-kernels組織の個別リポジトリで公開され、Apache-2.0ライセンスを採用する。
• 各カーネルはWGSLシェーダー単体ではなく、演算の契約、対応する型と形状、正確性テスト、ベンチマークケース、メタデータ、パラメーター化されたWGSLテンプレートを含む。
• JavaScriptパッケージはHubのリポジトリIDと契約バージョンを指定してカーネルを読み込み、型付きデータとテンソル形状を渡して実行できる。
• Fleetはブラウザー上でGPUをテスト・ベンチマークする仕組み。利用者が同意すれば、実際の端末から性能と正確性に関する証拠を集められる。

この構成の意味は、シェーダーを再利用可能なソフトウェア部品として扱えることにある。たとえばai.onnx.Addは、ブロードキャストを含む演算仕様だけでなく、同じ形状、ベクトル化されたブロードキャスト、スカラー処理、一般的なブロードキャストに対応する複数の実装経路を持つ。ランタイムは入力形状やデバイスに応じて経路を選択でき、アプリケーション側の呼び出し方は変えずに済む。

契約のバージョンは、ONNX opset、演算子のsince_version、モデルのリビジョンとは別に管理される。これにより、上位ランタイムは安定したJavaScript向けインターフェースに依存しながら、下位実装を更新できる。リポジトリにテストとベンチマークが同梱されるため、独自のWebGPUカーネルを作る開発者にとっても参考実装になり得る。

性能を見る際の注意点

Hugging FaceはApple M4 GPU上で、開発版のONNX Runtime Webに含まれるORT WebGPUと比較した。207演算、1,756ケースから始め、両方の出力が一致し、計測も信頼できる809ケースを残した。その範囲では、幾何平均で2.57倍、中央値で1.90倍の高速化を報告している。

ただし、これはすべての環境にそのまま適用できる数字ではない。WebGPUは共通APIを提供する一方、ワークグループサイズ、メモリアクセス、ベクトル化、データ型、演算融合、ブラウザー、ドライバー、GPU構成によって最適解は変わる。小さな入力ではGPUへの送受信コストが計算量を上回る場合もある。

意義

カーネルを上位ランタイムから切り分けることで、APIを維持したまま個別演算を継続的に改善できる。Fleetは実験室で網羅しにくい実機情報を集めるフィードバックループになるが、同意、サンプルの偏り、異なる端末間の比較方法には注意が必要だ。今回の発表は新しいモデルではなく、ブラウザー上のローカル推論を支えるGPU基盤を、発見しやすく、検証しやすく、バージョン管理しやすくするためのインフラ整備と位置付けられる。

Hugging Face Blog