mlx-bun 0.0.9 → 0.0.11

package/src/spec/two-model.ts

@@ -0,0 +1,171 @@
+// Two-model speculative drafting — mlx-lm parity (`mlx_lm.server
+// --draft-model`). A full second model drafts autoregressively; the serve
+// loop (src/spec/serve-loop.ts) verifies with the target. L1 oracle:
+// mlx_lm.server with the same target/draft pair, greedy, token-for-token
+// (spec-vs-spec — both batch the verify lm-head, so neither is bit-exact to
+// stock decode at knife-edges; see src/spec/generate.ts header).
+//
+// Faithfulness notes (read from the oracle venv's generate.py):
+//  - drafts are sampled with the REQUEST sampler (generate.py:593-601) —
+//    greedy drafting under a temperature>0 request is not parity;
+//  - draft-cache rewind is max(n - kAccept - 1, 0) (generate.py:589-591),
+//    with the all-accept re-feed of the last draft handled by the serve
+//    loop's `feed` (generate.py:645-648);
+//  - logits processors do NOT run on draft steps (they're target-side).
+
+import { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { clearCache } from "../mlx/ffi";
+import { flagOn } from "../flags";
+import { loadModelConfig } from "../config";
+import { Weights } from "../weights";
+import { createModel, type RuntimeModel } from "../model/factory";
+import type { Cache } from "../model/gemma4";
+import { toLogprobs } from "../sampler";
+import type { DraftProvider, DraftSource } from "./source";
+
+/** mlx-lm's prefill_step_size — the draft drain chunks at the same stride
+ *  as the target's (serve-loop.ts); chunking is numerically exact for
+ *  causal attention. */
+const PREFILL_CHUNK = 2048;
+
+export class TwoModelProvider implements DraftProvider {
+  readonly id: string;
+  readonly weightsBytes: number;
+
+  private constructor(
+    readonly model: RuntimeModel,
+    private readonly weights: Weights,
+    id: string,
+  ) {
+    this.id = id;
+    this.weightsBytes = [...weights.shards.files.values()].reduce(
+      (a, f) => a + f.mmap.size,
+      0,
+    );
+  }
+
+  /** Load the draft model from a snapshot dir. Vocab-size mismatch with the
+   *  target is a WARNING (mlx-lm parity, server.py:363-368); the tokenizer-
+   *  family hard check lives at the server (it owns both tokenizers). */
+  static async load(modelDir: string, targetVocabSize?: number): Promise<TwoModelProvider> {
+    const config = await loadModelConfig(modelDir);
+    const weights = await Weights.open(modelDir);
+    const model = createModel(weights, config);
+    if (
+      targetVocabSize !== undefined &&
+      config.text.vocabSize !== targetVocabSize
+    )
+      console.warn(
+        `draft model vocab size ${config.text.vocabSize} != target ${targetVocabSize} — ` +
+          `speculation will accept nothing useful (mlx-lm warns identically)`,
+      );
+    return new TwoModelProvider(model, weights, modelDir.split("/").filter(Boolean).at(-1)!);
+  }
+
+  open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
+    return new TwoModelSource(this.model, opts.sampler);
+  }
+
+  dispose(): void {
+    this.weights.dispose();
+  }
+}
+
+class TwoModelSource implements DraftSource {
+  readonly weightsBytes = 0; // provider-owned weights; per-request adds caches only
+  private caches: Cache[];
+
+  constructor(
+    private readonly model: RuntimeModel,
+    private readonly sampler: (lp: MlxArray, step: number) => MlxArray,
+  ) {
+    this.caches = this.model.makeCache();
+  }
+
+  prefill(promptIds: number[]): void {
+    // Oracle convention (mlx-lm speculative_generate_step._prefill, re-anchored
+    // 2026-07-07): the DRAFT model also drains only to len-1 (`while y.size >
+    // 1`) — the last prompt token stays unprocessed and is consumed by the
+    // FIRST draft() round's feed ([pending] = that same token under the serve
+    // loop's oracle shape), mirroring _draft_generate's first _step. Draft-side
+    // full-prefill was the residual knife-edge flipper in the 2026-07-07 live
+    // oracle gate (γ=2 haiku cell). MLX_BUN_PREFILL_TAIL_SPLIT=0 reverts.
+    // (mlx-lm chunks its drain at prefill_step_size=2048; chunking is
+    // numerically exact for causal attention — same KV, same positions —
+    // so we chunk too. Pre-2026-07-07 this was a SINGLE forward of the
+    // whole head: nothing enforced the "<2048 serve regime" the old
+    // comment assumed, and a 32k prompt through the spec lane ran one
+    // 32k-position draft forward — a large activation/mask transient on a
+    // 24 GB box, and untested numerics vs the chunked oracle.)
+    const tailSplit = flagOn("MLX_BUN_PREFILL_TAIL_SPLIT", true);
+    const upTo = tailSplit && promptIds.length > 1 ? promptIds.length - 1 : promptIds.length;
+    for (let at = 0; at < upTo; at += PREFILL_CHUNK) {
+      const end = Math.min(at + PREFILL_CHUNK, upTo);
+      const ids = ops.fromInt32(promptIds.slice(at, end), [1, end - at]);
+      const h = this.model.forwardHidden(ids, this.caches);
+      ids.dispose();
+      // The head's last-position logits are not needed: the serve loop's
+      // first `feed` is consumed at the top of draft().
+      h.dispose();
+      clearCache();
+    }
+  }
+
+  /** One forward over `feed`, sample the first draft from its last position,
+   *  then n-1 single-token steps — mirrors _draft_generate's autoregressive
+   *  chain (with async_eval pipelining upstream; eager here, the draft is
+   *  small and the serve loop overlaps nothing with it in v1). */
+  draft(feed: number[], n: number, stepBase: number): number[] {
+    const drafts: number[] = [];
+    let input = feed;
+    for (let k = 0; k < n; k++) {
+      const ids = ops.fromInt32(input, [1, input.length]);
+      const h = this.model.forwardHidden(ids, this.caches);
+      ids.dispose();
+      const L = h.shape[1]!;
+      const H = h.shape[2]!;
+      const hLast = h.slice([0, L - 1, 0], [1, L, H]);
+      h.dispose();
+      const logits = this.model.logitsFromHidden(hLast);
+      hLast.dispose();
+      const V = logits.shape[logits.shape.length - 1]!;
+      const flat = ops.reshape(logits, [1, V]);
+      logits.dispose();
+      const lp = toLogprobs(flat);
+      flat.dispose();
+      const tokArr = this.sampler(lp, stepBase + k);
+      lp.dispose();
+      const tok = ops.itemUint32(tokArr);
+      tokArr.dispose();
+      drafts.push(tok);
+      input = [tok];
+    }
+    return drafts;
+  }
+
+  commit(n: number, kAccept: number): void {
+    const rewind = Math.max(n - kAccept - 1, 0);
+    if (rewind > 0) {
+      for (const c of this.caches) {
+        if (!c.isTrimmable()) {
+          // The serve loop stops speculating on the TARGET-side trim check
+          // before this can matter; a wrapped draft ring here just means the
+          // source is done being useful. Recreate is the safe recovery.
+          this.resetCaches();
+          return;
+        }
+        c.trim(rewind);
+      }
+    }
+  }
+
+  private resetCaches(): void {
+    for (const c of this.caches) c.dispose();
+    this.caches = this.model.makeCache();
+  }
+
+  dispose(): void {
+    for (const c of this.caches) c.dispose();
+  }
+}

