cuda.core: allow updating Buffer deallocation streams

cuda_core/cuda/core/_memory/_buffer.pyi

@@ -127,6 +127,41 @@ class Buffer:
         stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional
             The stream object to use for asynchronous deallocation. If None,
             the deallocation stream stored in the handle is used.
+
+        See Also
+        --------
+        set_deallocation_stream
+            Change the deallocation stream without closing the buffer.
+        """
+
+    def set_deallocation_stream(self, stream: Stream | GraphBuilder) -> None:
+        """Change the stream that orders this buffer's eventual deallocation.
+
+        The buffer remains open and usable. A later :meth:`close` without a
+        stream, garbage collection, or release of the final retained device
+        pointer handle uses the replacement stream.
+
+        This method does not synchronize streams or establish dependencies.
+        The caller must ensure that allocation and all accesses are ordered
+        before the deallocation on ``stream``.
+
+        Parameters
+        ----------
+        stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
+            The stream to use for eventual asynchronous deallocation.
+
+        Raises
+        ------
+        RuntimeError
+            If the buffer is already closed, or if a default-stream token
+            cannot be bound because no CUDA context is current.
+        TypeError
+            If ``stream`` is ``None`` or is not an accepted stream object.
+
+        Notes
+        -----
+        Synchronizing concurrent mutation and destruction of the same buffer
+        is the caller's responsibility.
         """
 
     def __enter__(self):

cuda_core/cuda/core/_memory/_buffer.pyx

@@ -252,7 +252,7 @@ cdef class Buffer:
         # The parent process's stream is not portable across processes, so the
         # pickle path cannot thread an explicit stream through. Seed the
         # imported buffer's deallocation with the current context's default
-        # stream; the receiver can override via buffer.close(stream).
+        # stream; the receiver can override it before or during close.
         return Buffer.from_ipc_descriptor(mr, ipc_descriptor, stream=default_stream())
 
     def __reduce__(self) -> tuple[object, ...]:
@@ -347,9 +347,45 @@ cdef class Buffer:
         stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional
             The stream object to use for asynchronous deallocation. If None,
             the deallocation stream stored in the handle is used.
+
+        See Also
+        --------
+        set_deallocation_stream
+            Change the deallocation stream without closing the buffer.
         """
         Buffer_close(self, stream)
 
+    def set_deallocation_stream(self, stream: Stream | GraphBuilder) -> None:
+        """Change the stream that orders this buffer's eventual deallocation.
+
+        The buffer remains open and usable. A later :meth:`close` without a
+        stream, garbage collection, or release of the final retained device
+        pointer handle uses the replacement stream.
+
+        This method does not synchronize streams or establish dependencies.
+        The caller must ensure that allocation and all accesses are ordered
+        before the deallocation on ``stream``.
+
+        Parameters
+        ----------
+        stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`
+            The stream to use for eventual asynchronous deallocation.
+
+        Raises
+        ------
+        RuntimeError
+            If the buffer is already closed, or if a default-stream token
+            cannot be bound because no CUDA context is current.
+        TypeError
+            If ``stream`` is ``None`` or is not an accepted stream object.
+
+        Notes
+        -----
+        Synchronizing concurrent mutation and destruction of the same buffer
+        is the caller's responsibility.
+        """
+        Buffer_set_deallocation_stream(self, stream)
+
     def __enter__(self):
         return self
 
@@ -707,15 +743,21 @@ cdef tuple Buffer_coerce_batch(object buffers, str what, str single_hint):
     return tuple(out)
 
 
+cdef inline void Buffer_set_deallocation_stream(Buffer self, object stream):
+    """Validate and replace a live buffer's deallocation recipe."""
+    if not self._h_ptr:
+        raise RuntimeError("Cannot set the deallocation stream on a closed Buffer")
+    cdef Stream s = Stream_accept(stream)
+    _apply_deallocation_stream(self._h_ptr, s._h_stream)
+
+
 cdef inline void Buffer_close(Buffer self, object stream):
     """Close a buffer, freeing its memory."""
-    cdef Stream s
     if not self._h_ptr:
         return
     # Update deallocation stream if provided
     if stream is not None:
-        s = Stream_accept(stream)
-        _apply_deallocation_stream(self._h_ptr, s._h_stream)
+        Buffer_set_deallocation_stream(self, stream)
     # Reset handle - RAII deleter will free the memory (and release owner ref in C++)
     self._h_ptr.reset()
     self._size = 0

cuda_core/docs/source/api.rst

@@ -77,6 +77,14 @@ Memory management
    ManagedMemoryResourceOptions
    VirtualMemoryResourceOptions
 
+A :class:`Buffer` records the stream that will order its eventual deallocation.
+Use :meth:`Buffer.set_deallocation_stream` to replace that stream without
+closing the buffer. Changing the recorded stream does not synchronize streams;
+the caller must order allocation and every access before the deallocation,
+using events or other CUDA synchronization mechanisms as needed. See
+:cuda-core-example:`buffer_deallocation_stream.py <buffer_deallocation_stream.py>`
+for a complete example.
+
 
 CUDA compilation toolchain
 --------------------------

cuda_core/docs/source/examples.rst

@@ -39,6 +39,13 @@ Linking and graphs
 - :cuda-core-example:`cuda_graphs.py <cuda_graphs.py>`
   captures and replays a multi-kernel CUDA graph to reduce launch overhead.
 
