fix(cuda.core): pass numa_id to cuMemCreate for host_numa virtual memory

cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi

@@ -30,7 +30,9 @@ class VirtualMemoryResourceOptions:
     allocation_type: :obj:`~_memory.VirtualMemoryAllocationType` | str
         Controls the type of allocation.
     location_type: :obj:`~_memory.VirtualMemoryLocationType` | str
-        Controls the location of the allocation.
+        Controls the location of the allocation. ``"host_numa"`` requires
+        ``numa_id``. ``"host_numa_current"`` is rejected because the driver
+        does not accept it for physical allocations.
     handle_type: :obj:`~_memory.VirtualMemoryHandleType` | str
         Export handle type for the physical allocation, and the switch for
         sharing buffers with other processes (see
@@ -52,6 +54,12 @@ class VirtualMemoryResourceOptions:
         Access flags for the owning device.
     peer_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str
         Access flags for peers.
+    numa_id: int | None
+        Host NUMA node of the allocation. Required for
+        ``location_type="host_numa"`` and rejected for every other location
+        type. Must be >= 0; ``Device.properties.host_numa_id`` is the node
+        closest to a device. On a system without NUMA, ``host_numa_id`` is -1
+        and the driver expects ``numa_id=0``.
     """
     allocation_type: VirtualMemoryAllocationType = VirtualMemoryAllocationType.PINNED
     location_type: VirtualMemoryLocationType = VirtualMemoryLocationType.DEVICE
@@ -63,6 +71,7 @@ class VirtualMemoryResourceOptions:
     peers: Iterable[int] = field(default_factory=tuple)
     self_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE
     peer_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE
+    numa_id: int | None = None
     _a = driver.CUmemAccess_flags
     _access_flags = {VirtualMemoryAccessType.READ_WRITE: _a.CU_MEM_ACCESS_FLAGS_PROT_READWRITE, VirtualMemoryAccessType.READ: _a.CU_MEM_ACCESS_FLAGS_PROT_READ, None: 0}
     _h = driver.CUmemAllocationHandleType

cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx

@@ -99,7 +99,9 @@ class VirtualMemoryResourceOptions:
     allocation_type: :obj:`~_memory.VirtualMemoryAllocationType` | str
         Controls the type of allocation.
     location_type: :obj:`~_memory.VirtualMemoryLocationType` | str
-        Controls the location of the allocation.
+        Controls the location of the allocation. ``"host_numa"`` requires
+        ``numa_id``. ``"host_numa_current"`` is rejected because the driver
+        does not accept it for physical allocations.
     handle_type: :obj:`~_memory.VirtualMemoryHandleType` | str
         Export handle type for the physical allocation, and the switch for
         sharing buffers with other processes (see
@@ -121,6 +123,12 @@ class VirtualMemoryResourceOptions:
         Access flags for the owning device.
     peer_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str
         Access flags for peers.
+    numa_id: int | None
+        Host NUMA node of the allocation. Required for
+        ``location_type="host_numa"`` and rejected for every other location
+        type. Must be >= 0; ``Device.properties.host_numa_id`` is the node
+        closest to a device. On a system without NUMA, ``host_numa_id`` is -1
+        and the driver expects ``numa_id=0``.
     """
 
     allocation_type: VirtualMemoryAllocationType = VirtualMemoryAllocationType.PINNED
@@ -133,6 +141,7 @@ class VirtualMemoryResourceOptions:
     peers: Iterable[int] = field(default_factory=tuple)
     self_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE
     peer_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE
+    numa_id: int | None = None
 
     _a = driver.CUmemAccess_flags
     _access_flags = {  # noqa: RUF012
@@ -472,6 +481,29 @@ cdef class VirtualMemoryResource(MemoryResource):
                 f"config.location_type {str(cfg.location_type)!r} does not match the resource's "
                 f"{str(self.config.location_type)!r}; the location of a buffer cannot change"
             )
+        if cfg.numa_id != self.config.numa_id:
+            raise ValueError(
+                f"config.numa_id {cfg.numa_id!r} does not match the resource's "
+                f"{self.config.numa_id!r}; the location of a buffer cannot change"
+            )
+        if cfg.location_type == VirtualMemoryLocationType.HOST_NUMA_CURRENT:
+            raise ValueError(
+                'location_type="host_numa_current" is not supported: the driver does not accept it '
+                'for physical allocations; use location_type="host_numa" with numa_id'
+            )
+        if cfg.location_type == VirtualMemoryLocationType.HOST_NUMA:
+            if cfg.numa_id is None:
+                raise ValueError(
+                    'location_type="host_numa" requires numa_id; pass numa_id=<host NUMA node id>, '
+                    "e.g. device.properties.host_numa_id"
+                )
+            if cfg.numa_id < 0:
+                raise ValueError(f"numa_id must be >= 0, got {cfg.numa_id}")
+        elif cfg.numa_id is not None:
+            raise ValueError(
+                'numa_id is only valid with location_type="host_numa", '
+                f"got location_type={str(cfg.location_type)!r}"
+            )
         # The driver rejects an exportable handle type for HOST memory only;
         # the NUMA location types may carry one.
         if cfg.location_type == VirtualMemoryLocationType.HOST and cfg.handle_type is not None:
@@ -499,7 +531,10 @@ cdef class VirtualMemoryResource(MemoryResource):
         memset(prop, 0, sizeof(cydriver.CUmemAllocationProp))
         prop.type = <cydriver.CUmemAllocationType>alloc_type
         prop.location.type = <cydriver.CUmemLocationType>loc_type
-        prop.location.id = self.device.device_id if self.device is not None else -1
+        if self.config.location_type == VirtualMemoryLocationType.HOST_NUMA:
+            prop.location.id = self.config.numa_id
+        else:
+            prop.location.id = self.device.device_id if self.device is not None else -1
         prop.allocFlags.gpuDirectRDMACapable = 1 if cfg.gpu_direct_rdma else 0
         prop.requestedHandleTypes = <cydriver.CUmemAllocationHandleType>handle_type
         prop.win32HandleMetaData = NULL

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

@@ -303,3 +303,14 @@ Fixes and enhancements
   while operands from plain ``cuMemAlloc`` allocations accept it. The node's
   ``repr`` now shows ``U`` for unified operands.
   (`#2649 <https://github.com/NVIDIA/cuda-python/issues/2649>`__)
+
+- :class:`VirtualMemoryResource` can now allocate with
+  ``location_type="host_numa"``. Every such allocation failed with
+  ``CUDA_ERROR_INVALID_VALUE``, because the driver received ``-1`` as the NUMA
+  node. :class:`VirtualMemoryResourceOptions` has a new ``numa_id`` field that
+  names the node. It is required for ``"host_numa"`` and rejected for every
+  other location type; ``Device.properties.host_numa_id`` gives the node
+  closest to a device. ``location_type="host_numa_current"`` now raises
+  ``ValueError`` at construction, because the driver does not accept it for
+  physical allocations.
+  (`#2694 <https://github.com/NVIDIA/cuda-python/issues/2694>`__)

cuda_core/tests/memory_ipc/test_vmm_ipc.py

@@ -193,9 +193,8 @@ def _on_queue_feeder_error(e, _obj):
         pytest.param({"handle_type": "fabric"}, False, id="fabric"),
         pytest.param({"handle_type": "win32_kmt"}, False, id="win32_kmt"),
         pytest.param({"location_type": "host", "handle_type": None}, False, id="host"),
-        pytest.param({"location_type": "host_numa", "handle_type": "posix_fd"}, False, id="host_numa-posix_fd"),
         pytest.param(
-            {"location_type": "host_numa_current", "handle_type": "posix_fd"}, False, id="host_numa_current-posix_fd"
+            {"location_type": "host_numa", "numa_id": 0, "handle_type": "posix_fd"}, False, id="host_numa-posix_fd"
         ),
     ],
 )

