1#!/usr/bin/env python2"""Batch-1 (and batched) end-to-end decode latency per token for budget-0, forced budgets, dense writes and the3learned controller, with warm-up and device synchronisation. Reports median / p95 per token and peak memory.45Usage: python -m evaluation.latency --config configs/experiments/routing_p4.yaml --prompt-len 64 --new 32 [--device mps]6"""7from __future__ import annotations89import argparse, json, sys, time10from pathlib import Path1112import numpy as np13import torch1415sys.path.insert(0, str(Path(__file__).resolve().parents[1]))16from prisme.cache import RequestBatchCache # noqa: E40217from prisme.generation import step_logits # noqa: E40218from prisme.flops import action_cost_proxy, ops_per_token # noqa: E40219from training.common import load_checkpoint, load_yaml, pick_device, sync, peak_memory_bytes, git_sha # noqa: E40220from training.train_editor import build_model # noqa: E402212223@torch.no_grad()24def time_decode(model, cfg, device, B, prompt_len, n_new, budget, forced_k=None, cost=None, reps=3):25 times = []26 cache_bytes = 027 for rep in range(reps + 1):28 cache = RequestBatchCache(layers=model.backbone.new_cache())29 prompt = torch.randint(0, cfg.vocab_size, (B, prompt_len), device=device)30 fa = None31 if forced_k is not None:32 a = 0 if forced_k == 1 else cfg.n_slots33 fa = torch.full((B, prompt_len, budget), a, device=device, dtype=torch.long)34 step_logits(model, prompt, cache, budget, cost=cost, forced_actions=fa)35 sync(device)36 tok = prompt[:, -1:]37 per_tok = []38 for i in range(n_new):39 fa1 = None if forced_k is None else torch.full((B, 1, budget), 0 if forced_k == 1 else cfg.n_slots, device=device, dtype=torch.long)40 t0 = time.perf_counter()41 lg, _, _ = step_logits(model, tok, cache, budget, cost=cost, forced_actions=fa1)42 sync(device)43 per_tok.append(time.perf_counter() - t0)44 tok = lg[:, -1].argmax(-1, keepdim=True)45 if rep > 0:46 times.extend(per_tok)47 cache_bytes = cache.nbytes()48 t = np.array(times) * 100049 return {"median_ms": float(np.median(t)), "p95_ms": float(np.percentile(t, 95)), "mean_ms": float(t.mean()), "n": len(t), "cache_bytes": cache_bytes}505152def main(argv=None) -> int:53 ap = argparse.ArgumentParser(description=__doc__)54 ap.add_argument("--config", required=True)55 ap.add_argument("--device", default="auto")56 ap.add_argument("--prompt-len", type=int, default=64)57 ap.add_argument("--new", type=int, default=32)58 ap.add_argument("--batch-sizes", default="1,8")59 ap.add_argument("--out", default=None)60 args = ap.parse_args(argv)61 exp = load_yaml(args.config); device = pick_device(args.device)62 torch.manual_seed(0)63 cfg, model = build_model(exp); model = model.to(device).eval()64 meta = load_checkpoint(exp["checkpoint"], model, strict=False)65 T = args.prompt_len + args.new66 cost = torch.tensor(action_cost_proxy(cfg, T), device=device)67 recs = []68 for B in [int(b) for b in args.batch_sizes.split(",")]:69 def add(name, **kw):70 r = time_decode(model, cfg, device, B, args.prompt_len, args.new, **kw)71 rec = {"policy": name, "batch": B, **{k: v for k, v in kw.items() if k != "cost"}, **r}72 recs.append(rec); print(json.dumps(rec), flush=True)73 add("budget0", budget=0)74 for b in [1, 2, 4, 8]:75 add("forced_single", budget=b, forced_k=1)76 add("forced_pair", budget=b, forced_k=2)77 add("controller", budget=8, cost=cost)78 if not cfg.dense_writes:79 cfg.dense_writes = True80 for b in [1, 4, 8]:81 add("dense_forced", budget=b, forced_k=1)82 cfg.dense_writes = False83 analytic = {name: ops_per_token(cfg, T, n, k).total for name, n, k in [("budget0", 0, 1), ("4_k1", 4, 1), ("4_k2", 4, 2), ("8_k2", 8, 2), ("8_dense", 8, cfg.n_slots)]}84 out = {"config": args.config, "checkpoint_sha256": meta.get("weights_sha256"), "device": str(device), "torch": torch.__version__,85 "precision": "fp32", "prompt_len": args.prompt_len, "n_new": args.new, "code_commit": git_sha(), "peak_memory_bytes": peak_memory_bytes(device),86 "analytic_ops_per_token": analytic, "records": recs}87 path = Path(args.out or f"results/raw/latency_{Path(exp['checkpoint']).parent.name}_{device.type}.json")88 path.parent.mkdir(parents=True, exist_ok=True); path.write_text(json.dumps(out, indent=1)); print("wrote", path)89 return 0909192if __name__ == "__main__":93 raise SystemExit(main())94