ブラウザAIの最終ピースとなり得る「カーネル最適化」

ブラウザ上でAIモデルを実行した際、GPUを十分に活用できず、パフォーマンスが思うように出ないという課題は長らく存在していました。この問題に対し、Hugging Faceが重要な一手を打ち出しました。207個の最適化されたWebGPUカーネル群の公開です。

これは単なるカーネルの羅列ではなく、ブラウザ上でのAIモデル実行の性能を劇的に改善し得る基盤の構築です。本稿では、このカーネルコレクションの特徴と具体的な活用法を、元となった技術ブログの内容を基に詳細に分析します。

なぜ「カーネル」に焦点を当てたのか

ブラウザで実行されるAIモデルは、結局のところ、行列積、正規化、畳み込みなどの一連のGPU演算の連続です。WebGPUはこれらの演算をブラウザ上で実行するための標準APIであり、WGSLはGPU上で実行されるシェーダーを記述するための言語です。

しかし、ポータビリティ(移植性)がそのままパフォーマンス(性能)を保証するわけではありません。同じ演算を実装しても、GPUアーキテクチャによって性能は大きく変わり得ます。ワークグループサイズ、メモリアクセスパターン、ベクトル化、データ型、融合戦略など、性能に影響を与える要素は無数にあります。

Hugging Faceがカーネルから最適化を始めたのは、上位レベルのランタイムの効率性が、最終的にはディスパッチされる個々の演算の効率性に依存するためです。カーネルを個別に発見・テスト・ベンチマーク・バージョン管理できるようにすることで、上位レイヤーへの安定したコントラクト(契約)を維持しつつ、基盤となる部分を独立して改善できます。これは、建物の基礎工事をしっかりと行うことに似ています。

単なるシェーダーではなく「カーネルリポジトリ」という概念

今回公開された各カーネルは、単一のWGSLシェーダーファイルではありません。各カーネルは独自のリポジトリと詳細なドキュメント(カード)を持ち、そこには演算の意味、入力/出力、属性、サポートされるデータ型、ソースファイル、そしてすぐに実行可能な@huggingface/kernelsの例が含まれています。

例えば、ai.onnx.Addカーネルは、多次元ブロードキャストをサポートする要素ごとの加算演算です。これは、残差接続やバイアス加算など、ニューラルネットワークで最も基本的な演算の1つです。カーネルカードには、2つの入力、ブロードキャストされた出力形状、サポートされるデータ型、そして異なる形状やデバイス向けに利用可能なバリアントが文書化されています。

リポジトリ構造の詳細

カーネルリポジトリは以下のような構造で構成されており、実装を理解し評価するために必要な情報をすべて提供します。

  • manifest.json: 演算コントラクトの主要情報です。入力、出力、属性、型制約、形状導出ルールを定義します。
  • metadata.json: カーネルID、ダイジェスト、来歴情報を記録します。
  • test.json: 実装が期待される動作と一致するかを確認するための正しさのテストケースを含みます。
  • bench.json: カーネル評価に使用されるベンチマークおよびチューニングケースを含みます。
  • *.wgsl.jinja: 特定のリクエストとデバイス向けのシェーダーを生成するために使用される、パラメータ化されたWGSL実装ファイルです。

この構造により、単なるシェーダーが再利用可能なソフトウェア成果物へと昇華します。インターフェースはWGSLを読まなくても把握でき、正しさとパフォーマンスのテストケースは実装とともに提供され、公開されたバージョンはバージョン管理されていないファイルURLに依存せずに明示的にロードできます。

主要コード: Hubからのカーネル読み込みと実行

それでは、実際の使用方法をコードで確認しましょう。npmパッケージのインストール、WebGPUサポートの確認、カーネルの読み込みと実行のプロセスです。

// 1. npmパッケージインストール
// npm install @huggingface/kernels@preview

// WebGPUサポートの確認
if ("gpu" in navigator) {
  console.log("WebGPUをサポートしています。");
} else {
  console.log("WebGPUをサポートしていません。");
}

// 2. @huggingface/kernelsライブラリからgetKernel関数をインポート
import { getKernel } from "@huggingface/kernels";

// 3. Hubからカーネルをロード(リポジトリIDとバージョンを指定)
const add = await getKernel("webgpu-kernels/ai.onnx.Add", { version: 1 });

// 4. 型指定された入力データとテンソル形状を引数にカーネルを実行
const { c } = await add({
  a: {
    data: new Float32Array([1, 2, 3, 4, 5, 6]),
    shape: [2, 3],
  },
  b: {
    data: new Float32Array([10, 20, 30]),
    shape: [3],
  },
});

console.log("結果:", c.data); // Float32Array([11, 22, 33, 14, 25, 36])

2番目の入力bが最初の次元に渡ってブロードキャストされ、[2, 3]の形状の出力が生成されます。ローダーは、マニフェストコントラクトと入力から出力形状と論理データ型を導出し、cを自動的に割り当てます。この6つの浮動小数点の加算は意図的に最小のデモです。このサイズでは、GPUの往復コストが計算コストよりもはるかに大きくなります。重要なのは、呼び出しパターンが行列積(ai.onnx.MatMul)のような重い演算でも同じであることです。リポジトリIDと入力が変わるだけです。

