1"""Asynchronous prefetcher: a small pool of reader threads fetches predicted experts from the ExpertStore (SSD tier)2while the GPU computes the current layer; completed fetches are installed into the slot table on the main thread when the3layer is reached (`drain`). Predictors come from titan.predictors and only see information available at decision time.45Accounting per layer: prefetch_issued, prefetch_installed, prefetch_used (installed experts that were subsequently hit),6prefetch_wasted (installed but evicted before use), prefetch_late (needed before the fetch completed: counted as a demand7miss whose stall is the *remaining* latency).8"""9from __future__ import annotations1011import threading, time12from concurrent.futures import Future, ThreadPoolExecutor13from typing import Dict, List, Optional1415import mlx.core as mx16import numpy as np1718from .slots import SlotTable19from .store import ExpertStore202122class Prefetcher:23 def __init__(self, tables: List[SlotTable], store: ExpertStore, predictor, m: int = 8, workers: int = 2):24 self.tables = tables25 self.store = store26 self.pred = predictor27 self.m = m28 self.pool = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="titan-prefetch")29 self.inflight: List[Dict[int, Future]] = [dict() for _ in tables]30 self.lock = threading.Lock()31 self.issued = [0] * len(tables); self.installed = [0] * len(tables); self.late = [0] * len(tables); self.dup = [0] * len(tables)32 self.prev_topk: Optional[np.ndarray] = None33 self.prev_probs: Optional[np.ndarray] = None34 self.cur: Dict = {}35 self.L = len(tables)36 self.K = store.spec.top_k3738 # ---- reading ---------------------------------------------------------------39 def _fetch(self, layer: int, expert: int):40 t = time.perf_counter()41 tensors = self.store.read_expert(layer, expert, as_mx=False) # pread + numpy only; MLX arrays are created in drain()42 return tensors, time.perf_counter() - t4344 def issue(self, layer: int, experts: List[int]):45 tab = self.tables[layer]46 for e in experts[: self.m]:47 if e in tab.slot_of or e in self.inflight[layer]:48 self.dup[layer] += 149 continue50 self.inflight[layer][e] = self.pool.submit(self._fetch, layer, e)51 self.issued[layer] += 15253 def drain(self, layer: int, protected: List[int], wait_for: List[int] = ()):54 """Install completed prefetches for `layer`; optionally wait for the ones in `wait_for` (needed now)."""55 tab = self.tables[layer]56 inf = self.inflight[layer]57 for e in list(inf):58 fut = inf[e]59 if e in wait_for and not fut.done():60 t = time.perf_counter()61 fut.result()62 tab.stall_time += time.perf_counter() - t # remaining latency of a late prefetch63 self.late[layer] += 164 if fut.done():65 tensors, _ = fut.result()66 tensors = {n: self.store.to_mx(v) for n, v in tensors.items()}67 del inf[e]68 if tab.insert_prefetched(e, tensors, protected) is not None:69 self.installed[layer] += 17071 # ---- hooks called by the execution block ---------------------------------------72 def before_layer(self, layer: int, needed: List[int]):73 """Called with the materialised routing of `layer` before its experts are ensured."""74 self.drain(layer, needed, wait_for=[e for e in needed if e in self.inflight[layer]])75 if hasattr(self.pred, "set_true"):76 pass7778 def after_layer(self, layer: int, uniq: List[int], probs: Optional[np.ndarray] = None):79 self.cur[layer] = np.array(uniq)80 if probs is not None:81 self.cur[("probs", layer)] = probs82 if self.pred.horizon == "layer" and layer + 1 < self.L:83 self.issue(layer + 1, [int(e) for e in self.pred.predict(layer + 1, self.prev_topk, self.prev_probs, self.cur, self.m)])84 if layer == self.L - 1:85 self.prev_topk = np.array([self.cur.get(l, np.array([]))[: self.K] if len(self.cur.get(l, [])) >= self.K else np.pad(self.cur.get(l, np.array([])), (0, self.K - len(self.cur.get(l, [])))) for l in range(self.L)])86 self.prev_probs = np.stack([self.cur[("probs", l)] for l in range(self.L)]) if all(("probs", l) in self.cur for l in range(self.L)) else None87 self.cur = {}88 if self.pred.horizon == "token":89 for l in range(self.L):90 self.issue(l, [int(e) for e in self.pred.predict(l, self.prev_topk, self.prev_probs, {}, self.m)])9192 def stats(self) -> dict:93 return {"issued": sum(self.issued), "installed": sum(self.installed), "late": sum(self.late), "dup": sum(self.dup),94 "predictor": self.pred.name, "m": self.m}9596 def reset_stats(self):97 self.issued = [0] * self.L; self.installed = [0] * self.L; self.late = [0] * self.L; self.dup = [0] * self.L9899 def close(self):100 self.pool.shutdown(wait=True)101