From 55df22e2a574ee5ec72b4a87928a240a421fe2e1 Mon Sep 17 00:00:00 2001 From: Aryan Date: Tue, 28 Jul 2026 16:37:26 -0400 Subject: [PATCH 1/2] fix(core): restore access descriptors when a VMM grow rolls back The slow-path rollback in modify_allocation remapped the old physical memory to its original address but never reapplied the access descriptors that the preceding unmap dropped, so a buffer recovered by a failed grow faulted on its next access until cuMemSetAccess was run again. Hoist the descriptor list above the transaction so the rollback applies the same descriptors the forward path does, and reapply them to the old range after the remap succeeds. Defect 1 of #2388. Signed-off-by: Aryan --- .../core/_memory/_virtual_memory_resource.py | 6 +- cuda_core/docs/source/release/1.2.0-notes.rst | 6 + cuda_core/tests/test_memory.py | 103 ++++++++++++++++++ 3 files changed, 114 insertions(+), 1 deletion(-) diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py index f30e6e3838d..6d608ac265d 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py @@ -377,6 +377,8 @@ def _grow_allocation_slow_path( Returns: Buffer: The buffer object updated with the new pointer and size. """ + descs = self._build_access_descriptors(prop) + with Transaction() as trans: # Reserve a completely new, larger VA range res, new_ptr = driver.cuMemAddressReserve(total_aligned_size, addr_align, 0, 0) @@ -402,6 +404,9 @@ def _remap_old() -> None: try: (res,) = driver.cuMemMap(int(buf.handle), aligned_prev_size, 0, old_handle, 0) raise_if_driver_error(res) + if descs: + (res,) = driver.cuMemSetAccess(int(buf.handle), aligned_prev_size, descs, len(descs)) + raise_if_driver_error(res) except Exception: # noqa: S110 # TODO: consider logging this exception pass @@ -434,7 +439,6 @@ def _remap_old() -> None: ) # 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) diff --git a/cuda_core/docs/source/release/1.2.0-notes.rst b/cuda_core/docs/source/release/1.2.0-notes.rst index 6a047c9cfe8..a7133bb45bb 100644 --- a/cuda_core/docs/source/release/1.2.0-notes.rst +++ b/cuda_core/docs/source/release/1.2.0-notes.rst @@ -18,6 +18,12 @@ Fixes and enhancements (`#2357 `__, `#2371 `__) +- Rolling back a failed :meth:`VirtualMemoryResource.modify_allocation` now + restores access to the original range. Previously the rollback remapped the + old physical memory to its original address but left the access descriptors + unset, so the recovered buffer faulted on its next access. + (`#2388 `__) + Deprecation Notices ------------------- diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 5ef48919a50..c24a919ab9b 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1035,6 +1035,109 @@ def __init__(self, size): assert ("set_access", new_ptr, aligned_additional, 1) in calls +@pytest.mark.agent_authored(model="claude-opus-5") +def test_vmm_allocator_grow_slow_path_rollback_restores_access(init_cuda, monkeypatch): + """Fail the slow path after the old range is remapped and check the rollback. + + Rolling back a failed grow has to leave the original range usable, so + restoring the mapping is not enough on its own: the access descriptors that + the unmap dropped have to be reapplied to the same range. + """ + 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"), + ) + + prop = driver.CUmemAllocationProp() + prop.type = driver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + prop.location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + prop.location.id = device.device_id + + SUCCESS = driver.CUresult.CUDA_SUCCESS + FAILURE = driver.CUresult.CUDA_ERROR_INVALID_VALUE + OLD_PTR = 0x20_0000 + NEW_PTR = 0x40_0000 + OLD_HANDLE = 0xF00D + NEW_HANDLE = 0xBEEF + calls = [] + + def fake_addr_reserve(size, align, addr, flags): + return (SUCCESS, NEW_PTR) + + def fake_retain(ptr): + return (SUCCESS, OLD_HANDLE) + + def fake_create(size, p, flags): + return (SUCCESS, NEW_HANDLE) + + def fake_map(ptr, size, offset, handle, flags): + calls.append(("map", ptr, size, handle)) + return (SUCCESS,) + + def fake_unmap(ptr, size): + calls.append(("unmap", ptr, size)) + return (SUCCESS,) + + def fake_set_access(ptr, size, descs, count): + calls.append(("set_access", ptr, size, count)) + # Fail the forward call over the new range; the rollback call must still + # be allowed to succeed. + return (FAILURE,) if ptr == NEW_PTR else (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,) + + monkeypatch.setattr(driver, "cuMemAddressReserve", fake_addr_reserve) + monkeypatch.setattr(driver, "cuMemRetainAllocationHandle", fake_retain) + monkeypatch.setattr(driver, "cuMemCreate", fake_create) + monkeypatch.setattr(driver, "cuMemMap", fake_map) + monkeypatch.setattr(driver, "cuMemUnmap", fake_unmap) + monkeypatch.setattr(driver, "cuMemSetAccess", fake_set_access) + monkeypatch.setattr(driver, "cuMemRelease", fake_release) + monkeypatch.setattr(driver, "cuMemAddressFree", fake_addr_free) + + class FakeBuffer: + def __init__(self, ptr, size): + self.handle = ptr + self._size = size + + aligned_prev_size = 2 * 1024 * 1024 + aligned_additional_size = 2 * 1024 * 1024 + total_aligned_size = aligned_prev_size + aligned_additional_size + addr_align = 2 * 1024 * 1024 + buf = FakeBuffer(OLD_PTR, aligned_prev_size) + + with pytest.raises(CUDAError): + vmm_mr._grow_allocation_slow_path( + buf, + total_aligned_size, + prop, + aligned_additional_size, + total_aligned_size, + addr_align, + ) + + # The rollback remapped the old physical memory back to the original VA. + assert ("map", OLD_PTR, aligned_prev_size, OLD_HANDLE) in calls + # ... and reapplied access to that same range, which is the regression here. + assert ("set_access", OLD_PTR, aligned_prev_size, 1) in calls + # Access is restored only after the mapping exists again. + assert calls.index(("map", OLD_PTR, aligned_prev_size, OLD_HANDLE)) < calls.index( + ("set_access", OLD_PTR, aligned_prev_size, 1) + ) + # The failed reservation is still released. + assert ("addr_free", NEW_PTR, total_aligned_size) in calls + + def test_vmm_allocator_rdma_unsupported_exception(): """Test that VirtualMemoryResource throws an exception when RDMA is requested but device doesn't support it. From 7ee666246d2c4862c67f8a130a1c883111a2c75a Mon Sep 17 00:00:00 2001 From: Aryan Date: Tue, 28 Jul 2026 16:53:48 -0400 Subject: [PATCH 2/2] docs: point the release note at the PR Signed-off-by: Aryan --- cuda_core/docs/source/release/1.2.0-notes.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cuda_core/docs/source/release/1.2.0-notes.rst b/cuda_core/docs/source/release/1.2.0-notes.rst index a7133bb45bb..e26b0eda37b 100644 --- a/cuda_core/docs/source/release/1.2.0-notes.rst +++ b/cuda_core/docs/source/release/1.2.0-notes.rst @@ -22,7 +22,7 @@ Fixes and enhancements restores access to the original range. Previously the rollback remapped the old physical memory to its original address but left the access descriptors unset, so the recovered buffer faulted on its next access. - (`#2388 `__) + (`#2440 `__) Deprecation Notices -------------------