From 605c36e70b49204393b3a785c24eda012a09a1b4 Mon Sep 17 00:00:00 2001 From: Aryan Date: Wed, 16 Sep 2026 20:53:35 -0400 Subject: [PATCH 1/4] fix(cuda.core): order VMM unmaps on the stream Signed-off-by: Aryan --- .../core/_memory/_virtual_memory_resource.py | 9 +++-- cuda_core/tests/test_memory.py | 36 +++++++++++++++++++ 2 files changed, 40 insertions(+), 5 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py index 2c4f3f6867e..cc36593c197 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py @@ -504,8 +504,7 @@ def allocate(self, size: int, *, stream: Stream | GraphBuilder | None = None) -> size : int The size in bytes of the buffer to allocate. stream : Stream, optional - Keyword-only. Unused because virtual memory operations are - synchronous. + Keyword-only. Validated when provided but otherwise unused. Returns ------- @@ -591,15 +590,15 @@ def deallocate(self, ptr: DevicePointerType, size: int, *, stream: Stream | Grap size : int The size in bytes of the memory to deallocate. stream : Stream, optional - Keyword-only. Unused because virtual memory operations are - synchronous. + Keyword-only. If provided, ``stream.sync()`` is called before the + virtual memory is unmapped. """ ptr = 0 if ptr is None else int(ptr) if stream is not None: from cuda.core._stream import Stream_accept - Stream_accept(stream) + Stream_accept(stream).sync() result, handle = driver.cuMemRetainAllocationHandle(ptr) raise_if_driver_error(result) (result,) = driver.cuMemUnmap(ptr, size) diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index b61a2d5b2f3..228a4415de5 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1505,6 +1505,42 @@ def __init__(self, size): assert ("set_access", new_ptr, aligned_additional, 1) in calls +@pytest.mark.agent_authored(model="gpt-5.6-sol") +def test_vmm_deallocate_synchronizes_stream_before_unmap(monkeypatch): + """VMM deallocation orders outstanding stream work before unmapping.""" + mr = VirtualMemoryResource.__new__(VirtualMemoryResource) + events = [] + success = driver.CUresult.CUDA_SUCCESS + + class FakeStream: + def sync(self): + events.append("sync") + + monkeypatch.setattr("cuda.core._stream.Stream_accept", lambda stream: stream) + + def fake_retain(_ptr): + return success, 0xBEEF + + def fake_unmap(_ptr, _size): + events.append("unmap") + return (success,) + + def fake_address_free(_ptr, _size): + return (success,) + + def fake_release(_handle): + return (success,) + + monkeypatch.setattr(driver, "cuMemRetainAllocationHandle", fake_retain) + monkeypatch.setattr(driver, "cuMemUnmap", fake_unmap) + monkeypatch.setattr(driver, "cuMemAddressFree", fake_address_free) + monkeypatch.setattr(driver, "cuMemRelease", fake_release) + + mr.deallocate(0x1000, 4096, stream=FakeStream()) + + assert events == ["sync", "unmap"] + + def test_vmm_allocator_rdma_unsupported_exception(): """Test that VirtualMemoryResource throws an exception when RDMA is requested but device doesn't support it. From bdb3de6c6116d4e6501d9f7c2d7ad053b9f94776 Mon Sep 17 00:00:00 2001 From: Aryan Putta Date: Thu, 17 Sep 2026 20:34:55 -0400 Subject: [PATCH 2/4] fix(cuda.core): guard VMM stream sync for host resources --- cuda_core/cuda/core/_memory/_virtual_memory_resource.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py index cc36593c197..1db8c2e79f8 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py @@ -590,12 +590,12 @@ def deallocate(self, ptr: DevicePointerType, size: int, *, stream: Stream | Grap size : int The size in bytes of the memory to deallocate. stream : Stream, optional - Keyword-only. If provided, ``stream.sync()`` is called before the - virtual memory is unmapped. + Keyword-only. If provided for a device-located resource, + ``stream.sync()`` is called before the virtual memory is unmapped. """ ptr = 0 if ptr is None else int(ptr) - if stream is not None: + if stream is not None and self.is_device_accessible: from cuda.core._stream import Stream_accept Stream_accept(stream).sync() From c0dfe53effaa482bc4715f599fda2f8fb81a53bf Mon Sep 17 00:00:00 2001 From: Aryan Putta Date: Thu, 17 Sep 2026 20:35:12 -0400 Subject: [PATCH 3/4] test(cuda.core): cover VMM stream cleanup on device and host resources --- cuda_core/tests/test_memory.py | 45 +++++++++------------------------- 1 file changed, 12 insertions(+), 33 deletions(-) diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 228a4415de5..b44be2af19f 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1505,40 +1505,19 @@ def __init__(self, size): assert ("set_access", new_ptr, aligned_additional, 1) in calls -@pytest.mark.agent_authored(model="gpt-5.6-sol") -def test_vmm_deallocate_synchronizes_stream_before_unmap(monkeypatch): - """VMM deallocation orders outstanding stream work before unmapping.""" - mr = VirtualMemoryResource.__new__(VirtualMemoryResource) - events = [] - success = driver.CUresult.CUDA_SUCCESS - - class FakeStream: - def sync(self): - events.append("sync") - - monkeypatch.setattr("cuda.core._stream.Stream_accept", lambda stream: stream) - - def fake_retain(_ptr): - return success, 0xBEEF - - def fake_unmap(_ptr, _size): - events.append("unmap") - return (success,) - - def fake_address_free(_ptr, _size): - return (success,) - - def fake_release(_handle): - return (success,) - - monkeypatch.setattr(driver, "cuMemRetainAllocationHandle", fake_retain) - monkeypatch.setattr(driver, "cuMemUnmap", fake_unmap) - monkeypatch.setattr(driver, "cuMemAddressFree", fake_address_free) - monkeypatch.setattr(driver, "cuMemRelease", fake_release) - - mr.deallocate(0x1000, 4096, stream=FakeStream()) +@pytest.mark.parametrize("location_type", ["device", "host"]) +def test_vmm_deallocate_with_stream_has_no_warning(init_cuda, location_type): + """VMM buffers close cleanly with an explicit stream for both locations.""" + device = Device() + if not device.properties.virtual_memory_management_supported: + pytest.skip("Virtual memory management is not supported on this device") + device.set_current() + stream = device.create_stream() + mr = VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(location_type=location_type)) - assert events == ["sync", "unmap"] + with assert_no_cuda_warning(): + buffer = mr.allocate(2 * 1024 * 1024) + buffer.close(stream) def test_vmm_allocator_rdma_unsupported_exception(): From b4b5adb4e8fe20afcb9bed662ad83f2d50dabc51 Mon Sep 17 00:00:00 2001 From: Aryan Putta Date: Thu, 17 Sep 2026 20:35:16 -0400 Subject: [PATCH 4/4] docs(cuda.core): note VMM stream cleanup fix --- cuda_core/docs/source/release/1.3.0-notes.rst | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/cuda_core/docs/source/release/1.3.0-notes.rst b/cuda_core/docs/source/release/1.3.0-notes.rst index 413465d8b3c..0b7b9a40065 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -34,6 +34,11 @@ New features Fixes and enhancements ---------------------- +- :meth:`VirtualMemoryResource.deallocate` now synchronizes a supplied stream + before unmapping device-located virtual memory. Host-located resources skip + this synchronization because they do not retain a device context. + (`#2889 `__) + - ``Graph.__getitem__`` now declares an overload for each node type that has an executable view, so type checkers and editors see the precise view type: indexing with a :class:`~graph.KernelNode` yields an