fix(cuda.core): route pool-backed MR buffers through deallocate()

Pool-backed and graph memory resources now wrap raw allocations with
MR-owned device pointer handles, matching Buffer.from_handle(mr=...).
Subclasses can observe or customize teardown via deallocate().

Fixes #2615

Signed-off-by: Omar Atie <atiaomar1978-hub@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: atiaomar1978-hub

cuda_core/cuda/core/_memory/_buffer.pxd

@@ -2,10 +2,14 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from libc.stddef cimport size_t
+
 from libcpp cimport bool as cpp_bool
 from libcpp.atomic cimport atomic as std_atomic
 
+from cuda.bindings cimport cydriver
 from cuda.core._resource_handles cimport DevicePtrHandle
+from cuda.core._stream cimport Stream
 
 
 cdef struct _MemAttrs:
@@ -46,6 +50,15 @@ cdef Buffer Buffer_from_deviceptr_handle(
 )
 
 
+# Wrap a raw device pointer with MR-owned teardown and record the stream.
+cdef DevicePtrHandle deviceptr_create_owned_by_mr(
+    cydriver.CUdeviceptr ptr,
+    size_t size,
+    object mr,
+    Stream stream,
+) except *
+
+
 # Shared argument coercion for the batched free functions (copy_batch,
 # prefetch_batch, discard_batch, discard_prefetch_batch). `single_hint`
 # names the per-buffer API to use instead when a bare Buffer is passed.

cuda_core/cuda/core/_memory/_buffer.pyx

