mlx-bun 0.0.4

package/src/quantize/sensitivity.ts

@@ -0,0 +1,234 @@
+// Per-layer quantization sensitivity analysis (native TS port of
+// optiq/core/sensitivity.py `analyze_sensitivity_exact`, uniform-baseline path).
+//
+// For each quantizable Linear in a loaded (already-quantized) model:
+//   1. dequantize its packed weight back to bf16 — this is the bf16 "source"
+//      the OptIQ uniform_4bit path streams off disk; here it's self-contained
+//      because the running model already carries quantized weights.
+//   2. for each candidate bit-width, re-quantize that bf16 source at the new
+//      bits and SWAP it into the running QuantizedLinear (mutate w/scales/biases
+//      + the QuantSpec bits/groupSize).
+//   3. run a no-KV full-sequence forward on each calibration sample, compute
+//      KL(reference || current) (port of `_kl_from_ref`), average over samples.
+//   4. restore the layer's original quantized weight.
+//
+// Reference logits = the UNMODIFIED running model (taken once up front). When a
+// candidate bit equals the layer's baseline bit-width, re-quantizing reproduces
+// exactly what the layer already is, so KL is 0 by construction — we record 0.0
+// and skip the forward passes (port of the OptIQ baseline-bit short-circuit).
+//
+// Faithful to: `_simulate_quantize` (quantize→dequantize via MlxArray ops),
+// `_mutate_quantized_layer_to_bits`, `_kl_from_ref`, the uniform-4bit-baseline
+// path, and the bits==baseline ⇒ 0 short-circuit.
+
+import { Dtype } from "../mlx/ffi";
+import { cpuStream, MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { trainForward } from "../train/forward";
+import type { RuntimeModel } from "../model/factory";
+import type { QuantizedLinear } from "../model/gemma4-base";
+
+/** Per-layer KL sensitivity at each candidate bit-width (port of
+ *  optiq SensitivityResult). */
+export interface SensitivityResult {
+  layerName: string;
+  /** bit-width → mean KL(reference || that-bit) over calibration samples. */
+  sensitivities: Record<number, number>;
+  paramCount: number;
+}
+
+export interface AnalyzeSensitivityOptions {
+  candidateBits?: number[];
+  groupSize?: number;
+}
+
+// --------------------------------------------------------------------------
+// Minimal in-place mutation of a QuantizedLinear.
+//
+// gemma4-base.ts declares QuantizedLinear's w/scales/biases/spec as `readonly`
+// (they never change at inference time). The sensitivity sweep needs to swap a
+// layer's quantized weight transiently. Rather than edit the model file, we
+// poke the public fields here via a typed cast — a quantize-only mutation
+// helper. The fields are public, so this is a write to public state, not a
+// reach into private internals. Flagged in the port report.
+// --------------------------------------------------------------------------
+
+/** Mutable view of the public QuantizedLinear fields we swap. */
+type MutableQuantizedLinear = {
+  w: MlxArray;
+  scales: MlxArray;
+  biases: MlxArray | null;
+  spec: ops.QuantSpec;
+};
+
+/** Snapshot of a QuantizedLinear's quantized state (for restore). */
+interface LayerState {
+  w: MlxArray;
+  scales: MlxArray;
+  biases: MlxArray | null;
+  spec: ops.QuantSpec;
+}
+
+function captureState(layer: QuantizedLinear): LayerState {
+  return { w: layer.w, scales: layer.scales, biases: layer.biases, spec: layer.spec };
+}
+
+function restoreState(layer: QuantizedLinear, state: LayerState): void {
+  const m = layer as unknown as MutableQuantizedLinear;
+  // Dispose the transient swapped-in tensors before restoring the originals.
+  if (m.w !== state.w) m.w.dispose();
+  if (m.scales !== state.scales) m.scales.dispose();
+  if (m.biases && m.biases !== state.biases) m.biases.dispose();
+  m.w = state.w;
+  m.scales = state.scales;
+  m.biases = state.biases;
+  m.spec = state.spec;
+}
+
+/** Quantize `bf16Weight` at `newBits`/`groupSize` and swap it into `layer`
+ *  in place (port of `_mutate_quantized_layer_to_bits`). The previously
+ *  swapped-in tensors (if any beyond the captured original) are released by
+ *  the caller via restoreState; this only installs the new ones. */
+function mutateLayerToBits(
+  layer: QuantizedLinear,
+  bf16Weight: MlxArray,
+  newBits: number,
+  groupSize: number,
+  original: LayerState,
+): void {
+  const q = ops.quantize(bf16Weight, groupSize, newBits, "affine", cpuStream);
+  const m = layer as unknown as MutableQuantizedLinear;
+  // Free any previously swapped-in transient tensors (not the original).
+  if (m.w !== original.w) m.w.dispose();
+  if (m.scales !== original.scales) m.scales.dispose();
+  if (m.biases && m.biases !== original.biases) m.biases.dispose();
+  m.w = q.packed;
+  m.scales = q.scales;
+  m.biases = q.biases;
+  m.spec = { bits: newBits, groupSize, mode: "affine" };
+  ops.evalAll([q.packed, q.scales, q.biases]);
+}
+
+// --------------------------------------------------------------------------
+// KL primitive (port of _kl_from_ref)
+// --------------------------------------------------------------------------
+
+/** KL(reference || current) averaged across batch/seq (port of `_kl_from_ref`):
+ *    log_cur = cur - logsumexp(cur); log_ref = ref - logsumexp(ref)
+ *    ref_probs = softmax(ref)
+ *    mean( sum( ref_probs * (log_ref - log_cur), axis=-1 ) )
+ *  Returns the scalar KL as a JS number. Casts to f32 for stability. */
+export function klFromRef(curLogits: MlxArray, refLogits: MlxArray): number {
+  const cur = curLogits.dtype === Dtype.float32 ? curLogits : curLogits.astype(Dtype.float32);
+  const ref = refLogits.dtype === Dtype.float32 ? refLogits : refLogits.astype(Dtype.float32);
+
+  const lseCur = ops.logsumexpAxis(cur, -1, true);
+  const lseRef = ops.logsumexpAxis(ref, -1, true);
+  const logCur = ops.sub(cur, lseCur);
+  const logRef = ops.sub(ref, lseRef);
+  const refProbs = ops.softmaxAxis(ref, -1, true);
+
+  const diff = ops.sub(logRef, logCur);
+  const weighted = ops.mul(refProbs, diff);
+  const summed = ops.sumAxis(weighted, -1, false);
+  const meaned = ops.meanAll(summed, false);
+
+  const val = meaned.toFloat32()[0]!;
+
+  for (const a of [lseCur, lseRef, logCur, logRef, refProbs, diff, weighted, summed, meaned]) {
+    a.dispose();
+  }
+  if (cur !== curLogits) cur.dispose();
+  if (ref !== refLogits) ref.dispose();
+  return val;
+}
+
+// --------------------------------------------------------------------------
+// Public entry: analyzeSensitivityExact
+// --------------------------------------------------------------------------
+
+/** Per-layer KL sensitivity of an already-quantized MLX model.
+ *
+ *  @param model        loaded RuntimeModel (its QuantizedLinears are the
+ *                      uniform baseline; their dequantized weights are the
+ *                      self-contained bf16 source).
+ *  @param layers       the quantizable linears to probe, keyed by module path
+ *                      (typically `model.loraTargets()`).
+ *  @param calIds       calibration samples — each a row of int32 token ids.
+ *  @param options      candidateBits (default [4,8]) and groupSize (default 64).
+ */
+export function analyzeSensitivityExact(
+  model: RuntimeModel,
+  layers: Map<string, QuantizedLinear>,
+  calIds: number[][],
+  options: AnalyzeSensitivityOptions = {},
+  onProgress?: (done: number, total: number, layerName: string) => void,
+): SensitivityResult[] {
+  const candidateBits = [...(options.candidateBits ?? [4, 8])];
+  const groupSize = options.groupSize ?? 64;
+
+  // Calibration sample arrays (int32 [1, L]).
+  const calArrays = calIds.map((ids) =>
+    MlxArray.fromInt32(new Int32Array(ids), [1, ids.length]),
+  );
+
+  try {
+    // Reference logits from the unmodified running model (taken once).
+    const refLogits: MlxArray[] = calArrays.map((ids) => trainForward(model, ids));
+    ops.evalAll(refLogits);
+
+    const entries = [...layers.entries()];
+    const results: SensitivityResult[] = [];
+
+    try {
+      for (let li = 0; li < entries.length; li++) {
+        const [layerName, layer] = entries[li]!;
+        const original = captureState(layer);
+        const baselineBits = layer.spec.bits;
+
+        // Self-contained bf16 source: dequantize the layer's own packed weight.
+        const bf16Source = ops.dequantize(
+          original.w, original.scales, original.biases, original.spec, cpuStream,
+        ).astype(Dtype.bfloat16, cpuStream);
+        ops.evalAll([bf16Source]);
+
+        const sensitivities: Record<number, number> = {};
+
+        for (const bits of candidateBits) {
+          if (bits === baselineBits) {
+            // KL is 0 by construction (re-quantizing to the baseline bit-width
+            // reproduces the running model). Skip the forwards.
+            sensitivities[bits] = 0.0;
+            continue;
+          }
+          mutateLayerToBits(layer, bf16Source, bits, groupSize, original);
+
+          let klSum = 0;
+          for (let i = 0; i < calArrays.length; i++) {
+            const cur = trainForward(model, calArrays[i]!);
+            klSum += klFromRef(cur, refLogits[i]!);
+            cur.dispose();
+          }
+          sensitivities[bits] = klSum / Math.max(calArrays.length, 1);
+        }
+
+        // Restore the original quantized weight; release the bf16 source.
+        restoreState(layer, original);
+        bf16Source.dispose();
+
+        // param count = product of the dequantized (full) weight shape.
+        const fullShape = [layer.outFeatures, layer.inFeatures];
+        const paramCount = fullShape.reduce((a, b) => a * b, 1);
+
+        results.push({ layerName, sensitivities, paramCount });
+        onProgress?.(li + 1, entries.length, layerName);
+      }
+    } finally {
+      for (const r of refLogits) r.dispose();
+    }
+
+    return results;
+  } finally {
+    for (const a of calArrays) a.dispose();
+  }
+}

package/src/registry.ts

@@ -0,0 +1,336 @@
+// Model registry: bun:sqlite index over the HF cache. Answers questions
+// like "vision-capable models under 10 GB" without shell archaeology.
+//
+// scan() walks ~/.cache/huggingface/hub/models--*/snapshots/*, reading
+// only config.json + the safetensors index header (never tensor bytes).
+
+import { Database } from "bun:sqlite";
+import {
+  closeSync, existsSync, openSync, readdirSync, readFileSync, readSync, statSync,
+} from "node:fs";
+import { join } from "node:path";
+import { isDrafterModelType } from "./model/support";
+
+export interface ModelRecord {
+  path: string;
+  repoId: string;
+  modelType: string;
+  paramCount: number | null;
+  /** Language-model weight bytes (model-*.safetensors, sidecar excluded). */
+  sizeBytes: number;
+  /** optiq_vision.safetensors bytes (bf16 sidecar; loads only for vision). */
+  sidecarBytes: number;
+  /** Bytes of `.experts.` tensors (MoE; 0 for dense models). Per-token
+   *  decode only reads top_k/num_experts of these. */
+  expertsBytes: number;
+  quantBits: number | null;
+  quantGroupSize: number | null;
+  quantMode: string | null;
+  hasVisionSidecar: boolean;
+  hasKvConfig: boolean;
+  hasToolTemplate: boolean;
+  numLayers: number | null;
+  hiddenSize: number | null;
+  vocabSize: number | null;
+  /** SPDX-ish license id from the model card's README frontmatter
+   *  (e.g. "gemma" — custom terms — vs "apache-2.0"); null if absent.
+   *  Surfaces license obligations per model in `ls`/`fit` output. */
+  license: string | null;
+  scannedAt: number;
+}
+
+const SCHEMA = `
+CREATE TABLE IF NOT EXISTS models (
+  path TEXT PRIMARY KEY,
+  repo_id TEXT NOT NULL,
+  model_type TEXT NOT NULL,
+  param_count INTEGER,
+  size_bytes INTEGER NOT NULL,
+  sidecar_bytes INTEGER NOT NULL DEFAULT 0,
+  experts_bytes INTEGER NOT NULL DEFAULT 0,
+  quant_bits INTEGER,
+  quant_group_size INTEGER,
+  quant_mode TEXT,
+  has_vision_sidecar INTEGER NOT NULL,
+  has_kv_config INTEGER NOT NULL,
+  has_tool_template INTEGER NOT NULL,
+  num_layers INTEGER,
+  hidden_size INTEGER,
+  vocab_size INTEGER,
+  license TEXT,
+  scanned_at INTEGER NOT NULL
+);
+`;
+
+export const DEFAULT_HUB = `${process.env.HOME}/.cache/huggingface/hub`;
+export const DEFAULT_DB = `${process.env.HOME}/.cache/mlx-bun/registry.sqlite`;
+
+export class Registry {
+  readonly db: Database;
+
+  constructor(dbPath: string = DEFAULT_DB) {
+    if (dbPath !== ":memory:") {
+      const dir = dbPath.slice(0, dbPath.lastIndexOf("/"));
+      try { require("node:fs").mkdirSync(dir, { recursive: true }); } catch {}
+    }
+    this.db = new Database(dbPath);
+    this.db.exec(SCHEMA);
+    // The registry is a derived cache — on schema drift, rebuild from scratch.
+    const cols = (this.db.query("PRAGMA table_info(models)").all() as { name: string }[])
+      .map((c) => c.name);
+    if (
+      !cols.includes("sidecar_bytes") ||
+      !cols.includes("experts_bytes") ||
+      !cols.includes("license")
+    ) {
+      this.db.exec("DROP TABLE models");
+      this.db.exec(SCHEMA);
+    }
+  }
+
+  async scan(hubDir: string = DEFAULT_HUB): Promise<number> {
+    if (!existsSync(hubDir)) return 0;
+    // The cache is the source of truth and we only ever INSERT, so reap rows
+    // whose snapshot dir was deleted (else they linger as phantom matches).
+    const prune = this.db.prepare("DELETE FROM models WHERE path = $path");
+    for (const r of this.db.query("SELECT path FROM models").all() as { path: string }[])
+      if (!existsSync(r.path)) prune.run({ $path: r.path });
+    let count = 0;
+    const upsert = this.db.prepare(`
+      INSERT OR REPLACE INTO models VALUES
+      ($path, $repo, $type, $params, $size, $sidecar, $experts, $bits, $gs, $mode,
+       $vision, $kv, $tools, $layers, $hidden, $vocab, $license, $at)
+    `);
+    for (const entry of readdirSync(hubDir)) {
+      if (!entry.startsWith("models--")) continue;
+      const repoId = entry.slice("models--".length).replaceAll("--", "/");
+      const snapsDir = join(hubDir, entry, "snapshots");
+      if (!existsSync(snapsDir)) continue;
+      for (const snap of readdirSync(snapsDir)) {
+        const dir = join(snapsDir, snap);
+        const rec = await scanSnapshot(dir, repoId);
+        if (!rec) continue;
+        upsert.run({
+          $path: rec.path, $repo: rec.repoId, $type: rec.modelType,
+          $params: rec.paramCount, $size: rec.sizeBytes,
+          $sidecar: rec.sidecarBytes, $experts: rec.expertsBytes,
+          $bits: rec.quantBits, $gs: rec.quantGroupSize, $mode: rec.quantMode,
+          $vision: rec.hasVisionSidecar ? 1 : 0,
+          $kv: rec.hasKvConfig ? 1 : 0,
+          $tools: rec.hasToolTemplate ? 1 : 0,
+          $layers: rec.numLayers, $hidden: rec.hiddenSize, $vocab: rec.vocabSize,
+          $license: rec.license, $at: rec.scannedAt,
+        });
+        count++;
+      }
+    }
+    return count;
+  }
+
+  list(filter: { vision?: boolean; maxBytes?: number; query?: string } = {}): ModelRecord[] {
+    const clauses: string[] = [];
+    const params: Record<string, unknown> = {};
+    if (filter.vision !== undefined) {
+      clauses.push("has_vision_sidecar = $vision");
+      params.$vision = filter.vision ? 1 : 0;
+    }
+    if (filter.maxBytes !== undefined) {
+      clauses.push("size_bytes <= $max");
+      params.$max = filter.maxBytes;
+    }
+    if (filter.query) {
+      clauses.push("(repo_id LIKE $q OR model_type LIKE $q)");
+      params.$q = `%${filter.query}%`;
+    }
+    const where = clauses.length ? `WHERE ${clauses.join(" AND ")}` : "";
+    const rows = this.db
+      .query(`SELECT * FROM models ${where} ORDER BY size_bytes ASC`)
+      .all(params as never) as Record<string, unknown>[];
+    // Self-heal: never surface a row whose snapshot dir was deleted out from
+    // under us (scan() reaps these from the DB; this keeps reads correct even
+    // before the next scan).
+    return rows.map(rowToRecord).filter((r) => existsSync(r.path));
+  }
+
+  /** Resolve a fuzzy query to exactly one model (error listing candidates otherwise).
+   *  Speculative-decoding drafters are companion artifacts, never selectable on
+   *  their own, so they never count as candidates here. */
+  resolve(query: string): ModelRecord {
+    const matches = this.list({ query }).filter((m) => !isDrafterModelType(m.modelType));
+    if (matches.length === 0) throw new Error(`no model matching "${query}" — run \`mlx-bun scan\``);
+    // The HF cache can hold several revisions of one repo (snapshots/<hash>
+    // dirs), each a registry row. Resolving a repo name must not be "ambiguous"
+    // just because a stale revision lingers — collapse same-repo matches to the
+    // canonical snapshot. Genuinely different repos stay ambiguous.
+    const repos = [...new Set(matches.map((m) => m.repoId))];
+    if (repos.length === 1) return pickCanonicalRevision(matches);
+    throw new Error(`"${query}" is ambiguous:\n` + repos.map((r) => `  ${r}`).join("\n"));
+  }
+
+  close(): void {
+    this.db.close();
+  }
+}
+
+/** Of several cached revisions of ONE repo, pick the canonical snapshot: the
+ *  revision refs/main points at, else the most recently scanned (stable
+ *  tie-break on path). Paths are `<hub>/models--<repo>/snapshots/<hash>`. */
+function pickCanonicalRevision(matches: ModelRecord[]): ModelRecord {
+  if (matches.length === 1) return matches[0]!;
+  const root = matches[0]!.path.split("/snapshots/")[0]!;
+  try {
+    const head = readFileSync(join(root, "refs", "main"), "utf8").trim();
+    const onMain = matches.find((m) => m.path.includes(`/snapshots/${head}`));
+    if (onMain) return onMain;
+  } catch {
+    /* no refs/main — fall through to recency */
+  }
+  return [...matches].sort(
+    (a, b) => b.scannedAt - a.scannedAt || (a.path < b.path ? -1 : 1),
+  )[0]!;
+}
+
+function rowToRecord(r: Record<string, unknown>): ModelRecord {
+  return {
+    path: r.path as string,
+    repoId: r.repo_id as string,
+    modelType: r.model_type as string,
+    paramCount: r.param_count as number | null,
+    sizeBytes: r.size_bytes as number,
+    sidecarBytes: r.sidecar_bytes as number,
+    expertsBytes: r.experts_bytes as number,
+    quantBits: r.quant_bits as number | null,
+    quantGroupSize: r.quant_group_size as number | null,
+    quantMode: r.quant_mode as string | null,
+    hasVisionSidecar: !!r.has_vision_sidecar,
+    hasKvConfig: !!r.has_kv_config,
+    hasToolTemplate: !!r.has_tool_template,
+    numLayers: r.num_layers as number | null,
+    hiddenSize: r.hidden_size as number | null,
+    vocabSize: r.vocab_size as number | null,
+    license: r.license as string | null,
+    scannedAt: r.scanned_at as number,
+  };
+}
+
+/** License id from the model card's README.md YAML frontmatter
+ *  (`license:`; `license_name:` wins when license is "other"). */
+function readmeLicense(dir: string): string | null {
+  const p = join(dir, "README.md");
+  if (!existsSync(p)) return null;
+  try {
+    const head = readFileSync(p, "utf8").slice(0, 4096);
+    if (!head.startsWith("---")) return null;
+    const end = head.indexOf("\n---", 3);
+    const fm = end === -1 ? head : head.slice(0, end);
+    const strip = (s: string) => s.trim().replace(/^['"]|['"]$/g, "");
+    const lic = strip(/^license:\s*(.+)$/m.exec(fm)?.[1] ?? "");
+    const name = /^license_name:\s*(.+)$/m.exec(fm)?.[1];
+    if (lic === "other" && name) return strip(name);
+    return lic || null;
+  } catch {
+    return null;
+  }
+}
+
+/** Sum the byte sizes of `.experts.` tensors from a safetensors header
+ *  (header JSON only — never touches tensor data). */
+function expertTensorBytes(path: string): number {
+  const fd = openSync(path, "r");
+  try {
+    const lenBuf = Buffer.alloc(8);
+    readSync(fd, lenBuf, 0, 8, 0);
+    const headerLen = Number(lenBuf.readBigUInt64LE(0));
+    if (headerLen <= 0 || headerLen > 256 * 2 ** 20) return 0;
+    const headerBuf = Buffer.alloc(headerLen);
+    readSync(fd, headerBuf, 0, headerLen, 8);
+    const header = JSON.parse(headerBuf.toString("utf8")) as Record<
+      string, { data_offsets?: [number, number] }
+    >;
+    let bytes = 0;
+    for (const [name, entry] of Object.entries(header)) {
+      if (!name.includes(".experts.") || !entry.data_offsets) continue;
+      bytes += entry.data_offsets[1] - entry.data_offsets[0];
+    }
+    return bytes;
+  } finally {
+    closeSync(fd);
+  }
+}
+
+async function scanSnapshot(dir: string, repoId: string): Promise<ModelRecord | null> {
+  const configPath = join(dir, "config.json");
+  if (!existsSync(configPath)) return null;
+
+  // size: sum of model safetensors (resolved through HF's blob symlinks).
+  // The vision sidecar is its own line item — it loads only for vision
+  // requests and must never be folded into language-weight fit math.
+  let sizeBytes = 0;
+  let sidecarBytes = 0;
+  let expertsBytes = 0;
+  let hasWeights = false;
+  for (const f of readdirSync(dir)) {
+    if (!f.endsWith(".safetensors")) continue;
+    hasWeights = true;
+    let bytes = 0;
+    try { bytes = statSync(join(dir, f)).size; } catch {}
+    if (f === "optiq_vision.safetensors") {
+      sidecarBytes += bytes;
+    } else {
+      sizeBytes += bytes;
+      try { expertsBytes += expertTensorBytes(join(dir, f)); } catch {}
+    }
+  }
+  if (!hasWeights) return null;
+
+  let config: Record<string, any>;
+  try {
+    config = (await Bun.file(configPath).json()) as Record<string, any>;
+  } catch {
+    return null;
+  }
+  const text = (config.text_config ?? config) as Record<string, any>;
+  const quant = (config.quantization ?? config.quantization_config) as
+    | Record<string, any> | undefined;
+
+  let paramCount: number | null = null;
+  const idxPath = join(dir, "model.safetensors.index.json");
+  if (existsSync(idxPath)) {
+    try {
+      const idx = (await Bun.file(idxPath).json()) as any;
+      paramCount = idx.metadata?.total_parameters ?? null;
+    } catch {}
+  }
+
+  let hasToolTemplate = false;
+  for (const tf of ["chat_template.jinja", "tokenizer_config.json"]) {
+    const p = join(dir, tf);
+    if (existsSync(p)) {
+      try {
+        const body = await Bun.file(p).text();
+        if (body.includes("tool_call")) { hasToolTemplate = true; break; }
+      } catch {}
+    }
+  }
+
+  return {
+    path: dir,
+    repoId,
+    modelType: (config.model_type as string) ?? "unknown",
+    paramCount,
+    sidecarBytes,
+    expertsBytes,
+    sizeBytes,
+    quantBits: quant?.bits ?? null,
+    quantGroupSize: quant?.group_size ?? null,
+    quantMode: quant?.mode ?? (quant ? "affine" : null),
+    hasVisionSidecar: existsSync(join(dir, "optiq_vision.safetensors")),
+    hasKvConfig: existsSync(join(dir, "kv_config.json")),
+    hasToolTemplate,
+    numLayers: text.num_hidden_layers ?? null,
+    hiddenSize: text.hidden_size ?? null,
+    vocabSize: text.vocab_size ?? null,
+    license: readmeLicense(dir),
+    scannedAt: Date.now(),
+  };
+}