nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/memory/arena.py

@@ -0,0 +1,75 @@
+"""Triton Arena — NBXTensor slot storage for compiled execution.
+
+Memory layout: [weights | inputs | intermediates]
+All slots hold NBXTensor. Zero torch dependency.
+"""
+
+from operator import is_ as _is
+
+
+class Arena:
+    """O(1) tensor storage via integer-indexed list."""
+
+    __slots__ = ('_slots', '_num_weights', '_num_inputs', 'epoch', 'views')
+
+    def __init__(self, size: int, num_weights: int = 0, num_inputs: int = 0):
+        self._slots: list = [None] * size
+        self._num_weights = num_weights
+        self._num_inputs = num_inputs
+        # bumped whenever the weight slots may have been dropped or replaced (`clear_all`, `restore_from`): a bind
+        # memoized against an older epoch binds again (`TritonSequence.bind_weights`)
+        self.epoch = 0
+        # {tensor id: (source weight, its view)}: the views a bind derives from weights (a pre-transposed weight), one
+        # per source object so a re-bind of the same weights hands the same views back; dropped with the weights
+        self.views = {}
+
+    def __getitem__(self, idx: int):
+        return self._slots[idx]
+
+    def __setitem__(self, idx: int, value):
+        self._slots[idx] = value
+
+    def __len__(self) -> int:
+        return len(self._slots)
+
+    def restore_from(self, snapshot: list) -> None:
+        """Bulk-restore every slot from a same-length snapshot list.
+
+        C-level slice assign. The frozen-plan replay restores the full
+        arena after every replayed step; the per-slot Python loop it
+        replaces cost ~140k `__setitem__` calls per decode token on the
+        30B row (2026-08-23 host profile), most of one of the two
+        ~16 ms GPU-idle holes in the step.
+        """
+        if len(snapshot) != len(self._slots):
+            raise ValueError(
+                f"arena restore length mismatch: snapshot "
+                f"{len(snapshot)} vs arena {len(self._slots)}")
+        # The epoch moves only when the weight slots really change: a replayed step restores the very weight objects
+        # its bind put there, and an epoch bumped regardless made every next step re-bind them all (~20k weights,
+        # 24 ms a token on Qwen3-30B-A3B's expert-store decode, 2026-10-28). Identity, compared at C level.
+        nw = self._num_weights
+        same = all(map(_is, snapshot[:nw], self._slots[:nw]))
+        self._slots[:] = snapshot
+        if not same:
+            self.epoch += 1
+
+    def clear_intermediates(self):
+        """Free intermediate slots (keep weights + inputs)."""
+        start = self._num_weights + self._num_inputs
+        for i in range(start, len(self._slots)):
+            self._slots[i] = None
+
+    def clear_inputs(self) -> None:
+        """Clear input tensors (for next inference)."""
+        start = self._num_weights
+        end = self._num_weights + self._num_inputs
+        for i in range(start, end):
+            self._slots[i] = None
+
+    def clear_all(self) -> None:
+        """Clear all tensors."""
+        for i in range(len(self._slots)):
+            self._slots[i] = None
+        self.views.clear()
+        self.epoch += 1

nbxtensor/memory/device_memory.py

