Skip to content

Commit 5b7cd96

Browse files
committed
fix(cuda.core): infer NUMA pool support before CUDA 12.9
Signed-off-by: Aryan <[email protected]>
1 parent e52d932 commit 5b7cd96

4 files changed

Lines changed: 68 additions & 6 deletions

File tree

‎cuda_core/cuda/core/_device.pyi‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -447,7 +447,10 @@ class DeviceProperties:
447447
"""bool: Device supports HOST_NUMA location with the virtual memory management APIs like cuMemCreate, cuMemMap and related APIs."""
448448
@property
449449
def host_numa_memory_pools_supported(self) -> bool:
450-
"""bool: Device supports HOST_NUMA location with the cuMemAllocAsync and cuMemPool family of APIs."""
450+
"""bool: Device supports HOST_NUMA location with the cuMemAllocAsync and cuMemPool family of APIs.
451+
452+
The value is inferred for drivers before CUDA 12.9.
453+
"""
451454
@property
452455
def host_numa_multinode_ipc_supported(self) -> bool:
453456
"""bool: Device supports HOST_NUMA location IPC between nodes in a multi-node system."""

‎cuda_core/cuda/core/_device.pyx‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ from cuda.core._rt cimport (
2828
get_last_error,
2929
as_cu,
3030
)
31+
from cuda.core._utils.version cimport cy_driver_version
3132

3233
from cuda.core._stream import IsStreamType, Stream, StreamOptions
3334
from cuda.core._utils.clear_error_support import assert_type
@@ -910,10 +911,19 @@ cdef class DeviceProperties:
910911

911912
@property
912913
def host_numa_memory_pools_supported(self) -> bool:
913-
"""bool: Device supports HOST_NUMA location with the cuMemAllocAsync and cuMemPool family of APIs."""
914-
return bool(
915-
self._get_cached_attribute(driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_HOST_NUMA_MEMORY_POOLS_SUPPORTED)
916-
)
914+
"""bool: Device supports HOST_NUMA location with the cuMemAllocAsync and cuMemPool family of APIs.
915+
916+
The value is inferred for drivers before CUDA 12.9.
917+
"""
918+
cdef tuple drv_ver = cy_driver_version()
919+
if drv_ver >= (12, 9, 0):
920+
return bool(
921+
self._get_cached_attribute(
922+
driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_HOST_NUMA_MEMORY_POOLS_SUPPORTED
923+
)
924+
)
925+
# The attribute arrived in CUDA 12.9, HOST_NUMA pools in CUDA 12.6 (#2795).
926+
return drv_ver >= (12, 6, 0) and self.memory_pools_supported and not self.integrated
917927

918928
# Start of CUDA 13 device attributes
919929

‎cuda_core/cuda/core/_memory/_pinned_memory_resource.pyx‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,13 @@ cdef class PinnedMemoryResource(_MemPool):
121121
raise TypeError("Cannot allocate from a mapped IPC-enabled memory resource")
122122
cdef Stream s = Stream_accept(stream)
123123
device = s.device
124-
if self._numa_id < 0 and not device.properties.host_memory_pools_supported:
124+
cdef bint supported = (
125+
device.properties.host_numa_memory_pools_supported
126+
if self._numa_id >= 0
127+
else device.properties.host_memory_pools_supported
128+
)
129+
130+
if not supported:
125131
raise RuntimeError(
126132
f"CUDA device {device.device_id} does not support the requested "
127133
"host memory pool for PinnedMemoryResource. Use "

‎cuda_core/tests/test_device.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import contextlib
55
import subprocess
66
import sys
7+
import textwrap
78
from concurrent.futures import ThreadPoolExecutor
89

910
import pytest
@@ -471,6 +472,48 @@ def test_device_properties_complete():
471472
assert tab_props == live_props # Ensure exact match.
472473

473474

475+
@pytest.mark.agent_authored(model="gpt-6")
476+
@pytest.mark.parametrize("driver_version", [12050, 12060, 12080])
477+
def test_host_numa_memory_pools_supported_on_pre_12_9_driver(driver_version):
478+
script = textwrap.dedent(
479+
"""
480+
import sys
481+
482+
from cuda.bindings import driver
483+
484+
simulated_driver_version = int(sys.argv[1])
485+
driver.cuDriverGetVersion = lambda: (driver.CUresult.CUDA_SUCCESS, simulated_driver_version)
486+
(status,) = driver.cuInit(0)
487+
if int(status) == int(driver.CUresult.CUDA_ERROR_NO_DEVICE):
488+
print("No CUDA device is available")
489+
raise SystemExit(77)
490+
assert int(status) == int(driver.CUresult.CUDA_SUCCESS), status
491+
status, device = driver.cuDeviceGet(0)
492+
assert int(status) == int(driver.CUresult.CUDA_SUCCESS), status
493+
494+
from cuda.core._device import DeviceProperties
495+
496+
properties = DeviceProperties._init(device)
497+
expected = (
498+
False
499+
if simulated_driver_version < 12060
500+
else properties.memory_pools_supported and not properties.integrated
501+
)
502+
assert properties.host_numa_memory_pools_supported is expected
503+
"""
504+
)
505+
result = subprocess.run( # noqa: S603 - executes a repository-owned test script
506+
[sys.executable, "-c", script, str(driver_version)],
507+
capture_output=True,
508+
check=False,
509+
text=True,
510+
timeout=60,
511+
)
512+
if result.returncode == 77:
513+
pytest.skip(result.stdout + result.stderr)
514+
assert result.returncode == 0, result.stdout + result.stderr
515+
516+
474517
# ============================================================================
475518
# Device Equality Tests
476519
# ============================================================================

0 commit comments

Comments
 (0)