@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/state/ssd-cache.ts

@@ -0,0 +1,434 @@
+import type { CheckpointAttachment } from "../contracts/mlx/checkpoint";
+// SSD cold tier for the prompt/KV cache (`02d723a:docs/design/kv-cache.md` §5).
+//
+// Files ARE the database: no sidecar index. Layout
+//   <dir>/<configFingerprint>/<nsHash>/<uuid>.mlxkv
+// where the caller's configFingerprint includes the model identity, so two
+// models (or two kv-quant schemes, or two downloaded snapshot revisions)
+// share a directory without ever colliding, and a foreign 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 files unlinked, files whose header names another
+// identity SKIPPED (never deleted), `.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 { kvWriter, type KvWriteRequest } from "./persistence-worker";
+import { join, dirname } from "node:path";
+import { randomUUID } from "node:crypto";
+import { commonPrefixLength } from "./prefix-cache";
+import { saveKvCache, saveKvCacheAsync, loadKvCache, loadKvCacheAsync, readKvHeader, cacheHeadersTrimmable, cacheHeadersMinimumReusableOffset, legacyCacheCodecs } from "./persistence";
+import { type CacheCodecProvider, type KvSaveMeta, type KvLoadExpect, type LoadedKvCache } from "./persistence-types";
+import type { Cache } from "../contracts/mlx/cache";
+import { minimumReusableOffset } from "./views";
+
+export interface SsdIndexEntry {
+  blocks?: Array<{ hash: string; bytes: number }>;
+  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;
+  minimumReusableOffset?: number;
+  generationCheckpoint?: NonNullable<KvSaveMeta["generationCheckpoint"]>;
+}
+
+export interface SsdStoreOptions {
+  storage?: Pick<KvWriteRequest, "layout" | "blockBytes" | "segmented">;
+  codecs?: CacheCodecProvider;
+  dir: string;
+  maxBytes: number;
+  /** A live byte cap that replaces `maxBytes` when set: a caller sharing one
+   *  budget across several stores answers with what this store may hold now. */
+  limit?: () => 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 #codecs: CacheCodecProvider;
+  readonly #root: string; // <dir>/<configFingerprint>
+  #index: SsdIndexEntry[] = [];
+  #exact = new Map<string, SsdIndexEntry>();
+  #tokenKeys = new WeakMap<number[], string>();
+  #key(tokens: number[], ns: string): string {
+    let hash = this.#tokenKeys.get(tokens);
+    if (!hash) {
+      hash = Bun.hash(new Uint32Array(tokens)).toString(16);
+      this.#tokenKeys.set(tokens, hash);
+    }
+    return JSON.stringify([ns, tokens.length, hash]);
+  }
+  #addIndex(entry: SsdIndexEntry): void {
+    this.#index.push(entry);
+    this.#exact.set(this.#key(entry.tokens, entry.ns), entry);
+  }
+  /** Exact checkpoint selection for a known session, independent of residency. */
+  findExact(tokens: number[], ns = ""): { entry: SsdIndexEntry; prefixLen: number } | null {
+    const entry = this.#exact.get(this.#key(tokens, ns));
+    return entry && commonPrefixLength(entry.tokens, tokens) === tokens.length
+      ? { entry, prefixLen: tokens.length } : null;
+  }
+
+  #warnedWriteFailure = false;
+  stats = { restores: 0, spills: 0, restoreMsLast: 0 };
+
+  constructor(opts: SsdStoreOptions) {
+    this.#opts = opts;
+    this.#codecs = opts.codecs ?? legacyCacheCodecs;
+    this.#root = join(opts.dir, opts.configFingerprint);
+  }
+
+  get entries(): number {
+    return this.#index.length;
+  }
+
+  /** Header-only provenance for the serving tokenizer; no tensor restore. */
+  *tokenPrefixes(): Iterable<readonly number[]> {
+    for (const entry of this.#index) yield entry.tokens;
+  }
+
+  get totalBytes(): number {
+    const blocks = new Map<string, number>();
+    for (const entry of this.#index) for (const block of entry.blocks ?? [])
+      blocks.set(`${dirname(entry.path)}/${block.hash}`, block.bytes);
+    return this.#index.reduce((a, e) => a + e.bytes, 0) + [...blocks.values()].reduce((a, b) => a + b, 0);
+  }
+
+  get maxBytes(): number {
+    return this.#opts.limit?.() ?? this.#opts.maxBytes;
+  }
+
+  /** Longest token prefix whose atomic cache file is present in the index. */
+  get longestDurablePrefixTokens(): number {
+    return this.#index.reduce((best, entry) => Math.max(best, entry.tokens.length), 0);
+  }
+
+  /** True when the SSD index can seed every token in this prefix. A longer
+   * trimmable descendant is equivalent to an exact boundary snapshot. */
+  hasDurablePrefix(tokens: number[], ns = ""): boolean {
+    return this.#index.some((entry) => {
+      if (entry.ns !== ns || entry.tokens.length < tokens.length || tokens.length < (entry.minimumReusableOffset ?? 0)) return false;
+      if (commonPrefixLength(entry.tokens, tokens) !== tokens.length) return false;
+      return entry.tokens.length === tokens.length || entry.trimmable;
+    });
+  }
+
+  /** Startup recovery: rebuild the in-memory index from disk. Corrupt
+   *  headers are unlinked; `.tmp` orphans reaped; a file whose header names
+   *  another model/tokenizer/fingerprint is left in place and not indexed.
+   *  Only OUR fingerprint dir is touched. Returns entries indexed. */
+  scan(): number {
+    this.#index = []; this.#exact.clear();
+    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
+          // Another identity's entry: never ours to delete or serve.
+          if (
+            h.modelId !== this.#opts.modelId ||
+            h.configFingerprint !== this.#opts.configFingerprint ||
+            h.tokenizerHash !== this.#opts.tokenizerHash
+          ) continue;
+          if ((h.codecProvider ?? legacyCacheCodecs.id) !== this.#codecs.id) continue;
+          const st = statSync(path);
+          this.#addIndex({
+            blocks: [...h.caches, ...(h.attachments ?? [])].flatMap(e => e.tensors.flatMap(t => t.blocks ?? [])),
+            path, ns: h.ns ?? "", tokens: h.tokens, bytes: st.size, mtimeMs: st.mtimeMs,
+            trimmable: !h.attachments?.length && cacheHeadersTrimmable(h.caches, this.#codecs),
+            minimumReusableOffset: cacheHeadersMinimumReusableOffset(h.caches),
+            ...(h.generationCheckpoint
+              ? { generationCheckpoint: h.generationCheckpoint }
+              : {}),
+          });
+        } catch {
+          try { rmSync(path, { force: true }); } catch {}
+        }
+      }
+    }
+    void kvWriter.collect(this.#root).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 || p < (e.minimumReusableOffset ?? 0)) 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 packages/inference/src/state/persistence.ts'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 = {
+        modelId: this.#opts.modelId,
+        configFingerprint: this.#opts.configFingerprint,
+        tokenizerHash: this.#opts.tokenizerHash,
+        ns: entry.ns,
+        verify: this.#opts.verify,
+      };
+      const loaded = loadKvCache(entry.path, model, expect, this.#codecs);
+      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;
+    }
+  }
+
+  /** Read on the dedicated CPU queue, then let the model codec adopt state. */
+  async restoreAsync(entry: SsdIndexEntry, model: { makeCache(): Cache[] }): Promise<LoadedKvCache | null> {
+    const t0 = performance.now();
+    try {
+      const loaded = await loadKvCacheAsync(entry.path, model, {
+        ...this.#meta(entry.ns), verify: this.#opts.verify,
+      }, this.#codecs);
+      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 (error) {
+      console.warn(`[ssd-cache] restore failed, dropping ${entry.path}: ${error}`);
+      this.remove(entry.path); return null;
+    }
+  }
+
+  /** Persist an entry. Synchronous (the tier calls it on the idle serial
+   *  lane between requests); atomic via
+   *  packages/inference/src/state/persistence.ts'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 = "", attachments?: CheckpointAttachment[]): 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), attachments }, this.#codecs);
+      return this.#indexStored(path, tokens, caches, ns, undefined, attachments);
+    } 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.
+   *  `runStep` optionally wraps owner-thread snapshot preparation. Hashing,
+   *  packing and disk writes execute on the CPU worker. */
+  async storeAsync(
+    tokens: number[], caches: Cache[], ns = "",
+    runStep?: <T>(step: () => T) => Promise<T>,
+    attachments?: CheckpointAttachment[],
+  ): Promise<boolean> {
+    const dir = join(this.#root, nsHash(ns));
+    const path = join(dir, `${randomUUID()}.mlxkv`);
+    try {
+      await saveKvCacheAsync(path, tokens, caches, { ...this.#meta(ns), attachments }, runStep, this.#codecs, this.#opts.storage);
+      return this.#indexStored(path, tokens, caches, ns, undefined, attachments);
+    } catch (err) {
+      return this.#storeFailed(path, err);
+    }
+  }
+
+  /** Longest durable in-flight generation for an identical request. The
+   *  checkpoint namespace is separate from the ordinary prompt-cache tier,
+   *  so incomplete assistant output can never seed an unrelated prefill. */
+  findGenerationCheckpoint(
+    prompt: number[], key: string, cacheNs = "",
+  ): SsdIndexEntry | null {
+    let best: SsdIndexEntry | null = null;
+    for (const entry of this.#index) {
+      const checkpoint = entry.generationCheckpoint;
+      if (!checkpoint || checkpoint.key !== key || checkpoint.cacheNs !== cacheNs)
+        continue;
+      if (checkpoint.originalPromptTokens !== prompt.length) continue;
+      if (entry.tokens.length <= prompt.length) continue;
+      if (commonPrefixLength(entry.tokens, prompt) !== prompt.length) continue;
+      if (!best || entry.tokens.length > best.tokens.length) best = entry;
+    }
+    return best;
+  }
+
+  /** Persist one in-flight generation atomically. A completed newer
+   *  checkpoint supersedes older checkpoints for the same request only after
+   *  its rename succeeds, so a crash during a write leaves the prior point. */
+  async storeGenerationCheckpoint(
+    tokens: number[], caches: Cache[], checkpoint: NonNullable<KvSaveMeta["generationCheckpoint"]>,
+    runStep?: <T>(step: () => T) => Promise<T>,
+  ): Promise<boolean> {
+    const ns = `__generation_checkpoint__:${checkpoint.key}`;
+    const dir = join(this.#root, nsHash(ns));
+    const path = join(dir, `${randomUUID()}.mlxkv`);
+    try {
+      if (runStep) await runStep(() => mkdirSync(dir, { recursive: true }));
+      else mkdirSync(dir, { recursive: true });
+      await saveKvCacheAsync(path, tokens, caches, {
+        ...this.#meta(ns), generationCheckpoint: checkpoint,
+      }, runStep, this.#codecs, this.#opts.storage);
+      const index = () => {
+        const stored = this.#indexStored(path, tokens, caches, ns, checkpoint);
+        if (stored) {
+          for (const entry of [...this.#index]) {
+            if (entry.path === path) continue;
+            if (entry.generationCheckpoint?.key === checkpoint.key) this.remove(entry.path);
+          }
+        }
+        return stored;
+      };
+      return runStep ? await runStep(index) : index();
+    } catch (err) {
+      return runStep ? await runStep(() => this.#storeFailed(path, err)) : this.#storeFailed(path, err);
+    }
+  }
+
+  removeGenerationCheckpoints(key: string): void {
+    for (const entry of [...this.#index])
+      if (entry.generationCheckpoint?.key === key) this.remove(entry.path);
+  }
+
+  #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,
+    generationCheckpoint?: NonNullable<KvSaveMeta["generationCheckpoint"]>,
+    attachments?: CheckpointAttachment[],
+  ): boolean {
+    const st = statSync(path);
+    const header = readKvHeader(path);
+    const blocks = [...header.caches, ...(header.attachments ?? [])].flatMap(e => e.tensors.flatMap(t => t.blocks ?? []));
+    const unique = new Map(blocks.map(block => [block.hash, block.bytes]));
+    const entryBytes = st.size + [...unique.values()].reduce((a, b) => a + b, 0);
+    if (entryBytes > this.maxBytes) {
+      this.remove(path);
+      console.warn(`[ssd-cache] entry not stored: ${entryBytes} bytes exceeds the ${this.maxBytes}-byte cap`);
+      return false;
+    }
+    const trimmable = !attachments?.length && 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 || e.tokens.length < minimumReusableOffset(caches)) continue;
+        if (commonPrefixLength(e.tokens, tokens) === e.tokens.length) this.remove(e.path);
+      }
+    }
+    this.#addIndex({
+      blocks, path, ns, tokens, bytes: st.size, mtimeMs: Date.now(), trimmable,
+      minimumReusableOffset: minimumReusableOffset(caches),
+      ...(generationCheckpoint ? { generationCheckpoint } : {}),
+    });
+    this.stats.spills++;
+    this.evictToCap();
+    return true;
+  }
+
+  #storeFailed(path: string, err: unknown): boolean {
+    this.remove(path);
+    try { 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.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 {}
+    for (const entry of this.#index) if (entry.path === path) {
+      const key = this.#key(entry.tokens, entry.ns);
+      if (this.#exact.get(key) === entry) this.#exact.delete(key);
+    }
+    this.#index = this.#index.filter((e) => e.path !== path);
+    void kvWriter.collect(this.#root).catch(() => {});
+  }
+}

package/src/state/ssm-prefill-padding.ts

@@ -0,0 +1,39 @@
+import type { PrefillPadding } from "../contracts/mlx/cache";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+
+/** ArraysCache's remaining lengths and padding from mlx-lm (MIT).
+ * Recurrent state has no token axis to roll: masked tokens must not update it. */
+export class SsmPrefillPadding {
+  #left: number[];
+  #lengths: number[];
+  constructor(padding: PrefillPadding) {
+    this.#left = padding.leftPadding ? [...padding.leftPadding] : padding.lengths.map(() => 0);
+    this.#lengths = [...padding.lengths];
+  }
+  convolutionLengths(tokens: number): number[] {
+    return this.#lengths.map(length => Math.max(0, Math.min(tokens, length)));
+  }
+  makeMask(tokens: number): MlxArray {
+    const B = this.#lengths.length;
+    using arange = ops.arange(0, tokens, 1, Dtype.int32);
+    using positions = ops.reshape(arange, [1, tokens]);
+    using left = ops.fromInt32(this.#left, [B, 1]);
+    using ends = ops.fromInt32(this.#lengths, [B, 1]);
+    using after = ops.greaterEqual(positions, left);
+    using before = ops.less(positions, ends);
+    return ops.logicalAnd(after, before);
+  }
+  advance(tokens: number): number[] {
+    const counts = this.#lengths.map((length, row) => Math.max(0,
+      Math.min(tokens, length) - Math.max(0, Math.min(tokens, this.#left[row]!))));
+    this.#lengths = this.#lengths.map(length => length - tokens);
+    this.#left = this.#left.map(left => left - tokens);
+    return counts;
+  }
+  filter(keep: readonly number[]): void {
+    this.#left = keep.map(row => this.#left[row]!);
+    this.#lengths = keep.map(row => this.#lengths[row]!);
+  }
+}