實驗、觀察與想法：科技工具、人工智慧、書籍、媒體、日常系統
重點摘要：在 Apple Silicon 上，WebAssembly 模組的線性記憶體可以直接與 GPU 共用：無需拷貝、無需序列化、無需中間緩衝區。CPU 和 GPU 讀寫相同的實體位元組。端對端均可運作：Wasm 程式載入器在線性記憶體中填入矩陣，GPU 讀取、計算、寫回，程式載入器透過相同的指標、相同的記憶體看到結果，實現零拷貝。
通常 Wasm 和 GPU 之間存在昂貴的序列化邊界：在大多數硬體上，將資料從虛擬機器沙盒傳輸到加速器需要跨匯流排拷貝。Apple Silicon 的統一記憶體架構消除了這個邊界（沒有匯流排，相同的實體記憶體），從而產生了一個 Wasm 作為控制平面、GPU 作為計算平面的運行時，它們之間的開銷幾乎為零。
我正在建構一個名為 Driftwood 的專案，利用這一點實現狀態化的 AI 推論……這篇文章將介紹其基礎（零拷貝鏈如何運作、我測量了什麼、它開啟了哪些可能性）。目前仍處於早期階段，仍在探索中。
快速背景介紹，適用於不熟悉此技術堆疊的讀者：WebAssembly 提供了一個沙盒。您的模組有一個扁平的位元組陣列（線性記憶體），這就是它的全部世界……外部的一切都透過「主機」函數呼叫進行協調。其重點在於隔離性、可移植性和確定性。
GPU 也需要一個扁平的位元組陣列，但有特定要求：分頁對齊、釘選、DMA 引擎可存取。在獨立 GPU（例如 NVIDIA 或 AMD）上，該記憶體位於 CPU 的 PCIe 匯流排另一側，因此將資料從 Wasm 模組的線性記憶體傳輸到 GPU 需要：從沙盒拷貝到主機記憶體，然後跨匯流排拷貝到 GPU 記憶體。兩次拷貝、兩次延遲命中，以及「隔離虛擬機器」與「硬體加速器」之間的不協調。
Apple Silicon 改變了物理原理。CPU 和 GPU 共用相同的實體記憶體（Apple 的統一記憶體架構）……沒有匯流排！CPU 可以讀取的指標，GPU 也可以從相同的 DRAM 讀取。真正問題是：您能否在不經過任何防禦性拷貝的情況下，將該指標穿透抽象層（Wasm 運行時、GPU API）？
三個連結。我在嘗試組合它們之前，各自驗證了它們的有效性：這種情況下，如果您跳過隔離步驟，整個流程就會中斷，您將不知道「哪個環節出了問題」。
連結 1：mmap 提供分頁對齊的記憶體。在 ARM64 macOS 上，使用 MAP_ANON | MAP_PRIVATE 的 mmap 會返回 16 KB 對齊的位址。這不是巧合，碰巧是 ARM64 的分頁大小，且 mmap 依合約進行對齊。對齊很重要，因為 Metal 需要它。
連結 2：Metal 在不拷貝的情況下接受該指標。MTLDevice.makeBuffer(bytesNoCopy:length:) 將現有的指標包裝成 Metal 緩衝區。在 Apple Silicon 上，這是零拷貝路徑，即 GPU 存取與 CPU 相同的實體記憶體。我驗證了指標的身份：MTLBuffer.contents() 指標等於原始的 mmap 指標。我驗證了沒有隱藏拷貝：RSS 差異為 0.03 MB（測量雜訊），而明確拷貝路徑為 16.78 MB。而且，兩種情況下的計算延遲相同。
連結 3：Wasmtime 允許您自訂配置分配器。Wasmtime 的 MemoryCreator 特徵允許您控制線性記憶體的分配方式。您提供自己的後備記憶體，而不是讓 Wasmtime 在內部呼叫 mmap。我實現了 MemoryCreator 以返回我們自己的 mmap 區域，Wasmtime 的 memory.data_ptr() 返回了我傳遞給它的確切指標。Wasm 模組透過 Wasmtime 的記憶體 API 讀寫；GPU 透過 Metal 緩衝區讀寫；兩者都在操作相同的位元組。
組合：分配一個 mmap 區域，將其同時提供給 Wasmtime（作為 Actor 的線性記憶體）和 Metal（作為 GPU 緩衝區）。Wasm 模組在已知偏移量處寫入資料，GPU 在原地進行計算，結果出現在模組的線性記憶體中，沒有拷貝，也沒有明確的資料傳輸。
我使用 128×128 的矩陣乘法測試了整個鏈條：Wasm 模組填入矩陣 A 和 B，GPU 執行 GEMM 著色器，模組讀取結果 C。16,384 個元素沒有錯誤。測試很小，但這種情況下，要麼一切都對齊，要麼得到垃圾，所以零錯誤是我想要的訊號。
我關心的三件事：指標身份（是否真的零拷貝？）、記憶體開銷（是否有隱藏拷貝潛入？）、以及正確性（GPU 是否看到 Wasm 寫入的內容？）。
延遲的等效性是有道理的：在 UMA 上，計算本身在這兩種情況下都是相同的。記憶體方面是顯示差異的地方：零拷貝路徑在使資料可供 GPU 存取方面幾乎沒有開銷，而拷貝路徑則使記憶體佔用量加倍。
在小張量尺寸下，沒有人在意。在 Transformer 推論的 KV 快取規模（每次對話數百 MB）下，這決定了記憶體中能容納四個 Actor 還是兩個。這才是我真正想運作的範圍，所以記憶體部分很重要。
所以現在我有一個基本原件：Wasm 和 GPU 在沒有開銷的情況下共用記憶體。您可以用它做什麼？
我將此鏈條插入 Apple 的 MLX 框架，並從 Wasm Actor 運行 Llama 3.2 1B Instruct：一個用 Rust 編寫的完整 Transformer 解碼器，編譯為原生主機運行時，透過主機函數呼叫驅動 Apple Silicon GPU 的推論。（我太懶了，不想從頭開始連接自訂核心路徑，而且……MLX 就在那裡）
測量延遲，在 2021 M1 Macbook Pro 上運行 Llama 3.2 1B（4 位元量化，695 MB）（舊個人筆記型電腦，我將來有機會拿到 Mac Studio 時會重新評估 😄）：
主機函數邊界（Wasm 到 GPU 的調度）與推論成本相比，無法測量。任何處理過沙盒運行時的人，大概都對每次調度時跨越邊界的想法感到不安。在此硬體上，這不成問題。
Transformer 維護一個鍵值快取，該快取會在對話回合中累積上下文，這通常是暫時性的（結束處理程序，丟失快取，重新開始）。如果您嘗試過本地推論，您知道這種感覺。
因為快取位於我控制的 GPU 可存取記憶體中，所以我可以將其序列化。因此，我將 KV 快取轉存為 safetensors 格式（標準 ML 張量序列化，沒有什麼特別的），然後稍後恢復它，可以在同一台機器上，或不同的機器上，或可能針對不同機器上的不同模型！最後一種情況我還沒有在有意義的不同架構上進行測試……我們拭目以待。
在 24 個 token 時為 5.45 倍，且隨著上下文長度的增加，比例會改善：恢復時間幾乎恆定，重新預填充呈線性縮放。在 4,096 個 token 時，恢復時間將比重新計算快約 100 倍（我實際上還沒有推到 4,096；這是根據恆定與線性形狀推斷的紙上談兵）。
這是狀態化 Actor 可攜性的基礎：凍結對話的中途交換，將其移動到其他地方，並在完整上下文完好無損的情況下解凍。Wasm 模組的線性記憶體捕捉 Actor 的邏輯狀態；KV 快取捕捉推論引擎累積的上下文。結合起來：一個正在運行的 AI 對話的可攜式快照（或者，至少，這就是計劃 😅）。
Driftwood 是用於具有 GPU 推論的狀態化 Wasm Actor 的運行時。零拷貝鏈是基礎：在此基礎上，我將添加 Actor 快照（凍結和恢復任何對話）、檢查點可攜性（將推論狀態移動到機器之間）以及多模型支援（快照格式與模型無關，因此理論上 Actor 的身份可以在模型交換後保留……這可能會奏效，一旦測試後我會重新審視）。
這一切都還處於早期階段，仍在將它們縫合在一起。但「物理原理」是可行的：Wasm 和 GPU 可以在 Apple Silicon 上以零開銷共用記憶體，KV 快取是可攜的，並且一個完整的 Transformer 可以以原生速度從沙盒 Actor 運行。接下來我想探討的問題是：快照是否真的能在模型交換後保留，鏈條在更大的模型上是否仍然有效，以及我在規模化時是否遺漏了某些顯而易見的會導致失敗的原因。穩步前進……
在未來的一篇文章中，一旦我發布了超越「物理原理可行」階段的內容，將會更多地介紹 Actor 模型和快照架構。
獲取關於人工智慧、系統和意義的簡潔、分析性且人性化的見解。