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 class symmetric to Device. Host() (any host), Host(numa_id=N), Host.numa_current(). 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).
• *Options dataclasses — AdviseOptions, PrefetchOptions, DiscardOptions, DiscardPrefetchOptions. Frozen dataclasses reserved for future per-call flags; current ABI has no flags worth surfacing, but the dataclasses establish the contract so future flags land without an API break.

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: 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.

Free functions — advise / prefetch / discard / discard_prefetch

Each accepts a Buffer (or ManagedBuffer) or a sequence of them. Locations are expressed via Device, Host, or int (-1 → host, >=0 → device ordinal).

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)

# int shorthand: -1 = host, >=0 = device ordinal
prefetch(buf, -1, stream=stream)

# 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.
• 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 / int; _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.
• The defensive cuInit retry in _query_memory_attrs was removed; we don't auto-init CUDA elsewhere.

Tests

Managed-memory tests live in cuda_core/tests/memory/test_managed_ops.py: free-function dispatch (single + batched + mismatch + non-managed rejection), Host constructors and frozen-dataclass semantics, internal _coerce_location for Device | Host | int | None, full ManagedBuffer property roundtrips (read_mostly, preferred_location, accessed_by add/discard/assignment), and instance methods. The broader memory-tests reorg (buffer / managed_resource / pinned / vmm "siblings") is tracked as a separate cleanup PR.

Deferred follow-ups

