cuda.core: expose last prefetch location on ManagedBuffer

cuda_core/cuda/core/_host.py

@@ -23,7 +23,8 @@ class Host:
     ``Host`` is the symmetric counterpart of :class:`~cuda.core.Device`
     for managed-memory `prefetch`, `advise`, and `discard_prefetch`
     targets. Pass either a ``Device`` or a ``Host`` to those operations
-    and to ``ManagedBuffer.preferred_location`` / ``accessed_by``.
+    and to ``ManagedBuffer.preferred_location`` / ``accessed_by``. A
+    ``Host`` may also be returned by ``ManagedBuffer.last_prefetch_location``.
 
     ``Host`` is a singleton class, mirroring :class:`~cuda.core.Device`:
     constructor calls with the same arguments return the same instance,

cuda_core/cuda/core/_memory/_managed_buffer.py

@@ -15,6 +15,7 @@
     _do_single_discard_prefetch_py,
     _do_single_discard_py,
     _do_single_prefetch_py,
+    _read_last_prefetch_location_v2,
     _read_preferred_location_v2,
 )
 from cuda.core._utils.cuda_utils import driver, handle_return
@@ -43,6 +44,7 @@
 _ATTR_READ_MOSTLY = _RANGE.CU_MEM_RANGE_ATTRIBUTE_READ_MOSTLY
 _ATTR_PREFERRED = _RANGE.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION
 _ATTR_ACCESSED_BY = _RANGE.CU_MEM_RANGE_ATTRIBUTE_ACCESSED_BY
+_ATTR_LAST_PREFETCH = _RANGE.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION
 
 
 def _check_open(buf: Buffer) -> None:
@@ -147,11 +149,12 @@ class ManagedBuffer(Buffer):
 
     Note
     ----
-    On CUDA 13 builds, ``preferred_location`` round-trips full NUMA
-    information. On CUDA 12 builds, ``Host(numa_id=...)`` and
-    ``Host.numa_current()`` are rejected with ``TypeError`` at the call
-    boundary — only ``Device(...)`` and the generic ``Host()`` are
-    accepted. Use ``Host()`` to target the host on CUDA 12.
+    On CUDA 13 builds, ``preferred_location`` and
+    ``last_prefetch_location`` round-trip full NUMA information. On CUDA 12
+    builds, ``Host(numa_id=...)`` and ``Host.numa_current()`` are rejected
+    with ``TypeError`` at the call boundary — only ``Device(...)`` and the
+    generic ``Host()`` are accepted. Use ``Host()`` to target the host on
+    CUDA 12.
     """
 
     @classmethod
@@ -168,7 +171,8 @@ def from_handle(
 
         Use this when you have an externally-allocated managed pointer
         and want the property-style advice API (:attr:`read_mostly`,
-        :attr:`preferred_location`, :attr:`accessed_by`).
+        :attr:`preferred_location`, :attr:`last_prefetch_location`,
+        :attr:`accessed_by`).
 
         Parameters
         ----------
@@ -232,6 +236,28 @@ def preferred_location(self, value: Device | Host | None) -> None:
         else:
             _advise_one(self, _SET_PREFERRED, value)
 
+    @property
+    def last_prefetch_location(self) -> Device | Host | None:
+        """Location targeted by the most recent explicit prefetch.
+
+        Returns ``None`` if any page in the range has never been prefetched,
+        or if the pages do not share one last-prefetch location. This reports
+        the application's requested destination, not current residency or
+        completion of the asynchronous prefetch operation.
+
+        On CUDA 13 builds, fully round-trips ``Host(numa_id=N)``. On CUDA 12,
+        the legacy attribute carries only a device ordinal (or ``-1`` for
+        host), so host NUMA details are unavailable.
+        """
+        if binding_version() >= (13, 0, 0) and driver_version() >= (13, 0, 0):
+            return _read_last_prefetch_location_v2(self)
+        loc_id = _get_int_attr(self, _ATTR_LAST_PREFETCH)
+        if loc_id == -2:
+            return None
+        if loc_id == -1:
+            return Host()
+        return Device(loc_id)
+
     @property
     def accessed_by(self) -> AccessedBySetProxy:
         """Live set-like view of ``set_accessed_by`` locations."""

cuda_core/cuda/core/_memory/_managed_memory_ops.pyi

@@ -69,6 +69,7 @@ def _do_single_prefetch_py(buf: Buffer, location: Device | Host | None, stream:
     Uses cuMemPrefetchAsync (works on CUDA 12 and 13).
     """
 def _read_preferred_location_v2(buf: Buffer) -> Device | Host | None: ...