cuda_core/tests/test_memory.py

@@ -3117,32 +3117,98 @@ def test_vmm_options_handle_type_win32_raises():
 
 
 @pytest.mark.agent_authored(model="claude-fable-5-1")
-@pytest.mark.parametrize("location_type", ["host", "host_numa", "host_numa_current"])
+@pytest.mark.parametrize("location_type", ["host", "host_numa"])
 def test_vmm_host_location_types_report_host_accessible(init_cuda, location_type):
     """Every host-backed location type reports is_host_accessible; only "host" needs handle_type=None.
 
-    __init__ classifies "host", "host_numa" and "host_numa_current" alike when
+    __init__ classifies "host" and "host_numa" alike when
     deciding the resource is not bound to a device, so is_host_accessible must
     agree; otherwise a NUMA-located resource claims to be neither host- nor
     device-accessible. The driver rejects an exportable handle type for HOST
     memory only, so the default handle type is rejected at construction for
-    "host" and accepted for the NUMA location types.
+    "host" and accepted for "host_numa".
     """
     device = Device()
     device.set_current()
+    numa = {"numa_id": 0} if location_type == "host_numa" else {}
     if location_type == "host":
         with pytest.raises(ValueError, match="handle_type=None"):
             VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type=location_type))
     else:
-        VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type=location_type))
+        VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type=location_type, **numa))
     mr = VirtualMemoryResource(
-        device, config=VirtualMemoryResourceOptions(location_type=location_type, handle_type=None)
+        device, config=VirtualMemoryResourceOptions(location_type=location_type, handle_type=None, **numa)
     )
     assert mr.device is None
     assert mr.is_host_accessible is True
     assert mr.device_id == -1
 
 