パフォーマンスはどれくらい向上したのか

Hugging Faceは独自のベンチマーク結果を公開しています。Apple M4 GPU上でONNX Runtime Web(ORT WebGPU)と比較したところ、独自カーネルコレクションは幾何平均で2.57倍、中央値で1.90倍の高速化を達成しました。比較した全809ケースのうち、629ケースで勝利しました。

主要演算別のパフォーマンス比較

演算比較ケースHF WebGPUカーネルORT WebGPU速度向上
Add50.064 ms0.227 ms3.52倍
MatMul290.115 ms0.131 ms1.14倍
Softmax120.114 ms0.240 ms2.11倍
LayerNormalization60.061 ms0.135 ms2.22倍

一部の個別ケースでは、さらに劇的な結果が出ています。特に、サイズ4096の双線形Einsum演算は0.136msで実行され、ORT WebGPUの1,396msと比較して10,000倍以上高速でした。これは、一般的な実装が遅いパスに直面したとき、特殊化されたカーネルがどれほどの利点を提供できるかを示す好例です。ただし、これは極端なケースであり、すべてのアプリケーションで期待できるわけではありません。

Fleet: コミュニティベースのベンチマーキングの始まり

WebGPUのパフォーマンスは、GPU、ブラウザ、ドライバによって大きく異なります。したがって、1台のマシンでの結果は全体像を語るものではありません。Hugging Faceはこの問題を解決するために、Fleetと呼ばれるブラウザベースのベンチマーキングおよびテストスイートも公開しました。

Fleetを使用すると、誰でもブラウザ上で正しさとパフォーマンスのチェックを実行し、自分のハードウェア上でカーネルがどのように動作するかを確認できます。ユーザーの同意を得て、各実行はデバイス固有の障害を発見し、バリアントを比較し、選択ルールを改善するのに役立つ証拠を非公開で提供します。究極の目標は、広範で現実世界のカバレッジを利用して、すべてのユーザーにとってカーネルをより速く、より信頼性の高いものにすることです。

現実的な適用、限界、そして今後の方向性

今回のHugging Faceのアプローチは、ブラウザAIエコシステムに重要なマイルストーンを打ち立てました。カーネルを独立して公開し、明確なコントラクトと再現可能な証拠を提供することで、コミュニティが協力して基盤を一緒に改善する道を開きました。

ただし、いくつかの限界も存在します。

  • 測定範囲: ベンチマークはGPUで実行されるジョブの時間のみを測定しており、カーネルのロード、セッション作成、シェーダーコンパイルなどのオーバーヘッドは除外されています。実際のアプリケーションでは、これらの要素が大きな割合を占める可能性があります。
  • 個々の演算結果: 今回の結果は個々の演算に関するものであり、モデル全体のパフォーマンスはGPUやブラウザによって異なる可能性があります。
  • WebGPUサポート環境: WebGPUはまだすべてのブラウザやデバイスでサポートされているわけではありません。サポート範囲が広がるにつれて、このライブラリの有用性はさらに高まるでしょう。

しかし、それでもこの技術がもたらす変化は明白です。特に、複雑なAIモデルをサーバーなしでクライアント上でのみ実行しようとする試みが増えている中で、WebGPUカーネルの最適化は必須の要素となるでしょう。

AIエージェントの信頼性との接点

興味深いことに、今回のカーネル最適化の取り組みは、AIエージェントの信頼性を確保することとも関連しています。AIモデルがブラウザ上で高速かつ安定して実行されるほど、ユーザーにより良い体験を提供し、より複雑なタスクを安全に処理するための基盤が整います。AIエージェントの性能と安定性は、結局のところ、基盤となるインフラストラクチャの効率に大きく依存するからです。

また、今回公開されたアプローチは、大規模なAIシステムを設計する際の重要な原則を示しています。個々のコンポーネントを独立してテストし、バージョン管理し、パフォーマンスを測定することは、ブラウザAIに限らず、堅牢なAIシステム開発において核心的な要素です。

まとめ: ブラウザAIの新たな地平線

Hugging Faceが公開した207個のWebGPUカーネルとFleetツールは、ブラウザAIの状況を一変させる可能性を秘めています。これは単なるパフォーマンス向上を超え、コミュニティが一緒にパフォーマンスと正確性の証拠を蓄積し、改善していく新しい開発パラダイムを提示しています。

次のステップとしては、WebGPUカーネルコレクションを直接探索し、自分のハードウェアでFleetを実行して結果を確認することをお勧めします。また、@huggingface/kernelsライブラリを使用して簡単なプロジェクトに適用し、ブラウザAIの可能性を実感してみてください。これは、単なる技術紹介を超え、より高速で民主化されたAI環境のための基盤となるでしょう。

Hugging Face WebGPU kernel collection displayed on a laptop screen for local AI inference System Abstract Visual

本コンテンツは、信頼性の高い情報源をもとにAIツールを活用して作成され、編集者によるレビューを経て公開されています。専門家によるアドバイスの代替となるものではありません。