mlx-bun 0.0.4

package/src/eval/tasks/mmlu.ts

@@ -0,0 +1,145 @@
+// MMLU 5-shot — multi-domain multiple-choice knowledge. Port of
+// optiq/eval/mmlu.py: per-subject dev exemplars, A/B/C/D format, greedy
+// decode, extract the predicted answer letter, compare to the gold index.
+//
+// The Python reference scores via first-letter logit argmax (no generation).
+// The mlx-bun harness routes everything through generateText() (decoded
+// completion). For an instruct model the faithful argmax-on-raw-text path
+// underscores on a small noisy sample, so we drive the chat template: the
+// 5-shot exemplars become the priming context and an explicit instruction
+// asks for ONLY the single letter A/B/C/D. Extraction is robust to chatty
+// formatting (</think> wrappers, "Answer: X", "(X)", "**X**", bare letter).
+// Prompt format, header, exemplar selection, and gold mapping are verbatim.
+
+import { generateText, loadJsonl, sampleIndices, type TaskModel } from "../runner";
+
+interface MmluRow {
+  question: string;
+  subject: string;
+  choices: string[];
+  answer: number; // 0..3 → A/B/C/D
+}
+
+export interface MmluResult {
+  nCorrect: number;
+  nTotal: number;
+  accuracy: number; // 0..1
+}
+
+const LETTERS = ["A", "B", "C", "D"] as const;
+
+/** 'high_school_us_history' → 'high school us history'. */
+function formatSubject(subject: string): string {
+  return subject.replace(/_/g, " ");
+}
+
+/** Format one example; append the letter when `answerIdx` is given (demos). */
+function formatExample(question: string, choices: string[], answerIdx?: number): string {
+  let s = question.trim() + "\n";
+  for (let i = 0; i < choices.length; i++) s += `${LETTERS[i]}. ${choices[i]}\n`;
+  s += "Answer:";
+  if (answerIdx !== undefined) s += ` ${LETTERS[answerIdx]}\n`;
+  return s;
+}
+
+/** 5-shot prompt: subject header → n demos → open test question.
+ *
+ * Built for the chat template: the header + worked exemplars prime the model,
+ * and a trailing instruction pins the output to a single bare letter. The
+ * exemplars still close their "Answer: X" line so the model sees the expected
+ * shape, and the final question is left open ("Answer:") for the model. */
+function buildPrompt(row: MmluRow, devExamples: MmluRow[], nShots: number): string {
+  let prompt = `The following are multiple choice questions (with answers) about ${formatSubject(row.subject)}.\n\n`;
+  for (const ex of devExamples.slice(0, nShots)) {
+    prompt += formatExample(ex.question, ex.choices, ex.answer);
+    prompt += "\n";
+  }
+  prompt += formatExample(row.question, row.choices);
+  prompt += "\n\nRespond with ONLY the single letter (A, B, C, or D) of the correct answer. Do not explain.";
+  return prompt;
+}
+
+/** Model output → predicted letter index 0..3, else null. Robust to chatty
+ * instruct formatting: strips a leading <think>…</think> block, then prefers
+ * an explicit "Answer: X" / "(X)" / "**X**" marker, falling back to the first
+ * standalone A/B/C/D token (one not glued inside a longer word). */
+export function extractAnswer(textIn: string): number | null {
+  let text = textIn;
+  if (text.includes("</think>")) text = text.split("</think>").pop() ?? text;
+
+  const toIdx = (ch: string): number => LETTERS.indexOf(ch.toUpperCase() as (typeof LETTERS)[number]);
+
+  // Explicit answer markers, in priority order.
+  const markers = [
+    /(?:answer|correct answer|the answer is)\s*(?:is|:)?\s*\(?\*{0,2}([A-Da-d])\*{0,2}\)?/i,
+    /\*\*\s*([A-Da-d])\s*\*\*/, // **X**
+    /\(\s*([A-Da-d])\s*\)/, // (X)
+  ];
+  for (const re of markers) {
+    const m = text.match(re);
+    if (m) return toIdx(m[1]!);
+  }
+
+  // Fallback: first standalone letter (word-boundary, not embedded in a word).
+  const standalone = text.match(/\b([A-D])\b/);
+  if (standalone) return toIdx(standalone[1]!);
+
+  return null;
+}
+
+export async function evaluateMmlu(
+  tm: TaskModel,
+  opts: { nSamples?: number; nShots?: number; maxTokens?: number; seed?: number } = {},
+): Promise<MmluResult> {
+  const nShots = opts.nShots ?? 5;
+  const maxTokens = opts.maxTokens ?? 16;
+  const seed = opts.seed ?? 42;
+
+  const testRows = loadJsonl<MmluRow>("mmlu_test");
+  const devRows = loadJsonl<MmluRow>("mmlu_dev");
+
+  // Per-subject dev exemplars for prompt construction.
+  const devBySubject = new Map<string, MmluRow[]>();
+  for (const ex of devRows) {
+    const list = devBySubject.get(ex.subject) ?? [];
+    list.push(ex);
+    devBySubject.set(ex.subject, list);
+  }
+
+  // Stratified sample: roughly equal questions per subject (matches mmlu.py),
+  // deterministic via the shared seeded subsampler.
+  const testBySubject = new Map<string, number[]>();
+  for (let i = 0; i < testRows.length; i++) {
+    const subj = testRows[i]!.subject;
+    const list = testBySubject.get(subj) ?? [];
+    list.push(i);
+    testBySubject.set(subj, list);
+  }
+  const subjects = [...testBySubject.keys()].sort();
+  const nSamples = opts.nSamples ?? 1000;
+  const perSubject = Math.max(1, Math.floor(nSamples / subjects.length));
+
+  let sampled: number[] = [];
+  for (const subj of subjects) {
+    const pool = testBySubject.get(subj)!;
+    const take = Math.min(perSubject, pool.length);
+    for (const j of sampleIndices(pool.length, take, seed)) sampled.push(pool[j]!);
+  }
+  if (sampled.length > nSamples) {
+    const keep = sampleIndices(sampled.length, nSamples, seed);
+    sampled = keep.map((k) => sampled[k]!);
+  }
+
+  let nCorrect = 0;
+  for (let k = 0; k < sampled.length; k++) {
+    const item = testRows[sampled[k]!]!;
+    const devEx = devBySubject.get(item.subject) ?? [];
+    const out = await generateText(tm, buildPrompt(item, devEx, nShots), { maxTokens, useChat: true });
+    const pred = extractAnswer(out);
+    if (pred !== null && pred === item.answer) nCorrect++;
+    if ((k + 1) % 10 === 0 || k + 1 === sampled.length)
+      process.stderr.write(`\r  mmlu ${k + 1}/${sampled.length}  acc=${((nCorrect / (k + 1)) * 100).toFixed(1)}%`);
+  }
+  process.stderr.write("\n");
+  return { nCorrect, nTotal: sampled.length, accuracy: sampled.length ? nCorrect / sampled.length : 0 };
+}

