Add managed-memory advise, prefetch, and discard-prefetch free functions

[rparolin]

Summary

Adds managed-memory range operations to cuda.core:

• Free functions in cuda.core.utils: advise, prefetch, discard, discard_prefetch. Each accepts either a single Buffer or a sequence; N==1 dispatches to the per-range driver entry point and N>1 dispatches to the corresponding cuMem*BatchAsync (CUDA 13+).
• Host — new top-level singleton class symmetric to Device. Host() (any host), Host(numa_id=N), Host.numa_current(). Same-argument constructions are interned (Host() is Host()). Used together with Device to express managed-memory locations.
• ManagedBuffer — Buffer subclass returned by ManagedMemoryResource.allocate. Exposes a Pythonic property-style advice API on top of the same free functions. Wrap an external managed pointer with Buffer.from_handle(...) (now a @classmethod, so ManagedBuffer.from_handle(...) returns a ManagedBuffer).

Closes #1332. Addresses the managed-memory portion of #1333 (P1: cuMemPrefetchBatchAsync, cuMemDiscardBatchAsync, cuMemDiscardAndPrefetchBatchAsync). The P0 cuMemcpyBatchAsync from #1333 is intentionally out of scope and tracked separately; the holistic batched-API contract this PR commits to is documented in #issuecomment-4355502334 so the upcoming cuMemcpyBatchAsync work can mirror it.

Public API

ManagedBuffer — property-style advice on managed allocations

ManagedMemoryResource.allocate returns a ManagedBuffer (a Buffer subclass). All ManagedBuffer-specific behavior is layered on top of the free functions, so the two surfaces stay consistent.

from cuda.core import Device, Host, ManagedMemoryResource

mr = ManagedMemoryResource()
buf = mr.allocate(size)                # ManagedBuffer

# Driver-backed properties — getter queries the driver, setter calls cuMemAdvise.
buf.read_mostly = True
buf.preferred_location = Device(0)     # or Host(), or Host(numa_id=N)
buf.preferred_location = None          # unset

# Live set-like view of `set_accessed_by` advice.
buf.accessed_by.add(Device(1))
buf.accessed_by.discard(Device(1))
buf.accessed_by = {Device(0), Device(1)}   # diff vs current; advise only deltas

# Instance methods delegate to the matching free functions.
buf.prefetch(Device(0), stream=stream)
buf.discard(stream=stream)
buf.discard_prefetch(Device(0), stream=stream)

Note: on CUDA 13 builds, preferred_location round-trips full NUMA detail via the v2 attribute (Host(numa_id=N) and Host.numa_current() are preserved on read-back). On CUDA 12 builds, the legacy cuMemRangeGetAttribute query path returns integer device ordinals, so Host(numa_id=...) collapses to a generic Host() on read-back. Setters preserve full NUMA information when issuing advice on both toolkits.

Free functions — advise / prefetch / discard / discard_prefetch

Each accepts a Buffer (or ManagedBuffer) or a sequence of them. Locations are expressed via Device or Host.

from cuda.core import Device, Host
from cuda.core.utils import advise, prefetch, discard, discard_prefetch

# Stage to GPU, kernel, bring back to host
prefetch(buf, Device(0), stream=stream)
launch_my_kernel(buf, stream=stream)
prefetch(buf, Host(), stream=stream)
stream.sync()
result = bytes(buf)

# Advice
advise(weights, "set_read_mostly")
advise(activations, "set_preferred_location", Device(0))
advise(scratch, "set_accessed_by", Device(0))

# Discard / discard+prefetch (CUDA 13+)
discard(scratch, stream=stream)
for step in range(num_steps):
    discard_prefetch(activations, Device(0), stream=stream)
    run_forward(activations, stream=stream)

Batched form — same function, sequence of targets

When N>1, dispatch goes to the corresponding cuMem*BatchAsync. Sequence locations are paired by index; a scalar location broadcasts to every target.

# Pair-by-index: output → GPU 0, log_metrics → host
prefetch(
    [output, log_metrics],
    [Device(0), Host()],
    stream=stream,
)

# Scalar broadcast: every shard moves to GPU 0
prefetch([shard_a, shard_b, shard_c], Device(0), stream=stream)

