feat(gms): introduce VMMDevice abstraction for XPU (Phase 1)

lib/gpu_memory_service/GMS_MULTI_DEVICE.md

@@ -0,0 +1,185 @@
+# GMS Multiple Device Enablement — Design Proposal
+
+## 1. Overview
+
+GPU Memory Service (GMS) manages cross-process virtual memory on accelerators for weight /KV-cache sharing in Dynamo inference clusters.
+
+This design introduces a **device-agnostic VMM abstraction layer** so that Intel XPU can plug in alongside the existing CUDA path without touching the server, client, or snapshot logic.
+
+---
+
+## 2. Goals
+
+| # | Goal | Status |
+|---|------|--------|
+| G1 | Define a vendor-neutral `VMMDevice` ABC covering all device operations GMS needs | ✅ Phase 1 |
+| G2 | Wrap existing CUDA helpers in a `CudaVMM` class that inherits the ABC | ✅ Phase 1 |
+| G3 | Process-global VMM singleton via `init_vmm()` / `get_vmm()`, CLI `--device-type` | ✅ Phase 1 |
+| G4 | Implement `XpuVMM` | ⬜ Phase 2 |
+| G5 | Add XPU torch mempool dispatch via `torch.xpu` APIs | ⬜ Phase 2 |
+| G6 | Enable snapshot save/load for XPU | ⬜ Phase 2 |
+
+---
+
+## 3. Architecture (Phase 1 — Delivered)
+
+```
+┌──────────────────────────────────────────────────────┐
+│                  CLI / Supervisor                    │
+│  server.py  ──→  --device-type {cuda,xpu}            │
+└────────────────────────┬─────────────────────────────┘
+                         │ VMMDeviceType enum
+                         ▼
+┌──────────────────────────────────────────────────────┐
+│            common/vmm/__init__.py                    |
+│  Singleton instance:                                 |
+│  init_vmm(device_type)                               │
+│  get_vmm() → VMMDevice                               │
+│                                                      │
+│  ┌───────────────┐          ┌───────────────┐        │
+│  │  CudaVMM      │          │  XpuVMM       │        │
+│  │  (cuda_utils) │          │  (xpu_utils)  │        │
+│  └───────────────┘          └───────────────┘        │
+└──────────────────────────────────────────────────────┘
+                         │
+          ┌──────────────┼──────────────┐
+          ▼              ▼              ▼
+   GMSRPCServer       GMS          GMSAllocationManager
+   (server/rpc.py)  (server/gms.py) (server/allocations.py)
+          │
+          ▼
+   GMSClientMemoryManager  →  GMSStorageClient (snapshot)
+   (client/memory_manager.py)  (snapshot/storage_client.py)
+```
+
+### 3.1 New Module: `common/vmm/`
+
+| File | Purpose |
+|------|---------|
+| `__init__.py` | `VMMDeviceType` enum, `init_vmm()` singleton initializer, `get_vmm()` accessor, `get_vmm_device_type()` |
+| `device.py` | `VMMDevice` — `abc.ABC` with 24 `@abstractmethod` vendor-neutral methods |
+| `cuda_utils.py` | Module-level CUDA helpers (unchanged logic) + `CudaVMM(VMMDevice)` class |
+
+- `init_vmm(device_type)`: explicit process-level backend selection, e.g.
+   CLI `--device-type`. On a mixed node (e.g., CUDA and other device coexist),
+   --device-type device_name overrides the auto-detection because it calls
+   init_vmm(VMMDeviceType.DeviceName) at process startup — before any get_vmm() lazy path fires.
+- `get_vmm()`: returns the singleton and may lazily initialize the default/autodetected
+   backend for public client/integration compatibility via init_vmm(detected_device_type).
+- `vmm.ensure_initialized()`: initializes the selected backend driver/runtime, e.g. CUDA `cuInit`
+
+### 3.2 `VMMDevice` ABC Surface (24 methods)
+
+```
+Category                  Method
+───────────────────────── ─────────────────────────────────────
+Driver lifecycle          ensure_initialized()
+                          synchronize()
+Discovery / sizing        list_devices() → list[int]
+                          get_allocation_granularity(device) → int
+Physical memory           create_tolerate_oom(size, device) → (bool, int)
+                          release(handle)
+Shareable handles         export_to_shareable_handle(handle) → int (FD)
+                          import_shareable_handle_close_fd(fd) → int
+VA space + mapping        address_reserve(size, granularity) → int
+                          address_free(va, size)
+                          map(va, size, handle)
+                          unmap(va, size)
+                          set_access(va, size, device, access)
+Monitoring / sizing       device_memory_info(device) → (free_bytes, total_bytes)
+Pointer validation        validate_pointer(va)
+Runtime helpers           runtime_check_result(result, name)
+                          runtime_set_device(device)
+                          host_register(ptr, size)
+                          host_unregister(ptr)
+Stream management         stream_create_nonblocking() → opaque
+                          stream_destroy(stream)
+                          stream_synchronize(stream)
+Async copy                memcpy_h2d_async(dst, src, size, stream)
+                          memcpy_d2h_async(dst, src, size, stream)
+```
+
+### 3.3 Singleton Lifecycle
+
+```
+Process startup (cli/server.py, cli/runner.py, snapshot loader/saver):
+    init_vmm(VMMDeviceType.from_str(args.device_type))
+
+Any module needing VMM:
+    vmm = get_vmm()          # returns cached singleton
+    vmm.runtime_set_device(device)
+    vmm.list_devices()
+    ...
+```
+
+The singleton is immutable after initialization. Conflicting re-initialization
+raises `RuntimeError`. Thread-safe via `threading.Lock`.
+
+### 3.4 CLI Argument Propagation
+
+```
+gms-server (supervisor)        --device-type cuda|xpu
+  └─→ init_vmm(device_type)
+  └─→ gms-server (per-device)  --device-type  (forwarded to child process)
+       └─→ init_vmm(device_type)
+       └─→ get_vmm() used throughout server, allocations, client
+```
+
+No constructor threading — every module calls `get_vmm()` directly.
+
+## 4. Files Changed (Phase 1)
+
+| File | Change |
+|------|--------|
+| `cli/server.py` | Parse `--device-type`, call `init_vmm()`, use `get_vmm().list_devices()` |
+| `cli/runner.py` | Call `init_vmm(config.device_type)` at startup |
+| `cli/snapshot/loader.py` | Call `init_vmm(device_type)`, `vmm = get_vmm()` in helpers |
+| `cli/snapshot/saver.py` | Call `init_vmm(device_type)`, `vmm = get_vmm()` in helpers |
+| `client/memory_manager.py` | `self._vmm = get_vmm()`, property `device_type` via `get_vmm_device_type()` |
+| `client/torch/allocator.py` | CUDA-only guards via `get_vmm_device_type()` |
+| `common/vmm/__init__.py` | **NEW** — enum + singleton |
+| `common/vmm/device.py` | **NEW** — VMMDevice ABC |
+| `common/vmm/cuda_utils.py` | **NEW** (moved from `common/cuda_utils.py`) — CUDA helpers + CudaVMM class |
+| `common/utils.py` | Add `align_to_granularity()` utility |
+| `server/allocations.py` | `self._vmm = get_vmm()` |
+| `server/gms.py` | No longer forwards device params — singleton |
+| `server/rpc.py` | No longer forwards device params — singleton |
+| `snapshot/storage_client.py` | No longer accepts device_type — uses singleton |
+| `snapshot/disk.py` | Fix import path: `common.vmm` |
+| `snapshot/backends/nixl_staging.py` | Fix import path: `common.vmm` |
+| `snapshot/backends/pinned_host.py` | Fix import path: `common.vmm` |
+| `tests/report_pytest_markers.py` | Update stub module list |
+
+---
+
+## 5. Phase 2 — XPU Implementation
+
+### 5.1 `XpuVMM(VMMDevice)`
+
+Implement the methods on the OneAPI/Torch runtime.
+
+### 5.2 Torch Allocator Dispatch
+
+The PyTorch front door (`client/torch/`) routes torch tensor allocations through GMS via a pluggable allocator inside a `gms_use_mem_pool(tag)` context; the dispatch key is the existing `get_vmm_device_type()`.
+
+- **`extensions/allocator.cpp` + `setup.py` — no change.**  `my_malloc(ssize_t,  int, void* stream)` forwards the opaque stream/queue to Python *without dereferencing it*, so it is ABI-compatible with `XPUPluggableAllocator`'s
+  `void* alloc_fn(size_t, int, sycl::queue*)`. The **same** `_allocator_ext.so`
+  and `"my_malloc"` / `"my_free"` symbols serve both backends.
+- **`client/torch/allocator.py` — only change.** Swap the four CUDA-hardcoded
+  spots for device dispatch (`torch.cuda.*`, `torch.xpu.*`):
+
+  | Spot | CUDA | XPU |
+  |------|------|-----|
+  | `_ensure_callbacks_initialized` | `torch.cuda.CUDAPluggableAllocator` | `torch.xpu.memory.XPUPluggableAllocator` |
+  | `_create_mem_pool` | `torch.cuda.memory.MemPool` | `torch.xpu.memory.MemPool` |
+  | `gms_use_mem_pool` | `torch.cuda.use_mem_pool` | `torch.xpu.memory.use_mem_pool` |
+  | `prune_allocations` | `torch.cuda.synchronize` | `torch.xpu.synchronize` |
+
+  The torch APIs are symmetric, so a small accessor returning the active device's
+  `(PluggableAllocator, MemPool, use_mem_pool, synchronize)` is sufficient.
+
+---
+
+## 6. Backward Compatibility
+
+CUDA path remains unchanged. All existing `--device-type cuda` deployments work identically. The default value is `cuda` everywhere.

