mlx-bun 0.0.8 → 0.0.10

package/src/prompt-cache.ts

@@ -17,6 +17,11 @@ export interface PromptCacheEntry {
   /** Namespace key — adapter spec for LoRA requests ("" = base model).
    *  KV computed under one adapter must never seed another's prefill. */
   ns: string;
+  /** Release hook for entries restored from the SSD cold tier: unmaps the
+   *  backing COW file mapping. MUST run only after `caches` are disposed
+   *  (they may alias file pages); every dispose path here honors that, and
+   *  callers who take() an entry carry the thunk with it. */
+  retain?: () => void;
 }
 
 function cacheBytes(caches: Cache[]): number {
@@ -26,7 +31,7 @@ function cacheBytes(caches: Cache[]): number {
 }
 
 /** Length of the longest common prefix of two token sequences. */
-function commonPrefixLength(a: number[], b: number[]): number {
+export function commonPrefixLength(a: number[], b: number[]): number {
   const n = Math.min(a.length, b.length);
   let i = 0;
   while (i < n && a[i] === b[i]) i++;
@@ -39,9 +44,22 @@ export class PromptCache {
   #clock = 0;
   hits = 0;
   misses = 0;
+  /** Cold-tier spill hook: called with an entry ABOUT to be evicted (still
+   *  alive — the hook may serialize its caches); the cache disposes it after.
+   *  Failures are the hook's problem; eviction proceeds regardless. */
+  readonly #spill: ((entry: PromptCacheEntry) => void) | null;
 
-  constructor(maxBytes: number) {
+  constructor(maxBytes: number, spill: ((entry: PromptCacheEntry) => void) | null = null) {
     this.maxBytes = maxBytes;
+    this.#spill = spill;
+  }
+
+  #disposeEntry(entry: PromptCacheEntry, spillFirst: boolean): void {
+    if (spillFirst && this.#spill) {
+      try { this.#spill(entry); } catch {}
+    }
+    for (const c of entry.caches) c.dispose();
+    entry.retain?.();
   }
 
   get totalBytes(): number {
@@ -88,26 +106,54 @@ export class PromptCache {
     return entry;
   }
 
-  /** Insert (or reinsert) an entry; evicts LRU entries over the byte cap.
-   *  If the entry itself exceeds the cap it is disposed, not stored. */
-  put(tokens: number[], caches: Cache[], ns = ""): void {
+  /** Longest usable prefix length available for `prompt` WITHOUT taking
+   *  (the SSD tier compares this against its own best before restoring). */
+  peekPrefixLen(prompt: number[], ns = ""): number {
+    let best = 0;
+    for (const { entry: e } of this.#entries) {
+      if (e.ns !== ns) continue;
+      const p = Math.min(commonPrefixLength(e.tokens, prompt), prompt.length - 1);
+      if (p <= best) continue;
+      const trimNeeded = e.tokens.length - p; // same usability rule as take()
+      if (trimNeeded > 0 && !e.caches.every((c) => c.isTrimmable())) continue;
+      best = p;
+    }
+    return best;
+  }
+
+  /** Read-only exact-token lookup (the write-behind snapshot path: the
+   *  entry stays owned by the cache; the caller only reads array state,
+   *  under the gateway lock so no generation is mutating it). */
+  findExact(tokens: number[], ns = ""): PromptCacheEntry | null {
+    for (const { entry: e } of this.#entries) {
+      if (e.ns !== ns || e.tokens.length !== tokens.length) continue;
+      if (commonPrefixLength(e.tokens, tokens) === tokens.length) return e;
+    }
+    return null;
+  }
+
+  /** Insert (or reinsert) an entry; evicts LRU entries over the byte cap
+   *  (spilling each to the cold tier first, when one is attached). If the
+   *  entry itself exceeds the cap it is spilled + disposed, not stored. */
+  put(tokens: number[], caches: Cache[], ns = "", retain?: () => void): void {
     const bytes = cacheBytes(caches);
+    const entry: PromptCacheEntry = { tokens, caches, ns, retain };
     if (bytes > this.maxBytes) {
-      for (const c of caches) c.dispose();
+      this.#disposeEntry(entry, true);
       return;
     }
-    this.#entries.push({ entry: { tokens, caches, ns }, bytes, lastUsed: ++this.#clock });
+    this.#entries.push({ entry, bytes, lastUsed: ++this.#clock });
     while (this.totalBytes > this.maxBytes && this.#entries.length > 1) {
       let lruIdx = 0;
       for (let i = 1; i < this.#entries.length; i++)
         if (this.#entries[i]!.lastUsed < this.#entries[lruIdx]!.lastUsed) lruIdx = i;
       const [evicted] = this.#entries.splice(lruIdx, 1);
-      for (const c of evicted!.entry.caches) c.dispose();
+      this.#disposeEntry(evicted!.entry, true);
     }
   }
 
   clear(): void {
-    for (const e of this.#entries) for (const c of e.entry.caches) c.dispose();
+    for (const e of this.#entries) this.#disposeEntry(e.entry, false);
     this.#entries = [];
   }
 }

package/src/quantize/config-writer.ts

@@ -109,7 +109,7 @@ export async function writeQuantizedConfig(
 }
 
 /** Copy known aux files (and any `*.model`) from src → out when present. */
-async function copyAuxFiles(srcDir: string, outDir: string): Promise<void> {
+export async function copyAuxFiles(srcDir: string, outDir: string): Promise<void> {
   const seen = new Set<string>();
   for (const f of AUX_FILES) {
     const src = join(srcDir, f);

package/src/registry.ts

@@ -6,7 +6,8 @@
 
 import { Database } from "bun:sqlite";
 import {
-  closeSync, existsSync, openSync, readdirSync, readFileSync, readSync, statSync,
+  closeSync, existsSync, lstatSync, openSync, readdirSync, readFileSync,
+  readlinkSync, readSync, rmSync, statSync,
 } from "node:fs";
 import { join } from "node:path";
 import { isDrafterModelType } from "./model/support";
@@ -27,6 +28,12 @@ export interface ModelRecord {
   quantGroupSize: number | null;
   quantMode: string | null;
   hasVisionSidecar: boolean;
+  /** config.json's vision_config.model_type when it names a vision tower
+   *  (e.g. "gemma4_vision" — SigLIP, needs the bf16 sidecar — or
+   *  "gemma4_unified_vision" — encoder-free). null for text-only models
+   *  AND for configs whose nested vision_config is not a vision tower
+   *  (Qwen3.5 nests a copy of its own text config there). */
+  visionConfigType: string | null;
   hasKvConfig: boolean;
   hasToolTemplate: boolean;
   numLayers: number | null;
@@ -52,6 +59,7 @@ CREATE TABLE IF NOT EXISTS models (
   quant_group_size INTEGER,
   quant_mode TEXT,
   has_vision_sidecar INTEGER NOT NULL,
+  vision_config_type TEXT,
   has_kv_config INTEGER NOT NULL,
   has_tool_template INTEGER NOT NULL,
   num_layers INTEGER,
@@ -62,7 +70,14 @@ CREATE TABLE IF NOT EXISTS models (
 );
 `;
 
-export const DEFAULT_HUB = `${process.env.HOME}/.cache/huggingface/hub`;
+/** HF hub cache root, honoring the standard env overrides the same way
+ *  huggingface_hub (and our server.ts) do: HF_HUB_CACHE > HF_HOME/hub >
+ *  ~/.cache/huggingface/hub. */
+export const DEFAULT_HUB =
+  process.env.HF_HUB_CACHE ??
+  (process.env.HF_HOME
+    ? join(process.env.HF_HOME, "hub")
+    : `${process.env.HOME}/.cache/huggingface/hub`);
 export const DEFAULT_DB = `${process.env.HOME}/.cache/mlx-bun/registry.sqlite`;
 
 export class Registry {
@@ -81,7 +96,8 @@ export class Registry {
     if (
       !cols.includes("sidecar_bytes") ||
       !cols.includes("experts_bytes") ||
-      !cols.includes("license")
+      !cols.includes("license") ||
+      !cols.includes("vision_config_type")
     ) {
       this.db.exec("DROP TABLE models");
       this.db.exec(SCHEMA);
@@ -99,7 +115,7 @@ export class Registry {
     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)
+       $vision, $vtype, $kv, $tools, $layers, $hidden, $vocab, $license, $at)
     `);
     for (const entry of readdirSync(hubDir)) {
       if (!entry.startsWith("models--")) continue;
@@ -116,6 +132,7 @@ export class Registry {
           $sidecar: rec.sidecarBytes, $experts: rec.expertsBytes,
           $bits: rec.quantBits, $gs: rec.quantGroupSize, $mode: rec.quantMode,
           $vision: rec.hasVisionSidecar ? 1 : 0,
+          $vtype: rec.visionConfigType,
           $kv: rec.hasKvConfig ? 1 : 0,
           $tools: rec.hasToolTemplate ? 1 : 0,
           $layers: rec.numLayers, $hidden: rec.hiddenSize, $vocab: rec.vocabSize,
@@ -131,8 +148,12 @@ export class Registry {
     const clauses: string[] = [];
     const params: Record<string, unknown> = {};
     if (filter.vision !== undefined) {
-      clauses.push("has_vision_sidecar = $vision");
-      params.$vision = filter.vision ? 1 : 0;
+      // Vision-capable = SigLIP sidecar present OR the encoder-free unified
+      // tower declared in config.json (12B ships no sidecar upstream anymore).
+      // COALESCE: NULL vision_config_type must read as "no", not SQL-NULL
+      // (a bare NULL LIKE poisons the NOT branch of the filter).
+      const cap = "(has_vision_sidecar = 1 OR COALESCE(vision_config_type,'') LIKE '%unified_vision')";
+      clauses.push(filter.vision ? cap : `NOT ${cap}`);
     }
     if (filter.maxBytes !== undefined) {
       clauses.push("size_bytes <= $max");
@@ -152,6 +173,24 @@ export class Registry {
     return rows.map(rowToRecord).filter((r) => existsSync(r.path));
   }
 
+  /** list(), collapsed to ONE record per repo: the canonical snapshot
+   *  (refs/main, else most recently scanned). Upstream pushes create a new
+   *  snapshots/<commit> dir per revision and the old ones are never deleted,
+   *  so raw list() shows a row per revision — this is the user-facing view
+   *  (`ls`, /v1/models, /library). Order follows list() (size ascending). */
+  listCanonical(filter: Parameters<Registry["list"]>[0] = {}): ModelRecord[] {
+    const byRepo = new Map<string, ModelRecord[]>();
+    const order: string[] = [];
+    for (const m of this.list(filter)) {
+      if (!byRepo.has(m.repoId)) {
+        byRepo.set(m.repoId, []);
+        order.push(m.repoId);
+      }
+      byRepo.get(m.repoId)!.push(m);
+    }
+    return order.map((repo) => pickCanonicalRevision(byRepo.get(repo)!));
+  }
+
   /** 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. */
@@ -172,10 +211,17 @@ export class Registry {
   }
 }
 
+/** True when the model can answer vision requests: it either ships the bf16
+ *  SigLIP sidecar (gemma4_vision: e2b/e4b/26B/31B) or declares the
+ *  encoder-free unified tower in config.json (gemma4_unified_vision: 12B). */
+export function visionCapable(m: Pick<ModelRecord, "hasVisionSidecar" | "visionConfigType">): boolean {
+  return m.hasVisionSidecar || (m.visionConfigType?.endsWith("unified_vision") ?? false);
+}
+
 /** 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 {
+export function pickCanonicalRevision(matches: ModelRecord[]): ModelRecord {
   if (matches.length === 1) return matches[0]!;
   const root = matches[0]!.path.split("/snapshots/")[0]!;
   try {
@@ -203,6 +249,7 @@ function rowToRecord(r: Record<string, unknown>): ModelRecord {
     quantGroupSize: r.quant_group_size as number | null,
     quantMode: r.quant_mode as string | null,
     hasVisionSidecar: !!r.has_vision_sidecar,
+    visionConfigType: (r.vision_config_type as string | null) ?? null,
     hasKvConfig: !!r.has_kv_config,
     hasToolTemplate: !!r.has_tool_template,
     numLayers: r.num_layers as number | null,
@@ -258,7 +305,170 @@ function expertTensorBytes(path: string): number {
   }
 }
 
-async function scanSnapshot(dir: string, repoId: string): Promise<ModelRecord | null> {
+// ---------------------------------------------------------------------------
+// gc: reclaim superseded snapshots + dead blobs.
+//
+// The HF cache layout keeps one snapshots/<commit> dir per downloaded
+// revision and NEVER deletes old ones — every upstream push a `get` follows
+// strands the previous snapshot (and any blobs only it references). gc keeps
+// the snapshots refs/* point at, deletes the rest, then deletes blobs no
+// surviving snapshot symlink targets.
+
+export interface GcSkippedSnapshot {
+  path: string;
+  /** Files this snapshot has that NO kept snapshot has (relative names).
+   *  Deleting it would delete the only copy — needs an explicit --force. */
+  extraFiles: string[];
+}
+
+export interface GcRepoPlan {
+  repoId: string;
+  repoDir: string;
+  /** Snapshots referenced by refs/* — never deleted. */
+  keepSnapshots: string[];
+  /** Unreferenced snapshots safe to delete (file set ⊆ kept snapshots'). */
+  pruneSnapshots: string[];
+  /** Unreferenced snapshots with files the kept set lacks — skipped unless
+   *  force (their blobs count as live). */
+  skippedSnapshots: GcSkippedSnapshot[];
+  /** blobs/ entries no surviving snapshot links to. */
+  deadBlobs: string[];
+  reclaimBytes: number;
+}
+
+/** Relative paths of all files/symlinks under a snapshot dir. */
+function snapshotFiles(snapDir: string): string[] {
+  const out: string[] = [];
+  const walk = (dir: string, prefix: string) => {
+    for (const name of readdirSync(dir)) {
+      const p = join(dir, name);
+      const st = lstatSync(p);
+      if (st.isDirectory()) walk(p, `${prefix}${name}/`);
+      else out.push(`${prefix}${name}`);
+    }
+  };
+  try { walk(snapDir, ""); } catch {}
+  return out;
+}
+
+/** Blob basenames a snapshot's symlinks resolve to (within this repo). */
+function snapshotBlobTargets(snapDir: string): Set<string> {
+  const targets = new Set<string>();
+  for (const rel of snapshotFiles(snapDir)) {
+    const p = join(snapDir, rel);
+    try {
+      const t = readlinkSync(p); // throws for regular files — those aren't blobs
+      const base = t.split("/").pop();
+      if (base && t.includes("blobs")) targets.add(base);
+    } catch {}
+  }
+  return targets;
+}
+
+/** Plan gc for one models--* repo dir. Conservative by construction:
+ *  - no refs/ (or refs pointing at missing snapshots) → nothing is pruned;
+ *  - an unreferenced snapshot carrying files the kept set lacks is only
+ *    prunable under `force` (the live example: a stale snapshot whose
+ *    optiq_vision.safetensors the canonical revision dropped);
+ *  - blobs are dead only when NO surviving snapshot (kept or skipped)
+ *    links to them. `.incomplete`/`.lock` resume artifacts are never touched. */
+export function planRepoGc(repoDir: string, opts: { force?: boolean } = {}): GcRepoPlan {
+  const entry = repoDir.split("/").pop() ?? "";
+  const repoId = entry.slice("models--".length).replaceAll("--", "/");
+  const plan: GcRepoPlan = {
+    repoId, repoDir,
+    keepSnapshots: [], pruneSnapshots: [], skippedSnapshots: [],
+    deadBlobs: [], reclaimBytes: 0,
+  };
+  const snapsDir = join(repoDir, "snapshots");
+  const refsDir = join(repoDir, "refs");
+  if (!existsSync(snapsDir)) return plan;
+
+  const refCommits = new Set<string>();
+  if (existsSync(refsDir)) {
+    for (const r of readdirSync(refsDir)) {
+      try { refCommits.add(readFileSync(join(refsDir, r), "utf8").trim()); } catch {}
+    }
+  }
+
+  const candidates: string[] = [];
+  for (const snap of readdirSync(snapsDir)) {
+    const p = join(snapsDir, snap);
+    if (!lstatSync(p).isDirectory()) continue;
+    (refCommits.has(snap) ? plan.keepSnapshots : candidates).push(p);
+  }
+  // No ref resolves to an existing snapshot → we cannot tell what's live.
+  if (plan.keepSnapshots.length === 0) return plan;
+
+  const keptFiles = new Set<string>();
+  for (const k of plan.keepSnapshots) for (const f of snapshotFiles(k)) keptFiles.add(f);
+  for (const c of candidates) {
+    const extra = snapshotFiles(c).filter((f) => !keptFiles.has(f));
+    if (extra.length > 0 && !opts.force) plan.skippedSnapshots.push({ path: c, extraFiles: extra });
+    else plan.pruneSnapshots.push(c);
+  }
+
+  const live = new Set<string>();
+  for (const s of [...plan.keepSnapshots, ...plan.skippedSnapshots.map((x) => x.path)])
+    for (const b of snapshotBlobTargets(s)) live.add(b);
+  const blobsDir = join(repoDir, "blobs");
+  if (existsSync(blobsDir)) {
+    for (const b of readdirSync(blobsDir)) {
+      if (b.endsWith(".incomplete") || b.endsWith(".lock")) continue;
+      if (live.has(b)) continue;
+      const p = join(blobsDir, b);
+      plan.deadBlobs.push(p);
+      try { plan.reclaimBytes += lstatSync(p).size; } catch {}
+    }
+  }
+  return plan;
+}
+
+export function planGc(hubDir: string = DEFAULT_HUB, opts: { force?: boolean } = {}): GcRepoPlan[] {
+  if (!existsSync(hubDir)) return [];
+  const plans: GcRepoPlan[] = [];
+  for (const entry of readdirSync(hubDir)) {
+    if (!entry.startsWith("models--")) continue;
+    plans.push(planRepoGc(join(hubDir, entry), opts));
+  }
+  return plans;
+}
+
+/** Execute a plan: snapshot dirs first, then dead blobs. Returns totals. */
+export function executeGc(plans: GcRepoPlan[]): {
+  snapshots: number; blobs: number; reclaimedBytes: number;
+} {
+  let snapshots = 0, blobs = 0, reclaimedBytes = 0;
+  for (const p of plans) {
+    for (const s of p.pruneSnapshots) {
+      rmSync(s, { recursive: true, force: true });
+      snapshots++;
+    }
+    for (const b of p.deadBlobs) {
+      try {
+        const size = lstatSync(b).size;
+        rmSync(b, { force: true });
+        blobs++;
+        reclaimedBytes += size;
+      } catch {}
+    }
+  }
+  return { snapshots, blobs, reclaimedBytes };
+}
+
+/** vision_config.model_type, but only when it actually names a vision tower
+ *  (`*_vision`). Some text-only configs nest a serialized copy of themselves
+ *  under vision_config (Qwen3.5 puts model_type "qwen3_5" there) — presence
+ *  of the key alone is NOT a vision signal. */
+function visionConfigTypeOf(config: Record<string, any>): string | null {
+  const vt = (config.vision_config as Record<string, any> | undefined)?.model_type;
+  return typeof vt === "string" && vt.endsWith("_vision") ? vt : null;
+}
+
+/** Build a ModelRecord straight from a snapshot directory (also the scan()
+ *  workhorse). Exported for explicit-path loads (`serve --model <dir>`),
+ *  which must not depend on the registry DB having scanned that dir. */
+export async function scanSnapshot(dir: string, repoId: string): Promise<ModelRecord | null> {
   const configPath = join(dir, "config.json");
   if (!existsSync(configPath)) return null;
 
@@ -325,6 +535,7 @@ async function scanSnapshot(dir: string, repoId: string): Promise<ModelRecord |
     quantGroupSize: quant?.group_size ?? null,
     quantMode: quant?.mode ?? (quant ? "affine" : null),
     hasVisionSidecar: existsSync(join(dir, "optiq_vision.safetensors")),
+    visionConfigType: visionConfigTypeOf(config),
     hasKvConfig: existsSync(join(dir, "kv_config.json")),
     hasToolTemplate,
     numLayers: text.num_hidden_layers ?? null,

package/src/responses.ts

@@ -121,6 +121,19 @@ export interface ResponsesRequest {
   top_p?: number;
   top_k?: number;
   stream?: boolean;
+  /** mlx-lm sampler/penalty extensions + logit_bias — not part of the
+   *  Responses protocol proper; accepted as pass-through extras with the
+   *  same names as our /v1 chat surface. */
+  min_p?: number;
+  xtc_probability?: number;
+  xtc_threshold?: number;
+  logit_bias?: Record<string, number>;
+  repetition_penalty?: number;
+  repetition_context_size?: number;
+  presence_penalty?: number;
+  presence_context_size?: number;
+  frequency_penalty?: number;
+  frequency_context_size?: number;
   tools?: Array<Record<string, unknown>>;
   tool_choice?: unknown;
 }
@@ -235,6 +248,16 @@ export function responsesToChatBody(body: ResponsesRequest): Record<string, unkn
   if (body.top_p != null) oai.top_p = body.top_p;
   if (body.top_k != null) oai.top_k = body.top_k;
   if (body.stream) oai.stream = true;
+  // mlx-lm sampler/penalty extensions (+ OpenAI logit_bias): straight
+  // pass-through — same wire names on the chat surface.
+  for (const k of [
+    "min_p", "xtc_probability", "xtc_threshold", "logit_bias",
+    "repetition_penalty", "repetition_context_size",
+    "presence_penalty", "presence_context_size",
+    "frequency_penalty", "frequency_context_size",
+  ] as const) {
+    if (body[k] != null) oai[k] = body[k];
+  }
 
   // Flat Responses tool shape → nested chat shape; built-ins
   // (web_search, file_search, mcp, computer_use) dropped silently.