mlx-bun 0.0.4

package/src/eval/kl.ts

@@ -0,0 +1,294 @@
+// KL-divergence drift eval — port of optiq/eval/kl.py, adapted for the
+// fused-kernel head-to-head.
+//
+// KL(ref ‖ cand) per token = Σ_v softmax(ref)·(logsoftmax(ref) − logsoftmax(cand)),
+// reported as mean (primary), median, and p95 over N prompts × seq tokens.
+// Computation is TEACHER-FORCED: the same fixed prompt is fed to both
+// arms and the per-position distributions are compared — the repo's rule
+// (free-running greedy "measures chaos"; see teacher-forced-gating).
+//
+// Two reference modes:
+//   * self-flag (DEFAULT, the kernel-drift gate): one model, one weight
+//     load; forward each prompt twice with an env lever set to its
+//     reference value vs its candidate value. Directly measures the drift
+//     a perf lever introduces. Fits in RAM trivially (no second model).
+//   * two-model (optiq-style absolute quality): a separate reference
+//     model (bf16 if it fits, else uniform-4bit). Both resident.
+//
+// NOTE (M0 scope): this uses the compat forward() path (plain caches), so
+// it gates the shared-ops fusions (gelu / norm+add) and the tiled-vs-
+// unfused SDPA lever. A serving-path variant (quantized KV + the generated
+// class, teacher-forced prefill) is the next step — see eval README.
+
+import { existsSync } from "node:fs";
+import { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import * as ops from "../mlx/ops";
+import { Registry } from "../registry";
+import { Weights } from "../weights";
+import { loadModelConfig, type ModelConfig } from "../config";
+import { createModel, type RuntimeModel } from "../model/factory";
+import { loadTokenizer, type LoadedTokenizer } from "../tokenizer";
+import { maybeQuantizeKv } from "../generate";
+
+export interface KLResult {
+  nPrompts: number;
+  seqLen: number;
+  meanKl: number;   // mean of per-prompt mean KL — the primary scalar
+  medianKl: number; // median over all pooled per-token KLs
+  p95Kl: number;    // tail (long-distribution drift)
+  elapsedSec: number;
+  refLabel: string;
+}
+
+/** KL(p ‖ q) per token for logits [1, T, V]; returns a length-T array. */
+export function klPerToken(pLogits: MlxArray, qLogits: MlxArray): Float32Array {
+  const f32 = (a: MlxArray): { arr: MlxArray; owned: boolean } =>
+    a.dtype === Dtype.float32 ? { arr: a, owned: false } : { arr: a.astype(Dtype.float32), owned: true };
+  const pp = f32(pLogits);
+  const qq = f32(qLogits);
+  const p = pp.arr;
+  const q = qq.arr;
+
+  const lseP = ops.logsumexpAxis(p, -1, true);   // [1,T,1]
+  const lseQ = ops.logsumexpAxis(q, -1, true);
+  const logP = ops.sub(p, lseP);                  // [1,T,V]
+  const logQ = ops.sub(q, lseQ);
+  const pProbs = ops.softmaxAxis(p, -1, true);    // precise softmax
+  const diff = ops.sub(logP, logQ);
+  const prod = ops.mul(pProbs, diff);
+  const kl = ops.sumAxis(prod, -1, false);        // [1,T]
+  const out = kl.toFloat32();
+
+  for (const a of [lseP, lseQ, logP, logQ, pProbs, diff, prod, kl]) a.dispose();
+  if (pp.owned) p.dispose();
+  if (qq.owned) q.dispose();
+  return out;
+}
+
+const mean = (xs: number[]): number => (xs.length ? xs.reduce((a, b) => a + b, 0) / xs.length : 0);
+function percentile(xs: number[], pct: number): number {
+  if (!xs.length) return 0;
+  const s = [...xs].sort((a, b) => a - b);
+  const i = Math.min(s.length - 1, Math.max(0, Math.floor((pct / 100) * s.length)));
+  return s[i]!;
+}
+
+function aggregate(perPrompt: Float32Array[]): { mean: number; median: number; p95: number } {
+  const promptMeans = perPrompt.map((a) => mean(Array.from(a)));
+  const pooled: number[] = [];
+  for (const a of perPrompt) for (const v of a) pooled.push(v);
+  return { mean: mean(promptMeans), median: percentile(pooled, 50), p95: percentile(pooled, 95) };
+}
+
+/** A local model dir if `query` is a path, else resolved via the registry. */
+export function resolveModelDir(query: string): string {
+  if (existsSync(query) && existsSync(`${query}/config.json`)) return query;
+  return new Registry().resolve(query).path;
+}
+
+export interface Runnable {
+  model: RuntimeModel;
+  tokenizer: LoadedTokenizer;
+  config: ModelConfig;
+  dir: string;
+}
+
+export async function loadRunnable(query: string): Promise<Runnable> {
+  const dir = resolveModelDir(query);
+  const config = await loadModelConfig(dir);
+  const weights = await Weights.open(dir);
+  const model = createModel(weights, config);
+  const tokenizer = await loadTokenizer(dir);
+  return { model, tokenizer, config, dir };
+}
+
+/** Forward a fixed token sequence, return logits [1, L, V] (caller disposes). */
+function forwardLogits(model: RuntimeModel, tokens: number[]): MlxArray {
+  const cache = model.makeCache();
+  try {
+    return model.forward(tokens, cache);
+  } finally {
+    for (const c of cache) c.dispose();
+  }
+}
+
+function prepPrompts(prompts: string[], tok: LoadedTokenizer, seqLen: number, n: number): number[][] {
+  const tiled: string[] = [];
+  while (tiled.length < n) tiled.push(...prompts);
+  return tiled.slice(0, n).map((p) => tok.encode(p).slice(0, seqLen));
+}
+
+/** self-flag drift gate: ref = env[flag]=refValue, cand = env[flag]=candValue. */
+export async function evaluateKlSelfFlag(opts: {
+  candidate: string;
+  flag: string;
+  refValue: string;
+  candValue: string;
+  prompts: string[];
+  nPrompts?: number;
+  seqLen?: number;
+}): Promise<KLResult> {
+  const t0 = Date.now();
+  const n = opts.nPrompts ?? 64;
+  const seqLen = opts.seqLen ?? 256;
+  const { model, tokenizer } = await loadRunnable(opts.candidate);
+  const tokenized = prepPrompts(opts.prompts, tokenizer, seqLen, n);
+
+  const prev = process.env[opts.flag];
+  const perPrompt: Float32Array[] = [];
+  try {
+    for (const ids of tokenized) {
+      process.env[opts.flag] = opts.refValue;
+      const p = forwardLogits(model, ids);
+      process.env[opts.flag] = opts.candValue;
+      const q = forwardLogits(model, ids);
+      const T = Math.min(p.shape[1]!, q.shape[1]!);
+      const pT = p.slice([0, 0, 0], [1, T, p.shape[2]!]);
+      const qT = q.slice([0, 0, 0], [1, T, q.shape[2]!]);
+      perPrompt.push(klPerToken(pT, qT));
+      for (const a of [p, q, pT, qT]) a.dispose();
+    }
+  } finally {
+    if (prev === undefined) delete process.env[opts.flag];
+    else process.env[opts.flag] = prev;
+  }
+
+  const agg = aggregate(perPrompt);
+  return {
+    nPrompts: n, seqLen, meanKl: agg.mean, medianKl: agg.median, p95Kl: agg.p95,
+    elapsedSec: (Date.now() - t0) / 1000,
+    refLabel: `self:${opts.flag}=${opts.refValue}→${opts.candValue}`,
+  };
+}
+
+/** two-model absolute drift: KL(reference ‖ candidate) on shared prompts. */
+export async function evaluateKlTwoModel(opts: {
+  candidate: string;
+  reference: string;
+  prompts: string[];
+  nPrompts?: number;
+  seqLen?: number;
+}): Promise<KLResult> {
+  const t0 = Date.now();
+  const n = opts.nPrompts ?? 64;
+  const seqLen = opts.seqLen ?? 256;
+  const ref = await loadRunnable(opts.reference);
+  const cand = await loadRunnable(opts.candidate);
+  const tokenized = prepPrompts(opts.prompts, ref.tokenizer, seqLen, n);
+
+  const perPrompt: Float32Array[] = [];
+  for (const ids of tokenized) {
+    const p = forwardLogits(ref.model, ids);
+    const q = forwardLogits(cand.model, ids);
+    const T = Math.min(p.shape[1]!, q.shape[1]!);
+    const pT = p.slice([0, 0, 0], [1, T, p.shape[2]!]);
+    const qT = q.slice([0, 0, 0], [1, T, q.shape[2]!]);
+    perPrompt.push(klPerToken(pT, qT));
+    for (const a of [p, q, pT, qT]) a.dispose();
+  }
+
+  const agg = aggregate(perPrompt);
+  return {
+    nPrompts: n, seqLen, meanKl: agg.mean, medianKl: agg.median, p95Kl: agg.p95,
+    elapsedSec: (Date.now() - t0) / 1000,
+    refLabel: opts.reference,
+  };
+}
+
+// --- serving-path (decode) KL — the M0b gate -----------------------------
+// The compat forward() KL above runs plain bf16 caches → it can't see the
+// quantized-KV serving path. This variant reproduces serving exactly: bf16
+// prefill, then maybeQuantizeKv (the generate.ts hook) converts the populated
+// mixed-4/8-bit caches, then a TEACHER-FORCED decode over the prompt tail —
+// where the generated quantized SDPA (and the fused kernels) actually run.
+
+/** CPU KL(p ‖ q) for two logit vectors (same math as klPerToken, per step). */
+function klScalar(p: Float32Array, q: Float32Array): number {
+  const V = p.length;
+  let mp = -Infinity;
+  let mq = -Infinity;
+  for (let i = 0; i < V; i++) { if (p[i]! > mp) mp = p[i]!; if (q[i]! > mq) mq = q[i]!; }
+  let sp = 0;
+  let sq = 0;
+  for (let i = 0; i < V; i++) { sp += Math.exp(p[i]! - mp); sq += Math.exp(q[i]! - mq); }
+  const lsep = mp + Math.log(sp);
+  const lseq = mq + Math.log(sq);
+  let kl = 0;
+  for (let i = 0; i < V; i++) {
+    const pp = Math.exp(p[i]! - lsep);
+    if (pp > 0) kl += pp * ((p[i]! - lsep) - (q[i]! - lseq));
+  }
+  return kl;
+}
+
+/** Prefill (bf16) → quantize → teacher-forced decode; return per-step logit
+ *  distributions for the last `tokens.length - prefillLen` positions. */
+function decodeArm(
+  model: RuntimeModel, config: ModelConfig, tokens: number[], prefillLen: number,
+  flag: string, value: string,
+): Float32Array[] {
+  const prev = process.env[flag];
+  process.env[flag] = value;
+  try {
+    const cache = model.makeCache();
+    try {
+      const lp = model.forward(tokens.slice(0, prefillLen), cache); // bf16 prefill
+      lp.dispose();
+      // quantize the populated caches (mixed 4/8 per kv_config) — exactly
+      // what generate() does between prefill and decode.
+      maybeQuantizeKv(cache, { kvConfig: config.kvQuant ?? undefined, quantizedKvStart: 0 });
+      const out: Float32Array[] = [];
+      for (let i = prefillLen; i < tokens.length; i++) {
+        const logits = model.forward([tokens[i]!], cache); // [1,1,V] quantized decode
+        out.push(logits.toFloat32());
+        logits.dispose();
+      }
+      return out;
+    } finally {
+      for (const c of cache) c.dispose();
+    }
+  } finally {
+    if (prev === undefined) delete process.env[flag];
+    else process.env[flag] = prev;
+  }
+}
+
+/** Serving-path drift gate: KL between two arms of a perf lever, measured on
+ *  the teacher-forced quantized-decode path (the real e4b serving path). */
+export async function evaluateKlServingDecode(opts: {
+  candidate: string;
+  flag: string;
+  refValue: string;
+  candValue: string;
+  prompts: string[];
+  nPrompts?: number;
+  seqLen?: number;
+  decodeSteps?: number;
+}): Promise<KLResult> {
+  const t0 = Date.now();
+  const n = opts.nPrompts ?? 64;
+  const seqLen = opts.seqLen ?? 256;
+  const decodeSteps = opts.decodeSteps ?? 32;
+  const { model, tokenizer, config } = await loadRunnable(opts.candidate);
+  const tokenized = prepPrompts(opts.prompts, tokenizer, seqLen, n);
+
+  const perPrompt: Float32Array[] = [];
+  for (const tokens of tokenized) {
+    if (tokens.length < 2) continue;
+    const prefillLen = Math.max(1, tokens.length - decodeSteps);
+    const refSteps = decodeArm(model, config, tokens, prefillLen, opts.flag, opts.refValue);
+    const candSteps = decodeArm(model, config, tokens, prefillLen, opts.flag, opts.candValue);
+    const m = Math.min(refSteps.length, candSteps.length);
+    const kls = new Float32Array(m);
+    for (let s = 0; s < m; s++) kls[s] = klScalar(refSteps[s]!, candSteps[s]!);
+    perPrompt.push(kls);
+  }
+
+  const agg = aggregate(perPrompt);
+  return {
+    nPrompts: perPrompt.length, seqLen, meanKl: agg.mean, medianKl: agg.median, p95Kl: agg.p95,
+    elapsedSec: (Date.now() - t0) / 1000,
+    refLabel: `serving-decode self:${opts.flag}=${opts.refValue}→${opts.candValue} (${decodeSteps} steps)`,
+  };
+}

package/src/eval/quality-db.ts

@@ -0,0 +1,92 @@
+// Quality-eval DB — the capability/KL counterpart to evaldb.ts (which is
+// throughput-only: prefill_tps/decode_tps/peak_bytes). Lives in the SAME
+// sqlite file (~/.cache/mlx-bun/evals.sqlite) as a separate `quality_runs`
+// table so the throughput head-to-head and the quality gate sit together.
+//
+// One row = one (model, config) measured on one task (or the aggregate
+// "capability" / "kl" rows). `config_json` records the perf levers that
+// were active (compile / fused-decode / perf-kernel) so a row is
+// self-describing: which arm of the head-to-head produced this number.
+
+import { Database } from "bun:sqlite";
+import { mkdirSync } from "node:fs";
+import { DEFAULT_EVAL_DB } from "../evaldb";
+
+const SCHEMA = `
+CREATE TABLE IF NOT EXISTS quality_runs (
+  id INTEGER PRIMARY KEY AUTOINCREMENT,
+  ts INTEGER NOT NULL,
+  model_path TEXT NOT NULL,
+  commit_sha TEXT,
+  task TEXT NOT NULL,            -- 'kl' | 'gsm8k' | 'mmlu' | ... | 'capability'
+  config_json TEXT NOT NULL,     -- {compile,fusedDecode,perfKernel,...} active levers
+  n_samples INTEGER NOT NULL,
+  -- task accuracy (percent 0-100); null for kl-only rows
+  pct REAL,
+  -- KL drift block (null for task rows)
+  kl_mean REAL,
+  kl_median REAL,
+  kl_p95 REAL,
+  kl_ref TEXT,                   -- 'self:<FLAG>' | model path/id of the reference
+  -- aggregate capability score (null except on 'capability' rows)
+  capability_score REAL,
+  disk_gb REAL,
+  notes TEXT,
+  machine_state TEXT
+);
+`;
+
+export interface QualityRun {
+  modelPath: string;
+  commitSha?: string | null;
+  task: string;
+  /** Active perf levers when this number was produced. */
+  config: Record<string, string | boolean | number | null>;
+  nSamples: number;
+  pct?: number | null;
+  klMean?: number | null;
+  klMedian?: number | null;
+  klP95?: number | null;
+  klRef?: string | null;
+  capabilityScore?: number | null;
+  diskGb?: number | null;
+  notes?: string | null;
+  machineState?: string | null;
+}
+
+export class QualityDB {
+  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);
+  }
+
+  record(run: QualityRun): void {
+    this.db.prepare(`
+      INSERT INTO quality_runs (ts, model_path, commit_sha, task, config_json,
+        n_samples, pct, kl_mean, kl_median, kl_p95, kl_ref,
+        capability_score, disk_gb, notes, machine_state)
+      VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+    `).run(
+      Date.now(), run.modelPath, run.commitSha ?? null, run.task,
+      JSON.stringify(run.config), run.nSamples,
+      run.pct ?? null, run.klMean ?? null, run.klMedian ?? null, run.klP95 ?? null,
+      run.klRef ?? null, run.capabilityScore ?? null, run.diskGb ?? null,
+      run.notes ?? null, run.machineState ?? null,
+    );
+  }
+
+  recent(limit = 20): Record<string, unknown>[] {
+    return this.db
+      .query("SELECT * FROM quality_runs ORDER BY ts DESC LIMIT ?")
+      .all(limit) as Record<string, unknown>[];
+  }
+
+  close(): void {
+    this.db.close();
+  }
+}

