Zero-Copy GPU Inference with WebAssembly on Apple Silicon

On top of Apple Silicon's Unified Memory Architecture, WebAssembly linear memory and GPU memory are bound to the same physical memory, enabling inference without copying.

The core consists of three steps.

• Secure page-aligned memory with mmap. On ARM64 macOS, 16KB alignment is guaranteed.
• Wrap that pointer directly with Metal's MTLDevice.makeBuffer(bytesNoCopy:length:), so the GPU reads and writes the same memory. In practice, the pointer was identical, and RSS increase was about 0.03 MB, essentially none. In contrast, the explicit copy path increased by 16.78 MB.
• Control the allocation of Wasm linear memory itself with Wasmtime MemoryCreator, so that the memory.data_ptr() returned by Wasmtime exactly matches the mmap pointer I passed in.

Combining this chain, the GPU immediately reads values written to memory by the Wasm module, and Wasm reads the results written by the GPU back through the same pointer. In a 128×128 GEMM test, 0 errors across 16,384 elements were confirmed.

In terms of performance, the computation time itself is nearly the same as the copy path, but the difference is clear in the memory path. GEMM latency was about 6.75 ms, the same, and the value of zero-copy lies mainly in reducing memory footprint.

On top of this, MLX was attached to run Llama 3.2 1B Instruct inside a Wasm actor. On a 2021 M1 MacBook Pro, model load was 229 ms, prefill was 106 ms, and generation was about 9 ms per token, with the Wasm-to-GPU boundary cost essentially unmeasurable.

Also, the KV cache can be saved to and restored from a file. In safetensors format, a cache of 24 tokens was serialized in 1.1 ms / 1.58 MB, and restore took 1.4 ms. Re-prefilling the same content from scratch took 67.7 ms, making this 5.45× faster at this point. The restored result matched bit-identically.

This combination enables snapshotting and migration of stateful AI actors. Wasm linear memory holds the actor's logical state, while the KV cache holds the inference engine's accumulated context; saving both together forms a foundation for pausing an AI mid-conversation and resuming it elsewhere unchanged. It's still early days, but Wasm + GPU zero-copy inference on Apple Silicon actually works, and Driftwood, a stateful actor runtime, is being built on top of it.