package/src/evaldb.ts

@@ -0,0 +1,199 @@
+// Eval DB: every benchmark run recorded with model + config + commit
+// (bun:sqlite). The fit calculator's predictions are validated against
+// the measured peaks recorded here.
+
+import { Database } from "bun:sqlite";
+import { mkdirSync } from "node:fs";
+
+export const DEFAULT_EVAL_DB = `${process.env.HOME}/.cache/mlx-bun/evals.sqlite`;
+
+const SCHEMA = `
+CREATE TABLE IF NOT EXISTS runs (
+  id INTEGER PRIMARY KEY AUTOINCREMENT,
+  ts INTEGER NOT NULL,
+  model_path TEXT NOT NULL,
+  commit_sha TEXT,
+  prompt_tokens INTEGER NOT NULL,
+  cached_tokens INTEGER NOT NULL DEFAULT 0,
+  generated_tokens INTEGER NOT NULL,
+  prefill_tps REAL NOT NULL,
+  decode_tps REAL NOT NULL,
+  peak_bytes INTEGER NOT NULL,
+  predicted_peak_bytes INTEGER,
+  predicted_decode_tps REAL,
+  notes TEXT
+);
+
+-- Phase 18 P1 — parallel-load harness. ONE row per configuration point
+-- (a concurrency level or an arrival-rate target), not per request: the
+-- single-stream runs table can't hold p50/p95 distributions or a
+-- concurrency axis. Same additive-migration discipline as runs (these
+-- are real measurements, never dropped). stack + machine_state mirror the
+-- runs columns so cross-stack load rows read the same way.
+CREATE TABLE IF NOT EXISTS load_runs (
+  id INTEGER PRIMARY KEY AUTOINCREMENT,
+  ts INTEGER NOT NULL,
+  model_id TEXT NOT NULL,
+  base_url TEXT NOT NULL,
+  stack TEXT NOT NULL DEFAULT 'mlx-bun',
+  commit_sha TEXT,
+  mode TEXT NOT NULL,              -- 'closed-loop' | 'open-loop'
+  concurrency INTEGER NOT NULL,    -- in-flight cap (closed) or requester count (open)
+  target_rpm REAL,                 -- open-loop arrival target per requester (null = closed)
+  max_tokens INTEGER NOT NULL,
+  requests_ok INTEGER NOT NULL,
+  requests_err INTEGER NOT NULL,
+  duration_s REAL NOT NULL,        -- wall time of the measured window
+  ttft_p50_ms REAL NOT NULL,
+  ttft_p95_ms REAL NOT NULL,
+  e2e_p50_ms REAL NOT NULL,
+  e2e_p95_ms REAL NOT NULL,
+  e2e_p99_ms REAL NOT NULL,
+  agg_tps REAL NOT NULL,           -- generated tokens/s summed across streams
+  per_req_tps REAL NOT NULL,       -- median per-stream decode tok/s
+  achieved_rpm REAL NOT NULL,      -- completed requests/min over the window
+  peak_bytes INTEGER,              -- from /stats if exposed, else null
+  notes TEXT,
+  machine_state TEXT
+);
+`;
+
+export interface EvalRun {
+  modelPath: string;
+  commitSha?: string | null;
+  promptTokens: number;
+  cachedTokens?: number;
+  generatedTokens: number;
+  prefillTps: number;
+  decodeTps: number;
+  peakBytes: number;
+  predictedPeakBytes?: number | null;
+  predictedDecodeTps?: number | null;
+  notes?: string | null;
+  /** Which engine produced the number: mlx-bun (default) | mlx-lm | optiq. */
+  stack?: string;
+  /** JSON snapshot of preflight machine state (swap, free %, thermal). */
+  machineState?: string | null;
+}
+
+/** One configuration point of a parallel-load sweep (Phase 18 P1). */
+export interface LoadRun {
+  modelId: string;
+  baseUrl: string;
+  stack?: string;
+  commitSha?: string | null;
+  /** 'closed-loop' (fixed in-flight) or 'open-loop' (target arrival rate). */
+  mode: "closed-loop" | "open-loop";
+  /** In-flight cap (closed-loop) or requester count (open-loop). */
+  concurrency: number;
+  /** Per-requester target rpm (open-loop only; null for closed-loop). */
+  targetRpm?: number | null;
+  maxTokens: number;
+  requestsOk: number;
+  requestsErr: number;
+  durationS: number;
+  ttftP50Ms: number;
+  ttftP95Ms: number;
+  e2eP50Ms: number;
+  e2eP95Ms: number;
+  e2eP99Ms: number;
+  /** Generated tokens/s summed across all streams in the window. */
+  aggTps: number;
+  /** Median per-stream decode tok/s. */
+  perReqTps: number;
+  /** Completed requests per minute over the measured window. */
+  achievedRpm: number;
+  /** Peak resident bytes sampled from /stats, if the server exposes it. */
+  peakBytes?: number | null;
+  notes?: string | null;
+  machineState?: string | null;
+}
+
+export class EvalDB {
+  readonly db: Database;
+
+  constructor(dbPath: string = DEFAULT_EVAL_DB) {
+    if (dbPath !== ":memory:") {
+      try { mkdirSync(dbPath.slice(0, dbPath.lastIndexOf("/")), { recursive: true }); } catch {}
+    }
+    this.db = new Database(dbPath);
+    this.db.exec(SCHEMA);
+    // Additive migrations — eval rows are real data, never dropped
+    // (unlike the registry, which is a rebuildable cache).
+    const cols = (this.db.query("PRAGMA table_info(runs)").all() as { name: string }[])
+      .map((c) => c.name);
+    if (!cols.includes("stack"))
+      this.db.exec("ALTER TABLE runs ADD COLUMN stack TEXT NOT NULL DEFAULT 'mlx-bun'");
+    if (!cols.includes("machine_state"))
+      this.db.exec("ALTER TABLE runs ADD COLUMN machine_state TEXT");
+  }
+
+  record(run: EvalRun): void {
+    this.db.prepare(`
+      INSERT INTO runs (ts, model_path, commit_sha, prompt_tokens, cached_tokens,
+        generated_tokens, prefill_tps, decode_tps, peak_bytes,
+        predicted_peak_bytes, predicted_decode_tps, notes, stack, machine_state)
+      VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+    `).run(
+      Date.now(), run.modelPath, run.commitSha ?? null,
+      run.promptTokens, run.cachedTokens ?? 0, run.generatedTokens,
+      run.prefillTps, run.decodeTps, run.peakBytes,
+      run.predictedPeakBytes ?? null, run.predictedDecodeTps ?? null,
+      run.notes ?? null, run.stack ?? "mlx-bun", run.machineState ?? null,
+    );
+  }
+
+  /** Record one parallel-load configuration point (Phase 18 P1). */
+  recordLoad(run: LoadRun): void {
+    this.db.prepare(`
+      INSERT INTO load_runs (ts, model_id, base_url, stack, commit_sha, mode,
+        concurrency, target_rpm, max_tokens, requests_ok, requests_err,
+        duration_s, ttft_p50_ms, ttft_p95_ms, e2e_p50_ms, e2e_p95_ms, e2e_p99_ms,
+        agg_tps, per_req_tps, achieved_rpm, peak_bytes, notes, machine_state)
+      VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+    `).run(
+      Date.now(), run.modelId, run.baseUrl, run.stack ?? "mlx-bun",
+      run.commitSha ?? null, run.mode, run.concurrency, run.targetRpm ?? null,
+      run.maxTokens, run.requestsOk, run.requestsErr, run.durationS,
+      run.ttftP50Ms, run.ttftP95Ms, run.e2eP50Ms, run.e2eP95Ms, run.e2eP99Ms,
+      run.aggTps, run.perReqTps, run.achievedRpm, run.peakBytes ?? null,
+      run.notes ?? null, run.machineState ?? null,
+    );
+  }
+
+  recentLoad(limit = 20): Record<string, unknown>[] {
+    return this.db
+      .query("SELECT * FROM load_runs ORDER BY ts DESC LIMIT ?")
+      .all(limit) as Record<string, unknown>[];
+  }
+
+  /** Latest mlx-bun measurement for a model snapshot (status page:
+   *  measured numbers beat predictions when available). */
+  latestFor(modelPath: string): { decodeTps: number; ts: number; notes: string | null } | null {
+    const r = this.db
+      .query("SELECT decode_tps, ts, notes FROM runs WHERE model_path = ? AND stack = 'mlx-bun' ORDER BY ts DESC LIMIT 1")
+      .get(modelPath) as { decode_tps: number; ts: number; notes: string | null } | null;
+    return r ? { decodeTps: r.decode_tps, ts: r.ts, notes: r.notes } : null;
+  }
+
+  recent(limit = 20): Record<string, unknown>[] {
+    return this.db
+      .query("SELECT * FROM runs ORDER BY ts DESC LIMIT ?")
+      .all(limit) as Record<string, unknown>[];
+  }
+
+  close(): void {
+    this.db.close();
+  }
+}
+
+export function gitCommit(): string | null {
+  try {
+    const r = Bun.spawnSync(["git", "rev-parse", "--short", "HEAD"], {
+      cwd: import.meta.dir,
+    });
+    return r.exitCode === 0 ? r.stdout.toString().trim() : null;
+  } catch {
+    return null;
+  }
+}

