nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/backend/metal_device.py

@@ -0,0 +1,1274 @@
+"""The Metal device runtime, behind the seam `DeviceAllocator` already owns.
+
+Step 5 of the Metal adoption plan. `cuda` and `hip` are two symbol tables over
+one C ABI, so `_GPU_BACKENDS` can hold both as rows that differ only by name.
+**Metal cannot be a third row** — it is an Objective-C API with `MTLDevice`
+and `MTLBuffer` and there is no `metalMalloc` to name — and
+`test_device_backend_seam.py` pins exactly that, both in its contract
+assertion (`set(backends) == {"cuda", "hip"}`) and in its own docstring:
+*"it will be a second implementation behind this class rather than a third
+row in the dict."*
+
+This module is that second implementation. It presents the **call contract**
+the seam already uses — a runtime object whose entry points take `ctypes`
+arguments, write results through `ctypes.byref`, and return an integer status
+— so every one of the 49 runtime touches inside `DeviceAllocator` works
+unchanged. The ctypes shape is the seam's existing convention, not a pretence
+that Metal has a C ABI.
+
+## What unified memory deletes
+
+Apple Silicon has one physical memory shared by CPU and GPU, so a buffer made
+with `MTLResourceStorageModeShared` has a single address that is valid from
+both. `contents()` hands that address back, and it is the same address a
+kernel receives. Three consequences, and they are simplifications rather than
+work:
+
+* **H2D and D2H do not exist.** All four `cudaMemcpy` kinds are one byte copy
+  between two addresses in the same address space.
+* **Pinned host memory does not exist either**, because there is no DMA to
+  pin for. `malloc_host` returns a shared buffer, which genuinely satisfies
+  what `cudaMallocHost` is for: a host address the device can read with no
+  copy. That is a correct mapping, not a fallback.
+* **Peer access does not apply.** `MTLCopyAllDevices()` returns one device on
+  this hardware, and the key is simply absent, which makes
+  `DeviceAllocator.enable_peer_access` answer False — true, not silent.
+
+## Zero fallback
+
+Every failure is loud. If Metal is absent, `runtime()` raises. If a buffer
+cannot be allocated, the entry point returns a non-zero status and
+`DeviceAllocator.malloc_cuda` raises its OOM with full diagnostics. Nothing
+here ever hands back host memory pretending it is device memory, and nothing
+degrades to CPU.
+
+What Metal has no honest equivalent for is **omitted rather than stubbed**:
+events, `memcpy_async`, the stream-ordered allocator. The seam already turns
+a missing key into an explicit
+`RuntimeError("... unsupported on backend 'metal'")`, so omission is the
+loud refusal and a no-op stub would have been the silent wrong.
+
+R33 preserved — no torch, at the boundary included.
+R34 preserved — nothing here is keyed on a model name.
+"""
+
+from __future__ import annotations
+
+import bisect
+import ctypes
+import os
+import threading
+from typing import Dict, Optional, Tuple
+import nbxtensor.memory.allocator as _nbx_allocator  # the module the moved names live in (M2 cut of tensor.py)
+import nbxtensor.backend.runtime as _nbx_runtime  # the module the moved names live in (M2 cut of tensor.py)
+
+# Status codes. 0 is success, mirroring the C runtimes the seam was written
+# against; the non-zero values are ours and appear only in NeuroBrix errors.
+_OK = 0
+_ERR_ALLOC = 1
+_ERR_BAD_DEVICE = 2
+_ERR_UNKNOWN_POINTER = 3
+_ERR_BAD_ARGUMENT = 4
+
+
+
+class _MappedView:
+    """The registry's view slot for a buffer over a file mapping: it keeps the ctypes region the buffer was
+    made from; nothing to release (the mapping's owner unmaps it; the MTLBuffer is dropped with the entry)."""
+
+    def __init__(self, region=None):
+        self.region = region
+
+    def release(self):
+        pass
+
+class MetalUnavailableError(RuntimeError):
+    """Metal, or its Python bindings, are not usable on this machine.
+
+    Raised rather than returned so that no caller can mistake absence for an
+    empty device list and quietly continue on the CPU.
+    """
+
+
+def _drop_base(bases: list, address: int) -> None:
+    """Remove one base from the sorted index, in log time.
+
+    A freed base left in the index would let a later interior lookup resolve
+    into memory the allocator no longer owns — the exact opposite of what the
+    ownership rule is for.
+    """
+    index = bisect.bisect_left(bases, address)
+    if index < len(bases) and bases[index] == address:
+        del bases[index]
+
+
+def _deref(arg):
+    """The object behind a `ctypes.byref(...)` out-parameter.
+
+    The seam passes out-parameters exactly as it does to a C runtime. A
+    `CArgObject` carries the object it references as `_obj`; a real pointer
+    carries `contents`. Both are accepted so this works whichever the caller
+    used.
+    """
+    if hasattr(arg, "_obj"):
+        return arg._obj
+    if hasattr(arg, "contents"):
+        return arg.contents
+    raise TypeError(f"not an out-parameter: {type(arg).__name__}")
+
+
+def _declared_dispatch_fact(key: str) -> int:
+    """A `dispatch.<key>` the device file must state, positive — refused by name otherwise (ZERO FALLBACK: it was 32 / 512
+    written here under a bare `except`, the Dell's reading 2026-11-05 01:20; every Apple device file declares both)."""
+    from nbxtensor.ops.wrappers import _arch_dispatch_param
+    value = _arch_dispatch_param(key, None)
+    if value is None or int(value) <= 0:
+        raise RuntimeError(f"ZERO FALLBACK: the device file states no positive dispatch.{key} ({value!r}) — the Metal runtime "
+                           f"does not invent one; declare it in nbxtensor/backend/devices/apple/<arch>.yml")
+    return int(value)
+
+
+def _as_int(value) -> int:
+    """A Python int from a ctypes scalar or a plain int."""
+    return int(getattr(value, "value", value) or 0)
+
+
+#: The `new*` selectors this runtime calls through pyobjc. Every one returns an object the caller
+#: OWNS (+1, the Cocoa "new" family — for instance methods too, per clang's ownership rules).
+_NEW_FAMILY_SELECTORS = (b"newBufferWithLength:options:", b"newCommandQueue", b"newSharedEvent")
+_NEW_FAMILY_DECLARED = False
+
+
+def declare_new_family_ownership() -> None:
+    """Tell pyobjc that the `new*` Metal selectors return an already-retained object.
+
+    pyobjc up to 12.2.2 applies the "new" ownership rule to class methods only, treats the +1
+    result of an instance method such as `-[MTLDevice newBufferWithLength:options:]` as
+    autoreleased, retains it again, and never releases that retain: EVERY MTLBuffer this runtime
+    ever made outlived its free (upstream: ronaldoussoren/pyobjc#690, fixed in pyobjc 12.2.3,
+    not yet published when this was written; dipy/dipy#4190 met the same leak). Measured
+    2026-09-26 on an M4 Pro with pyobjc 12.2.2: a new buffer had retainCount 2, and after `del`
+    `MTLDevice.currentAllocatedSize()` did not move; declared as below, retainCount 1 and the
+    allocation returns to 0 at `del`. The process footprint grew by every freed buffer the GPU
+    had written — about 60 MB per key in a certifier sweep, into swap in a long render.
+
+    The declaration must be made after `import Metal` (loading the framework wrapper installs its
+    own metadata) and before the first call of each selector (pyobjc resolves a selector's
+    signature once). It is idempotent, and it stays correct under a pyobjc that fixes the rule
+    itself: both then say the same thing. `test_metal_new_family_is_not_leaked.py` measures it."""
+    global _NEW_FAMILY_DECLARED
+    if _NEW_FAMILY_DECLARED:
+        return
+    import objc
+    import Metal  # noqa: F401  the wrapper's own metadata first, ours over it
+    for selector in _NEW_FAMILY_SELECTORS:
+        objc.registerMetaDataForSelector(
+            b"NSObject", selector, {"retval": {"already_retained": True}, "arguments": {}})
+    _NEW_FAMILY_DECLARED = True
+
+
+class MetalRuntime:
+    """One Metal device, presented as the runtime object the seam expects.
+
+    Instantiated once per process by `runtime()`. Not thread-safe by
+    construction beyond the registry lock, which matches the rest of
+    `DeviceAllocator` ("single-threaded per component by convention").
+    """
+
+    def __init__(self) -> None:
+        try:
+            import Metal  # pyobjc-framework-Metal
+        except ImportError as exc:                      # pragma: no cover
+            raise MetalUnavailableError(
+                "The Metal Python bindings are not installed. NeuroBrix "
+                "reaches Apple GPUs through pyobjc:\n"
+                "    pip install pyobjc-framework-Metal\n"
+                "It is a dependency of the Metal Triton backend, so this "
+                "usually means the backend itself is missing — see "
+                "nbxtensor.backend.metal_backend."
+            ) from exc
+
+        self._Metal = Metal
+        declare_new_family_ownership()
+        device = Metal.MTLCreateSystemDefaultDevice()
+        if device is None:
+            raise MetalUnavailableError(
+                "No Metal device. `MTLCreateSystemDefaultDevice()` returned "
+                "nil, so this process cannot reach a GPU at all — a headless "
+                "session or a virtualised host without GPU passthrough will "
+                "do this. NeuroBrix refuses rather than computing on the CPU "
+                "under a device name."
+            )
+        if not device.hasUnifiedMemory():
+            raise MetalUnavailableError(
+                f"Metal device {device.name()!r} does not report unified "
+                f"memory. Every address this allocator hands out is a "
+                f"shared-storage buffer's `contents()`, which is only the "
+                f"SAME address the GPU sees when host and device share the "
+                f"memory. On a discrete Metal GPU that assumption is false "
+                f"and the pointers would be silently wrong, so this refuses "
+                f"instead. Apple Silicon reports unified memory; an Intel "
+                f"Mac with a discrete card does not."
+            )
+        self._device = device
+        # One queue owned by the allocator. See `sync` for exactly what it
+        # does and does not order against.
+        self._queue = device.newCommandQueue()
+        if self._queue is None:                         # pragma: no cover
+            raise MetalUnavailableError(
+                "Metal device present but `newCommandQueue()` returned nil.")
+
+        # The VM page size, read from the system. Every allocation is rounded
+        # up to a multiple of it, and that is not a micro-optimisation:
+        # `MTLDevice.newBufferWithBytesNoCopy:length:options:deallocator:` —
+        # which is how the Triton Metal backend turns one of our pointers
+        # into a buffer it can bind — requires a PAGE-ALIGNED pointer and a
+        # page-multiple length. Metal sub-allocates buffers smaller than a
+        # page inside one, so `contents()` for them lands mid-page. Measured
+        # 2026-09-05: of twelve sizes from 1 byte to 64 KB, the eight below
+        # one page all came back unaligned (offsets 1280, 1536, 1792, 2048,
+        # 2304, 3328, 7424, 15616) and every one of them was page-aligned
+        # once the length was rounded up. Handing a sub-page pointer to
+        # NoCopy is undefined behaviour, and it showed as
+        # "Failed to create Metal buffer" under memory pressure and, once, a
+        # SIGBUS.
+        self._page_size = int(os.sysconf("SC_PAGESIZE"))
+
+        self._lock = threading.RLock()
+        # address -> (MTLBuffer, memoryview, nbytes). The buffer entry is what
+        # keeps the allocation alive: dropping it is the free.
+        self._buffers: Dict[int, Tuple[object, memoryview, int]] = {}
+        self._host_buffers: Dict[int, Tuple[object, memoryview, int]] = {}
+        #: Every live base address, sorted, so an interior pointer resolves
+        #: by bisection instead of a scan. The engine sub-allocates arenas,
+        #: so interior pointers are the common case, not the exception.
+        self._bases: list = []
+        # No byte is counted here: the one account is the allocator's (`memory.allocator.DeviceAllocator`), and the
+        # device's budget is refused there (`memory.metal.MetalMemory`), never a second figure of the same fact.
+        # handle -> MTLCommandQueue, for the stream API.
+        self._queues: Dict[int, object] = {}
+        self._next_queue_handle = 1
+        # handle -> event record, for the event API.
+        self._events: Dict[int, dict] = {}
+        self._next_event_handle = 1
+        # Retired shared events, kept for reuse.
+        #
+        # `newSharedEvent` is not an unlimited resource: an autotune sweep
+        # times every candidate config of every op with a fresh pair, and a
+        # long run exhausted the device — `newSharedEvent` began returning
+        # nil, which surfaced as "event create failed (error 1)" in the middle
+        # of a matmul. Destroying them was not enough; the objects are
+        # released but a run that never returns to the event loop does not
+        # necessarily reclaim them.
+        #
+        # A shared event is reusable by construction — signalling is a
+        # monotonically increasing value on the same object — so a retired one
+        # goes back on this list with its counter intact and the next create
+        # takes it. The peak count is what a run actually holds at once, which
+        # for the timing path is two.
+        self._event_pool: list = []
+
+        # Command buffers committed and NOT yet waited on, oldest first.
+        #
+        # A launch used to commit its buffer and block on
+        # `waitUntilCompleted()` before returning, so every kernel cost a full
+        # host round trip. Measured 2026-09-07 on TinyLlama decode: 598 us per
+        # launch across 10,695 launches, 1,337 of them per decode step —
+        # 799 ms of an 857 ms step, 93% of it, spent waiting for a GPU that
+        # had already been given the work.
+        #
+        # Metal executes the buffers of one queue in the order they were
+        # committed, so dropping the wait keeps the ordering that made the
+        # engine correct. What it does NOT keep is the accident that made
+        # every host read of device memory safe: `memcpy` here is a memmove on
+        # shared memory and orders against nothing. So the wait moves to the
+        # places where the host actually looks — `memcpy`, `memset`, `sync`,
+        # `stream_sync` — and to `free`, because a buffer referenced by
+        # pending work must not be released.
+        #
+        # The list is bounded: Metal blocks in `commandBuffer()` once a queue
+        # holds 64 in flight, and blocking there is a stall nobody can see.
+        # `_drain_to` keeps it under the cap read from the hardware profile.
+        self._pending: Dict[int, list] = {}
+        self._pending_cap: Optional[int] = None
+
+        # The OPEN command buffer of each stream, and the encoder inside it.
+        #
+        # A NeuroBrix stream means ordered work that synchronises when it is
+        # asked to, and on Metal that is one command buffer with one
+        # serialised compute encoder — opened at the first launch, committed
+        # at a synchronisation point, an event, or a host read. It is the
+        # same contract a CUDA stream keeps; this is where it is kept here.
+        #
+        # It was one command buffer per LAUNCH until 2026-09-07, which put
+        # 1,505 of them in a single decode step: 501 ms of the step's 593 ms
+        # was draining them, and Metal's own GPUStartTime/GPUEndTime summed to
+        # 806 ms across those buffers — more than the step itself, because
+        # what that sum measures at this granularity is residency, not work.
+        #
+        # `_open` holds (command_buffer, encoder, kind, dispatches). A request
+        # for a different encoder kind ends the current one and opens the new
+        # one in the SAME buffer, so a blit and the kernels around it keep
+        # their order without a commit between them.
+        self._open: Dict[int, list] = {}
+        self._dispatch_cap: Optional[int] = None
+
+        # An encoder that is FREED while still open aborts the whole process:
+        #   -[_MTLCommandEncoder dealloc]:134: failed assertion
+        #   'Command encoder released without endEncoding'
+        # The encode paths already close their stream when encoding raises.
+        # What none of them covers is an error that escapes ALL the way out —
+        # a compile-time refusal, most of all — leaving a stream open with a
+        # live encoder for the interpreter to free at teardown.
+        #
+        # Measured 2026-09-09 on the whole `stt` family: after a `tt.dot`
+        # refusal the three models did not fall back, they ABORTED (rc=-6),
+        # and the assertion is what the log showed instead of the refusal that
+        # caused it. A loud, recoverable decline became a crash, and it took
+        # the family with it.
+        #
+        # So the encoders are ended at interpreter shutdown too. Ending is all
+        # that is needed — the buffers are not committed, because at exit
+        # there is nothing left to read their results.
+        import atexit as _atexit
+        _atexit.register(self._end_open_encoders)
+
+        # `DeviceAllocator.event_elapsed_ms` sets `.restype` on the entry
+        # point before calling it — a ctypes idiom that a bound method
+        # rejects. A function object accepts attributes, so this one entry
+        # point is exposed as a closure rather than a method. The seam's
+        # convention, honoured rather than worked around.
+        def event_elapsed(out_ms, start, end, _self=self):
+            return _self._event_elapsed(out_ms, start, end)
+
+        self.event_elapsed = event_elapsed
+
+        try:
+            self._device_count = len(Metal.MTLCopyAllDevices())
+        except Exception:                               # pragma: no cover
+            # MTLCopyAllDevices is macOS-only; a device exists, so it is one.
+            self._device_count = 1
+
+    # -- properties read from the device, never assumed ---------------------
+
+    @property
+    def device_name(self) -> str:
+        return str(self._device.name())
+
+    @property
+    def arch_name(self) -> str:
+        """The Triton target arch string for this GPU.
+
+        The Metal device name, lowercased and hyphenated — "apple-m4-pro".
+        Computed here rather than imported from the Metal backend's driver:
+        that module carries torch imports, and R33 does not make an exception
+        for Apple. The rule is one line of string handling, and the profile
+        matches it by prefix (`config/vendors/apple/apple_silicon.yml`).
+        """
+        from nbxtensor.backend.device_names import arch_string
+        return arch_string(self._device.name())
+
+    @property
+    def has_unified_memory(self) -> bool:
+        return bool(self._device.hasUnifiedMemory())
+
+    @property
+    def max_threadgroup_memory(self) -> int:
+        """Threadgroup memory per threadgroup, in bytes, from the device.
+
+        The one number the Apple profile calls its binding constraint, and
+        the profile has it `hand_curated_from_docs`. Exposed here so it can
+        be measured rather than read.
+        """
+        return int(self._device.maxThreadgroupMemoryLength())
+
+    # -- allocation ---------------------------------------------------------
+
+    def budget_bytes(self) -> int:
+        """The device's own working-set budget, read from the device."""
+        return int(self._device.recommendedMaxWorkingSetSize())
+
+    # Why the budget is checked against OUR live count and not against
+    # `MTLDevice.currentAllocatedSize()`, which is the obvious candidate:
+    # measured 2026-09-05, that counter does not come back down. Four cycles
+    # of allocating 2 GiB and releasing it read 2048 / 4096 / 6144 / 8192 MB
+    # while the process RSS stayed flat at ~2.3 GB — so the memory really is
+    # released and reused, and the counter is cumulative rather than live.
+    # Using it as a budget gauge would have produced an allocator that OOMs
+    # permanently after enough churn, on a device with nothing allocated.
+
+    def _new_shared_buffer(self, nbytes: int):
+        """A shared-storage MTLBuffer plus the address the CPU and GPU share.
+
+        The device's working-set budget is refused by the ONE account before it asks here
+        (`memory.allocator.DeviceAllocator`, `memory.metal.MetalMemory.budget`), and that is deliberate.
+        `recommendedMaxWorkingSetSize` is a recommendation: Metal will keep handing out buffers past it,
+        backed by a memory the whole machine is sharing, until the system starts paging or the process
+        is killed. An allocator that never says no turns "fill the device until it refuses" — which is
+        how the engine's own OOM-reclaim path is exercised — into an unbounded claim on the user's RAM.
+        Apple's number is the honest ceiling and it is read from the device, never guessed.
+        """
+        Metal = self._Metal
+        page = self._page_size
+        allocated = (nbytes + page - 1) // page * page
+        buffer = self._device.newBufferWithLength_options_(
+            allocated, Metal.MTLResourceStorageModeShared)
+        if buffer is None:
+            return None, None, 0, 0
+        view = buffer.contents().as_buffer(allocated)
+        # `from_buffer` holds an export on the view only for as long as the
+        # returned object lives; the address belongs to the MTLBuffer, which
+        # the registry keeps alive, so it stays valid once the keeper goes.
+        keeper = ctypes.c_char.from_buffer(view)
+        address = ctypes.addressof(keeper)
+        del keeper
+        if address % page:                              # pragma: no cover
+            # Rounding is supposed to guarantee this. If Metal ever stops
+            # honouring it, refuse rather than hand out a pointer the
+            # backend will wrap into a buffer with undefined behaviour.
+            return None, None, 0, 0
+        return buffer, view, address, allocated
+
+    def malloc(self, out_ptr, size) -> int:
+        """`cudaMalloc(&ptr, size)`."""
+        nbytes = _as_int(size)
+        slot = _deref(out_ptr)
+        if nbytes <= 0:
+            slot.value = 0
+            return _ERR_BAD_ARGUMENT if nbytes < 0 else _OK
+        buffer, view, address, allocated = self._new_shared_buffer(nbytes)
+        if buffer is None or not address:
+            slot.value = 0
+            return _ERR_ALLOC
+        with self._lock:
+            self._buffers[address] = (buffer, view, allocated)
+            bisect.insort(self._bases, address)
+        slot.value = address
+        return _OK
+
+    def free(self, ptr) -> int:
+        """`cudaFree(ptr)`. Dropping the MTLBuffer reference IS the free."""
+        address = _as_int(ptr)
+        if not address:
+            return _OK
+        # Releasing memory a committed command buffer still references is a
+        # use-after-free on the GPU's side of the fence.
+        self.flush()
+        with self._lock:
+            entry = self._buffers.pop(address, None)
+            if entry is not None:
+                _drop_base(self._bases, address)
+        if entry is None:
+            # Freeing a pointer this runtime never handed out is a bug in the
+            # caller, and it must not pass silently: on CUDA the same call
+            # returns cudaErrorInvalidValue.
+            return _ERR_UNKNOWN_POINTER
+        entry[1].release()
+        return _OK
+
+    def malloc_host(self, out_ptr, size) -> int:
+        """`cudaMallocHost(&ptr, size)` — a shared buffer, and that is exact.
+
+        `cudaMallocHost` exists to page-lock host memory so a DMA engine can
+        stream it to the device. Apple Silicon has no such transfer: the
+        address a shared buffer hands back is already readable by the GPU
+        with no copy, which is the whole property the caller wants.
+        """
+        nbytes = _as_int(size)
+        slot = _deref(out_ptr)
+        if nbytes <= 0:
+            slot.value = 0
+            return _ERR_BAD_ARGUMENT if nbytes < 0 else _OK
+        buffer, view, address, allocated = self._new_shared_buffer(nbytes)
+        if buffer is None or not address:
+            slot.value = 0
+            return _ERR_ALLOC
+        with self._lock:
+            self._host_buffers[address] = (buffer, view, allocated)
+            bisect.insort(self._bases, address)
+        slot.value = address
+        return _OK
+
+    def free_host(self, ptr) -> int:
+        address = _as_int(ptr)
+        if not address:
+            return _OK
+        with self._lock:
+            entry = self._host_buffers.pop(address, None)
+            if entry is not None:
+                _drop_base(self._bases, address)
+        if entry is None:
+            return _ERR_UNKNOWN_POINTER
+        entry[1].release()
+        return _OK
+
+    def bind_mapped(self, address: int, nbytes: int) -> int:
+        """Make `nbytes` of a file mapping at `address` bindable in place: a shared buffer over the pages that
+        hold them (`newBufferWithBytesNoCopy`, page-aligned as Metal requires), registered so a launch resolves
+        any address inside it to (buffer, offset) like every other allocation (`buffer_for_pointer`). No byte
+        is copied and none is counted against the working set: the pages are the file's, the page cache's
+        to keep or reclaim. The unified expert store reads a routed expert this way, the expert's own pages
+        only (one buffer per expert, never a whole shard: a command buffer keeps resident what it binds).
+        Returns the page base registered; a range already covered returns its base."""
+        page = self._page_size
+        start = int(address) // page * page
+        end = -(-(int(address) + int(nbytes)) // page) * page
+        with self._lock:
+            entry = self._host_buffers.get(start)
+            if entry is not None and start + entry[2] >= end:
+                return start
+            if entry is not None:                     # a neighbour's buffer starts on this page: cover both
+                end = max(end, start + entry[2])
+        # PyObjC takes the bytes as a buffer object, not an integer: a ctypes array AT the address (no copy —
+        # checked: the buffer's contents() is the mapping's own address)
+        region = (ctypes.c_char * (end - start)).from_address(start)
+        buffer = self._device.newBufferWithBytesNoCopy_length_options_deallocator_(
+            region, end - start, self._Metal.MTLResourceStorageModeShared, None)
+        if buffer is None:
+            raise RuntimeError(f"ZERO FALLBACK: Metal refused to bind {end - start} bytes of a file mapping at "
+                               f"{start:#x} in place (newBufferWithBytesNoCopy returned nil)")
+        with self._lock:
+            if start not in self._host_buffers:
+                bisect.insort(self._bases, start)
+            self._host_buffers[start] = (buffer, _MappedView(region), end - start)
+        return start
+
+    def unbind_mapped(self, start: int) -> bool:
+        """Drop the in-place buffer registered at page `start` (`bind_mapped`): its pages are the file's again,
+        nothing holds them resident. A command buffer still running retains the MTLBuffer it bound."""
+        with self._lock:
+            entry = self._host_buffers.pop(int(start), None)
+            if entry is not None:
+                _drop_base(self._bases, int(start))
+        return entry is not None
+
+    # -- copies -------------------------------------------------------------
+
+    def memcpy(self, dst, src, size, kind=None) -> int:
+        """`cudaMemcpy(dst, src, n, kind)` — one address space, so one copy.
+
+        `kind` is accepted and ignored on purpose: H2D, D2H, D2D and H2H are
+        the same operation when host and device share the memory. `memmove`
+        rather than `memcpy` because it is correct for every input the
+        caller can produce, including the overlap CUDA leaves undefined.
+        """
+        nbytes = _as_int(size)
+        if nbytes <= 0:
+            return _OK if nbytes == 0 else _ERR_BAD_ARGUMENT
+        dst_addr, src_addr = _as_int(dst), _as_int(src)
+        if not dst_addr or not src_addr:
+            return _ERR_BAD_ARGUMENT
+        # Device to device is STREAM-ORDERED by contract, not host-ordered:
+        # "for D2D the synchronous cudaMemcpy performs NO host-side
+        # synchronization — the copy is enqueued on the calling device's
+        # legacy stream and the host returns immediately" (the allocator says
+        # so in nbx_tensor.memcpy). Doing it as a host memmove made it
+        # host-ordered, and once the launch path stopped waiting that turned
+        # every one of them into a flush: measured 2026-09-07, 1,086 D2D
+        # copies in an 8-token decode went from 0.007 s to 4.08 s, taking back
+        # most of what dropping the per-launch wait had given. A blit on this
+        # allocator's queue is ordered by the GPU against the work on that
+        # queue and waits for nothing; the Triton kernels run on triton-ext's
+        # own queue, and the launcher drains this one before each of them.
+        if _as_int(kind) == 3:
+            blitted = self._blit(dst_addr, src_addr, nbytes)
+            if blitted is not None:
+                return blitted
+        # Host is involved on one side or the other: a memmove on shared
+        # memory orders against nothing, so the GPU work that produced these
+        # bytes has to have retired before it runs. It always had; it used to
+        # be true by accident because every launch waited.
+        self.flush()
+        ctypes.memmove(ctypes.c_void_p(dst_addr),
+                       ctypes.c_void_p(src_addr), nbytes)
+        return _OK
+
+    def memset(self, ptr, value, size) -> int:
+        nbytes = _as_int(size)
+        if nbytes <= 0:
+            return _OK if nbytes == 0 else _ERR_BAD_ARGUMENT
+        address = _as_int(ptr)
+        if not address:
+            return _ERR_BAD_ARGUMENT
+        self.flush()
+        ctypes.memset(ctypes.c_void_p(address),
+                      _as_int(value) & 0xFF, nbytes)
+        return _OK
+
+    # -- device queries -----------------------------------------------------
+
+    def set_device(self, index) -> int:
+        """Metal exposes one GPU here; anything but 0 is an error, not a
+        silent clamp to the only device."""
+        return _OK if _as_int(index) == 0 else _ERR_BAD_DEVICE
+
+    def get_device(self, out_index) -> int:
+        _deref(out_index).value = 0
+        return _OK
+
+    def device_count(self, out_count) -> int:
+        _deref(out_count).value = self._device_count
+        return _OK
+
+    def mem_get_info(self, out_free, out_total) -> int:
+        """`cudaMemGetInfo(&free, &total)`, read from the device.
+
+        `total` is `recommendedMaxWorkingSetSize()` — Apple's own budget for
+        this device, which on unified memory is the meaningful ceiling rather
+        than installed RAM, because that RAM is shared with everything else
+        running. `free` is that budget minus `currentAllocatedSize()`, so it
+        is a budget headroom rather than a physical free-page count. The
+        callers of this in NeuroBrix are feasibility probes and the pool's
+        parked-bytes cap, both of which want exactly that.
+        """
+        from nbxtensor.memory.metal import MetalMemory
+        rc, free, total = MetalMemory(self).mem_get_info()      # headroom against the ONE account
+        _deref(out_total).value = total
+        _deref(out_free).value = free
+        return rc
+
+    # -- ordering -----------------------------------------------------------
+
+    # -- deferred completion ------------------------------------------------
+
+    def _cap(self) -> int:
+        """How many committed command buffers may be in flight at once.
+
+        Read from the hardware profile (`dispatch.max_command_buffers_in_flight`)
+        so the number is the device's, not this file's. Metal's own ceiling is
+        64 per queue and it BLOCKS in `commandBuffer()` at that point, which is
+        a stall with no name attached; staying under it keeps the backpressure
+        ours and visible.
+        """
+        if self._pending_cap is None:
+            self._pending_cap = _declared_dispatch_fact("max_command_buffers_in_flight")
+        return self._pending_cap
+
+    def track_committed(self, queue_handle: int, command_buffer) -> None:
+        """Record a buffer committed without a wait, and keep the queue's
+        in-flight depth under the cap by waiting on the oldest."""
+        with self._lock:
+            pending = self._pending.setdefault(int(queue_handle or 0), [])
+            pending.append(command_buffer)
+            cap = self._cap()
+            while len(pending) > cap:
+                oldest = pending.pop(0)
+                self._await(oldest)
+
+    def _await(self, command_buffer) -> None:
+        command_buffer.waitUntilCompleted()
+        error = command_buffer.error()
+        if error is not None:
+            raise RuntimeError(f"Metal command buffer failed: {error}")
+
+    def _dispatches_per_buffer(self) -> int:
+        """How many dispatches one command buffer may carry before it is
+        committed, from `dispatch.max_dispatches_per_command_buffer` in the
+        hardware profile. It bounds latency and how much work a single
+        failure takes down with it; it is not a correctness knob."""
+        if self._dispatch_cap is None:
+            self._dispatch_cap = _declared_dispatch_fact("max_dispatches_per_command_buffer")
+        return self._dispatch_cap
+
+    def encoder_for(self, stream: int, kind: str = "compute"):
+        """The open encoder of this stream, opening one if there is none.
+
+        Returns `(command_buffer, encoder)`. The caller encodes and then calls
+        `note_dispatch`; it must not commit — the stream owns that.
+        """
+        key = int(stream or 0)
+        with self._lock:
+            entry = self._open.get(key)
+            if entry is not None and entry[2] != kind:
+                # Same buffer, different encoder: end the current one and open
+                # the kind asked for, so ordering is kept without a commit.
+                entry[1].endEncoding()
+                entry[1] = self._new_encoder(entry[0], kind)
+                entry[2] = kind
+                return entry[0], entry[1]
+            if entry is not None:
+                return entry[0], entry[1]
+            queue = self._resolve_queue(key)
+            if queue is None:
+                return None, None
+            command_buffer = queue.commandBuffer()
+            if command_buffer is None:                  # pragma: no cover
+                return None, None
+            encoder = self._new_encoder(command_buffer, kind)
+            if encoder is None:                         # pragma: no cover
+                command_buffer.commit()
+                return None, None
+            self._open[key] = [command_buffer, encoder, kind, 0]
+            return command_buffer, encoder
+
+    @staticmethod
+    def _new_encoder(command_buffer, kind: str):
+        if kind == "blit":
+            return command_buffer.blitCommandEncoder()
+        # `computeCommandEncoder()` is MTLDispatchTypeSerial: the dispatches
+        # inside it run in order with the barriers Metal inserts itself,
+        # which is what makes one encoder able to carry a whole graph.
+        return command_buffer.computeCommandEncoder()
+
+    def note_dispatch(self, stream: int) -> None:
+        """One more piece of work encoded on this stream; commit if the
+        buffer has carried enough of them."""
+        key = int(stream or 0)
+        with self._lock:
+            entry = self._open.get(key)
+            if entry is None:
+                return
+            entry[3] += 1
+            if entry[3] < self._dispatches_per_buffer():
+                return
+        self.close_stream(key)
+
+    def _end_open_encoders(self) -> None:
+        """End every open encoder, without committing. Teardown only.
+
+        `close_stream` ends AND commits AND tracks, which is right during a
+        run and wrong at exit: a commit at teardown queues work whose results
+        nobody will read, and tracking it adds a buffer nothing will await.
+        Here only the assertion is being avoided, so only `endEncoding` is
+        called. Never raises: this runs while the interpreter is going down,
+        and an exception there replaces one confusing abort with another.
+        """
+        try:
+            with self._lock:
+                entries = list(self._open.values())
+                self._open.clear()
+        except Exception:                               # pragma: no cover
+            return
+        for entry in entries:
+            try:
+                entry[1].endEncoding()
+            except Exception:                           # pragma: no cover
+                pass
+
+    def close_stream(self, stream: Optional[int] = None) -> None:
+        """End the open encoder and commit the buffer, without waiting."""
+        keys = [int(stream or 0)] if stream is not None else list(self._open)
+        for key in keys:
+            with self._lock:
+                entry = self._open.pop(key, None)
+            if entry is None:
+                continue
+            command_buffer, encoder = entry[0], entry[1]
+            try:
+                encoder.endEncoding()
+            except Exception:                           # pragma: no cover
+                pass
+            command_buffer.commit()
+            self.track_committed(key, command_buffer)
+
+    def abandon_stream(self, stream: int) -> None:
+        """Close the stream's buffer after an encoding failure.
+
+        A command buffer holds its queue's in-flight slot from creation until
+        it COMPLETES; one abandoned mid-encode never completes and the slot is
+        gone for the life of the process."""
+        self.close_stream(int(stream or 0))
+
+    def _blit(self, dst_addr: int, src_addr: int, nbytes: int):
+        """A device-to-device copy as GPU work on this allocator's queue (not the
+        Triton kernels' queue: triton-ext runs those on its own, and the launcher
+        drains this queue before every kernel).
+
+        Returns `_OK` when it was enqueued, or None when either end is not
+        memory this allocator handed out — in which case the caller falls
+        back to the host copy, because a blit can only address MTLBuffers.
+        """
+        dst = self.buffer_for_pointer(dst_addr)
+        src = self.buffer_for_pointer(src_addr)
+        if dst is None or src is None or dst[0] is None or src[0] is None:
+            return None
+        # Into the stream's own buffer, as a blit encoder between the compute
+        # encoders around it: one command buffer, order kept by Metal.
+        _cb, encoder = self.encoder_for(0, kind="blit")
+        if encoder is None:                             # pragma: no cover
+            return None
+        encoder.copyFromBuffer_sourceOffset_toBuffer_destinationOffset_size_(
+            src[0], src[1], dst[0], dst[1], nbytes)
+        self.note_dispatch(0)
+        return _OK
+
+    def flush(self, queue_handle: Optional[int] = None) -> int:
+        """Wait for everything committed on this queue (or all of them).
+
+        This is the seam the launch path traded its per-kernel wait for: the
+        host may read device memory only after it, and every entry point that
+        does — memcpy, memset, sync, stream_sync, free — goes through it.
+        """
+        # Work still being encoded has not been committed, so waiting on the
+        # pending list alone would return before it ran. The stream is closed
+        # first: that is what "synchronise" means for a stream that batches.
+        self.close_stream(None if queue_handle is None else int(queue_handle or 0))
+        with self._lock:
+            if queue_handle is None:
+                buffers = [cb for lst in self._pending.values() for cb in lst]
+                self._pending.clear()
+            else:
+                buffers = self._pending.pop(int(queue_handle or 0), [])
+        for command_buffer in buffers:
+            self._await(command_buffer)
+        return _OK
+
+    def has_pending_gpu_writes(self) -> bool:
+        """True when this queue holds work that could still be writing memory.
+
+        The launch path's drain exists for ONE crossing: our queue writes and
+        triton-ext's kernel reads. Only two things on this queue write memory a
+        kernel could read — a blit still being encoded (an open stream), and a
+        buffer already committed and not yet retired (`_pending`). Both are
+        counted here.
+
+        `event_record` and `stream_wait_event` also commit buffers and are NOT
+        counted, deliberately: they encode a signal and a wait, and neither
+        writes device memory. They are also not tracked in `_pending`, so a
+        caller that needs to wait for them must not use this — which is why this
+        is named for GPU WRITES and not for emptiness, and why `sync()` below is
+        left alone with its full device-synchronize semantics.
+        """
+        with self._lock:
+            return bool(self._open) or any(self._pending.values())
+
+    def sync(self) -> int:
+        """`cudaDeviceSynchronize()` on the allocator's own queue.
+
+        A real device round trip: an empty command buffer is committed and
+        waited on, so anything already enqueued on this queue has retired.
+
+        **What it does not cover, named rather than left silent:** the Triton
+        Metal backend dispatches kernels on a queue of its own, and this does
+        not order against that queue. Today nothing else enqueues work — the
+        copies above are CPU `memmove` on shared memory — so there is nothing
+        it misses. When the launch path lands it must either share this queue
+        or this method must be re-pointed at the backend's, and that is a
+        contract for that workstream, not an option.
+        """
+        # Everything the launch path committed without waiting, first: the
+        # contract this docstring used to describe as owed to the launch
+        # workstream is now kept here.
+        self.flush()
+        command_buffer = self._queue.commandBuffer()
+        if command_buffer is None:                      # pragma: no cover
+            return _ERR_ALLOC
+        command_buffer.commit()
+        command_buffer.waitUntilCompleted()
+        error = command_buffer.error()
+        return _OK if error is None else _ERR_ALLOC
+
+    def stream_create(self, out_handle) -> int:
+        """A Metal command queue, which is what a CUDA stream is here.
+
+        Handles start at 1: zero is the seam's NULL/default stream and must
+        never be produced by a create.
+        """
+        queue = self._device.newCommandQueue()
+        if queue is None:                               # pragma: no cover
+            _deref(out_handle).value = 0
+            return _ERR_ALLOC
+        with self._lock:
+            handle = self._next_queue_handle
+            self._next_queue_handle += 1
+            self._queues[handle] = queue
+        _deref(out_handle).value = handle
+        return _OK
+
+    def stream_destroy(self, handle) -> int:
+        key = _as_int(handle)
+        if not key:
+            return _OK
+        with self._lock:
+            return _OK if self._queues.pop(key, None) is not None \
+                else _ERR_UNKNOWN_POINTER
+
+    def stream_sync(self, handle) -> int:
+        key = _as_int(handle)
+        if not key:
+            return self.sync()
+        with self._lock:
+            queue = self._queues.get(key)
+        if queue is None:
+            return _ERR_UNKNOWN_POINTER
+        self.flush(key)
+        command_buffer = queue.commandBuffer()
+        if command_buffer is None:                      # pragma: no cover
+            return _ERR_ALLOC
+        command_buffer.commit()
+        command_buffer.waitUntilCompleted()
+        return _OK if command_buffer.error() is None else _ERR_ALLOC
+
+    # -- events -------------------------------------------------------------
+    #
+    # A CUDA event is a marker placed in a stream. Metal's equivalent is an
+    # `MTLSharedEvent` signalled from a command buffer, and the command
+    # buffer also carries `GPUEndTime`, which is a GPU-timeline timestamp —
+    # so both halves of the CUDA event API map onto real Metal machinery and
+    # neither is faked. `timing=False` events skip nothing on Metal (the
+    # timestamp comes free with the command buffer) but they still REFUSE to
+    # be read as a stopwatch, because the engine's contract is that an
+    # ordering handle is not a clock and code relying on that must break
+    # here exactly as it breaks on CUDA.
+
+    def _resolve_queue(self, handle: int):
+        """The queue a stream handle names; 0 is the allocator's own."""
+        if not handle:
+            return self._queue
+        with self._lock:
+            return self._queues.get(handle)
+
+    def event_create(self, out_handle) -> int:
+        return self._event_create(out_handle, timing=True)
+
+    def event_create_flags(self, out_handle, flags) -> int:
+        # 0x02 == cudaEventDisableTiming.
+        return self._event_create(out_handle, timing=not (_as_int(flags) & 0x02))
+
+    def _event_create(self, out_handle, timing: bool) -> int:
+        with self._lock:
+            retired = self._event_pool.pop() if self._event_pool else None
+        if retired is not None:
+            shared, value = retired["shared"], retired["value"]
+        else:
+            shared, value = self._device.newSharedEvent(), 0
+        if shared is None:                              # pragma: no cover
+            _deref(out_handle).value = 0
+            return _ERR_ALLOC
+        with self._lock:
+            handle = self._next_event_handle
+            self._next_event_handle += 1
+            self._events[handle] = {"shared": shared, "value": value,
+                                    "timing": timing, "buffer": None}
+        _deref(out_handle).value = handle
+        return _OK
+
+    #: How many retired shared events to keep for reuse. A run holds a couple
+    #: at a time; the rest are dropped so a pathological caller cannot turn
+    #: this into a leak of its own.
+    _EVENT_POOL_CAP = 64
+
+    def event_destroy(self, handle) -> int:
+        key = _as_int(handle)
+        if not key:
+            return _OK
+        with self._lock:
+            entry = self._events.pop(key, None)
+            if entry is None:
+                return _ERR_UNKNOWN_POINTER
+            # The command buffer goes; the shared event and its counter stay,
+            # because the next create can take them as they are.
+            entry["buffer"] = None
+            if len(self._event_pool) < self._EVENT_POOL_CAP:
+                self._event_pool.append({"shared": entry["shared"],
+                                         "value": entry["value"]})
+            return _OK
+
+    def event_record(self, handle, stream=None) -> int:
+        """Signal the event from a command buffer on `stream`'s queue."""
+        key = _as_int(handle)
+        with self._lock:
+            event = self._events.get(key)
+        if event is None:
+            return _ERR_UNKNOWN_POINTER
+        queue = self._resolve_queue(_as_int(stream))
+        if queue is None:
+            return _ERR_UNKNOWN_POINTER
+        command_buffer = queue.commandBuffer()
+        if command_buffer is None:                      # pragma: no cover
+            return _ERR_ALLOC
+        with self._lock:
+            event["value"] += 1
+            value = event["value"]
+        command_buffer.encodeSignalEvent_value_(event["shared"], value)
+        command_buffer.commit()
+        event["buffer"] = command_buffer
+        return _OK
+
+    def event_sync(self, handle) -> int:
+        key = _as_int(handle)
+        with self._lock:
+            event = self._events.get(key)
+        if event is None:
+            return _ERR_UNKNOWN_POINTER
+        command_buffer = event.get("buffer")
+        if command_buffer is None:
+            # Never recorded: nothing to wait for, exactly as on CUDA.
+            return _OK
+        command_buffer.waitUntilCompleted()
+        return _OK if command_buffer.error() is None else _ERR_ALLOC
+
+    def stream_wait_event(self, stream, handle, flags=None) -> int:
+        """Make a queue wait on the event, without blocking the host."""
+        key = _as_int(handle)
+        with self._lock:
+            event = self._events.get(key)
+        if event is None:
+            return _ERR_UNKNOWN_POINTER
+        queue = self._resolve_queue(_as_int(stream))
+        if queue is None:
+            return _ERR_UNKNOWN_POINTER
+        command_buffer = queue.commandBuffer()
+        if command_buffer is None:                      # pragma: no cover
+            return _ERR_ALLOC
+        command_buffer.encodeWaitForEvent_value_(event["shared"],
+                                                 event["value"])
+        command_buffer.commit()
+        return _OK
+
+    def _event_elapsed(self, out_ms, start, end) -> int:
+        """Milliseconds between two recorded events, on the GPU timeline."""
+        with self._lock:
+            a = self._events.get(_as_int(start))
+            b = self._events.get(_as_int(end))
+        if a is None or b is None:
+            return _ERR_UNKNOWN_POINTER
+        if not a["timing"] or not b["timing"]:
+            # The whole point of an ordering handle: refuse rather than
+            # return a plausible wrong number.
+            return _ERR_BAD_ARGUMENT
+        buf_a, buf_b = a.get("buffer"), b.get("buffer")
+        if buf_a is None or buf_b is None:
+            return _ERR_BAD_ARGUMENT
+        # GPUEndTime is only defined once the buffer has retired.
+        buf_a.waitUntilCompleted()
+        buf_b.waitUntilCompleted()
+        delta = (float(buf_b.GPUEndTime()) - float(buf_a.GPUEndTime())) * 1000.0
+        _deref(out_ms).value = max(0.0, delta)
+        return _OK
+
+    # -- binding ------------------------------------------------------------
+
+    def buffer_for_pointer(self, address: int):
+        """The `MTLBuffer` backing an address this allocator handed out.
+
+        How a launch path SHOULD bind our memory. The alternative — wrapping
+        the raw address with
+        `newBufferWithBytesNoCopy:length:options:deallocator:`, which is what
+        the Triton Metal driver does — asks Metal to take over pages it
+        already owns, requires page alignment to be defined at all, and was
+        observed returning nil under memory pressure. The buffer already
+        exists; the registry can simply hand it back.
+
+        Returns `(MTLBuffer, offset)`. The offset is not always 0: the engine
+        sub-allocates. `ComponentArena` takes ONE block per component and
+        bump-allocates every weight inside it, so a weight pointer is a base
+        plus an offset, and that is the normal path for every model on every
+        backend — not an Apple special case.
+
+        This looked up the exact address until 2026-09-06, which is the same
+        as answering "no" for every weight tensor in the engine. The first
+        TinyLlama --triton run refused at the embedding kernel: the address
+        was a perfectly good arena interior, and the ownership rule said the
+        allocator had never handed it out.
+
+        Returns `(None, 0)` for an address genuinely outside every live
+        allocation — the caller decides, and is not handed something that
+        merely looks right.
+        """
+        with self._lock:
+            entry = self._buffers.get(address) or self._host_buffers.get(address)
+            if entry is not None:
+                return entry[0], 0
+            # An interior address: the greatest base at or below it, if the
+            # allocation it names is long enough to contain it.
+            index = bisect.bisect_right(self._bases, address) - 1
+            if index < 0:
+                return None, 0
+            base = self._bases[index]
+            entry = self._buffers.get(base) or self._host_buffers.get(base)
+            if entry is None:                           # pragma: no cover
+                return None, 0
+            if address >= base + entry[2]:
+                return None, 0
+            return entry[0], address - base
+
+    # -- diagnostics --------------------------------------------------------
+
+    def live_allocation_count(self) -> int:
+        """Buffers this runtime is keeping alive. Tests read it; the engine
+        does not."""
+        with self._lock:
+            return len(self._buffers) + len(self._host_buffers)
+
+
+# The entry-point name table for this implementation. It is deliberately NOT
+# a row in `_GPU_BACKENDS`: that dict holds symbol tables over one C ABI and
+# `test_device_backend_seam.py` pins it to exactly `{"cuda", "hip"}`, because
+# a key present in one row and missing from the other is a crash on that
+# hardware alone. Metal is a different API, so it gets a different table.
+#
+# Keys ABSENT here are absent on purpose — Metal has no honest equivalent —
+# and the seam turns each into an explicit refusal naming the backend:
+#   memcpy_async
+#       there is no asynchronous copy to make: the copy is a CPU memmove
+#       over memory both processors already share.
+#   malloc_async / free_async
+#       the stream-ordered allocator has no counterpart, and it is disabled
+#       on CUDA here anyway.
+#   device_can_access_peer / device_enable_peer_access
+#       one GPU. `MTLCopyAllDevices()` returns 1, and the seam answers False
+#       for peer access when the key is missing, which is the truth.
+METAL_BACKEND = {
+    "rt_libs": [],              # nothing is dlopen'd: this is not a C library
+    "malloc": "malloc", "free": "free",
+    "memcpy": "memcpy", "memset": "memset",
+    "set_device": "set_device", "get_device": "get_device",
+    "device_count": "device_count",
+    "mem_get_info": "mem_get_info",
+    "sync": "sync",
+    "malloc_host": "malloc_host", "free_host": "free_host",
+    "stream_create": "stream_create",
+    "stream_destroy": "stream_destroy",
+    "stream_sync": "stream_sync",
+    "event_create": "event_create",
+    "event_create_flags": "event_create_flags",
+    "event_destroy": "event_destroy",
+    "event_record": "event_record",
+    "event_sync": "event_sync",
+    "event_elapsed": "event_elapsed",
+    "stream_wait_event": "stream_wait_event",
+}
+
+
+_RUNTIME: Optional[MetalRuntime] = None
+_RUNTIME_LOCK = threading.Lock()
+
+
+def runtime() -> MetalRuntime:
+    """The process-wide Metal runtime. Raises if Metal is unusable.
+
+    Under NBX_CENSUS=1 the Metal device is deliberately UNREACHABLE — the
+    honest analogue of the census's `CUDA_VISIBLE_DEVICES=` on a discrete
+    card. The census forms the keys the launcher WOULD form for a profile
+    (shapes, dtypes, block sizes, all data), and must touch no device; a
+    path that reaches for the real Metal device inside a shadow is one the
+    census does not cover, and it fails LOUDLY here rather than succeeding in
+    silence on the one card an Apple host always has. `device_count()` reads
+    0 through this (its own except path), so the census gate passes; the
+    backend NAME still resolves via NBX_GPU_BACKEND=metal, which needs no
+    device (`_pin_triton_backend`).
+    """
+    import os as _os_c
+    if _os_c.environ.get("NBX_CENSUS") == "1":
+        raise RuntimeError(
+            "census shadow: the Metal device is deliberately unreachable "
+            "(NBX_CENSUS=1). Key formation is pure data and needs no device; "
+            "a path that opened the Metal runtime here is uncovered by the "
+            "census and must be made pure.")
+    global _RUNTIME
+    if _RUNTIME is None:
+        with _RUNTIME_LOCK:
+            if _RUNTIME is None:
+                _RUNTIME = MetalRuntime()
+    return _RUNTIME
+
+
+def metal_device_available() -> bool:
+    """True when a Metal device can actually be opened.
+
+    A probe, so it answers False instead of raising — but it opens the real
+    device rather than checking for the import, because a machine with the
+    bindings and no usable GPU must not be reported as ready.
+
+    Under NBX_CENSUS the device is deliberately UNREACHABLE (`runtime()`
+    refuses), but the backend must still be NAMEABLE — the exact analogue of a
+    CUDA rack where `libcudart` loads and names the backend though
+    `CUDA_VISIBLE_DEVICES=` leaves no device visible. The backend resolver
+    (`nbx_tensor` `_resolve_backend`) calls this to decide whether to name
+    metal; answering False under census made every op that resolves the
+    backend — `aten::embedding` first — fail with "No GPU runtime found",
+    though the same census completes on CUDA (the Dell's 9ea81cd2 covers the
+    rack, not Apple). So answer True here (naming) while `runtime()` stays
+    unreachable (opening): keys are pure data and need no open device.
+    """
+    import os as _os_a
+    if _os_a.environ.get("NBX_CENSUS") == "1":
+        return True
+    try:
+        runtime()
+        return True
+    except Exception:
+        return False
+
+
+def reset_runtime_for_tests() -> None:
+    """Drop the cached runtime — and everything compiled against it.
+
+    A compiled kernel captures the runtime it was built with: that runtime's
+    device, its pipelines, and its registry of which addresses the allocator
+    handed out. Replacing the runtime while compiled kernels survive leaves
+    those kernels resolving buffer addresses through a registry that never
+    issued them, and the next launch refuses a perfectly good tensor with
+    "which the Metal allocator did not hand out".
+
+    Measured 2026-09-05: this made `test_cumsum_is_correct[128]` fail in the
+    full kernels suite and pass in every smaller selection — the contaminating
+    file was `test_metal_allocator.py`, the only caller of this function.
+    Tests only, but the invariant is general: the caches and the runtime have
+    the same lifetime.
+    """
+    global _RUNTIME
+
+    # BEFORE dropping it, hand everything the old runtime owns back to the
+    # old runtime — while it is still the one a free would reach.
+    #
+    #   * the allocator's free pool parks blocks for reuse. A block parked
+    #     under the old runtime and handed out after the reset is an address
+    #     the new runtime never issued, and the first launch that touches it
+    #     refuses a perfectly good tensor.
+    #   * `_INT64_ARRAY_CACHE` holds shape/stride tensors forever by design
+    #     (a frozen replay plan records their raw device pointers), which
+    #     also means forever across a reset.
+    from nbxtensor import tensor as nbx_tensor
+    import nbxtensor.memory.allocator as _nbx_allocator
+    import nbxtensor.backend.runtime as _nbx_runtime
+    try:
+        _nbx_allocator.DeviceAllocator.empty_cache_pool()
+    except Exception:                                   # pragma: no cover
+        pass                                            # nothing to release
+    nbx_tensor._INT64_ARRAY_CACHE.clear()
+
+    # Everything issued so far belongs to the runtime about to be dropped.
+    # The flush above released what was parked; a tensor still ALIVE right
+    # now is freed later, and without this that free would park its dead
+    # address for the next allocation to be served.
+    try:
+        _nbx_allocator.DeviceAllocator.new_allocation_epoch()
+    except Exception:                                   # pragma: no cover
+        pass
+
+    with _RUNTIME_LOCK:
+        _RUNTIME = None
+
+    # Everything that captured the old runtime goes with it. Lazily and
+    # after the lock, because both of these import this module.
+    #
+    # There are three caches, and they must have ONE lifetime:
+    #   * `_RUNTIME` here,
+    #   * the compiled kernels in `metal_driver`, which hold the device and
+    #     the pipelines they were built on,
+    #   * `_nbx_runtime.gpu_runtime`, an lru_cache the allocator resolves
+    #     every malloc through.
+    #
+    # Clearing any subset is worse than clearing none: the tensors then come
+    # from one runtime and the kernels from another, and a launch refuses a
+    # perfectly good address with "which the Metal allocator did not hand
+    # out". Measured 2026-09-05 as an order-dependent failure of
+    # `test_cumsum_is_correct[128]` in the full kernels suite.
+    # The fork's driver held a kernel/library cache that pinned the device;
+    # it was archived 2026-09-17 and the cache went with it. triton-ext
+    # holds its compiled functions in Triton's own cache, which this
+    # function's callers already clear.
+
+    # The launcher caches a driver-produced handle per specialisation — a
+    # Metal pipeline here — so it holds the runtime as surely as the driver
+    # does, one layer up.
+    from nbxtensor.backend import launcher
+    launcher.reset_caches()
+
+    for cached in (getattr(_nbx_runtime, "gpu_runtime", None),
+                   getattr(_nbx_runtime, "_metal_backend_table", None)):
+        clear = getattr(cached, "cache_clear", None)
+        if clear is not None:
+            clear()