moe_cache_sim.py

#!/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()