nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/dtype/policy.py

@@ -0,0 +1,155 @@
+"""Precision by policy (DESIGN rule 12): an explicit, immutable object that says in which dtype each op computes.
+
+Modelled on JAX's jmp (`Policy(param_dtype, compute_dtype, output_dtype)`, `get_policy("params=float32,
+compute=float16,output=float32")`), with the op classes of PyTorch's autocast as its default (`policy.yml`): a
+policy is plain data — three dtype names, a flag, and the op classes as a mapping of names — handed to the calls that
+use it. No hidden state, no context manager, no global switch: a call computes in the dtype the policy it was given
+says (rule 11: what a Rust core would take as a struct).
+
+- `param`: the dtype weights are kept in; `compute`: the dtype the "lower" class computes in; `output`: the dtype a
+  caller's results are returned in.
+- `compute_dtype(op, input_dtypes)`: the dtype op `op` computes in, from its class: `float32` (whatever the compute
+  dtype; or the compute dtype for a `float32_half_io` op when `float32_ops_half_io` is set), `lower` (the compute
+  dtype), `promote` (the widest floating input), none (its inputs promoted among themselves).
+- `fill_value(op, value, dtype)`: a scalar fill of a `scalar_fill` / `creation_fill` op clamped to `dtype`'s finite
+  range (a mask sentinel written for bfloat16 does not overflow float16).
+
+An application builds the policy from what it knows (NeuroBrix: its DtypeEngine, from the hardware profile and the
+container) and passes it; the library knows no model and no profile.
+"""
+from __future__ import annotations
+
+import dataclasses
+import functools
+from pathlib import Path
+from typing import Dict, FrozenSet, Mapping, Optional, Sequence, Tuple
+
+from nbxtensor.dtype.itemsize import canonical_name, finite_max, kind
+from nbxtensor.dtype.promotion import binary_result
+
+_FILE = Path(__file__).resolve().parent / "policy.yml"
+CLASSES = ("float32", "lower", "promote", "float32_half_io", "scalar_fill", "creation_fill")
+
+__all__ = ["Policy", "default_classes", "get_policy", "CLASSES"]
+
+
+@functools.lru_cache(maxsize=1)
+def _classes_from_file() -> Tuple[Tuple[str, FrozenSet[str]], ...]:
+    import yaml
+    data = (yaml.safe_load(_FILE.read_text()) or {}).get("classes") or {}
+    unknown = sorted(set(data) - set(CLASSES))
+    missing = sorted(set(CLASSES) - set(data))
+    if unknown or missing:
+        raise ValueError(f"{_FILE}: classes {unknown} unknown, {missing} missing (known: {', '.join(CLASSES)})")
+    return tuple((c, frozenset(str(o) for o in data[c])) for c in CLASSES)
+
+
+def default_classes() -> Dict[str, FrozenSet[str]]:
+    """The shipped op classes (`policy.yml`): {class: frozenset of op names}."""
+    return dict(_classes_from_file())
+
+
+def _result(names: Sequence[str]) -> str:
+    """The dtype inputs of `names` compute in together: the type promotion of every pair (`binary_result`)."""
+    return functools.reduce(binary_result, names)
+
+
+def _floating(dtype: str) -> bool:
+    return kind(dtype) == "float"
+
+
+@dataclasses.dataclass(frozen=True)
+class Policy:
+    param: str = "float32"
+    compute: str = "float32"
+    output: str = "float32"
+    float32_ops_half_io: bool = False
+    classes: Tuple[Tuple[str, FrozenSet[str]], ...] = dataclasses.field(default_factory=_classes_from_file)
+
+    def __post_init__(self):
+        for field in ("param", "compute", "output"):
+            name = canonical_name(getattr(self, field))
+            if not _floating(name):
+                raise ValueError(f"Policy.{field}: {name!r} is not a floating dtype")
+            object.__setattr__(self, field, name)
+        names = [c for c, _ in self.classes]
+        if sorted(names) != sorted(CLASSES):
+            raise ValueError(f"Policy.classes: {names} — a policy names every class ({', '.join(CLASSES)})")
+        object.__setattr__(self, "classes", tuple((c, frozenset(ops)) for c, ops in self.classes))
+
+    def _ops(self, cls: str) -> FrozenSet[str]:
+        return dict(self.classes)[cls]
+
+    def ops(self, cls: str) -> FrozenSet[str]:
+        """The op names of class `cls`."""
+        if cls not in CLASSES:
+            raise KeyError(f"Policy.ops: {cls!r} is not a class ({', '.join(CLASSES)})")
+        return self._ops(cls)
+
+    def op_class(self, op: str) -> Optional[str]:
+        """The compute class of `op` (`float32`, `lower`, `promote`), or None when it is in none."""
+        for cls in ("float32", "lower", "promote"):
+            if op in self._ops(cls):
+                return cls
+        return None
+
+    def compute_dtype(self, op: str, input_dtypes: Sequence[str]) -> Optional[str]:
+        """The dtype `op` computes in, given its inputs' dtypes; None when it has no input (a creation op takes the
+        dtype it is asked for). An op with no floating input computes in its inputs' promoted dtype."""
+        names = [canonical_name(d) for d in input_dtypes]
+        if not names:
+            return None
+        if not any(_floating(n) for n in names):
+            return _result(names)
+        cls = self.op_class(op)
+        if cls == "float32":
+            if self.float32_ops_half_io and op in self._ops("float32_half_io"):
+                return self.compute
+            return "float32"
+        if cls == "lower":
+            return self.compute
+        return _result([n for n in names if _floating(n)] if cls == "promote" else names)
+
+    def fill_value(self, op: str, value, dtype: str):
+        """`value`, a scalar a fill op converts to `dtype`, clamped to that dtype's finite range when `op` is a fill
+        op (a value of another kind, or of an op in no fill class, is returned unchanged)."""
+        if op not in self._ops("scalar_fill") and op not in self._ops("creation_fill"):
+            return value
+        if not isinstance(value, float):
+            return value
+        name = canonical_name(dtype)
+        if not _floating(name):
+            return value
+        top = finite_max(name)
+        return max(-top, min(top, value))
+
+    def replace(self, **changes) -> "Policy":
+        """A new policy with `changes` (the policy itself never changes)."""
+        return dataclasses.replace(self, **changes)
+
+    def as_dict(self) -> Dict[str, object]:
+        """The policy as plain data (names and lists only)."""
+        return {"param": self.param, "compute": self.compute, "output": self.output,
+                "float32_ops_half_io": self.float32_ops_half_io,
+                "classes": {c: sorted(ops) for c, ops in self.classes}}
+
+
+def get_policy(spec: str, classes: Optional[Mapping[str, Sequence[str]]] = None) -> Policy:
+    """A policy from jmp's spelling: `"params=float32,compute=float16,output=float32"` (also `p=`, `c=`, `o=`, and
+    `"float16"` alone for all three); `classes` replaces the shipped op classes."""
+    spec = spec.strip()
+    fields: Dict[str, str] = {}
+    if "=" not in spec:
+        fields = {"param": spec, "compute": spec, "output": spec}
+    else:
+        alias = {"p": "param", "params": "param", "param": "param", "c": "compute", "compute": "compute",
+                 "o": "output", "output": "output"}
+        for part in spec.split(","):
+            key, _, value = part.partition("=")
+            key = key.strip().lower()
+            if key not in alias:
+                raise ValueError(f"get_policy: {key!r} names no field (params, compute, output)")
+            fields[alias[key]] = value.strip()
+    if classes is not None:
+        fields["classes"] = tuple((c, frozenset(classes[c])) for c in CLASSES)
+    return Policy(**fields)

