nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/memory/slab.py

@@ -0,0 +1,273 @@
+"""Slab allocators: a run's device allocations carved out of one block per device, and the interceptor that serves
+them (moved from the engine's frozen-plan replayer, M2).
+
+A replay recording (the engine's `triton/replay.py`) runs a step once with every allocation it makes served from a slab,
+so that the addresses its launches recorded stay valid when they are replayed: the replay issues the same launches in
+the same order, so every reuse the recorded run made is safe at replay. `SlabInterceptor` stands in front of the
+allocator (`DeviceAllocator.install_interceptor`) and routes: while recording, to the slab of the requested device;
+between recordings, to the recycled per-step slab when one is in force; else to the allocator itself, measuring the
+watermark a slab will be sized from. Torch-free.
+"""
+from __future__ import annotations
+
+from typing import Dict, List, Optional, Tuple
+
+from nbxtensor.memory.allocator import AllocatorInterceptor, DeviceAllocator
+
+ALIGN = 512       # matches typical device allocation granularity
+
+# Slabs with outstanding carved allocations — the free path consults this PERMANENTLY (not only while recording): a
+# slab-carved tensor may die on any later run when its arena slot is overwritten, and its finalizer must return the
+# range to the slab, never free a slab-interior pointer to the driver.
+ACTIVE_SLABS: List["SlabAllocator"] = []
+
+
+def _carve(free_list: List[Tuple[int, int]], need: int) -> Optional[int]:
+    """First fit: the offset of the first free range holding `need` aligned bytes, the range shrunk; None if none."""
+    for i, (off, sz) in enumerate(free_list):
+        if sz >= need:
+            if sz == need:
+                free_list.pop(i)
+            else:
+                free_list[i] = (off + need, sz - need)
+            return off
+    return None
+
+
+def _give_back(free_list: List[Tuple[int, int]], off: int, sz: int) -> List[Tuple[int, int]]:
+    """The free list with [off, off + sz) returned, sorted and its neighbours merged."""
+    free_list.append((off, sz))
+    free_list.sort()
+    merged: List[Tuple[int, int]] = []
+    for o, s_ in free_list:
+        if merged and merged[-1][0] + merged[-1][1] == o:
+            merged[-1] = (merged[-1][0], merged[-1][1] + s_)
+        else:
+            merged.append((o, s_))
+    return merged
+
+
+def first_fit_extent(events: List[Tuple[int, int, int]]) -> int:
+    """The slab a run needs: the highest offset `SlabAllocator`'s first fit reaches when it serves `events` — the run's
+    pooled allocations and frees in order, (+1, key, aligned bytes) and (-1, key, 0) — from a slab without end. It is
+    at least the live peak and at most the sum: a first fit leaves holes a later, larger block cannot use. Sized from
+    the live peak alone, Allegro-TI2V's encoder tile (9.0 GB live) was short by 22.6 GB at its recording, the retry
+    then took 31.7 GB of a 32 GB card for a run that never holds more than 8.6 GB (2026-11-04)."""
+    end = 1 << 62
+    free_list: List[Tuple[int, int]] = [(0, end)]
+    at: Dict[int, Tuple[int, int]] = {}
+    top = 0
+    for sign, key, need in events:
+        if sign > 0:
+            off = _carve(free_list, need)
+            at[key] = (off, need)
+            top = max(top, off + need)
+        elif key in at:
+            free_list = _give_back(free_list, *at.pop(key))
+    return top
+
+
+class SlabAllocator:
+    """First-fit free-list allocator over ONE device slab.
+
+    The recording run's malloc/free sequence drives carving/reuse; the
+    replay issues the same launches in the same order, so every reuse
+    the recorded run performed is serialization-safe at replay.
+    """
+
+    def __init__(self, nbytes: int, dev_idx: Optional[int]) -> None:
+        self.size = int(nbytes)
+        self.dev_idx = dev_idx
+        self.base = DeviceAllocator.malloc_direct(self.size, dev_idx)
+        self.free_list: List[Tuple[int, int]] = [(0, self.size)]
+        self.allocs: Dict[int, Tuple[int, int]] = {}  # ptr -> (off, sz)
+        self.shortfall = 0  # aligned bytes we could not serve
+        # A retired slab (failed recording attempt / invalidated plan)
+        # must NOT cudaFree while tensors carved from it are alive:
+        # their finalizers free slab-interior pointers, which the seam
+        # routes back here. Actual release happens when the last
+        # outstanding alloc returns. Probe 2 proved the hazard: an
+        # eager release left the text-encoder outputs dangling
+        # ("Pointer argument cannot be accessed" + invalid cudaFree of
+        # a slab-interior pointer).
+        self.retired = False
+        ACTIVE_SLABS.append(self)
+
+    def malloc(self, nbytes: int) -> int:
+        need = (int(nbytes) + ALIGN - 1) // ALIGN * ALIGN
+        off = _carve(self.free_list, need)
+        if off is not None:
+            ptr = self.base + off
+            self.allocs[ptr] = (off, need)
+            return ptr
+        self.shortfall += need
+        return 0
+
+    def free(self, ptr: int) -> bool:
+        entry = self.allocs.pop(ptr, None)
+        if entry is None:
+            return False
+        self.free_list = _give_back(self.free_list, *entry)
+        if self.retired and not self.allocs:
+            self._release_now()
+        return True
+
+    def retire(self) -> None:
+        """Deferred release: cudaFree only after the last outstanding
+        carved allocation has been returned by its owner's finalizer."""
+        self.retired = True
+        if not self.allocs:
+            self._release_now()
+
+    def _release_now(self) -> None:
+        if self.base:
+            DeviceAllocator.free_direct(self.base)
+            self.base = 0
+        if self in ACTIVE_SLABS:
+            ACTIVE_SLABS.remove(self)
+
+
+class StepSlab(SlabAllocator):
+    """Per-sequence slab RECYCLED at every decode step (P-REPLAY-KV-
+    DECODE B1). reset() returns the whole range to the free list so
+    the next step's allocations carve IDENTICAL addresses (the decode
+    graph's allocation order is structurally constant — census-
+    proven); allocations from previous steps become 'stale' and their
+    later finalizer frees are absorbed silently (the range is already
+    recycled — safe because each step's consumers run before the next
+    step begins in the synchronous decode timeline)."""
+
+    def __init__(self, nbytes: int, dev_idx: Optional[int]) -> None:
+        super().__init__(nbytes, dev_idx)
+        self.stale: set = set()
+        self.carve_sizes: List[int] = []   # diagnosis: per-step size seq
+        self.prev_carve: Optional[List[int]] = None
+
+    def malloc(self, nbytes: int) -> int:
+        """PURE BUMP carving — mid-step frees are absorbed without
+        recycling (see free below), so placement is a pure function of
+        the allocation-size sequence (measured identical step-to-step)
+        and IMMUNE to free-timing nondeterminism (GC-cycle finalizers
+        fire at variable points; first-fit reuse turned that into
+        per-step placement churn — measured: identical size sequences,
+        differing offsets)."""
+        self.carve_sizes.append(int(nbytes))
+        need = (int(nbytes) + ALIGN - 1) // ALIGN * ALIGN
+        off, sz = self.free_list[0]
+        if sz < need:
+            self.shortfall += need
+            return 0
+        self.free_list[0] = (off + need, sz - need)
+        ptr = self.base + off
+        self.allocs[ptr] = (off, need)
+        return ptr
+
+    def reset(self) -> None:
+        self.prev_carve = self.carve_sizes
+        self.carve_sizes = []
+        self.stale.update(self.allocs.keys())
+        self.allocs.clear()
+        self.free_list = [(0, self.size)]
+
+    def free(self, ptr: int) -> bool:
+        """Absorb without recycling: the range stays consumed until the
+        next reset (bump purity beats intra-step reuse)."""
+        if ptr in self.stale:
+            self.stale.discard(ptr)
+            return True
+        if ptr in self.allocs:
+            self.allocs.pop(ptr)
+            return True
+        return False
+
+
+class SlabInterceptor(AllocatorInterceptor):
+    """The allocator's interceptor for slab-served runs. Its state is read and set by the replayer that owns it:
+
+    * `recording` with `slabs` (one per device) or `slab`: every pooled allocation is carved from the slab of the
+      REQUESTED device (serving a device-1 tensor from device 0's slab yields a pointer into the wrong card's memory —
+      a wrong result or an illegal access, never a clear error); a slab that cannot serve it, or a device without
+      one, refuses the recording through `slab_refused(reason)` and the allocator serves it.
+    * `stab_slab`: the recycled per-step slab in force; an allocation over it is served by the allocator (address
+      stability degrades for that tensor only).
+    * `measuring`: the allocator serves, and the aligned bytes are counted per pointer and device (`live`,
+      `high_water`, `live_by_dev`, `high_water_by_dev`, `total_measured`) — what a slab is sized from.
+
+    A block outside the pool (`pooled=False`, a weight arena) is never a slab's. A free returns to the slab that
+    carved the pointer, if any, before the allocator sees it."""
+
+    def __init__(self) -> None:
+        self.measuring = False
+        self.recording = False
+        self.slab: Optional[SlabAllocator] = None
+        self.slabs: Dict[int, SlabAllocator] = {}
+        self.stab_slab: Optional["StepSlab"] = None
+        # Per-device measurement: one high-water mark cannot size slabs for a run allocating on several cards without
+        # sizing each to the SUM.
+        self.live_by_dev: Dict[int, int] = {}
+        self.high_water_by_dev: Dict[int, int] = {}
+        self.measured_dev: Dict[int, int] = {}
+        # measuring counters (aligned bytes)
+        self.live = 0
+        self.high_water = 0
+        self.measured: Dict[int, int] = {}  # ptr -> aligned size
+        self.total_measured = 0             # cumulative (never decremented)
+        # the run's pooled allocations and frees in order, per device (`first_fit_extent`: the slab they need)
+        self.events_by_dev: Dict[int, List[Tuple[int, int, int]]] = {}
+        self._seq = 0
+        self._key_of_ptr: Dict[int, int] = {}
+
+    def slab_refused(self, reason: str) -> None:
+        """A recording the slabs cannot serve: the owner decides (a replayer abandons the plan, by name)."""
+        raise NotImplementedError(f"{type(self).__qualname__} does not say what a refused slab does: {reason}")
+
+    def malloc(self, direct, nbytes, dev_idx, pooled):
+        if not pooled:
+            return direct(nbytes, dev_idx, False)
+        if self.recording and (self.slab is not None or self.slabs):
+            _dev = dev_idx if dev_idx is not None else DeviceAllocator.get_device()
+            _slab = self.slabs.get(_dev, self.slab)
+            if _slab is not None:
+                ptr = _slab.malloc(nbytes)
+                if ptr:
+                    return ptr
+                self.slab_refused(f"slab exhausted on device {_dev} (short {_slab.shortfall} B)")
+            else:
+                self.slab_refused(f"no slab recorded for device {_dev}")
+            return direct(nbytes, dev_idx, True)
+        if self.stab_slab is not None:
+            ptr = self.stab_slab.malloc(nbytes)
+            if ptr:
+                return ptr
+        ptr = direct(nbytes, dev_idx, True)
+        if self.measuring and ptr:
+            need = (int(nbytes) + ALIGN - 1) // ALIGN * ALIGN
+            self.measured[ptr] = need
+            self.total_measured += need
+            self.live += need
+            if self.live > self.high_water:
+                self.high_water = self.live
+            _d = dev_idx if dev_idx is not None else DeviceAllocator.get_device()
+            self.measured_dev[ptr] = _d
+            self._seq += 1                    # a key per allocation: a pointer the allocator reuses is a new block
+            self._key_of_ptr[ptr] = self._seq
+            self.events_by_dev.setdefault(_d, []).append((1, self._seq, need))
+            self.live_by_dev[_d] = self.live_by_dev.get(_d, 0) + need
+            if self.live_by_dev[_d] > self.high_water_by_dev.get(_d, 0):
+                self.high_water_by_dev[_d] = self.live_by_dev[_d]
+        return ptr
+
+    def free(self, direct, ptr):
+        for slab in ACTIVE_SLABS:
+            if slab.free(ptr):
+                return  # returned to its slab (maybe releasing it)
+        if self.measuring:
+            need = self.measured.pop(ptr, 0)
+            self.live -= need
+            _d = self.measured_dev.pop(ptr, None)
+            if _d is not None:
+                self.live_by_dev[_d] = self.live_by_dev.get(_d, 0) - need
+                _k = self._key_of_ptr.pop(ptr, None)
+                if _k is not None:
+                    self.events_by_dev.setdefault(_d, []).append((-1, _k, 0))
+        return direct(ptr)

