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

[zxue2]

Overview:

GMS XPU Enablement Phase 1: add a device-agnostic VMM (Virtual Memory Management) abstraction layer. All existing CUDA paths should be not unchanged. All existing --device-kind cuda deployments should continue to work.
VMMDeviceType.XPU defined but raises NotImplementedError. The vendor-neutral VMMDevice Protocol abstracts the per-device virtual memory management surface.

Please refer to https://github.com/zxue2/dynamo/blob/enable_gms_xpu/lib/gpu_memory_service/GMS_MULTI_DEVICE.md

Details:

VMM abstraction layer:
- common/vmm/device.py: VMMDevice Protocol
- common/vmm/init.py: VMMDeviceType enum + get_vmm_device() factory
- common/vmm/cuda_utils.py: existing cuda_utils relocated with CudaVMM
class implementing VMMDevice via delegation to the module-level
CUDA driver helpers

Server-side:
- cli/args.py: --device-kind argument (default: "cuda")
- cli/server.py: forward --device-kind to spawned processes
- cli/runner.py: pass device_kind to GMSRPCServer
- server/rpc.py -> server/gms.py -> server/allocations.py: pass
device_kind through to GMSAllocationManager which instantiates
the appropriate VMMDevice.

Client-side:
- client/memory_manager.py: accept device_kind, store VMMDevice
instance, expose device_kind property.
- client/torch/allocator.py: _ensure_callbacks_initialized and
_create_mem_pool accept device_kind; gms_use_mem_pool guards
against non-CUDA backends with NotImplementedError.
get_or_create_gms_client_memory_manager and
get_or_create_scratch_manager forward device_kind.

XPU path correctly raises NotImplementedError at factory, allocator,
and mem_pool levels. XPU implementation will be added in Phase 2.

Where should the reviewer start?

Please check if VMM abstraction is suitable for CUDA and multi-device.

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

• closes GitHub issue: DEP: GMS Multiple Device(XPU) Enablement #10087

---

Summary by CodeRabbit

Release Notes

• New Features
  • Added --device-kind CLI option to configure GPU memory management backend selection
  • Extended Python client and server APIs to support device kind configuration across memory management components

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

🌿 Fern Docs Preview: https://nvidia-preview-019f5be3-1b37-766c-a3ac-75210de2b1c8.docs.buildwithfern.com/dynamo/dev

[coderabbitai]

Walkthrough

This PR extends the GPU Memory Service to support configurable VMM (Virtual Memory Management) device kinds. It introduces a vendor-neutral device abstraction protocol, defines a device-type enum with a factory pattern, and threads device-kind parameters through the server, client, CLI, and snapshot utilities, enabling selection between CUDA and other backends while enforcing CUDA-only behavior where required.

Changes

VMM Device Kind Support

