1"""Expert store: reads individual experts from the checkpoint file(s) with `pread`.23This is the SSD tier of the hierarchical memory. Every read is accounted (bytes, wall time) so that4stall time and SSD traffic can be reported per token. `nocache=True` sets F_NOCACHE on the file5descriptor (macOS) so that reads bypass the unified buffer cache for *new* pages; a cold measurement6still requires `purge` beforehand (see benchmarks/scripts/ssd_bandwidth.py).7"""8from __future__ import annotations910import fcntl11import os12import time13from typing import Dict, Tuple1415import mlx.core as mx16import numpy as np1718from concurrent.futures import ThreadPoolExecutor19from .registry import ModelSpec, TensorLoc2021F_NOCACHE = 4822NP_DTYPE = {"F32": np.float32, "F16": np.float16, "U32": np.uint32, "I32": np.int32, "U8": np.uint8, "BF16": np.uint16}23MX_DTYPE = {"F32": mx.float32, "F16": mx.float16, "U32": mx.uint32, "I32": mx.int32, "U8": mx.uint8, "BF16": mx.bfloat16}242526class ExpertStore:27 def __init__(self, spec: ModelSpec, nocache: bool = False, readers: int = 4):28 self.spec = spec29 self.nocache = nocache30 self._fds: Dict[str, int] = {}31 self.bytes_read = 032 self.reads = 033 self.read_time = 0.034 self.readers = readers35 self._pool = ThreadPoolExecutor(max_workers=readers, thread_name_prefix="titan-store") if readers > 1 else None36 # an expert whose tensors are contiguous in the file is read with ONE pread (OLMoE: all 1 024 experts are contiguous)37 self._contig: Dict[Tuple[int, int], Tuple[str, int, int, list]] = {}38 for key, es in spec.experts.items():39 locs = sorted(es.tensors.items(), key=lambda kv: kv[1].offset)40 if all(a[1].file == b[1].file and a[1].offset + a[1].nbytes == b[1].offset for a, b in zip(locs, locs[1:])):41 self._contig[key] = (locs[0][1].file, locs[0][1].offset, sum(t.nbytes for _, t in locs), locs)4243 def _fd(self, path: str) -> int:44 fd = self._fds.get(path)45 if fd is None:46 fd = os.open(path, os.O_RDONLY)47 if self.nocache:48 fcntl.fcntl(fd, F_NOCACHE, 1)49 self._fds[path] = fd50 return fd5152 def read_tensor(self, loc: TensorLoc, as_mx: bool = True):53 fd = self._fd(loc.file)54 t = time.perf_counter()55 buf = os.pread(fd, loc.nbytes, loc.offset)56 self.read_time += time.perf_counter() - t57 self.bytes_read += len(buf)58 self.reads += 159 if len(buf) != loc.nbytes:60 raise IOError(f"short read for {loc.key}: {len(buf)} != {loc.nbytes}")61 return self._wrap(loc, buf, as_mx)6263 def _wrap(self, loc: TensorLoc, buf, as_mx: bool = True):64 arr = np.frombuffer(buf, dtype=NP_DTYPE[loc.dtype]).reshape(loc.shape)65 if not as_mx:66 return (arr, loc.dtype)67 if loc.dtype == "BF16":68 return mx.array(arr).view(mx.bfloat16)69 return mx.array(arr)7071 @staticmethod72 def to_mx(t):73 """Convert a numpy (array, dtype) pair produced by a reader thread into an MLX array ON THE CALLING THREAD.74 MLX arrays must be created in the thread that evaluates them (streams are thread-local): creating them in reader75 threads raised 'There is no Stream(gpu, N) in current thread' on the worker (2026-11-01)."""76 if isinstance(t, mx.array):77 return t78 arr, dt = t79 return mx.array(arr).view(mx.bfloat16) if dt == "BF16" else mx.array(arr)8081 def read_expert(self, layer: int, expert: int, as_mx: bool = True):82 c = self._contig.get((layer, expert))83 if c is None:84 es = self.spec.experts[(layer, expert)]85 return {name: self.read_tensor(loc, as_mx) for name, loc in es.tensors.items()}86 path, off, n, locs = c87 fd = self._fd(path)88 t = time.perf_counter()89 buf = os.pread(fd, n, off)90 self.read_time += time.perf_counter() - t91 self.bytes_read += len(buf); self.reads += 192 if len(buf) != n:93 raise IOError(f"short read for expert ({layer},{expert}): {len(buf)} != {n}")94 out = {}95 for name, loc in locs:96 a = loc.offset - off97 out[name] = self._wrap(loc, memoryview(buf)[a:a + loc.nbytes], as_mx)98 return out99100 def read_experts(self, layer: int, experts) -> Dict[int, Dict[str, mx.array]]:101 """Read several experts, in parallel when a reader pool exists (NVMe benefits from queue depth)."""102 experts = list(experts)103 if self._pool is None or len(experts) == 1:104 return {e: self.read_expert(layer, e) for e in experts}105 t = time.perf_counter()106 raw = dict(zip(experts, self._pool.map(lambda e: self.read_expert(layer, e, as_mx=False), experts)))107 res = {e: {n: self.to_mx(v) for n, v in d.items()} for e, d in raw.items()} # MLX arrays created here, in the caller's thread108 self.read_time += 0 # per-read times already accumulated inside read_expert (sum of thread times, not wall)109 self.wall_batch_time = getattr(self, "wall_batch_time", 0.0) + (time.perf_counter() - t)110 return res111112 def read_expert_bytes(self, layer: int, expert: int) -> Dict[str, bytes]:113 """Raw bytes (for the network transport)."""114 es = self.spec.experts[(layer, expert)]115 out = {}116 for name, loc in es.tensors.items():117 fd = self._fd(loc.file)118 t = time.perf_counter()119 out[name] = os.pread(fd, loc.nbytes, loc.offset)120 self.read_time += time.perf_counter() - t121 self.bytes_read += loc.nbytes122 self.reads += 1123 return out124125 def stats(self) -> dict:126 return {"bytes_read": self.bytes_read, "reads": self.reads, "read_time_s": self.read_time}127128 def reset_stats(self):129 self.bytes_read = 0130 self.reads = 0131 self.read_time = 0.0132133 def close(self):134 for fd in self._fds.values():135 os.close(fd)136 self._fds.clear()137