mlx-bun 0.3.0 → 0.5.0

package/src/prompt-cache.ts

@@ -1,4 +1,5 @@
-import { leaseCacheState } from "./backends/mlx/state-views";
+import { type CheckpointAttachment, materializeCheckpoint, attachmentBytes, cloneAttachments, disposeAttachments } from "./backends/mlx/checkpoint-state";
+import { leaseCacheState, minimumReusableOffset } from "./backends/mlx/state-views";
 import { withResource, cleanupFailure, disposeResources, ownResource } from "./engine/resources";
 // Byte-capped LRU prompt cache — the RAM tier of the Layer-0 KV store.
 //
@@ -10,24 +11,22 @@ import { withResource, cleanupFailure, disposeResources, ownResource } from "./e
 //   const hit = cache.take(promptIds);   // longest-prefix match
 //   generate(model, promptIds, { cache: hit?.caches ?? fresh })
 //   cache.put([...promptIds, ...generated], caches, ns, hit?.retain);
-// take() is NON-CONSUMING for trimmable caches: it serves zero-copy clones
-// and leaves the donor in place, so N agents can share one prefill.
-// Qwen hybrid entries containing recurrent SSM state transfer ownership on
-// an exact hit, avoiding an extra resident recurrent snapshot. put() supersedes
-// same-ns prefix-ancestors (when the new entry is trimmable) so a conversation
-// stays one entry, not one per turn.
+// take() lends zero-copy views and leaves every RAM donor in place, including
+// recurrent and method state. SSD persistence runs independently of RAM reuse.
+// put() supersedes same-ns prefix-ancestors when the new entry is trimmable.
 //
 // TIERING (Layer 0, unified-engine plan): when a ColdTier is attached,
 // take() itself runs the two-tier dance — RAM peek vs cold find, restore
-// (zero-copy mmap) + trim when the cold tier holds a strictly longer
+// (streamed copy) + trim when the cold tier holds a strictly longer
 // prefix — so EVERY consumer (the serial lane, the batch scheduler, future
 // prefix sharing) gets SSD restores through the same take()/put() it
 // already calls. Eviction spills to the tier (the #spill hook); idle
-// entries DEMOTE to it (demoteIdle — free the GPU memory, keep the prefix
-// reachable); onPut lets the server schedule its debounced write-behind
-// snapshot for both lanes.
+// entries demote to it when policy permits. TieredPromptCache owns queued
+// persistence and promotes restored prefixes into RAM.
 
 import type { Cache } from "./model/gemma4";
+import type { PrefixCache, PrefixCacheHit } from "./contracts/prefix-cache";
+import { LruRetention, type RetentionPolicy, type RetentionCandidate } from "./storage/retention-policy";
 import { cloneKvCaches } from "./kv-store";
 
 /** Reference-counted release: wraps an entry's `retain` (e.g. an mmap
@@ -65,13 +64,17 @@ export interface ColdTier {
   /** Longest stored usable prefix for prompt/ns — index-only, no I/O.
    *  `handle` is the tier's opaque entry token, passed back to restore. */
   find(prompt: number[], ns: string): { prefixLen: number; handle: unknown } | null;
-  /** Materialize a found entry as GPU-visible caches (zero-copy COW mmap;
-   *  pages fault in lazily). `retain` must run after the caches are
-   *  disposed (it unmaps the backing file). Null on any failure. */
-  restore(handle: unknown): { tokens: number[]; caches: Cache[]; retain: () => void } | null;
+  /** Restore GPU-visible state. The default SSD backend copies into owned
+   *  arrays. An optional backing release runs after consumers dispose their
+   *  arrays. Null on any failure. */
+  restore(handle: unknown): { tokens: number[]; caches: Cache[]; attachments?: CheckpointAttachment[]; retain: () => void } | null;
+  findExact?(tokens: number[], ns: string): { prefixLen: number; handle: unknown } | null;
+  restoreAsync?(handle: unknown): Promise<ReturnType<ColdTier["restore"]>>;
+  /** Tensor objects have no target layers; restore with their own schema. */
+  restoreObjectAsync?(handle: unknown): Promise<ReturnType<ColdTier["restore"]>>;
   /** Borrow live caches and finish persistence before returning. A false
-   *  result prevents a required demotion or recurrent ownership transfer. */
-  store(tokens: number[], caches: Cache[], ns: string): boolean | void;
+   *  result prevents a required demotion. */
+  store(tokens: number[], caches: Cache[], ns: string, attachments?: CheckpointAttachment[]): boolean | void;
 }
 
 /** Spill sink for eviction/demotion (write-behind, 2026-07-06).
@@ -95,9 +98,7 @@ export interface SpillSink {
   spillSync?: (entry: PromptCacheEntry) => void;
 }
 
-export interface PromptCacheEntry {
-  tokens: number[];
-  caches: Cache[];
+export interface PromptCacheEntry extends PrefixCacheHit<Cache[], CheckpointAttachment[]> {
   /** Namespace key — adapter spec for LoRA requests ("" = base model).
    *  KV computed under one adapter must never seed another's prefill. */
   ns: string;