nbxtensor/dtype/policy.yml

@@ -0,0 +1,51 @@
+# The default precision policy's op classes (nbxtensor.dtype.policy, DESIGN rule 12): which ops compute in float32
+# whatever the policy's compute dtype, which ones compute in it, which ones take the widest of their inputs. An op in
+# no class keeps its inputs' dtypes (promoted among them). Op names are ATen's base names, the spelling a traced graph
+# carries.
+#
+# Source: PyTorch's autocast tables (aten/src/ATen/autocast_mode.h: AT_FORALL_FP32, AT_FORALL_FP32_SET_OPT_DTYPE,
+# AT_FORALL_LOWER_PRECISION_FP, AT_FORALL_PROMOTE), as NeuroBrix's DtypeEngine carried them
+# (moved here from its engine on 2026-10-29) with its documented deviations: polar and
+# view_as_complex compute in float32 (complex32 does not exist); div computes in the lower precision class; the
+# scaled-dot-product attention ops are in no class (the attention kernels choose their operand dtypes); the high-level
+# RNN ops (lstm, gru, rnn_tanh, rnn_relu) are lower precision, their kernels accumulate in float32.
+
+classes:
+  # compute in float32 (numerical range or accuracy): transcendental, power, norms, losses, reductions, softmax
+  float32: [
+    _softmax, _upsample_bicubic2d_aa, _upsample_bilinear2d_aa, _upsample_nearest_exact1d, _upsample_nearest_exact2d,
+    _upsample_nearest_exact3d, acos, asin, batch_norm, binary_cross_entropy_with_logits, cdist, cosh,
+    cosine_embedding_loss, cosine_similarity, cudnn_batch_norm, cumprod, cumsum, dist, erfinv, exp, expm1,
+    frobenius_norm, group_norm, hinge_embedding_loss, huber_loss, instance_norm, kl_div, l1_loss, layer_norm,
+    linalg_matrix_norm, linalg_vector_norm, log, log10, log1p, log2, log_softmax, logsumexp, margin_ranking_loss,
+    mse_loss, multi_margin_loss, multilabel_margin_loss, native_batch_norm, native_group_norm, native_layer_norm,
+    nll_loss, nll_loss2d, nuclear_norm, pdist, poisson_nll_loss, polar, pow, prod, reciprocal, renorm, rsqrt, sinh,
+    smooth_l1_loss, soft_margin_loss, softmax, softplus, sum, tan, triplet_margin_loss, upsample_bicubic2d,
+    upsample_bilinear2d, upsample_linear1d, upsample_nearest1d, upsample_nearest2d, upsample_nearest3d,
+    upsample_trilinear3d, view_as_complex
+  ]
+  # compute in the policy's compute dtype: convolutions, matrix products, RNN cells, some activations
+  lower: [
+    _convolution, _thnn_fused_gru_cell, _thnn_fused_lstm_cell, addbmm, addmm, addmv, addr, baddbmm, bmm,
+    chain_matmul, conv1d, conv2d, conv3d, conv_tbc, conv_transpose1d, conv_transpose2d, conv_transpose3d,
+    convolution, div, einsum, gru, gru_cell, linalg_multi_dot, linalg_vecdot, linear, lstm, lstm_cell, matmul, mm,
+    mv, prelu, rnn_relu, rnn_relu_cell, rnn_tanh, rnn_tanh_cell
+  ]
+  # compute in the widest floating dtype among their inputs (the inputs must agree)
+  promote: [
+    addcdiv, addcmul, atan2, bilinear, cross, dot, grid_sampler, grid_sampler_2d, grid_sampler_3d, index_add,
+    index_put, scatter_add, tensordot, vdot
+  ]
+  # float32-class ops whose kernels accumulate in float32 internally: a policy may let them take half-precision
+  # inputs and outputs (`Policy.float32_ops_half_io`)
+  float32_half_io: [
+    _log_softmax, _softmax, group_norm, layer_norm, log_softmax, native_group_norm, native_layer_norm, softmax
+  ]
+  # a scalar fill value converted to a tensor's dtype: clamped to that dtype's finite range (`Policy.fill_value`)
+  scalar_fill: [
+    fill, fill_, index_fill, index_fill_, masked_fill, masked_fill_
+  ]
+  # a creation op's scalar fill converted to its requested dtype: clamped likewise
+  creation_fill: [
+    full, full_like, new_full
+  ]

