[BUG]: Buffer cleanup must capture a complete deallocation recipe at creation

[Andy-Jost]

Summary

Buffer / device-pointer destruction can run with no suitable CUDA context current on the destroying thread. The free then fails (or was historically treated as success for CUDA_ERROR_INVALID_CONTEXT), leaking memory. This shows up especially for graph-retained buffers (deferred cleanup) and can also arise when the last reference is dropped off the allocating thread (e.g. cyclic GC under free-threading).

The fix is to capture a complete cleanup recipe at Buffer allocation / binding time, so teardown does not depend on ambient thread state for context.

Problem

Owning Buffers eventually call into the driver to free (pool / cuMemFreeAsync, MR deallocate, etc.). That call needs:

1. A stream that defines free ordering (for stream-ordered paths), and
2. A context that makes the driver call legal when destruction runs.

Destruction is not guaranteed to run on the allocating thread with that context still current:

• Graph attachments are released via deferred cleanup (CUDA user-object dtor cannot call CUDA; work is drained later via pending call / another thread).
• Ordinary last-ref / GC teardown can run on whichever thread drops the last reference.

Today, some owning paths store a deallocation StreamHandle, but default-stream tokens (CU_STREAM_LEGACY / CU_STREAM_PER_THREAD) are not real streams: their meaning depends on ambient state when the driver call runs. Separately, some adoption paths never record a stream at all and fall back to default_stream() only at destructor time (see Buffer.from_handle with mr=).

A StreamHandle already embeds an owning ContextHandle for created streams (resource-handles model). The cleanup recipe should be that handle (stream + context), not a side channel.

Proposed direction

At allocation / binding, record a complete deallocation recipe: a StreamHandle that carries the context needed to replay the free. At free time: make that context current if needed, free on that stream, restore the prior binding.

Concrete work includes:

• Ensure every owning Buffer path captures the recipe when the Buffer is created (pool allocate, async alloc, IPC import, MR-backed adoption, etc.).
• Add missing stream= (keyword-only) where adoption APIs take ownership of stream-ordered teardown — notably Buffer.from_handle / ManagedBuffer.from_handle when mr= is set — so the destructor is not the first place a stream is chosen.
• Bind context into default-stream tokens at capture when forming the stored StreamHandle (legacy and PTDS tokens both need a definite context for the free call to be legal).
• Fail loudly on cleanup errors; do not treat CUDA_ERROR_INVALID_CONTEXT as successful deallocation.

PTDS policy

CU_STREAM_PER_THREAD is local to both the host thread and the context. There is no public API to promote the token to a transferable real CUstream (driver-internal PTDS streams are not given public handles).

Decision: keep the PTDS token in the captured recipe (token + bound context). cuda.core will restore the context and pass CU_STREAM_PER_THREAD at free time. Applications are responsible for ordering if they use PTDS together with:

• free-threading / off-thread last-ref or GC destruction, or
• graph-retained buffers (deferred cleanup is not the allocating thread).

In those cases the free runs as the cleanup thread's PTDS in the bound context, which may not be the allocating thread's PTDS. Same-thread PTDS alloc/use/free remains coherent.

Optional UX: detect when a Buffer whose deallocation stream is PTDS is attached to a graph, and emit a warning that deferred release will not preserve the allocating thread's PTDS ordering. (Warning only — consistent with accepting the risk. Free-threading GC has the same hazard without graphs and should be called out in docs even if we only warn on graph attach.)

Related: bare default-stream query caching is #2485 / #2490 (ambient token semantics for .context / .device). Complementary; does not by itself fix captured free-ordering.

Non-goals

• Choosing the free stream only in the destructor via default_stream() when ownership was taken earlier.
• Replacing PTDS with a newly created stream at capture, or forbidding PTDS at Buffer creation (rejected; see policy above).
• A public "promote PTDS to real stream" API (does not exist).

Related

• [BUG]: default_stream() singletons permanently cache the first-seen context and device #2485 / cuda.core: resolve default-stream context per call instead of caching it #2490 — default-stream tokens must not sticky-cache context on the singleton.
• cuda.core: reject stream=None — require explicit stream everywhere #2001 (closed) — explicit stream= on allocate/deallocate paths; adoption / from_handle still incomplete for recorded teardown.
• Inconsistent lifetime management across from_handle / from_* APIs #1989 — from_handle / lifetime consistency more broadly.

Acceptance sketch

• Owning Buffer creation always records a complete cleanup recipe (StreamHandle with embedded context for the free).
• from_handle(..., mr=...) (and managed counterpart) accept stream= and store it for teardown; no destructor-only default-stream fallback for that path.
• Deferred / no-current-context destruction can activate the bound context and free.
• PTDS + cross-thread / graph deferred free is documented as app-ordered; optional warning on graph attach of PTDS-backed buffers.