package/src/expert-offload-build.ts

@@ -0,0 +1,77 @@
+// E1 expert offload — BUILD (PLAN Phase 19).
+//
+// Produces the page-aligned offload file (experts.bin + manifest.json): each
+// expert tensor's bytes are copied to a 16 KB-aligned offset so the GPU can
+// gather it directly from a read-only file mmap (safetensors packs tensors at
+// arbitrary offsets; the Phase-1 GPU-alignment rule needs page alignment).
+// Kept separate from the runtime (src/expert-offload.ts) so the hot model
+// path doesn't import the converter / safetensors-writer code.
+
+import {
+  openSync, writeSync, closeSync, writeFileSync, mkdirSync, existsSync, statSync, readFileSync,
+} from "node:fs";
+import { ShardedSafetensors } from "./safetensors";
+import type { OffloadEntry, OffloadManifest } from "./expert-offload";
+
+const PAGE = 16384;
+const alignUp = (n: number, a: number) => Math.ceil(n / a) * a;
+const isExpertTensor = (n: string) => n.includes("switch_glu") || n.includes(".experts.");
+const layerOf = (n: string) => Number(n.match(/layers\.(\d+)/)?.[1] ?? -1);
+
+/** Reuse `<modelPath>/.mlx-bun-offload` if present and valid for this model,
+ *  else build it. Returns the offload dir (pass to activateExpertOffload). */
+export async function ensureOffloadFile(
+  modelPath: string,
+  onProgress?: (msg: string) => void,
+  outDir = `${modelPath}/.mlx-bun-offload`,
+): Promise<string> {
+  const bin = `${outDir}/experts.bin`;
+  const man = `${outDir}/manifest.json`;
+  if (existsSync(bin) && existsSync(man)) {
+    try {
+      const m = JSON.parse(readFileSync(man, "utf8")) as OffloadManifest;
+      if (m.model === modelPath && statSync(bin).size === m.totalBytes) {
+        onProgress?.(`reusing ${bin}`);
+        return outDir;
+      }
+    } catch { /* fall through to rebuild */ }
+  }
+  await buildOffloadFile(modelPath, outDir, onProgress);
+  return outDir;
+}
+
+export async function buildOffloadFile(
+  modelPath: string,
+  outDir: string,
+  onProgress?: (msg: string) => void,
+  layerCap = Infinity,
+): Promise<OffloadManifest> {
+  const shards = await ShardedSafetensors.open(modelPath);
+  const names = shards.tensorNames
+    .filter(isExpertTensor)
+    .filter((n) => layerOf(n) < layerCap)
+    .sort((a, b) => layerOf(a) - layerOf(b) || a.localeCompare(b));
+  if (!names.length) throw new Error(`${modelPath}: no expert tensors (not an MoE model?)`);
+
+  mkdirSync(outDir, { recursive: true });
+  const fd = openSync(`${outDir}/experts.bin`, "w");
+  const pad = Buffer.alloc(PAGE);
+  let offset = 0;
+  let i = 0;
+  const tensors: OffloadEntry[] = [];
+  for (const name of names) {
+    const aligned = alignUp(offset, PAGE);
+    if (aligned > offset) { writeSync(fd, pad, 0, aligned - offset); offset = aligned; }
+    const info = shards.info(name);
+    const bytes = shards.view(name); // zero-copy view into the source mmap
+    writeSync(fd, bytes);
+    tensors.push({ name, offset, length: bytes.length, dtype: info.dtype, shape: info.shape });
+    offset += bytes.length;
+    if (++i % 30 === 0) onProgress?.(`packing experts ${(offset / 1e9).toFixed(1)} GB (${i}/${names.length})`);
+  }
+  closeSync(fd);
+  const manifest: OffloadManifest = { page: PAGE, model: modelPath, totalBytes: offset, tensors };
+  writeFileSync(`${outDir}/manifest.json`, JSON.stringify(manifest));
+  onProgress?.(`wrote ${(offset / 1e9).toFixed(2)} GB, ${tensors.length} tensors, ${new Set(names.map(layerOf)).size} layers`);
+  return manifest;
+}