Mismatched sequence lengths raise ValueError. On a CUDA 12 build of cuda.core, N>1 raises NotImplementedError (the *BatchAsync entry points are CUDA 13+); N==1 works on every supported toolkit.

Putting it together

weights = mr.allocate(weights_size)    # ManagedBuffer
inputs  = mr.allocate(inputs_size)
output  = mr.allocate(output_size)

# One-time hints (property API on ManagedBuffer)
weights.read_mostly = True
weights.preferred_location = Device(0)
output.preferred_location = Device(0)

# Per inference
inputs.prefetch(Device(0), stream=stream)
run_inference(weights, inputs, output, stream=stream)
output.prefetch(Host(), stream=stream)
stream.sync()

Implementation notes

• Cython implementation in cuda_core/cuda/core/_memory/_managed_memory_ops.pyx uses cimport cydriver for direct C-level driver calls.
• The CUDA 12 / 13 ABI split for cuMemAdvise and cuMemPrefetchAsync is handled at compile time with IF CUDA_CORE_BUILD_MAJOR >= 13: / ELSE:.
• Batched entry points (cuMemPrefetchBatchAsync, cuMemDiscardBatchAsync, cuMemDiscardAndPrefetchBatchAsync) are CUDA 13+ only. On CUDA 12 builds, N>1 calls raise NotImplementedError; single-buffer calls work everywhere.
• Host is a singleton with __slots__ and a __new__-based intern cache keyed by (numa_id, is_numa_current). Same-argument constructions return the same instance on both Python and Cython call paths.
• ManagedBuffer is a pure-Python subclass of the Cython Buffer cdef class. Buffer.from_handle is now a @classmethod (was @staticmethod) so MyBufferSubclass.from_handle(...) returns the typed instance via cls._init. Buffer_from_deviceptr_handle and _MP_allocate thread an optional cls parameter so ManagedMemoryResource.allocate materializes a ManagedBuffer.
• Internal _LocSpec (in _managed_location.py) carries the (kind, id) discriminator that the Cython layer maps to CUmemLocation (CUDA 13) or a legacy device ordinal (CUDA 12). Public callers see only Device / Host; _coerce_location produces the internal record.
• _buffer.pyx collapses out.is_managed = (is_managed != 0) to a single unconditional assignment and adds a TODO noting that HMM/ATS-mapped sysmem is not yet captured by CU_POINTER_ATTRIBUTE_IS_MANAGED.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[github-actions]

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

[rparolin]

/ok to test

[jrhemstad]

question: Does making these member functions of the Buffer type preclude this functionality for allocations that weren't created through the Buffer type? Did we consider making these free functions instead of member functions on the Buffer type?

[rparolin]

question: Does making these member functions of the Buffer type preclude this functionality for allocations that weren't created through the Buffer type? Did we consider making these free functions instead of member functions on the Buffer type?

I'm moving this back into draft. We discussed in our team meeting because I was already hesitant as Buffer is becoming a 'God object' with the functionality is gaining. We were going to explore alternatives. Free functions sounds like a good alternative to explore.

[leofang]

Approved, assuming all pending comments below will be either tracked or followed up.

[leofang]

This block (_MANAGED_ADVICE_ALIASES dict, _ADVICE_ENUM_TO_ALIAS reverse dict, and _normalize_managed_advice function, ~60 lines) appears to be dead code in the current design.

Since advise is now exclusively exposed through ManagedBuffer property setters (buf.read_mostly = True, buf.preferred_location = Device(0), buf.accessed_by.add(...)), all calls to _advise_one pass the CUmem_advise enum value directly. The string-alias lookup path (isinstance(advice, str) → dict lookup → attribute resolution) is never exercised by any public API.

If string advice names aren't part of the API contract going forward, this code should be simplified to only handle the enum path — or removed entirely.

[leofang]

This test should be moved to its own file (like Device)

[leofang]

The fact that these are needed at test time rings a bell. cuda.core tries hard to not leak the abstraction. This highlights a problem that we do not expose enough mem-range attributes for ManagedBuffer. We should have a follow-up PR on this.

[leofang]

These 8 standalone test functions are largely redundant with TestManagedBuffer methods — they test the same operations (prefetch, read_mostly, preferred_location, discard_prefetch) but with a different buffer source (pool-allocated vs external DummyUnifiedMemoryResource + from_handle).

The two axes being tested are operation × buffer source. This is a textbook @pytest.fixture(params=...) case. One parametrized fixture:

