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):58 self.M = np.zeros((L, E, E), dtype=np.float64)59 self.depth = depth60 self.name = "transition" if depth == 1 else f"transition{depth}"6162 def fit(self, traces):63 d = self.depth64 for topk, _ in traces:65 for l in range(d, topk.shape[1]):66 for t in range(topk.shape[0]):67 for a in topk[t, l - d]:68 self.M[l, a, topk[t, l]] += 169 return self7071 def predict(self, l, prev_topk, prev_probs, cur_routes, m):72 d = self.depth73 if l < d or (l - d) not in cur_routes:74 return list(prev_topk[l][:m]) if prev_topk is not None else []75 s = self.M[l][list(cur_routes[l - d])].sum(0)76 return list(np.argsort(-s)[:m])777879class Probe(Predictor):80 name = "probe"81 horizon = "layer"8283 def __init__(self, L, E, epochs=2000, lr=0.05, l2=1e-5, seed=0):84 self.L, self.E, self.epochs, self.lr, self.l2 = L, E, epochs, lr, l285 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)8687 def fit(self, traces):88 for l in range(1, self.L):89 X = np.concatenate([p[:, l - 1].astype(np.float32) for _, p in traces], 0)90 Y = np.zeros((X.shape[0], self.E), np.float32)91 rows = np.concatenate([np.repeat(np.arange(t.shape[0]), t.shape[2]) for t, _ in traces]); off = 0; cols = []92 for t, _ in traces:93 cols.append(t[:, l].reshape(-1)); 94 cols = np.concatenate(cols)95 Y[rows, cols] = 1.096 X = np.log(np.clip(X, 1e-6, None)) # log-probabilities as features (scale-free)97 self.mu[l] = X.mean(0); self.sd[l] = X.std(0) + 1e-698 Xb = np.concatenate([(X - self.mu[l]) / self.sd[l], np.ones((X.shape[0], 1), np.float32)], 1)99 W = np.zeros((self.E + 1, self.E)); V = np.zeros_like(W)100 for _ in range(self.epochs): # full-batch logistic regression with momentum (multi-label, independent sigmoids)101 P = 1 / (1 + np.exp(-(Xb @ W)))102 G = Xb.T @ (P - Y) / X.shape[0] + self.l2 * W103 V = 0.9 * V + G; W -= self.lr * V104 self.W[l] = W105 return self106107 def predict(self, l, prev_topk, prev_probs, cur_routes, m):108 if l == 0 or (l - 1) not in cur_routes or cur_routes.get(("probs", l - 1)) is None:109 return list(prev_topk[l][:m]) if prev_topk is not None else []110 x = (np.log(np.clip(cur_routes[("probs", l - 1)].astype(np.float32), 1e-6, None)) - self.mu[l]) / self.sd[l]111 return list(np.argsort(-(np.append(x, 1.0) @ self.W[l]))[:m])112113114class Oracle(Predictor):115 name = "oracle"116 horizon = "layer"117118 def predict(self, l, prev_topk, prev_probs, cur_routes, m):119 return list(cur_routes[("true", l)][:m])120121122def make(name, L, E):123 reg = {"none": Predictor, "prev_token": PrevToken, "prev_rank": PrevTokenRank, "transition": lambda: Transition(L, E),124 "transition2": lambda: Transition(L, E, 2), "transition3": lambda: Transition(L, E, 3), "transition4": lambda: Transition(L, E, 4),125 "probe": lambda: Probe(L, E), "oracle": Oracle}126 return reg[name]()127