From 24f3f6c20225a9b595962f0bc0f9f513be2d6881 Mon Sep 17 00:00:00 2001 From: Aryan Putta Date: Wed, 22 Jul 2026 11:29:29 -0400 Subject: [PATCH] fix(core): compare CUdeviceptr as int so the VMM grow fast path can run `modify_allocation()` asks `cuMemAddressReserve` for the address immediately after the current range, then checks whether the driver granted that exact address before taking the fast path. `cuMemAddressReserve` returns `new_ptr` as a `CUdeviceptr`. That type defines `__int__` and `__repr__` but no `__eq__` or `__richcmp__`, so `new_ptr != ` falls back to identity comparison and is unconditionally true. The guard is therefore always taken and `_grow_allocation_fast_path` is unreachable: every grow runs the slow path, re-reserving the whole range and remapping it, so the buffer's base pointer changes even when a contiguous extension was available. Convert with `int()` before comparing. Reported as defect 2 in #2388. `test_vmm_allocator_grow_dispatches_to_fast_path` stubs `cuMemAddressReserve` to grant the requested address and asserts the dispatch reaches the fast path. It fails on the current code, which records "slow". `test_vmm_allocator_grow_allocation_fast_path` already exercised `_grow_allocation_fast_path` directly, so the helper itself is covered; what was missing was coverage of the dispatch decision that reaches it. This makes a previously unreachable code path live, so it is a behavior change rather than a pure cleanup: successful contiguous grows now preserve the base pointer instead of moving it. Noted in the 1.2.0 release notes. Signed-off-by: Aryan Putta --- .../core/_memory/_virtual_memory_resource.py | 2 +- cuda_core/docs/source/release/1.2.0-notes.rst | 6 +++ cuda_core/tests/test_memory.py | 47 +++++++++++++++++++ 3 files changed, 54 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..82c0a486b48 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py @@ -265,7 +265,7 @@ def modify_allocation( 0, ) - if res != driver.CUresult.CUDA_SUCCESS or new_ptr != (int(buf.handle) + aligned_prev_size): + if res != driver.CUresult.CUDA_SUCCESS or int(new_ptr) != (int(buf.handle) + aligned_prev_size): # Check for specific errors that are not recoverable with the slow path if res in ( driver.CUresult.CUDA_ERROR_INVALID_VALUE, 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..cc0179fb1bd 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 `__) +- Growing a :class:`VirtualMemoryResource` allocation now preserves the buffer's + base pointer when the driver can extend the reservation contiguously. + Previously an address comparison never matched, so every grow re-reserved and + remapped the entire range. + (`#2407 `__) + Deprecation Notices ------------------- diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 5ef48919a50..bc9a051461c 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1035,6 +1035,53 @@ def __init__(self, size): assert ("set_access", new_ptr, aligned_additional, 1) in calls +@pytest.mark.agent_authored(model="claude-opus-4-8") +def test_vmm_allocator_grow_dispatches_to_fast_path(init_cuda, monkeypatch): + """A contiguous reservation must dispatch to the fast path, not the slow path. + + `test_vmm_allocator_grow_allocation_fast_path` covers the fast path itself by + calling it directly, so it stays green whether or not `modify_allocation` can + reach it. This covers the dispatch decision instead. + """ + 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"), + ) + try: + buffer = vmm_mr.allocate(2 * 1024 * 1024) + except NotImplementedError: + pytest.skip("handle_type not implemented on this platform") + + # Grant the reservation at exactly the address the caller asked for, which is + # the contiguous-extension case the fast path exists to serve. + def fake_addr_reserve(size, alignment, addr, flags): + return (driver.CUresult.CUDA_SUCCESS, driver.CUdeviceptr(addr)) + + chosen = [] + + def record_fast(*_args, **_kwargs): + chosen.append("fast") + return buffer + + def record_slow(*_args, **_kwargs): + chosen.append("slow") + return buffer + + monkeypatch.setattr(driver, "cuMemAddressReserve", fake_addr_reserve) + monkeypatch.setattr(vmm_mr, "_grow_allocation_fast_path", record_fast) + monkeypatch.setattr(vmm_mr, "_grow_allocation_slow_path", record_slow) + + try: + vmm_mr.modify_allocation(buffer, 4 * 1024 * 1024) + assert chosen == ["fast"] + finally: + buffer.close() + + def test_vmm_allocator_rdma_unsupported_exception(): """Test that VirtualMemoryResource throws an exception when RDMA is requested but device doesn't support it.