mlx-bun 0.0.4

package/src/download.ts

@@ -0,0 +1,283 @@
+// Resumable HF downloads with checksum verification (Phase 5).
+//
+// The Python downloader failure modes this replaces (PLAN "Context /
+// lore"): Xet stalls (we speak plain HTTPS resolve/CDN), no resume on
+// flaky links (we Range-resume partial blobs), and silent corruption
+// (every blob is verified — sha256 for LFS files against the API's
+// oid, git-blob sha1 for small files against blobId).
+//
+// Writes the exact huggingface_hub cache layout so the registry and
+// loaders find the result with zero changes:
+//   <hub>/models--{org}--{name}/blobs/<sha>            (verified bytes)
+//   <hub>/models--{org}--{name}/snapshots/<commit>/<file> → ../../blobs/<sha>
+//   <hub>/models--{org}--{name}/refs/<revision>        (commit sha)
+//
+// Files download sequentially (decode is bandwidth-bound, the network
+// rarely is; resumability matters more than parallelism here).
+
+import { mkdirSync, existsSync, statSync, createWriteStream, renameSync, rmSync, symlinkSync, readFileSync } from "node:fs";
+import { dirname, join } from "node:path";
+import { DEFAULT_HUB } from "./registry";
+
+const ENDPOINT = "https://huggingface.co";
+
+export interface RepoFile {
+  rfilename: string;
+  size: number;
+  /** git blob sha1 (non-LFS identity, and the blob filename for non-LFS). */
+  blobId: string;
+  /** LFS sha256 + size; present for large files. oid is the blob filename. */
+  lfs?: { oid: string; size: number };
+}
+
+export interface RepoListing {
+  /** Commit sha the revision resolves to (snapshot dir name). */
+  sha: string;
+  files: RepoFile[];
+}
+
+export interface DownloadOptions {
+  revision?: string;
+  cacheDir?: string;
+  endpoint?: string;
+  token?: string | null;
+  onProgress?: (file: string, received: number, total: number) => void;
+}
+
+/** HF auth: explicit > env > the file `hf auth login` writes. */
+export function hfToken(): string | null {
+  if (process.env.HF_TOKEN) return process.env.HF_TOKEN;
+  try {
+    return readFileSync(`${process.env.HOME}/.cache/huggingface/token`, "utf8").trim() || null;
+  } catch {
+    return null;
+  }
+}
+
+/** git blob identity: sha1("blob <size>\0" + content) — what blobId is. */
+export function gitBlobSha1(bytes: Uint8Array): string {
+  const hasher = new Bun.CryptoHasher("sha1");
+  hasher.update(`blob ${bytes.length}\0`);
+  hasher.update(bytes);
+  return hasher.digest("hex");
+}
+
+export async function listRepoFiles(
+  repoId: string, opts: DownloadOptions = {},
+): Promise<RepoListing> {
+  const endpoint = opts.endpoint ?? ENDPOINT;
+  const revision = opts.revision ?? "main";
+  const token = opts.token === undefined ? hfToken() : opts.token;
+  const url = `${endpoint}/api/models/${repoId}/revision/${encodeURIComponent(revision)}?blobs=true`;
+  const res = await fetch(url, {
+    headers: token ? { authorization: `Bearer ${token}` } : {},
+  });
+  if (!res.ok)
+    throw new Error(`HF API ${res.status} for ${repoId}@${revision}: ${(await res.text()).slice(0, 200)}`);
+  const body = (await res.json()) as {
+    sha: string;
+    siblings: { rfilename: string; size?: number; blobId?: string; lfs?: { oid: string; size: number } }[];
+  };
+  return {
+    sha: body.sha,
+    files: body.siblings.map((s) => ({
+      rfilename: s.rfilename,
+      size: s.lfs?.size ?? s.size ?? 0,
+      blobId: s.blobId ?? "",
+      lfs: s.lfs,
+    })),
+  };
+}
+
+/** Follow redirects manually so the Authorization header is DROPPED at
+ *  the CDN hop — presigned S3/CloudFront URLs reject requests that
+ *  carry it (the classic hub-client footgun). Range survives. */
+async function fetchBlob(
+  url: string, token: string | null, rangeStart: number,
+): Promise<Response> {
+  const range: Record<string, string> = {};
+  if (rangeStart > 0) range["range"] = `bytes=${rangeStart}-`;
+  const authed: Record<string, string> = { ...range };
+  if (token) authed["authorization"] = `Bearer ${token}`;
+  let res = await fetch(url, { redirect: "manual", headers: authed });
+  let hops = 0;
+  while (res.status >= 300 && res.status < 400 && hops < 5) {
+    const loc = res.headers.get("location");
+    if (!loc) break;
+    // no auth past the redirect — presigned URLs reject it
+    res = await fetch(new URL(loc, url).href, { redirect: "manual", headers: range });
+    hops++;
+  }
+  return res;
+}
+
+export async function downloadOne(
+  url: string, token: string | null, blobPath: string,
+  expected: { size: number; sha256?: string; sha1?: string; name: string },
+  onProgress?: DownloadOptions["onProgress"],
+): Promise<void> {
+  const partPath = `${blobPath}.incomplete`;
+  let offset = existsSync(partPath) ? statSync(partPath).size : 0;
+  if (offset > expected.size) {
+    rmSync(partPath); // stale junk; start over
+    offset = 0;
+  }
+
+  // Streaming verification — resume re-hashes the existing prefix so the
+  // final digest always covers every byte on disk.
+  const sha256 = new Bun.CryptoHasher("sha256");
+  const sha1 = new Bun.CryptoHasher("sha1");
+  sha1.update(`blob ${expected.size}\0`);
+  if (offset > 0) {
+    const existing = await Bun.file(partPath).arrayBuffer();
+    sha256.update(existing);
+    sha1.update(existing);
+  }
+
+  if (offset < expected.size) {
+    const res = await fetchBlob(url, token, offset);
+    if (res.status === 200 && offset > 0) {
+      // server ignored Range — start the file (and hashes) over
+      rmSync(partPath);
+      return downloadOne(url, token, blobPath, expected, onProgress);
+    }
+    if (res.status !== 200 && res.status !== 206)
+      throw new Error(`download failed (${res.status}) for ${expected.name}`);
+    const out = createWriteStream(partPath, { flags: offset > 0 ? "a" : "w" });
+    try {
+      for await (const chunk of res.body as ReadableStream<Uint8Array>) {
+        sha256.update(chunk);
+        sha1.update(chunk);
+        await new Promise<void>((resolve, reject) =>
+          out.write(chunk, (e) => (e ? reject(e) : resolve())),
+        );
+        offset += chunk.length;
+        onProgress?.(expected.name, offset, expected.size);
+      }
+    } finally {
+      await new Promise<void>((resolve) => out.end(() => resolve()));
+    }
+  }
+
+  if (expected.size === 0 && !existsSync(partPath)) await Bun.write(partPath, "");
+  if (offset !== expected.size) {
+    throw new Error(
+      `short download for ${expected.name}: ${offset} of ${expected.size} bytes ` +
+      `(.incomplete kept — rerun to resume)`,
+    );
+  }
+  const digestOk = expected.sha256
+    ? sha256.digest("hex") === expected.sha256
+    : expected.sha1
+      ? sha1.digest("hex") === expected.sha1
+      : true;
+  if (!digestOk) {
+    rmSync(partPath); // corrupt — never resume from it
+    throw new Error(`checksum mismatch for ${expected.name} — partial removed, rerun to retry`);
+  }
+  renameSync(partPath, blobPath);
+}
+
+/** Download (or complete) a model snapshot. Returns the snapshot path.
+ *  Idempotent: verified blobs are never re-fetched; partial blobs
+ *  resume with a Range request. */
+/** Process-global download tracker — the status page's /downloads
+ *  endpoint reads this. Only downloads performed by THIS process are
+ *  visible (a `mlx-bun get` in another terminal is not). */
+export interface DownloadStatus {
+  repoId: string;
+  state: "active" | "done" | "error";
+  currentFile: string | null;
+  receivedBytes: number;
+  totalBytes: number;
+  filesDone: number;
+  filesTotal: number;
+  /** Rolling-window transfer rate (server-measured, ~5 s window). */
+  bytesPerSec: number;
+  startedAt: number;
+  finishedAt: number | null;
+  error?: string;
+}
+const downloadLog: DownloadStatus[] = [];
+export function downloadsSnapshot(): DownloadStatus[] {
+  return downloadLog.slice(-5);
+}
+
+export async function downloadModel(
+  repoId: string, opts: DownloadOptions = {},
+): Promise<string> {
+  const endpoint = opts.endpoint ?? ENDPOINT;
+  const revision = opts.revision ?? "main";
+  const token = opts.token === undefined ? hfToken() : opts.token;
+  const hub = opts.cacheDir ?? DEFAULT_HUB;
+
+  const listing = await listRepoFiles(repoId, { ...opts, token });
+  const repoDir = join(hub, `models--${repoId.replaceAll("/", "--")}`);
+  const blobsDir = join(repoDir, "blobs");
+  const snapDir = join(repoDir, "snapshots", listing.sha);
+  mkdirSync(blobsDir, { recursive: true });
+  mkdirSync(join(repoDir, "refs"), { recursive: true });
+
+  const status: DownloadStatus = {
+    repoId, state: "active", currentFile: null,
+    receivedBytes: 0, totalBytes: listing.files.reduce((a, f) => a + f.size, 0),
+    filesDone: 0, filesTotal: listing.files.length,
+    bytesPerSec: 0, startedAt: Date.now(), finishedAt: null,
+  };
+  downloadLog.push(status);
+  let doneBytes = 0;
+  // Rolling ~5 s window of (time, receivedBytes) samples for the rate.
+  const samples: Array<[number, number]> = [];
+  const sampleRate = (bytes: number) => {
+    const now = performance.now();
+    samples.push([now, bytes]);
+    while (samples.length > 2 && now - samples[0]![0] > 5000) samples.shift();
+    const [t0, b0] = samples[0]!;
+    status.bytesPerSec = now > t0 ? ((bytes - b0) / (now - t0)) * 1000 : 0;
+  };
+
+  try {
+    for (const f of listing.files) {
+      const blobId = f.lfs?.oid ?? f.blobId;
+      if (!blobId) throw new Error(`no blob id for ${f.rfilename} (API response missing ?blobs=true data)`);
+      const blobPath = join(blobsDir, blobId);
+
+      if (!existsSync(blobPath) || statSync(blobPath).size !== f.size) {
+        const url = `${endpoint}/${repoId}/resolve/${listing.sha}/${f.rfilename}`;
+        status.currentFile = f.rfilename;
+        await downloadOne(url, token, blobPath, {
+          size: f.size, name: f.rfilename,
+          ...(f.lfs ? { sha256: f.lfs.oid } : { sha1: f.blobId }),
+        }, (file, received, total) => {
+          status.receivedBytes = doneBytes + received;
+          sampleRate(status.receivedBytes);
+          opts.onProgress?.(file, received, total);
+        });
+      }
+      doneBytes += f.size;
+      status.receivedBytes = doneBytes;
+      status.filesDone++;
+      status.currentFile = null;
+
+      const linkPath = join(snapDir, f.rfilename);
+      mkdirSync(dirname(linkPath), { recursive: true });
+      if (!existsSync(linkPath)) {
+        // relative target, depth-aware for nested rfilenames
+        const depth = f.rfilename.split("/").length - 1;
+        symlinkSync(join("../".repeat(depth + 2), "blobs", blobId), linkPath);
+      }
+    }
+    status.state = "done";
+    status.bytesPerSec = 0;
+    status.finishedAt = Date.now();
+  } catch (e) {
+    status.state = "error";
+    status.bytesPerSec = 0;
+    status.error = e instanceof Error ? e.message : String(e);
+    status.finishedAt = Date.now();
+    throw e;
+  }
+
+  await Bun.write(join(repoDir, "refs", revision), listing.sha);
+  return snapDir;
+}

