1#!/usr/bin/env python32"""Collect routing traces: for every prompt, record per token and per layer the router probabilities (float16, all experts)3and the top-k expert ids actually selected, during greedy generation with all experts resident (mlx-lm model).4Traces feed the offline cache/predictor simulator (benchmarks/scripts/simulate_cache.py) and predictor fitting.5Output: one .npz per prompt under --out, plus an index JSON. Prefill and decode tokens are both recorded (phase array).6"""7import argparse, json, os, sys, time8import numpy as np9import mlx.core as mx10import yaml11from pathlib import Path12sys.path.insert(0, str(Path(__file__).resolve().parents[2]))13from mlx_lm.utils import load_model, load_tokenizer14from mlx_lm.models.cache import make_prompt_cache151617def main():18 ap = argparse.ArgumentParser()19 ap.add_argument("--model", required=True)20 ap.add_argument("--workloads", default=str(Path(__file__).resolve().parents[1] / "configs" / "workloads.yaml"))21 ap.add_argument("--max-tokens", type=int, default=256)22 ap.add_argument("--out", required=True)23 a = ap.parse_args()24 model, cfg = load_model(Path(a.model)); tok = load_tokenizer(Path(a.model))25 wl = yaml.safe_load(open(a.workloads))26 layers = model.model.layers27 L = len(layers); E = cfg["num_experts"]; K = cfg["num_experts_per_tok"]28 probs_buf, ids_buf = [], []29 # hook the router: wrap each layer's mlp.__call__ to capture softmax probs and chosen ids30 for l, layer in enumerate(layers):31 mlp = layer.mlp32 def hooked(x, _mlp=mlp, _l=l):33 B, Ln, D = x.shape34 x_flat = x.reshape(-1, D)35 logits = _mlp.gate(x_flat)36 rw = mx.softmax(logits, axis=1, precise=True)37 idx = mx.argpartition(-rw, kth=_mlp.top_k - 1, axis=-1)[..., :_mlp.top_k]38 probs_buf.append((_l, rw.astype(mx.float16)))39 ids_buf.append((_l, idx))40 scores = mx.take_along_axis(rw, idx, axis=-1)41 if _mlp.norm_topk_prob:42 scores = scores / scores.sum(axis=-1, keepdims=True)43 y = _mlp.switch_mlp(x_flat, idx)44 y = (y * scores[..., None]).sum(axis=-2)45 return y.reshape(B, Ln, D)46 layer.mlp.__call__ = hooked # instance-level override is not seen by nn.Module.__call__ -> patch the class attribute per layer instead47 layer.mlp = _Wrapped(hooked)48 os.makedirs(a.out, exist_ok=True)49 index = []50 eos = tok.eos_token_id51 for fam, splits in wl["families"].items():52 for split, prompts in splits.items():53 for pi, p in enumerate(prompts):54 ids = mx.array(tok.encode(p))[None]55 cache = make_prompt_cache(model)56 probs_buf.clear(); ids_buf.clear()57 t0 = time.perf_counter()58 logits = model(ids, cache=cache); y = mx.argmax(logits[:, -1], -1); mx.eval(y)59 n_prompt = ids.shape[1]60 P = [np.array(pb) for _, pb in probs_buf]; I = [np.array(ib) for _, ib in ids_buf]61 probs = [np.stack([P[l] for l in range(L)], 1)] # [T, L, E]62 topk = [np.stack([I[l] for l in range(L)], 1)] # [T, L, K]63 gen = [int(y.item())]64 for step in range(a.max_tokens - 1):65 if gen[-1] == eos: break66 probs_buf.clear(); ids_buf.clear()67 logits = model(y[None], cache=cache); y = mx.argmax(logits[:, -1], -1); mx.eval(y)68 P = [np.array(pb) for _, pb in probs_buf]; I = [np.array(ib) for _, ib in ids_buf]69 probs.append(np.stack([P[l] for l in range(L)], 1)); topk.append(np.stack([I[l] for l in range(L)], 1))70 gen.append(int(y.item()))71 probs = np.concatenate(probs, 0).astype(np.float16); topk = np.concatenate(topk, 0).astype(np.int16)72 phase = np.array([0] * n_prompt + [1] * (probs.shape[0] - n_prompt), dtype=np.int8)73 name = f"{fam}_{split}_{pi}"74 np.savez_compressed(os.path.join(a.out, name + ".npz"), probs=probs, topk=topk, phase=phase, tokens=np.array(ids.reshape(-1).tolist() + gen))75 index.append({"name": name, "family": fam, "split": split, "prompt": p[:80], "n_prompt": n_prompt, "n_gen": len(gen),76 "seconds": time.perf_counter() - t0, "text": tok.decode(gen)[:120]})77 print(f"{name:24s} prompt={n_prompt:4d} gen={len(gen):4d} {index[-1]['seconds']:.1f}s")78 json.dump({"model": os.path.realpath(a.model), "L": L, "E": E, "K": K, "max_tokens": a.max_tokens, "ts": time.strftime("%Y-%m-%dT%H:%M:%S"),79 "cmd": " ".join(sys.argv), "prompts": index}, open(os.path.join(a.out, "index.json"), "w"), indent=2)808182class _Wrapped:83 def __init__(self, fn): self.fn = fn84 def __call__(self, x): return self.fn(x)858687if __name__ == "__main__":88 main()89