nbxtensor/dtype/promotion.py

@@ -0,0 +1,112 @@
+"""Type promotion: the one rule for the dtype an op's floating inputs are brought to.
+
+Two rules decide it, both read on dtype NAMES (plain data, R-NBXT 11) through the one table:
+
+- `promoted(dtypes)`: an op reading floating inputs of different widths reads them all at the WIDEST, when that
+  widest is wider than two bytes; among two-byte floats alone nothing is promoted (which half dtype an op computes
+  in is the caller's precision policy, not a promotion).
+- `binary_result(a, b)`: an element-wise op of two tensors computes in the higher of the two by kind, then by
+  width; two different dtypes of one kind and width (float16 and bfloat16) keep the first operand's.
+- `float_math_input(dtype)`: a floating-math op (rsqrt, exp, log ...) given an integer or bool input reads it as
+  float32, as PyTorch's type promotion does; a floating or complex input is read as it is. Without it rsqrt of an
+  int64 2 truncated 0.7071 to 0 (Kokoro's vocoder residual collapsed to zeros in NeuroBrix, 2026).
+"""
+from __future__ import annotations
+
+from typing import Optional, Sequence
+
+from .itemsize import itemsize, kind
+
+__all__ = ["promoted", "float_math_input", "binary_result", "scalar_result"]
+
+# The order of the kinds a binary result climbs (PyTorch type promotion: bool < unsigned < signed < float < complex).
+_KIND_ORDER = ("bool", "uint", "int", "float", "complex")
+
+
+def promoted(dtypes: Sequence[str]) -> Optional[str]:
+    """The dtype every floating input of an op is brought to, or None when no promotion applies: the widest
+    floating dtype among `dtypes` (the first of them in order on a tie), when it is wider than two bytes."""
+    widest, size = None, 0
+    for d in dtypes:
+        if kind(d) != "float":
+            continue
+        s = itemsize(d)
+        if s > size:
+            widest, size = d, s
+    return widest if widest is not None and size > 2 else None
+
+
+def float_math_input(dtype: str) -> str:
+    """The dtype a floating-math op reads an input of `dtype` at: float32 for an integer or bool input, the input's
+    own dtype otherwise."""
+    return dtype if kind(dtype) in ("float", "complex") else "float32"
+
+
+def _of_kind(k: str, nbytes: int) -> str:
+    """The table's dtype of kind `k` and `nbytes` bytes (the standard formats: intN, uintN, floatN, complexN)."""
+    name = {"int": "int", "uint": "uint", "float": "float", "complex": "complex"}[k] + str(8 * nbytes)
+    if kind(name) != k or itemsize(name) != nbytes:
+        raise ValueError(f"no {k} dtype of {nbytes} bytes in the dtype table")
+    return name
+
+
+def binary_result(a: str, b: str) -> str:
+    """The dtype an element-wise op of two tensors of dtypes `a` and `b` computes and writes in — the Array API
+    standard's type promotion (data-apis.org/array-api, "Type Promotion Rules"), the same whatever the operands' order:
+    within a kind the wider; a signed with an unsigned integer, the signed integer wide enough for both (none holds
+    uint64: refused by name); a real float with a complex, the complex carrying the wider float's precision. What the
+    standard leaves unspecified follows PyTorch, the engine's oracle: across kinds the higher kind (bool, unsigned,
+    signed, float, complex) with its own width; two floats of one width and different formats (float16, bfloat16)
+    meet in float32."""
+    ka, kb = kind(a), kind(b)
+    if ka == kb:
+        if itemsize(a) != itemsize(b):
+            return a if itemsize(a) > itemsize(b) else b
+        if a == b or ka not in ("float", "complex"):
+            return a
+        return _of_kind(ka, 2 * itemsize(a))                 # float16 with bfloat16: float32
+    if {ka, kb} == {"int", "uint"}:
+        signed, unsigned = (a, b) if ka == "int" else (b, a)
+        if itemsize(unsigned) >= 8:
+            raise ValueError(f"binary_result({a}, {b}): no signed integer holds every {unsigned} value")
+        return _of_kind("int", max(itemsize(signed), 2 * itemsize(unsigned)))
+    if {ka, kb} == {"float", "complex"}:
+        real, cplx = (a, b) if ka == "float" else (b, a)
+        return _of_kind("complex", max(itemsize(cplx), 2 * itemsize(real)))
+    return a if _KIND_ORDER.index(ka) > _KIND_ORDER.index(kb) else b
+
+
+def _int_bounds(dtype: str):
+    bits = 8 * itemsize(dtype)
+    return (0, (1 << bits) - 1) if kind(dtype) == "uint" else (-(1 << (bits - 1)), (1 << (bits - 1)) - 1)
+
+
+def scalar_result(dtype: str, scalar) -> str:
+    """The dtype an element-wise op of an array of `dtype` and a Python scalar computes in: the scalar is WEAK, it takes
+    the array's dtype (the Array API standard's "mixing arrays with Python scalars", NumPy's NEP 50) — a bool with a
+    bool array, an int with an integer array when the value fits it (refused by name otherwise: the standard leaves
+    the overflow open, a silent wrap is never taken), an int or a float with a real float array, an int, float or
+    complex with a complex array, a complex with a real float array as the complex of that precision (complex64 for a
+    half float: the table holds no complex32). Left open by the
+    standard and decided by NEP 50: a float with an integer or bool array is float64, a complex with one is
+    complex128, an int with a bool array is int64 (PyTorch differs on the first: float32; README, Array API)."""
+    k = kind(dtype)
+    if isinstance(scalar, bool):
+        return dtype
+    if isinstance(scalar, int):
+        if k in ("int", "uint"):
+            lo, hi = _int_bounds(dtype)
+            if not lo <= scalar <= hi:
+                raise OverflowError(f"scalar_result: {scalar} does not fit {dtype} [{lo}, {hi}]")
+            return dtype
+        return "int64" if k == "bool" else dtype
+    if isinstance(scalar, float):
+        return dtype if k in ("float", "complex") else "float64"
+    if isinstance(scalar, complex):
+        if k == "complex":
+            return dtype
+        if k != "float":
+            return "complex128"
+        # the complex of the float's precision; a half float has none in the table (no complex32): complex64 holds it
+        return _of_kind("complex", max(2 * itemsize(dtype), itemsize("complex64")))
+    raise TypeError(f"scalar_result: {type(scalar).__name__} is not a Python scalar (bool, int, float, complex)")