package/src/eval/runner.ts

@@ -0,0 +1,102 @@
+// Task-eval runner: load a model and greedily complete prompts through the
+// REAL serving path. Unlike the M0 KL (compat forward() + plain caches),
+// this drives generate() with the model's kv_config, so generation runs the
+// mixed-4/8-bit quantized-KV path the generated class actually serves.
+
+import { homedir } from "node:os";
+import { readFileSync } from "node:fs";
+import { generate, type GenerateOptions } from "../generate";
+import { ChatTemplate } from "../chat-template";
+import { loadModelConfig, type ModelConfig } from "../config";
+import { Weights } from "../weights";
+import { createModel, type RuntimeModel } from "../model/factory";
+import { loadTokenizer, type LoadedTokenizer } from "../tokenizer";
+import { resolveModelDir } from "./kl";
+
+export const EVAL_DATA_DIR = process.env.MLX_BUN_EVAL_DATA ?? `${homedir()}/.cache/mlx-bun/eval-data`;
+
+/** Read a jsonl dataset exported by scripts/eval/export-datasets.py. */
+export function loadJsonl<T = Record<string, unknown>>(name: string): T[] {
+  const path = `${EVAL_DATA_DIR}/${name}.jsonl`;
+  const out: T[] = [];
+  for (const line of readFileSync(path, "utf8").split("\n")) {
+    const t = line.trim();
+    if (t) out.push(JSON.parse(t) as T);
+  }
+  return out;
+}
+
+/** Deterministic seeded subsample of [0, total): mulberry32 shuffle, sorted. */
+export function sampleIndices(total: number, n: number, seed = 42): number[] {
+  let s = seed >>> 0;
+  const rand = (): number => {
+    s = (s + 0x6d2b79f5) | 0;
+    let t = Math.imul(s ^ (s >>> 15), 1 | s);
+    t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t;
+    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
+  };
+  const idx = Array.from({ length: total }, (_, i) => i);
+  for (let i = total - 1; i > 0; i--) {
+    const j = Math.floor(rand() * (i + 1));
+    [idx[i], idx[j]] = [idx[j]!, idx[i]!];
+  }
+  return idx.slice(0, Math.min(n, total)).sort((a, b) => a - b);
+}
+
+/** Sampler keys an eval arm may force across every task (greedy/HLG/curve). */
+type SamplerArm = Partial<Pick<GenerateOptions, "temperature" | "topP" | "topK" | "seed" | "hlg" | "curve">>;
+
+export interface TaskModel {
+  model: RuntimeModel;
+  tokenizer: LoadedTokenizer;
+  template: ChatTemplate | null;
+  config: ModelConfig;
+  dir: string;
+  /** When set, FORCES this sampler on every generateText call, overriding the
+   *  task's own (greedy) default — used to run the whole suite under one sampler
+   *  arm (e.g. the v2 curve vs the default chat recipe) for the degradation gate. */
+  samplerOverride?: SamplerArm;
+}
+
+export async function loadTaskModel(query: string): Promise<TaskModel> {
+  const dir = resolveModelDir(query);
+  const config = await loadModelConfig(dir);
+  const weights = await Weights.open(dir);
+  const model = createModel(weights, config);
+  const tokenizer = await loadTokenizer(dir);
+  let template: ChatTemplate | null = null;
+  try { template = await ChatTemplate.load(dir); } catch { template = null; }
+  return { model, tokenizer, template, config, dir };
+}
+
+export interface GenOpts {
+  maxTokens?: number;
+  /** Wrap `body` as a single user turn in the chat template (default true). */
+  useChat?: boolean;
+  /** Sampling overrides (default greedy). Lets the HLG/diversity evals vary the
+   *  sampler — temperature, top-p/k, seed, hlg — while sharing this prompt path. */
+  sampler?: Partial<Pick<GenerateOptions, "temperature" | "topP" | "topK" | "seed" | "hlg" | "curve">>;
+}
+
+/** Complete `body` via the model's real quantized-KV path. Greedy by default;
+ *  pass `opts.sampler` to drive temperature/HLG/etc. */
+export async function generateText(tm: TaskModel, body: string, opts: GenOpts = {}): Promise<string> {
+  const maxTokens = opts.maxTokens ?? 256;
+  const text = opts.useChat !== false && tm.template
+    ? tm.template.render([{ role: "user", content: body }], { addGenerationPrompt: true })
+    : body;
+  const ids = tm.tokenizer.encode(text);
+
+  const kv = tm.config.kvQuant?.length ? tm.config.kvQuant : undefined;
+  const gen = generate(tm.model, ids, {
+    maxTokens,
+    temperature: 0, // greedy default — deterministic head-to-head arms
+    ...(kv ? { kvConfig: kv, quantizedKvStart: 0 } : {}),
+    ...(opts.sampler ?? {}), // overrides temperature when provided
+    ...(tm.samplerOverride ?? {}), // arm override wins over the task's own sampler
+  });
+
+  const out: number[] = [];
+  for await (const { token } of gen) out.push(token);
+  return tm.tokenizer.decode(out, true);
+}