package/src/ssd-cache.ts

@@ -0,0 +1,285 @@
+// SSD cold tier for the prompt/KV cache (docs/design/ssd-kv-cold-tier.md).
+//
+// Files ARE the database: no sidecar index. Layout
+//   <dir>/<configFingerprint>/<nsHash>/<uuid>.mlxkv
+// so two models (or two kv-quant schemes of one model) share a directory
+// without ever colliding, and an incompatible fingerprint dir is simply
+// ignored (never deleted — it may belong to the user's other model).
+// Startup recovery = a header-only scan of OUR fingerprint dir: header hash
+// verified (cheap), corrupt or metadata-mismatched files unlinked, `.tmp`
+// write orphans reaped. LRU state = file mtime (bumped on hit); the byte
+// cap is enforced at write time by oldest-mtime eviction. A cache is always
+// droppable: every failure path degrades to "no hit" or "not stored" — the
+// tier must never take serving down (oMLX paged_ssd_cache lesson, Apache-2.0
+// idea port; ours spills whole prefix entries, not content-hashed blocks —
+// see the design doc's D1 for why).
+
+import {
+  existsSync, mkdirSync, readdirSync, rmSync, statSync, utimesSync,
+} from "node:fs";
+import { join } from "node:path";
+import { randomUUID } from "node:crypto";
+import { commonPrefixLength } from "./prompt-cache";
+import {
+  saveKvCache, saveKvCacheAsync, loadKvCache, readKvHeader,
+  type KvSaveMeta, type KvLoadExpect, type LoadedKvCache,
+} from "./kv-store";
+import type { Cache } from "./model/gemma4-base";
+
+export interface SsdIndexEntry {
+  path: string;
+  ns: string;
+  tokens: number[];
+  bytes: number;
+  mtimeMs: number;
+  /** Whether every serialized cache can trim (derived from kinds+state at
+   *  store/scan time). An untrimmable entry (wrapped ring, SSM) can only
+   *  seed prompts it matches WITHOUT a trim — find() must not prefer it
+   *  over a shorter usable one (the 2026-07-06 restart-0 defect: the
+   *  [prompt+gen] file always won, always got rejected after restore). */
+  trimmable: boolean;
+}
+
+/** Trimmability from a file header alone (mirror of the live caches'
+ *  isTrimmable(): kv/qkv/turboquant always trim; rotating kinds only pre-wrap
+ *  (offset < maxSize); ssm never). Keep in sync with gemma4-base/qwen3-delta. */
+const headerTrimmable = (caches: { kind: string; offset: number; maxSize?: number }[]): boolean =>
+  caches.every((c) =>
+    c.kind === "kv" || c.kind === "qkv" || c.kind === "turboquant"
+      ? true
+      : c.kind === "rotating" || c.kind === "rotating-qkv"
+        ? c.offset < (c.maxSize ?? 0)
+        : false);
+
+export interface SsdStoreOptions {
+  dir: string;
+  maxBytes: number;
+  /** Server identity — written into every file, enforced on scan/load. */
+  configFingerprint: string;
+  tokenizerHash: string;
+  modelId: string;
+  /** Verify every tensor hash on load (--ssd-cache-verify). */
+  verify?: boolean;
+}
+
+const nsHash = (ns: string): string =>
+  ns === "" ? "base" : Bun.hash(ns).toString(16);
+
+export class SsdCacheStore {
+  readonly #opts: SsdStoreOptions;
+  readonly #root: string; // <dir>/<configFingerprint>
+  #index: SsdIndexEntry[] = [];
+  #warnedWriteFailure = false;
+  stats = { restores: 0, spills: 0, restoreMsLast: 0 };
+
+  constructor(opts: SsdStoreOptions) {
+    this.#opts = opts;
+    this.#root = join(opts.dir, opts.configFingerprint);
+  }
+
+  get entries(): number {
+    return this.#index.length;
+  }
+
+  get totalBytes(): number {
+    return this.#index.reduce((a, e) => a + e.bytes, 0);
+  }
+
+  get maxBytes(): number {
+    return this.#opts.maxBytes;
+  }
+
+  /** Startup recovery: rebuild the in-memory index from disk. Corrupt
+   *  headers and metadata mismatches are unlinked; `.tmp` orphans reaped.
+   *  Only OUR fingerprint dir is touched. Returns entries indexed. */
+  scan(): number {
+    this.#index = [];
+    if (!existsSync(this.#root)) return 0;
+    for (const nsDir of readdirSync(this.#root)) {
+      const nsPath = join(this.#root, nsDir);
+      let files: string[];
+      try { files = readdirSync(nsPath); } catch { continue; }
+      for (const f of files) {
+        const path = join(nsPath, f);
+        if (f.endsWith(".tmp")) {
+          // A crash mid-write: the rename never happened. Ours by
+          // construction (uuid names), safe to reap.
+          try { rmSync(path, { force: true }); } catch {}
+          continue;
+        }
+        if (!f.endsWith(".mlxkv")) continue;
+        try {
+          const h = readKvHeader(path); // verifies the header hash
+          if (
+            h.configFingerprint !== this.#opts.configFingerprint ||
+            h.tokenizerHash !== this.#opts.tokenizerHash
+          )
+            throw new Error("metadata mismatch");
+          const st = statSync(path);
+          this.#index.push({
+            path, ns: h.ns ?? "", tokens: h.tokens, bytes: st.size, mtimeMs: st.mtimeMs,
+            trimmable: headerTrimmable(h.caches),
+          });
+        } catch {
+          try { rmSync(path, { force: true }); } catch {}
+        }
+      }
+    }
+    return this.#index.length;
+  }
+
+  /** Best USABLE stored entry for `prompt` in `ns`: the longest common
+   *  prefix (capped at prompt.length−1 — at least one token must forward).
+   *  An entry LONGER than the matched prefix is returned only when it can
+   *  actually seed the prompt: the divergent tail must be trimmable after
+   *  restore. Untrimmable entries (ring post-wrap, SSM) that would need a
+   *  trim are SKIPPED here — before this gate, the big [prompt+gen] file
+   *  always outranked the usable boundary snapshot, got restored, and was
+   *  thrown away by take()'s backstop (the 2026-07-06 restart-0 + wasted
+   *  84 s restore-then-re-prefill path). Pure index lookup; no I/O. */
+  find(prompt: number[], ns = ""): { entry: SsdIndexEntry; prefixLen: number } | null {
+    let best: SsdIndexEntry | null = null;
+    let bestLen = 0;
+    for (const e of this.#index) {
+      if (e.ns !== ns) continue;
+      const p = Math.min(commonPrefixLength(e.tokens, prompt), prompt.length - 1);
+      if (p === 0 || p <= bestLen) continue;
+      if (e.tokens.length - p > 0 && !e.trimmable) continue;
+      bestLen = p;
+      best = e;
+    }
+    return best ? { entry: best, prefixLen: bestLen } : null;
+  }
+
+  /** Restore an entry via kv-store's STREAMED COPY (bounded host
+   *  transient — live entry + one tensor; no mapping survives the call,
+   *  the caches own their bytes). Bumps LRU mtime. On ANY failure the
+   *  file is dropped and null returned. */
+  restore(entry: SsdIndexEntry, model: { makeCache(): Cache[] }): LoadedKvCache | null {
+    const t0 = performance.now();
+    try {
+      const expect: KvLoadExpect = {
+        configFingerprint: this.#opts.configFingerprint,
+        tokenizerHash: this.#opts.tokenizerHash,
+        ns: entry.ns,
+        verify: this.#opts.verify,
+      };
+      const loaded = loadKvCache(entry.path, model, expect);
+      const now = new Date();
+      try { utimesSync(entry.path, now, now); entry.mtimeMs = now.getTime(); } catch {}
+      this.stats.restores++;
+      this.stats.restoreMsLast = performance.now() - t0;
+      return loaded;
+    } catch (err) {
+      // Loud: a failed restore silently degraded to a full re-prefill for
+      // a month before the 2026-07-06 bench made it observable.
+      console.warn(`[ssd-cache] restore failed, dropping ${entry.path}: ${(err as Error).message}`);
+      this.remove(entry.path);
+      return null;
+    }
+  }
+
+  /** Persist an entry. Synchronous (the tier calls it on the idle serial
+   *  lane between requests); atomic via kv-store's tmp+fsync+rename. An
+   *  entry bigger than the cap is refused; disk/write failure is a warn-once
+   *  soft-fail — serving never stalls on the cold tier. Existing entries in
+   *  the same ns whose tokens are a PREFIX of the new entry are superseded
+   *  (the agent-conversation pattern: one growing entry, not N generations
+   *  of it). Returns true if stored. */
+  store(tokens: number[], caches: Cache[], ns = ""): boolean {
+    const dir = join(this.#root, nsHash(ns));
+    const path = join(dir, `${randomUUID()}.mlxkv`);
+    try {
+      mkdirSync(dir, { recursive: true });
+      saveKvCache(path, tokens, caches, this.#meta(ns));
+      return this.#indexStored(path, tokens, caches, ns);
+    } catch (err) {
+      return this.#storeFailed(path, err);
+    }
+  }
+
+  /** Non-blocking store (the write-behind persistence path): same file,
+   *  index, and supersede semantics as store(), but the tensor flush
+   *  yields the event loop between tensors so serving interleaves.
+   *  Caller passes zero-copy CLONES it owns (a consistent snapshot no
+   *  matter what the live entry does meanwhile) and disposes them after.
+   *  `waitTurn` gates every per-tensor step (see saveKvCacheAsync) — the
+   *  server passes the gateway's onIdle so the flush only progresses while
+   *  the engine is idle, never mid-decode. */
+  async storeAsync(tokens: number[], caches: Cache[], ns = "", waitTurn?: () => Promise<void>): Promise<boolean> {
+    const dir = join(this.#root, nsHash(ns));
+    const path = join(dir, `${randomUUID()}.mlxkv`);
+    try {
+      mkdirSync(dir, { recursive: true });
+      await saveKvCacheAsync(path, tokens, caches, this.#meta(ns), waitTurn);
+      return this.#indexStored(path, tokens, caches, ns);
+    } catch (err) {
+      return this.#storeFailed(path, err);
+    }
+  }
+
+  #meta(ns: string): KvSaveMeta {
+    return {
+      modelId: this.#opts.modelId,
+      configFingerprint: this.#opts.configFingerprint,
+      tokenizerHash: this.#opts.tokenizerHash,
+      ns,
+    };
+  }
+
+  #indexStored(path: string, tokens: number[], caches: Cache[], ns: string): boolean {
+    const st = statSync(path);
+    if (st.size > this.#opts.maxBytes) {
+      rmSync(path, { force: true });
+      console.warn(`[ssd-cache] entry not stored: ${st.size} bytes exceeds the ${this.#opts.maxBytes}-byte cap`);
+      return false;
+    }
+    const trimmable = caches.every((c) => c.isTrimmable());
+    // Exact duplicates (same tokens) are replaced regardless of
+    // trimmability — the new file serves exactly the old one's matches.
+    for (const e of [...this.#index]) {
+      if (e.ns !== ns || e.tokens.length !== tokens.length) continue;
+      if (commonPrefixLength(e.tokens, tokens) === tokens.length) this.remove(e.path);
+    }
+    // Supersede shadowed ancestors of this prefix (same ns) — ONLY when
+    // the new entry can serve their prefixes (all caches trimmable; the
+    // same guard as PromptCache.put()). An untrimmable [prompt+gen]
+    // entry deleting the boundary snapshot was the gemma restart-0
+    // defect (2026-07-06): the only restorable file died the moment the
+    // unrestorable one landed.
+    if (trimmable) {
+      for (const e of [...this.#index]) {
+        if (e.ns !== ns || e.tokens.length >= tokens.length) continue;
+        if (commonPrefixLength(e.tokens, tokens) === e.tokens.length) this.remove(e.path);
+      }
+    }
+    this.#index.push({ path, ns, tokens, bytes: st.size, mtimeMs: Date.now(), trimmable });
+    this.stats.spills++;
+    this.evictToCap();
+    return true;
+  }
+
+  #storeFailed(path: string, err: unknown): boolean {
+    try { rmSync(path, { force: true }); rmSync(`${path}.tmp`, { force: true }); } catch {}
+    if (!this.#warnedWriteFailure) {
+      this.#warnedWriteFailure = true;
+      console.warn(`[ssd-cache] store failed (disk full or unwritable?) — cold tier disabled for this entry: ${(err as Error).message}`);
+    }
+    return false;
+  }
+
+  /** Enforce the byte cap: unlink oldest-mtime entries until under it. */
+  evictToCap(): void {
+    while (this.totalBytes > this.#opts.maxBytes && this.#index.length > 0) {
+      let oldest = 0;
+      for (let i = 1; i < this.#index.length; i++)
+        if (this.#index[i]!.mtimeMs < this.#index[oldest]!.mtimeMs) oldest = i;
+      this.remove(this.#index[oldest]!.path);
+    }
+  }
+
+  remove(path: string): void {
+    try { rmSync(path, { force: true }); } catch {}
+    this.#index = this.#index.filter((e) => e.path !== path);
+  }
+}