1#!/usr/bin/env python32"""Single-node decode/prefill benchmark for TITAN (bounded slot tables) and the mlx-lm baseline.34Writes one JSON record per (backend, config, prompt, repeat) to --out (JSONL, append-only). Each record carries the exact5command, git commit, host, environment versions, timers and cache statistics. Warm-up runs are recorded with warmup=true.6"""7import argparse, json, os, platform, random, socket, subprocess, sys, time, statistics8import mlx.core as mx9import yaml10from pathlib import Path1112sys.path.insert(0, str(Path(__file__).resolve().parents[2]))13from mlx_lm import stream_generate14from mlx_lm.utils import load_model, load_tokenizer15from titan.runtime import TitanModel161718def git_commit():19 try:20 return subprocess.run(["git", "rev-parse", "--short", "HEAD"], capture_output=True, text=True,21 cwd=Path(__file__).resolve().parents[2]).stdout.strip()22 except Exception:23 return None242526def env_record():27 import mlx_lm, numpy28 return {"host": socket.gethostname(), "chip": subprocess.run(["sysctl", "-n", "machdep.cpu.brand_string"], capture_output=True, text=True).stdout.strip(),29 "ram_gb": int(subprocess.run(["sysctl", "-n", "hw.memsize"], capture_output=True, text=True).stdout) // 2**30,30 "macos": platform.mac_ver()[0], "python": platform.python_version(), "mlx": mx.__version__, "mlx_lm": mlx_lm.__version__,31 "numpy": numpy.__version__, "git": git_commit()}323334def run_one(model, tok, prompt, max_tokens):35 """Return timing record for one greedy generation."""36 t0 = time.perf_counter(); t_start_epoch = time.time()37 first = None38 stamps = []39 text = []40 last = None41 for r in stream_generate(model, tok, prompt=prompt, max_tokens=max_tokens, sampler=None):42 now = time.perf_counter()43 if first is None:44 first = now45 stamps.append(now)46 text.append(r.text)47 last = r48 itl = [(b - a) * 1e3 for a, b in zip(stamps, stamps[1:])]49 itl_s = sorted(itl)50 pct = lambda p: itl_s[min(len(itl_s) - 1, int(p * len(itl_s)))] if itl_s else None51 return {"prompt_tokens": last.prompt_tokens, "gen_tokens": last.generation_tokens,52 "ttft_ms": (first - t0) * 1e3, "prompt_tps": last.prompt_tps, "gen_tps": last.generation_tps,53 "decode_tps_wall": (len(stamps) - 1) / (stamps[-1] - stamps[0]) if len(stamps) > 1 else None,54 "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,55 "total_s": stamps[-1] - t0, "peak_mem_gb": last.peak_memory, "text": "".join(text)[:200],56 "t_start_epoch": t_start_epoch, "t_end_epoch": time.time(), "t_first_epoch": t_start_epoch + (first - t0)}575859def main():60 ap = argparse.ArgumentParser()61 ap.add_argument("--model", required=True)62 ap.add_argument("--backend", choices=["titan", "mlx_lm"], default="titan")63 ap.add_argument("--capacity", type=int, default=None)64 ap.add_argument("--policy", default="lru")65 ap.add_argument("--predictor", default="none")66 ap.add_argument("--m", type=int, default=8)67 ap.add_argument("--prefetch-workers", type=int, default=2)68 ap.add_argument("--traces", default=None, help="trace directory used to fit transition/probe predictors (fit split only)")69 ap.add_argument("--nocache", action="store_true", help="F_NOCACHE on the checkpoint (SSD reads bypass page cache)")70 ap.add_argument("--fused-upgate", action="store_true", help="H4: one gather_qmm for up+gate (stacked [C,2H,D] buffer)")71 ap.add_argument("--warm-all", action="store_true", help="preload every expert before timing (all-resident control; implied when capacity == num_experts)")72 ap.add_argument("--workloads", default=str(Path(__file__).resolve().parents[1] / "configs" / "workloads.yaml"))73 ap.add_argument("--families", default="general,code,reasoning,multilingual")74 ap.add_argument("--split", default="fit")75 ap.add_argument("--max-tokens", type=int, default=64)76 ap.add_argument("--repeats", type=int, default=3)77 ap.add_argument("--warmup", type=int, default=1)78 ap.add_argument("--seed", type=int, default=0)79 ap.add_argument("--tag", default="")80 ap.add_argument("--out", required=True)81 a = ap.parse_args()82 wl = yaml.safe_load(open(a.workloads))83 prompts = [(fam, p) for fam in a.families.split(",") for p in wl["families"][fam][a.split]]84 rng = random.Random(a.seed)85 if a.backend == "titan":86 tm = TitanModel(a.model, capacity=a.capacity, policy=a.policy, nocache=a.nocache, predictor=a.predictor, m=a.m,87 prefetch_workers=a.prefetch_workers, traces_dir=a.traces, fused_upgate=a.fused_upgate)88 model, tok = tm.model, tm.tokenizer89 if a.warm_all or tm.capacity >= tm.spec.num_experts:90 tm.warm_all()91 else:92 model, _ = load_model(Path(a.model)); tok = load_tokenizer(Path(a.model)); tm = None93 base = {"backend": a.backend, "capacity": a.capacity if a.backend == "titan" else None, "policy": a.policy if a.backend == "titan" else None,94 "predictor": a.predictor if a.backend == "titan" else None, "m": a.m, "prefetch_workers": a.prefetch_workers, "fused_upgate": a.fused_upgate,95 "nocache": a.nocache, "max_tokens": a.max_tokens, "split": a.split, "tag": a.tag, "model": os.path.realpath(a.model),96 "cmd": " ".join(sys.argv), "env": env_record(), "seed": a.seed}97 os.makedirs(os.path.dirname(a.out), exist_ok=True)98 order = []99 for rep in range(a.warmup + a.repeats):100 ps = prompts[:]101 rng.shuffle(ps)102 order += [(rep, fam, p) for fam, p in ps]103 with open(a.out, "a") as f:104 for rep, fam, p in order:105 if tm: tm.reset_stats()106 mx.reset_peak_memory()107 r = run_one(model, tok, p, a.max_tokens)108 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)109 if tm:110 cs = tm.cache_stats()111 rec.update(cache_hit_rate=cs["hit_rate"], cache_misses=cs["misses"], cache_hits=cs["hits"], stall_s=cs["stall_time_s"],112 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"],113 ssd_bytes=cs["store"]["bytes_read"], ssd_read_s=cs["store"]["read_time_s"])114 f.write(json.dumps(rec) + "\n"); f.flush()115 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"116 + (f" hit={rec['cache_hit_rate']:.3f} stall={rec['stall_s']:.2f}s" if tm else ""))117118119if __name__ == "__main__":120 main()121