fix(cuda.core): move VirtualMemoryResource onto the _rt handle layer

[Andy-Jost]

Summary

Moves VirtualMemoryResource onto the _rt handle layer, as planned in #2906 (design). Each physical allocation, address reservation and mapping is a std::shared_ptr handle with a deleter that knows the exact driver call to undo it. A buffer owns a range of mappings through its device pointer handle, so everything a buffer maps is released when the last buffer that maps it closes, and a failed multi-step operation unwinds by letting its local handles die. Ranges are immutable: a grow copies the input's mapping list and builds a new range for its result, so two buffers never share mutable state and each buffer owns exactly one range and one recorded stream. The module moves from Python to Cython. The design lands as cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md.

Behavior changes

• modify_allocation returns a new VirtualMemoryBuffer and leaves the buffer passed in open. The two alias the same physical memory, which is freed when the last of them closes. The pointer is preserved when the driver grants the adjacent address range. The result is always a new buffer, so closing it never closes the buffer passed in; a request the buffer already covers returns a full alias without a driver call.
• Buffer.size after a grow is the aligned total.
• config= applies to the chunk the call adds and is no longer stored on the resource. It must keep the resource's location and passes the constructor's option checks, including the RDMA support check.
• Buffers from allocate() free themselves when they close and no longer call deallocate(), which now serves pointers wrapped with Buffer.from_handle.
• A buffer records the stream passed to allocate() and synchronizes it when it closes, before releasing its share of the mappings; a mapping another buffer still holds stays mapped. An explicit close() on a capturing stream other than a default stream raises; a release ordered on a default stream that would disturb a capture in its context is reported as a CUDAWarning and unmaps without the synchronization. The synchronization runs in relaxed capture mode, so it does not invalidate a capture on an unrelated stream.
• modify_allocation accepts only buffers this resource returned.
• location_type="host" requires handle_type=None, which the driver requires.
• allocate(0) returns an empty buffer without a driver call.

Testing

• The mocked grow test is replaced by real-driver tests. Release is checked per address, by asking the driver that no mapping and no reservation remain at the exact addresses a buffer used, instead of comparing a device-wide free-memory count that other processes can move. Covered: release in three modes (allocate, grow, forced move), grow in place with the pointer preserved, grow that moves with contents preserved, aliasing and cross-alias visibility, unaligned and repeated grows, a failed grow that leaves the input intact, the capture check on close with and without a stream argument, a release from garbage collection during a capture, a close during an unrelated global-mode or thread-local capture, stream synchronization at close in both close orders (deterministic, with a sleep kernel and events), cuMemGetAccess on both chunks after a config= grow, the per-call config checks, the raw-pointer deallocate() path, host location without a current context, size zero with an inherited stream, buffers alive at interpreter shutdown, and four threads growing and closing aliases of one buffer at once, reduced from the reproduction in review.
• The graph in-flight retention test is parametrized over the device pool and the virtual memory resource: a memcpy node keeps a virtual memory buffer mapped after close() until the launch completes and the graph is gone. A second test keeps the range mapped across a grow and the close of every alias.
• On an H100 with CUDA 13.4 (driver 615), the full cuda_core suite passed: 4325 passed, 100 skipped, 1 xfailed, 0 failed. The VMM selection passed in fixed order (40 passed, 1 skipped for GPUDirect RDMA) and twice in random order, and the reproduction script from review runs to completion without a crash.

Issues

Fixes #2887
Fixes #2907
Fixes #2908
Fixes #2909
Fixes #2886
Fixes #2345
Fixes #2877

Addresses #2388 items 1, 2 and 3 (the rollback that lost access grants, the dead fast path, and the finalizer warnings; item 4 landed in #2418) and the size-0, misaligned-probe and host handle-type parts of #2910. Part of #2906. Supersedes #2880, #2889 and #2237: they edit the module this PR replaces, and the fixes they carried (#2877, #2886, #2345) are part of the redesign.

🤖 Generated with Claude Code

[github-actions]

Doc Preview CI
🚀 View preview at https://nvidia.github.io/cuda-python/pr-preview/pr-2917/
https://nvidia.github.io/cuda-python/pr-preview/pr-2917/cuda-core/
https://nvidia.github.io/cuda-python/pr-preview/pr-2917/cuda-bindings/
https://nvidia.github.io/cuda-python/pr-preview/pr-2917/cuda-pathfinder/
Preview will be ready when the GitHub Pages deployment is complete.

[sylvesterkaczmarek]

The no-access case can call this with count == 0 and descs.data() from an empty vector. Constructing access(descs, descs + count) then does pointer arithmetic/range construction on a possibly null pointer. Since self_access=None with no peers is supported, could this special-case zero descriptors before forming the range?

[Andy-Jost]

