#!/usr/bin/env python3 """Single-node decode/prefill benchmark for TITAN (bounded slot tables) and the mlx-lm baseline. Writes one JSON record per (backend, config, prompt, repeat) to --out (JSONL, append-only). Each record carries the exact command, git commit, host, environment versions, timers and cache statistics. Warm-up runs are recorded with warmup=true. """ import argparse, json, os, platform, random, socket, subprocess, sys, time, statistics import mlx.core as mx import yaml from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from mlx_lm import stream_generate from mlx_lm.utils import load_model, load_tokenizer from titan.runtime import TitanModel def git_commit(): try: return subprocess.run(["git", "rev-parse", "--short", "HEAD"], capture_output=True, text=True, cwd=Path(__file__).resolve().parents[2]).stdout.strip() except Exception: return None def env_record(): import mlx_lm, numpy return {"host": socket.gethostname(), "chip": subprocess.run(["sysctl", "-n", "machdep.cpu.brand_string"], capture_output=True, text=True).stdout.strip(), "ram_gb": int(subprocess.run(["sysctl", "-n", "hw.memsize"], capture_output=True, text=True).stdout) // 2**30, "macos": platform.mac_ver()[0], "python": platform.python_version(), "mlx": mx.__version__, "mlx_lm": mlx_lm.__version__, "numpy": numpy.__version__, "git": git_commit()} def run_one(model, tok, prompt, max_tokens): """Return timing record for one greedy generation.""" t0 = time.perf_counter() first = None stamps = [] text = [] last = None for r in stream_generate(model, tok, prompt=prompt, max_tokens=max_tokens, sampler=None): now = time.perf_counter() if first is None: first = now stamps.append(now) text.append(r.text) last = r itl = [(b - a) * 1e3 for a, b in zip(stamps, stamps[1:])] itl_s = sorted(itl) pct = lambda p: itl_s[min(len(itl_s) - 1, int(p * len(itl_s)))] if itl_s else None return {"prompt_tokens": last.prompt_tokens, "gen_tokens": last.generation_tokens, "ttft_ms": (first - t0) * 1e3, "prompt_tps": last.prompt_tps, "gen_tps": last.generation_tps, "decode_tps_wall": (len(stamps) - 1) / (stamps[-1] - stamps[0]) if len(stamps) > 1 else None, "itl_ms_p50": pct(0.5), "itl_ms_p95": pct(0.95), "itl_ms_p99": pct(0.99), "itl_ms_mean": statistics.mean(itl) if itl else None, "total_s": stamps[-1] - t0, "peak_mem_gb": last.peak_memory, "text": "".join(text)[:200]} def main(): ap = argparse.ArgumentParser() ap.add_argument("--model", required=True) ap.add_argument("--backend", choices=["titan", "mlx_lm"], default="titan") ap.add_argument("--capacity", type=int, default=None) ap.add_argument("--policy", default="lru") ap.add_argument("--predictor", default="none") ap.add_argument("--m", type=int, default=8) ap.add_argument("--prefetch-workers", type=int, default=2) ap.add_argument("--traces", default=None, help="trace directory used to fit transition/probe predictors (fit split only)") ap.add_argument("--nocache", action="store_true", help="F_NOCACHE on the checkpoint (SSD reads bypass page cache)") ap.add_argument("--fused-upgate", action="store_true", help="H4: one gather_qmm for up+gate (stacked [C,2H,D] buffer)") ap.add_argument("--warm-all", action="store_true", help="preload every expert before timing (all-resident control; implied when capacity == num_experts)") ap.add_argument("--workloads", default=str(Path(__file__).resolve().parents[1] / "configs" / "workloads.yaml")) ap.add_argument("--families", default="general,code,reasoning,multilingual") ap.add_argument("--split", default="fit") ap.add_argument("--max-tokens", type=int, default=64) ap.add_argument("--repeats", type=int, default=3) ap.add_argument("--warmup", type=int, default=1) ap.add_argument("--seed", type=int, default=0) ap.add_argument("--tag", default="") ap.add_argument("--out", required=True) a = ap.parse_args() wl = yaml.safe_load(open(a.workloads)) prompts = [(fam, p) for fam in a.families.split(",") for p in wl["families"][fam][a.split]] rng = random.Random(a.seed) if a.backend == "titan": tm = TitanModel(a.model, capacity=a.capacity, policy=a.policy, nocache=a.nocache, predictor=a.predictor, m=a.m, prefetch_workers=a.prefetch_workers, traces_dir=a.traces, fused_upgate=a.fused_upgate) model, tok = tm.model, tm.tokenizer if a.warm_all or tm.capacity >= tm.spec.num_experts: tm.warm_all() else: model, _ = load_model(Path(a.model)); tok = load_tokenizer(Path(a.model)); tm = None base = {"backend": a.backend, "capacity": a.capacity if a.backend == "titan" else None, "policy": a.policy if a.backend == "titan" else None, "predictor": a.predictor if a.backend == "titan" else None, "m": a.m, "prefetch_workers": a.prefetch_workers, "fused_upgate": a.fused_upgate, "nocache": a.nocache, "max_tokens": a.max_tokens, "split": a.split, "tag": a.tag, "model": os.path.realpath(a.model), "cmd": " ".join(sys.argv), "env": env_record(), "seed": a.seed} os.makedirs(os.path.dirname(a.out), exist_ok=True) order = [] for rep in range(a.warmup + a.repeats): ps = prompts[:] rng.shuffle(ps) order += [(rep, fam, p) for fam, p in ps] with open(a.out, "a") as f: for rep, fam, p in order: if tm: tm.reset_stats() mx.reset_peak_memory() r = run_one(model, tok, p, a.max_tokens) rec = dict(base, warmup=rep < a.warmup, repeat=rep - a.warmup, family=fam, prompt=p[:60], ts=time.strftime("%Y-%m-%dT%H:%M:%S"), **r) if tm: cs = tm.cache_stats() rec.update(cache_hit_rate=cs["hit_rate"], cache_misses=cs["misses"], cache_hits=cs["hits"], stall_s=cs["stall_time_s"], hits_device=sum(p["hits_device"] for p in cs["layers"]), layers_complete=sum(1 for t in tm.tables if t.complete()), prefetch=cs["prefetch"], host_time_s=cs["host_time_s"], ssd_bytes=cs["store"]["bytes_read"], ssd_read_s=cs["store"]["read_time_s"]) f.write(json.dumps(rec) + "\n"); f.flush() print(f"{a.backend} cap={a.capacity} {a.predictor} rep={rep} {fam:12s} ttft={r['ttft_ms']:.0f}ms prefill={r['prompt_tps']:.0f} decode={r['gen_tps']:.1f} tok/s" + (f" hit={rec['cache_hit_rate']:.3f} stall={rec['stall_s']:.2f}s" if tm else "")) if __name__ == "__main__": main()