package/src/eval/README.md

@@ -0,0 +1,86 @@
+# Quality evals — the acceptance gate for non-bit-exact perf work
+
+Fusing kernels (GELU, norm+add, online-softmax attention) reorders bf16
+float ops, so the model is **no longer bit-exact to the mlx-lm oracle**.
+Logit parity stops being a usable gate. This harness is the replacement —
+a faithful port of **optiq's eval framework**
+(`optiq/eval/` in the oracle venv): KL divergence as the continuous
+acceptance gate, plus the six-task capability score as the ship gate.
+
+## The gate, in two tiers (optiq methodology)
+
+1. **Smoketest (~5 min):** KL divergence (64 prompts × 256 tok, mean/median/p95)
+   + GSM8K-50. Fast triage — "did the distribution drift?"
+2. **Capability suite (~2 h):** MMLU · GSM8K · IFEval · BFCL · HumanEval ·
+   **HashHop@12k**. `Capability_Score = unweighted mean` of the six
+   (`score.ts`). HashHop@12k is the long-context guard — where online-softmax
+   error (opportunity A) is worst.
+
+KL math (`kl.ts`, port of `kl.py`): `KL(ref‖cand)` per token =
+`Σ softmax(ref)·(logsoftmax(ref) − logsoftmax(cand))`. **Teacher-forced** —
+same fixed prompt to both arms, compare per-position distributions (the
+repo rule: free-running greedy "measures chaos").
+
+## KL reference modes
+
+- **self-flag (default — the kernel-drift gate):** one weight load; forward
+  each prompt twice with a perf lever set to its ref vs cand value. Directly
+  measures the drift a lever introduces. Fits in 24 GB (no second model).
+  `KL(compat ‖ fused)` is literally "measure the output difference."
+- **two-model (optiq-style absolute quality):** a separate reference model
+  (bf16 if it fits, else uniform-4bit), both resident.
+
+## Run (on a cleared machine — loads models, not from an agent session)
+
+```sh
+# drift gate on an existing non-bit-exact lever (validates the harness today):
+bun scripts/eval.ts kl --candidate e4b
+
+# once a fused kernel is wired for e4b, its drift vs compat:
+bun scripts/eval.ts kl --candidate e4b --self MLX_BUN_PERF_KERNEL --ref-value 0 --cand-value 1
+
+# optiq-style absolute drift vs a reference model:
+bun scripts/eval.ts kl --candidate e4b --reference <bf16-or-4bit-id>
+
+# higher-fidelity prompts (optiq's calibration mix):
+bun scripts/eval.ts kl --candidate e4b --prompts-file \
+  /Users/joshrossi/Code/mlx-lm/.venv/.../optiq/calibration/data/optiq.jsonl
+```
+
+Rows land in `~/.cache/mlx-bun/evals.sqlite` → `quality_runs` (separate from
+the throughput `runs` table). Each row records the active perf levers
+(`config_json`) so it's self-describing: which head-to-head arm produced it.
+
+## Status / milestones
+
+- **M0 (done):** KL drift gate + capability score + `quality_runs` DB + CLI.
+- **M0b (done):** *serving-path* KL (`evaluateKlServingDecode` in `kl.ts`,
+  `bun scripts/eval.ts kl --decode`). Validated: identical arms → KL 0;
+  `MLX_BUN_FUSED_DECODE` tiled-decode → KL 0.27 (real drift on the quantized
+  path). Motivation below — the compat KL couldn't see it:
+  The compat `forward()` path with `makeCache()` uses **plain bf16 caches**. That is wrong
+  twice over: (1) e4b's KV is **per-layer mixed 4-bit/8-bit** (group 64, from
+  `kv_config.json` → `config.kvQuant`; 8-bit layers 3,4,9,10,11,14, 4-bit
+  elsewhere), NOT bf16; (2) bf16 caches fail the generated class's `#matches`
+  guard, so the model silently falls back to the monolith — a path that
+  doesn't exist in production. M0b must reproduce the serve setup: build
+  caches, then `cache[i] = c.toQuantized(e.groupSize, e.bits)` per
+  `config.kvQuant` (the logic already at `generate.ts:68-86`), so the
+  teacher-forced prefill runs the real mixed-4/8-bit quantized path and the
+  perf levers (NO_FUSED_SDPA, perf-kernel) actually bite.
+- **M1 (done):** all six capability tasks ported + runtime-validated, plus
+  `capability` aggregation and the `smoketest` (KL + GSM8K-50) gate. Datasets
+  exported by `scripts/eval/export-datasets.py` → `~/.cache/mlx-bun/eval-data/`.
+  Tasks run through the real quantized-KV path (`runner.ts`). HumanEval executes
+  in a sandbox-exec + venv-python subprocess; HashHop reconstructs Magic AI's
+  generator (`MultiHopEval` isn't on PyPI). **Open quality checks** (need a real
+  `--n 100+` run to confirm the port, not the harness): MMLU read ~chance at n=5
+  (raw-completion port may underperform for instruct models — consider chat
+  template or logit-scoring); BFCL produced no-calls (gemma-4 `<|tool_call>`
+  format vs the textual-fallback prompt).
+- **M2 (done):** compile-vs-fused paired A/B (`scripts/bench-compile-vs-fused.ts`).
+  Plumbing validated; the `fused` arm is a placeholder (compile-off) until the
+  kernels land — edit `ARMS.fused` to add each kernel's lever.
+- **M3–M4 (next):** the fused kernels themselves (E → B/C → D → A), each gated
+  + measured (KL/capability for drift, the A/B for speed) before it ships.
+- **M5:** head-to-head @600/@8k + KL/capability → decide.