+def _read_last_prefetch_location_v2(buf: Buffer) -> Device | Host | None: ...
 def discard_prefetch_batch(stream: Stream | GraphBuilder, buffers: Sequence[Buffer], locations: Device | Host | Sequence[Device | Host]) -> None:
     """Discard a batch of managed-memory ranges and prefetch them to target locations.
 

cuda_core/cuda/core/_memory/_managed_memory_ops.pyx

@@ -281,15 +281,13 @@ IF CUDA_CORE_BUILD_MAJOR >= 13:
     ) except ?cydriver.CUDA_ERROR_NOT_FOUND nogil
 
 
-    def _read_preferred_location_v2(Buffer buf) -> Device | Host | None:
-        """Internal: read preferred_location with full NUMA detail.
-
-        Bypasses cuda.bindings.driver.cuMemRangeGetAttribute (whose
-        attribute allowlist doesn't yet include the cu13 _TYPE / _ID
-        attributes) by calling cydriver directly.
-
-        Returns Device | Host | None.
-        """
+    cdef object _read_location_v2(
+        Buffer buf,
+        cydriver.CUmem_range_attribute type_attribute,
+        cydriver.CUmem_range_attribute id_attribute,
+    ):
+        # cuda.bindings.driver.cuMemRangeGetAttribute does not yet accept the
+        # CUDA 13 _TYPE / _ID attributes, so query them through cydriver.
         Buffer_check_open(buf)
         cdef cydriver.CUdeviceptr cu_ptr = as_cu(buf._h_ptr)
         cdef size_t nbytes = buf._size
@@ -298,12 +296,12 @@ IF CUDA_CORE_BUILD_MAJOR >= 13:
         with nogil:
             HANDLE_RETURN(cydriver.cuMemRangeGetAttribute(
                 <void*>&loc_type, sizeof(int),
-                cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_TYPE,
+                type_attribute,
                 cu_ptr, nbytes,
             ))
             HANDLE_RETURN(cydriver.cuMemRangeGetAttribute(
                 <void*>&loc_id, sizeof(int),
-                cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_ID,
+                id_attribute,
                 cu_ptr, nbytes,
             ))
         if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE:
@@ -315,7 +313,25 @@ IF CUDA_CORE_BUILD_MAJOR >= 13:
             return Host(numa_id=loc_id)
         if loc_type == <int>cydriver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT:
             return Host.numa_current()
-        return None  # CU_MEM_LOCATION_TYPE_INVALID — no preferred location
+        return None  # CU_MEM_LOCATION_TYPE_INVALID
+
+
+    def _read_preferred_location_v2(Buffer buf) -> Device | Host | None:
+        """Internal: read preferred_location with full NUMA detail."""
+        return _read_location_v2(
+            buf,
+            cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_TYPE,
+            cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION_ID,
+        )
+
+
+    def _read_last_prefetch_location_v2(Buffer buf) -> Device | Host | None:
+        """Internal: read last_prefetch_location with full NUMA detail."""
+        return _read_location_v2(
+            buf,
+            cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION_TYPE,
+            cydriver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION_ID,
+        )
 
 
     cdef void _do_batch_prefetch_op(tuple bufs, tuple locs, Stream s, _BatchPrefetchFn fn):
@@ -348,16 +364,21 @@ IF CUDA_CORE_BUILD_MAJOR >= 13:
             ))
 ELSE:
     def _read_preferred_location_v2(Buffer buf) -> Device | Host | None:
-        # Symbol exists so _managed_buffer.py can `from ... import
-        # _read_preferred_location_v2` unconditionally at module top.
-        # `ManagedBuffer.preferred_location` gates on both
-        # binding_version() and driver_version() >= (13, 0, 0) before
-        # calling, so this path is unreachable on a cu12 build.
+        # Symbols exist so _managed_buffer.py can import the v2 readers
+        # unconditionally. Their properties gate on both binding_version()
+        # and driver_version() >= (13, 0, 0), so these paths are unreachable
+        # on a CUDA 12 build.
         raise NotImplementedError(
             "_read_preferred_location_v2 requires a CUDA 13 build of cuda.core"
         )
 
 
+    def _read_last_prefetch_location_v2(Buffer buf) -> Device | Host | None:
+        raise NotImplementedError(
+            "_read_last_prefetch_location_v2 requires a CUDA 13 build of cuda.core"
+        )
+
+
 cdef void _do_batch_prefetch(tuple bufs, tuple locs, Stream s):
     IF CUDA_CORE_BUILD_MAJOR >= 13:
         _do_batch_prefetch_op(bufs, locs, s, cydriver.cuMemPrefetchBatchAsync)

cuda_core/cuda/core/_memory/_managed_memory_resource.pyi

@@ -93,9 +93,9 @@ class ManagedMemoryResource(_MemPool):
         ManagedBuffer
             A :class:`ManagedBuffer` (a :class:`Buffer` subclass) that
             exposes the property-style advice API
-            (``read_mostly``, ``preferred_location``, ``accessed_by``)
-            and instance methods (``prefetch``, ``discard``,
-            ``discard_prefetch``).
+            (``read_mostly``, ``preferred_location``,
+            ``last_prefetch_location``, ``accessed_by``) and instance methods
+            (``prefetch``, ``discard``, ``discard_prefetch``).
         """
     @property
     def device_id(self) -> int:

cuda_core/cuda/core/_memory/_managed_memory_resource.pyx

