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

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.

Per-Thread Default Stream (PTDS) policy

CU_STREAM_PER_THREAD is local to both the host thread and the context. There is no public API to promote the PTDS token to a transferable real CUstream (driver-internal PTDS streams are not given public handles). Therefore, when PTDS is used for destruction ordering it is impossible to capture a complete deallocation recipe. Deallocations running on a separate thread will be ordered relative to that thread's default stream rather than the default stream that was active at allocation time.

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).

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.