#!/usr/bin/env python3 """ MoE expert prefetch — a method-INDEPENDENT exploitability test ============================================================== moe_real_traces.py showed real co-activation locality-gain 1.4-2.1x with a shuffle null at ~1.0. But that metric still leans on a clustering algorithm + a chance model (the same shape of metric that lied to us before). This tests the SAME claim with zero clustering and zero chance model — pure next-token prefetch hit-rate on the real temporal trace, with an honest train/test split. Setup (per layer): the model fires top-k of N experts each token. Imagine a fast-memory cache that can hold a *prefetch budget* of B experts for the NEXT token. Two predictors, both trained on the first half of the trace, evaluated on the second half: - LFU (frequency): prefetch the B globally-most-frequent experts. Static. The "trivial caching" baseline the prior negative said was all that survives. - CO-ACT (temporal): given the experts that fired at token t, prefetch the B experts with the highest learned co-activation with that set, for t+1. hit-rate = fraction of the k experts actually fired at t+1 that were prefetched. If CO-ACT > LFU at the same budget B, temporal co-activation is exploitable for prefetch BEYOND marginals — a real, metric-independent win. If they tie, the only thing real is expert popularity (ordinary LFU) and the clustering gain was a mirage. Usage: python moe_cache_sim.py [--model allenai/OLMoE-1B-7B-0924] """ import argparse, json import numpy as np import torch from transformers import AutoModelForCausalLM, AutoTokenizer # Longer, contiguous real text per domain so the temporal (t -> t+1) signal is meaningful. CORPUS = [ "The mitochondrion is a double membrane-bound organelle found in most eukaryotic " "organisms. Mitochondria generate most of the cell's supply of adenosine triphosphate, " "used as a source of chemical energy. They were first described in the 1840s and the " "term was coined by Carl Benda in 1898. A mitochondrion contains its own genome, a " "small circular DNA molecule, and replicates independently of the cell. The number of " "mitochondria per cell varies widely by organism, tissue, and cell type; a red blood " "cell has none, while a liver cell can contain more than two thousand.", "def quicksort(a):\n if len(a) <= 1:\n return a\n pivot = a[len(a)//2]\n" " left = [x for x in a if x < pivot]\n mid = [x for x in a if x == pivot]\n" " right = [x for x in a if x > pivot]\n return quicksort(left) + mid + quicksort(right)\n\n" "def mergesort(a):\n if len(a) <= 1:\n return a\n m = len(a)//2\n" " l = mergesort(a[:m])\n r = mergesort(a[m:])\n out = []\n i = j = 0\n" " while i < len(l) and j < len(r):\n if l[i] <= r[j]:\n out.append(l[i]); i += 1\n" " else:\n out.append(r[j]); j += 1\n return out + l[i:] + r[j:]", "Theorem (Fermat's little theorem). For any prime p and integer a not divisible by p, " "a^(p-1) is congruent to 1 modulo p. Proof. Consider the residues a, 2a, ..., (p-1)a " "modulo p. No two are congruent, for if ia is congruent to ja then i is congruent to j " "since a is invertible modulo p. Hence these p-1 residues are a permutation of 1, 2, ..., " "p-1. Taking the product of both sides, a^(p-1) times (p-1)! is congruent to (p-1)! " "modulo p, and cancelling (p-1)! gives the result. This underlies the Fermat primality test.", "User: How do I reverse a linked list in place? Assistant: Walk the list with three " "pointers named previous, current, and next. Initialize previous to null and current to " "the head. At each step, store current.next in next, set current.next to previous, then " "advance previous to current and current to next. When current becomes null, previous " "points at the old tail, which is the new head, so return previous. This runs in linear " "time and constant space because no new nodes are allocated.", "El aprendizaje automatico es una rama de la inteligencia artificial que estudia " "algoritmos capaces de aprender de los datos. Un modelo se ajusta a partir de ejemplos " "de entrenamiento para hacer predicciones sobre datos nuevos sin ser programado de forma " "explicita para cada caso. Entre las familias principales estan el aprendizaje supervisado, " "el no supervisado y el aprendizaje por refuerzo. Se aplica en medicina, vision por " "computador, traduccion automatica y deteccion de fraude, entre muchos otros campos.", "Largest planets of the Solar System by equatorial radius, in order: Jupiter at about " "69,911 kilometres, Saturn at about 58,232, Uranus at about 25,362, and Neptune at about " "24,622. The terrestrial planets are far smaller: Earth at 6,371, Venus at 6,052, Mars at " "3,390, and Mercury at 2,440. The Sun accounts for about 99.86 percent of the total mass " "of the Solar System, and Jupiter accounts for most of the remainder.", ] def main(): ap = argparse.ArgumentParser() ap.add_argument("--model", default="allenai/OLMoE-1B-7B-0924") ap.add_argument("--budgets", default="12,16,24,32") args = ap.parse_args() budgets = [int(x) for x in args.budgets.split(",")] dev = "mps" if torch.backends.mps.is_available() else "cpu" print(f"loading {args.model} on {dev} ...") tok = AutoTokenizer.from_pretrained(args.model) model = AutoModelForCausalLM.from_pretrained( args.model, dtype=torch.float16 if dev != "cpu" else torch.float32).to(dev).eval() cfg = model.config n_exp, topk, n_layers = cfg.num_experts, cfg.num_experts_per_tok, cfg.num_hidden_layers print(f"{n_exp} experts, top-{topk}, {n_layers} layers\n") # capture the real per-token top-k selection sequence, per layer, per document # (kept as separate docs so t->t+1 never crosses a document boundary) seqs = [[] for _ in range(n_layers)] # seqs[L] = list of docs; each doc = [set,set,...] for text in CORPUS: ids = tok(text, return_tensors="pt").to(dev) with torch.no_grad(): out = model(**ids, output_router_logits=True) for L in range(n_layers): lg = out.router_logits[L].float().cpu().numpy() if lg.ndim != 2 or lg.shape[1] != n_exp: seqs[L].append([]); continue top = np.argpartition(-lg, topk, axis=1)[:, :topk] seqs[L].append([set(r.tolist()) for r in top]) layers_show = sorted(set([0, 1, n_layers//4, n_layers//2, 3*n_layers//4, n_layers-1])) print(f"{'='*78}\nNEXT-TOKEN EXPERT PREFETCH HIT-RATE (real trace, train/test split)") print(f"top-{topk} of {n_exp} | train=1st half, test=2nd half | hit = fired&prefetched / {topk}") print(f"{'='*78}") hdr = f"{'layer':>5s} " + " ".join(f"B={b:>2d}:LFU/COACT" for b in budgets) print(hdr); print("-" * len(hdr)) out_json = {} agg = {b: [[], []] for b in budgets} # budget -> [lfu_hits, coact_hits] for L in layers_show: docs = seqs[L] # split each doc in half: train on first half tokens, test on second half train, test = [], [] for d in docs: h = len(d) // 2 train.append(d[:h]); test.append(d[h:]) # learn frequency + co-activation on TRAIN only freq = np.zeros(n_exp) co = np.zeros((n_exp, n_exp)) for d in train: for s in d: for e in s: freq[e] += 1 sl = list(s) for i in range(len(sl)): for j in range(len(sl)): if i != j: co[sl[i], sl[j]] += 1 lfu_set = set(np.argsort(-freq)[:max(budgets)].tolist()) # ranked; slice per budget below lfu_rank = np.argsort(-freq) cells = [] for b in budgets: lfu_pred = set(lfu_rank[:b].tolist()) # static frequency prefetch lfu_hits, coact_hits, ncmp = [], [], 0 for d in test: for t in range(len(d) - 1): nxt = d[t + 1] if not nxt: continue # LFU: static top-b by freq lfu_hits.append(len(nxt & lfu_pred) / len(nxt)) # CO-ACT: score experts by co-activation with experts fired at token t score = co[list(d[t])].sum(axis=0) if d[t] else np.zeros(n_exp) # tie-break with global freq so cold-start falls back to LFU, not 0 score = score + 1e-6 * freq coact_pred = set(np.argsort(-score)[:b].tolist()) coact_hits.append(len(nxt & coact_pred) / len(nxt)) ncmp += 1 lh = float(np.mean(lfu_hits)) if lfu_hits else float('nan') ch = float(np.mean(coact_hits)) if coact_hits else float('nan') cells.append(f"{lh:.2f}/{ch:.2f}") agg[b][0].append(lh); agg[b][1].append(ch) out_json.setdefault(f"layer{L}", {})[f"B{b}"] = dict(lfu=lh, coact=ch, n=ncmp) print(f"{L:>5d} " + " ".join(f"{c:>13s}" for c in cells)) print("-" * len(hdr)) means = [] for b in budgets: lm = float(np.nanmean(agg[b][0])); cm = float(np.nanmean(agg[b][1])) means.append(f"{lm:.2f}/{cm:.2f}") out_json.setdefault("mean", {})[f"B{b}"] = dict(lfu=lm, coact=cm) print(f"{'mean':>5s} " + " ".join(f"{m:>13s}" for m in means)) print("=" * len(hdr)) print("Each cell = LFU / CO-ACT next-token prefetch hit-rate at budget B.") print("CO-ACT > LFU -> temporal co-activation is exploitable beyond popularity (real win).") print("CO-ACT ~ LFU -> only expert popularity matters; clustering gain was a mirage.") json.dump(out_json, open("/tmp/moe_cache_sim.json", "w"), indent=2) print("saved /tmp/moe_cache_sim.json") if __name__ == "__main__": main()