@@ -0,0 +1,150 @@
+"""The memory primitives of one device backend — what the allocator's account (`memory.allocator.DeviceAllocator`) asks
+a backend for, and nothing else: no byte is counted here (the account counts), no policy is decided here (the account
+decides: the pool, the rung, the OOM report).
+
+A primitive returns the backend's own status code beside its value (0 = success), as the runtime gives it: the account
+checks it and names the failure, exactly as it did when it called the runtime's symbols itself. One implementation
+serves every backend that exposes the C ABI of a CUDA-style runtime (`AbiMemory`: CUDA, HIP, and Metal's runtime
+emulation today); a backend with a native API implements `DeviceMemory` directly (Metal next, the Mac's M2).
+
+Built per call from the runtime in force (`backend.runtime.device_memory()`), never cached: a test that substitutes the
+runtime or its table is read at once. Torch-free, numpy-free.
+"""
+from __future__ import annotations
+
+import ctypes
+from typing import Tuple
+
+
+class DeviceMemory:
+    """The contract. Every method returns the backend's status code first (0 = success)."""
+
+    def supports(self, primitive: str) -> bool:                  # e.g. "host_alloc": a backend may not offer it
+        raise NotImplementedError
+
+    def malloc(self, nbytes: int) -> Tuple[int, int]:            # (rc, device pointer)
+        raise NotImplementedError
+
+    def free(self, ptr: int) -> int:                             # rc; a backend may release when its queue is done
+        raise NotImplementedError
+
+    def mem_get_info(self) -> Tuple[int, int, int]:              # (rc, free bytes, total bytes) of the current device
+        raise NotImplementedError
+
+    def memset(self, ptr: int, value: int, nbytes: int) -> int:
+        raise NotImplementedError
+
+    def memcpy(self, dst: int, src: int, nbytes: int, kind: int) -> int:   # kind: 0 H2H, 1 H2D, 2 D2H, 3 D2D
+        raise NotImplementedError
+
+    def host_alloc(self, nbytes: int) -> Tuple[int, int]:        # (rc, page-locked / shared host pointer)
+        raise NotImplementedError
+
+    def host_free(self, ptr: int) -> int:
+        raise NotImplementedError
+
+    def host_register(self, ptr: int, nbytes: int) -> int:       # rc; page-locks existing host bytes (mapped, portable)
+        raise NotImplementedError
+
+    def host_unregister(self, ptr: int) -> int:
+        raise NotImplementedError
+
+    def set_device(self, index: int) -> int:
+        raise NotImplementedError
+
+    def get_device(self) -> Tuple[int, int]:                     # (rc, current device index)
+        raise NotImplementedError
+
+    def device_count(self) -> Tuple[int, int]:                   # (rc, number of devices)
+        raise NotImplementedError
+
+    def budget(self) -> Tuple[int, int]:                         # (rc, bytes the device grants this process at most)
+        raise NotImplementedError
+
+    def map_host(self, ptr: int, nbytes: int) -> int:            # unified memory: a device-visible address of host bytes
+        raise NotImplementedError
+
+    def unmap_host(self, start: int) -> None:
+        raise NotImplementedError
+
+
+class AbiMemory(DeviceMemory):
+    """A runtime exposing the C ABI of a CUDA-style runtime, through its table of symbol names
+    (`backend.runtime._GPU_BACKENDS[...]`, or Metal's emulation table): the calls the account made itself before."""
+
+    __slots__ = ("_rt", "_table")
+
+    def __init__(self, rt, table: dict):
+        self._rt = rt
+        self._table = table
+
+    #: the contract's names that are table keys of another spelling
+    _KEYS = {"host_alloc": "malloc_host", "host_free": "free_host"}
+
+    def supports(self, primitive):
+        return self._KEYS.get(primitive, primitive) in self._table
+
+    def _fn(self, name: str):
+        return getattr(self._rt, self._table[name])
+
+    def malloc(self, nbytes):
+        p = ctypes.c_void_p()
+        rc = self._fn("malloc")(ctypes.byref(p), ctypes.c_size_t(int(nbytes)))
+        return rc, (p.value or 0)
+
+    def free(self, ptr):
+        return self._fn("free")(ctypes.c_void_p(int(ptr)))
+
+    def mem_get_info(self):
+        free_b, total_b = ctypes.c_size_t(), ctypes.c_size_t()
+        rc = getattr(self._rt, self._table.get("mem_get_info", "cudaMemGetInfo"))(ctypes.byref(free_b),
+                                                                                    ctypes.byref(total_b))
+        return rc, free_b.value, total_b.value
+
+    def memset(self, ptr, value, nbytes):
+        return self._fn("memset")(ctypes.c_void_p(int(ptr)), ctypes.c_int(int(value)), ctypes.c_size_t(int(nbytes)))
+
+    def memcpy(self, dst, src, nbytes, kind):
+        return self._fn("memcpy")(ctypes.c_void_p(int(dst)), ctypes.c_void_p(int(src)), ctypes.c_size_t(int(nbytes)),
+                                  ctypes.c_int(int(kind)))
+
+    def host_alloc(self, nbytes):
+        p = ctypes.c_void_p()
+        rc = self._fn("malloc_host")(ctypes.byref(p), ctypes.c_size_t(int(nbytes)))
+        return rc, (p.value or 0)
+
+    def host_free(self, ptr):
+        return self._fn("free_host")(ctypes.c_void_p(int(ptr)))
+
+    def host_register(self, ptr, nbytes):
+        # flags 0 (the Default flag): under unified addressing the bytes are mapped and portable, as the
+        # runtime's own page-locked allocation is
+        return self._fn("host_register")(ctypes.c_void_p(int(ptr)), ctypes.c_size_t(int(nbytes)), ctypes.c_uint(0))
+
+    def host_unregister(self, ptr):
+        return self._fn("host_unregister")(ctypes.c_void_p(int(ptr)))
+
+    def set_device(self, index):
+        return self._fn("set_device")(ctypes.c_int(int(index)))
+
+    def get_device(self):
+        d = ctypes.c_int()
+        rc = self._fn("get_device")(ctypes.byref(d))
+        return rc, d.value
+
+    def device_count(self):
+        n = ctypes.c_int()
+        rc = self._fn("device_count")(ctypes.byref(n))
+        return rc, n.value
+
+    def budget(self):
+        """The device's total, as the runtime reads it (a C-ABI runtime grants the whole device)."""
+        rc, _free, total = self.mem_get_info()
+        return rc, total
+
+    def map_host(self, ptr, nbytes):
+        """Metal's runtime emulation maps host bytes for the device (`bind_mapped`); a C-ABI runtime has no such call."""
+        return self._rt.bind_mapped(int(ptr), int(nbytes))
+
+    def unmap_host(self, start):
+        self._rt.unbind_mapped(start)