nbxtensor/host.py

@@ -0,0 +1,55 @@
+"""What the application that runs nbxtensor provides to it: the facts and policies the library does not own.
+
+nbxtensor computes; the application decides. The decisions still read here: the hardware vendor's files (tiles,
+budgets, the matrix unit) and where they live, the dtype policy that picks an attention's and a matrix unit's operand
+dtypes, and where a backend build may live. Each is moving to a fact the library reads itself or a setting with a
+default (docs/nbxtensor/API.md section 11); what left already: the memory limit (`memory.set_limit`), the certified
+store (`ops.autotune_settings.set_store`), the device's memory (`backend.device_facts`), the launch observer
+(`ops.launch_observer`). The application provides each service ONCE, by name (`provide`); the library reads it by name
+(`need`). A service no application provided is refused by name, never guessed.
+
+Each service is a callable. An application provides a function that reads its own state at call time, so a value it
+changes later is the one read.
+
+Torch-free, numpy-free, and nothing here imports the application.
+"""
+from __future__ import annotations
+
+from typing import Callable, Dict
+
+#: the services this library reads, and what each one answers — the contract an application fulfils
+SERVICES: Dict[str, str] = {
+    "vendors_dir": "() -> the directory of the application's own vendor files (pathlib.Path); optional: absent, the library's device files alone",
+}
+
+#: hooks an application MAY provide; absent, the library runs as the value here describes (never a guessed fact)
+_ABSENT: dict = {}
+
+_PROVIDED: Dict[str, Callable] = {}
+
+
+def provide(name: str, fn: Callable) -> None:
+    """The application provides service `name` (one of `SERVICES`). Providing it again replaces it."""
+    if name not in SERVICES:
+        raise KeyError(f"nbxtensor.host: {name!r} is not a service this library reads (known: {', '.join(sorted(SERVICES))})")
+    if not callable(fn):
+        raise TypeError(f"nbxtensor.host: service {name!r} must be callable, got {type(fn).__name__}")
+    _PROVIDED[name] = fn
+
+
+def need(name: str) -> Callable:
+    """Service `name` as the application provided it; refused by name when none did."""
+    try:
+        return _PROVIDED[name]
+    except KeyError:
+        if name in _ABSENT:
+            return _ABSENT[name]
+        if name not in SERVICES:
+            raise KeyError(f"nbxtensor.host: {name!r} is not a service this library reads") from None
+        raise RuntimeError(f"ZERO FALLBACK: nbxtensor needs the service {name!r} ({SERVICES[name]}) and the application "
+                           f"running it provided none (nbxtensor.host.provide)") from None
+
+
+def provided() -> Dict[str, Callable]:
+    """The services provided so far (a copy)."""
+    return dict(_PROVIDED)

