mlx-bun 0.3.0 → 0.5.0

package/src/spec/two-model.ts

@@ -16,13 +16,15 @@
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
 import { clearCache } from "../mlx/ffi";
-import { flagOn } from "../runtime-config";
+import { runtimeFlag } from "../runtime-config";
 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";
+import type { DraftProvider, DraftSource, GroupedDraftProvider } from "./source";
+import { standaloneDraftGroups, openStandaloneDraftRows } from "./two-model-rows";
+import { bindLegacyAutoregressiveModel } from "../backends/mlx/autoregressive";
+import { artifactIdentity } from "../model/artifact-identity";
+import { configFingerprint } from "../model/fingerprint";
 
 /** 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
@@ -32,13 +34,16 @@ const PREFILL_CHUNK = 2048;
 export class TwoModelProvider implements DraftProvider {
   readonly id: string;
   readonly weightsBytes: number;
+  readonly grouped: GroupedDraftProvider;
 
   private constructor(
     readonly model: RuntimeModel,
     private readonly weights: Weights,
     id: string,
+    private readonly namespace: string,
   ) {
     this.id = id;
+    this.grouped = standaloneDraftGroups(bindLegacyAutoregressiveModel(model), namespace);
     this.weightsBytes = [...weights.shards.files.values()].reduce(
       (a, f) => a + f.mmap.size,
       0,
@@ -51,20 +56,25 @@ export class TwoModelProvider implements DraftProvider {
   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)!);
+    try {
+      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)`,
+        );
+      const identity = await artifactIdentity(configFingerprint(config),
+        [...weights.shards.files].map(([name, shard]) => ({ name, path: shard.path })));
+      return new TwoModelProvider(model, weights, modelDir.split("/").filter(Boolean).at(-1)!,
+        `standalone-draft-v1:${identity}`);
+    } catch (error) { weights.dispose(); throw error; }
   }
 
   open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
-    return new TwoModelSource(this.model, opts.sampler);
+    return new TwoModelSource(this.model, opts.sampler, this.namespace);
   }
 
   dispose(): void {
@@ -72,100 +82,32 @@ export class TwoModelProvider implements DraftProvider {
   }
 }
 
+/** Compatibility adapter: the existing request loop also uses the shared
+ * draft graph at B=1. Proposal generation and state recovery have one owner. */
 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();
+  readonly weightsBytes = 0;
+  private readonly rows: ReturnType<typeof openStandaloneDraftRows>;
+  private processed = 0;
+  constructor(model: RuntimeModel, sampler: (lp: MlxArray, step: number) => MlxArray, namespace: string) {
+    this.rows = openStandaloneDraftRows(bindLegacyAutoregressiveModel(model), namespace,
+      { sample: (logprobs, steps) => sampler(logprobs, steps[0]!) }, [null]);
   }
-
-  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);
+  async prefill(promptIds: number[]): Promise<void> {
+    const tailSplit = runtimeFlag("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) {
+    for (let at = this.processed; 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];
+      using ids = ops.fromInt32(promptIds.slice(at, end), [1, end - at]);
+      await this.rows.prefill(ids); clearCache();
     }
-    return drafts;
+    this.processed = upTo;
   }
-
-  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();
+  async draft(feed: number[], count: number, step: number): Promise<number[]> {
+    return (await this.rows.draft([feed.at(-1)!], count, [step]))[0]!;
   }
-
-  dispose(): void {
-    for (const c of this.caches) c.dispose();
+  async commit(_drafted: number, accepted: number): Promise<void> {
+    await this.rows.commit([accepted]);
+    this.processed += accepted + 1;
   }
+  dispose(): void { this.rows.dispose(); }
 }

package/src/ssd-cache.ts

@@ -1,3 +1,4 @@
+import type { CheckpointAttachment } from "./backends/mlx/checkpoint-state";
 // SSD cold tier for the prompt/KV cache (docs/design/kv-cache.md).
 //
 // Files ARE the database: no sidecar index. Layout
@@ -17,17 +18,20 @@
 import {
   existsSync, mkdirSync, readdirSync, rmSync, statSync, utimesSync,
 } from "node:fs";
