1#!/usr/bin/env python32"""Offline, trace-driven simulation of per-layer slot tables with eviction policies and prefetch predictors.34Exact for hit/miss accounting (the router sequence is the recorded one). Stall time is *modelled* from measured constants:5t_read (seconds to fetch one expert from the tier) and t_layer (seconds of compute per layer), with a loader pool of `--loaders`6parallel readers: demand misses at a layer stall for ceil(misses/loaders) * t_read; a prefetch issued h layers ahead hides7h * t_layer of its latency. Outputs JSON + CSV under --out.8"""9import argparse, glob, json, math, os, sys, time10import numpy as np11from pathlib import Path12sys.path.insert(0, str(Path(__file__).resolve().parents[2]))13from titan.predictors import make as make_predictor14from titan.slots import EvictionPolicy, LFUPolicy1516POLICIES = {"lru": EvictionPolicy, "lfu": LFUPolicy}171819def load_traces(d):20 idx = json.load(open(os.path.join(d, "index.json")))21 out = []22 for p in idx["prompts"]:23 z = np.load(os.path.join(d, p["name"] + ".npz"))24 out.append(dict(p, topk=z["topk"].astype(np.int64), probs=z["probs"], phase=z["phase"]))25 return idx, out262728class LayerSim:29 def __init__(self, cap, policy, E):30 self.cap = cap; self.pol = POLICIES[policy](); self.res = set(); self.E = E31 self.hits = self.misses = self.pf_loads = self.pf_used = self.demand_loads = self.evict = 032 self.pf_pending = {} # expert -> True if prefetched and not yet used33 self.late_pending = 03435 def _insert(self, e, protected):36 if e in self.res: return37 if len(self.res) >= self.cap:38 try: v = self.pol.victim(protected)39 except RuntimeError: return40 self.res.discard(v); self.pol.remove(v); self.evict += 141 if self.pf_pending.pop(v, False): pass42 self.res.add(e)4344 def prefetch(self, experts, protected):45 for e in experts:46 if e in self.res: continue47 if len(self.res) >= self.cap and all(x in protected for x in self.res): break48 self._insert(e, protected); self.pol.touch(e); self.pf_loads += 1; self.pf_pending[e] = True4950 def access(self, needed):51 miss = 052 for e in needed:53 if e in self.res:54 self.hits += 155 if self.pf_pending.pop(e, False): self.pf_used += 156 else:57 self.misses += 1; miss += 1; self.demand_loads += 158 self._insert(e, set(needed))59 self.pol.touch(e)60 return miss616263_PRED_CACHE = {}6465def simulate(traces, L, E, K, cap, policy, pred_name, m, t_read, t_layer, loaders, fit_traces):66 if pred_name not in _PRED_CACHE: # fit each predictor once (on the fit split) and reuse across capacities/splits67 pred = make_predictor(pred_name, L, E)68 if fit_traces is not None: pred.fit([(t["topk"], t["probs"]) for t in fit_traces])69 _PRED_CACHE[pred_name] = pred70 pred = _PRED_CACHE[pred_name]71 depth = getattr(pred, "depth", 1)72 tot = {"tokens": 0, "hits": 0, "misses": 0, "pf_loads": 0, "pf_used": 0, "demand_loads": 0, "stall_s": 0.0, "pred_hits": 0, "pred_total": 0, "evictions": 0}73 per_family = {}74 for tr in traces:75 layers = [LayerSim(cap, policy, E) for _ in range(L)] # cold per prompt76 topk, probs = tr["topk"], tr["probs"]77 prev = None; prev_p = None78 fam = per_family.setdefault(tr["family"], {"tokens": 0, "hits": 0, "misses": 0, "stall_s": 0.0})79 for t in range(topk.shape[0]):80 cur = {("true", l): topk[t, l] for l in range(L)} # read ONLY by the Oracle upper bound81 stall = 0.082 if pred.horizon == "token":83 for l in range(L):84 layers[l].prefetch(pred.predict(l, prev, prev_p, cur, m), set())85 for l in range(L):86 need = list(topk[t, l])87 if pred.horizon == "layer" and l > 0:88 pass # already issued at l-189 miss = layers[l].access(need)90 if miss:91 # demand misses: a layer-horizon prefetch for this layer was issued `depth` layers earlier and overlapped depth*t_layer;92 # only misses that were *prefetched-but-late* benefit, so we model hidden = depth*t_layer for the late ones, 0 for unpredicted.93 late = min(miss, layers[l].late_pending)94 unpred = miss - late95 stall += (math.ceil(late / loaders) * max(0.0, t_read - depth * t_layer) + math.ceil(unpred / loaders) * t_read) if late + unpred else 0.096 layers[l].late_pending = 097 cur[l] = topk[t, l]; cur[("probs", l)] = probs[t, l]98 if pred.horizon == "layer" and l + depth < L:99 p = pred.predict(l + depth, prev, prev_p, cur, m)100 if pred_name != "none":101 tot["pred_hits"] += len(set(p[:K]) & set(topk[t, l + depth])); tot["pred_total"] += K102 layers[l + depth].prefetch(p, set())103 layers[l + depth].late_pending = len(set(p) & set(topk[t, l + depth]) - layers[l + depth].res)104 prev, prev_p = topk[t], probs[t]105 tot["tokens"] += 1; fam["tokens"] += 1; tot["stall_s"] += stall; fam["stall_s"] += stall106 h = sum(x.hits for x in layers); mi = sum(x.misses for x in layers)107 tot["hits"] += sum(x.hits for x in layers); tot["misses"] += sum(x.misses for x in layers)108 tot["pf_loads"] += sum(x.pf_loads for x in layers); tot["pf_used"] += sum(x.pf_used for x in layers)109 tot["demand_loads"] += sum(x.demand_loads for x in layers); tot["evictions"] += sum(x.evict for x in layers)110 fam["hits"] += sum(x.hits for x in layers); fam["misses"] += sum(x.misses for x in layers)111 n = tot["hits"] + tot["misses"]112 res = {"capacity": cap, "policy": policy, "predictor": pred_name, "m": m, "hit_rate": tot["hits"] / n, "misses_per_token": tot["misses"] / tot["tokens"],113 "stall_ms_per_token": 1e3 * tot["stall_s"] / tot["tokens"], "prefetch_loads_per_token": tot["pf_loads"] / tot["tokens"],114 "prefetch_precision": (tot["pf_used"] / tot["pf_loads"]) if tot["pf_loads"] else None,115 "loads_per_token": (tot["pf_loads"] + tot["demand_loads"]) / tot["tokens"],116 "pred_recall_at_k": (tot["pred_hits"] / tot["pred_total"]) if tot["pred_total"] else None,117 "tokens": tot["tokens"], "families": {f: {"hit_rate": v["hits"] / (v["hits"] + v["misses"]), "stall_ms_per_token": 1e3 * v["stall_s"] / v["tokens"]} for f, v in per_family.items()}}118 return res119120121def main():122 ap = argparse.ArgumentParser()123 ap.add_argument("--traces", required=True)124 ap.add_argument("--capacities", default="8,12,16,24,32,48")125 ap.add_argument("--policies", default="lru,lfu")126 ap.add_argument("--predictors", default="none,prev_rank,transition,transition2,transition3,probe,oracle")127 ap.add_argument("--m", type=int, default=8)128 ap.add_argument("--t-read-ms", type=float, default=0.43, help="measured per-expert fetch latency of the tier (laptop: 0.43 ms cold SSD, 0.27 ms page cache, 28 ms 1 GbE)")129 ap.add_argument("--t-layer-ms", type=float, default=0.128, help="per-layer decode compute (489 tok/s / 16 layers)")130 ap.add_argument("--loaders", type=int, default=1, help="parallel readers (current TITAN implementation: 1)")131 ap.add_argument("--out", required=True)132 a = ap.parse_args()133 idx, traces = load_traces(a.traces)134 L, E, K = idx["L"], idx["E"], idx["K"]135 fit = [t for t in traces if t["split"] == "fit"]; held = [t for t in traces if t["split"] == "heldout"]136 results = []137 for cap in [int(c) for c in a.capacities.split(",")]:138 for pol in a.policies.split(","):139 for pr in a.predictors.split(","):140 for split, trs in (("fit", fit), ("heldout", held)):141 r = simulate(trs, L, E, K, cap, pol, pr, a.m, a.t_read_ms / 1e3, a.t_layer_ms / 1e3, a.loaders, fit)142 r["split"] = split; results.append(r)143 print(f"cap={cap:2d} {pol} {pr:10s} {split:7s} hit={r['hit_rate']:.3f} miss/tok={r['misses_per_token']:5.1f} stall={r['stall_ms_per_token']:6.2f}ms "144 f"loads/tok={r['loads_per_token']:5.1f} pf_prec={r['prefetch_precision'] if r['prefetch_precision'] is None else round(r['prefetch_precision'],3)} recall@k={r['pred_recall_at_k'] if r['pred_recall_at_k'] is None else round(r['pred_recall_at_k'],3)}")145 os.makedirs(os.path.dirname(a.out), exist_ok=True)146 json.dump({"cmd": " ".join(sys.argv), "ts": time.strftime("%Y-%m-%dT%H:%M:%S"), "traces": a.traces, "L": L, "E": E, "K": K,147 "t_read_ms": a.t_read_ms, "t_layer_ms": a.t_layer_ms, "loaders": a.loaders, "results": results}, open(a.out, "w"), indent=1)148149150if __name__ == "__main__":151 main()152