mlx-bun 0.2.0 → 0.4.0

package/src/prompt-cache.ts

@@ -1,3 +1,6 @@
+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.
 //
 // The mlx-lm lesson (PLAN.md): a count-capped cache of multi-GB KV
@@ -5,14 +8,12 @@
 // per entry) and evicts least-recently-used until under the cap.
 //
 // Usage pattern (single generation queue, so take/put is race-free):
-//   const hit = cache.take(promptIds);   // longest-prefix match → CLONES
+//   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 (prefix sharing): it serves zero-copy clones and
-// leaves the donor entry in place — N agents sharing a system prompt all
-// clone from one donor, one prefill; 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
@@ -25,6 +26,7 @@
 // snapshot for both lanes.
 
 import type { Cache } from "./model/gemma4";
+import type { PrefixCache, PrefixCacheHit } from "./contracts/prefix-cache";
 import { cloneKvCaches } from "./kv-store";
 
 /** Reference-counted release: wraps an entry's `retain` (e.g. an mmap
@@ -55,8 +57,9 @@ function makeSharedRetain(retain: (() => void) | undefined): { acquire(): () =>
 
 /** The cold (SSD) tier the RAM cache tiers over. Structural — this module
  *  never imports ssd-cache (which imports us); the server binds
- *  SsdCacheStore + the model into this shape (src/server.ts). All methods
- *  must be failure-proof: any error degrades to "no hit"/"not stored". */
+ *  SsdCacheStore + the model into this shape (src/server.ts). Lookup and
+ *  restore failures return null; store returns false on failure. Operations
+ *  that require persistence retain their donor if storing fails. */
 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. */