The no-access case can call this with count == 0 and descs.data() from an empty vector. Constructing access(descs, descs + count) then does pointer arithmetic/range construction on a possibly null pointer. Since self_access=None with no peers is supported, could this special-case zero descriptors before forming the range?

@sylvesterkaczmarek This is well-defined C++ as written. The standard ([vector.data] C++17, N4659 section 26.3.11.4) guarantees that [data(), data() + size()] is a valid range. This is true whether or not data() returns a null pointer, since zero added to a null pointer yields a null pointer (see [expr.add] C++17, N4659 section 8.7). No workaround appears necessary here. This case is covered in test_vmm_allocate_without_access_descriptors.

[sylvesterkaczmarek]

Thanks for the correction. The C++17 vector::data contract explicitly requires [data(), data() + size()) to be a valid range, including the empty-vector case, and the dedicated no-access regression covers this path. My concern is resolved.

[0z5a]

Independent Jetson Thor / ARM64 validation of head 1f9c014b4c1f2d35fa404fd5bde2326b1ee96111, compared with base 530898e6448daf7e6d5a52dabdfd42abf38eb0f4 (CUDA Toolkit 13.2.78, driver 595.78, cuda-bindings 13.2.0). Both revisions were built from source in separate environments; the compiled head extension came from the head checkout. I am following #2906 and am not opening a competing VMM implementation PR.

Results on the real CUDA driver:

• Head pytest cuda_core/tests/test_memory.py -k vmm -vv -ra: 39 passed, 1 skipped, in each of three consecutive repeats after pausing this validation run's other build processes. Base's own selection: 16 passed, 1 skipped. The base suite does not have the head's additional ownership/grow/lifetime coverage.
• Head graph selection, pytest cuda_core/tests/graph/test_graph_definition_lifetime.py -k 'vmm or inflight_launch_retains_attachments' -vv -ra: 3 passed.
• Existing head tests executed the in-place and forced-move grow paths, content/alias preservation, repeated/unaligned grow, allocation release, graph/capture restrictions, and rollback/overflow checks. The sole VMM skip is the policy-configuration test requiring GPUDirect RDMA.
• Direct device queries: VMM, POSIX-FD handles, host VMM and host-NUMA VMM are supported; GPUDirect RDMA and fabric handles are not supported on this device.
• Direct cuMemGetAllocationGranularity queries returned 2,097,152 bytes for both minimum and recommended, for both non-exportable and POSIX-FD allocation properties. Sixteen supplemental allocations exercised g-1, g, g+1, and 2g+1 for both handle modes/policies, checking the rounded buffer size and actual driver writes/reads at both ends; all passed.

One qualification: the first non-isolated head run was 38 passed / 1 skipped / 1 failed at test_vmm_deallocate_frees_wrapped_pointer. The assertion observed a 2,228,224-byte free-memory delta against a 2,097,152-byte allocation while other builds/transfers were active. It did not recur in the three repeats; no threshold or implementation was changed. This observation does not establish a leak or its cause, but the global free-memory assertion may be sensitive to concurrent activity on unified-memory systems. The existing inference service remained active throughout.

This is platform-specific evidence, not an approval of every ownership/API design choice. No Windows/fabric/RDMA execution coverage is claimed.

[brandon-b-miller]

This appears to have the same purpose as

cuda-python/cuda_core/cuda/core/_stream.pyx

Lines 520 to 527 in f9ed2bd

	cdef inline bint Stream_is_default_token(Stream self) noexcept nogil:
	"""Return True for CU_STREAM_LEGACY and CU_STREAM_PER_THREAD.
	These tokens carry no context of their own; they refer to whatever context
	is current, so nothing resolved from one may be cached on the object.
	"""
	cdef uintptr_t h = <uintptr_t>as_cu(self._h_stream)
	return h == <uintptr_t>cydriver.CU_STREAM_LEGACY or h == <uintptr_t>cydriver.CU_STREAM_PER_THREAD

except the h==0 detail. We should only need one or at least one implementation both agree on right?

[brandon-b-miller]

we're also importing/using the above already in this file

[Andy-Jost]

I'll refactor to eliminate the duplication.

[brandon-b-miller]

A broad question about api.hpp. This pings a little like a public file / something that aspires towards true c++ interop. But it seems like it isn't yet, there's no statement about api stability or so forth. At the same time, it looks like we're shipping it today in our wheels tempting someone to maybe link against it. What's the end goal here?

[Andy-Jost]

A broad question about api.hpp. This pings a little like a public file / something that aspires towards true c++ interop. But it seems like it isn't yet, there's no statement about api stability or so forth. At the same time, it looks like we're shipping it today in our wheels tempting someone to maybe link against it. What's the end goal here?

