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 .registry import ModelSpec, TensorLoc1920F_NOCACHE = 4821NP_DTYPE = {"F32": np.float32, "F16": np.float16, "U32": np.uint32, "I32": np.int32, "U8": np.uint8, "BF16": np.uint16}22MX_DTYPE = {"F32": mx.float32, "F16": mx.float16, "U32": mx.uint32, "I32": mx.int32, "U8": mx.uint8, "BF16": mx.bfloat16}232425class ExpertStore:26 def __init__(self, spec: ModelSpec, nocache: bool = False):27 self.spec = spec28 self.nocache = nocache29 self._fds: Dict[str, int] = {}30 self.bytes_read = 031 self.reads = 032 self.read_time = 0.03334 def _fd(self, path: str) -> int:35 fd = self._fds.get(path)36 if fd is None:37 fd = os.open(path, os.O_RDONLY)38 if self.nocache:39 fcntl.fcntl(fd, F_NOCACHE, 1)40 self._fds[path] = fd41 return fd4243 def read_tensor(self, loc: TensorLoc) -> mx.array:44 fd = self._fd(loc.file)45 t = time.perf_counter()46 buf = os.pread(fd, loc.nbytes, loc.offset)47 self.read_time += time.perf_counter() - t48 self.bytes_read += len(buf)49 self.reads += 150 if len(buf) != loc.nbytes:51 raise IOError(f"short read for {loc.key}: {len(buf)} != {loc.nbytes}")52 arr = np.frombuffer(buf, dtype=NP_DTYPE[loc.dtype]).reshape(loc.shape)53 if loc.dtype == "BF16":54 return mx.array(arr).view(mx.bfloat16)55 return mx.array(arr)5657 def read_expert(self, layer: int, expert: int) -> Dict[str, mx.array]:58 es = self.spec.experts[(layer, expert)]59 return {name: self.read_tensor(loc) for name, loc in es.tensors.items()}6061 def read_expert_bytes(self, layer: int, expert: int) -> Dict[str, bytes]:62 """Raw bytes (for the network transport)."""63 es = self.spec.experts[(layer, expert)]64 out = {}65 for name, loc in es.tensors.items():66 fd = self._fd(loc.file)67 t = time.perf_counter()68 out[name] = os.pread(fd, loc.nbytes, loc.offset)69 self.read_time += time.perf_counter() - t70 self.bytes_read += loc.nbytes71 self.reads += 172 return out7374 def stats(self) -> dict:75 return {"bytes_read": self.bytes_read, "reads": self.reads, "read_time_s": self.read_time}7677 def reset_stats(self):78 self.bytes_read = 079 self.reads = 080 self.read_time = 0.08182 def close(self):83 for fd in self._fds.values():84 os.close(fd)85 self._fds.clear()86