lib/gpu_memory_service/cli/args.py

@@ -9,6 +9,7 @@
 from typing import Optional
 
 from gpu_memory_service.common.utils import GMS_TAGS, get_socket_path
+from gpu_memory_service.common.vmm import VMMDeviceType
 
 logger = logging.getLogger(__name__)
 
@@ -23,6 +24,7 @@ class Config:
     alloc_retry_interval: float
     alloc_retry_timeout: Optional[float]
     verbose: bool
+    device_type: VMMDeviceType
 
 
 def parse_args(argv: Optional[list[str]] = None) -> list[Config]:
@@ -71,6 +73,13 @@ def parse_args(argv: Optional[list[str]] = None) -> list[Config]:
         help="Max seconds to wait for allocation retries before failing (default: 60.0). "
         "Pass an explicit large value if you need essentially-unbounded retry.",
     )
+    parser.add_argument(
+        "--device-type",
+        type=str,
+        default=VMMDeviceType.CUDA.value,
+        choices=[d.value for d in VMMDeviceType],
+        help="VMM device type (vendor driver) to use (default: cuda).",
+    )
 
     args = parser.parse_args(argv)
 
@@ -94,6 +103,7 @@ def parse_args(argv: Optional[list[str]] = None) -> list[Config]:
             alloc_retry_interval=args.alloc_retry_interval,
             alloc_retry_timeout=args.alloc_retry_timeout,
             verbose=args.verbose,
