WebAssemblyとMetalでゼロコピーGPU推論をApple Siliconに実装する
目次
Abacus NoirのブログでDriftwoodという推論ランタイムの設計が公開された。
WebAssemblyモジュールとApple SiliconのGPUが同じ物理メモリを読み書きするゼロコピー推論チェーンの実装解説で、mmap・MetalのbytesNoCopyバッファ・Wasmtimeのカスタムアロケータを連結して「CPUとGPUが同一の物理バイトを読み書きする」状態を実測レベルで証明している。
地味だが、Apple Silicon上でWasmを制御プレーン、GPUを計算プレーンとして扱うときの基本構造にあたる話なので、仕組みを追ってみる。
UMAとディスクリートGPUでゼロコピーが違う意味になる
ディスクリートGPU(PCIe接続の外付けGPU)ではホスト側メモリとGPUメモリが物理的に分離されているため、推論時には必ずPCIeバス越しのデータ転送が発生する。
行列や埋め込みベクトルをGPUに渡すたびにコピーが入り、その分遅延とメモリ消費が増える。
Apple SiliconはUMA(Unified Memory Architecture、統合メモリアーキテクチャ)を採用していて、CPUとGPUが同一のDRAMチップを共有する。
理論的にはコピー不要になるはずだが、Metalで普通にmakeBuffer(length:options:)を呼ぶと内部でMetal管理のメモリ領域が別途確保される。
結果、ホスト側で用意したバッファからGPUバッファへコピーする手間がかかってしまう。
Driftwoodはここを突き、Wasmランタイムに渡す線形メモリと、MetalがGPUから触るバッファを、物理アドレスレベルで一致させる。
ポインタが同一であれば、ホスト関数呼び出しで引数をGPUに渡すオーバーヘッドが消える。
これは過去に解説したOllamaのMLXバックエンド移行やSwiftLMのMetalネイティブ実装のような、Apple Silicon専用最適化の延長線上にある。
違いはランタイム層にWasmを挟み、可搬な制御プレーンとして機能させる点だ。
3段チェーンの構造
実装は3本のリンクで構成される。
flowchart TD
A[Wasmモジュール<br/>linear memory] -->|同一ポインタ| B[mmap領域<br/>16KB alignedなページ]
B -->|同一ポインタ| C[MTLBuffer<br/>bytesNoCopy wrap]
C -->|GPU kernel実行| D[Metal GPU<br/>同一物理バイトを読み書き]
E[Wasmtime<br/>MemoryCreator trait] -->|linear memoryを差し替え| A
F[LLMカーネル<br/>GEMM等] -->|compute pipeline| D
Link 1: mmapでページアラインドな領域を確保
ARM64 macOSではmmap(MAP_ANON | MAP_PRIVATE)を使うと16KBアラインドなアドレスが返る。
この16KBはARM64のページサイズそのもので、mmapはページ境界にアラインして返すことがPOSIXの契約で保証される。
たまたまアラインされているのではなく、仕様上必ずアラインされる。
なぜアラインが要るかというと、次のステップで使うMetalのnewBufferWithBytesNoCopyがページアラインされたポインタしか受け付けないから。
mallocで確保したメモリは任意のバイト境界にあるため、この関数には渡せない。
必ずvm_allocateかmmapを使う必要がある。
Link 2: Metal bytesNoCopyで既存ポインタをラップ
MTLDevice.makeBuffer(bytesNoCopy:length:options:deallocator:)は、既存のホストポインタをコピーせずにそのままMTLBufferとしてラップするAPIだ。
内部でMetalが別領域を確保することはなく、GPUからのアクセスは渡したポインタが指す物理メモリ上で行われる。
DriftwoodはMTLBuffer.contents()で返されるポインタがmmapで取得した元ポインタと一致することを確認している。
これは「ラップ時にMetalが影でコピーしていない」ことを実測で示すための検証。
実装上は単純な==比較でよく、RSS(Resident Set Size、プロセス実使用物理メモリ量)の増分を計測することでも裏取りできる。
実測値として、16MBのバッファをbytesNoCopyでラップした場合のRSS増分が0.03MB(ノイズ相当)に対し、コピー経路を使うと16.78MB増える。
コピー経路では物理メモリが倍消費され、ホストで書いた内容をGPUに反映するたびに転送が要る。
ゼロコピー経路は数字上も明確な差がつく。
Link 3: Wasmtime MemoryCreatorで線形メモリを差し替える
WasmtimeにはMemoryCreatorというトレイト(Rustのインターフェース相当)があり、WebAssemblyの線形メモリ(モジュールがmemory.data_ptr()で参照するあのメモリ)をランタイム側で差し替えられる。
MemoryCreatorとLinearMemoryのペアを実装すると、Wasmtimeは内部でmmapを呼ばずに実装者が渡したメモリを使うようになる。
ここで先ほどのmmap領域を線形メモリの裏付けとして渡すと、Wasmモジュールが書き込んだバイトが即座にMetalから見える。
memory.data_ptr()で取得できるポインタは、最初にmmapで取ったアドレスそのもの。
Wasm・CPU・GPUが同一の物理バイトを指すという状態が成立する。
ただしMemoryCreatorには注意点がある。
guard_size_in_bytesというパラメータで、確保領域の末尾に何バイト分のガードページ(アクセスするとSEGVになる領域)を置くかを指定する。
WasmtimeのJITコードはこのガードサイズを前提に境界チェックを省略するので、ガード領域を確保せずに渡すとJITが不正メモリにアクセスしかねない。
実装側でmmap時に本体+ガードのサイズで確保し、ガード部分をmprotectでPROT_NONEにするのが定石。
パフォーマンス実測
Abacus Noir側が公開している数字を整理する。
| 項目 | 値 |
|---|---|
| GEMM(128x128)レイテンシ | 約6.75ms |
| ポインタ同一性 | mmap == MTLBuffer.contents() |
| 16MB確保時のRSS増分 | 0.03MB(ゼロコピー経路) |
| 正解性(16,384要素) | エラー0件 |
モデル推論の数字もある。
Llama 3.2 1BをM1 MacBook Proで動かすと、モデルロード229ms、5トークンのプリフィル106ms、トークンあたり生成約9ms。
ホスト関数呼び出し(WasmからMetalカーネルを起動する経路)のオーバーヘッドは「無視できる」レベルとされる。
1BクラスのモデルをWasm制御下で9ms/tokenで回せるのは、SwiftLMのようなネイティブSwift実装に遜色ない水準。
Wasmを挟んだことによる性能ペナルティが事実上ないことを示している。
KVキャッシュのsafetensors化
Driftwoodはこのゼロコピー基盤の上に、アクタースナップショット(会話状態の凍結・復元)を載せている。
KVキャッシュのシリアライズにsafetensors形式を使うのが設計の肝で、24トークン分(1.58MB)のシリアライズが1.1ms、復元が1.4ms。
注目は「再プリフィル比で5.45倍高速」という数字。
会話を復元するのに、すべてのトークンを頭から再計算する代わりにKVキャッシュを読み戻す。
復元後のトークン生成は10/10でbit-identicalになる。
長い文脈を持つエージェントの瞬間復帰や、マシン間の状態転送を成立させるために必要な精度だ。
safetensorsは機械学習界で一般化したテンソルシリアライズフォーマットで、pickleのような任意コード実行リスクがない。
KVキャッシュという実行時状態に同じフォーマットを使うのは合理的な選択。
何が嬉しいのか
この設計が成立すると、運用面・性能面・配布面それぞれに実利が出てくる。
可搬性の面から見ると、Wasmはサンドボックス化された実行環境なので異なるマシン・OSに配布しやすい。
推論ロジックをWasmで書いておけば、Apple Siliconマシン間での配布・実行を標準化できる。
オーバーヘッドの話では、通常Wasm/JIT/GPUを連携させるとシリアライズ境界でコストが出る。
UMA+ゼロコピーでこの境界が消えるため、Wasmを挟むコストが実質ゼロになる。
状態管理としても、KVキャッシュをsafetensors化してスナップショット可能にしておくことで、会話エージェントを一時停止・再開したり別マシンに移したりできる。
エージェント型アプリケーションの運用に効く観点だ。
一方、ディスクリートGPU(NVIDIA H100、AMD MI300等)ではbytesNoCopy相当のAPIはあるものの、物理メモリは本当に分離されているため「同一バイトを読み書き」にはならない。
UMA前提の設計で、他プラットフォームにそのまま移植できる話ではない。
実装上のハマりどころ
元記事と周辺ドキュメントから読み取れる注意点をまとめる。
malloc系ではbytesNoCopyに渡せない。
Apple公式のフォーラムでもvm_allocateかmmapを使う必要があると明記されている。
malloc実装は任意アラインで返すので、ページ境界にならない。
Wasmtimeのガードページ設計をスキップしない。
ガードを張らないとJITコードが境界外アクセスで予期せぬ挙動を示す。
mmapで余分に確保して末尾をmprotectする定型処理が必要になる。
MemoryCreatorのty引数はminimum/maximumをWasmページ単位(64KB)で受け取る。
ARM64のOSページ(16KB)と混同しない。4倍差がある。
デアロケータの責任も忘れない。
bytesNoCopyはdeallocatorクロージャを受けるので、MTLBufferがreleaseされたタイミングでmunmapを呼ぶ設計にする。
双方の所有権が交差するため、どちらが開放責任を持つか明示する必要がある。
Apple Silicon最適化の系譜
ゼロコピーチェーンの話自体は地味だが、AI推論スタックの最適化という文脈で見ると位置づけが見えてくる。
Ollama 0.19がMLXバックエンドに移行してApple SiliconのUMAを活かす方向に舵を切り、Flash-MoEが397BパラメータをMacBookで動かすためにMetalシェーダーとSSDエキスパートストリーミングを組み合わせ、SwiftLMがTurboQuant+SSDストリーミング+Metalをネイティブ統合する流れで、各レイヤーのオーバーヘッドを削っていく動きが続いている。
Driftwoodはこの流れの上で「Wasmを制御プレーンに挟んでも性能劣化しない」ことを示した。
可搬性と性能の両立という、普段はトレードオフになりがちな軸で、UMAというハードウェア特性を使って両方取りに行った設計。
HuggingFaceのtransformers-to-mlx PRのようなPythonエコシステム側のApple MPS対応とは別軸で、システムプログラミングレベルから推論ランタイムを再構築するアプローチが出てきている。
どちらも最終的にApple Siliconの物理特性を活かす方向に収束していくはず。
Driftwoodランタイム本体はまだ開発中でOSS公開されていないが、設計ドキュメントとして公開された今回の内容は、同種のシステムを自作するときの手順書として読める。
Wasmtime+Metalで推論系を組みたい場合の出発点になる。
LLM以外でもいけるのか
元記事のデモはLlama 3.2 1Bの推論だが、この3段チェーンはLLM固有の仕組みに依存していない。
必要なのは「ホストで確保したバッファをGPUカーネルが読み書きし、その結果をホストに戻す」というGPGPUの基本動作だけ。
LLM向けのGEMM(行列積)カーネルが走るなら、原理的には以下の用途にも同じ仕組みが使える。
画像・動画系の推論では、CNN系の物体検出(YOLO、DETR系)や画像分類が候補。
入力テンソルと出力テンソルをゼロコピーで渡すパターンそのものなので、モデル種別を問わず恩恵を受ける。
Stable DiffusionのようなUNetベースの拡散モデルも同じ。
OCR・文書解析では、NDLOCR-Liteのようなレイアウト検出+文字認識パイプラインや、PaddleOCR-VL-1.5・GLM-OCR系の0.9BクラスのVLMは内部がTransformerなので、LLMとほぼ同じ呼び出し経路で動く。
検出モデル(DEIMv2等のCNN)と認識モデル(PARSeq/VLM)を連結する場合も、バッファの受け渡しが最頻出の処理なのでゼロコピーの効きが大きい。
音声系だとWhisperやvoice cloning系(LuxTTSのような軽量TTS)も、メルスペクトログラム等のテンソルをGPUに渡す点は同じ構造。
埋め込み・再ランキング用途では、Sentence Transformers v5.4のマルチモーダル埋め込みのような短系列エンコーダはむしろ推論が軽い分、ホスト↔GPUの往復コストの比率が大きく、ゼロコピーの恩恵が相対的に大きい。
非生成のGPGPUとしては、行列計算・物理シミュレーション・画像フィルタなど、そもそも機械学習ですらない処理にも転用できる。
Metalカーネル自体はLLMに縛られないので、コンピュートシェーダーを書けば何でも載る。
要は「Wasmモジュールから大量のバイトをGPUに食わせたい」ユースケース全般が射程で、LLMはその中の派手な一例に過ぎない。
ゼロコピーの旨みが出るのは入出力テンソルがMB〜GB級に大きい処理で、小さすぎるとコピーコスト自体が微小で差が出にくい。
WASM+OCRの歴史との接続
このブログでは過去、ブラウザ上でWASM OCRを動かそうとして躓いた記録を残してきた。
@paddlejs-models/ocrがブラウザで動かなかった話は、EmscriptenビルドのWASMがNode.js前提のグローバル変数を参照していたという、ランタイム適合の問題だった。
2025年のOCR Web 2025でも、「大きめの辞書やWASMを配置したくない」という配信側の都合でTesseract.jsに落ち着いている。
今回のDriftwoodはブラウザではなくネイティブ(macOS/Apple Silicon)上のWasmtimeの話なので、配信上の制約やModuleグローバル問題とは直接ぶつからない。
ただし「Wasmランタイムに載せたモデルを、実マシンのGPUで高速に回す」という問題設定は、過去のブラウザWASM OCRの延長線上にある。
ブラウザ側ではWebGPU+WASMのゼロコピーが同じテーマで、仕様上MLBufferやWebGPUのGPUBufferをWasmリニアメモリと共有する議論が進んでいる。
将来的にNDLocr-Lite iOS実装のようなオンデバイスOCRを、Wasm制御プレーン+GPUで組み直す構成もあり得る。
過去記事との対比で見ると位置づけがはっきりする。
| 記事 | 主題 | 実行環境 | GPU連携 |
|---|---|---|---|
| @paddlejs-models/ocrが動かない | ブラウザOCR | ブラウザWASM | WebGL経由、失敗 |
| OCR Web 2025 | ブラウザOCR運用 | ブラウザWASM | Tesseract.jsのみ |
| NDLOCR-Lite iOS | オンデバイスOCR | ONNX Runtime Mobile | CoreML/Metal |
| Flash-MoE | 巨大モデル推論 | ネイティブ | Metal直叩き |
| SwiftLM | LLM推論 | ネイティブSwift | Metal直叩き |
| 今回(Driftwood) | LLM推論 | Wasmtime | Metal+ゼロコピー |
Wasmを挟むかどうかが軸の差で、性能面ではネイティブ実装に遜色なくなってきた。配布容易性を取りつつGPU性能を捨てない、という方向が実現可能になる。
非LLMタスクのほうが効く
個人的な関心としては、OCRや埋め込み・検索のような「生成しない」推論のほうが、この手の構造の恩恵を受けやすいと見ている。
生成系は出力トークン列を1本ずつ逐次生成する性質上、1回あたりの計算量に対してGPU↔CPUの往復が相対的に小さい。
対して、OCR・埋め込み・検出系はバッチでまとめて投げて結果を受け取るスループット型で、往復回数が性能に直結する。
Wasmサンドボックスで配布可能なOCRエンジンが、ネイティブ並みの速度でApple SiliconのGPUを叩ける、という構図になると、配布モデルの選択肢が広がる。
現状はPython+PyTorch/MLX、あるいはSwift+Core ML+Metalのどちらかに集約されがちだが、中間にWasmが入る余地が出てくる。
Driftwood本体のOSS公開を待ちつつ、自前で組むときの設計例としても十分参考になる話だった。
元記事: Zero-Copy GPU Inference from WebAssembly on Apple Silicon
関連ドキュメント: