| |
| """ |
| MoE expert cache — REAL miss-rate under real eviction policies |
| ============================================================== |
| moe_cache_sim.py measured a next-token prefetch *hit-rate* (a predictor-quality |
| proxy). This goes one level more concrete: an actual capacity-C expert cache with |
| real eviction, replaying the real token trace, counting real demand misses (cold |
| loads) — and, for the prefetch policy, also counting prefetch loads so its |
| bandwidth cost is not hidden. |
| |
| Per layer, test stream of tokens; each token fires top-k of N experts which must be |
| resident. Cache holds C experts. Policies: |
| - LRU : demand-load on miss, evict least-recently-used. |
| - LFU-static : cache = top-C experts by train frequency, pinned, never changes. |
| - CO-ACT prefetch: LRU, but after each token prefetch the B experts most co-activated |
| (learned on train) with the set that just fired, for the next token. |
| Counts BOTH demand misses and prefetch loads as bandwidth. |
| |
| Honest metric: demand miss-rate = cold loads / expert-activations. Lower = better. |
| total-load-rate (prefetch policy only) exposes the bandwidth it spends to get there. |
| If CO-ACT cuts demand misses below LFU/LRU *without* blowing up total loads, the |
| co-activation structure pays for itself. If it only wins by loading everything, it |
| doesn't. train/test split, same trace as moe_cache_sim.py. |
| |
| Usage: python moe_cache_miss.py [--model allenai/OLMoE-1B-7B-0924] [--caps 16,24,32] |
| """ |
| import argparse, json |
| import numpy as np |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| from moe_cache_sim import CORPUS |
|
|
|
|
| def lru_sim(test, cap, n_exp): |
| """Demand-paged LRU. Returns (demand_misses, activations).""" |
| cache = {} |
| clock = 0; misses = 0; acts = 0 |
| for d in test: |
| for s in d: |
| for e in s: |
| acts += 1; clock += 1 |
| if e in cache: |
| cache[e] = clock |
| else: |
| misses += 1 |
| if len(cache) >= cap: |
| ev = min(cache, key=cache.get); del cache[ev] |
| cache[e] = clock |
| return misses, acts |
|
|
|
|
| def lfu_static_sim(test, pinned, n_exp): |
| misses = 0; acts = 0 |
| for d in test: |
| for s in d: |
| for e in s: |
| acts += 1 |
| if e not in pinned: |
| misses += 1 |
| return misses, acts |
|
|
|
|
| def coact_prefetch_sim(test, cap, co, freq, B, n_exp): |
| """LRU + co-activation prefetch. Returns (demand_misses, prefetch_loads, acts).""" |
| cache = {}; clock = 0; misses = 0; pf_loads = 0; acts = 0 |
| score_base = 1e-6 * freq |
| for d in test: |
| prev = None |
| for s in d: |
| |
| if prev: |
| score = co[list(prev)].sum(axis=0) + score_base |
| for e in np.argsort(-score)[:B]: |
| e = int(e) |
| clock += 1 |
| if e not in cache: |
| pf_loads += 1 |
| if len(cache) >= cap: |
| ev = min(cache, key=cache.get); del cache[ev] |
| cache[e] = clock |
| else: |
| cache[e] = clock |
| for e in s: |
| acts += 1; clock += 1 |
| if e in cache: |
| cache[e] = clock |
| else: |
| misses += 1 |
| if len(cache) >= cap: |
| ev = min(cache, key=cache.get); del cache[ev] |
| cache[e] = clock |
| prev = s |
| return misses, pf_loads, acts |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", default="allenai/OLMoE-1B-7B-0924") |
| ap.add_argument("--caps", default="16,24,32") |
| args = ap.parse_args() |
| caps = [int(x) for x in args.caps.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") |
|
|
| seqs = [[] for _ in range(n_layers)] |
| 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"{'='*86}\nREAL EXPERT CACHE MISS-RATE (demand misses / activation; lower=better)") |
| print(f"top-{topk} of {n_exp} | train=1st half learns freq+coact, test=2nd half replayed") |
| print(f"{'='*86}") |
| hdr = f"{'layer':>5s} {'cap':>4s} {'LRU':>7s} {'LFU-st':>7s} {'COACT':>7s} {'COACT-tot':>9s}" |
| print(hdr); print("-" * len(hdr)) |
|
|
| out_json = {} |
| agg = {c: {'lru': [], 'lfu': [], 'coact': [], 'tot': []} for c in caps} |
| for L in layers_show: |
| docs = seqs[L] |
| train = [d[:len(d)//2] for d in docs] |
| test = [d[len(d)//2:] for d in docs] |
| 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 |
| rank = np.argsort(-freq) |
| for c in caps: |
| lm, la = lru_sim(test, c, n_exp) |
| fm, fa = lfu_static_sim(test, set(rank[:c].tolist()), n_exp) |
| cm, pf, ca = coact_prefetch_sim(test, c, co, freq, topk, n_exp) |
| lru_r = lm/la; lfu_r = fm/fa; co_r = cm/ca; tot_r = (cm+pf)/ca |
| print(f"{L:>5d} {c:>4d} {lru_r:>7.3f} {lfu_r:>7.3f} {co_r:>7.3f} {tot_r:>9.3f}") |
| agg[c]['lru'].append(lru_r); agg[c]['lfu'].append(lfu_r) |
| agg[c]['coact'].append(co_r); agg[c]['tot'].append(tot_r) |
| out_json.setdefault(f"layer{L}", {})[f"cap{c}"] = dict( |
| lru=lru_r, lfu_static=lfu_r, coact_demand=co_r, coact_total=tot_r) |
| print("-" * len(hdr)) |
| for c in caps: |
| m = {k: float(np.mean(v)) for k, v in agg[c].items()} |
| print(f"{'mean':>5s} {c:>4d} {m['lru']:>7.3f} {m['lfu']:>7.3f} {m['coact']:>7.3f} {m['tot']:>9.3f}") |
| out_json.setdefault("mean", {})[f"cap{c}"] = dict( |
| lru=m['lru'], lfu_static=m['lfu'], coact_demand=m['coact'], coact_total=m['tot']) |
| print("=" * len(hdr)) |
| print("LRU/LFU-st/COACT = demand miss-rate (cold loads / activation). COACT-tot adds") |
| print("the prefetch loads (bandwidth COACT spends). COACT < LFU on demand AND COACT-tot") |
| print("not far above it => co-activation prefetch pays for itself. Else it doesn't.") |
| json.dump(out_json, open("/tmp/moe_cache_miss.json", "w"), indent=2) |
| print("saved /tmp/moe_cache_miss.json") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
