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
Preview removed because the pull request was closed or merged.

[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?

[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]

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]

Looking very solid. Couple more Q's

[sylvesterkaczmarek]

Re-reviewed current head 818a440 after the latest review fixes. The two fresh blockers I was tracking are resolved:

• VirtualMemoryResource restores __weakref__, preserving weak-reference support after the move to a cdef class.
• MR-backed Buffer destruction now wraps Python deallocation with PendingExceptionGuard, and the new regression verifies that a temporary released during exception unwinding preserves the original TypeError while reporting the cleanup failure as CUDAWarning instead of producing SystemError.

The current cross-platform build/test matrix is green. The synchronous VMM-cleanup concern is explicitly tracked separately in #2989 and does not represent a newly introduced correctness regression on this head. No remaining blocker from my review.

[brandon-b-miller]

LGTM.