Layer / File(s)	Summary
VMM Device Abstraction Protocol and Factory lib/gpu_memory_service/common/vmm/device.py, lib/gpu_memory_service/common/vmm/__init__.py, lib/gpu_memory_service/common/vmm/cuda_utils.py	VMMDevice protocol defines vendor-neutral driver lifecycle, memory allocation, shareable handle, and access-control methods. VMMDeviceType enum with from_str() parser normalizes input. get_vmm_device() factory returns CudaVMM for CUDA, raises NotImplementedError for XPU, and ValueError for unknown kinds. CudaVMM wraps existing CUDA helper functions implementing the protocol.
Server-Side Device Kind Integration lib/gpu_memory_service/server/allocations.py, lib/gpu_memory_service/server/gms.py, lib/gpu_memory_service/server/rpc.py	GMSAllocationManager, GMS, and GMSRPCServer constructors extended with device_kind parameter (defaulting to CUDA). Allocation manager creates VMM instance via get_vmm_device(). All three layers log and forward device kind through the initialization chain.
Client-Side Device Kind Integration lib/gpu_memory_service/client/memory_manager.py, lib/gpu_memory_service/client/torch/allocator.py, lib/gpu_memory_service/snapshot/storage_client.py	GMSClientMemoryManager accepts device_kind, stores it, and initializes VMM via get_vmm_device(). PyTorch allocator helper functions enforce CUDA-only with NotImplementedError for non-CUDA devices. Public factory functions and gms_use_mem_pool() accept optional device_kind and validate CUDA constraints. GMSStorageClient accepts device_kind and passes it to memory manager in both save() and load_to_gms() paths.
CLI and Snapshot Device Kind Configuration lib/gpu_memory_service/cli/args.py, lib/gpu_memory_service/cli/runner.py, lib/gpu_memory_service/cli/server.py, lib/gpu_memory_service/cli/snapshot/saver.py, lib/gpu_memory_service/cli/snapshot/loader.py	CLI argument parser adds --device-kind option with enum-derived choices and VMMDeviceType.from_str() conversion. Runner logs and forwards device kind to GMSRPCServer. Server entry point parses --device-kind and forwards to child processes. Snapshot utilities read GMS_DEVICE_KIND environment variable (defaulting to CUDA), log it, and pass device_kind to helper functions and storage client.
Test Infrastructure tests/report_pytest_markers.py	Pytest marker collection adds gpu_memory_service.common.vmm submodule stubs to allowlist for test isolation.

---

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.93% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	PR description follows template structure with Overview, Details, Where to start, and Related Issues sections. All required information is present and clearly articulated.
Title check	✅ Passed	The title directly and accurately summarizes the main objective: introducing a VMM Device abstraction to support XPU (Phase 1). It matches the core changes across all modified files.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

Warning

Review ran into problems

🔥 Problems

Git: Failed to clone repository. Please run the @coderabbitai full review command to re-trigger a full review. If the issue persists, set path_filters to include or exclude specific files.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 4

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

lib/gpu_memory_service/client/memory_manager.py (1)

39-54: 🛠️ Refactor suggestion | 🟠 Major | 🏗️ Heavy lift

Use self._vmm as the execution path instead of keeping CUDA helpers as the active path.

The class now captures device_kind/_vmm, but core operations still execute via CUDA-specific functions. This keeps behavior CUDA-coupled and undercuts the new abstraction.

Proposed starting point

-        cuda_ensure_initialized()
-        self.granularity = cumem_get_allocation_granularity(device)
+        self._vmm.ensure_initialized()
+        self.granularity = self._vmm.get_allocation_granularity(device)

Then progressively switch map/unmap/import/release/access/sync/validate call sites to self._vmm.* methods to keep backend dispatch consistent end-to-end.

