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) -> mx.array: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 arr = np.frombuffer(buf, dtype=NP_DTYPE[loc.dtype]).reshape(loc.shape)62 if loc.dtype == "BF16":63 return mx.array(arr).view(mx.bfloat16)64 return mx.array(arr)6566 def _wrap(self, loc: TensorLoc, buf) -> mx.array:67 arr = np.frombuffer(buf, dtype=NP_DTYPE[loc.dtype]).reshape(loc.shape)68 if loc.dtype == "BF16":69 return mx.array(arr).view(mx.bfloat16)70 return mx.array(arr)7172 def read_expert(self, layer: int, expert: int) -> Dict[str, mx.array]:73 c = self._contig.get((layer, expert))74 if c is None:75 es = self.spec.experts[(layer, expert)]76 return {name: self.read_tensor(loc) for name, loc in es.tensors.items()}77 path, off, n, locs = c78 fd = self._fd(path)79 t = time.perf_counter()80 buf = os.pread(fd, n, off)81 self.read_time += time.perf_counter() - t82 self.bytes_read += len(buf); self.reads += 183 if len(buf) != n:84 raise IOError(f"short read for expert ({layer},{expert}): {len(buf)} != {n}")85 out = {}86 for name, loc in locs:87 a = loc.offset - off88 out[name] = self._wrap(loc, memoryview(buf)[a:a + loc.nbytes])89 return out9091 def read_experts(self, layer: int, experts) -> Dict[int, Dict[str, mx.array]]:92 """Read several experts, in parallel when a reader pool exists (NVMe benefits from queue depth)."""93 experts = list(experts)94 if self._pool is None or len(experts) == 1:95 return {e: self.read_expert(layer, e) for e in experts}96 t = time.perf_counter()97 res = dict(zip(experts, self._pool.map(lambda e: self.read_expert(layer, e), experts)))98 self.read_time += 0 # per-read times already accumulated inside read_expert (sum of thread times, not wall)99 self.wall_batch_time = getattr(self, "wall_batch_time", 0.0) + (time.perf_counter() - t)100 return res101102 def read_expert_bytes(self, layer: int, expert: int) -> Dict[str, bytes]:103 """Raw bytes (for the network transport)."""104 es = self.spec.experts[(layer, expert)]105 out = {}106 for name, loc in es.tensors.items():107 fd = self._fd(loc.file)108 t = time.perf_counter()109 out[name] = os.pread(fd, loc.nbytes, loc.offset)110 self.read_time += time.perf_counter() - t111 self.bytes_read += loc.nbytes112 self.reads += 1113 return out114115 def stats(self) -> dict:116 return {"bytes_read": self.bytes_read, "reads": self.reads, "read_time_s": self.read_time}117118 def reset_stats(self):119 self.bytes_read = 0120 self.reads = 0121 self.read_time = 0.0122123 def close(self):124 for fd in self._fds.values():125 os.close(fd)126 self._fds.clear()127