1"""Expert-demand predictors. Every predictor only uses information available at its decision point:2- `PrevToken`: at the start of token t, the experts used at layer l by token t-1 (all layers known after token t-1).3- `Transition`: at layer l-1 of token t (its routing just materialised), a transition-count matrix fitted on *fit* traces,4 M_l[e', e] = #(e' chosen at l-1 and e at l) -> score(e) = sum_{e' in R(t,l-1)} M_l[e', e].5- `Probe`: a per-layer linear multi-label probe on the layer-(l-1) router probabilities (64 -> 64 logits), fitted on fit traces.6- `Oracle`: the true R(t,l) (upper bound only; never a method).7No predictor sees routes of layer >= l at token t, nor any future token.8"""9from __future__ import annotations10import numpy as np111213class Predictor:14 name = "none"15 horizon = 0 # 0 = no prefetch; "token" = at token start; "layer" = one layer ahead1617 def fit(self, traces): # traces: list of (topk[T,L,K], probs[T,L,E])18 return self1920 def predict(self, l, prev_topk, prev_probs, cur_routes, m):21 """Return up to m expert ids for layer l.22 prev_topk: [L,K] routes of the previous token (or None); cur_routes: dict l' -> routes for l' < l at the current token;23 prev_probs: [L,E] probs of the previous token."""24 return []252627class PrevToken(Predictor):28 name = "prev_token"29 horizon = "token"3031 def predict(self, l, prev_topk, prev_probs, cur_routes, m):32 if prev_topk is None:33 return []34 return list(prev_topk[l][:m])353637class PrevTokenRank(Predictor):38 """At token start: for layer l, the experts ranked m highest by the *previous token's* router probabilities at layer l39 (ranks 1..K are normally already resident, so this effectively prefetches the next-most-likely experts)."""40 name = "prev_rank"41 horizon = "token"4243 def __init__(self, K=8):44 self.K = K4546 def predict(self, l, prev_topk, prev_probs, cur_routes, m):47 if prev_probs is None:48 return []49 return list(np.argsort(-prev_probs[l].astype(np.float32))[self.K:self.K + m]) # ranks K+1..K+m (top-K are already resident)505152class Transition(Predictor):53 """Transition counts from layer l-d to layer l (d = lookahead depth); issued when layer l-d's routing is known."""54 name = "transition"55 horizon = "layer"5657 def __init__(self, L, E, depth=1, threshold=0.0):58 self.M = np.zeros((L, E, E), dtype=np.float64)59 self.depth = depth60 self.threshold = threshold # keep only experts whose score >= threshold * max score (precision/recall knob)61 self.name = ("transition" if depth == 1 else f"transition{depth}") + (f"t{int(threshold*100)}" if threshold else "")6263 def fit(self, traces):64 d = self.depth65 for topk, _ in traces:66 for l in range(d, topk.shape[1]):67 for t in range(topk.shape[0]):68 for a in topk[t, l - d]:69 self.M[l, a, topk[t, l]] += 170 return self7172 def predict(self, l, prev_topk, prev_probs, cur_routes, m):73 d = self.depth74 if l < d or (l - d) not in cur_routes:75 return list(prev_topk[l][:m]) if prev_topk is not None else []76 s = self.M[l][list(cur_routes[l - d])].sum(0)77 order = np.argsort(-s)[:m]78 if self.threshold:79 order = [e for e in order if s[e] >= self.threshold * s[order[0]]]80 return list(order)818283class Probe(Predictor):84 name = "probe"85 horizon = "layer"8687 def __init__(self, L, E, epochs=2000, lr=0.05, l2=1e-5, seed=0, depth=1):88 self.L, self.E, self.epochs, self.lr, self.l2 = L, E, epochs, lr, l289 self.depth = depth90 self.name = "probe" if depth == 1 else f"probe{depth}"91 self.W = np.zeros((L, E + 1, E)); self.mu = np.zeros((L, E)); self.sd = np.ones((L, E)); self.rng = np.random.default_rng(seed)9293 def fit(self, traces):94 d = self.depth95 for l in range(d, self.L):96 X = np.concatenate([p[:, l - d].astype(np.float32) for _, p in traces], 0)97 Y = np.zeros((X.shape[0], self.E), np.float32)98 rows, cols, off = [], [], 099 for t, _ in traces: # row offsets accumulate across traces (bug fixed 2026-11-01: labels were misaligned)100 rows.append(off + np.repeat(np.arange(t.shape[0]), t.shape[2])); cols.append(t[:, l].reshape(-1)); off += t.shape[0]101 Y[np.concatenate(rows), np.concatenate(cols)] = 1.0102 X = np.log(np.clip(X, 1e-6, None)) # log-probabilities as features (scale-free)103 self.mu[l] = X.mean(0); self.sd[l] = X.std(0) + 1e-6104 Xb = np.concatenate([(X - self.mu[l]) / self.sd[l], np.ones((X.shape[0], 1), np.float32)], 1)105 W = np.zeros((self.E + 1, self.E)); V = np.zeros_like(W)106 for _ in range(self.epochs): # full-batch logistic regression with momentum (multi-label, independent sigmoids)107 P = 1 / (1 + np.exp(-(Xb @ W)))108 G = Xb.T @ (P - Y) / X.shape[0] + self.l2 * W109 V = 0.9 * V + G; W -= self.lr * V110 self.W[l] = W111 return self112113 def predict(self, l, prev_topk, prev_probs, cur_routes, m):114 d = self.depth115 if l < d or cur_routes.get(("probs", l - d)) is None:116 return list(prev_topk[l][:m]) if prev_topk is not None else []117 x = (np.log(np.clip(cur_routes[("probs", l - d)].astype(np.float32), 1e-6, None)) - self.mu[l]) / self.sd[l]118 return list(np.argsort(-(np.append(x, 1.0) @ self.W[l]))[:m])119120121class Oracle(Predictor):122 name = "oracle"123 horizon = "layer"124125 def predict(self, l, prev_topk, prev_probs, cur_routes, m):126 return list(cur_routes[("true", l)][:m])127128129def make(name, L, E):130 reg = {"none": Predictor, "prev_token": PrevToken, "prev_rank": PrevTokenRank, "transition": lambda: Transition(L, E),131 "transition2": lambda: Transition(L, E, 2), "transition3": lambda: Transition(L, E, 3), "transition4": lambda: Transition(L, E, 4),132 "transition2t50": lambda: Transition(L, E, 2, 0.5), "transition2t70": lambda: Transition(L, E, 2, 0.7), "transition1t50": lambda: Transition(L, E, 1, 0.5),133 "probe": lambda: Probe(L, E), "probe2": lambda: Probe(L, E, depth=2), "probe3": lambda: Probe(L, E, depth=3), "oracle": Oracle}134 return reg[name]()135