mlx-bun 0.0.9 → 0.0.11

package/src/preflight.ts

@@ -5,6 +5,16 @@
 //   - cold page cache: prefill timer dominated by SSD page-in
 // The harness refuses to record headline numbers when checks fail;
 // the state snapshot is stored in every eval-DB row either way.
+//
+// LIMITS of this instrument (2026-07-05 finding): the snapshot is taken
+// BETWEEN legs, when the bench process has exited and memory is quiet —
+// free_pct reads a near-constant ~92% on an otherwise-idle box and load1m
+// sat under every threshold while a 40-minute mid-pass slow window cut
+// decode rates ~40% across every stack. Preflight ok ≠ clean measurements
+// and preflight-failed ≠ dirty ones (the 07-04 pass was ‡-tagged yet
+// stable). The DIRECT quality signal is per-cell run spread — see the
+// run-stability policy in scripts/bench-h2h.ts (SPREAD_TOL retries +
+// `unstable`/`stabilized` tags). Preflight remains the coarse pre-gate.
 
 export interface MachineState {
   swapUsedMB: number;
@@ -77,21 +87,30 @@ export function checkMachine(limits: PreflightLimits = {}): MachineState {
   if (loadAvg1m > 8)
     problems.push(`1-min load ${loadAvg1m.toFixed(1)} — post-boot indexing still running; wait a minute`);
 
-  // big foreign processes (rss in KB from ps)
+  // big foreign processes: high RSS *or* high CPU (pcpu + rss in KB from
+  // ps). The CPU arm is the 2026-07-05 lesson: Apple's
+  // knowledgeconstructiond (Apple Intelligence indexing) burned 87% CPU
+  // with a tiny RSS during the pass's 40-min slow window — a background
+  // ML workload the RSS-only check could never flag. CPU is bursty, so
+  // one clean sample doesn't prove absence (the run-spread stability
+  // policy in bench-h2h.ts is the direct signal); a hot sample here is
+  // still a cheap early warning.
   const SELF_PATTERN = /(bun|\.venv\/bin\/python|mlx_lm|optiq)/;
   const SYSTEM_PATTERN = /(kernel_task|WindowServer|launchd|mds|spotlight)/i;
+  const BIG_CPU_PCT = 50;
   const bigProcesses: { rssMB: number; command: string }[] = [];
-  for (const line of sh(["ps", "axo", "rss=,command="]).split("\n")) {
-    const m = line.match(/^\s*(\d+)\s+(.*)$/);
+  for (const line of sh(["ps", "axo", "pcpu=,rss=,command="]).split("\n")) {
+    const m = line.match(/^\s*([\d.]+)\s+(\d+)\s+(.*)$/);
     if (!m) continue;
-    const rssMB = Number(m[1]) / 1024;
-    const command = m[2]!.slice(0, 80);
-    if (rssMB < bigRss) continue;
+    const cpuPct = Number(m[1]);
+    const rssMB = Number(m[2]) / 1024;
+    const command = m[3]!.slice(0, 80);
+    if (rssMB < bigRss && cpuPct < BIG_CPU_PCT) continue;
     if (SELF_PATTERN.test(command) || SYSTEM_PATTERN.test(command)) continue;
-    bigProcesses.push({ rssMB: Math.round(rssMB), command });
+    bigProcesses.push({ rssMB: Math.round(rssMB), command: cpuPct >= BIG_CPU_PCT ? `${command} (${cpuPct.toFixed(0)}% cpu)` : command });
   }
   for (const p of bigProcesses)
-    problems.push(`big process resident: ${p.rssMB} MB — ${p.command}`);
+    problems.push(`big process: ${p.rssMB} MB — ${p.command}`);
 
   return {
     swapUsedMB, freePercent, cpuSpeedLimit, gpuWiredLimitMB, loadAvg1m,

package/src/prompt-cache.ts

@@ -1,15 +1,93 @@
-// Byte-capped LRU prompt cache.
+// 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
 // entries is an OOM footgun. Ours accounts bytes (sum of KV array bytes
 // per entry) and evicts least-recently-used until under the cap.
 //
-// Usage pattern (single generation queue, so take/reinsert is race-free):
-//   const hit = cache.take(promptIds);     // longest strict-prefix match
+// Usage pattern (single generation queue, so take/put is race-free):
+//   const hit = cache.take(promptIds);   // longest-prefix match → CLONES
 //   generate(model, promptIds, { cache: hit?.caches ?? fresh })
-//   cache.put([...promptIds, ...generated], caches);  // extended entry
+//   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.
+//
+// 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
+// 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.
+
+import { TurboQuantKVCache, type Cache } from "./model/gemma4";
+import { cloneKvCaches } from "./kv-store";
+
+/** Reference-counted release: wraps an entry's `retain` (e.g. an mmap
+ *  unmap) so DONOR + CLONES can each hold a share — the underlying hook
+ *  runs once, after the LAST holder releases. Prefix sharing hands out
+ *  zero-copy VIEWS of a donor's buffers; if the donor's unmap ran while a
+ *  clone still referenced the mapped pages, the clone would read freed
+ *  memory. */
+function makeSharedRetain(retain: (() => void) | undefined): { acquire(): () => void } {
+  let count = 0;
+  let released = false;
+  return {
+    acquire() {
+      count++;
+      let mine = false;
+      return () => {
+        if (mine) return; // idempotent per share
+        mine = true;
+        count--;
+        if (count === 0 && !released) {
+          released = true;
+          retain?.();
+        }
+      };
+    },
+  };
+}
 
-import type { Cache } from "./model/gemma4";
+/** 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". */
+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;
+  store(tokens: number[], caches: Cache[], ns: string): void;
+}
+
+/** Spill sink for eviction/demotion (write-behind, 2026-07-06).
+ *
+ *  `spillOwned` (preferred, NON-BLOCKING): the cache hands the sink an
+ *  entry it OWNS — zero-copy CLONES of the caches (made via the same
+ *  injectable cloner take() uses, BEFORE the donor is disposed; entries
+ *  are immutable so the clones stay consistent forever) plus a copy of
+ *  the tokens. The sink must dispose entry.caches when done, on BOTH
+ *  fulfill and reject paths (the server chains storeAsync -> dispose on
+ *  ssdWriteChain). A sink that throws SYNCHRONOUSLY has not taken
+ *  ownership — the cache disposes the clones and the spill degrades to
+ *  a no-op; eviction/demotion never unwinds.
+ *
+ *  `spillSync` (fallback): called with the LIVE entry about to be
+ *  disposed; must finish reading before returning. This is the old
+ *  contract — a full synchronous serialize+write under the generation
+ *  lock — kept for tests and simple embedders. */
+export interface SpillSink {
+  spillOwned?: (entry: PromptCacheEntry) => void;
+  spillSync?: (entry: PromptCacheEntry) => void;
+}
 
 export interface PromptCacheEntry {
   tokens: number[];
@@ -17,31 +95,115 @@ export interface PromptCacheEntry {
   /** Namespace key — adapter spec for LoRA requests ("" = base model).
    *  KV computed under one adapter must never seed another's prefill. */
   ns: string;
+  /** Release hook for entries restored from the SSD cold tier: unmaps the
+   *  backing COW file mapping. MUST run only after `caches` are disposed
+   *  (they may alias file pages); every dispose path here honors that, and
+   *  callers who take() an entry carry the thunk with it. */
+  retain?: () => void;
 }
 
-function cacheBytes(caches: Cache[]): number {
+export function cacheBytes(caches: Cache[]): number {
   let total = 0;
-  for (const c of caches) for (const a of c.state()) total += a.nbytes;
+  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();
+  }
   return total;
 }
 
 /** Length of the longest common prefix of two token sequences. */
-function commonPrefixLength(a: number[], b: number[]): number {
+export function commonPrefixLength(a: number[], b: number[]): number {
   const n = Math.min(a.length, b.length);
   let i = 0;
   while (i < n && a[i] === b[i]) i++;
   return i;
 }
 
+interface EntryRecord {
+  entry: PromptCacheEntry;
+  bytes: number;
+  lastUsed: number;
+  lastUsedMs: number;
+  /** Ref-counted retain: the donor holds one share (entry.retain); every
+   *  clone handed out by take() holds another. */
+  share: { acquire(): () => void };
+}
+
 export class PromptCache {
   readonly maxBytes: number;
-  #entries: { entry: PromptCacheEntry; bytes: number; lastUsed: number }[] = [];
+  #entries: EntryRecord[] = [];
   #clock = 0;
   hits = 0;
   misses = 0;
+  demotions = 0;
+  /** Cold-tier spill sinks (see SpillSink): spillOwned gets clones and
+   *  runs the write off-lock; spillSync gets the live entry inline.
+   *  Failures are the sink's problem; eviction proceeds regardless. */
+  readonly #spillOwned: ((entry: PromptCacheEntry) => void) | null;
+  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. */
+  onPut: ((tokens: number[], ns: string) => void) | null = null;
 
-  constructor(maxBytes: number) {
+  /** Zero-copy view cloner — injectable so model-free tests can stub it. */
+  readonly #clone: (caches: Cache[]) => Cache[];
+
+  constructor(
+    maxBytes: number,
+    // Bare function = legacy sync spill (tests, simple embedders).
+    spill: SpillSink | ((entry: PromptCacheEntry) => void) | null = null,
+    cold: ColdTier | null = null,
+    clone: (caches: Cache[]) => Cache[] = cloneKvCaches,
+  ) {
     this.maxBytes = maxBytes;
+    this.#spillOwned = typeof spill === "function" ? null : spill?.spillOwned ?? null;
+    this.#spillSync = typeof spill === "function" ? spill : spill?.spillSync ?? null;
+    this.#cold = cold;
+    this.#clone = clone;
+  }
+
+  #disposeEntry(entry: PromptCacheEntry, spillFirst: boolean): void {
+    if (spillFirst) {
+      if (this.#spillOwned) {
+        // NON-BLOCKING spill: clone BEFORE disposing (clones are zero-copy
+        // views of the live arrays) and hand ownership to the sink, which
+        // flushes off the generation lock. The underlying GPU buffers stay
+        // alive until the sink disposes the clones — for demoteIdle that
+        // means the memory is freed one bounded write-chain hop later, not
+        // instantly; acceptable because the chain is serial and each write
+        // disposes its clones on settle (fulfill AND reject).
+        // Retain contract: entry.retain still runs below, possibly while
+        // the clones are in flight — safe because cold-restored entries
+        // OWN their bytes (streamed copy-restore, 2026-07-07): there is
+        // no backing mmap for retain to unmap, so nothing the clones
+        // could alias goes away.
+        let clones: Cache[] | null = null;
+        try {
+          clones = this.#clone(entry.caches);
+          this.#spillOwned({ tokens: [...entry.tokens], caches: clones, ns: entry.ns });
+        } catch {
+          // Failure degrades to no-spill; eviction/demotion never unwinds.
+          if (clones) for (const c of clones) try { c.dispose(); } catch {}
+        }
+      } else {
+        const spill =
+          this.#spillSync ??
+          (this.#cold
+            ? (e: PromptCacheEntry) => this.#cold!.store(e.tokens, e.caches, e.ns)
+            : null);
+        if (spill) try { spill(entry); } catch {}
+      }
+    }
+    for (const c of entry.caches) c.dispose();
+    entry.retain?.();
   }
 
   get totalBytes(): number {
@@ -52,15 +214,26 @@ export class PromptCache {
     return this.#entries.length;
   }
 
-  /** Find the entry with the longest usable common prefix of `prompt`,
-   *  remove it, trim its caches to that prefix if needed (and possible),
-   *  and hand over ownership (caller must put() it back — possibly
-   *  extended — or dispose it).
+  /** 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.
    *
    *  Usable prefix = common prefix capped at prompt.length - 1 (at least
    *  one token must be forwarded to produce logits). Entries longer than
    *  the prefix need cache.trim(); ring caches lose trimability once
-   *  wrapped — those entries only match in full. */
+   *  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
+   *  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 {
     let bestIdx = -1;
     let bestLen = 0;
@@ -74,40 +247,173 @@ export class PromptCache {
       bestLen = p;
       bestIdx = i;
     }
-    if (bestIdx === -1) {
-      this.misses++;
-      return null;
+    if (bestIdx === -1) this.misses++;
+    else this.hits++;
+
+    // Cold tier: strictly longer usable prefix ⇒ restore + trim; the RAM
+    // 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);
+      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,
+            };
+          }
+        }
+      }
     }
-    this.hits++;
-    const { entry } = this.#entries.splice(bestIdx, 1)[0]!;
-    const trimNeeded = entry.tokens.length - bestLen;
-    if (trimNeeded > 0) {
-      for (const c of entry.caches) c.trim(trimNeeded);
-      entry.tokens = entry.tokens.slice(0, bestLen);
+
+    if (bestIdx === -1) return null;
+    // PREFIX SHARING (non-consuming serve): the caller gets zero-copy
+    // CLONES of the donor's caches, trimmed to the matched prefix; the
+    // donor stays in the cache untouched, ready to serve the next agent /
+    // session with the same prefix. Safe because mlx cache updates are
+    // FUNCTIONAL (updateAndFetch reassigns fresh arrays; buffer donation
+    // only fires at refcount 1, and the donor always holds a ref), so a
+    // clone extending itself can never mutate the donor's bytes. The old
+    // 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]!;
+    rec.lastUsed = ++this.#clock;
+    rec.lastUsedMs = Date.now();
+    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);
+    return {
+      tokens: rec.entry.tokens.slice(0, bestLen),
+      caches: clones,
+      ns,
+      retain: rec.share.acquire(),
+    };
+  }
+
+  /** Longest usable prefix length available for `prompt` WITHOUT taking
+   *  (the SSD tier compares this against its own best before restoring). */
+  peekPrefixLen(prompt: number[], ns = ""): number {
+    let best = 0;
+    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;
+      const trimNeeded = e.tokens.length - p; // same usability rule as take()
+      if (trimNeeded > 0 && !e.caches.every((c) => c.isTrimmable())) continue;
+      best = p;
     }
-    return entry;
+    return best;
   }
 
-  /** Insert (or reinsert) an entry; evicts LRU entries over the byte cap.
-   *  If the entry itself exceeds the cap it is disposed, not stored. */
-  put(tokens: number[], caches: Cache[], ns = ""): void {
+  /** 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). */
+  findExact(tokens: number[], ns = ""): PromptCacheEntry | null {
+    for (const { entry: e } of this.#entries) {
+      if (e.ns !== ns || e.tokens.length !== tokens.length) continue;
+      if (commonPrefixLength(e.tokens, tokens) === tokens.length) return e;
+    }
+    return null;
+  }
+
+  /** 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 };
     if (bytes > this.maxBytes) {
-      for (const c of caches) c.dispose();
+      this.#disposeEntry(entry, true);
       return;
     }
-    this.#entries.push({ entry: { tokens, caches, ns }, bytes, lastUsed: ++this.#clock });
+    // The donor's retain becomes ref-counted so take()'s clones can hold
+    // 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 };
+    this.#entries.push(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
+    // repeat left another full-size wrapped-ring entry behind (measured
+    // ~0.4 GB retained per repeat on the 12B bench).
+    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) !== tokens.length) continue;
+      this.#entries = this.#entries.filter((r) => r !== old);
+      this.#disposeEntry(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
+    // entries whose tokens are a prefix of the new entry are redundant
+    // and disposed WITHOUT spill (the new entry serves their prefixes and
+    // its own write-behind persists them) — but ONLY when the new entry's
+    // caches are all trimmable. An untrimmable new entry (wrapped ring)
+    // 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())) {
+      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);
+      }
+    }
     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);
-      for (const c of evicted!.entry.caches) c.dispose();
+      this.#disposeEntry(evicted!.entry, true);
+    }
+    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. */
+  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;
+      this.#entries = this.#entries.filter((r) => r !== rec);
+      this.#disposeEntry(rec.entry, true); // spill-first, then dispose
+      this.demotions++;
+      n++;
     }
+    return n;
   }
 
   clear(): void {
-    for (const e of this.#entries) for (const c of e.entry.caches) c.dispose();
+    for (const e of this.#entries) this.#disposeEntry(e.entry, false);
     this.#entries = [];
   }
 }

package/src/quantize/quantizer.ts

@@ -49,6 +49,13 @@ export interface QuantizeOptions {
   groupSize: 32 | 64;
   /** Quantization scheme; v1 only exercises "affine". */
   mode?: string;
+  /** Extra per-module gate on top of the shape-eligibility check: return
+   *  false to keep a module full-precision even though its shape qualifies
+   *  (e.g. the DeepSpec drafter's confidence_head — threshold-sensitive and
+   *  tiny). Excluded modules are recorded as `false` in the quantization
+   *  block (mlx's "not quantized" convention). Shape-ineligible modules
+   *  never reach this. */
+  quantizePredicate?: (base: string, fullShape: number[]) => boolean;
 
   // --- mixed-precision (OptiQ sensitivity + knapsack) -------------------
   // Setting targetBpw switches quantizeModelDir to the mixed path: it loads the
@@ -225,7 +232,11 @@ export async function quantizeModelDir(
       // re-quantized weight it is the dequantized array's shape.
       const fullShape = alreadyQuant ? materialized!.shape : shape;
 
-      if (!isQuantizable(fullShape, groupSize)) {
+      const predicateExcluded =
+        isQuantizable(fullShape, groupSize) &&
+        opts.quantizePredicate !== undefined &&
+        !opts.quantizePredicate(base, fullShape);
+      if (!isQuantizable(fullShape, groupSize) || predicateExcluded) {
         // Not eligible: pass the full-precision weight through unchanged
         // (re-materialized to bf16 if it had been quantized).
         if (alreadyQuant) {
@@ -235,6 +246,10 @@ export async function quantizeModelDir(
         } else {
           out.push({ name: weightName, array: fullWeight });
         }
+        // Policy-excluded (shape-eligible) modules get an explicit `false`
+        // in the quantization block — the loader must not apply the block's
+        // default spec to them (mlx's "not quantized" convention).
+        if (predicateExcluded) perLayer.set(base, false);
         bumpProgress();
         continue;
       }

package/src/registry.ts

@@ -34,6 +34,14 @@ export interface ModelRecord {
    *  AND for configs whose nested vision_config is not a vision tower
    *  (Qwen3.5 nests a copy of its own text config there). */
   visionConfigType: string | null;
+  /** config.json declares an `audio_config` block (gemma-4 Conformer). */
+  hasAudioConfig: boolean;
+  /** The optiq_vision.safetensors sidecar actually ships `audio_tower.*`
+   *  tensors (header-only check). Distinct from hasAudioConfig: the 12B
+   *  OptiQ snapshot pairs an audio_config with a STUB sidecar holding only
+   *  embed_audio.embedding_projection.weight, so config presence alone
+   *  is not an audio-capability signal. */
+  hasAudioTower: boolean;
   hasKvConfig: boolean;
   hasToolTemplate: boolean;
   numLayers: number | null;
@@ -60,6 +68,8 @@ CREATE TABLE IF NOT EXISTS models (
   quant_mode TEXT,
   has_vision_sidecar INTEGER NOT NULL,
   vision_config_type TEXT,
+  has_audio_config INTEGER NOT NULL DEFAULT 0,
+  has_audio_tower INTEGER NOT NULL DEFAULT 0,
   has_kv_config INTEGER NOT NULL,
   has_tool_template INTEGER NOT NULL,
   num_layers INTEGER,
@@ -97,7 +107,8 @@ export class Registry {
       !cols.includes("sidecar_bytes") ||
       !cols.includes("experts_bytes") ||
       !cols.includes("license") ||
-      !cols.includes("vision_config_type")
+      !cols.includes("vision_config_type") ||
+      !cols.includes("has_audio_tower")
     ) {
       this.db.exec("DROP TABLE models");
       this.db.exec(SCHEMA);
@@ -115,7 +126,8 @@ export class Registry {
     const upsert = this.db.prepare(`
       INSERT OR REPLACE INTO models VALUES
       ($path, $repo, $type, $params, $size, $sidecar, $experts, $bits, $gs, $mode,
-       $vision, $vtype, $kv, $tools, $layers, $hidden, $vocab, $license, $at)
+       $vision, $vtype, $audiocfg, $audiotower, $kv, $tools, $layers, $hidden,
+       $vocab, $license, $at)
     `);
     for (const entry of readdirSync(hubDir)) {
       if (!entry.startsWith("models--")) continue;
@@ -133,6 +145,8 @@ export class Registry {
           $bits: rec.quantBits, $gs: rec.quantGroupSize, $mode: rec.quantMode,
           $vision: rec.hasVisionSidecar ? 1 : 0,
           $vtype: rec.visionConfigType,
+          $audiocfg: rec.hasAudioConfig ? 1 : 0,
+          $audiotower: rec.hasAudioTower ? 1 : 0,
           $kv: rec.hasKvConfig ? 1 : 0,
           $tools: rec.hasToolTemplate ? 1 : 0,
           $layers: rec.numLayers, $hidden: rec.hiddenSize, $vocab: rec.vocabSize,
@@ -218,6 +232,15 @@ export function visionCapable(m: Pick<ModelRecord, "hasVisionSidecar" | "visionC
   return m.hasVisionSidecar || (m.visionConfigType?.endsWith("unified_vision") ?? false);
 }
 
+/** True when the model can answer audio (`input_audio`) requests: config.json
+ *  declares an `audio_config` AND the sidecar actually ships the Conformer
+ *  (`audio_tower.*`) tensors. Both legs matter — the 12B OptiQ snapshot has
+ *  audio_config but a stub sidecar (embed_audio only), so config presence
+ *  alone would advertise a capability that 400s at request time. */
+export function audioCapable(m: Pick<ModelRecord, "hasAudioConfig" | "hasAudioTower">): boolean {
+  return m.hasAudioConfig && m.hasAudioTower;
+}
+
 /** Of several cached revisions of ONE repo, pick the canonical snapshot: the
  *  revision refs/main points at, else the most recently scanned (stable
  *  tie-break on path). Paths are `<hub>/models--<repo>/snapshots/<hash>`. */
@@ -250,6 +273,8 @@ function rowToRecord(r: Record<string, unknown>): ModelRecord {
     quantMode: r.quant_mode as string | null,
     hasVisionSidecar: !!r.has_vision_sidecar,
     visionConfigType: (r.vision_config_type as string | null) ?? null,
+    hasAudioConfig: !!r.has_audio_config,
+    hasAudioTower: !!r.has_audio_tower,
     hasKvConfig: !!r.has_kv_config,
     hasToolTemplate: !!r.has_tool_template,
     numLayers: r.num_layers as number | null,
@@ -280,6 +305,34 @@ function readmeLicense(dir: string): string | null {
   }
 }
 
+/** True when a safetensors file's header names any `audio_tower.*` tensor
+ *  (header JSON only — never touches tensor data). Separates a real
+ *  audio-capable sidecar (e4b: 752 audio tensors) from the 12B stub, whose
+ *  only audio entry is embed_audio.embedding_projection.weight. Missing or
+ *  unparseable files read as false. */
+export function sidecarShipsAudioTower(path: string): boolean {
+  let fd: number;
+  try {
+    fd = openSync(path, "r");
+  } catch {
+    return false;
+  }
+  try {
+    const lenBuf = Buffer.alloc(8);
+    readSync(fd, lenBuf, 0, 8, 0);
+    const headerLen = Number(lenBuf.readBigUInt64LE(0));
+    if (headerLen <= 0 || headerLen > 256 * 2 ** 20) return false;
+    const headerBuf = Buffer.alloc(headerLen);
+    readSync(fd, headerBuf, 0, headerLen, 8);
+    const header = JSON.parse(headerBuf.toString("utf8")) as Record<string, unknown>;
+    return Object.keys(header).some((n) => n.startsWith("audio_tower."));
+  } catch {
+    return false;
+  } finally {
+    closeSync(fd);
+  }
+}
+
 /** Sum the byte sizes of `.experts.` tensors from a safetensors header
  *  (header JSON only — never touches tensor data). */
 function expertTensorBytes(path: string): number {
@@ -465,7 +518,10 @@ function visionConfigTypeOf(config: Record<string, any>): string | null {
   return typeof vt === "string" && vt.endsWith("_vision") ? vt : null;
 }
 
-async function scanSnapshot(dir: string, repoId: string): Promise<ModelRecord | null> {
+/** Build a ModelRecord straight from a snapshot directory (also the scan()
+ *  workhorse). Exported for explicit-path loads (`serve --model <dir>`),
+ *  which must not depend on the registry DB having scanned that dir. */
+export async function scanSnapshot(dir: string, repoId: string): Promise<ModelRecord | null> {
   const configPath = join(dir, "config.json");
   if (!existsSync(configPath)) return null;
 
@@ -533,6 +589,8 @@ async function scanSnapshot(dir: string, repoId: string): Promise<ModelRecord |
     quantMode: quant?.mode ?? (quant ? "affine" : null),
     hasVisionSidecar: existsSync(join(dir, "optiq_vision.safetensors")),
     visionConfigType: visionConfigTypeOf(config),
+    hasAudioConfig: !!config.audio_config,
+    hasAudioTower: sidecarShipsAudioTower(join(dir, "optiq_vision.safetensors")),
     hasKvConfig: existsSync(join(dir, "kv_config.json")),
     hasToolTemplate,
     numLayers: text.num_hidden_layers ?? null,