• HMM/ATS-aware is_managed semantics — flagged as a TODO in _buffer.pyx, tracked alongside the broader HMM/ATS work.
• cuMemcpyBatchAsync (P0 of Support batched memory movement #1333) — different family, separate PR; will mirror the contract in #issuecomment-4355502334.
• Concrete fields on the *Options dataclasses — they're empty today; concrete options land when CUDA introduces per-call flags worth surfacing.
• CUDA 13 split-attribute read-back for preferred_location / accessed_by — currently uses the legacy combined attribute (Python binding limitation), which loses NUMA fidelity on round-trip. Setters preserve full NUMA info.

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

I need to run. Will try to revisit tonight. I haven't done reviewing (too many lines).

[leofang]

I dunno if this is a collective agentic illusion or what, in recent PRs I've seen many data classes. Why do we need one here?

[rparolin]

When I asked claude why he selected a dataclass... I'll remove that decorator.

"A. Reply only — defend the dataclass

The dataclass is doing real work, not decoration:

• eq/hash are tested behaviors — tests/memory/test_managed_ops.py:323-327 asserts Host() == Host(), hash(Host(numa_id=1)) ==
  hash(Host(numa_id=1)), etc.
• Host is used in set/list comparisons in _managed_buffer.py:49 ([Host() if v == -1 else Device(v) for v in raw...]) — needs hashability if
  it ever lands in a set.
• frozen=True ensures users can't mutate a Host after stashing it on ManagedBuffer.preferred_location."

[leofang]

My bot reviewed and raised this idea: Host should follow Device and be a singleton class. @Andy-Jost thoughts?

[Andy-Jost]

None of my comments are blocking. Looks ready to me.

[leofang]

Code-level review focusing on DRY, verbosity, and test coverage — complementing the design-level comments already on the PR.

Summary of inline comments:

1. _do_batch_prefetch / _do_batch_discard_prefetch are copy-pasted (~40 lines) — parameterize into one function
2. Options isinstance check repeated 4x — extract a one-liner helper; also consolidate the near-identical prefetch / discard_prefetch preambles
3. CUDA 12 batch fallback — loop over singles instead of NotImplementedError (the batch semantics are documented as equivalent to individual calls)
4. _normalize_managed_advice over-engineered — the alias dict + lazy reverse dict can be replaced with a getattr on the naming convention (~15 lines saved)
5. Test setup boilerplate — ~25 tests repeat the same 5-line preamble; a pytest fixture would save ~75 lines
6. Test helper duplicated — _get_mem_range_attr in the test is identical to _get_int_attr in production code
7. Test coverage gaps — CUDA 12 batch fallback, AccessedBySet iteration, stream=None, error message assertions

[leofang]

nit: Nearly every test function in this file repeats the same 5-line preamble:

device = Device()
_skip_if_managed_location_ops_unsupported(device)  # or variant
device.set_current()
mr = create_managed_memory_resource_or_skip()
buffer = mr.allocate(_MANAGED_TEST_ALLOCATION_SIZE)

This appears in ~25 tests. A @pytest.fixture would eliminate this boilerplate and save ~75 lines:

@pytest.fixture
def managed_buffer(init_cuda):
    device = Device()
    _skip_if_managed_location_ops_unsupported(device)
    device.set_current()
    mr = create_managed_memory_resource_or_skip()
    buf = mr.allocate(_MANAGED_TEST_ALLOCATION_SIZE)
    yield buf
    buf.close()
    mr.close()

Tests that need a different skip level (e.g. _skip_if_managed_discard_prefetch_unsupported) could use a second fixture or parametrize. The local _skip_if_* helpers also partially overlap with conftest's skip_if_managed_memory_unsupported — worth consolidating.

[rparolin]

@leofang addressed all your feedback.

[rparolin]

Code review

Found 2 issues:

1. cuda_core/docs/source/api.rst documents 8 names that don't exist anywhere in the codebase: free functions advise, prefetch, discard, discard_prefetch and dataclasses AdviseOptions, PrefetchOptions, DiscardOptions, DiscardPrefetchOptions. The actual exports from cuda.core.utils are only prefetch_batch, discard_batch, discard_prefetch_batch (plus the pre-existing StridedMemoryView and args_viewable_as_strided_memory). Sphinx will fail to resolve these autosummary entries. Looks like leftover from the earlier API surface that was removed under R9/R11; release notes were updated, api.rst was missed.

https://github.com/NVIDIA/cuda-python/blob/b0d1a216e3932468d3801da5d83449465b3f8faf/cuda_core/docs/source/api.rst#L248-L268

2. The CUDA 12 ELSE fallbacks for the single-buffer advise and prefetch paths in _managed_memory_ops.pyx reference cydriver.cuMemAdvise and cydriver.cuMemPrefetchAsync with the legacy v1 4-argument int-device signature, but those public cydriver wrappers are only emitted when cuMemAdvise_v2 / cuMemPrefetchAsync_v2 are present in the toolkit headers (CUDA 13+). On a cu12 build the symbols don't exist (see cuda_bindings/cuda/bindings/cydriver.pxd.in:3840 and cydriver.pyx.in:1237-1247, both gated on 'cuMem*_v2' in found_functions), so the ELSE branch will fail to compile / import. The batched paths handle cu12 via a NotImplementedError + per-range loop, but the single-buffer paths assume a v1 wrapper that cuda-python never exposes.

cuda-python/cuda_core/cuda/core/_memory/_managed_memory_ops.pyx

Lines 255 to 325 in b0d1a21

	cu_loc.type = cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST
	cu_loc.id = 0
	else:
	cu_loc = _to_cumemlocation(loc)
	with nogil:
	HANDLE_RETURN(cydriver.cuMemAdvise(cu_ptr, nbytes, advice_enum, cu_loc))
	ELSE:
	cdef int dev_int = -1 if loc is None else _to_legacy_device(loc)
	with nogil:
	HANDLE_RETURN(cydriver.cuMemAdvise(cu_ptr, nbytes, advice_enum, dev_int))
	def prefetch_batch(buffers, locations, *, stream):
	"""Prefetch a batch of managed-memory ranges to target locations.
	Requires CUDA 13+. For a single buffer, use
	:meth:`ManagedBuffer.prefetch` instead.
	Parameters
	----------
	buffers : Sequence[:class:`Buffer`]
	Two or more managed allocations to operate on.
	locations : :class:`~cuda.core.Device` | :class:`~cuda.core.Host` | Sequence[...]
	Target location(s). A single location applies to all buffers; a
	sequence must match ``len(buffers)``.
	stream : :class:`~_stream.Stream` | :class:`~graph.GraphBuilder`
	Stream for the asynchronous prefetch (keyword-only).
	Notes
	-----
	On a CUDA 12 build, falls back to a Python-level loop calling
	``cuMemPrefetchAsync`` per buffer (no batched driver entry point on
	CUDA 12). CUDA 13 builds use ``cuMemPrefetchBatchAsync`` directly.
	"""
	cdef tuple bufs = _coerce_batch_buffers(buffers, "prefetch_batch")
	cdef Py_ssize_t n = len(bufs)
	cdef tuple locs = _broadcast_locations(locations, n, False, "prefetch_batch")
	cdef Stream s = Stream_accept(stream)
	cdef Buffer buf
	for buf in bufs:
	_require_managed_buffer(buf, "prefetch_batch")
	_do_batch_prefetch(bufs, locs, s)
	def _do_single_prefetch_py(Buffer buf, location, stream):
	"""Internal: single-buffer prefetch for ManagedBuffer.prefetch().
	Uses cuMemPrefetchAsync (works on CUDA 12 and 13).
	"""
	_require_managed_buffer(buf, "prefetch")
	cdef object loc = _coerce_location(location, allow_none=False)
	cdef Stream s = Stream_accept(stream)
	_do_single_prefetch(buf, loc, s)
	cdef void _do_single_prefetch(Buffer buf, object loc, Stream s):
	cdef cydriver.CUdeviceptr cu_ptr = as_cu(buf._h_ptr)
	cdef size_t nbytes = buf._size
	cdef cydriver.CUstream hstream = as_cu(s._h_stream)
	IF CUDA_CORE_BUILD_MAJOR >= 13:
	cdef cydriver.CUmemLocation cu_loc = _to_cumemlocation(loc)
	with nogil:
	HANDLE_RETURN(cydriver.cuMemPrefetchAsync(cu_ptr, nbytes, cu_loc, 0, hstream))
	ELSE:
	cdef int dev_int = _to_legacy_device(loc)
	with nogil:
	HANDLE_RETURN(cydriver.cuMemPrefetchAsync(cu_ptr, nbytes, dev_int, hstream))

🤖 Generated with Claude Code

_{- If this code review was useful, please react with 👍. Otherwise, react with 👎.}