nbxtensor/memory/device_transfer.py

@@ -0,0 +1,154 @@
+"""Triton-side device utilities — device-string parsing and cross-device
+NBXTensor transfer, shared by the compiled (TritonSequence) and op-by-op
+(triton_sequential) execution paths.
+
+Pipeline-parallel placement shards a component's weights across GPUs (block N
+on cuda:i, block N+k on cuda:j). At a stage boundary the activation produced on
+the previous stage's device must be moved to the next stage's device before the
+op runs — `NBXTensor.to('cuda:N')` is intentionally a no-op (it only casts
+dtype), so the move is a real D2D `memcpy` here. R33-pure (DeviceAllocator +
+NBXTensor only; no torch). Extracted from `sequence.py:_transfer_tensor` /
+`_needs_move` so both modes share ONE implementation (R30 parity by
+construction). Model-agnostic.
+
+`parse_device_idx` / `parse_device_idxs` are the SINGLE triton-side device-
+string parsers (brick-consolidation E1). They replace six divergent
+`_parse_device_idx` copies that lived in the triton flow handlers, which
+disagreed on every non-trivial form of the Prism device string:
+
+  - the fast path `device.split(":")[-1]` returned the LAST ordinal for a
+    bare comma list ("cuda:2,cuda:3" → 3) while the compound branch of the
+    same copies returned the FIRST for the fgp-prefixed form
+    ("fgp:cuda:0,cuda:1" → 0);
+  - three copies had no compound branch at all and silently returned 0 for
+    every "fgp:"/"tp:" string (wrong device whenever the primary was not
+    cuda:0);
+  - one naive copy (`int(s.split(":")[1])`) CRASHED with ValueError on any
+    compound string.
+
+Canonical semantics here: the primary device is the FIRST CUDA ordinal named
+in the string ("fgp:cuda:1,cuda:2" → 1, "cuda:2,cuda:3" → 2), a bare integer
+string is accepted as an ordinal, and anything unparseable (bare "cuda",
+"cpu") resolves to 0 — the behaviour every flow copy converged on for the
+degenerate forms.
+"""
+
+import re
+from typing import List
+
+from nbxtensor.tensor import NBXTensor
+from nbxtensor.memory.allocator import DeviceAllocator
+
+_CUDA_ORDINAL_RE = re.compile(r"cuda:(\d+)")
+
+
+def parse_device_idxs(device) -> List[int]:
+    """All CUDA ordinals named in a device string, in order of appearance.
+
+    Handles every Prism device-string form: "cuda:2" → [2],
+    "fgp:cuda:0,cuda:1" → [0, 1], "tp:cuda:2,cuda:3" → [2, 3],
+    "cuda" / "cpu" / None → [].
+    """
+    return [int(m) for m in _CUDA_ORDINAL_RE.findall(str(device))]
+
+
+def parse_device_idx(device) -> int:
+    """Primary (FIRST) CUDA ordinal of a Prism device string.
+
+    "cuda:2" → 2, "fgp:cuda:1,cuda:2" → 1, "cuda:2,cuda:3" → 2, "3" → 3,
+    "cuda" / "cpu" / unparseable → 0. Never raises — the flow handlers feed
+    this straight into `DeviceAllocator.set_device` at setup time and every
+    historical copy resolved degenerate strings to device 0.
+    """
+    idxs = parse_device_idxs(device)
+    if idxs:
+        return idxs[0]
+    try:
+        return int(str(device).strip())
+    except (ValueError, TypeError):
+        return 0
+
+
+def needs_move(t: NBXTensor, target_dev: int) -> bool:
+    """True iff `t` must be transferred to land on cuda:target_dev.
+
+    1. CPU source (zero3 offload) → must H2D (even if `_device_idx` is 0).
+    2. Different CUDA device → must D2D.
+    3. Already on target → no-op.
+    """
+    if getattr(t, "_device", "cuda") == "cpu":
+        return True
+    return t._device_idx != target_dev
+
+
+def transfer_tensor(tensor: NBXTensor, target_dev: int) -> NBXTensor:
+    """Copy an NBXTensor to cuda:target_dev, preserving shape/strides/dtype.
+
+    kind=1 (H2D) for a CPU source (zero3), kind=3 (D2D) for a cross-GPU move.
+
+    Stride handling — dense-window rule (P-TRITON-MLA root fix): the flat
+    ``memcpy(tensor._nbytes)`` from ``data_ptr()`` is only meaningful when
+    the view's elements tile exactly one dense span from data_ptr()
+    (`NBXTensor.spans_densely`, the one definition — it also refuses an
+    overlapping view whose span happens to equal its numel). For such views (contiguous tensors, the zero3
+    pre-transposed weight ``.t()``, dim-0 narrows) the strides are carried
+    over unchanged — the historical contract, preserved byte-for-byte.
+    Every OTHER view — interior narrow/slice (span > numel, offset
+    window), step>1 de-interleave slices, broadcast/expand (stride-0,
+    span < numel) — is materialised via ``.contiguous()`` on the SOURCE
+    device first. The old code only materialised the expand case; an
+    interior/strided view was flat-copied and its original strides
+    re-attached over an nbytes-sized allocation — the strides address a
+    window FAR LARGER than the allocation, so the first downstream kernel
+    reading the view walks past the end of the new buffer — async illegal
+    memory access (error 700) surfacing a few launches later. Proven at
+    the DeepSeek-Coder-V2-Lite pipeline_parallel boundary (block.13 MLA
+    rope q/k de-interleave: a [1,16,23,64] narrow-of-transpose view,
+    strides addressing a 70528-element span over a 23552-element
+    allocation → 2/3 of reads OOB → strided_copy fault surfacing at
+    aten.cat::91, BOTH triton modes — this helper is shared by
+    _run_multi_device and the triton_sequential per-op transfer block).
+    """
+    if not tensor.spans_densely():
+        tensor = tensor.contiguous()
+    src_device = getattr(tensor, "_device", "cuda")
+    kind = 1 if src_device == "cpu" else 3
+    if kind == 3:
+        # D2D read barrier: the memcpy runs on the TARGET device's
+        # legacy stream, which does NOT wait for the SOURCE device's
+        # triton (non-default) stream — reading a tensor a kernel is
+        # still writing violates the "never read a weight still being
+        # copied" contract the zero3 H2D path enforces with events.
+        # Sync the source device before the peer copy (this is the
+        # declared slow path; an event-based overlap is the perf-layer
+        # upgrade). Proven on the Qwen3-Omni audio triton leg:
+        # deterministic error-700 poison at step-2 entry on the
+        # block-scatter placement, gone under CUDA_LAUNCH_BLOCKING,
+        # allocator ledger clean, flow tail fully quiesced — only the
+        # run-entry transfer region remained.
+        _src_idx = getattr(tensor, "_device_idx", None)
+        if _src_idx is not None:
+            # Let the copy use the direct card-to-card link. Without this the
+            # driver stages every peer copy through HOST memory — it never
+            # touches NVLink/xGMI whatever the topology reports. Memoised, so
+            # this is one dict lookup per hand-off after the first.
+            # Measured 2026-09-03: 8.7x on a 34 MB stage boundary
+            # (48 ms -> 5.5 ms), 2.4x at 4 KB.
+            DeviceAllocator.ensure_peer_access(_src_idx, target_dev)
+            _prev = DeviceAllocator.get_device()
+            if _prev != _src_idx:
+                DeviceAllocator.set_device(_src_idx)
+                DeviceAllocator.sync_device()
+                DeviceAllocator.set_device(_prev)
+            else:
+                DeviceAllocator.sync_device()
+    DeviceAllocator.set_device(target_dev)
+    if tensor._nbytes > 0:
+        dst_raw_ptr = DeviceAllocator.malloc_cuda(tensor._nbytes)
+        DeviceAllocator.memcpy(dst_raw_ptr, tensor.data_ptr(),
+                               tensor._nbytes, kind=kind)
+    else:
+        dst_raw_ptr = 0
+    return NBXTensor(
+        dst_raw_ptr, tensor._shape, tensor._strides, tensor._dtype,
+        "cuda", owns_data=True, device_idx=target_dev, offset=0)