package/src/expert-offload.ts

@@ -0,0 +1,61 @@
+// E1 expert offload — RUNTIME (PLAN Phase 19).
+//
+// Routes expert WEIGHT tensors through a read-only file mmap (clean,
+// page-aligned pages that cost ~0 phys_footprint and that the GPU gathers
+// directly — probes confirmed bit-exact + non-wiring) instead of the
+// anonymous mlx_load_safetensors copy that counts against memory pressure.
+//
+// Activated two ways, both before model construction:
+//   - env MLX_BUN_EXPERT_OFFLOAD=<dir>  (scripts / direct runs)
+//   - activateExpertOffload(dir)        (CLI `--expert-offload`)
+// The file is produced by src/expert-offload-build.ts. Tensors not in the
+// manifest (or when never activated) fall back to the resident path, so this
+// is inert by default and safe on a partially-converted file.
+
+import { readFileSync } from "node:fs";
+import { MmapFile } from "./mmap";
+import { MlxArray, SAFETENSORS_TO_MLX } from "./mlx/array";
+import type { SafetensorsDtype } from "./safetensors";
+
+export interface OffloadEntry {
+  name: string;
+  offset: number;
+  length: number;
+  dtype: SafetensorsDtype;
+  shape: number[];
+}
+export interface OffloadManifest {
+  page: number;
+  model: string;
+  totalBytes: number;
+  tensors: OffloadEntry[];
+}
+
+let mm: MmapFile | null = null;
+let manifest: Map<string, OffloadEntry> | null = null;
+
+/** Open the offload file at `dir` (experts.bin + manifest.json) and route
+ *  expert weights through it. Call BEFORE the model is constructed. */
+export function activateExpertOffload(dir: string): void {
+  const parsed = JSON.parse(readFileSync(`${dir}/manifest.json`, "utf8")) as OffloadManifest;
+  manifest = new Map(parsed.tensors.map((t) => [t.name, t]));
+  mm = MmapFile.open(`${dir}/experts.bin`, "ro");
+  process.stderr.write(`[expert-offload] mmap ${dir}/experts.bin (${parsed.tensors.length} tensors)\n`);
+}
+
+export function isExpertOffload(): boolean {
+  return mm !== null;
+}
+
+/** Zero-copy mmap-backed array for `name` if offload is active and the tensor
+ *  is in the manifest; else null. The converter aligns every tensor to 16 KB
+ *  and the mmap base is page-aligned, so the pointer is page-aligned —
+ *  required for the GPU to gather from it. */
+export function expertOffloadArray(name: string): MlxArray | null {
+  if (!mm || !manifest) return null;
+  const e = manifest.get(name);
+  if (!e) return null;
+  return MlxArray.fromView(mm.view(e.offset, e.length), e.shape, SAFETENSORS_TO_MLX[e.dtype]);
+}
+
+if (process.env.MLX_BUN_EXPERT_OFFLOAD) activateExpertOffload(process.env.MLX_BUN_EXPERT_OFFLOAD);