@@ -92,6 +92,22 @@ cdef inline void _apply_deallocation_stream(
     HANDLE_RETURN(status)
 
 
+cdef DevicePtrHandle deviceptr_create_owned_by_mr(
+    cydriver.CUdeviceptr ptr,
+    size_t size,
+    object mr,
+    Stream stream,
+) except *:
+    """Create an MR-owned device pointer handle with a recorded deallocation stream."""
+    cdef DevicePtrHandle h_ptr = deviceptr_create_with_mr(ptr, size, mr)
+    try:
+        _apply_deallocation_stream(h_ptr, stream._h_stream)
+    except BaseException:
+        h_ptr.reset()
+        raise
+    return h_ptr
+
+
 __all__ = ['Buffer', 'MemoryResource']
 
 

cuda_core/cuda/core/_memory/_graph_memory_resource.pyx

@@ -7,11 +7,14 @@ from __future__ import annotations
 from libc.stdint cimport intptr_t
 
 from cuda.bindings cimport cydriver
-from cuda.core._memory._buffer cimport Buffer, Buffer_from_deviceptr_handle, MemoryResource
+from cuda.core._memory._buffer cimport (
+    Buffer,
+    Buffer_from_deviceptr_handle,
+    MemoryResource,
+    deviceptr_create_owned_by_mr,
+)
 from cuda.core._resource_handles cimport (
     DevicePtrHandle,
-    deviceptr_alloc_async,
-    get_last_error,
     as_cu,
 )
 
@@ -210,18 +213,12 @@ cdef inline int check_capturing(cydriver.CUstream s) except?-1 nogil:
 
 cdef inline Buffer GMR_allocate(cyGraphMemoryResource self, size_t size, Stream stream):
     cdef cydriver.CUstream s = as_cu(stream._h_stream)
+    cdef cydriver.CUdeviceptr ptr
     cdef DevicePtrHandle h_ptr
     with nogil:
         check_capturing(s)
-        h_ptr = deviceptr_alloc_async(size, stream._h_stream)
-    if not h_ptr:
-        HANDLE_RETURN(get_last_error())
-        raise RuntimeError(
-            f"Failed to allocate {size} bytes from GraphMemoryResource: "
-            "cuda-core returned an empty allocation handle without recording a CUDA error. "
-            "This is an internal cuda-core error; please report it with your CUDA driver, "
-            "CUDA Toolkit, and cuda-python versions."
-        )
+        HANDLE_RETURN(cydriver.cuMemAllocAsync(&ptr, size, s))
+    h_ptr = deviceptr_create_owned_by_mr(ptr, size, self, stream)
     return Buffer_from_deviceptr_handle(h_ptr, size, self, None)
 
 

cuda_core/cuda/core/_memory/_memory_pool.pyx

@@ -10,14 +10,18 @@ from libc.stdint cimport uintptr_t
 from libc.string cimport memset
 
 from cuda.bindings cimport cydriver
-from cuda.core._memory._buffer cimport Buffer, Buffer_from_deviceptr_handle, MemoryResource
+from cuda.core._memory._buffer cimport (
+    Buffer,
+    Buffer_from_deviceptr_handle,
+    MemoryResource,
+    deviceptr_create_owned_by_mr,
+)
 from cuda.core._memory cimport _ipc
 from cuda.core._stream cimport Stream_accept, Stream
 from cuda.core._resource_handles cimport (
     MemoryPoolHandle,
     DevicePtrHandle,
     create_mempool_handle,
-    deviceptr_alloc_from_pool,
     get_last_error,
     as_cu,
     as_py,
@@ -330,18 +334,12 @@ cdef inline int check_not_capturing(cydriver.CUstream s) except?-1 nogil:
 
 cdef Buffer _MP_allocate(_MemPool self, size_t size, Stream stream, type cls = Buffer):
     cdef cydriver.CUstream s = as_cu(stream._h_stream)
+    cdef cydriver.CUdeviceptr ptr
     cdef DevicePtrHandle h_ptr
     with nogil:
         check_not_capturing(s)
-        h_ptr = deviceptr_alloc_from_pool(size, self._h_pool, stream._h_stream)
-    if not h_ptr:
-        HANDLE_RETURN(get_last_error())
-        raise RuntimeError(
-            f"Failed to allocate {size} bytes from {self.__class__.__name__}: "
-            "cuda-core returned an empty allocation handle without recording a CUDA error. "
-            "This is an internal cuda-core error; please report it with your CUDA driver, "
-            "CUDA Toolkit, and cuda-python versions."
-        )
+        HANDLE_RETURN(cydriver.cuMemAllocFromPoolAsync(&ptr, size, as_cu(self._h_pool), s))
+    h_ptr = deviceptr_create_owned_by_mr(ptr, size, self, stream)
     return Buffer_from_deviceptr_handle(h_ptr, size, self, None, cls)
 
 

cuda_core/tests/test_memory.py

@@ -515,6 +515,94 @@ def deallocate(self, ptr, size, *, stream=None):
     assert received["stream"].handle == stream.handle
 
 
+class _RecordingPoolMR(DeviceMemoryResource):
+    """Pool-backed MR that records deallocate() invocations."""
+
+    def __init__(self, device, options=None):
+        super().__init__(device, options)
+        self.dealloc_calls = []
+
+    def deallocate(self, ptr, size, *, stream):
+        self.dealloc_calls.append((ptr, size, stream))
+        super().deallocate(ptr, size, stream=stream)
+
+
+@pytest.mark.agent_authored(model="cursor-grok-4.6")
+@pytest.mark.parametrize("pinned", [False, True], ids=["device", "pinned"])
+def test_pool_backed_mr_deallocate_called_on_close(mempool_device, pinned):
+    """Pool-backed mr.allocate() honors overridden deallocate() on close (#2615)."""
+    dev = mempool_device
+    stream = dev.default_stream
+    if pinned:
+        skip_if_pinned_memory_unsupported(dev)
+
+        class RecordingMR(PinnedMemoryResource):
+            def __init__(self, options=None):
+                super().__init__(options)
+                self.dealloc_calls = []
+
+            def deallocate(self, ptr, size, *, stream):
+                self.dealloc_calls.append((ptr, size, stream))
+                super().deallocate(ptr, size, stream=stream)
+
+        mr = RecordingMR(PinnedMemoryResourceOptions(max_size=POOL_SIZE))
+    else:
+
+        class RecordingMR(DeviceMemoryResource):
+            def __init__(self, device, options=None):
+                super().__init__(device, options)
+                self.dealloc_calls = []
+
+            def deallocate(self, ptr, size, *, stream):
+                self.dealloc_calls.append((ptr, size, stream))
+                super().deallocate(ptr, size, stream=stream)
+
+        mr = RecordingMR(dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE))
+
+    buf = mr.allocate(1024, stream=stream)
+    assert buf.memory_resource is mr
+    assert len(mr.dealloc_calls) == 0
+    buf.close(stream=stream)
+    stream.sync()
+    assert len(mr.dealloc_calls) == 1
+    assert mr.dealloc_calls[0][1] == 1024
+
+
+@pytest.mark.agent_authored(model="cursor-grok-4.6")
+def test_pool_backed_mr_deallocate_called_on_gc(mempool_device):
+    """Pool-backed mr.allocate() honors overridden deallocate() on GC (#2615)."""
+    import gc
+
+    dev = mempool_device
+    stream = dev.default_stream
+    mr = _RecordingPoolMR(dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE))
+    buf = mr.allocate(1024, stream=stream)
+    assert len(mr.dealloc_calls) == 0
+    del buf
+    gc.collect()
+    stream.sync()
+    assert len(mr.dealloc_calls) == 1
+
+
+@pytest.mark.agent_authored(model="cursor-grok-4.6")
+def test_pool_backed_mr_deallocate_receives_stream(mempool_device):
+    """Pool-backed mr.allocate() forwards close(stream) to deallocate() (#2615)."""
+    dev = mempool_device
+    stream = dev.create_stream()
+    received = {}
+
+    class StreamCapturePoolMR(_RecordingPoolMR):
+        def deallocate(self, ptr, size, *, stream):
+            received["stream"] = stream
+            super().deallocate(ptr, size, stream=stream)
+
+    mr = StreamCapturePoolMR(dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE))
+    buf = mr.allocate(1024, stream=stream)
+    buf.close(stream=stream)
+    stream.sync()
+    assert received["stream"].handle == stream.handle
+
+
 @pytest.mark.parametrize("buffer_type", [Buffer, ManagedBuffer])
 def test_from_handle_mr_records_default_stream(buffer_type):
     """When a Buffer/ManagedBuffer is minted via :meth:`from_handle` with ``mr``
@@ -2304,25 +2392,36 @@ def test_dmr_handle_and_ownership(mempool_device):
 
 @pytest.mark.agent_authored(model="claude-opus-4.8")
 def test_dmr_deallocate_frees_pool_pointer(mempool_device):
-    """Closing a Buffer.from_handle(..., mr=mr) view frees the pointer via the Python
-    _MemPool.deallocate path; the pool's in-use bytes drop back."""
+    """Closing a buffer from mr.allocate() frees via deallocate() and returns pool bytes."""
+    dev = mempool_device
+    stream = dev.default_stream
+    mr = DeviceMemoryResource(dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE))
+    size = 256
+    buf = mr.allocate(size, stream=stream)
+    stream.sync()
+    used_after_alloc = mr.attributes.used_mem_current
+    assert used_after_alloc >= size
+    buf.close(stream=stream)
+    stream.sync()
+    assert int(buf.handle) == 0
+    assert mr.attributes.used_mem_current < used_after_alloc
+
+
+@pytest.mark.agent_authored(model="cursor-grok-4.6")
+def test_dmr_from_handle_deallocate_frees_pool_pointer(mempool_device):
+    """Buffer.from_handle(..., mr=mr) also routes teardown through deallocate()."""
     dev = mempool_device
     stream = dev.default_stream
     mr = DeviceMemoryResource(dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE))
     size = 256
-    # Raw pool allocation owned by nobody else, so exactly one owner frees it (no
-    # double free); a Buffer.from_handle view then routes teardown through the
-    # Python deallocate path that mr.allocate()'s C++-direct free would skip.
     ptr = handle_return(driver.cuMemAllocFromPoolAsync(size, mr.handle, stream.handle))
     stream.sync()
     used_after_alloc = mr.attributes.used_mem_current
     assert used_after_alloc >= size
     buf = Buffer.from_handle(int(ptr), size, mr=mr)
-    buf.close(stream)
+    buf.close(stream=stream)
     stream.sync()
     assert int(buf.handle) == 0
-    # In-use bytes fell back, so the pointer was actually returned (buf.handle == 0
-    # alone wouldn't prove it: the deleter callback swallows a failed free).
     assert mr.attributes.used_mem_current < used_after_alloc