mlx-bun 0.0.4

package/src/jobs/types.ts

@@ -0,0 +1,58 @@
+// Shared job-system contract. Both the job runner (src/jobs/*) and every
+// job producer (src/dataset/*, src/quantize/*, src/train/*) speak this
+// vocabulary. Mirrors the optiq lab SSE event grammar so the web wizards
+// port over mechanically (started → log/stage/metric → done|failed).
+
+/** A single streamed job event, serialized one-per-line into the job's
+ *  NDJSON log and re-emitted to the browser as SSE `data:` frames. */
+export type JobEvent =
+  | { type: "started"; ts: number }
+  | { type: "log"; line: string }
+  | {
+      type: "stage";
+      stage: string;
+      progress?: number; // 0..1
+      message?: string;
+      [k: string]: unknown;
+    }
+  | {
+      type: "metric";
+      kind: "train" | "val";
+      step: number;
+      loss: number;
+      [k: string]: unknown; // grad_norm, learning_rate, tokens_per_sec, accuracy, margin, progress, message
+    }
+  | { type: "done"; ts: number; output_dir?: string; summary?: unknown; [k: string]: unknown }
+  | { type: "failed"; error: string; ts: number };
+
+/** Sink a job calls to report progress. Implementations append to the log
+ *  file and (for stage/metric events carrying `progress`/`message`) update
+ *  the SQLite row. Never throws — a logging failure must not kill a job. */
+export type Emit = (e: JobEvent) => void;
+
+/** The unit of work. Returns an optional output path recorded on the row.
+ *  Receives the parsed `config` from the submit request. */
+export type JobRunner = (
+  emit: Emit,
+  config: Record<string, unknown>,
+) => Promise<{ outputPath?: string } | void>;
+
+/** Job kinds in the system. */
+export type JobKind = "quantize" | "finetune" | "dataset" | "noop" | "crash";
+
+export type JobStatus = "queued" | "running" | "done" | "failed" | "zombie";
+
+/** A persisted job row (shape returned by the HTTP `GET /api/jobs/:id`). */
+export interface JobRow {
+  id: string;
+  kind: string;
+  status: JobStatus;
+  config_json: string;
+  progress: number;
+  message: string | null;
+  log_path: string;
+  output_path: string | null;
+  error: string | null;
+  started_at: string;
+  ended_at: string | null;
+}

package/src/kv-store.ts

