Allow UMP via managed memory with NumPy allocators by seberg · Pull Request #10232 · cupy/cupy

seberg · 2026-08-27T16:20:03Z

This is towards #8711 although as discussed earlier: it is already possible to "proper" UMP on normal hardware with new linux kernels, it generally "just works".

@leofang if you have some time, this does:

Create new custom NumPy allocators and machinery to set them, cleaning up a few things along the way as well.
Accepting these custom NumPy allocator allocated numpy arrays for zero-copy from NumPy.
A few fixups along the way.

Locally the test seemed passing (but not with -n16 hinting to potential stream order issues), but since I suspect there may be a few things to figure out I'll keep this draft. For one thing, there are some issues around stream-order and others (malloc_managed() might not be using the pool, the owndata logic isn't nice, and maybe most importantly, stream order of a no-copy NumPy arrays deallocation is probably wrong -- we probably need to delay the full deallocate there).

leofang · 2026-08-31T14:12:23Z

+            # If the system supports direct system memory access (hardware or
+            # software) we default to the system allocator.
+            if runtime.deviceGetAttribute(
+                runtime.cudaDevAttrPageableMemoryAccess,


One quick drive-by question: This does not differentiate between ATS (both cudaDevAttrPageableMemoryAccess and cudaDevAttrPageableMemoryAccessUsesHostPageTables are 1) and HMM (only the former is 1). Is it the intent?

I suppose it's safe since one can always explicitly overwrite.

Hmmm, I got HMM not ATS (cudaDevAttrPageableMemoryAccessUsesHostPageTables is false) and things work fine.

But the question here is about defaults and after a (possibly too brief) search around it seems like whatever the system capabilities managed may well be a reasonable default if we assume narrow use with immediate GPU interaction (but it becomes less of an issue with ATS).

So maybe the default should just remain managed or we even just force passing it?

(I am not quite sure about the device check but if it's about default it gets complicated, I guess... Since one could modify the device within this context.)

yeah both HMM and ATS would work, and I agree it seems better to set default to managed

leofang · 2026-08-31T14:13:41Z

+            # software) we default to the system allocator.
+            if runtime.deviceGetAttribute(
+                runtime.cudaDevAttrPageableMemoryAccess,
+                device.get_device_id(),


Another question: Since ATS/HMM support is device-dependent, I am wondering if we should forbid the heterogeneous envs from turning on the special allocator. Seems safer to me to be conservative?

leofang · 2026-09-01T14:30:07Z

+    elif arr.data.mem.size == 0:
+        return


Q: Do we want to special-case for 0-size arrays?

leofang · 2026-09-01T14:32:46Z

+    if allocator == "managed":
+        return True


ditto, if we don't ban heterogeneous systems this technically isn't always True

leofang · 2026-09-03T01:41:43Z

+        # SystemMemory / ManagedMemory and have no such attribute.
+        mem = self.data.mem
+        if isinstance(mem, (memory.SystemMemory, memory.ManagedMemory)):
+            return True


leofang · 2026-09-03T01:44:28Z

+**Managed memory:** The above is also possible on systems without HMM support. In this
+case ``CuPyNumPyAllocator()`` defaults to ``CuPyNumPyAllocator("managed")``, backing the
+NumPy allocations with managed rather than system memory (otherwise the default is
+``CuPyNumPyAllocator("system")``).


ditto, update default

leofang · 2026-09-03T01:46:11Z

+You may also set the environment variable ``CUPY_ENABLE_UMP=1`` to relax checks further.
+CuPy will accept NumPy arrays not strictly backed by its own allocator as well as support
+system memory allocations in ``UnownedMemory``.


I don't follow how CUPY_ENABLE_UMP becomes optional in this PR?

Well, optional in the sense that with this PR the following is a no-copy:

with CuPyNumPyAllocator(): a = np.arange(100) cupy_a = cp.asarray(a)

On the other hand, CUPY_ENABLE_UMP=1 means you don't have to set a NumPy allocator at all.
(This is also assuming that there is no perf regression associated with checking for the allocator use.)

Which also answers the CI question: We may not be able to run CUPY_ENABLE_UMP=1 in a CI run, but we can install the managed memory allocator for NumPy on the global run (although it may slow things down).

leofang

A few additional comments inline.

Separately, nothing exercises the managed allocator path, ManagedMemory.from_external, is_device_accessible, the 0-size peer-access skip, or the reworked is_host_accessible. Given the subtlety here, some smoke coverage before merge would help.

-- Leo's bot

leofang · 2026-09-03T02:04:46Z

+
+    # TODO: Should managed memory also use hint_hugepages?
+    #       (Although if the case, that should be in memory.pyx!)
+    _managed_blocks[mem.ptr] = mem


Only _realloc_managed is marked with gil, with the note that NumPy sometimes calls realloc without the GIL. But _malloc_managed / _calloc_managed / _free_managed also touch Python objects (this dict, memory.malloc_managed, memory.MemoryPointer) and are only noexcept, so Cython assumes the GIL is held. If NumPy ever calls any of them without the GIL for the same reason it does with realloc, it's UB. Either mark all four with gil, or add a comment explaining why realloc is uniquely GIL-less.

leofang · 2026-09-03T02:04:49Z

+        new_ptr = _malloc_managed(ctx, size)
+        if new_ptr == NULL:
+            return NULL
+        memcpy(new_ptr, ptr, old_mem.mem.size)


Same class of stream-order issue you flag in the PR body for no-copy NumPy dealloc: this host-side memcpy on managed memory races with any in-flight kernel still using the old buffer. Probably needs a runtime.streamSynchronize(...) before the copy.

Yeah, this, the memory pool, and new tests are all still needed here. We also may want to bike-shed the name of the allocator once (maybe just cupy...numpy_allocator() even) since we probably want to make it public (and I had not done that yet).

leofang · 2026-09-03T02:04:52Z

+            self._c_contiguous, self._f_contiguous,
+            (self.base is None
+                # TODO(seberg): MemoryPointer should have a flag for "owning"?
+                and getattr(self.data.mem, '_owner', None) is None))


This handles the new managed/system wraps, but the pre-existing UnownedMemory has no _owner field either, so owndata still reports True for it. Your own TODO above notes MemoryPointer needs a real "owning" flag — worth doing here rather than leaving another latent case?

leofang · 2026-09-03T02:04:56Z

+        cpython.PyErr_WriteUnraisable("CuPy malloc_managed")
+        return NULL
+
+    # TODO: Should managed memory also use hint_hugepages?


Two items stacked here that read as blockers rather than TODOs to me:

memory.malloc_managed(size) bypasses the CuPy memory pool — every NumPy allocation hits cudaMallocManaged directly (also called out in the PR body).

No hugepage hint on managed allocations.

Fine to defer if there's a plan for a follow-up, but worth naming as such.

leofang · 2026-09-03T02:04:58Z

+    if allocator == "managed":
+        return True
+    if allocator == "system":
+        return runtime.deviceGetAttribute(


Every cp.asarray(np_arr) on the system path pays a runtime API call here. Worth caching per device (or hoisting into is_ump_supported alongside _is_ump_enabled).

leofang · 2026-09-03T02:05:00Z

+    cpython.PyObject* PyArray_HANDLER(cnp.ndarray arr)
+
+
+# Similar to, but simplfied from, the NumPy alloc.cpp code


nit: simplfied → simplified

leofang · 2026-09-03T02:05:02Z

+
+# Define an allocator for aligned memory (only useful with UMP supported)
+cdef cnp.PyDataMem_Handler _aligned_handler
+# The init.pyd file uses `char *` rather than `char[127]`, making awkward:


nit: .pyd → .pxd (same typo on L202)

leofang · 2026-09-03T02:05:04Z

+    cdef readonly str kind
+
+    def __cinit__(self, kind=None):
+        """Initialize the NumPy allcator.


nit: allcator → allocator

leofang · 2026-09-03T02:05:06Z

 cpdef inline _check_peer_access(_ndarray_base arr, int device_id):
    if arr.data.device_id == device_id:
        return
+    elif arr.data.mem.size == 0:


nit: elif after an unconditional return — plain if reads cleaner.

leofang · 2026-09-03T02:06:22Z

Another thing, non-blocking: We probably should figure out a way to test this in the CI, at least nightly...

mergify · 2026-10-13T06:50:35Z

This pull request is now in conflicts. Could you fix it @seberg? 🙏

seberg added 2 commits August 27, 2026 15:51

Improve numpy aligned allocator and add managed version

a609036

Further fixups (some divergence to main)

985976f

seberg marked this pull request as draft August 27, 2026 16:20

leofang self-assigned this Aug 27, 2026

leofang added cat:enhancement Improvements to existing features prio:medium labels Aug 27, 2026

leofang added this to the v14.3.0 milestone Aug 27, 2026

leofang reviewed Aug 31, 2026

View reviewed changes

leofang reviewed Sep 3, 2026

View reviewed changes

leofang added prio:high and removed prio:medium labels Sep 3, 2026

seberg mentioned this pull request Sep 16, 2026

ENH: Help with custom alignments for (structured) dtypes numpy/numpy#28318

Open

		cpython.PyObject* PyArray_HANDLER(cnp.ndarray arr)


		# Similar to, but simplfied from, the NumPy alloc.cpp code

Conversation

seberg commented Aug 27, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

leofang left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

leofang commented Sep 3, 2026

Uh oh!

mergify Bot commented Oct 13, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants