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"), fused_upgate: bool = False):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 self.fused = fused_upgate82 for name, loc in es.tensors.items():83 if self.fused and name.startswith(("up.", "gate.")):84 continue85 self.buf[name] = mx.zeros((capacity,) + tuple(loc.shape), dtype=MX_DTYPE[loc.dtype])86 if self.fused: # [C, 2H, ...]: rows 0..H-1 = up, H..2H-1 = gate87 for part in ("weight", "scales", "biases"):88 loc = es.tensors.get(f"up.{part}")89 if loc is not None:90 self.buf[f"upgate.{part}"] = mx.zeros((capacity, 2 * loc.shape[0]) + tuple(loc.shape[1:]), dtype=MX_DTYPE[loc.dtype])91 mx.eval(*self.buf.values())92 # accounting93 self.hits = 094 self.misses = 095 self.evictions = 096 self.stall_time = 0.097 self.loads = 098 self.prefetched_used = 099 self.hits_device = 0 # accesses served on the fast path (table complete, no host sync)100 self._complete = False101 self.slot_map = None # device array expert -> slot, rebuilt when the table becomes complete102103 # ---- residency management -------------------------------------------------104 def _install(self, expert: int, tensors: Dict[str, mx.array], protected: Iterable[int]) -> int:105 if self.free:106 slot = self.free.pop()107 else:108 victim = self.policy.victim(protected)109 slot = self.slot_of.pop(victim)110 self.policy.remove(victim)111 self.evictions += 1112 if self.fused:113 H = tensors["up.weight"].shape[0]114 for part in ("weight", "scales", "biases"):115 if f"up.{part}" in tensors:116 self.buf[f"upgate.{part}"][slot, :H] = tensors[f"up.{part}"]117 self.buf[f"upgate.{part}"][slot, H:] = tensors[f"gate.{part}"]118 for name, arr in tensors.items():119 if name.startswith("down."):120 self.buf[name][slot] = arr121 else:122 for name, arr in tensors.items():123 self.buf[name][slot] = arr124 self.slot_of[expert] = slot125 self.expert_in[slot] = expert126 self.policy.touch(expert) # newly installed experts are MRU (prefetched ones included), so they are always evictable127 self.loads += 1128 self._complete = False129 return slot130131 def complete(self) -> bool:132 """True iff every expert of the layer is resident (then no host synchronisation is needed)."""133 if not self._complete and len(self.slot_of) == self.spec.num_experts:134 self.slot_map = mx.array([self.slot_of[e] for e in range(self.spec.num_experts)], dtype=mx.uint32)135 mx.eval(self.slot_map)136 self._complete = True137 return self._complete138139 def load(self, expert: int, protected: Iterable[int] = ()) -> int:140 """Load an expert from the store into a slot (demand miss). The slot write is a lazy slice assignment that the141 layer's gather_qmm will consume: no per-expert GPU synchronisation (the stall counted here is the read time)."""142 t = time.perf_counter()143 tensors = self.store.read_expert(self.layer, expert)144 self.stall_time += time.perf_counter() - t145 return self._install(expert, tensors, protected)146147 def insert_prefetched(self, expert: int, tensors: Dict[str, mx.array], protected: Iterable[int] = ()) -> Optional[int]:148 """Install tensors that were fetched asynchronously (no stall accounted)."""149 if expert in self.slot_of:150 return self.slot_of[expert]151 if not self.free and all(e in set(protected) for e in self.slot_of):152 return None153 return self._install(expert, tensors, protected)154155 def ensure(self, experts: List[int]) -> List[int]:156 """Make every expert resident (demand-loading misses, read as a batch) and return their slot indices."""157 needed = set(experts)158 missing = [e for e in experts if e not in self.slot_of]159 self.hits += len(experts) - len(missing)160 self.misses += len(missing)161 if missing:162 t = time.perf_counter()163 tensors = self.store.read_experts(self.layer, missing)164 self.stall_time += time.perf_counter() - t165 for e in missing:166 self._install(e, tensors[e], needed)167 for e in experts:168 self.policy.touch(e)169 return [self.slot_of[e] for e in experts]170171 def resident(self) -> List[int]:172 return list(self.slot_of)173174 def stats(self) -> dict:175 n = self.hits + self.misses176 return {"layer": self.layer, "capacity": self.capacity, "hits": self.hits, "misses": self.misses, "hits_device": self.hits_device,177 "hit_rate": (self.hits / n) if n else None, "evictions": self.evictions, "loads": self.loads,178 "stall_time_s": self.stall_time, "policy": self.policy.name}179180 def reset_stats(self):181 self.hits = self.misses = self.evictions = self.loads = self.hits_device = 0182 self.stall_time = 0.0183