@pytest.fixture(params=["pool", "external"], ids=["pool", "external"])
def managed_buffer(request, location_ops_device):
    if request.param == "pool":
        mr = create_managed_memory_resource_or_skip()
        buf = mr.allocate(_MANAGED_TEST_ALLOCATION_SIZE, stream=location_ops_device.default_stream)
        yield buf
        buf.close()
        mr.close()
    else:
        plain = DummyUnifiedMemoryResource(location_ops_device).allocate(_MANAGED_TEST_ALLOCATION_SIZE)
        buf = ManagedBuffer.from_handle(plain.handle, plain.size, owner=plain)
        yield buf
        plain.close()

Then each TestManagedBuffer test runs twice — pool and external — and these 8 standalone functions (~150 lines) can be deleted entirely. conftest already uses this factory pattern with memory_resource_factory (parametrized over DeviceMR/PinnedMR/ManagedMR).

[leofang]

9 fixtures is a lot for what amounts to 3 skip-tiers × {device, mr, buffer}. The three *_mr fixtures are byte-for-byte identical (create_managed_memory_resource_or_skip()), and the three *_device fixtures differ only in which skip function they call.

Suggestion: keep only location_ops_device + location_ops_mr (the most common tier) and a managed_buffer fixture parametrized over buffer source (see comment on standalone tests). The 2–3 tests needing memory_pool_* or discard_prefetch_* can call their specific skip function inline — they already know their requirement. This would reduce 9 fixtures (~50 lines) to 3–4 (~20 lines).

Also, conftest already has an init_cuda fixture that calls Device(0).set_current(). The *_device fixtures duplicate that work (they call Device(); device.set_current() inside). They could depend on init_cuda and just add the skip check.

[leofang]

nit: _get_int_attr(buf, driver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION) is repeated 8 times verbatim in this file. A one-liner helper would improve readability:

def _last_prefetch_location(buf):
    return _get_int_attr(buf, driver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION)

Note: this concern becomes moot once ManagedBuffer exposes mem-range attributes directly (per the pending comment on line 17 about leaking the cuda.bindings.driver abstraction). Once buf.last_prefetch_location or similar exists, tests would use that instead of reaching into driver internals.

[leofang]

test_rejects_single_buffer is copy-pasted across all three batch test classes (TestPrefetchBatch, TestDiscardBatch, TestDiscardPrefetchBatch) — identical except for the function name. test_length_mismatch is also duplicated across two. These can be parametrized into standalone tests:

@pytest.mark.parametrize("fn_name,needs_loc", [
    ("prefetch_batch", True),
    ("discard_batch", False),
    ("discard_prefetch_batch", True),
])
def test_batch_rejects_single_buffer(memory_pool_device, memory_pool_mr, fn_name, needs_loc):
    from cuda.core import utils
    fn = getattr(utils, fn_name)
    buf = memory_pool_mr.allocate(_MANAGED_TEST_ALLOCATION_SIZE, stream=memory_pool_device.default_stream)
    stream = memory_pool_device.create_stream()
    args = (stream, buf, Host()) if needs_loc else (stream, buf)
    with pytest.raises(TypeError, match="sequence of Buffers"):
        fn(*args)
    buf.close()

With this + parametrization of test_length_mismatch, the three batch test classes shrink significantly — or could be merged into one TestBatchOps class.

[leofang]

The local _skip_if_managed_allocation_unsupported vs conftest's skip_if_managed_memory_unsupported have confusingly similar names but check different things:

• conftest: "can I create a ManagedMemoryResource?" (pool infra + runtime check)
• local: "does the device support cuMemAllocManaged?" (raw managed alloc capability)

The local version is needed because DummyUnifiedMemoryResource uses cuMemAllocManaged directly, not ManagedMemoryResource. But the naming makes it look like a copy-paste bug.

Suggestions:

1. Rename to _skip_if_cuMemAllocManaged_unsupported to make the distinction obvious.
2. Better: inline this check into DummyUnifiedMemoryResource itself (it already knows it needs managed memory — it could skip/raise at allocation time), eliminating the need for callers to remember which skip function to use.
3. _skip_if_managed_allocation_unsupported is only ever called by _skip_if_managed_location_ops_unsupported — could be inlined there (it's 4 lines).