@@ -64,8 +67,10 @@ export interface ColdTier {
   /** 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;
-  store(tokens: number[], caches: Cache[], ns: string): void;
+  restore(handle: unknown): { tokens: number[]; caches: Cache[]; attachments?: CheckpointAttachment[]; retain: () => void } | null;
+  /** Borrow live caches and finish persistence before returning. A false
+   *  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).
@@ -89,9 +94,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;
@@ -105,14 +108,7 @@ export interface PromptCacheEntry {
 export function cacheBytes(caches: Cache[]): number {
   let total = 0;
   for (const c of caches) {
-    const state = c.state();
-    for (const a of state) total += a.nbytes;
-    // Most Cache kinds' state() returns their own live-owned arrays (safe
-    // to just read), but TurboQuantKVCache allocates fresh trimmed slice
-    // views per call (kv-store.ts snapshotCache/cloneKvCaches contract) —
-    // those must be disposed here or they leak (see generate.ts
-    // evalCacheState for the same hazard on the prefill path).
-    if (c.stateNeedsDispose) for (const a of state) a.dispose();
+    total += withResource(leaseCacheState(c), (state) => state.reduce((bytes, a) => bytes + a.nbytes, 0));
   }
   return total;
 }
@@ -135,7 +131,7 @@ interface EntryRecord {
   share: { acquire(): () => void };
 }
 
-export class PromptCache {
+export class PromptCache implements PrefixCache<Cache[], CheckpointAttachment[]> {
   readonly maxBytes: number;
   #entries: EntryRecord[] = [];
   #clock = 0;
@@ -162,6 +158,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;
@@ -186,24 +183,27 @@ 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 {}
       }
     }
-    for (const c of entry.caches) c.dispose();
-    entry.retain?.();
+    disposeResources([...entry.caches, { dispose: () => disposeAttachments(entry.attachments) }]);
+    entry.retain?.(); // backing release requires successful cache disposal
   }
 
   get totalBytes(): number {
@@ -214,14 +214,36 @@ export class PromptCache {
     return this.#entries.length;
   }
 
+  /** Release LRU entries until the caller's live-memory check is satisfied.
+   * Unlike background eviction, pressure relief must finish the SSD write
+   * and release the buffers now, not enqueue clones that keep them alive.
+   * Call under the engine lock, after taking the request's reusable prefix. */
+  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]!;
+      // 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, entry.attachments) === false) break;
+      }
+      else this.#spillSync?.(entry);
+      this.#entries.splice(oldest, 1);
+      this.#disposeEntry(entry, false);
+      this.demotions++;
+      count++;
+    }
+    return count;
+  }
+
   /** Find the entry with the longest usable common prefix of `prompt`
-   *  across BOTH tiers and serve it NON-CONSUMINGLY (prefix sharing,
-   *  2026-07-05): the caller receives zero-copy CLONES trimmed to the
-   *  matched prefix (plus a ref-counted retain share); the donor entry
-   *  stays in the cache, ready for the next agent/session with the same
-   *  prefix. The caller owns the returned caches — dispose them or put()
-   *  an extended entry (which supersedes prefix-ancestors), honoring
-   *  `retain` either way.
+   *  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
    *  one token must be forwarded to produce logits). Entries longer than
@@ -241,9 +263,9 @@ export class PromptCache {
       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;
     }
@@ -258,27 +280,23 @@ export class PromptCache {
       if (hit && hit.prefixLen > bestLen) {
         const loaded = this.#cold.restore(hit.handle);
         if (loaded) {
-          const trimNeeded = loaded.tokens.length - hit.prefixLen;
-          if (trimNeeded > 0 && !loaded.caches.every((c) => c.isTrimmable())) {
-            // Divergent tail on an untrimmable kind: this file can't seed
-            // the prompt — fall back to the RAM candidate. Loud: this used
-            // to be the silent 84 s "restore-then-re-prefill" path (the
-            // whole cold tier degrading with zero observables, 2026-07-06).
-            console.warn(
-              `[prompt-cache] cold entry unusable: ${loaded.tokens.length} tokens, ` +
-              `needs trim ${trimNeeded} but untrimmable (wrapped ring/SSM) — re-prefilling`,
-            );
-            for (const c of loaded.caches) c.dispose();
-            loaded.retain();
-          } else {
-            if (trimNeeded > 0) for (const c of loaded.caches) c.trim(trimNeeded);
-            return {
-              tokens: loaded.tokens.slice(0, hit.prefixLen),
-              caches: loaded.caches,
-              ns,
-              retain: loaded.retain,
-            };
-          }
+          const owner = ownResource(loaded, (entry) =>
+            this.#disposeEntry({ ...entry, ns }, false));
+          const result = withResource(owner, (entry) => {
+            const trimNeeded = entry.tokens.length - hit.prefixLen;
+            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`,
+              );
+              return null;
+            }
+            if (trimNeeded > 0) for (const cache of entry.caches) cache.trim(trimNeeded);
+            const tokens = entry.tokens.slice(0, hit.prefixLen);
+            return { ...owner.transfer(), tokens, ns };
+          });
+          if (result) return result;
         }
       }
     }
@@ -296,12 +314,22 @@ export class PromptCache {
     const rec = this.#entries[bestIdx]!;
     rec.lastUsed = ++this.#clock;
     rec.lastUsedMs = Date.now();
+    // 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;
-    if (trimNeeded > 0) for (const c of clones) c.trim(trimNeeded);
+    try {
+      attachments = cloneAttachments(rec.entry.attachments);
+      if (trimNeeded > 0) for (const c of clones) c.trim(trimNeeded);
+    } catch (error) {
+      cleanupFailure(error, () => disposeResources([...clones, { dispose: () => disposeAttachments(attachments) }]));
+    }
     return {
       tokens: rec.entry.tokens.slice(0, bestLen),
       caches: clones,
+      attachments,
       ns,
       retain: rec.share.acquire(),
     };
@@ -314,9 +342,9 @@ 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;
@@ -336,11 +364,23 @@ 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 {
-    const bytes = cacheBytes(caches);
-    const entry: PromptCacheEntry = { tokens, caches, ns, retain };
+  put(tokens: number[], caches: Cache[], ns = "", retain?: () => void, attachments?: CheckpointAttachment[]): void {
+    // All cache-provided validation runs before ownership is adopted.
+    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) {
+        // The new entry is already owned here. Throwing back to its previous
+        // owner would invite disposal of state still held by this store.
+        console.warn(`[prompt-cache] entry cleanup failed: ${error}`);
+      }
+    };
+    const entry: PromptCacheEntry = { tokens, caches, ns, retain, attachments };
     if (bytes > this.maxBytes) {
-      this.#disposeEntry(entry, true);
+      discard(entry, true);
       return;
     }
     // The donor's retain becomes ref-counted so take()'s clones can hold
@@ -359,7 +399,7 @@ export class PromptCache {
       if (old.entry.tokens.length !== tokens.length) continue;
       if (commonPrefixLength(old.entry.tokens, tokens) !== tokens.length) continue;
       this.#entries = this.#entries.filter((r) => r !== old);
-      this.#disposeEntry(old.entry, false);
+      discard(old.entry, false);
     }
     // SUPERSEDE strict prefix-ancestors (the SSD store's rule, brought to
     // RAM — prefix sharing would otherwise grow an entry per turn): same-ns
@@ -370,13 +410,13 @@ export class PromptCache {
     // can only serve exact-length matches, so shorter ancestors — e.g.
     // the prompt-boundary snapshot that makes drifted multi-turn agents
     // hit — must survive.
-    if (caches.every((c) => c.isTrimmable())) {
+    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);
-        this.#disposeEntry(old.entry, false);
+        discard(old.entry, false);
       }
     }
     while (this.totalBytes > this.maxBytes && this.#entries.length > 1) {
@@ -384,7 +424,7 @@ export class PromptCache {
       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);
-      this.#disposeEntry(evicted!.entry, true);
+      discard(evicted!.entry, true);
     }
     try { this.onPut?.(tokens, ns); } catch {}
   }
@@ -413,7 +453,8 @@ export class PromptCache {
   }
 
   clear(): void {
-    for (const e of this.#entries) this.#disposeEntry(e.entry, false);
+    const entries = this.#entries;
     this.#entries = [];
+    disposeResources(entries.map(({ entry }) => ({ dispose: () => this.#disposeEntry(entry, false) })));
   }
 }

package/src/quantize/allocator.ts

@@ -11,7 +11,7 @@
 //   - NaN-sensitivity handling (treat as maximally sensitive)
 //   - post-greedy block-run guard
 //
-// Verified bit-for-bit against the Python via tests/quantize-allocator.test.ts.
+// Verified bit-for-bit against the Python via tests/parity/quantize-allocator.test.ts.
 //
 // The latency-aware variants (`optimize_latency_aware`,
 // `optimize_for_latency_budget`) are intentionally NOT ported — they pull in

package/src/quantize/config-writer.ts

@@ -13,7 +13,9 @@ import { existsSync } from "node:fs";
 import { join } from "node:path";
 
 /** A per-module quantization override. `false` means "left unquantized". */