package/src/eval/kl-prompts.ts

@@ -0,0 +1,29 @@
+// Default KL calibration prompts. A small, domain-diverse set (code,
+// prose, math, QA, instructions) so the drift signal isn't specialized to
+// one register. Tiled to `n` prompts by the KL eval, matching optiq's
+// repeat-to-fill behavior.
+//
+// For higher fidelity against the distribution the quantizer was tuned
+// for, pass --prompts-file pointing at optiq's bundled calibration mix
+// (optiq/calibration/data/optiq.jsonl in the oracle venv) — jsonl with
+// {"text": ...} or {"messages": [...]} per line. We do NOT vendor that
+// file (provenance), so this bundled set is the self-contained default.
+
+export const DEFAULT_KL_PROMPTS: string[] = [
+  "The history of the printing press begins in the fifteenth century, when Johannes Gutenberg combined movable type with the screw press. Explain how this changed the spread of information across Europe.",
+  "Write a Python function that takes a list of integers and returns the longest strictly increasing contiguous subsequence. Include a docstring and handle the empty-list case.",
+  "A train leaves Chicago travelling east at 60 miles per hour. Two hours later a second train leaves the same station travelling east at 80 miles per hour. How long after the second train departs does it catch the first?",
+  "Summarize the central argument of the following claim and then give one objection: economic growth and ecological sustainability are fundamentally incompatible over the long run.",
+  "Describe, step by step, how a TCP connection is established and torn down, including the purpose of each flag in the three-way handshake and the role of TIME_WAIT.",
+  "In the kitchen, photosynthesis of ideas rarely occurs, yet a careful cook learns to balance acid, fat, salt, and heat. Continue this passage in the same reflective tone for several sentences.",
+  "Given a 4x4 matrix of integers, write pseudocode to rotate it 90 degrees clockwise in place, and explain why the in-place approach avoids allocating a second matrix.",
+  "Explain the difference between supervised, unsupervised, and reinforcement learning to a curious high-school student, using a single running example for all three.",
+  "Translate the following sentence into formal written English and then into a casual spoken register: 'the meeting got moved cause half the team was out sick.'",
+  "What are the trade-offs between depth-first and breadth-first search for finding the shortest path in an unweighted graph? When does each fail or waste work?",
+  "A baker has 3 cups of flour and each loaf needs 3/4 of a cup. After making as many loaves as possible, how much flour is left, and how many loaves were made? Show the arithmetic.",
+  "Outline the plot of a short story in which a lighthouse keeper discovers that the light has been guiding something other than ships. Keep it to four sentences.",
+  "Explain what a hash collision is, why it matters for hash tables, and how separate chaining and open addressing each respond to one. Be precise about average-case behavior.",
+  "Compare the causes commonly cited for the fall of the Western Roman Empire, and explain why historians treat single-cause explanations with suspicion.",
+  "Refactor the idea of 'technical debt' for a non-engineering executive: define it, give a concrete example, and describe one healthy and one unhealthy way teams accumulate it.",
+  "Derive the formula for the sum of the first n positive integers two different ways: once by pairing terms, and once by induction. State the result clearly.",
+];