Apple Silicon에서 WebAssembly로 하는 Zero-Copy GPU Inference

Apple Silicon의 Unified Memory Architecture 위에서 WebAssembly linear memory와 GPU 메모리를 같은 물리 메모리로 묶어, 복사 없이 추론을 수행한다.

핵심은 세 단계다.

• mmap 으로 페이지 정렬된 메모리를 확보한다. ARM64 macOS에서는 16KB 정렬이 보장된다.
• Metal 의 MTLDevice.makeBuffer(bytesNoCopy:length:) 로 그 포인터를 그대로 감싸면, GPU가 같은 메모리를 읽고 쓴다. 실제로 포인터가 동일했고, RSS 증가도 0.03 MB 수준으로 사실상 없었다. 반면 명시적 복사 경로는 16.78 MB가 늘었다.
• Wasmtime MemoryCreator 로 Wasm linear memory의 할당 자체를 제어해, Wasmtime이 반환하는 memory.data_ptr() 가 내가 넘긴 mmap 포인터와 정확히 같게 만든다.

이 체인을 합치면, Wasm 모듈이 메모리에 쓴 값을 GPU가 바로 읽고, GPU가 쓴 결과를 다시 같은 포인터로 Wasm이 읽는다. 128×128 GEMM 테스트에서는 16,384개 요소에서 0 오류를 확인했다.

성능상으로는 계산 시간 자체는 복사 경로와 거의 같지만, 메모리 경로에서 차이가 뚜렷하다. GEMM latency는 약 6.75 ms로 동일했고, 제로카피의 가치는 주로 메모리 footprint 절감에 있다.

이 위에 MLX 를 붙여 Llama 3.2 1B Instruct 를 Wasm actor에서 돌렸다. 2021년형 M1 MacBook Pro 기준으로 model load 229 ms, prefill 106 ms, token당 생성 약 9 ms였고, Wasm-to-GPU 경계 비용은 사실상 측정되지 않았다.

또한 KV cache 를 파일로 저장·복원할 수 있다. safetensors 형식으로 24 tokens의 캐시를 1.1 ms / 1.58 MB에 serialize했고, restore는 1.4 ms였다. 같은 내용을 scratch에서 다시 prefill하면 67.7 ms가 걸려, 이 시점에서 5.45× 빠르다. 복원 결과는 bit-identical 하게 일치했다.

이 조합은 상태를 가진 AI actor의 스냅샷과 이동을 가능하게 한다. Wasm linear memory는 actor의 논리 상태를, KV cache는 추론 엔진의 누적 컨텍스트를 담고, 둘을 함께 저장하면 대화 중인 AI를 멈췄다가 다른 곳에서 그대로 재개할 수 있는 기반이 된다. 지금은 아직 초기 단계지만, Apple Silicon에서 Wasm + GPU zero-copy inference 는 실제로 동작하며, 그 위에 stateful actor runtime인 Driftwood 를 쌓고 있다.