mlx-bun 0.3.0 → 0.5.0

package/src/storage/kv-writer.ts

@@ -0,0 +1,83 @@
+import { Worker } from "node:worker_threads";
+import workerSource from "./kv-writer.worker.js" with { type: "text" };
+import type { KvFileHeader } from "../kv-store";
+
+/** Immutable, evaluated host-visible storage. The owner retains the native
+ * buffers until write() settles. No MLX handles cross the worker boundary. */
+export interface StoredTensorView {
+  pointer: number;
+  shape: number[];
+  strides: number[];
+  itemSize: number;
+}
+
+export interface KvWriteRequest {
+  layout?: "whole" | "blocks";
+  blockBytes?: number;
+  segmented?: boolean;
+  path: string;
+  header: KvFileHeader;
+  tensors: StoredTensorView[];
+}
+
+/** Hashing, layout packing and atomic disk writes run on a dedicated CPU
+ * worker. Only metadata and pointers are posted; tensor bytes are not copied
+ * through structured clone. The RAM cache remains an independent owner. */
+export class KvWriter {
+  #worker: Worker | null = null;
+  #next = 0;
+  #pending = new Map<number, { resolve(): void; reject(error: Error): void }>();
+
+  lastMetrics: { elapsedMs: number; writtenBytes?: number; reusedBytes?: number; copiedBytes?: number; scratchPeak?: number } | null = null;
+
+  write(request: KvWriteRequest): Promise<void> { return this.#send(request); }
+
+  read(request: { path: string; header: KvFileHeader & { dataStart: number }; pointers: number[]; verify?: boolean }): Promise<void> {
+    return this.#send({ ...request, operation: "read" });
+  }
+
+  collect(root: string): Promise<void> { return this.#send({ root, operation: "collect" }); }
+
+  #send(request: object): Promise<void> {
+    const worker = this.#worker ?? this.#start();
+    worker.ref();
+    const id = ++this.#next;
+    return new Promise<void>((resolve, reject) => {
+      this.#pending.set(id, { resolve, reject });
+      try { worker.postMessage({ id, ...request }); }
+      catch (error) {
+        this.#pending.delete(id);
+        if (!this.#pending.size) worker.unref();
+        reject(error);
+      }
+    });
+  }
+
+  #start(): Worker {
+    // Embed the worker in the single-file executable as well as source runs.
+    const worker = new Worker(workerSource, { eval: true });
+    worker.on("message", (message: { id: number; error?: string; metrics?: KvWriter["lastMetrics"] }) => {
+      const result = this.#pending.get(message.id);
+      if (!result) return;
+      this.#pending.delete(message.id);
+      if (message.error) result.reject(new Error(message.error));
+      else { this.lastMetrics = message.metrics ?? null; result.resolve(); }
+      if (!this.#pending.size) worker.unref();
+    });
+    // Job failures are acknowledged by the worker. A worker-level failure
+    // settles only at exit, after it can no longer read borrowed buffers.
+    worker.on("error", () => {});
+    worker.on("exit", () => {
+      for (const job of this.#pending.values()) job.reject(new Error("KV persistence worker exited"));
+      this.#pending.clear();
+      this.#worker = null;
+    });
+    this.#worker = worker;
+    return worker;
+  }
+}
+
+export const kvWriter = new KvWriter();
+
+/** Reads have their own queue: a restore never waits behind persistence. */
+export const kvReader = new KvWriter();

package/src/storage/kv-writer.worker.js

@@ -0,0 +1,264 @@
+// CPU-only persistence worker. It deliberately imports no model, MLX,
+// scheduler or cache manager. The sender owns all buffers through completion.
+import { parentPort } from "node:worker_threads";
+import { toArrayBuffer, ptr } from "bun:ffi";
+import { openSync, writeSync, closeSync, fsyncSync, renameSync, rmSync, mkdirSync, existsSync, readSync, readFileSync, readdirSync } from "node:fs";
+import { createHash } from "node:crypto";
+import { dirname, join } from "node:path";
+
+const MAGIC = "MLXBUNKV2\n";
+const PREFIX_LEN = MAGIC.length + 16;
+const alignUp = n => Math.ceil(n / 16384) * 16384;
+const hash64 = bytes => Bun.hash(bytes).toString(16).padStart(16, "0");
+let scratch = new Uint8Array(0);
+
+function bytesOf(tensor) {
+  const { pointer, shape, strides, itemSize } = tensor;
+  const size = shape.reduce((a, b) => a * b, 1);
+  if (!size) return new Uint8Array(0);
+  // Coalesce the contiguous suffix. Most KV layouts need only a copy per
+  // head to omit capacity padding, never a GPU contiguous() operation.
+  let run = 1, axis = shape.length - 1;
+  while (axis >= 0 && (shape[axis] === 1 || strides[axis] === run)) {
+    run *= shape[axis--];
+  }
+  if (axis < 0) return new Uint8Array(toArrayBuffer(pointer, 0, size * itemSize));
+  let lo = 0, hi = 0;
+  for (let i = 0; i < shape.length; i++) {
+    const end = (shape[i] - 1) * strides[i];
+    lo += Math.min(0, end); hi += Math.max(0, end);
+  }
+  const source = new Uint8Array(toArrayBuffer(pointer + lo * itemSize, 0, (hi - lo + 1) * itemSize));
+  if (scratch.length < size * itemSize) scratch = new Uint8Array(size * itemSize);
+  const bytes = scratch.subarray(0, size * itemSize);
+  const runBytes = run * itemSize;
+  for (let row = 0; row < size / run; row++) {
+    let index = row, offset = -lo;
+    for (let i = axis; i >= 0; i--) {
+      offset += (index % shape[i]) * strides[i];
+      index = Math.floor(index / shape[i]);
+    }
+    bytes.set(source.subarray(offset * itemSize, offset * itemSize + runBytes), row * runBytes);
+  }
+  return bytes;
+}
+
+function writeAll(fd, bytes, position) {
+  let offset = 0;
+  while (offset < bytes.length) {
+    const n = writeSync(fd, bytes, offset, bytes.length - offset, position + offset);
+    if (!n) throw new Error("KV write made no progress");
+    offset += n;
+  }
+}
+
+function writeSnapshot({ path, header, tensors }) {
+  const slots = [...header.caches, ...(header.attachments ?? [])].flatMap(entry => entry.tensors);
+  const headerLen = new TextEncoder().encode(JSON.stringify(header)).length;
+  const dataStart = alignUp(PREFIX_LEN + headerLen);
+  mkdirSync(dirname(path), { recursive: true });
+  const tmp = `${path}.tmp`;
+  const fd = openSync(tmp, "w");
+  try {
+    for (let i = 0; i < slots.length; i++) {
+      const bytes = bytesOf(tensors[i]);
+      if (bytes.length !== slots[i].bytes) throw new Error("KV tensor byte-length drift");
+      slots[i].hash = hash64(bytes);
+      writeAll(fd, bytes, dataStart + slots[i].off);
+    }
+    const json = new TextEncoder().encode(JSON.stringify(header));
+    if (json.length !== headerLen) throw new Error("KV header byte-length drift");
+    const prefix = new Uint8Array(dataStart);
+    prefix.set(new TextEncoder().encode(MAGIC));
+    const view = new DataView(prefix.buffer);
+    view.setUint32(MAGIC.length, headerLen, true);
+    view.setUint32(MAGIC.length + 4, dataStart, true);
+    view.setBigUint64(MAGIC.length + 8, BigInt(Bun.hash(json)), true);
+    prefix.set(json, PREFIX_LEN);
+    writeAll(fd, prefix, 0);
+    fsyncSync(fd);
+  } catch (error) {
+    closeSync(fd);
+    rmSync(tmp, { force: true });
+    throw error;
+  }
+  closeSync(fd);
+  try { renameSync(tmp, path); }
+  catch (error) { rmSync(tmp, { force: true }); throw error; }
+}
+
+// The block boundary resets for each leading row (e.g. attention head).
+// Growing a token axis therefore leaves preceding blocks byte-identical.
+function* segments(tensor, blockBytes, metrics) {
+  const { pointer, shape, strides, itemSize } = tensor;
+  const size = shape.reduce((a, b) => a * b, 1);
+  const rowSize = shape.slice(-2).reduce((a, b) => a * b, 1);
+  if (!size) return;
+  let run = 1, axis = shape.length - 1;
+  while (axis >= 0 && (shape[axis] === 1 || strides[axis] === run)) run *= shape[axis--];
+  const limit = Math.max(1, Math.floor(blockBytes / itemSize));
+  let packing = null;
+  for (let start = 0; start < size;) {
+    const count = Math.min(limit, rowSize - start % rowSize, size - start);
+    let index = start, offset = 0;
+    for (let i = shape.length - 1; i >= 0; i--) {
+      offset += index % shape[i] * strides[i]; index = Math.floor(index / shape[i]);
+    }
+    if (count <= run - start % run) {
+      yield new Uint8Array(toArrayBuffer(pointer + offset * itemSize, 0, count * itemSize));
+    } else {
+      if (!packing) packing = new Uint8Array(limit * itemSize);
+      metrics.scratchPeak = Math.max(metrics.scratchPeak, packing.byteLength);
+      const bytes = packing.subarray(0, count * itemSize);
+      for (let copied = 0; copied < count;) {
+        index = start + copied; offset = 0;
+        for (let i = shape.length - 1; i >= 0; i--) {
+          offset += index % shape[i] * strides[i]; index = Math.floor(index / shape[i]);
+        }
+        const n = Math.min(count - copied, run - (start + copied) % run);
+        bytes.set(new Uint8Array(toArrayBuffer(pointer + offset * itemSize, 0, n * itemSize)), copied * itemSize);
+        copied += n;
+      }
+      metrics.copiedBytes += bytes.length;
+      yield bytes;
+    }
+    start += count;
+  }
+}
+
+function writeHeader(path, header) {
+  const json = new TextEncoder().encode(JSON.stringify(header));
+  const prefix = new Uint8Array(PREFIX_LEN + json.length);
+  prefix.set(new TextEncoder().encode(MAGIC));
+  const view = new DataView(prefix.buffer);
+  view.setUint32(MAGIC.length, json.length, true);
+  view.setUint32(MAGIC.length + 4, prefix.length, true);
+  view.setBigUint64(MAGIC.length + 8, BigInt(Bun.hash(json)), true);
+  prefix.set(json, PREFIX_LEN);
+  const tmp = path + ".tmp", fd = openSync(tmp, "w");
+  try { writeAll(fd, prefix, 0); fsyncSync(fd); }
+  finally { closeSync(fd); }
+  renameSync(tmp, path);
+  return prefix.length;
+}
+
+function writeBlocks({ path, header, tensors, blockBytes = 1024 * 1024, segmented = true }) {
+  const metrics = { writtenBytes: 0, reusedBytes: 0, copiedBytes: 0, scratchPeak: 0 };
+  const root = join(dirname(path), "blocks");
+  mkdirSync(root, { recursive: true });
+  const slots = [...header.caches, ...(header.attachments ?? [])].flatMap(entry => entry.tensors);
+  for (let i = 0; i < slots.length; i++) {
+    const slot = slots[i]; slot.blocks = [];
+    let bytes = 0;
+    let source = tensors[i];
+    let packed;
+    if (!segmented) {
+      packed = bytesOf(source);
+      metrics.copiedBytes += packed.buffer === scratch.buffer ? packed.length : 0;
+      metrics.scratchPeak = Math.max(metrics.scratchPeak, scratch.length);
+      // Contiguous strides, but retain row boundaries for identical blocks.
+      let stride = 1;
+      const strides = Array(source.shape.length);
+      for (let j = strides.length - 1; j >= 0; j--) { strides[j] = stride; stride *= source.shape[j]; }
+      // No second copy; the pointer remains live until this tensor finishes.
+      source = { ...source, pointer: packed.length ? Number(ptr(packed)) : 0, strides };
+    }
+    for (const chunk of segments(source, blockBytes, metrics)) {
+      const hash = createHash("sha256").update(chunk).digest("hex");
+      bytes += chunk.length;
+      slot.blocks.push({ hash, bytes: chunk.length });
+      const destination = join(root, hash);
+      if (existsSync(destination)) { metrics.reusedBytes += chunk.length; continue; }
+      const tmp = destination + ".tmp", fd = openSync(tmp, "w");
+      try { writeAll(fd, chunk, 0); fsyncSync(fd); }
+      finally { closeSync(fd); }
+      renameSync(tmp, destination);
+      metrics.writtenBytes += chunk.length;
+    }
+    if (bytes !== slot.bytes) throw new Error("KV tensor byte-length drift");
+    // The header authenticates block order; each block authenticates bytes.
+    // Avoid hashing the same tensor a second time.
+    slot.hash = "";
+  }
+  header.formatVersion = 5;
+  metrics.writtenBytes += writeHeader(path, header);
+  return metrics;
+}
+
+function readAll(fd, bytes, position) {
+  let offset = 0;
+  while (offset < bytes.length) {
+    const n = readSync(fd, bytes, offset, bytes.length - offset, position + offset);
+    if (!n) throw new Error("truncated KV tensor");
+    offset += n;
+  }
+}
+
+// Destinations are unpublished, independently owned host allocations. No
+// MLX API is called here, and the owner retains them until acknowledgement.
+function readSnapshot({ path, header, pointers, verify }) {
+  const slots = [...header.caches, ...(header.attachments ?? [])].flatMap(entry => entry.tensors);
+  const fd = header.formatVersion === 5 ? null : openSync(path, "r");
+  try {
+    for (let i = 0; i < slots.length; i++) {
+      const slot = slots[i];
+      const bytes = slot.bytes ? new Uint8Array(toArrayBuffer(pointers[i], 0, slot.bytes)) : new Uint8Array(0);
+      if (header.formatVersion === 5) {
+        let offset = 0;
+        for (const block of slot.blocks ?? []) {
+          if (!/^[a-f0-9]{64}$/.test(block.hash) || block.bytes <= 0 || offset + block.bytes > bytes.length)
+            throw new Error("invalid KV block reference");
+          const blockFd = openSync(join(dirname(path), "blocks", block.hash), "r");
+          const target = bytes.subarray(offset, offset + block.bytes);
+          try { readAll(blockFd, target, 0); } finally { closeSync(blockFd); }
+          if (verify && createHash("sha256").update(target).digest("hex") !== block.hash)
+            throw new Error("KV block hash mismatch");
+          offset += block.bytes;
+        }
+        if (offset !== bytes.length) throw new Error("incomplete KV tensor blocks");
+      } else {
+        readAll(fd, bytes, header.dataStart + slot.off);
+        if (verify && hash64(bytes) !== slot.hash) throw new Error("KV tensor hash mismatch");
+      }
+    }
+  } finally { if (fd !== null) closeSync(fd); }
+}
+
+// Serialized with writes: an in-flight manifest cannot lose its blocks.
+// Only unreferenced content is collected; shared prefixes remain on SSD.
+function collectBlocks({ root }) {
+  if (!existsSync(root)) return;
+  for (const ns of readdirSync(root, { withFileTypes: true })) {
+    if (!ns.isDirectory()) continue;
+    const dir = join(root, ns.name), blocks = join(dir, "blocks");
+    if (!existsSync(blocks)) continue;
+    const live = new Set();
+    for (const file of readdirSync(dir)) {
+      if (!file.endsWith(".mlxkv")) continue;
+      const fd = openSync(join(dir, file), "r");
+      try {
+        const prefix = new Uint8Array(PREFIX_LEN); readAll(fd, prefix, 0);
+        const view = new DataView(prefix.buffer), json = new Uint8Array(view.getUint32(MAGIC.length, true));
+        readAll(fd, json, PREFIX_LEN);
+        if (BigInt(Bun.hash(json)) !== view.getBigUint64(MAGIC.length + 8, true)) continue;
+        const header = JSON.parse(new TextDecoder().decode(json));
+        for (const entry of [...header.caches, ...(header.attachments ?? [])])
+          for (const slot of entry.tensors) for (const block of slot.blocks ?? []) live.add(block.hash);
+      } finally { closeSync(fd); }
+    }
+    for (const file of readdirSync(blocks))
+      if (!live.has(file)) rmSync(join(blocks, file), { force: true });
+  }
+}
+
+parentPort.on("message", data => {
+  try {
+    const started = performance.now();
+    const metrics = data.operation === "read" ? readSnapshot(data)
+      : data.operation === "collect" ? collectBlocks(data)
+      : data.layout === "blocks" ? writeBlocks(data) : writeSnapshot(data);
+    parentPort.postMessage({ id: data.id, metrics: { ...metrics, elapsedMs: performance.now() - started } });
+  } catch (error) {
+    parentPort.postMessage({ id: data.id, error: String(error) });
+  }
+});

package/src/storage/retention-policy.ts

@@ -0,0 +1,37 @@
+/** RAM policy sees reuse metadata, never tensors, storage or scheduling. */
+export interface RetentionCandidate {
+  bytes: number;
+  cost: number;
+  lastUsed: number;
+  uses: number;
+  priority: number;
+}
+export interface RetentionPolicy {
+  readonly name: string;
+  accessed(entry: RetentionCandidate): void;
+  victim<T extends RetentionCandidate>(entries: readonly T[]): T;
+  evicted(entry: RetentionCandidate): void;
+}
+export class LruRetention implements RetentionPolicy {
+  readonly name = "lru";
+  accessed(_entry: RetentionCandidate): void {}
+  victim<T extends RetentionCandidate>(entries: readonly T[]): T {
+    return entries.reduce((a, b) => a.lastUsed < b.lastUsed ? a : b);
+  }
+  evicted(_entry: RetentionCandidate): void {}
+}
+
+/** GreedyDual size/frequency: saved prefill tokens are the cost proxy.
+ * Advancing age on eviction lets old popularity decay without wall timers. */
+export class CostSizeRetention implements RetentionPolicy {
+  readonly name = "cost-size";
+  #age = 0;
+  accessed(entry: RetentionCandidate): void {
+    entry.priority = this.#age + entry.cost * entry.uses / Math.max(1, entry.bytes);
+  }
+  victim<T extends RetentionCandidate>(entries: readonly T[]): T {
+    return entries.reduce((a, b) => a.priority < b.priority ||
+      (a.priority === b.priority && a.lastUsed < b.lastUsed) ? a : b);
+  }
+  evicted(entry: RetentionCandidate): void { this.#age = Math.max(this.#age, entry.priority); }
+}

package/src/tiered-prompt-cache.ts

@@ -0,0 +1,45 @@
+import { PromptCache, cacheBytes, type ColdTier } from "./prompt-cache";
+import { cloneKvCaches, SpillQueue } from "./kv-store";
+import { SsdDurabilityCoordinator } from "./ssd-durability";
+import { disposeResources } from "./engine/resources";
+import type { SsdCacheStore } from "./ssd-cache";
+
+/** One cache with RAM and SSD storage. Inference sees the same put/take
+ * interface. Recency and capacity choose RAM residents; a pending write is
+ * never treated as a completed demotion. */
+export class TieredPromptCache extends PromptCache {
+  protected override promoteRestores = true;
+  readonly spillQueue: SpillQueue;
+  readonly durability: SsdDurabilityCoordinator;
+
+  constructor(
+    maxBytes: number,
+    storage: Pick<SsdCacheStore, "storeAsync" | "hasDurablePrefix">,
+    cold: ColdTier,
+    clone = cloneKvCaches,
+    readonly writeBehind = true,
+  ) {
+    super(maxBytes, null, cold, clone);
+    this.spillQueue = new SpillQueue(Infinity, cacheBytes, async item => {
+      return storage.hasDurablePrefix(item.tokens, item.ns) ||
+        await storage.storeAsync(item.tokens, item.caches, item.ns, undefined, item.attachments);
+    }, disposeResources, () => false);
+    this.durability = new SsdDurabilityCoordinator(this, this.spillQueue, clone,
+      (tokens, ns) => storage.hasDurablePrefix(tokens, ns), 0, 5_000, () => {
+        // Writing alone does not remove a RAM entry. Apply capacity and
+        // pressure policy after the completed writer releases its views.
+        this.evictToBudget();
+        this.reclaim();
+      });
+    this.canEvict = entry => {
+      if (storage.hasDurablePrefix(entry.tokens, entry.ns)) return true;
+      this.durability.schedule(entry.tokens, entry.ns, false);
+      return false;
+    };
+  }
+
+  override put(...args: Parameters<PromptCache["put"]>): void {
+    super.put(...args);
+    if (this.writeBehind) this.durability.schedule(args[0], args[2]);
+  }
+}

package/src/vision/embedder.ts

@@ -7,6 +7,7 @@
 //
 // Sidecar weights stay bf16; features come out bf16 like text embeddings.
 
+import { tensorFingerprint } from "../model/fingerprint";
 import { ptr, read } from "bun:ffi";
 import { MlxArray, cpuStream } from "../mlx/array";
 import { C } from "../mlx/ffi";
@@ -18,6 +19,13 @@ const cstr = (s: string) => Buffer.from(s + "\0", "utf8");
 
 export class VisionTower {
   #weights = new Map<string, MlxArray>();
+  #cacheIdentity?: string;
+  /** Immutable encoder weights and computation settings, hashed once on use. */
+  get cacheIdentity(): string {
+    return this.#cacheIdentity ??= new Bun.CryptoHasher("sha256")
+      .update("gemma-vision-v1").update(JSON.stringify([this.embedScale, this.rmsNormEps]))
+      .update(tensorFingerprint(this.#weights)).digest("hex");
+  }
   readonly embedScale: number;
   readonly rmsNormEps: number;