+@pytest.mark.thread_unsafe(reason="records process-global warnings")
+@pytest.mark.agent_authored(model="claude-opus-5-5")
+def test_vmm_host_numa_allocates_and_copies(init_cuda):
+    """A "host_numa" resource allocates on its NUMA node, and the host can use the memory.
+
+    The driver rejects CU_MEM_LOCATION_TYPE_HOST_NUMA with location id -1, so
+    the allocation works only if numa_id reaches the allocation properties
+    (#2694). The resource grants access to its own location only, so the data
+    goes through the host. A per-call config cannot move the buffer to another
+    node.
+    """
+    device = Device()
+    device.set_current()
+    props = device.properties
+    if not props.host_numa_virtual_memory_management_supported:
+        pytest.skip("Host NUMA virtual memory management is not supported on this device")
+    if props.host_numa_id < 0:
+        pytest.skip("The system reports no host NUMA node")
+    mr = VirtualMemoryResource(
+        device,
+        config=VirtualMemoryResourceOptions(location_type="host_numa", numa_id=props.host_numa_id, handle_type=None),
+    )
+    with assert_no_cuda_warning():
+        buf = mr.allocate(4096)
+        try:
+            assert buf.is_host_accessible
+            ctypes.memset(int(buf.handle), 0xAB, buf.size)
+            assert ctypes.string_at(int(buf.handle), buf.size) == b"\xab" * buf.size
+            other_node = VirtualMemoryResourceOptions(
+                location_type="host_numa", numa_id=props.host_numa_id + 1, handle_type=None
+            )
+            with pytest.raises(ValueError, match="config.numa_id"):
+                mr.modify_allocation(buf, 2 * buf.size, other_node)
+        finally:
+            buf.close()
+
+
+@pytest.mark.agent_authored(model="claude-opus-5-5")
+def test_vmm_host_numa_current_rejected(init_cuda):
+    """The driver does not accept "host_numa_current" for physical allocations, so construction rejects it."""
+    device = Device()
+    device.set_current()
+    with pytest.raises(ValueError, match="does not accept it for physical allocations"):
+        VirtualMemoryResource(
+            device, config=VirtualMemoryResourceOptions(location_type="host_numa_current", handle_type=None)
+        )
+
+
+@pytest.mark.agent_authored(model="claude-opus-5-5")
+@pytest.mark.parametrize(
+    ("options", "match"),
+    [
+        pytest.param({"location_type": "host_numa"}, "requires numa_id", id="host_numa-missing"),
+        pytest.param({"location_type": "host_numa", "numa_id": -1}, "numa_id must be >= 0", id="host_numa-negative"),
+        pytest.param({"location_type": "host", "numa_id": 0}, "only valid with", id="host-misplaced"),
+        pytest.param({"location_type": "device", "numa_id": 0}, "only valid with", id="device-misplaced"),
+    ],
+)
+def test_vmm_numa_id_validation(init_cuda, options, match):
+    """numa_id is required for "host_numa", must be non-negative, and is rejected for other location types."""
+    device = _vmm_device_or_skip()
+    with pytest.raises(ValueError, match=match):
+        VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(handle_type=None, **options))
+
+
 def test_device_memory_resource_peer_accessible_by_non_owned(mempool_device):
     """peer_accessible_by on a non-owned (default) DMR queries the driver live."""
     dev = mempool_device