@@ -0,0 +1,185 @@
+// KV-cache persistence: save prompt caches to disk, reload zero-copy.
+//
+// File layout (every tensor PAGE-ALIGNED — the Phase 1 corollary: files
+// we write can be mmap'd and handed to the GPU without copies):
+//   [magic "MLXBUNKV1\n"][u32 LE header length][JSON header][padding]
+//   [tensor data at 16 KiB-aligned offsets]
+// Header: { tokens, caches: [{ kind, offset, idx?, maxSize?, dtype,
+//           kShape, vShape, kOffset, vOffset }] }
+//
+// Reload mmaps copy-on-write (MAP_PRIVATE): if mlx ever donates one of
+// these buffers, writes hit private pages, never the file.
+
+import { openSync, writeSync, closeSync } from "node:fs";
+import { MmapFile } from "./mmap";
+import { MlxArray } from "./mlx/array";
+import type { Dtype } from "./mlx/ffi";
+import * as ops from "./mlx/ops";
+import {
+  type Cache, KVCache, RotatingKVCache, Gemma4Model,
+} from "./model/gemma4";
+
+/** Contiguous raw bytes of (possibly a view of) an array. */
+function contiguousBytes(a: MlxArray): Uint8Array {
+  const c = ops.contiguous(a);
+  const bytes = c.rawBytes();
+  c.dispose();
+  return bytes;
+}
+
+const MAGIC = "MLXBUNKV1\n";
+const ALIGN = 16384;
+
+interface CacheHeaderEntry {
+  kind: "kv" | "rotating";
+  offset: number;
+  idx?: number;
+  maxSize?: number;
+  dtype: number;
+  kShape: number[];
+  vShape: number[];
+  kOffset: number;
+  vOffset: number;
+  kBytes: number;
+  vBytes: number;
+}
+
+export interface KvFileHeader {
+  tokens: number[];
+  caches: CacheHeaderEntry[];
+}
+
+const alignUp = (n: number) => Math.ceil(n / ALIGN) * ALIGN;
+
+export function saveKvCache(path: string, tokens: number[], caches: Cache[]): void {
+  const entries: CacheHeaderEntry[] = [];
+  const blobs: Uint8Array[] = [];
+
+  // compact each cache to its live region and snapshot bytes
+  let dataOffset = 0; // relative; rebased after header is sized
+  for (const c of caches) {
+    let keys: MlxArray;
+    let values: MlxArray;
+    let kind: "kv" | "rotating";
+    let idx: number | undefined;
+    let maxSize: number | undefined;
+    if (c instanceof RotatingKVCache) {
+      kind = "rotating";
+      idx = c.ringIdx;
+      maxSize = c.maxSize;
+      if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
+      keys = c.keys;
+      values = c.values;
+    } else if (c instanceof KVCache) {
+      kind = "kv";
+      if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
+      // store only the live [.., :offset, :] region
+      const [B, H, , D] = c.keys.shape as [number, number, number, number];
+      const vD = c.values.shape[3]!;
+      keys = c.keys.slice([0, 0, 0, 0], [B, H, c.offset, D]);
+      values = c.values.slice([0, 0, 0, 0], [B, H, c.offset, vD]);
+    } else {
+      throw new Error("unknown cache type");
+    }
+
+    const kBytesArr = contiguousBytes(keys);
+    const vBytesArr = contiguousBytes(values);
+    const kOffset = alignUp(dataOffset);
+    const vOffset = alignUp(kOffset + kBytesArr.length);
+    dataOffset = vOffset + vBytesArr.length;
+
+    entries.push({
+      kind, offset: c.offset, idx, maxSize,
+      dtype: keys.dtype,
+      kShape: keys.shape, vShape: values.shape,
+      kOffset, vOffset,
+      kBytes: kBytesArr.length, vBytes: vBytesArr.length,
+    });
+    blobs.push(kBytesArr, vBytesArr);
+    if (c instanceof KVCache) {
+      keys.dispose();
+      values.dispose();
+    }
+  }
+
+  const header: KvFileHeader = { tokens, caches: entries };
+  const headerJson = new TextEncoder().encode(JSON.stringify(header));
+  const dataStart = alignUp(MAGIC.length + 8 + headerJson.length);
+
+  const fd = openSync(path, "w");
+  try {
+    const pre = new Uint8Array(dataStart);
+    pre.set(new TextEncoder().encode(MAGIC), 0);
+    const dv = new DataView(pre.buffer);
+    dv.setUint32(MAGIC.length, headerJson.length, true);
+    dv.setUint32(MAGIC.length + 4, dataStart, true);
+    pre.set(headerJson, MAGIC.length + 8);
+    writeSync(fd, pre, 0, pre.length, 0);
+
+    let blobIdx = 0;
+    for (const e of entries) {
+      writeSync(fd, blobs[blobIdx]!, 0, e.kBytes, dataStart + e.kOffset);
+      writeSync(fd, blobs[blobIdx + 1]!, 0, e.vBytes, dataStart + e.vOffset);
+      blobIdx += 2;
+    }
+  } finally {
+    closeSync(fd);
+  }
+}
+
+/** Read only the header (cheap — for prefix matching across many files). */
+export function readKvHeader(path: string): KvFileHeader & { dataStart: number } {
+  const fd = openSync(path, "r");
+  try {
+    const head = new Uint8Array(MAGIC.length + 8);
+    require("node:fs").readSync(fd, head, 0, head.length, 0);
+    if (new TextDecoder().decode(head.subarray(0, MAGIC.length)) !== MAGIC)
+      throw new Error(`${path}: not an mlx-bun KV cache file`);
+    const dv = new DataView(head.buffer);
+    const len = dv.getUint32(MAGIC.length, true);
+    const dataStart = dv.getUint32(MAGIC.length + 4, true);
+    const body = new Uint8Array(len);
+    require("node:fs").readSync(fd, body, 0, len, MAGIC.length + 8);
+    const header = JSON.parse(new TextDecoder().decode(body)) as KvFileHeader;
+    return { ...header, dataStart };
+  } finally {
+    closeSync(fd);
+  }
+}
+
+export interface LoadedKvCache {
+  tokens: number[];
+  caches: Cache[];
+  /** Keep referenced as long as the caches live. */
+  mmap: MmapFile;
+}
+
+export function loadKvCache(path: string, model: Gemma4Model): LoadedKvCache {
+  const header = readKvHeader(path);
+  const mmap = MmapFile.open(path, "cow");
+  const dataStart = header.dataStart;
+
+  const caches: Cache[] = [];
+  for (const e of header.caches) {
+    const keys = MlxArray.fromPointer(
+      mmap.pointer(dataStart + e.kOffset), e.kShape, e.dtype as Dtype,
+    );
+    const values = MlxArray.fromPointer(
+      mmap.pointer(dataStart + e.vOffset), e.vShape, e.dtype as Dtype,
+    );
+    if (e.kind === "rotating") {
+      const c = new RotatingKVCache(e.maxSize!);
+      c.restoreState(keys, values, e.offset, e.idx!);
+      caches.push(c);
+    } else {
+      const c = new KVCache();
+      c.restoreState(keys, values, e.offset);
+      caches.push(c);
+    }
+  }
+  if (caches.length !== model.layers.length)
+    throw new Error(
+      `${path}: ${caches.length} cached layers but model has ${model.layers.length}`,
+    );
+  return { tokens: header.tokens, caches, mmap };
+}

