mlx-bun 0.0.13 → 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,
@@ -24,7 +27,7 @@
 // reachable); onPut lets the server schedule its debounced write-behind
 // snapshot for both lanes.
 
-import { TurboQuantKVCache, type Cache } from "./model/gemma4";
+import type { Cache } from "./model/gemma4";
 import { cloneKvCaches } from "./kv-store";
 
 /** Reference-counted release: wraps an entry's `retain` (e.g. an mmap
@@ -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 instanceof TurboQuantKVCache) 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/atomic-output.ts

@@ -0,0 +1,38 @@
+import {
+  existsSync,
+  mkdirSync,
+  mkdtempSync,
+  readdirSync,
+  renameSync,
+  rmdirSync,
+  rmSync,
+  statSync,
+} from "node:fs";
+import { basename, dirname, join, resolve } from "node:path";
+
+/** Build a complete directory beside its destination, then publish it with a
+ * single same-filesystem rename. Failed writers leave no partial destination. */
+export async function writeAtomicDirectory<T>(
+  outDir: string,
+  write: (stagingDir: string) => Promise<T>,
+): Promise<T> {
+  const destination = resolve(outDir);
+  const replaceEmptyDestination = existsSync(destination);
+  if (
+    replaceEmptyDestination &&
+    (!statSync(destination).isDirectory() || readdirSync(destination).length !== 0)
+  ) throw new Error(`output directory already exists: ${outDir}`);
+
+  const parent = dirname(destination);
+  mkdirSync(parent, { recursive: true });
+  const staging = mkdtempSync(join(parent, `.${basename(destination)}.tmp-`));
+  try {
+    const result = await write(staging);
+    if (replaceEmptyDestination) rmdirSync(destination);
+    renameSync(staging, destination);
+    return result;
+  } catch (error) {
+    rmSync(staging, { recursive: true, force: true });
+    throw error;
+  }
+}

package/src/quantize/calibration.ts

@@ -20,6 +20,7 @@
 
 import { existsSync } from "node:fs";
 import type { LoadedTokenizer } from "../tokenizer";
+import { runtimeValue } from "../runtime-config";
 
 /** One JSONL calibration sample (subset of the OptIQ schema we can consume). */
 interface CalibrationSample {
@@ -34,7 +35,7 @@ function bundledMixCandidates(): string[] {
   const home = process.env.HOME ?? "";
   const rel = "lib/python3.14/site-packages/optiq/calibration/data/optiq.jsonl";
   return [
-    process.env.MLX_BUN_CALIBRATION_JSONL,
+    runtimeValue("MLX_BUN_CALIBRATION_JSONL"),
     `${home}/Code/mlx-lm-example/.venv/${rel}`,
     `${home}/Code/mlx-lm/.venv/${rel}`,
   ].filter((p): p is string => !!p);

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;
 }
@@ -76,6 +83,12 @@ export interface OptiqMetadata {
   group_size: number;
   achieved_bpw: number;
   per_layer_count: number;
+  weight_transforms?: Array<{
+    id: string;
+    seed: number;
+    family: string;
+    deviations: string[];
+  }>;
 }
 
 /**

package/src/quantize/index.ts

@@ -2,13 +2,40 @@
 // import from here: quantize a model directory, write sharded safetensors,
 // build the config block, and resolve/inspect models for the wizard UI.
 
-export { quantizeModelDir, isQuantizable } from "./quantizer";
+export { quantizeModelDir, isQuantizable, withPreparedProbe } from "./quantizer";
 export type {
+  PreparedProbe,
+  ProbeSource,
   QuantizeOptions,
   QuantizeResult,
   ProgressEvent,
 } from "./quantizer";
 
+export {
+  planQwen35Fold,
+  planQwenMtpFold,
+} from "./rotate";
+export type {
+  FoldOptions,
+  QwenFoldOptions,
+} from "./rotate";
+
+export {
+  automaticRotationWeightTransform,
+  llamaWeightTransform,
+  planLlamaWeightTransform,
+  qwen35WeightTransform,
+  qwenMtpWeightTransform,
+} from "./weight-transform";
+export type {
+  LlamaWeightTransformPlan,
+  QwenWeightTransformPlan,
+  WeightTransform,
+  WeightTransformContext,
+  WeightTransformMetadata,
+  WeightTransformPlan,
+} from "./weight-transform";
+
 export {
   writeShardedSafetensors,
   DEFAULT_SHARD_BYTES,

package/src/quantize/job.ts

@@ -13,6 +13,7 @@ import { registerRunner } from "../jobs/runner";
 import { Registry } from "../registry";
 import { loadModelConfig } from "../config";
 import { quantizeModelDir, type QuantizeOptions } from "./quantizer";
+import { automaticRotationWeightTransform } from "./weight-transform";
 
 /** Resolve the source model directory from a job config. Accepts an explicit
  *  filesystem path (src_dir) or a registry query / model id (model_id). */
@@ -57,6 +58,9 @@ export const quantizeRunner: JobRunner = async (emit: Emit, config) => {
   // MUST be forwarded into opts or quantizeModelDir silently runs uniform.
   const targetBpw = config.target_bpw != null ? Number(config.target_bpw) : undefined;
   const mixed = targetBpw !== undefined;
+  const rotationSeed = Number(config.rotation_seed ?? 42);
+  if (config.rotate_weights && !Number.isInteger(rotationSeed))
+    throw new Error(`quantize job: rotation_seed must be an integer (got ${String(config.rotation_seed)})`);
 
   emit({
     type: "stage",
@@ -74,6 +78,13 @@ export const quantizeRunner: JobRunner = async (emit: Emit, config) => {
     ...(config.reference ? { reference: String(config.reference) } : {}),
     ...(config.calibration_mix ? { calibrationMix: String(config.calibration_mix) } : {}),
     ...(config.n_calibration != null ? { nCalibration: Number(config.n_calibration) } : {}),
+    ...(config.rotate_weights
+      ? {
+          weightTransform: automaticRotationWeightTransform({
+            seed: rotationSeed,
+          }),
+        }
+      : {}),
   };
 
   const r = await quantizeModelDir(srcDir, outDir, opts, (e) =>