cuda.core: DeviceMemoryResource setup is refused under non-relaxed stream capture, and tearing down the invalidated GraphBuilder segfaults

從 cuda_core/cuda/core/_memory/_device_memory_resource.pyx 和 _memory_pool.pyx 開始，重點關注 _DMR_init 和 MP_raise_release_threshold。接著檢查 _graph_builder.pyx 和 cuda/core/_cpp/resource_handles.cpp 中 GraphBuilder 的 begin/end 路徑以及 create_graph_handle。完成的標準是實現 capture-safe 的資源設定，使失效的 captures 結束時不會銷毀過時的 graph，並補充回歸覆蓋，其中包括 subprocess no-crash 檢查。

Summary

Two related defects surfaced while checking whether cuda.core is affected by the libcu++ issue fixed in NVIDIA/cccl#11360 (default memory pool resolution fails under stream capture).

1. DeviceMemoryResource(device) with no options wraps the driver's current pool for the device and raises its release threshold. That step calls cuMemPoolGetAttribute, and cuMemPoolSetAttribute when the threshold is zero. The driver treats both as potentially unsafe calls while the calling thread is inside a global or thread-local capture: the constructor raises and the capture is invalidated. Device.memory_resource constructs this resource lazily on first access, so a first allocation through Device.allocate can invalidate a capture that is in progress. cuda.core's own GraphBuilder.begin_building() defaults to "relaxed", which passes the check, so this hits code that asks for "global" or "thread_local", and captures started by other libraries on the same thread (PyTorch's graph capture defaults to global mode). A global capture on another thread is invalidated too.

2. After any invalidated capture, tearing down the GraphBuilder segfaults inside cuGraphDestroy. The builder holds an owning handle to the capture graph it obtained from cuStreamGetCaptureInfo. When the capture was invalidated, cuStreamEndCapture returns a NULL graph and the driver releases the capture graph itself, so the builder's later cuGraphDestroy is a use-after-free. This is independent of item 1: any invalidation triggers it, in all three capture modes, and the minimal reproducer crashes deterministically.

What was observed

cuda.core main at 4357f375, cuda.bindings 13.4.1, Python 3.14.7, Linux x86_64. Reproduced on an H100 PCIe (driver 610.57.04) and an H200 (driver 595.58.03). Each scenario below runs in a fresh process.

Item 1, minimal reproducer:

from cuda.core import Device, DeviceMemoryResource

dev = Device()
dev.set_current()
gb = dev.create_stream().create_graph_builder().begin_building(mode="global")  # or "thread_local"
DeviceMemoryResource(dev)  # raises CUDAError: CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED
gb.is_building             # RuntimeError: the build process has been invalidated

Item 2, minimal reproducer. Device.sync() is refused whenever a stream in the context is capturing, which is expected and is used here only to invalidate the capture. The crash afterwards is the bug. It reproduces in all three modes:

from cuda.core import Device

dev = Device()
dev.set_current()
gb = dev.create_stream().create_graph_builder().begin_building()
try:
    dev.sync()          # refused under capture; the capture is now invalidated
except Exception:
    pass
try:
    gb.end_building()   # RuntimeError: invalidated; the capture is not ended
except Exception:
    pass
del gb                  # segfault inside cuGraphDestroy

Backtrace of the crash (trimmed):

#3  cuGraphDestroy () from libcuda.so.1
#5  operator() at cuda/core/_cpp/resource_handles.cpp:1958   # GraphHierarchy deleter from create_graph_handle
#13 __Pyx_call_destructor<std::shared_ptr<CUgraph_st* const>> at _graph_builder.cpp
#14 __pyx_tp_dealloc ... GraphBuilder
#15 _Py_Dealloc

Analysis

Item 1 lives in _DMR_init (cuda_core/cuda/core/_memory/_device_memory_resource.pyx), which calls MP_raise_release_threshold (_memory_pool.pyx) on the no-options path. The pool lookup itself (cuDeviceGetMemPool) is a query and is legal under capture; only the attribute read and write are refused. The pinned and managed resources do not raise the threshold and are unaffected. Other pool calls (owned-pool create and destroy, attributes reads, peer_accessible_by, IPC export and import) are refused under the same rule, but they are explicit user actions and are out of scope here.

Item 2: begin_building calls cuStreamGetCaptureInfo after cuStreamBeginCapture and wraps the capturing graph with create_graph_handle, whose deleter calls cuGraphDestroy. GB_end_capture_if_needed (called from close() and __dealloc__) then calls cuStreamEndCapture; for an invalidated capture the driver returns the invalidation error and a NULL graph, and the builder's owning handle still destroys the stale graph afterwards. Two related gaps make the situation unrecoverable from Python: end_building() raises from is_building before it ever calls cuStreamEndCapture, and close() raises from the end-capture error before it resets its handles. #2776 is a sibling: a fork left mid-capture also crashes at collection.

Suggested direction

Item 1: run the threshold read and write inside a relaxed-capture scope, mirroring the CCCL fix. cuThreadExchangeStreamCaptureMode swaps the calling thread's mode to relaxed and the scope restores the previous mode on exit. It is per-thread, needs no stream, costs two cheap driver calls, and has no observable effect when the thread is not capturing. The attribute write executes immediately rather than being recorded, which is the intent for a process-wide setting. Add a test that constructs DeviceMemoryResource and touches Device.memory_resource under each capture mode and checks that the capture stays valid.

Item 2: the builder must not own the capture graph while the capture is in progress. Hold a non-owning handle from begin_building, and take ownership only of the graph returned by a successful cuStreamEndCapture; equivalently, release _h_graph without cuGraphDestroy whenever end-capture returns NULL or an error. Make end_building() and close() end an invalidated capture and raise the invalidation error, leaving the builder closed and safe to collect. Add a regression test, run in a subprocess, that invalidates a capture, drops the builder, and expects no crash.