api.hpp is the public interface to a private layer that might become shared in the future. We are covered in terms of policy risk; cuda-core is protected because that file appears in a private module (_rt) and the support policy reserves the right to change those unconditionally. For now, this is just cuda-core's runtime layer, implemented in C++.

[Andy-Jost]

Since the last round of review, one design change was applied, prompted by the segfault Brandon found, and a few smaller changes came with it.

The design change: VmmRange is now an immutable list of mapping handles that lives in the buffer's own box, a VmmDevicePtrBox that inherits DevicePtrBox with no virtual functions. A grow copies the input's list and builds a new range for its result, so two buffers never share mutable state, and each buffer owns one range and one recorded stream. That removed the stream union in the range deleter, the range mutex, the base-address registry, and vmm_range.hpp.

What came with it:

• modify_allocation always returns a new buffer, so closing the result never closes the input.
• A new test grows and closes aliases of one buffer from four threads, reduced from Brandon's script.
• The release tests ask the driver about the exact addresses a buffer used instead of comparing device-wide free memory, which other processes can move.
• device and config are read-only, and the default-stream check is shared with the stream module.

The PR body and the range section of VMM_DESIGN.md are current. Full suite on an H100 with CUDA 13.4: 4325 passed, 0 failed.

[brandon-b-miller]

I think this may under state what could manifest as somewhat difficult to reason about user behavior. One thing I didn't realize about VMM in general (which I guess makes sense in retrospect) is that the free is not a stream ordered operation, e.g. cuMemunmap doesn't take a stream parameter).

The result is that the caller has to use its own bookkeeping to decide when to issue the real free. This PR does that correctly it appears in virtual_memory.cpp, and blocks:

status = p_cuStreamSynchronize(s);

The potential weirdness here is that depending on how much work is queued, this might happen inside a reference cycle or some other weird part of the users python code that doesn't even relate to CUDA at all. How can I debug why my main thread is waiting for 20 seconds of GPU work at a random point?

This PR intends to move the entire VMM mechanism to c++ and I am +1 on that generally, but I wonder if this particular issue might be ameliorated by moving just the blocking, synching part of the process into some piece of the python layer?

[brandon-b-miller]

Moving this class to a cdef class without adding __weakref__ here actually drops the __weakref__ the original class inherited from the python object class. I think it's probably worth keeping since this is a resource and it looks like most other cuda core resources inherit it or get it elsewhere.

[brandon-b-miller]

This sync let's a deleter raise which violates the contract.

mr = VirtualMemoryResource(dev)
pool_buf = DeviceMemoryResource(dev).allocate(4096, stream=dev.default_stream)
# A: bound to a name -- destructor runs after the error propagates
named = Buffer.from_handle(int(pool_buf.handle), 4096, mr)
mr.modify_allocation(named, 4 << 20)
#   TypeError: buf must be a VirtualMemoryBuffer, got Buffer


# B: passed as a temporary -- destructor runs during unwinding
mr.modify_allocation(Buffer.from_handle(int(pool_buf.handle), 4096, mr), 4 << 20)
#   SystemError: error return without exception set
#   __context__: None

What's happening is modify_allocation raises TypeError, and while it propagates the temporary's deleter in deviceptr_create_with_mr acquires the GIL and calls mr.deallocate without saving the pending exception.

I asked an agent to suggest what to do about this, I think it's solution is pretty clean

Suggested fix. report_message already does the right thing at py_report.cpp:30 ("Deleters can run while a Python exception is propagating; keep it"). Factoring that into an RAII guard, which allocates nothing and is noexcept as the pw_ section requires:

struct PendingExceptionGuard {
    PendingExceptionGuard() noexcept {
#if PY_VERSION_HEX >= 0x030C0000
        pending_ = PyErr_GetRaisedException();
#else
        PyErr_Fetch(&type_, &value_, &tb_);
#endif
    }
    ~PendingExceptionGuard() noexcept {
        PyErr_Clear();  // drop anything the cleanup call itself raised
#if PY_VERSION_HEX >= 0x030C0000
        PyErr_SetRaisedException(pending_);
#else
        PyErr_Restore(type_, value_, tb_);
#endif
    }
    PendingExceptionGuard(const PendingExceptionGuard&) = delete;
    PendingExceptionGuard& operator=(const PendingExceptionGuard&) = delete;
private:
#if PY_VERSION_HEX >= 0x030C0000
    PyObject* pending_ = nullptr;
#else
    PyObject *type_ = nullptr, *value_ = nullptr, *tb_ = nullptr;
#endif
};

makes it one line in the deviceptr_create_with_mr deleter in memory.cpp:

            GILAcquireGuard gil;
            if (gil.acquired()) {
                PendingExceptionGuard keep_exception;
                if (mr_dealloc_cb) {

[brandon-b-miller]

Looking very solid. Couple more Q's