+Memory management
+-----------------
+
+- :cuda-core-example:`buffer_deallocation_stream.py <buffer_deallocation_stream.py>`
+  transfers a buffer between streams and safely changes the stream that orders
+  its deallocation.
+
 Interoperability and memory access
 ----------------------------------
 

cuda_core/docs/source/release/1.2.0-notes.rst

@@ -35,6 +35,10 @@ Fixes and enhancements
   free recipe can pin that context.
   (`#2497 <https://github.com/NVIDIA/cuda-python/issues/2497>`__)
 
+- Added :meth:`Buffer.set_deallocation_stream` to change the stream that orders
+  a buffer's eventual deallocation without closing the buffer.
+  (`#2600 <https://github.com/NVIDIA/cuda-python/issues/2600>`__)
+
 - Explicit calls to ``deallocate()`` on pool-backed memory resources and
   :class:`GraphMemoryResource` now propagate errors from the underlying CUDA
   free operation. Previously, these errors could be suppressed. Automatic

cuda_core/examples/buffer_deallocation_stream.py

@@ -0,0 +1,76 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# ################################################################################
+#
+# This example transfers a buffer from a producer stream to a consumer stream.
+# An event orders the consumer after the producer. The buffer then records the
+# consumer stream for its eventual deallocation.
+#
+# ################################################################################
+
+# /// script
+# dependencies = ["cuda_bindings", "cuda_core"]
+# ///
+
+import ctypes
+
+from cuda.core import Device, LegacyPinnedMemoryResource
+
+
+def produce_data(device, stream, size, value):
+    """Allocate and fill a buffer on the producer stream."""
+    buffer = device.allocate(size, stream=stream)
+    buffer.fill(value, stream=stream)
+    ready = stream.record()
+    return buffer, ready
+
+
+def consume_data(buffer, ready, output, stream):
+    """Submit consumer work and transfer the deallocation stream."""
+    stream.wait(ready)
+    buffer.set_deallocation_stream(stream)
+    buffer.copy_to(output, stream=stream)
+
+
+def main():
+    device = Device()
+    device.set_current()
+    producer_stream = device.create_stream()
+    consumer_stream = device.create_stream()
+    pinned_mr = LegacyPinnedMemoryResource()
+
+    size = 4096
+    value = 42
+    buffer = None
+    ready = None
+    output = None
+
+    try:
+        output = pinned_mr.allocate(size)
+        buffer, ready = produce_data(device, producer_stream, size, value)
+        consume_data(buffer, ready, output, consumer_stream)
+
+        # No stream argument is needed. The buffer now records consumer_stream.
+        # The free operation runs after the copy on that stream.
+        buffer.close()
+        buffer = None
+        consumer_stream.sync()
+
+        result = ctypes.string_at(int(output.handle), output.size)
+        assert result == bytes([value]) * size
+        print("Buffer deallocation stream transfer completed.")
+    finally:
+        if buffer is not None:
+            buffer.close()
+        if output is not None:
+            output.close()
+        if ready is not None:
+            ready.close()
+        consumer_stream.close()
+        producer_stream.close()
+
+
+if __name__ == "__main__":
+    main()

cuda_core/tests/test_memory.py

@@ -516,6 +516,62 @@ def test_mr_deallocate_receives_stream():
     assert telemetry["deallocations"][-1]["stream"].handle == stream.handle
 
 
+@pytest.mark.agent_authored(model="gpt-5.6")
+@pytest.mark.parametrize(
+    ("configuration", "destruction"),
+    [
+        ("initialization", "close"),
+        ("initialization", "gc"),
+        ("setter", "close"),
+        ("setter", "gc"),
+        ("close", "close"),
+    ],
+)
+def test_buffer_deallocation_stream_configuration_paths(configuration, destruction):
+    """Creation, mutation, and close overrides use the requested stream."""
+    import gc
+
+    device = Device()
+    device.set_current()
+    initial_stream = device.create_stream()
+    target_stream = device.create_stream()
+    CapturingMR, telemetry = make_instrumented_memory_resource(record_streams=True)
+    mr = CapturingMR(device)
+
+    stream = target_stream if configuration == "initialization" else initial_stream
+    buf = Buffer.from_handle(1, 1024, mr=mr, stream=stream)
+    if configuration == "setter":
+        handle = buf.handle
+        buf.set_deallocation_stream(target_stream)
+        assert buf.handle == handle
+        assert buf.size == 1024
+
+    if destruction == "close":
+        buf.close(stream=target_stream if configuration == "close" else None)
+    else:
+        del buf
+        gc.collect()
+
+    assert len(telemetry["deallocations"]) == 1
+    assert telemetry["deallocations"][0]["stream"].handle == target_stream.handle
+
+
+@pytest.mark.agent_authored(model="gpt-5.6")
+def test_set_deallocation_stream_rejects_none_and_closed_buffer():
+    device = Device()
+    device.set_current()
+    stream = device.create_stream()
+    mr = StubMemoryResource(device)
+    buf = Buffer.from_handle(1, 1024, mr=mr, stream=stream)
+
+    with pytest.raises(TypeError, match="stream is required"):
+        buf.set_deallocation_stream(None)
+
+    buf.close()
+    with pytest.raises(RuntimeError, match="closed Buffer"):
+        buf.set_deallocation_stream(stream)
+
+
 @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``