From 7fb6e61b953e6598aa6b0bffad4c9db0435e734e Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Wed, 12 Aug 2026 06:16:34 +0800 Subject: [PATCH] Fix CUDA 13.0 VMM handle type compatibility (#34431) --- python/sglang/srt/cuda_vmm_utils.py | 22 ++++++++---- test/registered/unit/test_cuda_vmm_utils.py | 40 +++++++++++++++++++++ 2 files changed, 55 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/cuda_vmm_utils.py b/python/sglang/srt/cuda_vmm_utils.py index 06dbc8abd..79c14a762 100644 --- a/python/sglang/srt/cuda_vmm_utils.py +++ b/python/sglang/srt/cuda_vmm_utils.py @@ -281,20 +281,28 @@ def make_device_allocation_prop( elif not isinstance(handle_types, int): raise ValueError("handle_types must be 'auto', an integer, or None") - handle_types = int(handle_types) + handle_type_value = int(handle_types) valid_handle_types = { - int(drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE), - int(drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR), - int(drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC), + int(drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE): ( + drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE + ), + int(drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR): ( + drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR + ), + int(drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC): ( + drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC + ), } - if handle_types not in valid_handle_types: - raise ValueError(f"invalid CUDA handle-type value: {handle_types}") + if handle_type_value not in valid_handle_types: + raise ValueError(f"invalid CUDA handle-type value: {handle_type_value}") prop = drv.CUmemAllocationProp() prop.type = drv.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED prop.location.type = drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE prop.location.id = int(device_id) - prop.requestedHandleTypes = handle_types + # cuda-bindings 13.0.x requires the generated enum here; newer releases + # also accept a plain int, which previously hid this compatibility issue. + prop.requestedHandleTypes = valid_handle_types[handle_type_value] prop.allocFlags.gpuDirectRDMACapable = int(gpu_direct_rdma) return prop diff --git a/test/registered/unit/test_cuda_vmm_utils.py b/test/registered/unit/test_cuda_vmm_utils.py index a11086b8c..776898f53 100644 --- a/test/registered/unit/test_cuda_vmm_utils.py +++ b/test/registered/unit/test_cuda_vmm_utils.py @@ -159,6 +159,46 @@ def test_default_handle_type_fallback(monkeypatch, rejected, expected) -> None: get_device_allocation_handle_type.cache_clear() +@pytest.mark.parametrize( + ("handle_types", "expected"), + [ + (None, drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE), + (0, drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE), + (_POSIX_FD, _POSIX_FD), + (_FABRIC, _FABRIC), + ], +) +def test_allocation_prop_assigns_handle_type_enum( + monkeypatch, handle_types, expected +) -> None: + """CUDA bindings 13.0.x reject integer VMM handle types at assignment.""" + + class Fields: + pass + + class StrictAllocationProp: + def __init__(self): + self.location = Fields() + self.allocFlags = Fields() + self._requested_handle_types = None + + @property + def requestedHandleTypes(self): + return self._requested_handle_types + + @requestedHandleTypes.setter + def requestedHandleTypes(self, value): + if not isinstance(value, drv.CUmemAllocationHandleType): + raise TypeError("requestedHandleTypes requires a CUDA enum") + self._requested_handle_types = value + + monkeypatch.setattr(drv, "CUmemAllocationProp", StrictAllocationProp) + + prop = make_device_allocation_prop(0, handle_types=handle_types) + + assert prop.requestedHandleTypes is expected + + def test_granularity_defaults_to_recommended(monkeypatch) -> None: prop = make_device_allocation_prop(0, handle_types=None) seen = []