-import { join } from "node:path";
+import { kvWriter, type KvWriteRequest } from "./storage/kv-writer";
+import { join, dirname } from "node:path";
 import { randomUUID } from "node:crypto";
 import { commonPrefixLength } from "./prompt-cache";
 import {
-  saveKvCache, saveKvCacheAsync, loadKvCache, readKvHeader,
-  cacheHeadersTrimmable, legacyCacheCodecs, type CacheCodecProvider,
+  saveKvCache, saveKvCacheAsync, loadKvCache, loadKvCacheAsync, readKvHeader,
+  cacheHeadersTrimmable, cacheHeadersMinimumReusableOffset, legacyCacheCodecs, type CacheCodecProvider,
   type KvSaveMeta, type KvLoadExpect, type LoadedKvCache,
 } from "./kv-store";
 import type { Cache } from "./model/gemma4-base";
+import { minimumReusableOffset } from "./backends/mlx/state-views";
 
 export interface SsdIndexEntry {
+  blocks?: Array<{ hash: string; bytes: number }>;
   path: string;
   ns: string;
   tokens: number[];
@@ -39,10 +43,12 @@ export interface SsdIndexEntry {
    *  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;
@@ -62,6 +68,27 @@ export class SsdCacheStore {
   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 };
 
@@ -75,8 +102,16 @@ export class SsdCacheStore {
     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 {
-    return this.#index.reduce((a, e) => a + e.bytes, 0);
+    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 {
@@ -92,7 +127,7 @@ export class SsdCacheStore {
    * 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) return false;
+      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;
     });
@@ -102,7 +137,7 @@ export class SsdCacheStore {
    *  headers and metadata mismatches are unlinked; `.tmp` orphans reaped.
    *  Only OUR fingerprint dir is touched. Returns entries indexed. */
   scan(): number {
-    this.#index = [];
+    this.#index = []; this.#exact.clear();
     if (!existsSync(this.#root)) return 0;
     for (const nsDir of readdirSync(this.#root)) {
       const nsPath = join(this.#root, nsDir);
@@ -127,9 +162,11 @@ export class SsdCacheStore {
             throw new Error("metadata mismatch");
           if ((h.codecProvider ?? legacyCacheCodecs.id) !== this.#codecs.id) continue;
           const st = statSync(path);
-          this.#index.push({
+          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: cacheHeadersTrimmable(h.caches, this.#codecs),
+            trimmable: !h.attachments?.length && cacheHeadersTrimmable(h.caches, this.#codecs),
+            minimumReusableOffset: cacheHeadersMinimumReusableOffset(h.caches),
             ...(h.generationCheckpoint
               ? { generationCheckpoint: h.generationCheckpoint }
               : {}),
@@ -139,6 +176,7 @@ export class SsdCacheStore {
         }
       }
     }
+    void kvWriter.collect(this.#root).catch(() => {});
     return this.#index.length;
   }
 
@@ -157,7 +195,7 @@ export class SsdCacheStore {
     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 (p === 0 || p <= bestLen || p < (e.minimumReusableOffset ?? 0)) continue;
       if (e.tokens.length - p > 0 && !e.trimmable) continue;
       bestLen = p;
       best = e;
@@ -194,6 +232,23 @@ export class SsdCacheStore {
     }
   }
 
+  /** 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 kv-store's tmp+fsync+rename. An
    *  entry bigger than the cap is refused; disk/write failure is a warn-once
@@ -201,13 +256,13 @@ export class SsdCacheStore {
    *  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 {
+  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), this.#codecs);
-      return this.#indexStored(path, tokens, caches, ns);
+      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);
     }
@@ -218,19 +273,18 @@ export class SsdCacheStore {
    *  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` owns every per-tensor step (see saveKvCacheAsync) — the
-   *  server passes the gateway's exclusive runner so a request cannot start
-   *  between an idle check and the blocking MLX readback. */
+   *  `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 {
-      mkdirSync(dir, { recursive: true });
-      await saveKvCacheAsync(path, tokens, caches, this.#meta(ns), runStep, this.#codecs);
-      return this.#indexStored(path, tokens, caches, ns);
+      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);
     }
@@ -260,25 +314,30 @@ export class SsdCacheStore {
    *  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 {
-      mkdirSync(dir, { recursive: true });
+      if (runStep) await runStep(() => mkdirSync(dir, { recursive: true }));
+      else mkdirSync(dir, { recursive: true });
       await saveKvCacheAsync(path, tokens, caches, {
         ...this.#meta(ns), generationCheckpoint: checkpoint,
-      }, undefined, this.#codecs);
-      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);
+      }, 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 stored;
+      };
+      return runStep ? await runStep(index) : index();
     } catch (err) {
-      return this.#storeFailed(path, err);
+      return runStep ? await runStep(() => this.#storeFailed(path, err)) : this.#storeFailed(path, err);
     }
   }
 
@@ -299,14 +358,19 @@ export class SsdCacheStore {
   #indexStored(
     path: string, tokens: number[], caches: Cache[], ns: string,
     generationCheckpoint?: NonNullable<KvSaveMeta["generationCheckpoint"]>,
+    attachments?: CheckpointAttachment[],
   ): 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`);
+    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.#opts.maxBytes) {
+      this.remove(path);
+      console.warn(`[ssd-cache] entry not stored: ${entryBytes} bytes exceeds the ${this.#opts.maxBytes}-byte cap`);
       return false;
     }
-    const trimmable = caches.every((c) => c.isTrimmable());
+    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]) {
@@ -321,12 +385,13 @@ export class SsdCacheStore {
     // unrestorable one landed.
     if (trimmable) {
       for (const e of [...this.#index]) {
-        if (e.ns !== ns || e.tokens.length >= tokens.length) continue;
+        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.#index.push({
-      path, ns, tokens, bytes: st.size, mtimeMs: Date.now(), trimmable,
+    this.#addIndex({
+      blocks, path, ns, tokens, bytes: st.size, mtimeMs: Date.now(), trimmable,
+      minimumReusableOffset: minimumReusableOffset(caches),
       ...(generationCheckpoint ? { generationCheckpoint } : {}),
     });
     this.stats.spills++;
@@ -335,7 +400,8 @@ export class SsdCacheStore {
   }
 
   #storeFailed(path: string, err: unknown): boolean {
-    try { rmSync(path, { force: true }); rmSync(`${path}.tmp`, { force: true }); } catch {}
+    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}`);
@@ -355,6 +421,11 @@ export class SsdCacheStore {
 
   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/ssd-durability.ts

@@ -1,12 +1,9 @@
+import { cloneAttachments } from "./backends/mlx/checkpoint-state";
+import { cleanupFailure, disposeResources } from "./engine/resources";
 import type { Cache } from "./model/gemma4";
 import type { PromptCacheEntry } from "./prompt-cache";
 import type { SpillItem, SpillQueue } from "./kv-store";
 
-export interface DurabilityGateway {
-  readonly busy: boolean;
-  runExclusive<T>(fn: () => Promise<T>): Promise<T>;
-}
-
 export interface DurabilityPromptCache {
   findExact(tokens: number[], ns?: string): PromptCacheEntry | null;
 }
@@ -48,21 +45,23 @@ export class SsdDurabilityCoordinator {
   #flush: Promise<DurabilityFlushResult> | null = null;
 
   constructor(
-    readonly gateway: DurabilityGateway,
     readonly promptCache: DurabilityPromptCache,
     readonly spillQueue: SpillQueue,
     readonly cloneCaches: (caches: Cache[]) => Cache[],
     readonly isAlreadyDurable: (tokens: number[], ns: string) => boolean = () => false,
-    readonly debounceMs = 1_000,
-    readonly busyRetryMs = 5_000,
+    readonly debounceMs = 0,
+    readonly retryMs = 5_000,
+    readonly onStored: () => void = () => {},
   ) {}
 
-  schedule(tokens: number[], ns = ""): void {
-    if (tokens.length === 0) return;
+  schedule(tokens: number[], ns = "", replace = true): void {
+    // Exact tensor objects use an empty history and still need persistence.
+    if (tokens.length === 0 && !this.promptCache.findExact(tokens, ns)?.attachments?.length) return;
     // PromptCache can hold unrelated entries with the same namespace and
     // length. Include the tokens so one conversation cannot cancel another
     // conversation's pending durability record.
     const key = `${ns.length}:${ns}:${tokens.join(",")}`;
+    if (!replace && this.#dirty.has(key)) return;
     const rec: DirtySnapshot = { key, tokens: [...tokens], ns };
     this.#dirty.set(key, rec);
     this.#arm(key, this.debounceMs);
@@ -106,10 +105,6 @@ export class SsdDurabilityCoordinator {
 
     const rec = this.#dirty.get(key);
     if (!rec) return "stored";
-    if (!force && this.gateway.busy) {
-      this.#arm(key, this.busyRetryMs);
-      return "failed";
-    }
 
     const task = this.#store(rec);
     this.#attempts.set(key, task);
@@ -118,7 +113,7 @@ export class SsdDurabilityCoordinator {
       if (outcome === "stored" && this.#dirty.get(key) === rec)
         this.#dirty.delete(key);
       if (outcome !== "stored" && !force && this.#dirty.get(key) === rec)
-        this.#arm(key, this.busyRetryMs);
+        this.#arm(key, this.retryMs);
       return outcome;
     } finally {
       if (this.#attempts.get(key) === task) this.#attempts.delete(key);
@@ -128,22 +123,28 @@ export class SsdDurabilityCoordinator {
   async #store(rec: DirtySnapshot): Promise<StoreOutcome> {
     let snap: SpillItem | null = null;
     try {
-      snap = await this.gateway.runExclusive(async () => {
+      snap = (() => {
         if (this.isAlreadyDurable(rec.tokens, rec.ns)) return null;
         const entry = this.promptCache.findExact(rec.tokens, rec.ns);
         if (!entry) return null;
-        return {
-          tokens: [...entry.tokens],
-          caches: this.cloneCaches(entry.caches),
-          ns: rec.ns,
-        };
-      });
+        const caches = this.cloneCaches(entry.caches);
+        try {
+          return { tokens: [...entry.tokens], caches, ns: rec.ns,
+            attachments: cloneAttachments(entry.attachments) };
+        } catch (error) {
+          return cleanupFailure(error, () => disposeResources(caches));
+        }
+      })();
     } catch {
       return "failed";
     }
     if (!snap)
       return this.isAlreadyDurable(rec.tokens, rec.ns) ? "stored" : "missing";
-    return await this.spillQueue.enqueue(snap) ? "stored" : "failed";
+    const stored = await this.spillQueue.enqueue(snap);
+    // The queue releases its views before this continuation runs. Residency
+    // policy now sees memory without the completed writer's ownership.
+    if (stored) this.onStored();
+    return stored ? "stored" : "failed";
   }
 
   async #flushInner(started: number): Promise<DurabilityFlushResult> {
@@ -155,12 +156,13 @@ export class SsdDurabilityCoordinator {
     const droppedBefore = this.spillQueue.droppedCount;
     const failedBefore = this.spillQueue.failedCount;
 
-    if (this.#attempts.size > 0)
-      await Promise.allSettled([...this.#attempts.values()]);
+    if (this.#attempts.size > 0) {
+      const results = await Promise.allSettled([...this.#attempts.values()]);
+      flushedSnapshots += results.filter(result => result.status === "fulfilled" && result.value === "stored").length;
+    }
     await this.spillQueue.drain();
 
-    // Flush one entry at a time. This prevents the queue cap from dropping a
-    // boundary snapshot while a large final snapshot is already in flight.
+    // Flush dirty entries serially without a second queue of snapshots.
     // Attempt each record version once. Missing state stays dirty, so a
     // second flush cannot report success without a real durable snapshot.
     const attempted = new Set<DirtySnapshot>();

package/src/storage/host-buffer.ts

@@ -0,0 +1,24 @@
+import { dlopen, ptr, read, type Pointer } from "bun:ffi";
+
+// Page alignment permits Metal to use the restored allocation directly.
+// The destructor is native free, so final release may occur on any thread.
+const libc = dlopen("/usr/lib/libSystem.B.dylib", {
+  posix_memalign: { args: ["ptr", "u64", "u64"], returns: "i32" },
+  free: { args: ["ptr"], returns: "void" },
+});
+export const hostBufferDestructor = Number(Reflect.get(libc.symbols.free, "ptr"));
+
+/** Unpublished CPU storage. A reader fills it before ownership moves to MLX. */
+export class HostBuffer {
+  #pointer: number;
+  constructor(readonly bytes: number) {
+    const slot = new BigUint64Array(1), address = ptr(slot);
+    const status = libc.symbols.posix_memalign(address, 16384n, BigInt(Math.max(16384, Math.ceil(bytes / 16384) * 16384)));
+    if (status) throw new Error(`KV restore allocation failed (${status})`);
+    this.#pointer = Number(read.u64(address));
+  }
+  get pointer(): number { return this.#pointer; }
+  transfer(): number { const p = this.#pointer; this.#pointer = 0; return p; }
+  dispose(): void { if (this.#pointer) libc.symbols.free(this.transfer() as Pointer); }
+  [Symbol.dispose](): void { this.dispose(); }
+}