#!/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 argparse
import json
import math
import struct
from 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()