Driftwood: Zero-copy GPU inference từ WebAssembly trên chip Apple Silicon

Trên các máy tính sử dụng chip Apple Silicon, một mô-đun WebAssembly (Wasm) giờ đây có thể chia sẻ bộ nhớ trực tiếp với GPU mà không cần bất kỳ bước sao chép hay tuần tự hóa dữ liệu nào. Điều này có nghĩa là CPU và GPU có thể đọc và ghi cùng một vùng nhớ vật lý, tạo ra một cơ chế hoạt động "zero-copy" (không sao chép) hoàn toàn từ đầu đến cuối.

Driftwood Zero-Copy

Thông thường, Wasm và GPU hoạt động tách biệt bởi một ranh giới tuần tự hóa tốn kém. Trên phần lớn các phần cứng khác, việc đưa dữ liệu từ sandbox của máy ảo (VM) sang bộ tăng tốc (accelerator) như GPU đòi hỏi phải sao chép dữ liệu qua các bus (ví dụ: PCIe). Tuy nhiên, kiến trúc Bộ nhớ Thống nhất (Unified Memory Architecture - UMA) của Apple Silicon đã xóa bỏ ranh giới này. Không còn bus tách biệt, cả CPU và GPU đều sử dụng cùng một bộ nhớ DRAM vật lý.

Dự án Driftwood đang khai thác khả năng này để tạo ra một runtime nơi Wasm đóng vai trò là mặt điều khiển (control plane) và GPU là mặt tính toán (compute plane) với chi phí trung gian gần như bằng không.

Tại sao vấn đề này thường khó giải quyết?

WebAssembly cung cấp một sandbox an toàn, nơi mô-đun nhận được một mảng byte phẳng (linear memory) và mọi thứ bên ngoài đều được trung gian hóa bởi các lệnh gọi hàm của "host". Mục đích chính là sự cô lập, tính di động và tính xác định.

Ngược lại, GPU cũng muốn một mảng byte phẳng, nhưng với các yêu cầu cụ thể: căn chỉnh theo trang (page-aligned), được ghim (pinned) và có thể truy cập bởi động cơ DMA. Trên các GPU rời rời (như NVIDIA hoặc AMD), bộ nhớ này nằm ở phía bên kia bus PCIe khỏi CPU. Do đó, việc chuyển dữ liệu từ bộ nhớ tuyến tính của Wasm sang GPU thường yêu cầu: sao chép ra khỏi sandbox vào bộ nhớ host, sau đó sao chép qua bus vào bộ nhớ GPU. Đó là hai lần sao chép, hai lần đánh đổi độ trễ và sự không tương thích về trở kháng giữa "VM cô lập" và "bộ tăng tốc phần cứng".

Apple Silicon đã thay đổi quy luật vật lý này. Vì CPU và GPU chia sẻ cùng một bộ nhớ vật lý, một con trỏ mà CPU có thể đọc thì GPU cũng có thể đọc. Vấn đề đặt ra là: liệu có thể luồn con trỏ đó qua các lớp trừu tượng (runtime Wasm, GPU API) mà không có bất kỳ cơ chế phòng vệ nào tạo ra bản sao chép dọc theo đường đi?

Chuỗi ba liên kết (Three-link chain)

Giải pháp này bao gồm ba liên kết chính:

1. mmap cung cấp bộ nhớ căn chỉnh trang: Trên ARM64 macOS, mmap với cờ MAP_ANON | MAP_PRIVATE trả về các địa chỉ căn chỉnh 16 KB. Đây không phải là ngẫu nhiên mà là kích thước trang của ARM64, điều mà Metal yêu cầu.
2. Metal chấp nhận con trỏ mà không sao chép: Phương thức MTLDevice.makeBuffer(bytesNoCopy:length:) bao bọc một con trỏ hiện có dưới dạng Metal buffer. Trên Apple Silicon, đây là đường dẫn zero-copy, nơi GPU truy cập cùng một bộ nhớ vật lý với CPU. Việc xác minh danh tính con trỏ cho thấy MTLBuffer.contents() bằng với con trỏ mmap gốc.
3. Wasmtime cho phép bạn tự cấp phát bộ nhớ: Trait MemoryCreator của Wasmtime cho phép kiểm soát cách cấp phát bộ nhớ tuyến tính. Thay vì để Wasmtime gọi mmap nội bộ, bạn có thể tự cung cấp bộ nhớ hỗ trợ. Điều này cho phép Wasmtime sử dụng vùng mmap của chúng ta, và GPU sử dụng Metal buffer, cả hai đều hoạt động trên cùng một byte.

Hiệu suất và ứng dụng thực tế

Tác giả đã kiểm tra chuỗi hoàn chỉnh với phép nhân ma trận 128×128. Mô-đun Wasm điền ma trận A và B, GPU chạy shader GEMM, và mô-đun đọc kết quả C trở lại. Kết quả cho thấy không có lỗi nào trên 16.384 phần tử.

Về hiệu suất, đường dẫn zero-copy về cơ bản không có chi phí khi làm cho dữ liệu có thể truy cập bởi GPU, trong khi đường dẫn sao chép sẽ làm tăng gấp đôi dấu chân bộ nhớ (memory footprint). Với các kích thước tensor nhỏ, điều này không đáng kể, nhưng ở quy mô của bộ đệm KV (KV cache) trong suy luận transformer (hàng trăm megabyte mỗi cuộc hội thoại), đây là sự khác biệt giữa việc vừa đủ bộ nhớ cho 4 tác nhân hay chỉ 2.

Apple Silicon GPU

Kết hợp chuỗi này với khung MLX của Apple, tác giả đã chạy thành công mô hình Llama 3.2 1B Instruct từ một tác nhân Wasm trên MacBook Pro M1 năm 2021. Ranh giới hàm host (kích hoạt Wasm sang GPU) là không đáng kể so với chi phí suy luận.

Khả năng di chuyển của bộ đệm KV (KV Cache)

Một lợi ích lớn khác là khả năng di chuyển trạng thái. Transformers duy trì bộ đệm key-value tích lũy ngữ cảnh qua các lượt hội thoại. Vì bộ đệm này nằm trong bộ nhớ có thể truy cập GPU mà ta kiểm soát, nó có thể được tuần tự hóa.

Tác giả đã thử nghiệm việc đổ bộ đệm KV sang định dạng safetensors và khôi phục lại sau đó. Việc khôi phục từ đĩa nhanh hơn 5,45 lần so với việc tính toán lại từ đầu (re-prefill) với 24 token. Tỷ lệ này còn cải thiện hơn nữa khi độ dài ngữ cảnh tăng lên.

Đây là cơ sở cho tính di chuyển của tác nhân có trạng thái (stateful actor mobility): đóng băng một cuộc hội thoại giữa chừng, di chuyển nó đến nơi khác, và giải đông với ngữ cảnh đầy đủ còn nguyên vẹn.

Driftwood hiện đang được phát triển như một runtime cho các tác nhân Wasm có trạng thái với suy luận GPU. Mặc dù còn ở giai đoạn sơ khai, nhưng "vật lý" của nó đã hoạt động: Wasm và GPU có thể chia sẻ bộ nhớ trên Apple Silicon với chi phí bằng không, bộ đệm KV có thể di chuyển, và một transformer hoàn chỉnh có thể chạy từ một tác nhân sandbox với tốc độ gốc.