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.loads += 1109 self._complete = False110 return slot111112 def complete(self) -> bool:113 """True iff every expert of the layer is resident (then no host synchronisation is needed)."""114 if not self._complete and len(self.slot_of) == self.spec.num_experts:115 self.slot_map = mx.array([self.slot_of[e] for e in range(self.spec.num_experts)], dtype=mx.uint32)116 mx.eval(self.slot_map)117 self._complete = True118 return self._complete119120 def load(self, expert: int, protected: Iterable[int] = ()) -> int:121 """Synchronously load an expert from the store into a slot (demand miss)."""122 t = time.perf_counter()123 tensors = self.store.read_expert(self.layer, expert)124 slot = self._install(expert, tensors, protected)125 mx.eval(*self.buf.values())126 self.stall_time += time.perf_counter() - t127 return slot128129 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) and return their slot indices."""139 slots = []140 needed = set(experts)141 for e in experts:142 if e in self.slot_of:143 self.hits += 1144 else:145 self.misses += 1146 self.load(e, protected=needed)147 self.policy.touch(e)148 slots.append(self.slot_of[e])149 return slots150151 def resident(self) -> List[int]:152 return list(self.slot_of)153154 def stats(self) -> dict:155 n = self.hits + self.misses156 return {"layer": self.layer, "capacity": self.capacity, "hits": self.hits, "misses": self.misses, "hits_device": self.hits_device,157 "hit_rate": (self.hits / n) if n else None, "evictions": self.evictions, "loads": self.loads,158 "stall_time_s": self.stall_time, "policy": self.policy.name}159160 def reset_stats(self):161 self.hits = self.misses = self.evictions = self.loads = self.hits_device = 0162 self.stall_time = 0.0163