1#!/usr/bin/env python32"""Gate D profiling: where does a decode step spend its time with all experts resident? Times (synchronous, median of N) the3full step, one layer's attention, one layer's MoE block (3 gather_qmm + swiglu, top-8), and the three gather_qmm kernels alone,4and computes the effective bytes moved per second against the expert bytes touched (bandwidth-boundedness). Writes JSON.5"""6import argparse, json, os, platform, socket, statistics, subprocess, sys, time7import mlx.core as mx8from pathlib import Path9sys.path.insert(0, str(Path(__file__).resolve().parents[2]))10from mlx_lm.utils import load_model, load_tokenizer11from mlx_lm.models.cache import make_prompt_cache12from mlx_lm.models.activations import swiglu13from titan.registry import load_spec1415def timeit(f, n=100, warm=10):16 for _ in range(warm): f()17 mx.synchronize(); ts = []18 for _ in range(n):19 t = time.perf_counter(); f(); mx.synchronize(); ts.append((time.perf_counter() - t) * 1e3)20 return {"median_ms": statistics.median(ts), "p10_ms": sorted(ts)[n // 10], "p90_ms": sorted(ts)[9 * n // 10]}2122def main():23 ap = argparse.ArgumentParser(); ap.add_argument("--model", required=True); ap.add_argument("--out", required=True); ap.add_argument("--n", type=int, default=100)24 a = ap.parse_args()25 spec = load_spec(a.model); model, cfg = load_model(Path(a.model)); tok = load_tokenizer(Path(a.model))26 L, E, K, D = spec.num_layers, spec.num_experts, spec.top_k, spec.hidden_size27 ids = mx.array(tok.encode("The capital of Canada is"))[None]28 cache = make_prompt_cache(model); y = mx.argmax(model(ids, cache=cache)[:, -1], -1); mx.eval(y)29 rec = {"host": socket.gethostname(), "chip": subprocess.run(["sysctl", "-n", "machdep.cpu.brand_string"], capture_output=True, text=True).stdout.strip(),30 "macos": platform.mac_ver()[0], "mlx": mx.__version__, "model": os.path.realpath(a.model), "L": L, "E": E, "K": K, "D": D,31 "bytes_per_expert": spec.bytes_per_expert, "expert_bytes_per_token": L * K * spec.bytes_per_expert, "n": a.n, "cmd": " ".join(sys.argv)}32 # 1. full decode step (sync)33 def step():34 out = model(y[None], cache=cache); mx.eval(out)35 rec["decode_step"] = timeit(step, a.n)36 # 2. one layer's blocks in isolation (layer 4), on a realistic hidden state37 layer = model.model.layers[4]38 h = mx.random.normal((1, 1, D)).astype(mx.float16); mx.eval(h)39 from mlx_lm.models.base import create_attention_mask40 c2 = make_prompt_cache(model); model(ids, cache=c2); mx.eval(c2[4].state)41 rec["attention_layer"] = timeit(lambda: mx.eval(layer.self_attn(layer.input_layernorm(h), None, c2[4])), a.n)42 rec["moe_block_layer"] = timeit(lambda: mx.eval(layer.mlp(layer.post_attention_layernorm(h))), a.n)43 sw = layer.mlp.switch_mlp44 xf = h.reshape(-1, D); idx = mx.array([[list(range(K))]], dtype=mx.uint32)45 xe = mx.expand_dims(xf, (-2, -3))46 def qmm3():47 up = sw.up_proj(xe, idx); g = sw.gate_proj(xe, idx); d = sw.down_proj(swiglu(g, up), idx); mx.eval(d)48 rec["gather_qmm_x3_layer"] = timeit(qmm3, a.n)49 rec["gather_qmm_up_only"] = timeit(lambda: mx.eval(sw.up_proj(xe, idx)), a.n)50 rec["router_layer"] = timeit(lambda: mx.eval(mx.argpartition(-mx.softmax(layer.mlp.gate(xf), axis=1, precise=True), kth=K - 1, axis=-1)[..., :K]), a.n)51 # 3. derived52 t_step = rec["decode_step"]["median_ms"]; t_moe = rec["moe_block_layer"]["median_ms"] * L; t_att = rec["attention_layer"]["median_ms"] * L53 rec["derived"] = {"moe_fraction_of_step_est": t_moe / t_step, "attention_fraction_est": t_att / t_step,54 "expert_GBps_effective_in_moe_block": rec["expert_bytes_per_token"] / (t_moe / 1e3) / 1e9,55 "expert_GBps_effective_full_step": rec["expert_bytes_per_token"] / (t_step / 1e3) / 1e9,56 "tok_s_sync": 1e3 / t_step}57 os.makedirs(os.path.dirname(a.out), exist_ok=True); json.dump(rec, open(a.out, "w"), indent=2); print(json.dumps(rec, indent=1))5859if __name__ == "__main__":60 main()61