Skip to content

Commit 4b9fc22

Browse files
committed
Merge remote-tracking branch 'origin/main' into ajost/vmm-fabric
2 parents 196eeaa + 8d2ec27 commit 4b9fc22

7 files changed

Lines changed: 176 additions & 12 deletions

File tree

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

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,29 @@ class PinnedMemoryResource(_MemPool):
8484
"""
8585
def __init__(self, options: PinnedMemoryResourceOptions | None=None) -> None: ...
8686
def allocate(self, size: int, *, stream: Stream | GraphBuilder) -> Buffer:
87-
"""Allocate a host-pinned buffer asynchronously on the supplied stream."""
87+
"""Allocate a host-pinned buffer asynchronously on the supplied stream.
88+
89+
Parameters
90+
----------
91+
size : int
92+
The size of the buffer to allocate, in bytes.
93+
stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
94+
Keyword-only. The stream on which to perform the allocation
95+
asynchronously. Must be passed explicitly; pass
96+
``device.default_stream`` to use the default stream.
97+
98+
Returns
99+
-------
100+
Buffer
101+
The allocated buffer object, which is accessible from the host.
102+
103+
Raises
104+
------
105+
RuntimeError
106+
If the stream's device does not support the requested host memory
107+
pool. Use :class:`LegacyPinnedMemoryResource` when stream-ordered
108+
allocation is not required.
109+
"""
88110
def __reduce__(self) -> tuple[object, ...]: ...
89111
@staticmethod
90112
def from_registry(uuid: uuid.UUID) -> PinnedMemoryResource:

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

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -113,7 +113,29 @@ cdef class PinnedMemoryResource(_MemPool):
113113
_PMR_init(self, options)
114114

115115
def allocate(self, size_t size, *, stream: Stream | GraphBuilder) -> Buffer:
116-
"""Allocate a host-pinned buffer asynchronously on the supplied stream."""
116+
"""Allocate a host-pinned buffer asynchronously on the supplied stream.
117+
118+
Parameters
119+
----------
120+
size : int
121+
The size of the buffer to allocate, in bytes.
122+
stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
123+
Keyword-only. The stream on which to perform the allocation
124+
asynchronously. Must be passed explicitly; pass
125+
``device.default_stream`` to use the default stream.
126+
127+
Returns
128+
-------
129+
Buffer
130+
The allocated buffer object, which is accessible from the host.
131+
132+
Raises
133+
------
134+
RuntimeError
135+
If the stream's device does not support the requested host memory
136+
pool. Use :class:`LegacyPinnedMemoryResource` when stream-ordered
137+
allocation is not required.
138+
"""
117139
MP_check_open(self)
118140
if self.is_mapped:
119141
raise TypeError("Cannot allocate from a mapped IPC-enabled memory resource")

‎cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,9 @@ class VirtualMemoryResourceOptions:
3232
allocation_type: :obj:`~_memory.VirtualMemoryAllocationType` | str
3333
Controls the type of allocation.
3434
location_type: :obj:`~_memory.VirtualMemoryLocationType` | str
35-
Controls the location of the allocation.
35+
Controls the location of the allocation. ``"host_numa"`` requires
36+
``numa_id``. ``"host_numa_current"`` is rejected because the driver
37+
does not accept it for physical allocations.
3638
handle_type: :obj:`~_memory.VirtualMemoryHandleType` | str
3739
Export handle type for the physical allocation, and the switch for
3840
sharing buffers with other processes (see
@@ -64,6 +66,12 @@ class VirtualMemoryResourceOptions:
6466
Access flags for the owning device.
6567
peer_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str
6668
Access flags for peers.
69+
numa_id: int | None
70+
Host NUMA node of the allocation. Required for
71+
``location_type="host_numa"`` and rejected for every other location
72+
type. Must be >= 0; ``Device.properties.host_numa_id`` is the node
73+
closest to a device. On a system without NUMA, ``host_numa_id`` is -1
74+
and the driver expects ``numa_id=0``.
6775
"""
6876
allocation_type: VirtualMemoryAllocationType = VirtualMemoryAllocationType.PINNED
6977
location_type: VirtualMemoryLocationType = VirtualMemoryLocationType.DEVICE
@@ -75,6 +83,7 @@ class VirtualMemoryResourceOptions:
7583
peers: Iterable[int] = field(default_factory=tuple)
7684
self_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE
7785
peer_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE
86+
numa_id: int | None = None
7887
_a = driver.CUmemAccess_flags
7988
_access_flags = {VirtualMemoryAccessType.READ_WRITE: _a.CU_MEM_ACCESS_FLAGS_PROT_READWRITE, VirtualMemoryAccessType.READ: _a.CU_MEM_ACCESS_FLAGS_PROT_READ, None: 0}
8089
_h = driver.CUmemAllocationHandleType

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

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,9 @@ class VirtualMemoryResourceOptions:
105105
allocation_type: :obj:`~_memory.VirtualMemoryAllocationType` | str
106106
Controls the type of allocation.
107107
location_type: :obj:`~_memory.VirtualMemoryLocationType` | str
108-
Controls the location of the allocation.
108+
Controls the location of the allocation. ``"host_numa"`` requires
109+
``numa_id``. ``"host_numa_current"`` is rejected because the driver
110+
does not accept it for physical allocations.
109111
handle_type: :obj:`~_memory.VirtualMemoryHandleType` | str
110112
Export handle type for the physical allocation, and the switch for
111113
sharing buffers with other processes (see
@@ -137,6 +139,12 @@ class VirtualMemoryResourceOptions:
137139
Access flags for the owning device.
138140
peer_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str
139141
Access flags for peers.
142+
numa_id: int | None
143+
Host NUMA node of the allocation. Required for
144+
``location_type="host_numa"`` and rejected for every other location
145+
type. Must be >= 0; ``Device.properties.host_numa_id`` is the node
146+
closest to a device. On a system without NUMA, ``host_numa_id`` is -1
147+
and the driver expects ``numa_id=0``.
140148
"""
141149

142150
allocation_type: VirtualMemoryAllocationType = VirtualMemoryAllocationType.PINNED
@@ -149,6 +157,7 @@ class VirtualMemoryResourceOptions:
149157
peers: Iterable[int] = field(default_factory=tuple)
150158
self_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE
151159
peer_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE
160+
numa_id: int | None = None
152161

153162
_a = driver.CUmemAccess_flags
154163
_access_flags = { # noqa: RUF012
@@ -541,6 +550,29 @@ cdef class VirtualMemoryResource(MemoryResource):
541550
f"config.location_type {str(cfg.location_type)!r} does not match the resource's "
542551
f"{str(self.config.location_type)!r}; the location of a buffer cannot change"
543552
)
553+
if cfg.numa_id != self.config.numa_id:
554+
raise ValueError(
555+
f"config.numa_id {cfg.numa_id!r} does not match the resource's "
556+
f"{self.config.numa_id!r}; the location of a buffer cannot change"
557+
)
558+
if cfg.location_type == VirtualMemoryLocationType.HOST_NUMA_CURRENT:
559+
raise ValueError(
560+
'location_type="host_numa_current" is not supported: the driver does not accept it '
561+
'for physical allocations; use location_type="host_numa" with numa_id'
562+
)
563+
if cfg.location_type == VirtualMemoryLocationType.HOST_NUMA:
564+
if cfg.numa_id is None:
565+
raise ValueError(
566+
'location_type="host_numa" requires numa_id; pass numa_id=<host NUMA node id>, '
567+
"e.g. device.properties.host_numa_id"
568+
)
569+
if cfg.numa_id < 0:
570+
raise ValueError(f"numa_id must be >= 0, got {cfg.numa_id}")
571+
elif cfg.numa_id is not None:
572+
raise ValueError(
573+
'numa_id is only valid with location_type="host_numa", '
574+
f"got location_type={str(cfg.location_type)!r}"
575+
)
544576
# The driver rejects an exportable handle type for HOST memory only;
545577
# the NUMA location types may carry one.
546578
if cfg.location_type == VirtualMemoryLocationType.HOST and cfg.handle_type is not None:
@@ -568,7 +600,10 @@ cdef class VirtualMemoryResource(MemoryResource):
568600
memset(prop, 0, sizeof(cydriver.CUmemAllocationProp))
569601
prop.type = <cydriver.CUmemAllocationType>alloc_type
570602
prop.location.type = <cydriver.CUmemLocationType>loc_type
571-
prop.location.id = self.device.device_id if self.device is not None else -1
603+
if self.config.location_type == VirtualMemoryLocationType.HOST_NUMA:
604+
prop.location.id = self.config.numa_id
605+
else:
606+
prop.location.id = self.device.device_id if self.device is not None else -1
572607
prop.allocFlags.gpuDirectRDMACapable = 1 if cfg.gpu_direct_rdma else 0
573608
prop.requestedHandleTypes = <cydriver.CUmemAllocationHandleType>handle_type
574609
prop.win32HandleMetaData = NULL

‎cuda_core/docs/source/release/1.3.0-notes.rst‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -338,6 +338,17 @@ Fixes and enhancements
338338
``repr`` now shows ``U`` for unified operands.
339339
(`#2649 <https://github.com/NVIDIA/cuda-python/issues/2649>`__)
340340

341+
- :class:`VirtualMemoryResource` can now allocate with
342+
``location_type="host_numa"``. Every such allocation failed with
343+
``CUDA_ERROR_INVALID_VALUE``, because the driver received ``-1`` as the NUMA
344+
node. :class:`VirtualMemoryResourceOptions` has a new ``numa_id`` field that
345+
names the node. It is required for ``"host_numa"`` and rejected for every
346+
other location type; ``Device.properties.host_numa_id`` gives the node
347+
closest to a device. ``location_type="host_numa_current"`` now raises
348+
``ValueError`` at construction, because the driver does not accept it for
349+
physical allocations.
350+
(`#2694 <https://github.com/NVIDIA/cuda-python/issues/2694>`__)
351+
341352
Deprecations
342353
------------
343354

‎cuda_core/tests/memory_ipc/test_vmm_ipc.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -180,9 +180,8 @@ def _on_queue_feeder_error(e, _obj):
180180
pytest.param({"handle_type": None}, False, id="none"),
181181
pytest.param({"handle_type": "win32_kmt"}, False, id="win32_kmt"),
182182
pytest.param({"location_type": "host", "handle_type": None}, False, id="host"),
183-
pytest.param({"location_type": "host_numa", "handle_type": "posix_fd"}, False, id="host_numa-posix_fd"),
184183
pytest.param(
185-
{"location_type": "host_numa_current", "handle_type": "posix_fd"}, False, id="host_numa_current-posix_fd"
184+
{"location_type": "host_numa", "numa_id": 0, "handle_type": "posix_fd"}, False, id="host_numa-posix_fd"
186185
),
187186
],
188187
)

‎cuda_core/tests/test_memory.py‎

Lines changed: 71 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3124,32 +3124,98 @@ def test_vmm_options_store_peers_as_tuple():
31243124

31253125

31263126
@pytest.mark.agent_authored(model="claude-fable-5-1")
3127-
@pytest.mark.parametrize("location_type", ["host", "host_numa", "host_numa_current"])
3127+
@pytest.mark.parametrize("location_type", ["host", "host_numa"])
31283128
def test_vmm_host_location_types_report_host_accessible(init_cuda, location_type):
31293129
"""Every host-backed location type reports is_host_accessible; only "host" needs handle_type=None.
31303130
3131-
__init__ classifies "host", "host_numa" and "host_numa_current" alike when
3131+
__init__ classifies "host" and "host_numa" alike when
31323132
deciding the resource is not bound to a device, so is_host_accessible must
31333133
agree; otherwise a NUMA-located resource claims to be neither host- nor
31343134
device-accessible. The driver rejects an exportable handle type for HOST
31353135
memory only, so the default handle type is rejected at construction for
3136-
"host" and accepted for the NUMA location types.
3136+
"host" and accepted for "host_numa".
31373137
"""
31383138
device = Device()
31393139
device.set_current()
3140+
numa = {"numa_id": 0} if location_type == "host_numa" else {}
31403141
if location_type == "host":
31413142
with pytest.raises(ValueError, match="handle_type=None"):
31423143
VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type=location_type))
31433144
else:
3144-
VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type=location_type))
3145+
VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type=location_type, **numa))
31453146
mr = VirtualMemoryResource(
3146-
device, config=VirtualMemoryResourceOptions(location_type=location_type, handle_type=None)
3147+
device, config=VirtualMemoryResourceOptions(location_type=location_type, handle_type=None, **numa)
31473148
)
31483149
assert mr.device is None
31493150
assert mr.is_host_accessible is True
31503151
assert mr.device_id == -1
31513152

31523153

3154+
@pytest.mark.thread_unsafe(reason="records process-global warnings")
3155+
@pytest.mark.agent_authored(model="claude-opus-5-5")
3156+
def test_vmm_host_numa_allocates_and_copies(init_cuda):
3157+
"""A "host_numa" resource allocates on its NUMA node, and the host can use the memory.
3158+
3159+
The driver rejects CU_MEM_LOCATION_TYPE_HOST_NUMA with location id -1, so
3160+
the allocation works only if numa_id reaches the allocation properties
3161+
(#2694). The resource grants access to its own location only, so the data
3162+
goes through the host. A per-call config cannot move the buffer to another
3163+
node.
3164+
"""
3165+
device = Device()
3166+
device.set_current()
3167+
props = device.properties
3168+
if not props.host_numa_virtual_memory_management_supported:
3169+
pytest.skip("Host NUMA virtual memory management is not supported on this device")
3170+
if props.host_numa_id < 0:
3171+
pytest.skip("The system reports no host NUMA node")
3172+
mr = VirtualMemoryResource(
3173+
device,
3174+
config=VirtualMemoryResourceOptions(location_type="host_numa", numa_id=props.host_numa_id, handle_type=None),
3175+
)
3176+
with assert_no_cuda_warning():
3177+
buf = mr.allocate(4096)
3178+
try:
3179+
assert buf.is_host_accessible
3180+
ctypes.memset(int(buf.handle), 0xAB, buf.size)
3181+
assert ctypes.string_at(int(buf.handle), buf.size) == b"\xab" * buf.size
3182+
other_node = VirtualMemoryResourceOptions(
3183+
location_type="host_numa", numa_id=props.host_numa_id + 1, handle_type=None
3184+
)
3185+
with pytest.raises(ValueError, match="config.numa_id"):
3186+
mr.modify_allocation(buf, 2 * buf.size, other_node)
3187+
finally:
3188+
buf.close()
3189+
3190+
3191+
@pytest.mark.agent_authored(model="claude-opus-5-5")
3192+
def test_vmm_host_numa_current_rejected(init_cuda):
3193+
"""The driver does not accept "host_numa_current" for physical allocations, so construction rejects it."""
3194+
device = Device()
3195+
device.set_current()
3196+
with pytest.raises(ValueError, match="does not accept it for physical allocations"):
3197+
VirtualMemoryResource(
3198+
device, config=VirtualMemoryResourceOptions(location_type="host_numa_current", handle_type=None)
3199+
)
3200+
3201+
3202+
@pytest.mark.agent_authored(model="claude-opus-5-5")
3203+
@pytest.mark.parametrize(
3204+
("options", "match"),
3205+
[
3206+
pytest.param({"location_type": "host_numa"}, "requires numa_id", id="host_numa-missing"),
3207+
pytest.param({"location_type": "host_numa", "numa_id": -1}, "numa_id must be >= 0", id="host_numa-negative"),
3208+
pytest.param({"location_type": "host", "numa_id": 0}, "only valid with", id="host-misplaced"),
3209+
pytest.param({"location_type": "device", "numa_id": 0}, "only valid with", id="device-misplaced"),
3210+
],
3211+
)
3212+
def test_vmm_numa_id_validation(init_cuda, options, match):
3213+
"""numa_id is required for "host_numa", must be non-negative, and is rejected for other location types."""
3214+
device = _vmm_device_or_skip()
3215+
with pytest.raises(ValueError, match=match):
3216+
VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(handle_type=None, **options))
3217+
3218+
31533219
def test_device_memory_resource_peer_accessible_by_non_owned(mempool_device):
31543220
"""peer_accessible_by on a non-owned (default) DMR queries the driver live."""
31553221
dev = mempool_device

0 commit comments

Comments
 (0)