mlx-bun 0.2.0 → 0.3.0

package/src/prompt-cache.ts

@@ -1,3 +1,5 @@
+import { leaseCacheState } 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,13 +7,14 @@
 // 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
+// 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.
 //
 // TIERING (Layer 0, unified-engine plan): when a ColdTier is attached,
@@ -55,8 +58,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. */
@@ -65,7 +69,9 @@ export interface ColdTier {
    *  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;
+  /** 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;
 }
 
 /** Spill sink for eviction/demotion (write-behind, 2026-07-06).
@@ -105,14 +111,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;
 }
@@ -202,8 +201,8 @@ export class PromptCache {
         if (spill) try { spill(entry); } catch {}
       }
     }
-    for (const c of entry.caches) c.dispose();
-    entry.retain?.();
+    disposeResources(entry.caches);
+    entry.retain?.(); // backing release requires successful cache disposal
   }
 
   get totalBytes(): number {
@@ -214,14 +213,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) === 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. 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 —
+   *  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
@@ -258,27 +279,22 @@ 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 (trimNeeded > 0 && !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,9 +312,25 @@ export class PromptCache {
     const rec = this.#entries[bestIdx]!;
     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;
+    }
     const clones = this.#clone(rec.entry.caches);
     const trimNeeded = rec.entry.tokens.length - bestLen;
-    if (trimNeeded > 0) for (const c of clones) c.trim(trimNeeded);
+    try {
+      if (trimNeeded > 0) for (const c of clones) c.trim(trimNeeded);
+    } catch (error) {
+      cleanupFailure(error, () => disposeResources(clones));
+    }
     return {
       tokens: rec.entry.tokens.slice(0, bestLen),
       caches: clones,
@@ -337,10 +369,20 @@ export class PromptCache {
    *  (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 {
+    // All cache-provided validation runs before ownership is adopted.
     const bytes = cacheBytes(caches);
+    const trimmable = caches.every((cache) => cache.isTrimmable());
+    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 };
     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 +401,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 +412,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 (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 +426,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 +455,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;