4月18日、著者が「Zero-Copy GPU Inference from WebAssembly on Apple Silicon」と題した記事を公開した。この記事では、Apple SiliconでWebAssemblyモジュールがGPUとメモリを直接共有し、コピーゼロでAI推論を実行する手法について詳しく紹介されている。
AI推論の高速化競争で注目されるWebAssembly
現在、AIアプリケーションの普及に伴い、推論処理の高速化と効率化が急務となっている。特に、サーバーレス環境やエッジコンピューティングでは、セキュリティを保ちながら高性能を実現する必要がある。WebAssembly（Wasm）は、サンドボックス化された安全な実行環境を提供する一方で、GPU推論との組み合わせでは性能上の課題があった。
著者が発見したのは、Apple SiliconのUnified Memory Architecture（UMA）を活用することで、この根本的な制約を解決できるということだった。結果として、KVキャッシュの復元が67倍高速化を実現している。
従来の「2回コピー問題」とApple Siliconの解決策
通常、WebAssemblyとGPUの間には高コストな境界が存在する。WasmモジュールはWebアプリケーションやサーバーレス環境で安全にコードを実行するための仮想マシンで、線形メモリという隔離された領域で動作する。一方、GPUは独自のメモリ空間を持つため、データのやり取りには2回のコピーが必要だった。まずWasmのサンドボックスからホストメモリへコピーし、次にPCIeバス経由でGPUメモリへコピーするという具合だ。
しかし、Apple SiliconのUnified Memory Architecture（UMA）は、この物理的制約を取り除いた。CPUとGPUが同一の物理メモリを共有するため、バスを介したコピーが不要になる。Apple Siliconは、M1チップ以降でこのアーキテクチャを採用し、従来のx86系プロセッサとは根本的に異なるメモリ管理を実現している。
著者が挑戦したのは、この共有ポインターを抽象化レイヤー（Wasmランタイム、GPU API）を通して「防御的コピー」なしに渡すことだった。そして、その答えは「できる」だった。
3つのリンクで構成するゼロコピーチェーン
著者は3つの技術的要素を組み合わせてゼロコピーチェーンを構築した。
リンク1：mmapによるページ対齊メモリ確保
ARM64 macOSでは、mmapシステムコールが16KB境界に整列されたアドレスを返す。これはARM64のページサイズに一致しており、Metalが要求する条件を満たす。
リンク2：Metalがコピーなしでポインターを受け入れ
Apple Metalは、Apple独自のGPU計算フレームワークで、MTLDevice.makeBuffer(bytesNoCopy:length:)メソッドを提供している。これは既存のポインターをMetalバッファとしてラップし、Apple Siliconでは真のゼロコピー経路となる。
リンク3：Wasmtimeで独自アロケータを使用
Wasmtimeは、Bytecode Alliance（Mozilla、Fastly、Intel等が参加）が開発するRustベースのWebAssemblyランタイムで、高いパフォーマンスと安全性を提供する。そのMemoryCreatorトレイトを使用することで、線形メモリの割り当て方法を制御できる。
実測データが証明するゼロコピーの効果
著者は128×128行列乗算でフルチェーンをテストし、以下の結果を得た。
測定項目 ゼロコピー経路 コピー経路
────────────────────────────────────────────────────
ポインター同一性 mmap == MTLBuffer 異なるアドレス
RSS増分（16MB領域） 0.03 MB 16.78 MB
GEMM遅延（128×128） ~6.75 ms ~6.75 ms
正確性（16K要素） 0エラー 0エラー
計算遅延は同等だが、メモリオーバーヘッドは劇的に改善した。小さなテンソルでは差は見えないが、大規模言語モデル（LLM）の推論で使われるKVキャッシュ（会話あたり数百MB）の規模では、「メモリに4つのアクターを収容できるか、2つしか収容できないか」の違いを生む。
Llama 3.2での実証実験と67倍の高速化
著者はこのゼロコピーチェーンをApple MLXフレームワークに接続し、WasmアクターからLlama 3.2 1B Instructを実行した。MLXは、Apple SiliconのGPUで機械学習を効率的に実行するためにApple自身が開発したオープンソースフレームワークだ。
2021年M1 MacBook Proでの測定結果：
操作 遅延
───────────────────────────────────
モデルロード 229 ms （1回のみ）
プリフィル（5トークン） 106 ms
トークン生成あたり ~9 ms
Wasm-GPU境界オーバーヘッド 測定不能
さらに注目すべきはKVキャッシュの可搬性だ。制御可能なGPUアクセス可能メモリにKVキャッシュが存在するため、これをsafetensors形式でシリアライズできる。safetensorsは、Hugging Faceが開発したテンソルデータの安全なシリアライゼーション形式で、メタデータの検証機能を持つ。
操作 遅延 サイズ
──────────────────────────────────────────────
シリアライズ（24トークン） 1.1 ms 1.58 MB（~66 KB/トークン）
ディスクから復元 1.4 ms
スクラッチから再プリフィル 67.7 ms （従来手法）
──────────────────────────────────────────────
復元による高速化: 48.3× （67.7ms → 1.4ms）
往復忠実度: ビット同一（10/10トークン一致）
24トークンで48.3倍の高速化を実現し、この比率はコンテキスト長に比例して改善する。4,096トークンでは理論上約100倍の高速化が期待できる。
「Driftwood」：ステートフルAIアクターランタイムの未来
この成果を基に、著者はDriftwoodと呼ばれるランタイムを開発している。これはGPU推論機能を持つステートフルWasmアクター向けのシステムで、以下の革新的機能を実現する。
- アクタースナップショット：任意の会話状態を瞬時に凍結・復元
- チェックポイント可搬性：マシン間での推論状態完全移動
- マルチモデルサポート：モデル非依存のスナップショット形式
これによりステートフルアクターモビリティが可能になる。会話を途中で凍結し、他の場所に移動して、完全なコンテキストと共に復元できる。Wasmモジュールの線形メモリがアクターの論理状態を、KVキャッシュが推論エンジンの蓄積コンテキストを捉える組み合わせだ。
AI推論の新たなパラダイム
Apple SiliconでWasmとGPUがゼロオーバーヘッドでメモリを共有でき、KVキャッシュは可搬で、完全なトランスフォーマーがサンドボックス化されたアクターからネイティブ速度で動作することが実証された。
この技術は、サーバーレスAI推論やエッジコンピューティング環境での効率的なモデル実行において革新的だ。特に、ChatGPT的な会話AIサービスでは、ユーザーセッションの状態管理が課題となっているが、この手法により会話の「冷凍保存」と「瞬時復元」が現実的になる。また、リソース制約のある環境でのマルチテナント推論において、メモリ効率の67倍改善は直接的なコスト削減につながる。
Wasmがブラウザを超えてサーバーサイドで注目される理由の一つが、このようなセキュリティとパフォーマンスの両立にある。今回の成果は、Apple Siliconという特定のハードウェアに依存するものの、統合メモリアーキテクチャの利点を最大化する具体的な実装例として価値が高い。
詳細はZero-Copy GPU Inference from WebAssembly on Apple Siliconを参照していただきたい。