-export type PerLayerEntry = { bits: number; groupSize: number } | false;
+export type PerLayerEntry =
+  | { bits: number; groupSize: number; mode?: string; trellis?: { L: number; code: string; axis: 0 | 1 } }
+  | false;
 
 /** The default quantization definition for the model. */
 export interface QuantDef {
@@ -28,7 +30,7 @@ export interface QuantizationBlock {
   bits: number;
   mode: string;
   // plus per-module entries: { [modulePath]: {bits, group_size} | false }
-  [modulePath: string]: number | string | { bits: number; group_size: number } | false;
+  [modulePath: string]: number | string | { bits: number; group_size: number; mode?: string; trellis?: { L: number; code: string; axis: 0 | 1 } } | false;
 }
 
 /**
@@ -47,7 +49,12 @@ export function buildQuantizationBlock(
     mode: def.mode ?? "affine",
   };
   for (const [path, entry] of perLayer) {
-    block[path] = entry === false ? false : { bits: entry.bits, group_size: entry.groupSize };
+    if (entry === false) { block[path] = false; continue; }
+    const v: { bits: number; group_size: number; mode?: string; trellis?: { L: number; code: string; axis: 0 | 1 } } =
+      { bits: entry.bits, group_size: entry.groupSize };
+    if (entry.mode) v.mode = entry.mode;
+    if (entry.trellis) v.trellis = entry.trellis;
+    block[path] = v;
   }
   return block;
 }

package/src/quantize/rotate.ts

@@ -3,7 +3,7 @@
 // Folds an orthogonal rotation offline into producer/consumer weight pairs so
 // weights are quantized in a basis where outlier channels are smeared into
 // near-Gaussian marginals. Recipe, fold table, and the decided deviations from
-// the QuaRot/SpinQuant references live in docs/design/turboquant-weights.md —
+// the QuaRot/SpinQuant references live in docs/design/turboquant.md —
 // notably: γ-fold FIRST (R1 only commutes with gain-free RMSNorm), SpinQuant's
 // fully-offline per-head R2 (NOT QuaRot's, which needs a runtime op), NO R4
 // half-fold on down_proj's input, NO embedding mean-centering (not an exact
@@ -47,7 +47,7 @@ function assertPow2(n: number, what: string): void {
   if (n < 2 || (n & (n - 1)) !== 0)
     throw new Error(
       `rotation fold: ${what}=${n} is not a power of two — the Kronecker ` +
-      `Hadamard path is not implemented (docs/design/turboquant-weights.md)`,
+      `Hadamard path is not implemented (docs/design/turboquant.md)`,
     );
 }
 
@@ -323,7 +323,7 @@ export function foldLlamaWeights(
 // elementwise in head space) does not commute with a per-head rotation, so R2
 // is architecturally off for this family — recorded in the design doc.
 //
-// Residual corridors (docs/design/turboquant-weights.md, W1 map):
+// Residual corridors (docs/design/turboquant.md, W1 map):
 //   readers  (@R1 input dim, γ premultiplied): q/k/v_proj,
 //     linear_attn.in_proj_{qkv,z,b,a}, mlp.gate/up_proj, lm_head
 //   writers  (R1ᵀ output dim): self_attn.o_proj, linear_attn.out_proj,
@@ -370,7 +370,7 @@ export interface QwenFoldPlan {
  * FoldOp — pure name analysis, no arrays. R1-ONLY (+γ): the attention output
  * gate (`o_proj(out · σ(gate))`, elementwise in head space) does not commute
  * with a per-head rotation, so R2 is architecturally off for this family.
- * Corridor map: docs/design/turboquant-weights.md §W1.
+ * Corridor map: docs/design/turboquant.md §W1.
  */
 export function planQwen35Fold(names: readonly string[], prefix = "language_model."): QwenFoldPlan {
   const P = prefix;