Something went wrong. Try again.
strata but optimized for m5 (and maybe more later)
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324#!/usr/bin/env python3"""Offline MoE expert cache & MoCE closed-loop replayer for Qwen3.8-Flash-Next (qwen4exp).Implements: - Open-loop and MoCE cache-prior reranking: z'_e = z_e + lambda_moce * delta_avg * 1[e in cache or e in original top-J] - Optional --keep-map pruning filter (filters top-64 trace candidates to kept experts) - Eviction policies: * lru: Least Recently Used * lfu: Least Frequently Used * ema: Popularity EMA on selected top-K experts (all n_experts decay every token) * near_pick_ema: Popularity EMA crediting unbiased router top-N (ds4-recipe doc 06) * inverse_router_ema: EMA of normalized router rank/logit utility (dawn's inverse-router idea) * belady_min: Optimal offline oracle (Belady's MIN algorithm, theoretical ceiling) - M5 Air latency prediction using measured hardware constants from bench/README.md."""import argparseimport jsonimport mathimport structfrom pathlib import Path
TRACE_MAGIC = b"Q4RT" # Qwen4exp Routing Trace v1
def parse_trace_file(path: str, keep_map_path: str | None = None) -> dict: raw = Path(path).read_bytes() magic, ver, n_layers, n_experts, top_m, n_tokens = struct.unpack_from("<4sHHHH I", raw, 0) if magic != TRACE_MAGIC: raise ValueError(f"Bad trace magic: {magic}") rec_bytes = top_m * 2 + top_m * 2 actual_tokens = (len(raw) - 16) // (n_layers * rec_bytes) if n_tokens == 0 or actual_tokens < n_tokens: n_tokens = actual_tokens
keep_sets = None remap = None if keep_map_path: km = json.loads(Path(keep_map_path).read_text()) keep_sets = [] remap = [] for l in range(n_layers): kept = km[str(l)] keep_sets.append(set(kept)) remap.append({old_id: new_id for new_id, old_id in enumerate(kept)}) n_experts = len(km["0"])
layers_data = [[] for _ in range(n_layers)] off = 16 ids_fmt = f"<{top_m}H" logits_fmt = f"<{top_m}e" id_bytes = top_m * 2 for _t in range(n_tokens): for l in range(n_layers): ids = struct.unpack_from(ids_fmt, raw, off) logits = struct.unpack_from(logits_fmt, raw, off + id_bytes) off += rec_bytes if keep_sets is not None: ks = keep_sets[l] rm = remap[l] f_ids = [] f_logits = [] for eid, z in zip(ids, logits): if eid in ks: f_ids.append(rm[eid]) f_logits.append(z) ids = tuple(f_ids) logits = tuple(f_logits) layers_data[l].append((ids, logits))
return { "version": ver, "n_layers": n_layers, "n_experts": n_experts, "top_m": top_m, "n_tokens": n_tokens, "layers_data": layers_data, }
def select_moce_top_k( ids: tuple[int, ...], logits: tuple[float, ...], cache_set: set[int], k_used: int = 10, lambda_moce: float = 0.0, protect_top_j: int = 3,) -> list[int]: if not ids: return [] if lambda_moce <= 0.0 or len(ids) <= k_used: return list(ids[:k_used])
delta_avg = max(1e-4, logits[0] - logits[min(k_used - 1, len(logits) - 1)]) bonus = lambda_moce * delta_avg scored = [] for rank, (eid, z) in enumerate(zip(ids, logits)): if rank < protect_top_j or eid in cache_set: scored.append((z + bonus, -rank, eid)) else: scored.append((z, -rank, eid)) scored.sort(reverse=True) return [eid for _, _, eid in scored[:k_used]]
def replay_layer( layer_trace: list[tuple[tuple[int, ...], tuple[float, ...]]], n_experts: int, cache_cap: int, policy: str = "lru", k_used: int = 10, lambda_moce: float = 0.0, protect_top_j: int = 3, near_pick_n: int = 48, ema_decay: float = 0.99, warmup_tokens: int = 64,) -> dict: if cache_cap < k_used: raise ValueError(f"cache_cap ({cache_cap}) must be >= k_used ({k_used})") if len(layer_trace) <= warmup_tokens: raise ValueError(f"trace length ({len(layer_trace)}) must exceed warmup_tokens ({warmup_tokens})") if cache_cap >= n_experts: eval_tokens = max(1, len(layer_trace) - warmup_tokens) return {"misses": 0, "eval_tokens": eval_tokens, "misses_per_tok": 0.0, "hit_rate": 1.0}
next_use_lists: list[list[int]] = [[] for _ in range(n_experts)] next_use_ptr = [0] * n_experts if policy == "belady_min": for t_idx, (ids, _logits) in enumerate(layer_trace): for eid in ids[:k_used]: next_use_lists[eid].append(t_idx)
cache_set: set[int] = set() last_used = [0] * n_experts freq = [0] * n_experts ema_score = [0.0] * n_experts ema_mult = 1.0 inv_decay = 1.0 / ema_decay
misses = 0 total_picks = 0
for t_idx, (ids, logits) in enumerate(layer_trace): chosen = select_moce_top_k( ids, logits, cache_set, k_used=k_used, lambda_moce=lambda_moce, protect_top_j=protect_top_j ) if policy == "belady_min": for eid in chosen: while next_use_ptr[eid] < len(next_use_lists[eid]) and next_use_lists[eid][next_use_ptr[eid]] <= t_idx: next_use_ptr[eid] += 1
chosen_set = set(chosen) if t_idx >= warmup_tokens: missing_slots = max(0, k_used - len(chosen)) total_picks += len(chosen) + missing_slots misses += missing_slots for eid in chosen: if eid not in cache_set: if t_idx >= warmup_tokens: misses += 1 if len(cache_set) >= cache_cap: candidates = [c for c in cache_set if c not in chosen_set] if not candidates: candidates = list(cache_set) if policy == "lru": victim = min(candidates, key=lambda c: (last_used[c], c)) elif policy == "lfu": victim = min(candidates, key=lambda c: (freq[c], last_used[c], c)) elif policy in ("ema", "near_pick_ema", "inverse_router_ema"): victim = min(candidates, key=lambda c: (ema_score[c], last_used[c], c)) elif policy == "belady_min": def next_time(c: int) -> tuple[int, int]: ptr = next_use_ptr[c] lst = next_use_lists[c] return (lst[ptr] if ptr < len(lst) else 10**9, -c) victim = max(candidates, key=next_time) else: raise ValueError(f"Unknown policy: {policy}") cache_set.remove(victim) cache_set.add(eid)
last_used[eid] = t_idx + 1 freq[eid] += 1
if policy in ("ema", "near_pick_ema", "inverse_router_ema"): # Lazy global decay: multiplying added credits by inv_decay^t is identical to # decaying all n_experts every step, and avoids O(n_experts) inner loops. ema_mult *= inv_decay if ema_mult > 1e100: scale = 1.0 / ema_mult for e in range(n_experts): ema_score[e] *= scale ema_mult = 1.0
if policy == "ema": for eid in chosen: ema_score[eid] += ema_mult elif policy == "near_pick_ema": top_n = min(near_pick_n, len(ids)) denom = max(1, top_n - k_used) for rank in range(top_n): eid = ids[rank] credit = 1.0 if rank < k_used else 0.35 * (1.0 - (rank - k_used) / denom) ema_score[eid] += credit * ema_mult elif policy == "inverse_router_ema" and logits: z_max = logits[0] z_min = logits[-1] span = max(1e-4, z_max - z_min) for rank, (eid, z) in enumerate(zip(ids, logits)): util = (z - z_min) / span if rank < k_used: util += 1.0 ema_score[eid] += util * ema_mult
eval_tokens = max(1, len(layer_trace) - warmup_tokens) misses_per_tok = misses / eval_tokens hit_rate = 1.0 - (misses / max(1, total_picks)) return { "misses": misses, "eval_tokens": eval_tokens, "misses_per_tok": misses_per_tok, "hit_rate": hit_rate, }
def run_self_test() -> None: import random rng = random.Random(42) n_experts = 64 top_m = 24 n_tokens = 400 weights = [1.0 / ((i + 1) ** 0.85) for i in range(n_experts)] trace = [] hot_burst = [0, 1, 2, 3, 4] for t in range(n_tokens): if t % 20 == 0: hot_burst = rng.sample(range(n_experts), 6) scored = [] for e in range(n_experts): base = math.log(weights[e]) + rng.gauss(0, 0.8) if e in hot_burst: base += 1.5 scored.append((base, e)) scored.sort(reverse=True) ids = tuple(e for _, e in scored[:top_m]) logits = tuple(float(z) for z, _ in scored[:top_m]) trace.append((ids, logits))
res = {} for pol in ["lru", "lfu", "ema", "near_pick_ema", "inverse_router_ema", "belady_min"]: r = replay_layer(trace, n_experts=n_experts, cache_cap=24, policy=pol, k_used=6, warmup_tokens=40) res[pol] = r["misses"]
assert res["belady_min"] <= min(res["lru"], res["lfu"], res["ema"], res["near_pick_ema"], res["inverse_router_ema"]), res r_moce = replay_layer(trace, n_experts=n_experts, cache_cap=24, policy="near_pick_ema", k_used=6, lambda_moce=0.15, warmup_tokens=40) assert r_moce["misses"] < res["near_pick_ema"], (r_moce["misses"], res["near_pick_ema"]) print("replay_moe_cache.py self-test: PASS", res, "moce:", r_moce["misses"])
def main(): p = argparse.ArgumentParser(description="Replay MoE routing traces across cache policies and MoCE.") p.add_argument("--self-test", action="store_true") p.add_argument("--trace", help="Path to Q4RT binary trace file.") p.add_argument("--keep-map", default=None, help="Optional keep_map.json for pruned expert replay.") p.add_argument("--cache-cap", type=int, default=113, help="Resident experts per layer (113 = M5 Air 7.5 GB Q2_0).") p.add_argument("--k-used", type=int, default=10) p.add_argument("--lambda-moce", type=float, default=0.0) p.add_argument("--protect-top-j", type=int, default=3) p.add_argument("--ema-decay", type=float, default=0.99) p.add_argument("--near-pick-n", type=int, default=48) p.add_argument("--policies", default="lru,lfu,ema,near_pick_ema,inverse_router_ema,belady_min") args = p.parse_args()
if args.self_test: run_self_test() return
if not args.trace: p.error("--trace is required unless --self-test is set")
tr = parse_trace_file(args.trace, keep_map_path=args.keep_map) policies = [x.strip() for x in args.policies.split(",") if x.strip()] summary = { "n_tokens": tr["n_tokens"], "n_experts": tr["n_experts"], "cache_cap": args.cache_cap, "lambda_moce": args.lambda_moce, "protect_top_j": args.protect_top_j, "policies": {}, } for pol in policies: if pol == "belady_min" and args.lambda_moce > 0.0: continue total_misses_per_tok = 0.0 hit_rates = [] for l in range(tr["n_layers"]): r = replay_layer( tr["layers_data"][l], n_experts=tr["n_experts"], cache_cap=args.cache_cap, policy=pol, k_used=args.k_used, lambda_moce=args.lambda_moce, protect_top_j=args.protect_top_j, near_pick_n=args.near_pick_n, ema_decay=args.ema_decay, ) total_misses_per_tok += r["misses_per_tok"] hit_rates.append(r["hit_rate"]) mean_hit = sum(hit_rates) / len(hit_rates) miss_mb_per_tok = total_misses_per_tok * 1.3824 summary["policies"][pol] = { "total_misses_per_token_48L": round(total_misses_per_tok, 2), "mean_hit_rate": round(mean_hit, 4), "ssd_mb_per_token_q2_0": round(miss_mb_per_tok, 2), "est_m5_air_decode_tok_s_1th_3.9GBs": round(1.0 / (miss_mb_per_tok / 3900.0 + 0.012), 2), "est_m5_air_decode_tok_s_4th_6.8GBs": round(1.0 / (miss_mb_per_tok / 6800.0 + 0.012), 2), } print(json.dumps(summary, indent=2))
if __name__ == "__main__": main()