cuda.core: allow Buffer deallocation stream to be updated before teardown

[Andy-Jost]

Summary

Add a public method for changing the stream that will order a Buffer's eventual deallocation without closing the buffer:

Buffer.set_deallocation_stream(stream: Stream | GraphBuilder) -> None

This is a follow-up to #2526.

Problem

A Buffer records a deallocation recipe when it acquires an allocation. Buffer.close(stream=...) can replace that recipe, but it immediately closes the Buffer.

Applications may transfer responsibility for a live allocation from work on stream A to work on stream B while retaining the Buffer or relying on automatic/deferred destruction. After establishing the necessary stream dependency, there is currently no public way to make stream B order the eventual deallocation without also closing the buffer.

This matters when the final release occurs through garbage collection or through another retained DevicePtrHandle, such as a graph attachment.

Proposed API

buffer.set_deallocation_stream(stream)

The method should:

• Accept the same Stream | GraphBuilder values as Buffer.close(stream=...).
• Replace the stored deallocation recipe without freeing or otherwise invalidating the buffer.
• Use the context-binding and PTDS handling introduced by cuda.core: capture complete Buffer deallocation recipe at creation #2526.
• Reject None.
• Raise if the buffer is already closed.
• Return None.

Buffer.close(stream=stream) should remain equivalent to updating the deallocation stream and then closing the buffer.

Ordering semantics

This method does not enqueue synchronization, transfer ownership to a CUDA stream, or establish dependencies between streams. The caller must ensure that the new deallocation stream is ordered after allocation and all uses of the memory, using events or other CUDA synchronization mechanisms as appropriate.

Concurrent mutation and destruction of the same buffer remain the caller's responsibility.

Reflection

Do not add a public getter in this initial change.

The stored DeallocationStream is a teardown recipe rather than an ordinary cuda.core.Stream: default-stream recipes include a bound context, and PTDS recipes also retain host-thread metadata. Returning only a Stream would not faithfully represent or round-trip the complete recipe.

A separate immutable diagnostic representation can be considered later if a concrete introspection use case arises.

Acceptance criteria

• A live buffer can have its deallocation stream changed without being closed.
• Subsequent close() or automatic destruction uses the replacement recipe.
• The buffer remains usable after set_deallocation_stream().
• Context-bound default-stream and PTDS behavior matches cuda.core: capture complete Buffer deallocation recipe at creation #2526.
• Existing close(stream=...) behavior remains unchanged.

Related

• [BUG]: Buffer cleanup must capture a complete deallocation recipe at creation #2497
• cuda.core: capture complete Buffer deallocation recipe at creation #2526
• cuda.core: reject stream=None — require explicit stream everywhere #2001
• Inconsistent lifetime management across from_handle / from_* APIs #1989