+            device_type=VMMDeviceType.from_str(args.device_type),
         )
         for tag in tags
     ]

lib/gpu_memory_service/cli/runner.py

@@ -22,6 +22,7 @@
 from collections.abc import Sequence
 
 import uvloop
+from gpu_memory_service.common.vmm import init_vmm
 from gpu_memory_service.server.rpc import GMSRPCServer
 
 from .args import Config, parse_args
@@ -53,11 +54,13 @@ async def serve_configs(configs: Sequence[Config]) -> None:
         logging.getLogger().setLevel(logging.DEBUG)
         logging.getLogger("gpu_memory_service").setLevel(logging.DEBUG)
 
+    init_vmm(configs[0].device_type)
     servers = []
     for config in configs:
         logger.info("Starting GPU Memory Service Server for device %d", config.device)
         logger.info("GMS tag: %s", config.tag)
         logger.info("Socket path: %s", config.socket_path)
+        logger.info("VMM device type: %s", config.device_type.value)
         logger.info(
             "Allocation retry config: interval=%ss timeout=%s",
             config.alloc_retry_interval,

lib/gpu_memory_service/cli/server.py

@@ -11,13 +11,14 @@
 
 from __future__ import annotations
 
+import argparse
 import logging
 import signal
 import subprocess
 import sys
 import time
 
-from gpu_memory_service.common.cuda_utils import list_devices
+from gpu_memory_service.common.vmm import VMMDeviceType, get_vmm, init_vmm
 
 logging.basicConfig(
     level=logging.INFO,
@@ -26,9 +27,17 @@
 logger = logging.getLogger(__name__)
 
 
-def _child_command(device: int) -> list[str]:
+def _child_command(device: int, device_type: str) -> list[str]:
     """Command for one child process serving every production tag on one GPU."""
-    return [sys.executable, "-m", "gpu_memory_service", "--device", str(device)]
+    return [
+        sys.executable,
+        "-m",
+        "gpu_memory_service",
+        "--device",
+        str(device),
+        "--device-type",
+        device_type,
+    ]
 
 
 def _terminate_all(processes: list[subprocess.Popen]) -> None:
@@ -50,10 +59,31 @@ def _supervise(processes: list[subprocess.Popen]) -> int:
 
 
 def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="GPU Memory Service supervisor (one server per (device, tag))."
+    )
+    parser.add_argument(
+        "--device-type",
+        type=str,
+        default=VMMDeviceType.CUDA.value,
+        choices=[d.value for d in VMMDeviceType],
+        help="VMM device type forwarded to server (default: cuda).",
+    )
+    args = parser.parse_args()
+
+    init_vmm(VMMDeviceType.from_str(args.device_type))
+    vmm = get_vmm()
+    vmm.ensure_initialized()
+    devices = vmm.list_devices()
     processes = []
-    for device in list_devices():
-        proc = subprocess.Popen(_child_command(device))
-        logger.info("Started GMS device=%d pid=%d", device, proc.pid)
+    for device in devices:
+        proc = subprocess.Popen(_child_command(device, args.device_type))
+        logger.info(
+            "Started GMS device=%d device_type=%s pid=%d",
+            device,
+            args.device_type,
+            proc.pid,
+        )
         processes.append(proc)
 
     def terminate(*_args) -> None:

lib/gpu_memory_service/cli/snapshot/loader.py

@@ -18,8 +18,8 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
 
-from gpu_memory_service.common import cuda_utils
 from gpu_memory_service.common.utils import get_socket_path
+from gpu_memory_service.common.vmm import VMMDeviceType, get_vmm, init_vmm
 from gpu_memory_service.snapshot.backends.sharded_ssd import parse_sharded_ssd_roots
 from gpu_memory_service.snapshot.storage_client import GMSStorageClient
 from gpu_memory_service.snapshot.transfer import TransferBackendKind
@@ -52,7 +52,9 @@ def _load_device(
     # GMSStorageClient still publishes the restored layout from this thread.
     # Ensure the loader's main per-device thread has a current CUDA context for
     # the final synchronize/unmap/commit path.
-    cuda_utils.cuda_runtime_set_device(device)
+    vmm = get_vmm()
+    vmm.ensure_initialized()
+    vmm.runtime_set_device(device)
     client = GMSStorageClient(
         socket_path=get_socket_path(device),
         device=device,
@@ -105,11 +107,22 @@ def _build_parser() -> argparse.ArgumentParser:
         default=2,
         help="Number of independent sharded-ssd restore queues per SSD root.",
     )
+    parser.add_argument(
+        "--device-type",
+        type=str,
+        default=VMMDeviceType.CUDA.value,
+        choices=[d.value for d in VMMDeviceType],
+        help="VMM device type (default: cuda).",
+    )
     return parser
 
 
-def _list_checkpoint_devices(checkpoint_dir: str | None) -> list[int]:
-    devices = cuda_utils.list_devices()
+def _list_checkpoint_devices(
+    checkpoint_dir: str | None,
+) -> list[int]:
+    vmm = get_vmm()
+    vmm.ensure_initialized()
+    devices = vmm.list_devices()
     if not checkpoint_dir:
         return devices
 
@@ -155,6 +168,8 @@ def main(argv: list[str] | None = None) -> None:
         parser.error("--sharded-ssd-queues-per-root must be a positive integer")
     checkpoint_dir = args.checkpoint_dir
     max_workers = args.max_workers
+    device_type = VMMDeviceType.from_str(args.device_type)
+    init_vmm(device_type)
     transfer_backend = args.transfer_backend
     sharded_ssd_roots = parse_sharded_ssd_roots(args.sharded_ssd_roots)
     sharded_ssd_queues_per_root = args.sharded_ssd_queues_per_root