@@ -124,7 +125,7 @@ export function commonPrefixLength(a: number[], b: number[]): number {
   return i;
 }
 
-interface EntryRecord {
+interface EntryRecord extends RetentionCandidate {
   entry: PromptCacheEntry;
   bytes: number;
   lastUsed: number;
@@ -132,12 +133,98 @@ interface EntryRecord {
   /** Ref-counted retain: the donor holds one share (entry.retain); every
    *  clone handed out by take() holds another. */
   share: { acquire(): () => void };
+  resident: boolean;
+  sessions: number;
 }
 
-export class PromptCache {
+export class PromptCache implements PrefixCache<Cache[], CheckpointAttachment[]> {
+  /** Exact tensor objects share this store's byte budget, eviction and writer.
+   * Empty token history denotes no conversation; the private namespace is the
+   * complete object identity. No fabricated language tokens enter the index. */
+  readonly #objects: import("./contracts/object-cache").ObjectCache<CheckpointAttachment[]> = {
+    take: key => this.#takeObject(`object:${key}`),
+    put: (key, value) => this.put([], [], `object:${key}`, undefined, value),
+  };
+
+  get objects(): import("./contracts/object-cache").ObjectCache<CheckpointAttachment[]> | undefined {
+    return this.maxBytes > 0 || this.#cold ? this.#objects : undefined;
+  }
+
+  async #takeObject(ns: string): Promise<{ value: CheckpointAttachment[]; dispose(): void } | null> {
+    let entry = this.findExact([], ns);
+    if (entry) this.objectHits++;
+    else this.objectMisses++;
+    if (!entry) {
+      const cold = this.#cold;
+      const hit = cold?.findExact?.([], ns);
+      if (!hit || !cold?.restoreObjectAsync) return null;
+      const epoch = this.#epoch;
+      const loaded = await cold.restoreObjectAsync(hit.handle);
+      if (!loaded) return null;
+      this.objectRestores++;
+      if (epoch !== this.#epoch) { this.#disposeEntry({ ...loaded, ns }, false); return null; }
+      // Another request may have published this immutable object during IO.
+      entry = this.findExact([], ns);
+      if (entry) this.#disposeEntry({ ...loaded, ns }, false);
+      else {
+        this.#publishingPrefetch = true;
+        try { this.put([], loaded.caches, ns, loaded.retain, loaded.attachments); }
+        catch (error) { this.#disposeEntry({ ...loaded, ns }, false); throw error; }
+        finally { this.#publishingPrefetch = false; }
+        entry = this.findExact([], ns);
+      }
+    }
+    if (!entry) return null;
+    const borrowed = this.#borrow(this.#records.get(entry)!, 0, ns);
+    try { this.evictToBudget(); this.reclaim(); }
+    catch (error) { this.#disposeEntry(borrowed, false); throw error; }
+    let released = false;
+    return { value: borrowed.attachments ?? [], dispose: () => {
+      if (released) return; released = true;
+      this.#disposeEntry(borrowed, false);
+    } };
+  }
+
   readonly maxBytes: number;
   #entries: EntryRecord[] = [];
   #clock = 0;
+  #sessions = new Map<string, Map<string, EntryRecord>>();
+  #records = new WeakMap<PromptCacheEntry, EntryRecord>();
+  sessionHits = 0;
+  sessionMisses = 0;
+  prefixScans = 0;
+  objectHits = 0;
+  objectMisses = 0;
+  objectRestores = 0;
+
+  closeSession(sessionId: string): void {
+    for (const record of this.#sessions.get(sessionId)?.values() ?? []) record.sessions--;
+    this.#sessions.delete(sessionId);
+  }
+
+  #associate(sessionId: string | undefined, ns: string, record: EntryRecord): void {
+    if (!sessionId) return;
+    let entries = this.#sessions.get(sessionId);
+    if (!entries) this.#sessions.set(sessionId, entries = new Map());
+    const old = entries.get(ns);
+    if (old === record) return;
+    if (old) old.sessions--;
+    entries.set(ns, record); record.sessions++;
+  }
+
+  #sessionRecord(prompt: number[], ns: string, sessionId?: string): EntryRecord | undefined {
+    const record = sessionId ? this.#sessions.get(sessionId)?.get(ns) : undefined;
+    // A session labels affinity. Full-history callers may edit or branch;
+    // incompatible history uses ordinary lookup, never a request refusal.
+    return record && record.entry.tokens.length < prompt.length &&
+      commonPrefixLength(record.entry.tokens, prompt) === record.entry.tokens.length ? record : undefined;
+  }
+
+  #coldHit(prompt: number[], ns: string, sessionId?: string) {
+    const known = this.#sessionRecord(prompt, ns, sessionId);
+    return (known && this.#cold?.findExact?.(known.entry.tokens, ns)) || this.#cold?.find(prompt, ns) || null;
+  }
+
   hits = 0;
   misses = 0;
   demotions = 0;
@@ -148,9 +235,98 @@ export class PromptCache {
   readonly #spillSync: ((entry: PromptCacheEntry) => void) | null;
   /** Optional cold tier: take() tiers over it, demoteIdle() spills into it. */
   readonly #cold: ColdTier | null;
-  /** Fired after every successful put() — the server hangs its debounced
-   *  write-behind SSD snapshot here so BOTH lanes' entries persist. */
+  /** Observer after publication, independent of persistence policy. */
   onPut: ((tokens: number[], ns: string) => void) | null = null;
+  /** The backend measures pressure; this store chooses which optional
+   * snapshots to release. No storage I/O runs at a reclaim boundary. */
+  pressure: { overBudget(): boolean; releasePending?(): void } | null = null;
+  /** Tier policy: a pending write is not a lower-tier copy. This controls
+   * eviction only; persistence completion does not itself remove an entry. */
+  canEvict: ((entry: PromptCacheEntry) => boolean) | null = null;
+  protected promoteRestores = false;
+  retention: RetentionPolicy = new LruRetention();
+  #prefetches = new Map<unknown, Promise<PromptCacheEntry | null>>();
+  #pins = new Map<PromptCacheEntry, number>();
+  #publishingPrefetch = false;
+  #epoch = 0;
+
+  /** Start cold reads during request preparation, coalescing identical IO.
+   * Active request interest keeps a just-restored donor resident until the
+   * caller has consumed it. Cancellation releases interest through the lease. */
+  async prefetch(prompt: number[], ns = "", sessionId?: string): Promise<() => void> {
+    const known = this.#sessionRecord(prompt, ns, sessionId);
+    if (known?.resident) return this.#pin(known.entry);
+    const cold = this.#cold;
+    const hit = cold?.restoreAsync ? this.#coldHit(prompt, ns, sessionId) : null;
+    if (!cold?.restoreAsync || !hit || (!known && hit.prefixLen <= this.peekPrefixLen(prompt, ns))) return () => {};
+    let pending = this.#prefetches.get(hit.handle);
+    if (!pending) {
+      const epoch = this.#epoch;
+      pending = (async () => {
+        const loaded = await cold.restoreAsync!(hit.handle);
+        if (!loaded) return null;
+        if (epoch !== this.#epoch) { this.#disposeEntry({ ...loaded, ns }, false); return null; }
+        this.#publishingPrefetch = true;
+        try { this.put(loaded.tokens, loaded.caches, ns, loaded.retain, loaded.attachments); }
+        catch (error) { this.#disposeEntry({ ...loaded, ns }, false); throw error; }
+        finally { this.#publishingPrefetch = false; }
+        return this.findExact(loaded.tokens, ns);
+      })();
+      this.#prefetches.set(hit.handle, pending);
+    }
+    let entry: PromptCacheEntry | null;
+    try { entry = await pending; }
+    finally { if (this.#prefetches.get(hit.handle) === pending) this.#prefetches.delete(hit.handle); }
+    if (!entry) return () => {};
+    this.#associate(sessionId, ns, this.#records.get(entry)!);
+    return this.#pin(entry);
+  }
+
+  #pin(entry: PromptCacheEntry): () => void {
+    this.#pins.set(entry, (this.#pins.get(entry) ?? 0) + 1);
+    let released = false;
+    return () => {
+      if (released) return; released = true;
+      const count = this.#pins.get(entry!) ?? 1;
+      if (count === 1) this.#pins.delete(entry!); else this.#pins.set(entry!, count - 1);
+      this.evictToBudget(); this.reclaim();
+    };
+  }
+
+  #victim(): EntryRecord | undefined {
+    const candidates = this.#entries.filter(record => !this.#pins.has(record.entry));
+    const unrelated = candidates.filter(record => record.sessions === 0);
+    return candidates.length ? this.retention.victim(unrelated.length ? unrelated : candidates) : undefined;
+  }
+
+  evictToBudget(): void {
+    if (this.#publishingPrefetch) return;
+    while (this.totalBytes > this.maxBytes) {
+      if (!this.#entries.length) break;
+      const oldest = this.#victim();
+      if (!oldest) break;
+      if (this.canEvict && !this.canEvict(oldest.entry)) break;
+      this.retention.evicted(oldest);
+      this.#entries = this.#entries.filter(r => r !== oldest);
+      try { this.#disposeEntry(oldest.entry, !this.canEvict); }
+      catch (error) { console.warn(`[prompt-cache] entry cleanup failed: ${error}`); }
+    }
+  }
+
+  reclaim(): void {
+    const pressure = this.pressure;
+    if (!pressure?.overBudget()) return;
+    pressure.releasePending?.();
+    while (this.#entries.length && pressure.overBudget()) {
+      const oldest = this.#victim();
+      if (!oldest) break;
+      if (this.canEvict && !this.canEvict(oldest.entry)) break;
+      this.retention.evicted(oldest);
+      this.#entries = this.#entries.filter(r => r !== oldest);
+      this.#disposeEntry(oldest.entry, false);
+      this.demotions++;
+    }
+  }
 
   /** Zero-copy view cloner — injectable so model-free tests can stub it. */
   readonly #clone: (caches: Cache[]) => Cache[];
@@ -161,6 +337,7 @@ export class PromptCache {
     spill: SpillSink | ((entry: PromptCacheEntry) => void) | null = null,
     cold: ColdTier | null = null,
     clone: (caches: Cache[]) => Cache[] = cloneKvCaches,
+    private readonly materialize: (caches: Cache[], attachments?: CheckpointAttachment[]) => void = materializeCheckpoint,
   ) {
     this.maxBytes = maxBytes;
     this.#spillOwned = typeof spill === "function" ? null : spill?.spillOwned ?? null;
@@ -170,6 +347,8 @@ export class PromptCache {
   }
 
   #disposeEntry(entry: PromptCacheEntry, spillFirst: boolean): void {
+    const record = this.#records.get(entry);
+    if (record) record.resident = false;
     if (spillFirst) {
       if (this.#spillOwned) {
         // NON-BLOCKING spill: clone BEFORE disposing (clones are zero-copy
@@ -185,23 +364,26 @@ export class PromptCache {
         // no backing mmap for retain to unmap, so nothing the clones
         // could alias goes away.
         let clones: Cache[] | null = null;
+        let attachments: CheckpointAttachment[] | undefined;
         try {
           clones = this.#clone(entry.caches);
-          this.#spillOwned({ tokens: [...entry.tokens], caches: clones, ns: entry.ns });
+          attachments = cloneAttachments(entry.attachments);
+          this.#spillOwned({ tokens: [...entry.tokens], caches: clones, ns: entry.ns, attachments });
         } catch {
           // Failure degrades to no-spill; eviction/demotion never unwinds.
           if (clones) for (const c of clones) try { c.dispose(); } catch {}
+          try { disposeAttachments(attachments); } catch {}
         }
       } else {
         const spill =
           this.#spillSync ??
           (this.#cold
-            ? (e: PromptCacheEntry) => this.#cold!.store(e.tokens, e.caches, e.ns)
+            ? (e: PromptCacheEntry) => this.#cold!.store(e.tokens, e.caches, e.ns, e.attachments)
             : null);
         if (spill) try { spill(entry); } catch {}
       }
     }
-    disposeResources(entry.caches);
+    disposeResources([...entry.caches, { dispose: () => disposeAttachments(entry.attachments) }]);
     entry.retain?.(); // backing release requires successful cache disposal
   }
 
@@ -220,16 +402,17 @@ export class PromptCache {
   relievePressure(overBudget: () => boolean): number {
     let count = 0;
     while (this.#entries.length && overBudget()) {
-      let oldest = 0;
-      for (let i = 1; i < this.#entries.length; i++)
-        if (this.#entries[i]!.lastUsed < this.#entries[oldest]!.lastUsed) oldest = i;
-      const { entry } = this.#entries[oldest]!;
+      const victim = this.#victim();
+      if (!victim) break;
+      const oldest = this.#entries.indexOf(victim);
+      const { entry } = victim;
       // Keep the entry resident if persistence throws. A failed SSD write
       // must not be presented as a successful demotion.
       if (this.#cold) {
-        if (this.#cold.store(entry.tokens, entry.caches, entry.ns) === false) break;
+        if (this.#cold.store(entry.tokens, entry.caches, entry.ns, entry.attachments) === false) break;
       }
       else this.#spillSync?.(entry);
+      this.retention.evicted(victim);
       this.#entries.splice(oldest, 1);
       this.#disposeEntry(entry, false);
       this.demotions++;
@@ -239,9 +422,9 @@ export class PromptCache {
   }
 
   /** Find the entry with the longest usable common prefix of `prompt`
-   *  across BOTH tiers. Trimmable RAM entries are served non-consumingly as
-   *  zero-copy clones; exact hits on Qwen hybrid entries containing recurrent
-   *  SSM state transfer the donor itself. The caller owns the returned caches —
+   *  across BOTH tiers. Every RAM entry lends zero-copy views and retains its
+   *  donor, including exact-boundary recurrent state and method attachments.
+   *  The caller owns the returned views —
    *  dispose them or put() an extended entry, honoring `retain` either way.
    *
    *  Usable prefix = common prefix capped at prompt.length - 1 (at least
@@ -250,21 +433,39 @@ export class PromptCache {
    *  wrapped — those entries only match in full.
    *
    *  Tier order: the cold tier wins only with a STRICTLY longer usable
-   *  prefix (its restore costs a mmap + lazy fault-in; RAM is free). A
+   *  prefix (restore reads storage and materializes native arrays). A
    *  cold entry with an untrimmable divergent tail (ring post-wrap, SSM)
    *  is dropped and the RAM candidate (or a fresh prefill) serves.
    *  hits/misses count RAM candidacy only (cold restores are counted by
    *  the tier itself), preserving the /stats meaning. */
-  take(prompt: number[], ns = ""): PromptCacheEntry | null {
+  take(prompt: number[], ns = "", sessionId?: string): PromptCacheEntry | null {
+    const hit = this.#take(prompt, ns, sessionId);
+    // Keep the chosen state's owned views before releasing optional donors.
+    // This also covers a cold restore materializing its bytes at admission.
+    try { this.reclaim(); return hit; }
+    catch (error) {
+      if (hit) this.#disposeEntry(hit, false);
+      throw error;
+    }
+  }
+
+  #take(prompt: number[], ns: string, sessionId?: string): PromptCacheEntry | null {
+    const known = this.#sessionRecord(prompt, ns, sessionId);
+    if (known?.resident) {
+      this.sessionHits++; this.hits++;
+      return this.#borrow(known, known.entry.tokens.length, ns);
+    }
+    if (sessionId) this.sessionMisses++;
+    this.prefixScans++;
     let bestIdx = -1;
     let bestLen = 0;
     for (let i = 0; i < this.#entries.length; i++) {
       const e = this.#entries[i]!.entry;
       if (e.ns !== ns) continue;
       const p = Math.min(commonPrefixLength(e.tokens, prompt), prompt.length - 1);
-      if (p <= bestLen) continue;
+      if (p <= bestLen || p < minimumReusableOffset(e.caches)) continue;
       const trimNeeded = e.tokens.length - p;
-      if (trimNeeded > 0 && !e.caches.every((c) => c.isTrimmable())) continue;
+      if (trimNeeded > 0 && (!!e.attachments?.length || !e.caches.every((c) => c.isTrimmable()))) continue;
       bestLen = p;
       bestIdx = i;
     }
@@ -275,7 +476,7 @@ export class PromptCache {
     // candidate stays put UNTOUCHED (the old serial-lane flow trimmed it
     // before comparing — a needless entry degradation, gone now).
     if (this.#cold) {
-      const hit = this.#cold.find(prompt, ns);
+      const hit = this.#coldHit(prompt, ns, sessionId);
       if (hit && hit.prefixLen > bestLen) {
         const loaded = this.#cold.restore(hit.handle);
         if (loaded) {
@@ -283,7 +484,8 @@ export class PromptCache {
             this.#disposeEntry({ ...entry, ns }, false));
           const result = withResource(owner, (entry) => {
             const trimNeeded = entry.tokens.length - hit.prefixLen;
-            if (trimNeeded > 0 && !entry.caches.every((cache) => cache.isTrimmable())) {
+            if (hit.prefixLen < minimumReusableOffset(entry.caches)) return null;
+            if (trimNeeded > 0 && (!!entry.attachments?.length || !entry.caches.every((cache) => cache.isTrimmable()))) {
               console.warn(
                 `[prompt-cache] cold entry unusable: ${entry.tokens.length} tokens, ` +
                 `needs trim ${trimNeeded} but untrimmable (wrapped ring/SSM) — re-prefilling`,
@@ -294,7 +496,27 @@ export class PromptCache {
             const tokens = entry.tokens.slice(0, hit.prefixLen);
             return { ...owner.transfer(), tokens, ns };
           });
-          if (result) return result;
+          if (result) {
+            if (this.promoteRestores) {
+              // Cache the restored prefix immediately. A second request can
+              // reuse RAM even before this borrower finishes its generation.
+              const share = makeSharedRetain(result.retain);
+              result.retain = share.acquire();
+              const retained = share.acquire();
+              let clones: Cache[] = [];
+              let attachments: CheckpointAttachment[] | undefined;
+              try {
+                clones = this.#clone(result.caches);
+                attachments = cloneAttachments(result.attachments);
+                this.put([...result.tokens], clones, ns, retained, attachments, sessionId);
+              } catch (error) {
+                cleanupFailure(error, () => disposeResources([...clones,
+                  { dispose: () => disposeAttachments(attachments) },
+                  { dispose: retained }, { dispose: () => this.#disposeEntry(result, false) }]));
+              }
+            }
+            return result;
+          }
         }
       }
     }
@@ -310,30 +532,30 @@ export class PromptCache {
     // consume-and-trim semantics CANNIBALIZED donors: agent B borrowing
     // agent A's 2k system prompt destroyed A's 10k entry to do it.
     const rec = this.#entries[bestIdx]!;
+    this.#associate(sessionId, ns, rec);
+    return this.#borrow(rec, bestLen, ns);
+  }
+
+  #borrow(rec: EntryRecord, bestLen: number, ns: string): PromptCacheEntry {
     rec.lastUsed = ++this.#clock;
     rec.lastUsedMs = Date.now();
-    // Qwen recurrent caches are untrimmable. Transfer the exact-boundary
-    // entry instead of retaining a second recurrent snapshot in RAM.
-    // This reduced memory pressure in the recorded multi-turn reproduction.
-    // Save the boundary before consumption: the extended recurrent state
-    // cannot be trimmed back, and a delayed writer can no longer find it
-    // in RAM after ownership moves. An already durable boundary is cheap.
-    if (rec.entry.caches.some((c) => c.signature() === "ssm")) {
-      if (this.#cold?.store(rec.entry.tokens, rec.entry.caches, ns) === false)
-        throw new Error("prompt-cache: failed to persist recurrent boundary before transfer");
-      this.#entries.splice(bestIdx, 1);
-      return rec.entry;
-    }
+    rec.uses++; rec.cost = bestLen; this.retention.accessed(rec);
+    // Lend immutable views for every backend state. RAM retention and SSD
+    // durability remain cache responsibilities while the caller extends its
+    // own views. This keeps synchronous persistence out of a RAM hit.
     const clones = this.#clone(rec.entry.caches);
+    let attachments: CheckpointAttachment[] | undefined;
     const trimNeeded = rec.entry.tokens.length - bestLen;
     try {
+      attachments = cloneAttachments(rec.entry.attachments);
       if (trimNeeded > 0) for (const c of clones) c.trim(trimNeeded);
     } catch (error) {
-      cleanupFailure(error, () => disposeResources(clones));
+      cleanupFailure(error, () => disposeResources([...clones, { dispose: () => disposeAttachments(attachments) }]));
     }
     return {
       tokens: rec.entry.tokens.slice(0, bestLen),
       caches: clones,
+      attachments,
       ns,
       retain: rec.share.acquire(),
     };
@@ -346,17 +568,15 @@ export class PromptCache {
     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;
+      if (p <= best || p < minimumReusableOffset(e.caches)) continue;
       const trimNeeded = e.tokens.length - p; // same usability rule as take()
-      if (trimNeeded > 0 && !e.caches.every((c) => c.isTrimmable())) continue;
+      if (trimNeeded > 0 && (!!e.attachments?.length || !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). */
+  /** Read-only lookup of an immutable, cache-owned checkpoint. */
   findExact(tokens: number[], ns = ""): PromptCacheEntry | null {
     for (const { entry: e } of this.#entries) {
       if (e.ns !== ns || e.tokens.length !== tokens.length) continue;
@@ -368,10 +588,12 @@ export class PromptCache {
   /** 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 {
+  put(tokens: number[], caches: Cache[], ns = "", retain?: () => void, attachments?: CheckpointAttachment[], sessionId?: string): void {
     // All cache-provided validation runs before ownership is adopted.
-    const bytes = cacheBytes(caches);
-    const trimmable = caches.every((cache) => cache.isTrimmable());
+    const bytes = cacheBytes(caches) + attachmentBytes(attachments);
+    const trimmable = !attachments?.length && caches.every((cache) => cache.isTrimmable());
+    // Storage takes resolved snapshots, independently of their producer.
+    this.materialize(caches, attachments);
     const discard = (entry: PromptCacheEntry, spill: boolean) => {
       try { this.#disposeEntry(entry, spill); }
       catch (error) {
@@ -380,8 +602,8 @@ export class PromptCache {
         console.warn(`[prompt-cache] entry cleanup failed: ${error}`);
       }
     };
-    const entry: PromptCacheEntry = { tokens, caches, ns, retain };
-    if (bytes > this.maxBytes) {
+    const entry: PromptCacheEntry = { tokens, caches, ns, retain, attachments };
+    if (bytes > this.maxBytes && !this.canEvict) {
       discard(entry, true);
       return;
     }
@@ -389,8 +611,11 @@ export class PromptCache {
     // shares; the entry's own retain is the donor's share.
     const share = makeSharedRetain(retain);
     entry.retain = share.acquire();
-    const rec: EntryRecord = { entry, bytes, lastUsed: ++this.#clock, lastUsedMs: Date.now(), share };
+    const rec: EntryRecord = { entry, bytes, cost: tokens.length, uses: 1, priority: 0, lastUsed: ++this.#clock, lastUsedMs: Date.now(), share, resident: true, sessions: 0 };
+    this.retention.accessed(rec);
     this.#entries.push(rec);
+    this.#records.set(entry, rec);
+    this.#associate(sessionId, ns, rec);
     // Exact duplicates (equal tokens) are redundant REGARDLESS of
     // trimmability — the new entry serves exactly the matches the old one
     // did (mlx-lm's trie replaces in place). Without this, every gemma ctx
@@ -400,6 +625,12 @@ export class PromptCache {
       if (old === rec || old.entry.ns !== ns) continue;
       if (old.entry.tokens.length !== tokens.length) continue;
       if (commonPrefixLength(old.entry.tokens, tokens) !== tokens.length) continue;
+      // Equivalent publication replaces the cache object, not the checkpoint
+      // identity. Every session sharing it follows the replacement.
+      if (old.sessions) for (const sessions of this.#sessions.values()) {
+        if (sessions.get(ns) !== old) continue;
+        sessions.set(ns, rec); old.sessions--; rec.sessions++;
+      }
       this.#entries = this.#entries.filter((r) => r !== old);
       discard(old.entry, false);
     }
@@ -415,39 +646,26 @@ export class PromptCache {
     if (trimmable) {
       for (const old of [...this.#entries]) {
         if (old === rec || old.entry.ns !== ns) continue;
-        if (old.entry.tokens.length >= tokens.length) continue;
+        if (old.entry.tokens.length >= tokens.length || old.entry.tokens.length < minimumReusableOffset(caches)) continue;
         if (commonPrefixLength(old.entry.tokens, tokens) !== old.entry.tokens.length) continue;
         this.#entries = this.#entries.filter((r) => r !== old);
         discard(old.entry, false);
       }
     }
-    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);
-      discard(evicted!.entry, true);
-    }
+    this.evictToBudget();
     try { this.onPut?.(tokens, ns); } catch {}
   }
 
-  /** Layer-0 idle demotion: spill every entry unused for `idleMs` to the
-   *  cold tier and FREE its GPU memory. The prefix stays reachable — the
-   *  next take() restores it via zero-copy mmap (~0.25 s for a 13.7k-token
-   *  entry, vs a 12 s re-prefill) — but between bursts the RAM tier drains
-   *  toward empty, returning unified memory to the system. Caller must
-   *  hold the generation lock (entries' arrays are disposed here) — with
-   *  a spillOwned sink the lock covers only the zero-copy clone; the
-   *  write runs off-lock and the buffers free when it settles. No-op
-   *  without a cold tier — demotion without a place to demote TO would
-   *  just be data loss. Returns entries demoted. */
+  /** Evict entries older than idleMs. Tiered residency policy first ensures
+   * SSD coverage. No-op without a cold tier. Returns entries removed from RAM. */
   demoteIdle(idleMs: number, now = Date.now()): number {
     if (!this.#cold) return 0;
     let n = 0;
     for (const rec of [...this.#entries]) {
-      if (now - rec.lastUsedMs < idleMs) continue;
+      if (now - rec.lastUsedMs < idleMs || this.#pins.has(rec.entry)) continue;
+      if (this.canEvict && !this.canEvict(rec.entry)) continue;
       this.#entries = this.#entries.filter((r) => r !== rec);
-      this.#disposeEntry(rec.entry, true); // spill-first, then dispose
+      this.#disposeEntry(rec.entry, !this.canEvict); // persisted entries need no duplicate write
       this.demotions++;
       n++;
     }
@@ -455,6 +673,8 @@ export class PromptCache {
   }
 
   clear(): void {
+    this.#sessions.clear();
+    this.#epoch++;
     const entries = this.#entries;
     this.#entries = [];
     disposeResources(entries.map(({ entry }) => ({ dispose: () => this.#disposeEntry(entry, false) })));

package/src/runtime-config.ts

@@ -52,11 +52,7 @@ export function runtimeValue(name: RuntimeKey): string | undefined {
   return runtimeConfig().value(name);
 }
 
-/** On/off env flag: an explicit "1"/"0" always wins; otherwise `defaultOn`.
- *  (Formerly src/flags.ts `flagOn` — same resolver, one module.) */
-export function flagOn(name: RuntimeKey, defaultOn: boolean): boolean {
-  return runtimeFlag(name, defaultOn);
-}
+/** An explicit "1"/"0" wins; other values retain the caller's default. */
 export function runtimeFlag(name: RuntimeKey, defaultOn: boolean): boolean {
   return runtimeConfig().flag(name, defaultOn);
 }

package/src/safetensors.ts

@@ -127,8 +127,15 @@ export class ShardedSafetensors {
         self.tensorToFile.set(tensor, sf);
       }
     } else {
-      const sf = SafetensorsFile.open(`${modelDir}/model.safetensors`);
-      self.files.set("model.safetensors", sf);
+      // Single-file artifacts: HF's model.safetensors, or mlx-whisper's
+      // weights.safetensors (mlx-community/whisper-*).
+      const file = (await Bun.file(`${modelDir}/model.safetensors`).exists())
+        ? "model.safetensors"
+        : (await Bun.file(`${modelDir}/weights.safetensors`).exists())
+          ? "weights.safetensors"
+          : "model.safetensors";
+      const sf = SafetensorsFile.open(`${modelDir}/${file}`);
+      self.files.set(file, sf);
       for (const name of sf.tensors.keys()) self.tensorToFile.set(name, sf);
     }
     return self;