nbxtensor/memory/layout.py

@@ -0,0 +1,43 @@
+"""Where tensors placed one after another in one block land, and how many bytes the block takes: the layout every arena
+of the library allocates (`memory_pool.ComponentArena`, `memory_pool.HostPinnedArena`, an expert store's slot arena),
+and the very figure a planner prices an arena at — one function, so the plan and the allocation cannot differ (M3,
+docs/nbxtensor/COST.md section 2). Pure arithmetic: no device, torch-free, numpy-free.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable, Sequence, Tuple
+
+from nbxtensor.dtype.itemsize import itemsize
+from nbxtensor.shape import numel
+
+#: every tensor of an arena starts on this boundary (the widest vector load of any backend; CUDA's allocation granule
+#: for a sub-allocation)
+ARENA_ALIGNMENT = 256
+
+
+def aligned(nbytes: int, alignment: int = ARENA_ALIGNMENT) -> int:
+    """`nbytes` rounded up to `alignment`."""
+    return -(-int(nbytes) // alignment) * alignment
+
+
+def tensor_nbytes(shape: Sequence[int], dtype) -> int:
+    """The bytes a dense tensor of `shape` holds at `dtype` (a name, a torch spelling or an NBXDtype, as `itemsize`)."""
+    return numel(tuple(int(d) for d in shape)) * itemsize(dtype)
+
+
+@dataclass(frozen=True)
+class ArenaLayout:
+    """`total_bytes`: the block; `offsets`: where each tensor starts in it, in placement order."""
+    total_bytes: int
+    offsets: Tuple[int, ...]
+
+
+def arena_layout(sizes: Iterable[int], alignment: int = ARENA_ALIGNMENT) -> ArenaLayout:
+    """The layout of tensors of `sizes` bytes (placement order) in one arena: each at the next `alignment` boundary, a
+    zero-byte tensor taking no room — the bump allocation `ComponentArena.alloc` performs."""
+    offsets, at = [], 0
+    for n in sizes:
+        offsets.append(at)
+        at += aligned(n, alignment) if int(n) else 0
+    return ArenaLayout(at, tuple(offsets))

nbxtensor/memory/memory_pool.py

@@ -0,0 +1,191 @@
+"""Triton Memory Pool — arena-based GPU memory allocator.
+
+Pre-allocates one contiguous block per component, sub-allocates inside
+via bump pointer. No individual free, no fragmentation.
+
+Zero torch. Uses DeviceAllocator (hardware-agnostic: CUDA/ROCm).
+"""
+
+import gc
+
+import numpy as np
+
+from nbxtensor.memory.allocator import DeviceAllocator
+from nbxtensor.memory.layout import ARENA_ALIGNMENT, aligned
+from nbxtensor.memory.device_transfer import parse_device_idxs
+
+
+def release_flow_memory(device=None) -> None:
+    """Flow-boundary memory release — triton-side mirror of the compiled
+    `MemoryManager.release_flow_memory` discipline (R30 naming parity,
+    separate implementation: R33-pure, DeviceAllocator only, no torch).
+
+    Consolidates the bare `gc.collect()` calls the triton flow handlers
+    ran at stage boundaries. Those calls skipped the sync-before-free
+    protection: `gc.collect()` triggers NBXTensor / ComponentArena
+    finalizers (`cudaFree`) on cyclic or lingering references, and freeing
+    memory an in-flight kernel is still reading is the err-700 class
+    documented in `core/memory/manager.py`. Ordering here:
+
+      1. sync every CUDA device named in `device` (the flow's Prism
+         primary-device string — may be compound, e.g. "fgp:cuda:0,cuda:1")
+         BEFORE any reference can be finalized;
+      2. gc.collect() — drop refs, run finalizers;
+      3. DeviceAllocator.empty_cache_pool() — return pool-cached blocks to
+         the driver (no-op under NBX_ALLOC_POOL=0; individual frees then
+         already went straight to cudaFree).
+    """
+    idxs = parse_device_idxs(device)
+    if idxs:
+        prev = DeviceAllocator.get_device()
+        for i in idxs:
+            DeviceAllocator.set_device(i)
+            DeviceAllocator.sync_device()
+        DeviceAllocator.set_device(prev)
+    gc.collect()
+    DeviceAllocator.empty_cache_pool()
+
+
+class ComponentArena:
+    """Single contiguous GPU memory block for one component's weights.
+
+    Allocates one big block via DeviceAllocator.malloc_cuda, then
+    sub-allocates inside via bump pointer with 256-byte alignment.
+
+    Usage:
+        arena = ComponentArena(total_bytes, device_idx)
+        ptr1 = arena.alloc(1024)   # sub-allocate 1024 bytes
+        ptr2 = arena.alloc(4096)   # sub-allocate 4096 bytes
+        arena.free()               # release entire block at once
+    """
+
+    ALIGNMENT = ARENA_ALIGNMENT  # the library's one arena layout (memory.layout)
+
+    def __init__(self, total_bytes: int, device_idx: int):
+        # Initialize _base_ptr BEFORE the malloc call so __del__ is safe
+        # even if malloc_cuda raises (OOM on one device shouldn't leave
+        # a half-constructed object whose __del__ then crashes on
+        # attribute access). This closes a secondary issue flagged
+        # during the zero3 leak investigation (see tests/scratch/
+        # zero3_leak_investigation/REPORT.md §0.5).
+        self.device_idx = device_idx
+        self._base_ptr = 0
+        self._total = self._align(total_bytes)
+        self._offset = 0
+
+        DeviceAllocator.set_device(device_idx)
+        # never pooled: once freed, its bytes go back to the driver whole (DeviceAllocator.malloc_cuda `pooled`)
+        self._base_ptr = DeviceAllocator.malloc_cuda(self._total, pooled=False)
+        if self._base_ptr == 0 and self._total > 0:
+            raise RuntimeError(
+                f"ComponentArena: failed to allocate {self._total / 1e9:.2f} GB "
+                f"on device {device_idx}")
+
+    def alloc(self, nbytes: int) -> int:
+        """Sub-allocate nbytes from the arena. Returns GPU pointer."""
+        if nbytes == 0:
+            return 0
+        aligned = self._align(nbytes)
+        if self._offset + aligned > self._total:
+            raise RuntimeError(
+                f"ComponentArena: out of memory. "
+                f"Requested {nbytes} bytes, "
+                f"used {self._offset}/{self._total} bytes "
+                f"on device {self.device_idx}")
+        ptr = self._base_ptr + self._offset
+        self._offset += aligned
+        return ptr
+
+    def free(self):
+        """Release the entire arena block."""
+        if self._base_ptr:
+            DeviceAllocator.set_device(self.device_idx)
+            DeviceAllocator.free_cuda(self._base_ptr)
+            self._base_ptr = 0
+            self._offset = 0
+            self._total = 0
+
+    def reset(self):
+        """Reset bump pointer without freeing GPU memory.
+
+        Allows reuse of the same arena for a new component load
+        (e.g., lifecycle reload in lazy_sequential strategy).
+        """
+        self._offset = 0
+
+    @property
+    def used(self) -> int:
+        return self._offset
+
+    @property
+    def capacity(self) -> int:
+        return self._total
+
+    def _align(self, n: int) -> int:
+        return aligned(n, self.ALIGNMENT)
+
+    def __del__(self):
+        # Guard: __init__ may raise before _base_ptr is set (rare but
+        # happens when malloc_cuda throws OOM immediately). Let __del__
+        # remain a no-op in that case.
+        if getattr(self, '_base_ptr', 0):
+            self.free()
+
+
+class HostPinnedArena:
+    """One page-locked host block for a component's host-held weights (zero3's block weights, an expert store's
+    routed experts), sub-allocated like `ComponentArena`. A pinned allocation per weight cost Qwen3-30B-A3B 18 864
+    page-locking calls — about 40 % of a 447 s load (py-spy, V100 rack, 2026-10-27); one block costs one. Each
+    weight holds the arena as its `base`, so the block is released when the last of them is."""
+
+    ALIGNMENT = ARENA_ALIGNMENT
+
+    def __init__(self, total_bytes: int):
+        """On a backend that page-locks existing memory (`host_register`), the block is ordinary host memory, written
+        by the load and page-locked by `pin()` in one call before any transfer reads it: a page-locked allocation
+        ran at 1.2 GiB/s on one thread, its threads serialised by the runtime (16 GiB in 13.0 s; 4 or 8 threads,
+        16 s), while 16 GiB written by 16 threads and then registered took 4.9 s (V100 rack, 2026-10-30) — and the
+        load's own writes are that touch. Elsewhere the block is the backend's page-locked allocation, as before."""
+        self._base_ptr = 0
+        self._total = aligned(total_bytes, self.ALIGNMENT)
+        self._offset = 0
+        self._pageable = None          # the host bytes until `pin()` page-locks them (a register-capable backend)
+        self._registered = False
+        if self._total:
+            from nbxtensor.backend import runtime as _rt
+            if _rt.device_memory().supports("host_register"):
+                self._pageable = np.empty(self._total + self.ALIGNMENT, dtype=np.uint8)
+                self._base_ptr = aligned(self._pageable.ctypes.data, self.ALIGNMENT)
+            else:
+                self._base_ptr = DeviceAllocator.malloc_host_pinned(self._total)
+
+    def pin(self) -> None:
+        """Page-lock the block (once; a no-op where it was allocated page-locked). Called by the loader after the
+        last weight is written and before any transfer reads one."""
+        if self._pageable is not None and not self._registered:
+            DeviceAllocator.register_host_pinned(self._base_ptr, self._total)
+            self._registered = True
+
+    def alloc(self, nbytes: int) -> int:
+        if nbytes == 0:
+            return 0
+        size = aligned(nbytes, self.ALIGNMENT)
+        if self._offset + size > self._total:
+            raise RuntimeError(f"HostPinnedArena: {nbytes} bytes asked, {self._offset}/{self._total} used — the "
+                               f"sizing pass and the load disagree on a host weight's bytes")
+        ptr = self._base_ptr + self._offset
+        self._offset += size
+        return ptr
+
+    @property
+    def used(self) -> int:
+        return self._offset
+
+    def __del__(self):
+        if getattr(self, "_base_ptr", 0):
+            if getattr(self, "_pageable", None) is None:
+                DeviceAllocator.free_host_pinned(self._base_ptr)
+            elif self._registered:
+                DeviceAllocator.unregister_host_pinned(self._base_ptr)
+            self._pageable = None
+            self._base_ptr = 0