1"""Per-layer slot table: a bounded set of resident experts held in stacked MLX buffers.23Compute always runs `gather_qmm` on the slot buffers with expert indices remapped to slot indices,4so the kernel is the same whether an expert was resident or just loaded. Loading an expert into a5slot is an in-place slice assignment (measured 0.3 ms for 9 tensors / 3.4 MB on M5 Max, see6benchmarks/raw/gateB/design/). Eviction policy is pluggable (Gate C); the default is LRU.7Every miss is accounted as a stall (wall time between the request and the slot being ready).8"""9from __future__ import annotations1011import time12from collections import OrderedDict13from typing import Dict, Iterable, List, Optional1415import mlx.core as mx1617from .store import ExpertStore181920class EvictionPolicy:21 """Interface: called on every access; must return a victim expert id when the table is full."""22 name = "lru"2324 def __init__(self):25 self.order: "OrderedDict[int, None]" = OrderedDict()2627 def touch(self, expert: int):28 self.order.pop(expert, None)29 self.order[expert] = None3031 def victim(self, protected: Iterable[int]) -> int:32 prot = set(protected)33 for e in self.order:34 if e not in prot:35 return e36 raise RuntimeError("no evictable expert")3738 def remove(self, expert: int):39 self.order.pop(expert, None)404142class LFUPolicy(EvictionPolicy):43 name = "lfu"4445 def __init__(self):46 super().__init__()47 self.count: Dict[int, int] = {}4849 def touch(self, expert: int):50 super().touch(expert)51 self.count[expert] = self.count.get(expert, 0) + 15253 def victim(self, protected: Iterable[int]) -> int:54 prot = set(protected)55 cands = [(self.count.get(e, 0), i, e) for i, e in enumerate(self.order) if e not in prot]56 if not cands:57 raise RuntimeError("no evictable expert")58 return min(cands)[2]5960 def remove(self, expert: int):61 super().remove(expert)62 self.count.pop(expert, None)636465class SlotTable:66 def __init__(self, layer: int, capacity: int, store: ExpertStore, policy: Optional[EvictionPolicy] = None,67 matrices=("gate", "up", "down")):68 self.layer = layer69 self.capacity = capacity70 self.store = store71 self.policy = policy or EvictionPolicy()72 self.matrices = matrices73 self.spec = store.spec74 self.slot_of: Dict[int, int] = {}75 self.expert_in: List[Optional[int]] = [None] * capacity76 self.free: List[int] = list(range(capacity - 1, -1, -1))77 self.buf: Dict[str, mx.array] = {}78 # allocate stacked buffers from the first expert's shapes79 es = self.spec.experts[(layer, 0)]80 from .store import MX_DTYPE81 for name, loc in es.tensors.items():82 self.buf[name] = mx.zeros((capacity,) + tuple(loc.shape), dtype=MX_DTYPE[loc.dtype])83 mx.eval(*self.buf.values())84 # accounting85 self.hits = 086 self.misses = 087 self.evictions = 088 self.stall_time = 0.089 self.loads = 090 self.prefetched_used = 091 self.hits_device = 0 # accesses served on the fast path (table complete, no host sync)92 self._complete = False93 self.slot_map = None # device array expert -> slot, rebuilt when the table becomes complete9495 # ---- residency management -------------------------------------------------96 def _install(self, expert: int, tensors: Dict[str, mx.array], protected: Iterable[int]) -> int:97 if self.free:98 slot = self.free.pop()99 else:100 victim = self.policy.victim(protected)101 slot = self.slot_of.pop(victim)102 self.policy.remove(victim)103 self.evictions += 1104 for name, arr in tensors.items():105 self.buf[name][slot] = arr106 self.slot_of[expert] = slot107 self.expert_in[slot] = expert108 self.policy.touch(expert) # newly installed experts are MRU (prefetched ones included), so they are always evictable109 self.loads += 1110 self._complete = False111 return slot112113 def complete(self) -> bool:114 """True iff every expert of the layer is resident (then no host synchronisation is needed)."""115 if not self._complete and len(self.slot_of) == self.spec.num_experts:116 self.slot_map = mx.array([self.slot_of[e] for e in range(self.spec.num_experts)], dtype=mx.uint32)117 mx.eval(self.slot_map)118 self._complete = True119 return self._complete120121 def load(self, expert: int, protected: Iterable[int] = ()) -> int:122 """Load an expert from the store into a slot (demand miss). The slot write is a lazy slice assignment that the123 layer's gather_qmm will consume: no per-expert GPU synchronisation (the stall counted here is the read time)."""124 t = time.perf_counter()125 tensors = self.store.read_expert(self.layer, expert)126 self.stall_time += time.perf_counter() - t127 return self._install(expert, tensors, protected)128129 def insert_prefetched(self, expert: int, tensors: Dict[str, mx.array], protected: Iterable[int] = ()) -> Optional[int]:130 """Install tensors that were fetched asynchronously (no stall accounted)."""131 if expert in self.slot_of:132 return self.slot_of[expert]133 if not self.free and all(e in set(protected) for e in self.slot_of):134 return None135 return self._install(expert, tensors, protected)136137 def ensure(self, experts: List[int]) -> List[int]:138 """Make every expert resident (demand-loading misses, read as a batch) and return their slot indices."""139 needed = set(experts)140 missing = [e for e in experts if e not in self.slot_of]141 self.hits += len(experts) - len(missing)142 self.misses += len(missing)143 if missing:144 t = time.perf_counter()145 tensors = {e: self.store.read_expert(self.layer, e) for e in missing}146 self.stall_time += time.perf_counter() - t147 for e in missing:148 self._install(e, tensors[e], needed)149 for e in experts:150 self.policy.touch(e)151 return [self.slot_of[e] for e in experts]152153 def resident(self) -> List[int]:154 return list(self.slot_of)155156 def stats(self) -> dict:157 n = self.hits + self.misses158 return {"layer": self.layer, "capacity": self.capacity, "hits": self.hits, "misses": self.misses, "hits_device": self.hits_device,159 "hit_rate": (self.hits / n) if n else None, "evictions": self.evictions, "loads": self.loads,160 "stall_time_s": self.stall_time, "policy": self.policy.name}161162 def reset_stats(self):163 self.hits = self.misses = self.evictions = self.loads = self.hits_device = 0164 self.stall_time = 0.0165