From c74fe0533327025085b47aa3d88589213ec70b33 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Wed, 16 Sep 2026 12:37:38 -0700 Subject: [PATCH 1/3] cuda.core: close the old buffer when a VMM grow moves the mapping The slow path of VirtualMemoryResource.modify_allocation unmapped the old VA range, remapped the physical memory into the new range, freed the old reservation by hand and then reset the old buffer's handle with Buffer._clear(). That reset runs the handle's deleter, which calls mr.deallocate() on the range that was just freed. The failing cuMemRetainAllocationHandle is reported as a CUDAWarning since #2759 and was silently swallowed before. Map the old physical memory into the new range as a second mapping instead, which the VMM APIs allow (virtual aliasing), and then close the old buffer. Closing runs deallocate() on a range that is still mapped, so the old range is unmapped, its reservation freed and its handle reference released through the normal path. The remap-on-rollback undo step, which swallowed its own errors, is no longer needed because the old mapping is never removed before the transaction commits. Issue #2877 Co-authored-by: Claude Fable 5.1 --- .../core/_memory/_virtual_memory_resource.py | 44 +++++++------------ cuda_core/docs/source/release/1.3.0-notes.rst | 8 ++++ cuda_core/tests/test_memory.py | 39 ++++++++++++++++ 3 files changed, 63 insertions(+), 28 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..f7580b6975a 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py @@ -372,10 +372,11 @@ def _grow_allocation_slow_path( Slow path for growing a virtual memory allocation when the new region cannot be reserved contiguously after the existing buffer. - This function reserves a new, larger virtual address (VA) range, remaps the old - physical memory to the beginning of the new VA range, creates and maps new physical - memory for the additional size, sets access permissions, and updates the buffer's - pointer and size. + This function reserves a new, larger virtual address (VA) range, maps the old + physical memory to the beginning of the new VA range as a second mapping, creates + and maps new physical memory for the additional size, sets access permissions, and + then closes the old buffer, which releases the old VA range through + :meth:`deallocate`. Args: buf (Buffer): The buffer to grow. @@ -386,8 +387,9 @@ def _grow_allocation_slow_path( addr_align (int): The required address alignment for the new VA range. Returns: - Buffer: The buffer object updated with the new pointer and size. + Buffer: A new buffer for the new VA range. ``buf`` is closed. """ + aligned_prev_size = total_aligned_size - aligned_additional_size with Transaction() as trans: # Reserve a completely new, larger VA range res, new_ptr = driver.cuMemAddressReserve(total_aligned_size, addr_align, 0, 0) @@ -403,23 +405,10 @@ def _grow_allocation_slow_path( # Register undo for old_handle trans.append(lambda h=old_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) - # Unmap the old VA range (aligned previous size) - aligned_prev_size = total_aligned_size - aligned_additional_size - (result,) = driver.cuMemUnmap(int(buf.handle), aligned_prev_size) - raise_if_driver_error(result) - - def _remap_old() -> None: - # Try to remap the old physical memory back to the original VA range - try: - (res,) = driver.cuMemMap(int(buf.handle), aligned_prev_size, 0, old_handle, 0) - raise_if_driver_error(res) - except Exception: # noqa: S110 - # TODO: consider logging this exception - pass - - trans.append(_remap_old) - - # Remap the old physical memory to the new VA range (aligned previous size) + # Map the old physical memory to the new VA range (aligned previous size). + # The old VA range stays mapped too (virtual aliasing), so the old buffer + # is untouched if anything below fails and is released as a whole by + # buf.close() once the new mapping is complete. (res,) = driver.cuMemMap(int(new_ptr), aligned_prev_size, 0, old_handle, 0) raise_if_driver_error(res) @@ -453,12 +442,11 @@ def _remap_old() -> None: # All succeeded, cancel undo actions trans.commit() - # Free the old VA range (aligned previous size) - (res2,) = driver.cuMemAddressFree(int(buf.handle), aligned_prev_size) - raise_if_driver_error(res2) - - # Invalidate the old buffer so its destructor won't try to free again - buf._clear() + # Release the old VA range through the resource: closing the buffer runs + # deallocate(), which unmaps the old range, frees its reservation and + # releases the handle reference it retains. The physical memory stays + # alive through the new mapping. + buf.close() # Return a new Buffer for the new mapping return Buffer.from_handle(ptr=new_ptr, size=new_size, mr=self) 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..723d4c31aa0 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,14 @@ New features Fixes and enhancements ---------------------- +- Growing a :class:`~_memory.VirtualMemoryResource` allocation with + :meth:`~_memory.VirtualMemoryResource.modify_allocation` when the address + range cannot be extended in place no longer reports a spurious + :class:`CUDAWarning`. The old buffer is now closed through the resource's + :meth:`~_memory.VirtualMemoryResource.deallocate` instead of having its + range freed by hand and then released a second time. + (`#2877 `__) + - ``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 diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index b61a2d5b2f3..637208b5c95 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1505,6 +1505,45 @@ def __init__(self, size): assert ("set_access", new_ptr, aligned_additional, 1) in calls +@pytest.mark.agent_authored(model="claude-fable-5-1") +@pytest.mark.thread_unsafe(reason="warning capture is process-global") +def test_vmm_allocator_grow_allocation_slow_path_closes_old_buffer(init_cuda): + """The slow grow path closes the old buffer without a spurious CUDAWarning (#2877). + + It used to free the old VA range by hand and then reset the old buffer's + handle, whose deleter called deallocate() on the freed range a second time. + """ + device = Device() + if not device.properties.virtual_memory_management_supported: + pytest.skip("Virtual memory management is not supported on this device") + + vmm_mr = VirtualMemoryResource( + device, + config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), + ) + buf = vmm_mr.allocate(2 * 1024 * 1024) + old_ptr, old_size = int(buf.handle), buf.size + handle_return(driver.cuMemsetD8(old_ptr, 7, old_size)) + + # Occupy the address range right after buf so the adjacent reservation cannot + # be honored and modify_allocation has to take the slow path. + decoy = handle_return(driver.cuMemAddressReserve(old_size, 0, old_ptr + old_size, 0)) + try: + with assert_no_cuda_warning(): + grown = vmm_mr.modify_allocation(buf, 2 * old_size) + finally: + handle_return(driver.cuMemAddressFree(decoy, old_size)) + + assert buf.is_closed + assert int(grown.handle) != old_ptr + assert grown.size == 2 * old_size + # The old contents are reachable through the new mapping. + host = (ctypes.c_ubyte * old_size)() + handle_return(driver.cuMemcpyDtoH(ctypes.addressof(host), int(grown.handle), old_size)) + assert bytes(host) == bytes([7]) * old_size + grown.close() + + def test_vmm_allocator_rdma_unsupported_exception(): """Test that VirtualMemoryResource throws an exception when RDMA is requested but device doesn't support it. From 9e7578df5bca89a8867f74ca5636c3b0665fc3ea Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Wed, 16 Sep 2026 13:55:05 -0700 Subject: [PATCH 2/3] test(cuda.core): wait for the memset before growing the VMM buffer The slow-path regression test fills the buffer with cuMemsetD8 and grows it at once. Memset is asynchronous with respect to the host, and on WDDM the batched kernel can still be pending when the grow unmaps the old range, so it faults with a sticky CUDA_ERROR_ILLEGAL_ADDRESS that took every later test in the job down with it. Synchronize before the grow. Issue #2877 Co-authored-by: Claude Fable 5.1 --- cuda_core/tests/test_memory.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 637208b5c95..288797832f8 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1524,6 +1524,9 @@ def test_vmm_allocator_grow_allocation_slow_path_closes_old_buffer(init_cuda): buf = vmm_mr.allocate(2 * 1024 * 1024) old_ptr, old_size = int(buf.handle), buf.size handle_return(driver.cuMemsetD8(old_ptr, 7, old_size)) + # The memset is asynchronous and the grow unmaps the old range, so let it + # finish first; on WDDM the batched kernel otherwise faults after the unmap. + device.sync() # Occupy the address range right after buf so the adjacent reservation cannot # be honored and modify_allocation has to take the slow path. From 9ad8672cd03e099e053a8613a7a8b428375cf6f9 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Wed, 16 Sep 2026 13:58:26 -0700 Subject: [PATCH 3/3] cuda.core: release VMM allocation handles once they are mapped VirtualMemoryResource never dropped the reference that cuMemCreate returns. allocate() registered the release only as a rollback action, which commit discards, and deallocate() releases only the reference it retains itself. The driver frees an allocation only once every mapping is unmapped and every handle reference is released, so every buffer's physical memory stayed allocated after close() until the process exited. Both grow paths leaked the new chunk the same way, and the slow path also kept the reference it retained on the old handle. Release the creation reference right after the mapping is attempted, in allocate() and in both grow paths, and release the retained old handle after it is mapped into the new range. The mapping holds its own reference, so the memory stays alive while mapped and is freed by deallocate(), which unmaps the whole range and releases the one reference it retains. Issue #2882 Co-authored-by: Claude Fable 5.1 --- .../core/_memory/_virtual_memory_resource.py | 151 ++++++++++-------- cuda_core/docs/source/release/1.3.0-notes.rst | 7 + cuda_core/tests/test_memory.py | 61 ++++++- 3 files changed, 143 insertions(+), 76 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py index f7580b6975a..7d05fc87b09 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py @@ -328,31 +328,34 @@ def _grow_allocation_fast_path( Buffer: The same buffer object with its size updated to `new_size`. """ with Transaction() as trans: - # Create new physical memory for the additional size + # The caller reserved the extension range; free it unless the grow commits. trans.append( lambda np=new_ptr, s=aligned_additional_size: raise_if_driver_error(driver.cuMemAddressFree(np, s)[0]) ) + # Create new physical memory for the additional size res, new_handle = driver.cuMemCreate(aligned_additional_size, prop, 0) raise_if_driver_error(res) - # Register undo for creation - trans.append(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) - - # Map the new physical memory to the extended VA range - (res,) = driver.cuMemMap(new_ptr, aligned_additional_size, 0, new_handle, 0) - raise_if_driver_error(res) - # Register undo for mapping - trans.append( - lambda np=new_ptr, s=aligned_additional_size: raise_if_driver_error(driver.cuMemUnmap(np, s)[0]) - ) - - # Set access permissions for the new portion - descs = self._build_access_descriptors(prop) - if descs: - (res,) = driver.cuMemSetAccess(new_ptr, aligned_additional_size, descs, len(descs)) + try: + # Map the new physical memory to the extended VA range + (res,) = driver.cuMemMap(new_ptr, aligned_additional_size, 0, new_handle, 0) raise_if_driver_error(res) + # Register undo for mapping + trans.append( + lambda np=new_ptr, s=aligned_additional_size: raise_if_driver_error(driver.cuMemUnmap(np, s)[0]) + ) + + # Set access permissions for the new portion + descs = self._build_access_descriptors(prop) + if descs: + (res,) = driver.cuMemSetAccess(new_ptr, aligned_additional_size, descs, len(descs)) + raise_if_driver_error(res) - # All succeeded, cancel undo actions - trans.commit() + # All succeeded, cancel undo actions + trans.commit() + finally: + # The mapping holds its own reference; drop the one from cuMemCreate + # on success and failure alike (#2882). + raise_if_driver_error(driver.cuMemRelease(new_handle)[0]) # Update the buffer size (pointer stays the same). `Buffer.size` has # no public setter, so this reaches into the private attribute. @@ -399,18 +402,20 @@ def _grow_allocation_slow_path( lambda np=new_ptr, s=total_aligned_size: raise_if_driver_error(driver.cuMemAddressFree(np, s)[0]) ) - # Get the old allocation handle for remapping + # Retain the old allocation handle to map it a second time. The old + # mapping keeps the memory alive, so the retained reference is dropped + # again right after the new mapping is attempted (#2882). result, old_handle = driver.cuMemRetainAllocationHandle(buf.handle) raise_if_driver_error(result) - # Register undo for old_handle - trans.append(lambda h=old_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) - - # Map the old physical memory to the new VA range (aligned previous size). - # The old VA range stays mapped too (virtual aliasing), so the old buffer - # is untouched if anything below fails and is released as a whole by - # buf.close() once the new mapping is complete. - (res,) = driver.cuMemMap(int(new_ptr), aligned_prev_size, 0, old_handle, 0) - raise_if_driver_error(res) + try: + # Map the old physical memory to the new VA range (aligned previous size). + # The old VA range stays mapped too (virtual aliasing), so the old buffer + # is untouched if anything below fails and is released as a whole by + # buf.close() once the new mapping is complete. + (res,) = driver.cuMemMap(int(new_ptr), aligned_prev_size, 0, old_handle, 0) + raise_if_driver_error(res) + finally: + raise_if_driver_error(driver.cuMemRelease(old_handle)[0]) # Register undo for mapping trans.append(lambda np=new_ptr, s=aligned_prev_size: raise_if_driver_error(driver.cuMemUnmap(np, s)[0])) @@ -418,29 +423,30 @@ def _grow_allocation_slow_path( # Create new physical memory for the additional size res, new_handle = driver.cuMemCreate(aligned_additional_size, prop, 0) raise_if_driver_error(res) + try: + # Map the new physical memory to the extended portion (aligned offset) + (res,) = driver.cuMemMap(int(new_ptr) + aligned_prev_size, aligned_additional_size, 0, new_handle, 0) + raise_if_driver_error(res) - # Register undo for new physical memory - trans.append(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) - - # Map the new physical memory to the extended portion (aligned offset) - (res,) = driver.cuMemMap(int(new_ptr) + aligned_prev_size, aligned_additional_size, 0, new_handle, 0) - raise_if_driver_error(res) - - # Register undo for mapping - trans.append( - lambda base=int(new_ptr), offs=aligned_prev_size, s=aligned_additional_size: raise_if_driver_error( - driver.cuMemUnmap(base + offs, s)[0] + # Register undo for mapping + trans.append( + lambda base=int(new_ptr), offs=aligned_prev_size, s=aligned_additional_size: raise_if_driver_error( + driver.cuMemUnmap(base + offs, s)[0] + ) ) - ) - # Set access permissions for the entire new range - descs = self._build_access_descriptors(prop) - if descs: - (res,) = driver.cuMemSetAccess(new_ptr, total_aligned_size, descs, len(descs)) - raise_if_driver_error(res) + # Set access permissions for the entire new range + descs = self._build_access_descriptors(prop) + if descs: + (res,) = driver.cuMemSetAccess(new_ptr, total_aligned_size, descs, len(descs)) + raise_if_driver_error(res) - # All succeeded, cancel undo actions - trans.commit() + # All succeeded, cancel undo actions + trans.commit() + finally: + # The mapping holds its own reference; drop the one from cuMemCreate + # on success and failure alike (#2882). + raise_if_driver_error(driver.cuMemRelease(new_handle)[0]) # Release the old VA range through the resource: closing the buffer runs # deallocate(), which unmaps the old range, frees its reservation and @@ -537,32 +543,37 @@ def allocate(self, size: int, *, stream: Stream | GraphBuilder | None = None) -> addr_align = config.addr_align or gran # ---- Transactional allocation ---- - with Transaction() as trans: - # ---- Create physical memory ---- - res, handle = driver.cuMemCreate(aligned_size, prop, 0) - raise_if_driver_error(res) - # Register undo for physical memory - trans.append(lambda h=handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) - - # ---- Reserve VA space ---- - # Potentially, use a separate size for the VA reservation from the physical allocation size - res, ptr = driver.cuMemAddressReserve(aligned_size, addr_align, config.addr_hint, 0) - raise_if_driver_error(res) - # Register undo for VA reservation - trans.append(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemAddressFree(p, s)[0])) - - # ---- Map physical memory into VA ---- - (res,) = driver.cuMemMap(ptr, aligned_size, 0, handle, 0) - trans.append(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemUnmap(p, s)[0])) - raise_if_driver_error(res) + # ---- Create physical memory ---- + res, handle = driver.cuMemCreate(aligned_size, prop, 0) + raise_if_driver_error(res) + try: + with Transaction() as trans: + # ---- Reserve VA space ---- + # Potentially, use a separate size for the VA reservation from the physical allocation size + res, ptr = driver.cuMemAddressReserve(aligned_size, addr_align, config.addr_hint, 0) + raise_if_driver_error(res) + # Register undo for VA reservation + trans.append(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemAddressFree(p, s)[0])) - # ---- Set access for owner + peers ---- - descs = self._build_access_descriptors(prop) - if descs: - (res,) = driver.cuMemSetAccess(ptr, aligned_size, descs, len(descs)) + # ---- Map physical memory into VA ---- + (res,) = driver.cuMemMap(ptr, aligned_size, 0, handle, 0) + trans.append(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemUnmap(p, s)[0])) raise_if_driver_error(res) - trans.commit() + # ---- Set access for owner + peers ---- + descs = self._build_access_descriptors(prop) + if descs: + (res,) = driver.cuMemSetAccess(ptr, aligned_size, descs, len(descs)) + raise_if_driver_error(res) + + trans.commit() + finally: + # The mapping holds its own reference to the physical allocation, so the + # reference returned by cuMemCreate is dropped here whether or not the + # mapping succeeded. Keeping it made deallocate() unable to free the + # memory: the driver frees an allocation only once every mapping is + # unmapped and every handle reference is released (#2882). + raise_if_driver_error(driver.cuMemRelease(handle)[0]) # Done — return a Buffer that tracks this VA range buf = Buffer.from_handle(ptr=ptr, size=aligned_size, mr=self) 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 723d4c31aa0..fbb9c128260 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,13 @@ New features Fixes and enhancements ---------------------- +- :class:`~_memory.VirtualMemoryResource` now releases the physical allocation + handle once it is mapped, so closing a buffer returns its memory to the + device. Previously every allocation, and every chunk added by + :meth:`~_memory.VirtualMemoryResource.modify_allocation`, stayed allocated + until the process exited. + (`#2882 `__) + - Growing a :class:`~_memory.VirtualMemoryResource` allocation with :meth:`~_memory.VirtualMemoryResource.modify_allocation` when the address range cannot be extended in place no longer reports a spurious diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 288797832f8..d4a4f2e6f7f 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1465,14 +1465,16 @@ def fake_unmap(ptr, size): calls.append(("unmap", ptr, size)) return (SUCCESS,) - def fake_release(handle): - calls.append(("release", handle)) - return (SUCCESS,) - def fake_addr_free(ptr, size): calls.append(("addr_free", ptr, size)) return (SUCCESS,) + # Runs on success as well: the mapping keeps the memory alive, so the + # reference returned by cuMemCreate must be dropped once it is mapped (#2882). + def fake_release(handle): + calls.append(("release", handle)) + return (SUCCESS,) + monkeypatch.setattr(driver, "cuMemCreate", fake_create) monkeypatch.setattr(driver, "cuMemMap", fake_map) monkeypatch.setattr(driver, "cuMemSetAccess", fake_set_access) @@ -1498,11 +1500,58 @@ def __init__(self, size): assert result is buf assert buf._size == new_size - # Successful commit: create, map, set access, and no rollback calls. - assert [c[0] for c in calls] == ["create", "map", "set_access"] + # Successful commit: create, map, set access, release the creation + # reference, and no rollback calls. + assert [c[0] for c in calls] == ["create", "map", "set_access", "release"] assert ("create", aligned_additional) in calls assert ("map", new_ptr, aligned_additional, NEW_HANDLE) in calls assert ("set_access", new_ptr, aligned_additional, 1) in calls + assert ("release", NEW_HANDLE) in calls + + +@pytest.mark.agent_authored(model="claude-fable-5-1") +@pytest.mark.thread_unsafe(reason="measures device-wide free memory") +def test_vmm_allocator_close_returns_physical_memory(init_cuda): + """Closing VMM buffers returns their physical memory to the device (#2882). + + allocate() and both grow paths kept the reference that cuMemCreate returns, + so the memory stayed allocated after the range was unmapped and freed. + """ + device = Device() + if not device.properties.virtual_memory_management_supported: + pytest.skip("Virtual memory management is not supported on this device") + + vmm_mr = VirtualMemoryResource( + device, + config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), + ) + chunk = 32 * 1024 * 1024 + rounds = 8 + + def free_memory(): + device.sync() + return handle_return(driver.cuMemGetInfo())[0] + + # Plain allocate and close: the leak was one chunk per round. + before = free_memory() + for _ in range(rounds): + vmm_mr.allocate(chunk).close() + retained = before - free_memory() + assert retained < rounds * chunk // 2, f"{retained >> 20} MiB still allocated after close" + + # Grow through the slow path (forced by a decoy reservation), then close: + # the leak was two chunks per round. + before = free_memory() + for _ in range(rounds): + buf = vmm_mr.allocate(chunk) + decoy = handle_return(driver.cuMemAddressReserve(chunk, 0, int(buf.handle) + buf.size, 0)) + try: + grown = vmm_mr.modify_allocation(buf, 2 * chunk) + finally: + handle_return(driver.cuMemAddressFree(decoy, chunk)) + grown.close() + retained = before - free_memory() + assert retained < rounds * chunk, f"{retained >> 20} MiB still allocated after grow and close" @pytest.mark.agent_authored(model="claude-fable-5-1")