package/src/expert-trace.ts

@@ -0,0 +1,67 @@
+// E0 expert-routing tracer (PLAN Phase 19 — expert offload).
+//
+// Records every MoE router decision as JSONL — one record per router call:
+//   { c: callSeq, l: layerIdx, s: [B, L, k], i: [flattened uint32 expert ids] }
+// Analyse with scripts/analyze-expert-trace.ts (coverage curve, working-set
+// size, within-task stability, cross-task specialisation).
+//
+// Two ways to drive it:
+//   - env: `MLX_BUN_EXPERT_TRACE=<path> mlx-bun serve ...` (auto-starts here)
+//   - programmatic: beginExpertTrace(path) / endExpertTrace() per domain
+//     (scripts/run-expert-trace.ts uses this to trace several domains from
+//     one model load).
+//
+// Recording adds a per-call GPU->host sync (it reads the routed indices back),
+// so a traced run decodes slower — it's a MEASUREMENT tool, not a serving
+// path. When no trace is active, isExpertTracing() is false and the hook in
+// Router.forward is a single branch (zero overhead) — inert by default, same
+// discipline as the slots flag.
+
+import { Dtype } from "./mlx/ffi";
+import type { MlxArray } from "./mlx/array";
+
+let sink: Bun.FileSink | null = null;
+let seq = 0;
+let written = 0;
+
+export function beginExpertTrace(path: string): void {
+  if (sink) endExpertTrace();
+  sink = Bun.file(path).writer();
+  seq = 0;
+  written = 0;
+  sink.write(JSON.stringify({ meta: "mlx-bun-expert-trace-v1", note: "one record per router call" }) + "\n");
+  process.stderr.write(`[expert-trace] recording -> ${path}\n`);
+}
+
+export function endExpertTrace(): void {
+  if (!sink) return;
+  sink.flush();
+  sink.end();
+  process.stderr.write(`[expert-trace] wrote ${written} records (${seq} calls)\n`);
+  sink = null;
+}
+
+/** Alias kept for the process-exit handler / env callers. */
+export const finishExpertTrace = endExpertTrace;
+
+export function isExpertTracing(): boolean {
+  return sink !== null;
+}
+
+/** Record one router decision. `indices` is [B, L, k] uint32 (expert ids). */
+export function recordRouting(layer: number, indices: MlxArray): void {
+  if (!sink) return;
+  const shape = indices.shape; // [B, L, k]
+  const f = indices.astype(Dtype.float32); // ids <=127 are exact in f32
+  const data = f.toFloat32();
+  f.dispose();
+  const ids = new Array<number>(data.length);
+  for (let i = 0; i < data.length; i++) ids[i] = Math.round(data[i]!);
+  sink.write(JSON.stringify({ c: seq++, l: layer, s: shape, i: ids }) + "\n");
+  if (++written % 2000 === 0) sink.flush();
+}
+
+if (process.env.MLX_BUN_EXPERT_TRACE) beginExpertTrace(process.env.MLX_BUN_EXPERT_TRACE);
+
+process.on("exit", () => { try { endExpertTrace(); } catch { /* best-effort */ } });
+process.on("SIGINT", () => { endExpertTrace(); process.exit(0); });