package/src/lora.ts

@@ -0,0 +1,314 @@
+// Hot-swap mounted LoRA adapters — port of optiq/adapters/{mount,
+// registry,resolver}.py (serving side; lora/apply.py is the training-
+// side rank logic). Mount N adapters on one quantized base, select per
+// request by id, never reload the base.
+//
+// Deviations from the reference, both deliberate (PLAN Phase 8):
+// - No ContextVar/serve-pin: our generation queue is serialized, so the
+//   active adapter is a plain field (LoraState) set by generate().
+// - Residual composition is mlx-lm LoRALinear / optiq apply.py
+//   (`y + (scale·z).astype(x.dtype)`), not mount.py's uncast f32 add —
+//   the cast form is what the adapters were trained behind.
+
+import { existsSync, readdirSync, statSync } from "node:fs";
+import { resolve } from "node:path";
+import { ptr, read } from "bun:ffi";
+import { MlxArray, cpuStream } from "./mlx/array";
+import { C } from "./mlx/ffi";
+import { SafetensorsFile } from "./safetensors";
+import type { LoraWeights } from "./model/gemma4";
+import type { RuntimeModel } from "./model/factory";
+
+const cstr = (s: string) => Buffer.from(s + "\0", "utf8");
+
+export interface AdapterInfo {
+  id: string;
+  path: string;
+  rank: number | null;
+  scale: number;
+  sizeBytes: number;
+  mountedLayers: number;
+  /** Adapter tensors for modules outside the 7 target suffixes
+   *  (per-layer-input projections etc.) — skipped, like the reference. */
+  skippedTensors: number;
+}
+
+/** Read every tensor from an adapter safetensors file (small, f32) into
+ *  materialized mlx arrays. Tensor names come from our header parser;
+ *  arrays from mlx's native loader via the proven map-get pattern
+ *  (weights.ts). The map and header mmap are freed before returning. */
+function loadAdapterTensors(file: string): Map<string, MlxArray> {
+  const sf = SafetensorsFile.open(file);
+  const names = [...sf.tensors.keys()];
+  sf.mmap.unmap();
+
+  const arrMapSlot = new BigUint64Array([C.mlx_map_string_to_array_new()]);
+  const metaMapSlot = new BigUint64Array([C.mlx_map_string_to_string_new()]);
+  const arrMapPtr = ptr(arrMapSlot);
+  const metaMapPtr = ptr(metaMapSlot);
+  const status = C.mlx_load_safetensors(arrMapPtr, metaMapPtr, ptr(cstr(file)), cpuStream);
+  C.mlx_map_string_to_string_free(read.u64(metaMapPtr, 0));
+  const mapHandle = read.u64(arrMapPtr, 0);
+  if (status !== 0) {
+    C.mlx_map_string_to_array_free(mapHandle);
+    throw new Error(`mlx_load_safetensors(${file}) failed`);
+  }
+  const out = new Map<string, MlxArray>();
+  try {
+    for (const name of names) {
+      const slot = new BigUint64Array([C.mlx_array_new()]);
+      const slotPtr = ptr(slot);
+      if (C.mlx_map_string_to_array_get(slotPtr, mapHandle, ptr(cstr(name))) !== 0)
+        throw new Error(`adapter tensor ${name} missing from native map`);
+      const arr = new MlxArray(read.u64(slotPtr, 0));
+      arr.eval(); // materialize — the map (and its Load refs) is freed below
+      out.set(name, arr);
+    }
+  } catch (e) {
+    for (const a of out.values()) a.dispose();
+    C.mlx_map_string_to_array_free(mapHandle);
+    throw e;
+  }
+  C.mlx_map_string_to_array_free(mapHandle);
+  return out;
+}
+
+/** Adapter scale from its config: mlx-lm writes lora_parameters.scale;
+ *  PEFT writes lora_alpha + r (scale = alpha / r). Reference: mount.py. */
+async function readAdapterScale(dir: string): Promise<{ scale: number; rank: number | null }> {
+  for (const name of ["optiq_lora_config.json", "adapter_config.json"]) {
+    const f = Bun.file(`${dir}/${name}`);
+    if (await f.exists()) {
+      const cfg = (await f.json()) as Record<string, any>;
+      const lp = cfg.lora_parameters;
+      if (lp && typeof lp === "object")
+        return { scale: Number(lp.scale ?? 20.0), rank: lp.rank ?? null };
+      const alpha = Number(cfg.lora_alpha ?? 16);
+      const r = Number(cfg.r ?? 8);
+      return { scale: r ? alpha / r : 1.0, rank: r || null };
+    }
+  }
+  throw new Error(`no adapter_config.json in ${dir}`);
+}
+
+function adapterWeightsFile(dir: string): string {
+  for (const name of ["adapters.safetensors", "adapter_model.safetensors"]) {
+    const p = `${dir}/${name}`;
+    if (existsSync(p)) return p;
+  }
+  throw new Error(`no adapter weights at ${dir}/adapters.safetensors or adapter_model.safetensors`);
+}
+
+/** A mountable adapter found on disk (not yet mounted). */
+export interface AvailableAdapter {
+  id: string; // directory basename — the handle to mount it under
+  path: string; // absolute adapter dir
+  scale: number;
+  rank: number | null;
+  /** Base model repo id (org/name) the adapter was trained on, or null if the
+   *  config doesn't record it. Lets the chat selector hide adapters that don't
+   *  fit the served model (a MiniCPM5 adapter can't mount on Gemma). */
+  baseModel: string | null;
+}
+
+/** The base-model repo id an adapter was trained on, from its config's
+ *  source_model / base_model_name_or_path. These are stored as machine-local
+ *  snapshot PATHS (…/models--ORG--NAME/snapshots/HASH), so recover the stable
+ *  repo id (ORG/NAME) from the path; returns null if unrecorded. */
+async function readAdapterBaseRepoId(dir: string): Promise<string | null> {
+  for (const [name, key] of [
+    ["optiq_lora_config.json", "source_model"],
+    ["adapter_config.json", "base_model_name_or_path"],
+  ] as const) {
+    const f = Bun.file(`${dir}/${name}`);
+    if (await f.exists()) {
+      const cfg = (await f.json()) as Record<string, any>;
+      const src = cfg[key];
+      if (typeof src === "string" && src) {
+        const m = src.match(/models--([^/]+)/);
+        return m ? m[1]!.replace(/--/g, "/") : src;
+      }
+    }
+  }
+  return null;
+}
+
+/** Scan adapter stores for mountable adapters: directories holding an adapter
+ *  weights file. id = dir basename; scale/rank/baseModel read from each adapter's
+ *  config. Dirs without weights (dataset folders) and unreadable stores are skipped.
+ *  Backs GET /v1/adapters/available, which populates the chat adapter selector. */
+export async function listAvailableAdapters(stores: string[]): Promise<AvailableAdapter[]> {
+  const out: AvailableAdapter[] = [];
+  const seen = new Set<string>();
+  for (const store of stores) {
+    let entries;
+    try {
+      entries = readdirSync(store, { withFileTypes: true });
+    } catch {
+      continue; // store missing → skip
+    }
+    for (const e of entries) {
+      if (!e.isDirectory()) continue;
+      const dir = resolve(store, e.name);
+      try {
+        adapterWeightsFile(dir);
+      } catch {
+        continue; // no weights → not an adapter
+      }
+      if (seen.has(dir)) continue;
+      seen.add(dir);
+      let scale = 1;
+      let rank: number | null = null;
+      try {
+        ({ scale, rank } = await readAdapterScale(dir));
+      } catch {
+        // keep defaults if the config is unreadable
+      }
+      const baseModel = await readAdapterBaseRepoId(dir).catch(() => null);
+      out.push({ id: e.name, path: dir, scale, rank, baseModel });
+    }
+  }
+  return out;
+}
+
+/** Parse an adapter selection spec: "sft" | "sft+dpo" | "sft,dpo"
+ *  (reference mount.py stacking syntax — residuals sum in order). */
+export function parseAdapterSpec(spec: string): string[] {
+  const sep = spec.includes("+") ? "+" : ",";
+  return spec.split(sep).map((s) => s.trim()).filter(Boolean);
+}
+
+export class AdapterManager {
+  readonly #model: RuntimeModel;
+  readonly #mounted = new Map<string, AdapterInfo>();
+  /** Owned adapter arrays per id (disposed on unmount). */
+  readonly #arrays = new Map<string, MlxArray[]>();
+
+  constructor(model: RuntimeModel) {
+    this.#model = model;
+  }
+
+  /** Mount an adapter directory under `id`. Validates layer-name and
+   *  rank/shape compatibility against the base BEFORE any weights are
+   *  attached — a bad adapter fails here, never at request time
+   *  (reference: resolver.py + mount.py add_adapter checks). */
+  async mount(id: string, dir: string): Promise<AdapterInfo> {
+    if (this.#mounted.has(id)) return this.#mounted.get(id)!;
+    if (id.includes("+") || id.includes(","))
+      throw new Error(`adapter id ${JSON.stringify(id)} may not contain '+' or ','`);
+    const path = resolve(dir);
+    if (!existsSync(path)) throw new Error(`adapter dir not found: ${path}`);
+
+    const weightsFile = adapterWeightsFile(path);
+    const { scale, rank: configRank } = await readAdapterScale(path);
+    const tensors = loadAdapterTensors(weightsFile);
+    const targets = this.#model.loraTargets();
+
+    // Group adapter tensors into (modulePath → {a, b}), probing both the
+    // pure-LLM and VLM-wrapped prefixes like the reference _find_weight_pair.
+    const pairs = new Map<string, { a?: MlxArray; b?: MlxArray }>();
+    let skipped = 0;
+    const ourPrefix = `${this.#model.prefixBase}.layers.`;
+    const altPrefix = ourPrefix.startsWith("language_model.")
+      ? ourPrefix.slice("language_model.".length)
+      : `language_model.${ourPrefix}`;
+    for (const [name, arr] of tensors) {
+      const m = name.match(/^(.*)\.(lora_a|lora_A|lora_b|lora_B)(\.weight)?$/);
+      if (!m) { skipped++; arr.dispose(); continue; }
+      let modulePath = m[1]!;
+      if (modulePath.startsWith(altPrefix))
+        modulePath = ourPrefix + modulePath.slice(altPrefix.length);
+      if (!targets.has(modulePath)) { skipped++; arr.dispose(); continue; }
+      const slot = pairs.get(modulePath) ?? {};
+      if (m[2]!.toLowerCase() === "lora_a") slot.a = arr; else slot.b = arr;
+      pairs.set(modulePath, slot);
+    }
+
+    // Validate every pair against the base linear's dims before mounting
+    // anything (all-or-nothing: a bad adapter must not half-mount).
+    const validated: { linear: import("./model/gemma4").QuantizedLinear; lw: LoraWeights }[] = [];
+    const dispose = () => {
+      for (const { a, b } of pairs.values()) { a?.dispose(); b?.dispose(); }
+    };
+    try {
+      for (const [modulePath, { a, b }] of pairs) {
+        if (!a || !b)
+          throw new Error(`${modulePath}: adapter has only one of lora_a/lora_b`);
+        const linear = targets.get(modulePath)!;
+        const [aIn, rank] = a.shape as [number, number];
+        const [bRank, bOut] = b.shape as [number, number];
+        if (aIn !== linear.inFeatures || bRank !== rank || bOut !== linear.outFeatures)
+          throw new Error(
+            `${modulePath}: shape mismatch — lora_a [${a.shape}] / lora_b [${b.shape}] ` +
+            `vs base [in ${linear.inFeatures}, out ${linear.outFeatures}]; ` +
+            `was this adapter trained for a different base model?`,
+          );
+        validated.push({ linear, lw: { a, b, scale, rank } });
+      }
+      if (validated.length === 0)
+        throw new Error(
+          `failed to mount adapter ${JSON.stringify(id)}: no tensors match the ` +
+          `target modules (q/k/v/o/gate/up/down_proj). Check that the adapter ` +
+          `was trained for this base model.`,
+        );
+    } catch (e) {
+      dispose();
+      throw e;
+    }
+
+    for (const { linear, lw } of validated) {
+      (linear.adapters ??= new Map()).set(id, lw);
+      linear.loraState = this.#model.loraState;
+    }
+
+    const info: AdapterInfo = {
+      id,
+      path,
+      rank: configRank,
+      scale,
+      sizeBytes: statSync(weightsFile).size,
+      mountedLayers: validated.length,
+      skippedTensors: skipped,
+    };
+    this.#mounted.set(id, info);
+    this.#arrays.set(id, validated.flatMap(({ lw }) => [lw.a, lw.b]));
+    return info;
+  }
+
+  /** Remove `id` from every linear (frees its arrays). The mount points
+   *  stay in place for other adapters, like the reference. */
+  unmount(id: string): number {
+    if (!this.#mounted.delete(id)) return 0;
+    let removed = 0;
+    for (const linear of this.#model.loraTargets().values()) {
+      if (linear.adapters?.delete(id)) removed++;
+    }
+    for (const a of this.#arrays.get(id) ?? []) a.dispose();
+    this.#arrays.delete(id);
+    if (this.#model.loraState.active.includes(id))
+      this.#model.loraState.active = this.#model.loraState.active.filter((x) => x !== id);
+    return removed;
+  }
+
+  list(): AdapterInfo[] {
+    return [...this.#mounted.values()];
+  }
+
+  get(id: string): AdapterInfo | undefined {
+    return this.#mounted.get(id);
+  }
+
+  /** Resolve a request's adapter spec to validated ids ([] for none).
+   *  Unknown ids are an error — selection must fail loudly, not no-op. */
+  resolveSpec(spec: string | null | undefined): string[] {
+    if (!spec || spec.toLowerCase() === "none") return [];
+    const ids = parseAdapterSpec(spec);
+    for (const id of ids)
+      if (!this.#mounted.has(id))
+        throw new Error(
+          `unknown adapter ${JSON.stringify(id)} — mounted: ` +
+          `${[...this.#mounted.keys()].join(", ") || "(none)"}`,
+        );
+    return ids;
+  }
+}

package/src/md-modules.d.ts

@@ -0,0 +1,7 @@
+// Ambient type for `import md from "./foo.md" with { type: "text" }`.
+// Bun reads the file as a string, but bun-types only declares .txt/.html,
+// not .md — so we declare it here (used by src/web/skills.ts).
+declare module "*.md" {
+  const content: string;
+  export default content;
+}