在 Apple Silicon 上,WebAssembly 模組的線性記憶體可以直接與 GPU 共享——無複製、無序列化、無中間緩衝。CPU 與 GPU 讀寫同一份實體位元組。這不是理論,這是實際驗證過的結果。

背景:為何這件事通常很困難

WebAssembly 給你一個沙盒。模組取得一個平面的位元組陣列(線性記憶體),這就是它的宇宙——所有外部操作都必須透過「宿主」函數呼叫。隔離性、可移植性、確定性是它的核心價值。

GPU 同樣需要一個平面的位元組陣列,但有特定限制:必須是頁面對齊的、固定的、且 DMA 引擎可以存取的。在離散顯示卡上(NVIDIA 或 AMD),那塊記憶體位於 PCIe 匯流排的另一端,因此將資料從 Wasm 模組的線性記憶體傳到 GPU 意味著:先從沙盒複製到宿主記憶體,再透過匯流排複製到 GPU 記憶體。兩次複製、兩次延遲,還有「隔離 VM」與「硬體加速器」之間的阻抗不匹配問題。

Apple Silicon 改變了這個物理限制。CPU 與 GPU 共享同一塊實體記憶體(Unified Memory Architecture)——沒有匯流排。CPU 可讀取的指標,GPU 同樣可以從同一塊 DRAM 讀取。真正的問題是:你能將這個指標穿透各層抽象(Wasm 執行時、GPU API)而途中沒有人做防禦性複製嗎?

結果是可以的。

三環鏈結

三個環節,每個都先獨立驗證才嘗試組合:

環節一:mmap 給你頁面對齊的記憶體

在 ARM64 macOS 上,mmap 搭配 MAP_ANON | MAP_PRIVATE 會回傳 16 KB 對齊的位址。這不是運氣巧合,而是 ARM64 頁面大小,mmap 按合約對齊。對齊很重要,因為 Metal 需要它。

環節二:Metal 直接接受該指標而不複製

MTLDevice.makeBuffer(bytesNoCopy:length:options:deallocator:) 將現有指標包裝成 Metal buffer。在 Apple Silicon 上,這是零複製路徑——GPU 存取與 CPU 相同的實體記憶體。驗證方式:

  • 指標一致性:MTLBuffer.contents() 等於原始 mmap 指標
  • 無隱藏複製:RSS 增量為 0.03 MB(量測雜訊),對比明確複製路徑的 16.78 MB
  • 計算延遲相同

環節三:Wasmtime 讓你帶自己的記憶體配置器

Wasmtime 的 MemoryCreator trait 讓你控制線性記憶體的配置方式。不讓 Wasmtime 內部呼叫 mmap,而是自己提供備援記憶體。實作 MemoryCreator 回傳我們自己的 mmap 區域,Wasmtime 的 memory.data_ptr() 就會回傳我們當初給它的指標。

組合結果:配置一個 mmap 區域,同時交給 Wasmtime(作為演員的線性記憶體)和 Metal(作為 GPU buffer)。Wasm 模組在已知偏移量寫入資料,GPU 就地計算,結果直接出現在模組的線性記憶體中,零複製、零明確資料傳輸。

用 128×128 矩陣乘法測試完整鏈結:Wasm 模組填入矩陣 A 和 B,GPU 執行 GEMM 著色器,模組讀回結果 C。16,384 個元素零錯誤。

測量結果

驗證三件事:指標一致性、記憶體開銷、正確性。

測量項目 零複製路徑 複製路徑
指標一致性 mmap == MTLBuffer 不同位址
RSS 增量(16 MB 區域) 0.03 MB 16.78 MB
GEMM 延遲(128×128) ~6.75 ms ~6.75 ms
正確性(16K 元素) 0 錯誤 0 錯誤

延遲相同是合理的:在 UMA 架構上,計算本身是相同的。差異在記憶體端:零複製路徑幾乎沒有讓資料可被 GPU 存取的額外開銷,而複製路徑則讓記憶體使用量翻倍。

在小張量大小下,沒人在意。但在 transformer 推論中 KV cache 的規模(每次對話數百 MB),這是能否在記憶體中容納四個演員還是只有兩個的差別。

整合 MLX 執行 Llama 3.2 1B

將這條鏈結接入 Apple 的 MLX 框架,從 Wasm 演員執行 Llama 3.2 1B Instruct:一個完整的 transformer 解碼器,以 Rust 撰寫、編譯為原生宿主執行時,透過宿主函數呼叫在 Apple Silicon GPU 上驅動推論。

測量環境:2021 M1 Macbook Pro(舊款個人筆電),Llama 3.2 1B(4-bit 量化,695 MB)。

操作 延遲
模型載入(safetensors) 229 ms(一次性)
Prefill(5 tokens) 106 ms
每 token 生成 ~9 ms
宿主函數邊界 可忽略

宿主函數邊界(Wasm 到 GPU 的調度)相對於推論成本幾乎無法測量。任何曾經處理過沙盒執行時的人都可能對「每次調度都跨越那個邊界」的想法感到擔憂,但在這塊硬體上,這不是問題。

KV Cache 可攜性

Transformer 維護一個 key-value cache,在對話過程中累積上下文,這通常是暫時的(殺掉程序就失去 cache,從頭開始)。如果你嘗試過本地推論,你一定知道那種感覺。

因為 cache 存在我控制的 GPU 可存取記憶體中,所以可以序列化。將 KV cache 傾印到 safetensors 格式,稍後再恢復——在同一台機器上、不同機器上,甚至可能在不同模型上。

操作 延遲 大小
序列化(24 tokens) 1.1 ms 1.58 MB (~66 KB/token)
從磁碟恢復 1.4 ms
從零重新 Prefill 67.7 ms
恢復加速比 5.45×
往返精確度 位元相同(10/10 tokens 匹配)

24 tokens 下加速 5.45 倍,隨上下文長度增加,比率會更好:恢復時間幾乎是常數,而重新 Prefill 是線性成長。推估 4,096 tokens 時,恢復比重新計算快約 100 倍。

這是狀態演員移動性的基礎:在對話中途凍結、移到其他地方、完整上下文恢復。Wasm 模組的線性記憶體捕捉演員的邏輯狀態;KV cache 捕捉推論引擎的累積上下文。兩者合一:一個可攜的運行中 AI 對話快照。

Driftwood:正在建構的東西

Driftwood 是一個為狀態化 Wasm 演員搭配 GPU 推論的執行時。零複製鏈結是基礎:在此之上將陸續加入演員快照(凍結並恢復任何對話)、checkpoint 可攜性(跨機器移動推論狀態)、以及多模型支援(快照格式是模型無關的,理論上演員身份在模型交換後仍可延續)。

目前一切都還很早期。物理層面已經驗證可行:Wasm 與 GPU 在 Apple Silicon 上可以零開銷共享記憶體,KV cache 是可攜的,完整的 transformer 在沙盒化演員中以原生速度運行。

接下來要探索的方向:快照是否真的能在模型交換後存活、鏈結在更大模型上是否依然成立、以及是否忽略了什麼在規模化時會導致崩潰的明顯原因。

穩扎穩打。

關於演員模型和快照架構的更多細節,留待真正將東西推進到「物理層面可行」階段之後再寫。