nbxtensor/io/__init__.py

@@ -0,0 +1 @@
+"""Reading tensors from their files into NBXTensor (`io.safetensors`)."""

nbxtensor/io/read_rate.py

@@ -0,0 +1,197 @@
+"""How fast a machine reads its weight files, by number of parallel reads — the measurement behind a hardware profile's
+I/O worker count (`measure`). The reads are the loaders' own: pieces of `PIECE` bytes through `os.preadv` from a pool
+of threads (`nbxtensor.io.safetensors.read_bytes`). Each trial evicts the files' cached pages first and bypasses the
+cache while it reads (macOS `msync(MS_INVALIDATE)` + `F_NOCACHE`; Linux `posix_fadvise(DONTNEED)`), so the rate is the
+storage path's (a disk, a RAID, an NFS or RDMA mount), never the page cache's."""
+from __future__ import annotations
+
+import os
+import sys
+import time
+from concurrent.futures import ThreadPoolExecutor
+from typing import Dict, List, Optional, Sequence
+
+__all__ = ["PIECE", "WITNESS_BYTES", "START_BYTES", "mount_of", "pieces", "trial", "choose", "measure", "workers_for"]
+
+PIECE = 16 << 20                       # the loaders' read piece
+
+
+def _no_cache(fd: int) -> None:
+    if sys.platform == "darwin":
+        import fcntl
+        fcntl.fcntl(fd, 48, 1)         # F_NOCACHE: this descriptor's reads add nothing to the unified buffer cache
+    elif hasattr(os, "posix_fadvise"):
+        os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
+
+
+def _evict(path: str, size: int) -> None:
+    """The file's pages out of the cache before a trial: F_NOCACHE alone still SERVES pages cached earlier (measured
+    2026-10-30 on the Mac's NFS cache: 19 064 MiB/s cached, 32 MiB/s after this). macOS: msync(MS_INVALIDATE) on a
+    shared read-only mapping; Linux: posix_fadvise(DONTNEED)."""
+    fd = os.open(path, os.O_RDONLY)
+    try:
+        if sys.platform == "darwin":
+            import ctypes
+            import ctypes.util
+            libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)
+            libc.mmap.restype = ctypes.c_void_p
+            libc.mmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int,
+                                  ctypes.c_longlong]
+            libc.msync.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
+            libc.munmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
+            addr = libc.mmap(None, size, 1, 1, fd, 0)                    # PROT_READ, MAP_SHARED
+            if libc.msync(addr, size, 2) != 0:                           # MS_INVALIDATE
+                raise OSError(ctypes.get_errno(), f"msync(MS_INVALIDATE) {path}")
+            libc.munmap(addr, size)
+        elif hasattr(os, "posix_fadvise"):
+            os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
+    finally:
+        os.close(fd)
+
+
+def pieces(files, max_bytes: int):
+    """(path, offset, bytes) pieces of the loaders' size over `files`, in order, until `max_bytes`; and their total."""
+    out, total = [], 0
+    for path in files:
+        size = os.path.getsize(path)
+        for off in range(0, size, PIECE):
+            n = min(PIECE, size - off)
+            out.append((path, off, n))
+            total += n
+            if total >= max_bytes:
+                return out, total
+    return out, total
+
+
+def trial(pieces, workers: int) -> float:
+    """Bytes per second reading `pieces` with `workers` threads, the cache bypassed."""
+    for path in {p for p, _, _ in pieces}:
+        _evict(path, os.path.getsize(path))
+    fds = {}
+    for path, _, _ in pieces:
+        if path not in fds:
+            fd = os.open(path, os.O_RDONLY)
+            _no_cache(fd)
+            fds[path] = fd
+    bufs = [bytearray(PIECE) for _ in range(workers)]
+
+    def read(i):
+        path, off, n = pieces[i]
+        buf = memoryview(bufs[i % workers])[:n]
+        got = 0
+        while got < n:
+            got += os.preadv(fds[path], [buf[got:]], off + got)
+        return n
+
+    t = time.perf_counter()
+    with ThreadPoolExecutor(workers) as pool:
+        total = sum(pool.map(read, range(len(pieces))))
+    dt = time.perf_counter() - t
+    for fd in fds.values():
+        os.close(fd)
+    return total / dt
+
+
+def choose(rates, within: float) -> int:
+    """The smallest worker count whose rate is within `within` of the best: a worker past it holds a shard in host
+    memory for nothing."""
+    best = max(rates.values())
+    return min(c for c, r in rates.items() if r >= (1 - within) * best)
+
+
+def measure(files: Sequence[str], max_bytes: int, max_workers: int, within: float = 0.05) -> Dict[str, object]:
+    """Rates (bytes/s) at 1, 2, 4, ... `max_workers` parallel reads of `files` (at most `max_bytes` per trial), the
+    worker count `choose` takes, and whether the storage was steady: the one-worker trial is run first and last, and a
+    gap past `within` between them makes `steady` False — such a measure is not to be recorded."""
+    ps, total = pieces(list(files), max_bytes)
+    counts: List[int] = []
+    n = 1
+    while n <= max_workers:
+        counts.append(n)
+        n *= 2
+    rates = {c: trial(ps, c) for c in counts}
+    last = trial(ps, 1)
+    steady = abs(last - rates[1]) <= within * max(last, rates[1])
+    return {"bytes": total, "rates": rates, "workers": choose(rates, within), "steady": steady, "last_one_worker": last}
+
+
+WITNESS_BYTES = 1 << 30                # the witness file written when the storage holds no weight file yet
+START_BYTES = 2 << 30                  # the first trial: one shard's read, at most this
+
+
+def mount_of(path) -> str:
+    """The mount point the file `path` lives on: the storage a read of it goes through."""
+    p = os.path.realpath(path)
+    while not os.path.ismount(p):
+        p = os.path.dirname(p)
+    return p
+
+
+def workers_for(directory, witness_dir=None, max_workers: Optional[int] = None,
+                pattern: str = "*/components/*/weights/*.safetensors") -> Dict[str, object]:
+    """The parallel reads to keep in flight on the storage a weight cache is read from, MEASURED there:
+    {"workers": n, "storage": the mount point read, "read": "shards" or "witness"}.
+
+    The files read are `directory`'s largest matching `pattern` (the cache's shards). When it holds none yet, a witness
+    of `WITNESS_BYTES` is written in `witness_dir` (a local scratch directory the caller names, on the local cache's
+    storage when there is one) and removed — NEVER in `directory`, which may be a cache shared over a network; with no
+    `witness_dir`, refused by name. The answer says which storage it read. Counts 1, 2, 4, ... up to `max_workers`
+    (the machine's `os.cpu_count()` when not given).
+
+    The first trial reads one shard (the largest file, at most `START_BYTES`): a loader reads whole shards, and a
+    shorter trial undervalues parallel reads (64 MiB picked 4 on the rack's RDMA cache, 2 GiB 8/16/16 — the count
+    the planner prices host memory by is the one a load sees). The trials then grow while the measure is not stable:
+    at a given size, up to three measures; when two of them are steady (`measure`'s first/last one-worker check) the
+    answer is their median — one draw fluctuates and a plan prices the count. Otherwise the trial size is quadrupled,
+    up to everything the files hold. When no size gives a stable measure, refused by name: never a silent None, never
+    a default."""
+    from pathlib import Path
+    root = Path(directory)
+    max_workers = max_workers or os.cpu_count() or 1
+    files = [str(p) for p in sorted(root.glob(pattern), key=lambda p: -p.stat().st_size)[:4]]
+    witness = None
+    if not files:
+        if witness_dir is None:
+            raise RuntimeError(f"ZERO FALLBACK: {root} holds no weight file to measure its reads on, and no local "
+                               f"directory was named for a witness file (it is never written into the cache itself)")
+        Path(witness_dir).mkdir(parents=True, exist_ok=True)
+        witness = Path(witness_dir) / f".read_rate_witness_{os.getpid()}"
+        with open(witness, "wb") as f:
+            for _ in range(WITNESS_BYTES // PIECE):
+                f.write(os.urandom(PIECE))
+        files = [str(witness)]
+    try:
+        available = sum(os.path.getsize(p) for p in files)
+        # A loader reads whole shards: a trial shorter than one is a rate no load sees.
+        size, tried = max(4 * PIECE, min(max(os.path.getsize(p) for p in files), START_BYTES)), []
+        while True:
+            size = min(size, available)
+            picks = []
+            for _ in range(3):
+                m = measure(files, size, max_workers)
+                if m["steady"]:
+                    picks.append(int(m["workers"]))
+                if len(picks) == 2:
+                    break
+            tried.append(f"{size >> 20} MiB: {len(picks)}/3 steady")
+            if len(picks) >= 2:
+                n = picks[0] if picks[0] == picks[1] else _median_of_three(files, size, max_workers, picks)
+                return {"workers": n, "storage": mount_of(files[0]), "read": "witness" if witness else "shards"}
+            if size >= available:
+                raise RuntimeError(
+                    f"ZERO FALLBACK: no stable read-rate measure on {mount_of(files[0])} up to {available >> 20} MiB per "
+                    f"trial ({'; '.join(tried)}): another reader on the storage? measure again on a quiet machine "
+                    f"(tools/measure_io_workers.py) — no worker count is assumed")
+            size *= 4
+    finally:
+        if witness is not None:
+            witness.unlink(missing_ok=True)
+
+
+def _median_of_three(files, size, max_workers, picks):
+    """Two steady draws disagree: a third steady one decides (their median)."""
+    for _ in range(3):
+        m = measure(files, size, max_workers)
+        if m["steady"]:
+            return sorted(picks + [int(m["workers"])])[1]
+    return sorted(picks)[0]