nbxtensor/memory/trace_replay.py

@@ -0,0 +1,122 @@
+"""The expandable range's chunk, measured by replaying a REAL run's allocations (the supervisor 2026-11-02 17:15).
+
+A run's mapping cost is what its refusals' reclaims map again (SANA-Video_2B_720p on a 16 GB V100: 89 GB mapped over a
+run at a 2 MiB chunk — registres/benchmarks/vmm_v100.md), and only the run's own allocation order produces those
+refusals; a synthetic churn does not (it always favours the largest chunk). So the library carries recorded traces as
+data (`traces/*.tsv`: every allocation and free of one gate cell, in order, by id and bytes — no model, no pointer),
+and the chunk probe replays one through THIS allocator, on the device, under the card size the trace was recorded on
+(`card_mb`, the imposed limit — the trace presses that class on any larger card), once with the expandable range off
+(the driver segments: R16 OFF) and once per candidate chunk, each in its own process (a clean allocator).
+
+The chunk is the smallest candidate whose replay costs no more than R16 OFF's replay plus OFF's own spread (max -
+median over the repetitions): the expandable range then holds what it saves (the 32g replay: 8 GB less at its peak)
+without paying for it in time. None when no candidate does — said, and the caller keeps the driver segments.
+
+Torch-free, numpy-free."""
+from __future__ import annotations
+
+import json
+import os
+import statistics
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Dict, List, Optional, Sequence, Tuple
+
+TRACES_DIR = Path(__file__).parent / "traces"
+_MARK = "NBX-TRACE-REPLAY "      # the child's result line, whatever else it prints
+
+
+def traces() -> List[Path]:
+    """The traces the library carries, by name."""
+    return sorted(TRACES_DIR.glob("*.tsv"))
+
+
+def load_trace(path) -> Tuple[Dict[str, str], List[Tuple[str, int, int]]]:
+    """(the trace's header facts `# key: value`, its events (kind, id, nbytes)). A trace without `card_mb` or without
+    one event is refused by name."""
+    meta: Dict[str, str] = {}
+    events: List[Tuple[str, int, int]] = []
+    for line in open(path):
+        if line.startswith("#"):
+            key, sep, value = line[1:].strip().partition(": ")
+            if sep and " " not in key:
+                meta[key] = value.strip()
+            continue
+        parts = line.rstrip("\n").split("\t")
+        if len(parts) == 3 and parts[0] in ("M", "A", "F"):
+            events.append((parts[0], int(parts[1]), int(parts[2])))
+    if "card_mb" not in meta or not events:
+        raise RuntimeError(f"ZERO FALLBACK: the allocation trace {path} states no card_mb or holds no event")
+    return meta, events
+
+
+def replay_trace(events: Sequence[Tuple[str, int, int]], device: int, chunk: int, limit_bytes: int) -> Dict[str, float]:
+    """Replay `events` through this process's allocator on `device` under `limit_bytes`, the large class in an
+    expandable range of `chunk` bytes (0: the driver segments). The wall seconds, the mapping seconds, the reclaims and
+    the peak held."""
+    from nbxtensor.memory.allocator import DeviceAllocator as DA
+    DA.set_device(device)
+    DA.use_expandable(device, chunk, range_bytes=limit_bytes)
+    DA.set_limit(limit_bytes, device)
+    live: Dict[int, int] = {}
+    t0 = time.perf_counter()
+    for kind, key, nbytes in events:
+        if kind == "F":
+            ptr = live.pop(key, None)
+            if ptr is not None:
+                DA.free_cuda(ptr)
+        elif nbytes:
+            live[key] = DA.malloc_cuda(nbytes, device, pooled=(kind == "M"))
+    DA.device_synchronize(device)
+    seconds = time.perf_counter() - t0
+    st = DA._POOL.stats
+    return {"seconds": seconds, "map_seconds": st.get("range_map_us", 0) / 1e6,
+            "reclaims": st.get("range_reclaims", 0), "peak_bytes": DA._cuda_peak_bytes.get(device, 0)}
+
+
+def _child(trace: str, device: int, chunk: int) -> Dict[str, float]:
+    """One replay in a fresh process (no allocator state carried between candidates)."""
+    code = ("import json, sys; from nbxtensor.memory import trace_replay as T; "
+            "m, ev = T.load_trace(sys.argv[1]); "
+            "print('" + _MARK + "' + json.dumps(T.replay_trace(ev, int(sys.argv[2]), int(sys.argv[3]), "
+            "int(m['card_mb']) << 20)))")
+    # the replay measures the allocator, never its diagnostics: the parent's run-level switches stay out of it
+    env = {k: v for k, v in os.environ.items() if not k.startswith(("NBX_ALLOC_STATS", "NBX_MALLOC_TRACE"))}
+    r = subprocess.run([sys.executable, "-c", code, trace, str(device), str(chunk)], capture_output=True, text=True,
+                       env=env)
+    line = next((l for l in r.stdout.splitlines() if l.startswith(_MARK)), None)
+    if r.returncode != 0 or line is None:
+        raise RuntimeError(f"the allocation-trace replay of {Path(trace).name} at chunk {chunk} failed: "
+                           f"{(r.stderr or r.stdout)[-400:]}")
+    return json.loads(line[len(_MARK):])
+
+
+def chunk_by_replay(device: int, granularity: int, ks: Sequence[int], reps: int,
+                    trace: Optional[str] = None) -> Tuple[Optional[int], Dict[int, List[float]]]:
+    """(the chunk, every candidate's replay seconds — 0 is R16 OFF). The smallest granularity x k whose median replay
+    costs no more than OFF's median plus OFF's spread; None when none does. `trace`: one of `traces()` (the only one
+    the library carries when it carries one; several = the caller names it)."""
+    if trace is None:
+        carried = traces()
+        if len(carried) != 1:
+            raise RuntimeError(f"ZERO FALLBACK: the library carries {len(carried)} allocation traces "
+                               f"({', '.join(p.name for p in carried)}): the chunk probe names the one it replays")
+        trace = str(carried[0])
+    # every replay runs in a child sized at the recorded card: what THIS process caches on the device is given back
+    # first, or the child finds the card short by it (a test process's pools after the kernels suite: 16 GB, the
+    # replay out of memory at chunk 0 — order-dependent, 2026-11-03)
+    from nbxtensor.memory.allocator import DeviceAllocator as DA
+    DA.device_synchronize(device)
+    if DA._range_vmm(device) is not None:
+        DA._range_reclaim(device, None)
+    DA._pool_flush(device)
+    times: Dict[int, List[float]] = {0: [_child(trace, device, 0)["seconds"] for _ in range(reps)]}
+    off = times[0]
+    bound = statistics.median(off) + (max(off) - statistics.median(off))
+    for chunk in [granularity * int(k) for k in sorted(ks)]:     # ascending: the first within OFF's noise is the chunk
+        times[chunk] = [_child(trace, device, chunk)["seconds"] for _ in range(reps)]
+        if statistics.median(times[chunk]) <= bound:
+            return chunk, times
+    return None, times