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 = 09192 # ---- residency management -------------------------------------------------93 def _install(self, expert: int, tensors: Dict[str, mx.array], protected: Iterable[int]) -> int:94 if self.free:95 slot = self.free.pop()96 else:97 victim = self.policy.victim(protected)98 slot = self.slot_of.pop(victim)99 self.policy.remove(victim)100 self.evictions += 1101 for name, arr in tensors.items():102 self.buf[name][slot] = arr103 self.slot_of[expert] = slot104 self.expert_in[slot] = expert105 self.loads += 1106 return slot107108 def load(self, expert: int, protected: Iterable[int] = ()) -> int:109 """Synchronously load an expert from the store into a slot (demand miss)."""110 t = time.perf_counter()111 tensors = self.store.read_expert(self.layer, expert)112 slot = self._install(expert, tensors, protected)113 mx.eval(*self.buf.values())114 self.stall_time += time.perf_counter() - t115 return slot116117 def insert_prefetched(self, expert: int, tensors: Dict[str, mx.array], protected: Iterable[int] = ()) -> Optional[int]:118 """Install tensors that were fetched asynchronously (no stall accounted)."""119 if expert in self.slot_of:120 return self.slot_of[expert]121 if not self.free and all(e in set(protected) for e in self.slot_of):122 return None123 return self._install(expert, tensors, protected)124125 def ensure(self, experts: List[int]) -> List[int]:126 """Make every expert resident (demand-loading misses) and return their slot indices."""127 slots = []128 needed = set(experts)129 for e in experts:130 if e in self.slot_of:131 self.hits += 1132 else:133 self.misses += 1134 self.load(e, protected=needed)135 self.policy.touch(e)136 slots.append(self.slot_of[e])137 return slots138139 def resident(self) -> List[int]:140 return list(self.slot_of)141142 def stats(self) -> dict:143 n = self.hits + self.misses144 return {"layer": self.layer, "capacity": self.capacity, "hits": self.hits, "misses": self.misses,145 "hit_rate": (self.hits / n) if n else None, "evictions": self.evictions, "loads": self.loads,146 "stall_time_s": self.stall_time, "policy": self.policy.name}147148 def reset_stats(self):149 self.hits = self.misses = self.evictions = self.loads = 0150 self.stall_time = 0.0151