Also applies to: 166-190

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/gpu_memory_service/client/memory_manager.py` around lines 39 - 54, The
code imports and directly calls CUDA-specific helpers (cumem_map, cumem_unmap,
cumem_import_from_shareable_handle_close_fd, cumem_release, cumem_set_access,
cuda_synchronize, cuda_validate_pointer, etc.) instead of routing operations
through the VMM abstraction captured on the instance; update all call sites in
memory_manager.py (including the block around lines 166–190) to invoke the
corresponding methods on self._vmm (e.g., self._vmm.map, self._vmm.unmap,
self._vmm.import_from_shareable_handle_close_fd, self._vmm.release,
self._vmm.set_access, self._vmm.synchronize, self._vmm.validate_pointer or
whatever the VMM API provides) so backend dispatch is consistent end-to-end and
remove or stop using the CUDA-specific helpers in this module.

lib/gpu_memory_service/client/torch/allocator.py (1)

125-131: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Validate device_kind when reusing an existing tag state.

Both reuse paths check socket/device but not backend kind, so a caller can request a different device_kind and silently get the previous manager.

Suggested fix

-        if state.socket_path != socket_path or state.device != device:
+        if (
+            state.socket_path != socket_path
+            or state.device != device
+            or state.manager.device_kind != device_kind
+        ):
             raise RuntimeError(
                 f"GMS allocator tag={tag} was initialized for "
-                f"{state.socket_path} on device {state.device}, not {socket_path} "
-                f"on device {device}"
+                f"{state.socket_path} on device {state.device} "
+                f"(device_kind={state.manager.device_kind.value}), not {socket_path} "
+                f"on device {device} (device_kind={device_kind.value})"
             )

(Apply the same check in get_or_create_scratch_manager.)

Also applies to: 203-209

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/gpu_memory_service/client/torch/allocator.py` around lines 125 - 131, The
reuse check only compares socket_path and device but not backend kind, so update
the reuse guard to also validate state.device_kind matches the requested
device_kind and raise a RuntimeError with the same style if it differs; apply
this change in the allocator reuse branch shown (the block that raises for tag
mismatch) and make the identical addition in get_or_create_scratch_manager (the
reuse path around lines 203-209) so callers cannot silently get a manager for a
different device_kind.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@lib/gpu_memory_service/cli/server.py`:
- Around line 21-23: The CLI currently calls the CUDA-only helper list_devices
to discover devices (references: list_devices) even when the user passed a
non-CUDA --device-kind (references: VMMDeviceType and the CLI device-kind
parsing logic), causing incorrect/empty spawn behavior; change the discovery to
be backend-aware by switching on the parsed device-kind (use VMMDeviceType
enum/value) and calling the appropriate backend discovery routine (or returning
an empty/placeholder list and letting the backend report the unsupported-kind
error) instead of list_devices for non-CUDA kinds, and update all occurrences
where list_devices is used (the device discovery blocks around the CLI parsing
and the spawn/pre-check paths) to this backend-switching approach so device
discovery matches --device-kind.

In `@lib/gpu_memory_service/cli/snapshot/loader.py`:
- Around line 21-23: The loader currently always calls the CUDA-specific
list_devices despite parsing GMS_DEVICE_KIND; update the device enumeration to
branch on the parsed device kind (VMMDeviceType) instead of unconditionally
calling list_devices. In loader.py replace direct list_devices usage with a
dispatch that: checks the device_kind value (VMMDeviceType.CUDA,
VMMDeviceType.ROCm/ROCM, etc.), calls the appropriate enumeration helper for
each kind (add or import a ROCm/other list function if missing), and raises/logs
a clear error for unsupported kinds; apply this change to all places where
list_devices is used in this module (imports at top and usages around the
current discovery logic referenced by VMMDeviceType and GMS_DEVICE_KIND).

In `@lib/gpu_memory_service/cli/snapshot/saver.py`:
- Around line 18-20: The code parses GMS_DEVICE_KIND but always calls
list_devices (CUDA); update the discovery logic in saver.py to branch on
GMS_DEVICE_KIND (the parsed value) using the VMMDeviceType enum: if
GMS_DEVICE_KIND == VMMDeviceType.CUDA, call list_devices as currently done; for
other VMMDeviceType values call the appropriate backend discovery helper (or a
generic discovery API provided by your VMM layer) instead of list_devices, and
if no backend-discovery exists, raise an explicit error/NotImplementedError so
the failure is clear. Ensure you change the code paths where list_devices is
used (references near the current call) to use this conditional logic and keep
VMMDeviceType and GMS_DEVICE_KIND as the selectors.

In `@lib/gpu_memory_service/server/allocations.py`:
- Around line 72-77: The init currently calls CUDA helpers directly
(cuda_ensure_initialized and cumem_get_allocation_granularity) instead of using
the VMM abstraction; update the constructor to route these operations through
the VMM returned by get_vmm_device(device_kind) (self._vmm) so non-CUDA backends
can implement their own init/granularity logic—replace direct calls to
cuda_ensure_initialized() and cumem_get_allocation_granularity(device) with
self._vmm.initialize() (or equivalent VMM init method) and
self._vmm.get_allocation_granularity(device) (or VMM property/method) and make
the same change for the allocation-related calls referenced around lines 81-85
to consistently use self._vmm instead of CUDA helpers.

---

Outside diff comments:
In `@lib/gpu_memory_service/client/memory_manager.py`:
- Around line 39-54: The code imports and directly calls CUDA-specific helpers
(cumem_map, cumem_unmap, cumem_import_from_shareable_handle_close_fd,
cumem_release, cumem_set_access, cuda_synchronize, cuda_validate_pointer, etc.)
instead of routing operations through the VMM abstraction captured on the
instance; update all call sites in memory_manager.py (including the block around
lines 166–190) to invoke the corresponding methods on self._vmm (e.g.,
self._vmm.map, self._vmm.unmap, self._vmm.import_from_shareable_handle_close_fd,
self._vmm.release, self._vmm.set_access, self._vmm.synchronize,
self._vmm.validate_pointer or whatever the VMM API provides) so backend dispatch
is consistent end-to-end and remove or stop using the CUDA-specific helpers in
this module.

In `@lib/gpu_memory_service/client/torch/allocator.py`:
- Around line 125-131: The reuse check only compares socket_path and device but
not backend kind, so update the reuse guard to also validate state.device_kind
matches the requested device_kind and raise a RuntimeError with the same style
if it differs; apply this change in the allocator reuse branch shown (the block
that raises for tag mismatch) and make the identical addition in
get_or_create_scratch_manager (the reuse path around lines 203-209) so callers
cannot silently get a manager for a different device_kind.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 642eaa55-5176-420f-a960-98912628dd7e

📥 Commits

Reviewing files that changed from the base of the PR and between cd3c322 and b658cc4.

📒 Files selected for processing (15)

• lib/gpu_memory_service/cli/args.py
• lib/gpu_memory_service/cli/runner.py
• lib/gpu_memory_service/cli/server.py
• lib/gpu_memory_service/cli/snapshot/loader.py
• lib/gpu_memory_service/cli/snapshot/saver.py
• lib/gpu_memory_service/client/memory_manager.py
• lib/gpu_memory_service/client/torch/allocator.py
• lib/gpu_memory_service/common/vmm/__init__.py
• lib/gpu_memory_service/common/vmm/cuda_utils.py
• lib/gpu_memory_service/common/vmm/device.py
• lib/gpu_memory_service/server/allocations.py
• lib/gpu_memory_service/server/gms.py
• lib/gpu_memory_service/server/rpc.py
• lib/gpu_memory_service/snapshot/storage_client.py
• tests/report_pytest_markers.py

[zxue2]

@dzier, per meeting discussion pls help involve gms team to check if VMMDevice abstraction proposal is OK for cuda. Then we can plug in xpu.

[dagil-nvidia]

/ok to test 578e84d

[GuanLuo]

/ok to test d01df1c

[hhzhang16]

nit: I find it a bit confusing how the "initialization"/what it represents in init_vmm and vmm.ensure_initialized are different

[GuanLuo]

/ok to test 4decbe1

[zxue2]

@GuanLuo @hhzhang16 pls kindly have a look at two commits as well. Please also trigger ok test if no concern. thx

replace torch.cuda.synchronize with get_vmm().synchronize

add torch_device helper for integration

[hhzhang16]

/ok to test e11e37a

[hhzhang16]

/ok to test 1819dc7

[zxue2]

/ok to test 1819dc7

Hi @GuanLuo @hhzhang16 , it seems CI passed with 32 skipped ,could you pls help review and check the next step? thx

[hhzhang16]

LGTM

[zxue2]

LGTM

Hi @hhzhang16, I rebased it against latest main again. Please help check /ok test. thx

[hhzhang16]

/ok to test cf51474

[zxue2]

Hi @hhzhang16 @GuanLuo , could you pls share if any issue blocks merging this PR? As to reduce the continuous rebase efforts, could you pls check if we can merge it now? thx

[nnshah1]

unblocking - changes focused on gpu memory service - and @hhzhang16 has approved.