From ada5e7f0034b5e4294bdf77c1c2da8e3091e6fc9 Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Tue, 3 Nov 2026 23:53:25 -0800 Subject: [PATCH] feat(kvbm): bind typed CUDA and NVML fabric observation Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../vllm/tests/test_vllm_mnnvl_memory.py | 187 ++++++++++++ .../python/kvbm/vllm_integration/gpu_mnnvl.py | 273 ++++++++++++++++++ 2 files changed, 460 insertions(+) diff --git a/components/src/dynamo/vllm/tests/test_vllm_mnnvl_memory.py b/components/src/dynamo/vllm/tests/test_vllm_mnnvl_memory.py index ac19f299dc9e..68110cd36cec 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_mnnvl_memory.py +++ b/components/src/dynamo/vllm/tests/test_vllm_mnnvl_memory.py @@ -2,8 +2,10 @@ # SPDX-License-Identifier: Apache-2.0 """No-GPU VMM ownership, admission and rollback tests with a fake CUDA driver.""" +import ctypes import importlib.util import sys +from contextlib import contextmanager from pathlib import Path import pytest @@ -45,6 +47,185 @@ ] +class FakeDriver: + def __init__(self): + self.calls = [] + self.fail = None + self.identity = m.FabricIdentity( + "GPU-00000000-0000-0000-0000-000000000001", "1" * 32, 0 + ) + self.properties_from_handle = self.properties(0) + + def record(self, name, *args): + self.calls.append((name, *args)) + if self.fail == name: + self.fail = None + raise RuntimeError("injected " + name) + + def device(self): + return 0 + + def gpu_uuid(self, device): + return self.identity.gpu_uuid + + def retain_context(self, device): + self.record("retain", device) + return 77 + + @contextmanager + def context(self, context): + self.record("push", context) + try: + yield + finally: + self.record("pop", context) + + def release_context(self, device): + self.record("release-context", device) + + def properties(self, device): + return m._CudaDriver.properties(None, device) + + def granularity(self, properties): + self.record( + "granularity", + properties.type, + properties.requestedHandleTypes, + properties.location.type, + ) + return 8 + + def create(self, size, properties): + self.record("create", size, properties.requestedHandleTypes) + return 90 + + def allocation_properties(self, handle): + self.record("properties", handle) + return self.properties_from_handle + + def reserve(self, size, granularity): + self.record("reserve", size, granularity) + return 4096 + + def map(self, pointer, size, handle): + self.record("map", pointer, size, handle) + + def set_access(self, pointer, size, device): + self.record("access", pointer, size, device) + + def zero(self, pointer, size): + self.record("zero", pointer, size) + + def synchronize(self): + self.record("sync") + + def export(self, handle): + self.record("export", handle) + return b"h" * 64 + + def import_handle(self, value): + if type(value) is not bytes or len(value) != 64: + raise ValueError("exact64 bytes") + self.record("import", value) + return 91 + + def unmap(self, pointer, size): + self.record("unmap", pointer, size) + + def free_address(self, pointer, size): + self.record("free", pointer, size) + + def release(self, handle): + self.record("release", handle) + + +class SignatureFunction: + """Model ctypes default-int conversion when a fresh function has no ABI.""" + + def __init__(self, library, name): + self.library, self.name = library, name + self.argtypes = None + self.restype = ctypes.c_int + + def __call__(self, *arguments): + converted = [] + for index, value in enumerate(arguments): + if isinstance(value, (int, ctypes.c_void_p)): + scalar = value if isinstance(value, int) else value.value + kind = ctypes.c_int if self.argtypes is None else self.argtypes[index] + converted.append(kind(scalar).value) + else: + converted.append(value) + self.library.calls.append((self.name, self, converted)) + if self.name == "nvmlDeviceGetHandleByUUID": + arguments[1]._obj.value = self.library.high_handle + elif self.name == "cuMemImportFromShareableHandle": + arguments[0]._obj.value = self.library.high_handle + elif self.name == "nvmlDeviceGetGpuFabricInfoV": + info = arguments[1]._obj + info.state, info.status, info.cliqueId = 3, 0, 0 + info.clusterUuid[:] = b"\x11" * 16 + return 0 + + +class SignatureLibrary: + """CDLL indexing returns fresh functions; attributes cache each function.""" + + high_handle = 0x123456789ABC + + def __init__(self): + self.functions, self.calls, self.index_lookups = {}, [], [] + + def __getitem__(self, name): + self.index_lookups.append(name) + return SignatureFunction(self, name) + + def __getattr__(self, name): + if name not in self.functions: + self.functions[name] = SignatureFunction(self, name) + return self.functions[name] + + +def test_cuda_driver_reuses_declared_signatures_for_high_64bit_addresses(monkeypatch): + library = SignatureLibrary() + monkeypatch.setattr(m.ctypes, "CDLL", lambda name: library) + driver = m._CudaDriver() + pointer, handle, extent = 0x123456789ABC, 0x23456789ABCD, 1 << 33 + driver.map(pointer, extent, handle) + driver.free_address(pointer, extent) + driver.release(handle) + assert not library.index_lookups + by_name = { + name: (function, arguments) for name, function, arguments in library.calls + } + assert by_name["cuMemMap"][1] == [pointer, extent, 0, handle, 0] + assert by_name["cuMemAddressFree"][1] == [pointer, extent] + assert by_name["cuMemRelease"][1] == [handle] + for name, (function, _) in by_name.items(): + assert function is library.functions[name] + assert function.argtypes is not None + + +def test_nvml_observation_uses_configured_cached_functions_for_device_handle( + monkeypatch, +): + library = SignatureLibrary() + monkeypatch.setattr(m.ctypes, "CDLL", lambda name: library) + driver = FakeDriver() + identity = m.observe_fabric(device=0, driver=driver) + assert identity == driver.identity + assert not library.index_lookups + query = [ + (function, values) + for name, function, values in library.calls + if name == "nvmlDeviceGetGpuFabricInfoV" + ] + assert len(query) == 1 and query[0][1][0] == library.high_handle + for name, function, _ in library.calls: + assert function is library.functions[name] + assert function.argtypes is not None + + def test_effective_domain_matches_same_clique_not_gpu_uuid_and_separates_other_profiles(): first = m.FabricIdentity("GPU-00000000-0000-0000-0000-000000000001", "1" * 32, 0) second = m.FabricIdentity("GPU-00000000-0000-0000-0000-000000000002", "1" * 32, 0) @@ -72,6 +253,12 @@ def test_uncanonical_or_missing_fabric_domain_is_not_a_capability(domain): m.effective_domain("layout-v1", domain) +def test_ctypes_layout_matches_pinned_cuda_and_nvml_struct_declarations(): + assert ctypes.sizeof(m._AllocationProperties) == 32 + assert ctypes.sizeof(m._Access) == 12 + assert ctypes.sizeof(m._FabricInfo) == 36 + + @pytest.mark.parametrize("layout_domain", ["bad\x00layout", "", " padded ", None]) def test_invalid_raw_layout_domain_cannot_be_hidden_by_compatibility_hash( layout_domain, diff --git a/lib/bindings/kvbm/python/kvbm/vllm_integration/gpu_mnnvl.py b/lib/bindings/kvbm/python/kvbm/vllm_integration/gpu_mnnvl.py index 0c0c24174d98..8e67de9dcc6d 100644 --- a/lib/bindings/kvbm/python/kvbm/vllm_integration/gpu_mnnvl.py +++ b/lib/bindings/kvbm/python/kvbm/vllm_integration/gpu_mnnvl.py @@ -7,11 +7,69 @@ the owner must outlive every local view and all native transfer leases. """ +import ctypes import hashlib import json import uuid +from contextlib import contextmanager from dataclasses import dataclass +_FABRIC = 8 + + +_DEVICE = 1 + + +_PINNED = 1 + + +_READ_WRITE = 3 + + +_FABRIC_HANDLE_BYTES = 64 + + +_MAX_BYTES = min((1 << (ctypes.sizeof(ctypes.c_size_t) * 8)) - 1, (1 << 63) - 1) + + +class _Location(ctypes.Structure): + _fields_ = [("type", ctypes.c_int), ("id", ctypes.c_int)] + + +class _AllocationFlags(ctypes.Structure): + _fields_ = [ + ("compressionType", ctypes.c_ubyte), + ("gpuDirectRDMACapable", ctypes.c_ubyte), + ("usage", ctypes.c_ushort), + ("reserved", ctypes.c_ubyte * 4), + ] + + +class _AllocationProperties(ctypes.Structure): + _fields_ = [ + ("type", ctypes.c_int), + ("requestedHandleTypes", ctypes.c_int), + ("location", _Location), + ("win32HandleMetaData", ctypes.c_void_p), + ("allocFlags", _AllocationFlags), + ] + + +class _Access(ctypes.Structure): + _fields_ = [("location", _Location), ("flags", ctypes.c_int)] + + +class _FabricInfo(ctypes.Structure): + # NVML nvmlGpuFabricInfo_v2_t, also present in nvidia-ml-py13.580.65. + _fields_ = [ + ("version", ctypes.c_uint), + ("clusterUuid", ctypes.c_ubyte * 16), + ("status", ctypes.c_uint), + ("cliqueId", ctypes.c_uint), + ("state", ctypes.c_uint), + ("healthMask", ctypes.c_uint), + ] + def _gpu_uuid(value): if not isinstance(value, str): @@ -78,3 +136,218 @@ def effective_domain(layout_domain, fabric_domain): json.dumps(body, sort_keys=True, separators=(",", ":")).encode() ).hexdigest() ) + + +class _CudaDriver: + def __init__(self): + if ctypes.sizeof(ctypes.c_void_p) != 8: + raise RuntimeError("MNNVL VMM requires a64-bit CUDA process") + self.lib = ctypes.CDLL("libcuda.so.1") + ptr, size, handle, device = ( + ctypes.c_ulonglong, + ctypes.c_size_t, + ctypes.c_ulonglong, + ctypes.c_int, + ) + signatures = { + "cuInit": [ctypes.c_uint], + "cuDeviceGet": [ctypes.POINTER(device), ctypes.c_int], + "cuCtxGetDevice": [ctypes.POINTER(device)], + "cuDeviceGetUuid_v2": [ctypes.c_void_p, device], + "cuDevicePrimaryCtxRetain": [ctypes.POINTER(ctypes.c_void_p), device], + "cuDevicePrimaryCtxRelease_v2": [device], + "cuCtxPushCurrent_v2": [ctypes.c_void_p], + "cuCtxPopCurrent_v2": [ctypes.POINTER(ctypes.c_void_p)], + "cuCtxSynchronize": [], + "cuMemGetAllocationGranularity": [ + ctypes.POINTER(size), + ctypes.POINTER(_AllocationProperties), + ctypes.c_int, + ], + "cuMemCreate": [ + ctypes.POINTER(handle), + size, + ctypes.POINTER(_AllocationProperties), + ctypes.c_ulonglong, + ], + "cuMemGetAllocationPropertiesFromHandle": [ + ctypes.POINTER(_AllocationProperties), + handle, + ], + "cuMemAddressReserve": [ + ctypes.POINTER(ptr), + size, + size, + ptr, + ctypes.c_ulonglong, + ], + "cuMemMap": [ptr, size, size, handle, ctypes.c_ulonglong], + "cuMemSetAccess": [ptr, size, ctypes.POINTER(_Access), size], + "cuMemsetD8_v2": [ptr, ctypes.c_ubyte, size], + "cuMemExportToShareableHandle": [ + ctypes.c_void_p, + handle, + ctypes.c_int, + ctypes.c_ulonglong, + ], + "cuMemImportFromShareableHandle": [ + ctypes.POINTER(handle), + ctypes.c_void_p, + ctypes.c_int, + ], + "cuMemUnmap": [ptr, size], + "cuMemAddressFree": [ptr, size], + "cuMemRelease": [handle], + } + for name, arguments in signatures.items(): + function = getattr(self.lib, name) + function.argtypes, function.restype = arguments, ctypes.c_int + self.call("cuInit", 0) + + def call(self, name, *args): + result = getattr(self.lib, name)(*args) + if result != 0: + raise RuntimeError( + f"MNNVL CUDA {name} failed ({result}); FABRIC/IMEX access is required; no fallback" + ) + + def device(self): + value = ctypes.c_int() + result = self.lib.cuCtxGetDevice(ctypes.byref(value)) + if result == 201: # No current context in the TP1 registration process. + self.call("cuDeviceGet", ctypes.byref(value), 0) + elif result != 0: + raise RuntimeError(f"MNNVL current CUDA device query failed ({result})") + return value.value + + def gpu_uuid(self, device): + value = (ctypes.c_ubyte * 16)() + self.call("cuDeviceGetUuid_v2", ctypes.byref(value), device) + return "GPU-" + str(uuid.UUID(bytes=bytes(value))) + + def retain_context(self, device): + value = ctypes.c_void_p() + self.call("cuDevicePrimaryCtxRetain", ctypes.byref(value), device) + return value.value + + @contextmanager + def context(self, context): + self.call("cuCtxPushCurrent_v2", context) + try: + yield + finally: + previous = ctypes.c_void_p() + self.call("cuCtxPopCurrent_v2", ctypes.byref(previous)) + + def release_context(self, device): + self.call("cuDevicePrimaryCtxRelease_v2", device) + + def properties(self, device): + value = _AllocationProperties() + value.type, value.requestedHandleTypes = _PINNED, _FABRIC + value.location.type, value.location.id = _DEVICE, device + return value + + def granularity(self, properties): + value = ctypes.c_size_t() + self.call( + "cuMemGetAllocationGranularity", + ctypes.byref(value), + ctypes.byref(properties), + 0, + ) + return value.value + + def create(self, size, properties): + value = ctypes.c_ulonglong() + self.call("cuMemCreate", ctypes.byref(value), size, ctypes.byref(properties), 0) + return value.value + + def allocation_properties(self, handle): + value = _AllocationProperties() + self.call("cuMemGetAllocationPropertiesFromHandle", ctypes.byref(value), handle) + return value + + def reserve(self, size, granularity): + value = ctypes.c_ulonglong() + self.call("cuMemAddressReserve", ctypes.byref(value), size, granularity, 0, 0) + return value.value + + def map(self, pointer, size, handle): + self.call("cuMemMap", pointer, size, 0, handle, 0) + + def set_access(self, pointer, size, device): + value = _Access(_Location(_DEVICE, device), _READ_WRITE) + self.call("cuMemSetAccess", pointer, size, ctypes.byref(value), 1) + + def zero(self, pointer, size): + self.call("cuMemsetD8_v2", pointer, 0, size) + + def synchronize(self): + self.call("cuCtxSynchronize") + + def export(self, handle): + value = (ctypes.c_ubyte * _FABRIC_HANDLE_BYTES)() + self.call( + "cuMemExportToShareableHandle", ctypes.byref(value), handle, _FABRIC, 0 + ) + return bytes(value) + + def import_handle(self, value): + if not isinstance(value, bytes) or len(value) != _FABRIC_HANDLE_BYTES: + raise ValueError("Exactly64 opaque FABRIC handle bytes required") + raw = (ctypes.c_ubyte * _FABRIC_HANDLE_BYTES).from_buffer_copy(value) + handle = ctypes.c_ulonglong() + self.call( + "cuMemImportFromShareableHandle", + ctypes.byref(handle), + ctypes.byref(raw), + _FABRIC, + ) + return handle.value + + def unmap(self, pointer, size): + self.call("cuMemUnmap", pointer, size) + + def free_address(self, pointer, size): + self.call("cuMemAddressFree", pointer, size) + + def release(self, handle): + self.call("cuMemRelease", handle) + + +def observe_fabric(device=None, *, driver=None): + """Observe actual assigned GPU/NVML fabric; device presence is not IMEX proof.""" + driver = _CudaDriver() if driver is None else driver + device = driver.device() if device is None else device + gpu = driver.gpu_uuid(device) + library = ctypes.CDLL("libnvidia-ml.so.1") + signatures = { + "nvmlInit_v2": [], + "nvmlShutdown": [], + "nvmlDeviceGetHandleByUUID": [ctypes.c_char_p, ctypes.POINTER(ctypes.c_void_p)], + "nvmlDeviceGetGpuFabricInfoV": [ctypes.c_void_p, ctypes.POINTER(_FabricInfo)], + } + for name, arguments in signatures.items(): + function = getattr(library, name) + function.argtypes, function.restype = arguments, ctypes.c_uint + if library.nvmlInit_v2() != 0: + raise RuntimeError("NVML initialization is required for MNNVL") + try: + handle = ctypes.c_void_p() + if library.nvmlDeviceGetHandleByUUID(gpu.encode(), ctypes.byref(handle)) != 0: + raise RuntimeError("Actual assigned GPU UUID is absent from NVML") + info = _FabricInfo() + info.version = 0x02000024 + if ( + library.nvmlDeviceGetGpuFabricInfoV(handle, ctypes.byref(info)) != 0 + or info.state != 3 + or info.status != 0 + ): + raise RuntimeError("MNNVL requires actual NVML fabric Completed/Success") + return FabricIdentity( + gpu, bytes(info.clusterUuid).hex(), int(info.cliqueId), int(info.healthMask) + ) + finally: + if library.nvmlShutdown() != 0: + raise RuntimeError("NVML shutdown failed after MNNVL fabric observation")