Fix CUDA 13.0 VMM handle type compatibility (#34431)

This commit is contained in:
Mohammad Miadh Angkad
2026-08-11 15:16:34 -07:00
committed by GitHub
parent 9ced8d0981
commit 7fb6e61b95
2 changed files with 55 additions and 7 deletions
+15 -7
View File
@@ -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
@@ -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 = []