package/src/eval/score.ts

@@ -0,0 +1,49 @@
+// Capability score — port of optiq/eval/score.py.
+//
+// A single scalar: the simple UNWEIGHTED mean of whichever of the six
+// benchmarks (MMLU, GSM8K, IFEval, BFCL, HumanEval, HashHop) were run.
+// Disk size is reported alongside but NEVER enters the score — no hidden
+// quality/disk tradeoff baked into the number (optiq's explicit design
+// choice). Missing components are skipped, not zero-filled.
+
+export interface CapabilityScore {
+  /** 0–100, simple mean of the provided benchmark percents. */
+  score: number;
+  /** Per-benchmark contribution to the mean (only the ones that ran). */
+  components: Record<string, number>;
+  /** Reported separately; does NOT affect `score`. */
+  diskGb: number;
+}
+
+export type TaskName = "MMLU" | "GSM8K" | "IFEval" | "BFCL" | "HumanEval" | "HashHop";
+
+/** Unweighted mean of the provided task percents (each 0–100). */
+export function computeCapabilityScore(
+  percents: Partial<Record<TaskName, number>>,
+  diskGb: number,
+): CapabilityScore {
+  const components: Record<string, number> = {};
+  for (const [name, val] of Object.entries(percents))
+    if (val !== undefined && val !== null && Number.isFinite(val))
+      components[name] = val;
+
+  const keys = Object.keys(components);
+  const score = keys.length === 0
+    ? 0
+    : keys.reduce((a, k) => a + components[k]!, 0) / keys.length;
+
+  return { score, components, diskGb };
+}
+
+/** Sum of `*.safetensors` shard sizes for a local model dir, in GB. */
+export function diskGbForDir(modelDir: string): number {
+  try {
+    const { readdirSync, statSync } = require("node:fs") as typeof import("node:fs");
+    let total = 0;
+    for (const f of readdirSync(modelDir))
+      if (f.endsWith(".safetensors")) total += statSync(`${modelDir}/${f}`).size;
+    return total / 1024 ** 3;
+  } catch {
+    return 0;
+  }
+}