@@ -119,9 +119,9 @@ cdef class ManagedMemoryResource(_MemPool):
         ManagedBuffer
             A :class:`ManagedBuffer` (a :class:`Buffer` subclass) that
             exposes the property-style advice API
-            (``read_mostly``, ``preferred_location``, ``accessed_by``)
-            and instance methods (``prefetch``, ``discard``,
-            ``discard_prefetch``).
+            (``read_mostly``, ``preferred_location``,
+            ``last_prefetch_location``, ``accessed_by``) and instance methods
+            (``prefetch``, ``discard``, ``discard_prefetch``).
         """
         MP_check_open(self)
         assert isinstance(stream, Stream), "Only Stream is supported for managed memory allocations"

cuda_core/tests/memory/test_managed_ops.py

@@ -9,7 +9,6 @@
 
 from cuda.bindings import driver
 from cuda.core import Device, Host, ManagedBuffer
-from cuda.core._memory._managed_buffer import _get_int_attr
 
 # Managed-memory prefetch and CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION
 # operate at physical-page granularity. Test buffers must each occupy a full
@@ -20,16 +19,9 @@
 _PAGE_SIZE = mmap.PAGESIZE
 _MANAGED_TEST_ALLOCATION_SIZE = _PAGE_SIZE
 _READ_MOSTLY_ENABLED = 1
-_HOST_LOCATION_ID = -1
 _INVALID_HOST_DEVICE_ORDINAL = 0
 
 
-# TODO(#2109): replace with ``buf.last_prefetch_location`` once
-# ``ManagedBuffer`` exposes mem-range attributes directly.
-def _last_prefetch_location(buf):
-    return _get_int_attr(buf, driver.CUmem_range_attribute.CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION)
-
-
 def _page_base(buf):
     # Page-aligned base of the buffer's start address; two buffers sharing a
     # page cannot be prefetched to different locations independently.
@@ -224,8 +216,7 @@ def test_same_location(self, location_ops_device, location_ops_mr):
         stream.sync()
 
         for buf in bufs:
-            last = _last_prefetch_location(buf)
-            assert last == device.device_id
+            assert buf.last_prefetch_location == device
             buf.close()
 
     def test_per_buffer_location(self, location_ops_device, location_ops_mr):
@@ -242,10 +233,8 @@ def test_per_buffer_location(self, location_ops_device, location_ops_mr):
         prefetch_batch(stream, bufs, [Host(), device])
         stream.sync()
 
-        last0 = _last_prefetch_location(bufs[0])
-        last1 = _last_prefetch_location(bufs[1])
-        assert last0 == _HOST_LOCATION_ID
-        assert last1 == device.device_id
+        assert bufs[0].last_prefetch_location == Host()
+        assert bufs[1].last_prefetch_location == device
         for buf in bufs:
             buf.close()
 
@@ -285,8 +274,7 @@ def test_same_location(self, location_ops_device, location_ops_mr):
         discard_prefetch_batch(stream, bufs, device)
         stream.sync()
         for buf in bufs:
-            last = _last_prefetch_location(buf)
-            assert last == device.device_id
+            assert buf.last_prefetch_location == device
             buf.close()
 
 
@@ -366,6 +354,23 @@ def test_from_handle(self, init_cuda):
         finally:
             plain.close()
 
+    @pytest.mark.agent_authored(model="gpt-5")
+    def test_last_prefetch_location_initially_none(self, external_managed_buffer):
+        assert external_managed_buffer.last_prefetch_location is None
+
+    @pytest.mark.agent_authored(model="gpt-5")
+    def test_last_prefetch_location_roundtrip_host_numa(self, location_ops_device, managed_buffer):
+        from cuda.core._utils.version import binding_version, driver_version
+
+        if binding_version() < (13, 0, 0) or driver_version() < (13, 0, 0):
+            pytest.skip("Host NUMA last-prefetch location requires CUDA 13")
+
+        stream = location_ops_device.create_stream()
+        location = Host(numa_id=0)
+        managed_buffer.prefetch(location, stream=stream)
+        stream.sync()
+        assert managed_buffer.last_prefetch_location == location
+
     @pytest.mark.thread_unsafe(reason="external_managed_buffer is shared between threads")
     def test_read_mostly_roundtrip(self, external_managed_buffer):
         buf = external_managed_buffer
@@ -478,7 +483,7 @@ def test_instance_prefetch(self, location_ops_device, managed_buffer):
         stream = device.create_stream()
         buf.prefetch(device, stream=stream)
         stream.sync()
-        assert _last_prefetch_location(buf) == device.device_id
+        assert buf.last_prefetch_location == device
 
     def test_instance_discard(self, location_ops_device, managed_buffer):
         if not hasattr(driver, "cuMemDiscardBatchAsync"):
@@ -501,7 +506,7 @@ def test_instance_discard_prefetch(self, discard_prefetch_device):
             stream.sync()
             buf.discard_prefetch(device, stream=stream)
             stream.sync()
-            assert _last_prefetch_location(buf) == device.device_id
+            assert buf.last_prefetch_location == device
         finally:
             buf.close()