mlx-bun 0.0.10 → 0.0.11

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (173) hide show
  1. package/README.md +66 -21
  2. package/THIRD_PARTY_LICENSES.md +1 -0
  3. package/docs/archive/mlx-bun-lab-report.html +275 -0
  4. package/docs/archive/wwdc-mlx-bun.html +848 -0
  5. package/docs/design/audio-input-plan.md +356 -0
  6. package/docs/design/batching-perf-path.md +58 -35
  7. package/docs/design/batching-v2-plan.md +23 -9
  8. package/docs/design/bucketing-stage.md +2 -1
  9. package/docs/design/compat-cli-surface-design.md +4 -3
  10. package/docs/design/decode-speed-program.md +172 -0
  11. package/docs/design/docs-reorg-plan.md +1 -1
  12. package/docs/design/dspark-serving-program.md +371 -0
  13. package/docs/design/dspark-speculative-decoding.md +167 -5
  14. package/docs/design/faithful-l1-consolidation.md +310 -0
  15. package/docs/design/generic-model-support.md +8 -5
  16. package/docs/design/grammar-spec-batching-integration.md +301 -0
  17. package/docs/design/minicpm5-decode-megakernel.md +16 -11
  18. package/docs/design/mlx-lm-tool-parity-plan.md +11 -1
  19. package/docs/design/omlx-adoption-map.md +14 -6
  20. package/docs/design/orpo-training.md +34 -30
  21. package/docs/design/paged-kv-cache.md +214 -0
  22. package/docs/design/parallel-slots.md +82 -49
  23. package/docs/design/parity-tier-dag.md +43 -33
  24. package/docs/design/repo-cleanup-plan.md +189 -0
  25. package/docs/design/runtime-isolation.md +19 -3
  26. package/docs/design/segmented-backward-training.md +3 -2
  27. package/docs/design/ssd-kv-cold-tier.md +130 -0
  28. package/docs/design/structured-output.md +484 -0
  29. package/docs/design/the-dreaming-master-plan.md +2 -3
  30. package/docs/design/turboquant-kv.md +235 -0
  31. package/docs/design/unified-engine-frontier-plan.md +774 -0
  32. package/docs/design/web-chat-beat-matrix.md +315 -0
  33. package/docs/design/web-chat-redesign.md +1112 -0
  34. package/docs/investigations/dspark-handoff.md +182 -16
  35. package/docs/investigations/lab-build-journal.md +2 -2
  36. package/docs/investigations/openwiki-evaluation.md +249 -0
  37. package/docs/planning/IDEAS.md +14 -0
  38. package/docs/planning/PRODUCT_ROADMAP.md +2 -3
  39. package/docs/planning/chat-ui-vision-memory-first.md +199 -0
  40. package/docs/planning/chat-ui-vision.md +50 -0
  41. package/docs/planning/memory-docs-and-dag-plan.md +15 -6
  42. package/docs/planning/release-notes-v0.0.11.md +239 -0
  43. package/docs/reference/cli.md +35 -16
  44. package/docs/reference/distribution.md +13 -6
  45. package/docs/reference/embedding.md +35 -14
  46. package/docs/reference/features-matrix.md +97 -0
  47. package/docs/reference/library-api.md +30 -9
  48. package/docs/reference/memory.md +26 -6
  49. package/docs/reference/models.md +11 -1
  50. package/docs/reference/orpo-quickstart.md +3 -4
  51. package/docs/reference/server-api.md +634 -22
  52. package/docs/reference/server-config.md +250 -121
  53. package/docs/reference/training.md +14 -15
  54. package/package.json +5 -3
  55. package/src/anthropic.ts +13 -0
  56. package/src/audio/conformer.ts +647 -0
  57. package/src/audio/decode.ts +133 -0
  58. package/src/audio/features.ts +202 -0
  59. package/src/audio/transcode.ts +67 -0
  60. package/src/cli.ts +356 -105
  61. package/src/config.ts +49 -2
  62. package/src/download.ts +40 -0
  63. package/src/eval/README.md +7 -6
  64. package/src/eval/kl.ts +59 -3
  65. package/src/eval/quality-db.ts +2 -2
  66. package/src/eval/runner.ts +3 -3
  67. package/src/expert-offload.ts +15 -2
  68. package/src/flags.ts +13 -0
  69. package/src/generate.ts +330 -24
  70. package/src/grammar.ts +585 -0
  71. package/src/hub-rest.ts +242 -0
  72. package/src/kv-store.ts +442 -88
  73. package/src/lora.ts +8 -0
  74. package/src/md-modules.d.ts +42 -0
  75. package/src/media-fetch.ts +284 -0
  76. package/src/memory/rest.ts +245 -0
  77. package/src/memory/vault.ts +74 -0
  78. package/src/mlx/array.ts +106 -19
  79. package/src/mlx/autograd.ts +2 -2
  80. package/src/mlx/ffi.ts +27 -3
  81. package/src/mlx/metal-kernel.ts +7 -4
  82. package/src/mlx/ops.ts +153 -23
  83. package/src/mlx/turboquant-ops.ts +831 -0
  84. package/src/mlx/turboquant-tables.ts +237 -0
  85. package/src/model/batched-mask.ts +62 -1
  86. package/src/model/batched-quant.ts +206 -0
  87. package/src/model/batched-rotating-quant.ts +346 -0
  88. package/src/model/batched-rotating.ts +46 -1
  89. package/src/model/compiled-decode.ts +4 -15
  90. package/src/model/factory.ts +10 -4
  91. package/src/model/gemma4-base.ts +363 -39
  92. package/src/model/gemma4.ts +122 -33
  93. package/src/model/generated/gemma4-12b.ts +9 -13
  94. package/src/model/generated/gemma4-26b.ts +10 -5
  95. package/src/model/generated/gemma4-e4b.ts +14 -9
  96. package/src/model/minicpm5.ts +58 -102
  97. package/src/model/paged-kv.ts +253 -0
  98. package/src/model/qwen3-delta.ts +76 -15
  99. package/src/model/qwen3-moe.ts +404 -0
  100. package/src/model/qwen3.ts +9 -8
  101. package/src/model/qwen3_5.ts +88 -23
  102. package/src/model/support.ts +12 -1
  103. package/src/model/universal/dense.ts +40 -5
  104. package/src/model/universal/rope.ts +36 -21
  105. package/src/pi-web.ts +1299 -82
  106. package/src/preflight.ts +27 -8
  107. package/src/prompt-cache.ts +293 -33
  108. package/src/quantize/quantizer.ts +16 -1
  109. package/src/registry.ts +57 -2
  110. package/src/serve/batch-scheduler.ts +965 -72
  111. package/src/serve/generation-gateway.ts +214 -11
  112. package/src/serve/isolate.ts +417 -0
  113. package/src/serve/lane-registry.ts +44 -0
  114. package/src/serve/session-search.ts +326 -0
  115. package/src/server.ts +1380 -135
  116. package/src/spec/assistant-source.ts +150 -0
  117. package/src/spec/deepspec-source.ts +143 -0
  118. package/src/spec/dflash-source.ts +124 -0
  119. package/src/spec/dspark/ab-stats.ts +136 -0
  120. package/src/spec/dspark/calibration.ts +112 -0
  121. package/src/spec/dspark/deepspec-module.ts +801 -0
  122. package/src/spec/dspark/generate-dflash.ts +72 -20
  123. package/src/spec/dspark/loader.ts +23 -0
  124. package/src/spec/dspark/module-dflash.ts +278 -36
  125. package/src/spec/dspark/quantize-drafter.ts +91 -0
  126. package/src/spec/generate.ts +59 -25
  127. package/src/spec/ngram-source.ts +132 -0
  128. package/src/spec/serve-loop.ts +536 -0
  129. package/src/spec/source.ts +93 -0
  130. package/src/spec/two-model.ts +171 -0
  131. package/src/ssd-cache.ts +107 -35
  132. package/src/tokenizer.ts +288 -2
  133. package/src/tool-approvals.ts +131 -0
  134. package/src/tool-call.ts +408 -20
  135. package/src/train/flash-cce.ts +2 -2
  136. package/src/train/job.ts +0 -2
  137. package/src/train/trainer.ts +1 -4
  138. package/src/vision/prompt.ts +217 -48
  139. package/src/web/app.html +1184 -1789
  140. package/src/web/app.js +5667 -0
  141. package/src/web/icon.svg +13 -0
  142. package/src/web/manifest.webmanifest +24 -0
  143. package/src/web/src/adapters-panel.ts +288 -0
  144. package/src/web/src/api.ts +57 -0
  145. package/src/web/src/assistant.ts +485 -0
  146. package/src/web/src/chat.ts +1050 -0
  147. package/src/web/src/composer.ts +959 -0
  148. package/src/web/src/dataset.ts +132 -0
  149. package/src/web/src/finetune.ts +267 -0
  150. package/src/web/src/hub.ts +428 -0
  151. package/src/web/src/main.ts +98 -0
  152. package/src/web/src/markdown.ts +0 -0
  153. package/src/web/src/memory-panel.ts +572 -0
  154. package/src/web/src/model-picker.ts +159 -0
  155. package/src/web/src/palette.ts +346 -0
  156. package/src/web/src/protocol.ts +62 -0
  157. package/src/web/src/quantize.ts +167 -0
  158. package/src/web/src/rag.ts +274 -0
  159. package/src/web/src/sessions.ts +469 -0
  160. package/src/web/src/shell.ts +804 -0
  161. package/src/web/src/status.ts +252 -0
  162. package/src/web/src/ui-catalog.ts +92 -0
  163. package/src/web/sw.js +71 -0
  164. package/src/web/vendor/README.md +61 -0
  165. package/src/web/vendor/hljs-LICENSE +29 -0
  166. package/src/web/vendor/hljs-theme.css +59 -0
  167. package/src/web/vendor/hljs.js +26 -0
  168. package/src/model/fused-decode-kernel.ts +0 -312
  169. package/src/model/fused-geglu-kernel.ts +0 -198
  170. package/src/model/fused-mlp-kernel.ts +0 -169
  171. package/src/model/fused-swiglu-kernel.ts +0 -240
  172. package/src/model/steel-linear-kernel.ts +0 -96
  173. package/src/status-page.html +0 -497
package/src/preflight.ts CHANGED
@@ -5,6 +5,16 @@
5
5
  // - cold page cache: prefill timer dominated by SSD page-in
6
6
  // The harness refuses to record headline numbers when checks fail;
7
7
  // the state snapshot is stored in every eval-DB row either way.
8
+ //
9
+ // LIMITS of this instrument (2026-07-05 finding): the snapshot is taken
10
+ // BETWEEN legs, when the bench process has exited and memory is quiet —
11
+ // free_pct reads a near-constant ~92% on an otherwise-idle box and load1m
12
+ // sat under every threshold while a 40-minute mid-pass slow window cut
13
+ // decode rates ~40% across every stack. Preflight ok ≠ clean measurements
14
+ // and preflight-failed ≠ dirty ones (the 07-04 pass was ‡-tagged yet
15
+ // stable). The DIRECT quality signal is per-cell run spread — see the
16
+ // run-stability policy in scripts/bench-h2h.ts (SPREAD_TOL retries +
17
+ // `unstable`/`stabilized` tags). Preflight remains the coarse pre-gate.
8
18
 
9
19
  export interface MachineState {
10
20
  swapUsedMB: number;
@@ -77,21 +87,30 @@ export function checkMachine(limits: PreflightLimits = {}): MachineState {
77
87
  if (loadAvg1m > 8)
78
88
  problems.push(`1-min load ${loadAvg1m.toFixed(1)} — post-boot indexing still running; wait a minute`);
79
89
 
80
- // big foreign processes (rss in KB from ps)
90
+ // big foreign processes: high RSS *or* high CPU (pcpu + rss in KB from
91
+ // ps). The CPU arm is the 2026-07-05 lesson: Apple's
92
+ // knowledgeconstructiond (Apple Intelligence indexing) burned 87% CPU
93
+ // with a tiny RSS during the pass's 40-min slow window — a background
94
+ // ML workload the RSS-only check could never flag. CPU is bursty, so
95
+ // one clean sample doesn't prove absence (the run-spread stability
96
+ // policy in bench-h2h.ts is the direct signal); a hot sample here is
97
+ // still a cheap early warning.
81
98
  const SELF_PATTERN = /(bun|\.venv\/bin\/python|mlx_lm|optiq)/;
82
99
  const SYSTEM_PATTERN = /(kernel_task|WindowServer|launchd|mds|spotlight)/i;
100
+ const BIG_CPU_PCT = 50;
83
101
  const bigProcesses: { rssMB: number; command: string }[] = [];
84
- for (const line of sh(["ps", "axo", "rss=,command="]).split("\n")) {
85
- const m = line.match(/^\s*(\d+)\s+(.*)$/);
102
+ for (const line of sh(["ps", "axo", "pcpu=,rss=,command="]).split("\n")) {
103
+ const m = line.match(/^\s*([\d.]+)\s+(\d+)\s+(.*)$/);
86
104
  if (!m) continue;
87
- const rssMB = Number(m[1]) / 1024;
88
- const command = m[2]!.slice(0, 80);
89
- if (rssMB < bigRss) continue;
105
+ const cpuPct = Number(m[1]);
106
+ const rssMB = Number(m[2]) / 1024;
107
+ const command = m[3]!.slice(0, 80);
108
+ if (rssMB < bigRss && cpuPct < BIG_CPU_PCT) continue;
90
109
  if (SELF_PATTERN.test(command) || SYSTEM_PATTERN.test(command)) continue;
91
- bigProcesses.push({ rssMB: Math.round(rssMB), command });
110
+ bigProcesses.push({ rssMB: Math.round(rssMB), command: cpuPct >= BIG_CPU_PCT ? `${command} (${cpuPct.toFixed(0)}% cpu)` : command });
92
111
  }
93
112
  for (const p of bigProcesses)
94
- problems.push(`big process resident: ${p.rssMB} MB — ${p.command}`);
113
+ problems.push(`big process: ${p.rssMB} MB — ${p.command}`);
95
114
 
96
115
  return {
97
116
  swapUsedMB, freePercent, cpuSpeedLimit, gpuWiredLimitMB, loadAvg1m,
@@ -1,15 +1,93 @@
1
- // Byte-capped LRU prompt cache.
1
+ // Byte-capped LRU prompt cache — the RAM tier of the Layer-0 KV store.
2
2
  //
3
3
  // The mlx-lm lesson (PLAN.md): a count-capped cache of multi-GB KV
4
4
  // entries is an OOM footgun. Ours accounts bytes (sum of KV array bytes
5
5
  // per entry) and evicts least-recently-used until under the cap.
6
6
  //
7
- // Usage pattern (single generation queue, so take/reinsert is race-free):
8
- // const hit = cache.take(promptIds); // longest strict-prefix match
7
+ // Usage pattern (single generation queue, so take/put is race-free):
8
+ // const hit = cache.take(promptIds); // longest-prefix match → CLONES
9
9
  // generate(model, promptIds, { cache: hit?.caches ?? fresh })
10
- // cache.put([...promptIds, ...generated], caches); // extended entry
10
+ // cache.put([...promptIds, ...generated], caches, ns, hit?.retain);
11
+ // take() is NON-CONSUMING (prefix sharing): it serves zero-copy clones and
12
+ // leaves the donor entry in place — N agents sharing a system prompt all
13
+ // clone from one donor, one prefill; put() supersedes same-ns
14
+ // prefix-ancestors (when the new entry is trimmable) so a conversation
15
+ // stays one entry, not one per turn.
16
+ //
17
+ // TIERING (Layer 0, unified-engine plan): when a ColdTier is attached,
18
+ // take() itself runs the two-tier dance — RAM peek vs cold find, restore
19
+ // (zero-copy mmap) + trim when the cold tier holds a strictly longer
20
+ // prefix — so EVERY consumer (the serial lane, the batch scheduler, future
21
+ // prefix sharing) gets SSD restores through the same take()/put() it
22
+ // already calls. Eviction spills to the tier (the #spill hook); idle
23
+ // entries DEMOTE to it (demoteIdle — free the GPU memory, keep the prefix
24
+ // reachable); onPut lets the server schedule its debounced write-behind
25
+ // snapshot for both lanes.
26
+
27
+ import { TurboQuantKVCache, type Cache } from "./model/gemma4";
28
+ import { cloneKvCaches } from "./kv-store";
29
+
30
+ /** Reference-counted release: wraps an entry's `retain` (e.g. an mmap
31
+ * unmap) so DONOR + CLONES can each hold a share — the underlying hook
32
+ * runs once, after the LAST holder releases. Prefix sharing hands out
33
+ * zero-copy VIEWS of a donor's buffers; if the donor's unmap ran while a
34
+ * clone still referenced the mapped pages, the clone would read freed
35
+ * memory. */
36
+ function makeSharedRetain(retain: (() => void) | undefined): { acquire(): () => void } {
37
+ let count = 0;
38
+ let released = false;
39
+ return {
40
+ acquire() {
41
+ count++;
42
+ let mine = false;
43
+ return () => {
44
+ if (mine) return; // idempotent per share
45
+ mine = true;
46
+ count--;
47
+ if (count === 0 && !released) {
48
+ released = true;
49
+ retain?.();
50
+ }
51
+ };
52
+ },
53
+ };
54
+ }
55
+
56
+ /** The cold (SSD) tier the RAM cache tiers over. Structural — this module
57
+ * never imports ssd-cache (which imports us); the server binds
58
+ * SsdCacheStore + the model into this shape (src/server.ts). All methods
59
+ * must be failure-proof: any error degrades to "no hit"/"not stored". */
60
+ export interface ColdTier {
61
+ /** Longest stored usable prefix for prompt/ns — index-only, no I/O.
62
+ * `handle` is the tier's opaque entry token, passed back to restore. */
63
+ find(prompt: number[], ns: string): { prefixLen: number; handle: unknown } | null;
64
+ /** Materialize a found entry as GPU-visible caches (zero-copy COW mmap;
65
+ * pages fault in lazily). `retain` must run after the caches are
66
+ * disposed (it unmaps the backing file). Null on any failure. */
67
+ restore(handle: unknown): { tokens: number[]; caches: Cache[]; retain: () => void } | null;
68
+ store(tokens: number[], caches: Cache[], ns: string): void;
69
+ }
11
70
 
12
- import type { Cache } from "./model/gemma4";
71
+ /** Spill sink for eviction/demotion (write-behind, 2026-07-06).
72
+ *
73
+ * `spillOwned` (preferred, NON-BLOCKING): the cache hands the sink an
74
+ * entry it OWNS — zero-copy CLONES of the caches (made via the same
75
+ * injectable cloner take() uses, BEFORE the donor is disposed; entries
76
+ * are immutable so the clones stay consistent forever) plus a copy of
77
+ * the tokens. The sink must dispose entry.caches when done, on BOTH
78
+ * fulfill and reject paths (the server chains storeAsync -> dispose on
79
+ * ssdWriteChain). A sink that throws SYNCHRONOUSLY has not taken
80
+ * ownership — the cache disposes the clones and the spill degrades to
81
+ * a no-op; eviction/demotion never unwinds.
82
+ *
83
+ * `spillSync` (fallback): called with the LIVE entry about to be
84
+ * disposed; must finish reading before returning. This is the old
85
+ * contract — a full synchronous serialize+write under the generation
86
+ * lock — kept for tests and simple embedders. */
87
+ export interface SpillSink {
88
+ spillOwned?: (entry: PromptCacheEntry) => void;
89
+ spillSync?: (entry: PromptCacheEntry) => void;
90
+ }
13
91
 
14
92
  export interface PromptCacheEntry {
15
93
  tokens: number[];
@@ -24,9 +102,18 @@ export interface PromptCacheEntry {
24
102
  retain?: () => void;
25
103
  }
26
104
 
27
- function cacheBytes(caches: Cache[]): number {
105
+ export function cacheBytes(caches: Cache[]): number {
28
106
  let total = 0;
29
- for (const c of caches) for (const a of c.state()) total += a.nbytes;
107
+ for (const c of caches) {
108
+ const state = c.state();
109
+ for (const a of state) total += a.nbytes;
110
+ // Most Cache kinds' state() returns their own live-owned arrays (safe
111
+ // to just read), but TurboQuantKVCache allocates fresh trimmed slice
112
+ // views per call (kv-store.ts snapshotCache/cloneKvCaches contract) —
113
+ // those must be disposed here or they leak (see generate.ts
114
+ // evalCacheState for the same hazard on the prefill path).
115
+ if (c instanceof TurboQuantKVCache) for (const a of state) a.dispose();
116
+ }
30
117
  return total;
31
118
  }
32
119
 
@@ -38,25 +125,82 @@ export function commonPrefixLength(a: number[], b: number[]): number {
38
125
  return i;
39
126
  }
40
127
 
128
+ interface EntryRecord {
129
+ entry: PromptCacheEntry;
130
+ bytes: number;
131
+ lastUsed: number;
132
+ lastUsedMs: number;
133
+ /** Ref-counted retain: the donor holds one share (entry.retain); every
134
+ * clone handed out by take() holds another. */
135
+ share: { acquire(): () => void };
136
+ }
137
+
41
138
  export class PromptCache {
42
139
  readonly maxBytes: number;
43
- #entries: { entry: PromptCacheEntry; bytes: number; lastUsed: number }[] = [];
140
+ #entries: EntryRecord[] = [];
44
141
  #clock = 0;
45
142
  hits = 0;
46
143
  misses = 0;
47
- /** Cold-tier spill hook: called with an entry ABOUT to be evicted (still
48
- * alive — the hook may serialize its caches); the cache disposes it after.
49
- * Failures are the hook's problem; eviction proceeds regardless. */
50
- readonly #spill: ((entry: PromptCacheEntry) => void) | null;
144
+ demotions = 0;
145
+ /** Cold-tier spill sinks (see SpillSink): spillOwned gets clones and
146
+ * runs the write off-lock; spillSync gets the live entry inline.
147
+ * Failures are the sink's problem; eviction proceeds regardless. */
148
+ readonly #spillOwned: ((entry: PromptCacheEntry) => void) | null;
149
+ readonly #spillSync: ((entry: PromptCacheEntry) => void) | null;
150
+ /** Optional cold tier: take() tiers over it, demoteIdle() spills into it. */
151
+ readonly #cold: ColdTier | null;
152
+ /** Fired after every successful put() — the server hangs its debounced
153
+ * write-behind SSD snapshot here so BOTH lanes' entries persist. */
154
+ onPut: ((tokens: number[], ns: string) => void) | null = null;
51
155
 
52
- constructor(maxBytes: number, spill: ((entry: PromptCacheEntry) => void) | null = null) {
156
+ /** Zero-copy view cloner — injectable so model-free tests can stub it. */
157
+ readonly #clone: (caches: Cache[]) => Cache[];
158
+
159
+ constructor(
160
+ maxBytes: number,
161
+ // Bare function = legacy sync spill (tests, simple embedders).
162
+ spill: SpillSink | ((entry: PromptCacheEntry) => void) | null = null,
163
+ cold: ColdTier | null = null,
164
+ clone: (caches: Cache[]) => Cache[] = cloneKvCaches,
165
+ ) {
53
166
  this.maxBytes = maxBytes;
54
- this.#spill = spill;
167
+ this.#spillOwned = typeof spill === "function" ? null : spill?.spillOwned ?? null;
168
+ this.#spillSync = typeof spill === "function" ? spill : spill?.spillSync ?? null;
169
+ this.#cold = cold;
170
+ this.#clone = clone;
55
171
  }
56
172
 
57
173
  #disposeEntry(entry: PromptCacheEntry, spillFirst: boolean): void {
58
- if (spillFirst && this.#spill) {
59
- try { this.#spill(entry); } catch {}
174
+ if (spillFirst) {
175
+ if (this.#spillOwned) {
176
+ // NON-BLOCKING spill: clone BEFORE disposing (clones are zero-copy
177
+ // views of the live arrays) and hand ownership to the sink, which
178
+ // flushes off the generation lock. The underlying GPU buffers stay
179
+ // alive until the sink disposes the clones — for demoteIdle that
180
+ // means the memory is freed one bounded write-chain hop later, not
181
+ // instantly; acceptable because the chain is serial and each write
182
+ // disposes its clones on settle (fulfill AND reject).
183
+ // Retain contract: entry.retain still runs below, possibly while
184
+ // the clones are in flight — safe because cold-restored entries
185
+ // OWN their bytes (streamed copy-restore, 2026-07-07): there is
186
+ // no backing mmap for retain to unmap, so nothing the clones
187
+ // could alias goes away.
188
+ let clones: Cache[] | null = null;
189
+ try {
190
+ clones = this.#clone(entry.caches);
191
+ this.#spillOwned({ tokens: [...entry.tokens], caches: clones, ns: entry.ns });
192
+ } catch {
193
+ // Failure degrades to no-spill; eviction/demotion never unwinds.
194
+ if (clones) for (const c of clones) try { c.dispose(); } catch {}
195
+ }
196
+ } else {
197
+ const spill =
198
+ this.#spillSync ??
199
+ (this.#cold
200
+ ? (e: PromptCacheEntry) => this.#cold!.store(e.tokens, e.caches, e.ns)
201
+ : null);
202
+ if (spill) try { spill(entry); } catch {}
203
+ }
60
204
  }
61
205
  for (const c of entry.caches) c.dispose();
62
206
  entry.retain?.();
@@ -70,15 +214,26 @@ export class PromptCache {
70
214
  return this.#entries.length;
71
215
  }
72
216
 
73
- /** Find the entry with the longest usable common prefix of `prompt`,
74
- * remove it, trim its caches to that prefix if needed (and possible),
75
- * and hand over ownership (caller must put() it back — possibly
76
- * extended — or dispose it).
217
+ /** Find the entry with the longest usable common prefix of `prompt`
218
+ * across BOTH tiers and serve it NON-CONSUMINGLY (prefix sharing,
219
+ * 2026-07-05): the caller receives zero-copy CLONES trimmed to the
220
+ * matched prefix (plus a ref-counted retain share); the donor entry
221
+ * stays in the cache, ready for the next agent/session with the same
222
+ * prefix. The caller owns the returned caches — dispose them or put()
223
+ * an extended entry (which supersedes prefix-ancestors), honoring
224
+ * `retain` either way.
77
225
  *
78
226
  * Usable prefix = common prefix capped at prompt.length - 1 (at least
79
227
  * one token must be forwarded to produce logits). Entries longer than
80
228
  * the prefix need cache.trim(); ring caches lose trimability once
81
- * wrapped — those entries only match in full. */
229
+ * wrapped — those entries only match in full.
230
+ *
231
+ * Tier order: the cold tier wins only with a STRICTLY longer usable
232
+ * prefix (its restore costs a mmap + lazy fault-in; RAM is free). A
233
+ * cold entry with an untrimmable divergent tail (ring post-wrap, SSM)
234
+ * is dropped and the RAM candidate (or a fresh prefill) serves.
235
+ * hits/misses count RAM candidacy only (cold restores are counted by
236
+ * the tier itself), preserving the /stats meaning. */
82
237
  take(prompt: number[], ns = ""): PromptCacheEntry | null {
83
238
  let bestIdx = -1;
84
239
  let bestLen = 0;
@@ -92,18 +247,64 @@ export class PromptCache {
92
247
  bestLen = p;
93
248
  bestIdx = i;
94
249
  }
95
- if (bestIdx === -1) {
96
- this.misses++;
97
- return null;
98
- }
99
- this.hits++;
100
- const { entry } = this.#entries.splice(bestIdx, 1)[0]!;
101
- const trimNeeded = entry.tokens.length - bestLen;
102
- if (trimNeeded > 0) {
103
- for (const c of entry.caches) c.trim(trimNeeded);
104
- entry.tokens = entry.tokens.slice(0, bestLen);
250
+ if (bestIdx === -1) this.misses++;
251
+ else this.hits++;
252
+
253
+ // Cold tier: strictly longer usable prefix ⇒ restore + trim; the RAM
254
+ // candidate stays put UNTOUCHED (the old serial-lane flow trimmed it
255
+ // before comparing — a needless entry degradation, gone now).
256
+ if (this.#cold) {
257
+ const hit = this.#cold.find(prompt, ns);
258
+ if (hit && hit.prefixLen > bestLen) {
259
+ const loaded = this.#cold.restore(hit.handle);
260
+ if (loaded) {
261
+ const trimNeeded = loaded.tokens.length - hit.prefixLen;
262
+ if (trimNeeded > 0 && !loaded.caches.every((c) => c.isTrimmable())) {
263
+ // Divergent tail on an untrimmable kind: this file can't seed
264
+ // the prompt — fall back to the RAM candidate. Loud: this used
265
+ // to be the silent 84 s "restore-then-re-prefill" path (the
266
+ // whole cold tier degrading with zero observables, 2026-07-06).
267
+ console.warn(
268
+ `[prompt-cache] cold entry unusable: ${loaded.tokens.length} tokens, ` +
269
+ `needs trim ${trimNeeded} but untrimmable (wrapped ring/SSM) — re-prefilling`,
270
+ );
271
+ for (const c of loaded.caches) c.dispose();
272
+ loaded.retain();
273
+ } else {
274
+ if (trimNeeded > 0) for (const c of loaded.caches) c.trim(trimNeeded);
275
+ return {
276
+ tokens: loaded.tokens.slice(0, hit.prefixLen),
277
+ caches: loaded.caches,
278
+ ns,
279
+ retain: loaded.retain,
280
+ };
281
+ }
282
+ }
283
+ }
105
284
  }
106
- return entry;
285
+
286
+ if (bestIdx === -1) return null;
287
+ // PREFIX SHARING (non-consuming serve): the caller gets zero-copy
288
+ // CLONES of the donor's caches, trimmed to the matched prefix; the
289
+ // donor stays in the cache untouched, ready to serve the next agent /
290
+ // session with the same prefix. Safe because mlx cache updates are
291
+ // FUNCTIONAL (updateAndFetch reassigns fresh arrays; buffer donation
292
+ // only fires at refcount 1, and the donor always holds a ref), so a
293
+ // clone extending itself can never mutate the donor's bytes. The old
294
+ // consume-and-trim semantics CANNIBALIZED donors: agent B borrowing
295
+ // agent A's 2k system prompt destroyed A's 10k entry to do it.
296
+ const rec = this.#entries[bestIdx]!;
297
+ rec.lastUsed = ++this.#clock;
298
+ rec.lastUsedMs = Date.now();
299
+ const clones = this.#clone(rec.entry.caches);
300
+ const trimNeeded = rec.entry.tokens.length - bestLen;
301
+ if (trimNeeded > 0) for (const c of clones) c.trim(trimNeeded);
302
+ return {
303
+ tokens: rec.entry.tokens.slice(0, bestLen),
304
+ caches: clones,
305
+ ns,
306
+ retain: rec.share.acquire(),
307
+ };
107
308
  }
108
309
 
109
310
  /** Longest usable prefix length available for `prompt` WITHOUT taking
@@ -142,7 +343,42 @@ export class PromptCache {
142
343
  this.#disposeEntry(entry, true);
143
344
  return;
144
345
  }
145
- this.#entries.push({ entry, bytes, lastUsed: ++this.#clock });
346
+ // The donor's retain becomes ref-counted so take()'s clones can hold
347
+ // shares; the entry's own retain is the donor's share.
348
+ const share = makeSharedRetain(retain);
349
+ entry.retain = share.acquire();
350
+ const rec: EntryRecord = { entry, bytes, lastUsed: ++this.#clock, lastUsedMs: Date.now(), share };
351
+ this.#entries.push(rec);
352
+ // Exact duplicates (equal tokens) are redundant REGARDLESS of
353
+ // trimmability — the new entry serves exactly the matches the old one
354
+ // did (mlx-lm's trie replaces in place). Without this, every gemma ctx
355
+ // repeat left another full-size wrapped-ring entry behind (measured
356
+ // ~0.4 GB retained per repeat on the 12B bench).
357
+ for (const old of [...this.#entries]) {
358
+ if (old === rec || old.entry.ns !== ns) continue;
359
+ if (old.entry.tokens.length !== tokens.length) continue;
360
+ if (commonPrefixLength(old.entry.tokens, tokens) !== tokens.length) continue;
361
+ this.#entries = this.#entries.filter((r) => r !== old);
362
+ this.#disposeEntry(old.entry, false);
363
+ }
364
+ // SUPERSEDE strict prefix-ancestors (the SSD store's rule, brought to
365
+ // RAM — prefix sharing would otherwise grow an entry per turn): same-ns
366
+ // entries whose tokens are a prefix of the new entry are redundant
367
+ // and disposed WITHOUT spill (the new entry serves their prefixes and
368
+ // its own write-behind persists them) — but ONLY when the new entry's
369
+ // caches are all trimmable. An untrimmable new entry (wrapped ring)
370
+ // can only serve exact-length matches, so shorter ancestors — e.g.
371
+ // the prompt-boundary snapshot that makes drifted multi-turn agents
372
+ // hit — must survive.
373
+ if (caches.every((c) => c.isTrimmable())) {
374
+ for (const old of [...this.#entries]) {
375
+ if (old === rec || old.entry.ns !== ns) continue;
376
+ if (old.entry.tokens.length >= tokens.length) continue;
377
+ if (commonPrefixLength(old.entry.tokens, tokens) !== old.entry.tokens.length) continue;
378
+ this.#entries = this.#entries.filter((r) => r !== old);
379
+ this.#disposeEntry(old.entry, false);
380
+ }
381
+ }
146
382
  while (this.totalBytes > this.maxBytes && this.#entries.length > 1) {
147
383
  let lruIdx = 0;
148
384
  for (let i = 1; i < this.#entries.length; i++)
@@ -150,6 +386,30 @@ export class PromptCache {
150
386
  const [evicted] = this.#entries.splice(lruIdx, 1);
151
387
  this.#disposeEntry(evicted!.entry, true);
152
388
  }
389
+ try { this.onPut?.(tokens, ns); } catch {}
390
+ }
391
+
392
+ /** Layer-0 idle demotion: spill every entry unused for `idleMs` to the
393
+ * cold tier and FREE its GPU memory. The prefix stays reachable — the
394
+ * next take() restores it via zero-copy mmap (~0.25 s for a 13.7k-token
395
+ * entry, vs a 12 s re-prefill) — but between bursts the RAM tier drains
396
+ * toward empty, returning unified memory to the system. Caller must
397
+ * hold the generation lock (entries' arrays are disposed here) — with
398
+ * a spillOwned sink the lock covers only the zero-copy clone; the
399
+ * write runs off-lock and the buffers free when it settles. No-op
400
+ * without a cold tier — demotion without a place to demote TO would
401
+ * just be data loss. Returns entries demoted. */
402
+ demoteIdle(idleMs: number, now = Date.now()): number {
403
+ if (!this.#cold) return 0;
404
+ let n = 0;
405
+ for (const rec of [...this.#entries]) {
406
+ if (now - rec.lastUsedMs < idleMs) continue;
407
+ this.#entries = this.#entries.filter((r) => r !== rec);
408
+ this.#disposeEntry(rec.entry, true); // spill-first, then dispose
409
+ this.demotions++;
410
+ n++;
411
+ }
412
+ return n;
153
413
  }
154
414
 
155
415
  clear(): void {
@@ -49,6 +49,13 @@ export interface QuantizeOptions {
49
49
  groupSize: 32 | 64;
50
50
  /** Quantization scheme; v1 only exercises "affine". */
51
51
  mode?: string;
52
+ /** Extra per-module gate on top of the shape-eligibility check: return
53
+ * false to keep a module full-precision even though its shape qualifies
54
+ * (e.g. the DeepSpec drafter's confidence_head — threshold-sensitive and
55
+ * tiny). Excluded modules are recorded as `false` in the quantization
56
+ * block (mlx's "not quantized" convention). Shape-ineligible modules
57
+ * never reach this. */
58
+ quantizePredicate?: (base: string, fullShape: number[]) => boolean;
52
59
 
53
60
  // --- mixed-precision (OptiQ sensitivity + knapsack) -------------------
54
61
  // Setting targetBpw switches quantizeModelDir to the mixed path: it loads the
@@ -225,7 +232,11 @@ export async function quantizeModelDir(
225
232
  // re-quantized weight it is the dequantized array's shape.
226
233
  const fullShape = alreadyQuant ? materialized!.shape : shape;
227
234
 
228
- if (!isQuantizable(fullShape, groupSize)) {
235
+ const predicateExcluded =
236
+ isQuantizable(fullShape, groupSize) &&
237
+ opts.quantizePredicate !== undefined &&
238
+ !opts.quantizePredicate(base, fullShape);
239
+ if (!isQuantizable(fullShape, groupSize) || predicateExcluded) {
229
240
  // Not eligible: pass the full-precision weight through unchanged
230
241
  // (re-materialized to bf16 if it had been quantized).
231
242
  if (alreadyQuant) {
@@ -235,6 +246,10 @@ export async function quantizeModelDir(
235
246
  } else {
236
247
  out.push({ name: weightName, array: fullWeight });
237
248
  }
249
+ // Policy-excluded (shape-eligible) modules get an explicit `false`
250
+ // in the quantization block — the loader must not apply the block's
251
+ // default spec to them (mlx's "not quantized" convention).
252
+ if (predicateExcluded) perLayer.set(base, false);
238
253
  bumpProgress();
239
254
  continue;
240
255
  }
package/src/registry.ts CHANGED
@@ -34,6 +34,14 @@ export interface ModelRecord {
34
34
  * AND for configs whose nested vision_config is not a vision tower
35
35
  * (Qwen3.5 nests a copy of its own text config there). */
36
36
  visionConfigType: string | null;
37
+ /** config.json declares an `audio_config` block (gemma-4 Conformer). */
38
+ hasAudioConfig: boolean;
39
+ /** The optiq_vision.safetensors sidecar actually ships `audio_tower.*`
40
+ * tensors (header-only check). Distinct from hasAudioConfig: the 12B
41
+ * OptiQ snapshot pairs an audio_config with a STUB sidecar holding only
42
+ * embed_audio.embedding_projection.weight, so config presence alone
43
+ * is not an audio-capability signal. */
44
+ hasAudioTower: boolean;
37
45
  hasKvConfig: boolean;
38
46
  hasToolTemplate: boolean;
39
47
  numLayers: number | null;
@@ -60,6 +68,8 @@ CREATE TABLE IF NOT EXISTS models (
60
68
  quant_mode TEXT,
61
69
  has_vision_sidecar INTEGER NOT NULL,
62
70
  vision_config_type TEXT,
71
+ has_audio_config INTEGER NOT NULL DEFAULT 0,
72
+ has_audio_tower INTEGER NOT NULL DEFAULT 0,
63
73
  has_kv_config INTEGER NOT NULL,
64
74
  has_tool_template INTEGER NOT NULL,
65
75
  num_layers INTEGER,
@@ -97,7 +107,8 @@ export class Registry {
97
107
  !cols.includes("sidecar_bytes") ||
98
108
  !cols.includes("experts_bytes") ||
99
109
  !cols.includes("license") ||
100
- !cols.includes("vision_config_type")
110
+ !cols.includes("vision_config_type") ||
111
+ !cols.includes("has_audio_tower")
101
112
  ) {
102
113
  this.db.exec("DROP TABLE models");
103
114
  this.db.exec(SCHEMA);
@@ -115,7 +126,8 @@ export class Registry {
115
126
  const upsert = this.db.prepare(`
116
127
  INSERT OR REPLACE INTO models VALUES
117
128
  ($path, $repo, $type, $params, $size, $sidecar, $experts, $bits, $gs, $mode,
118
- $vision, $vtype, $kv, $tools, $layers, $hidden, $vocab, $license, $at)
129
+ $vision, $vtype, $audiocfg, $audiotower, $kv, $tools, $layers, $hidden,
130
+ $vocab, $license, $at)
119
131
  `);
120
132
  for (const entry of readdirSync(hubDir)) {
121
133
  if (!entry.startsWith("models--")) continue;
@@ -133,6 +145,8 @@ export class Registry {
133
145
  $bits: rec.quantBits, $gs: rec.quantGroupSize, $mode: rec.quantMode,
134
146
  $vision: rec.hasVisionSidecar ? 1 : 0,
135
147
  $vtype: rec.visionConfigType,
148
+ $audiocfg: rec.hasAudioConfig ? 1 : 0,
149
+ $audiotower: rec.hasAudioTower ? 1 : 0,
136
150
  $kv: rec.hasKvConfig ? 1 : 0,
137
151
  $tools: rec.hasToolTemplate ? 1 : 0,
138
152
  $layers: rec.numLayers, $hidden: rec.hiddenSize, $vocab: rec.vocabSize,
@@ -218,6 +232,15 @@ export function visionCapable(m: Pick<ModelRecord, "hasVisionSidecar" | "visionC
218
232
  return m.hasVisionSidecar || (m.visionConfigType?.endsWith("unified_vision") ?? false);
219
233
  }
220
234
 
235
+ /** True when the model can answer audio (`input_audio`) requests: config.json
236
+ * declares an `audio_config` AND the sidecar actually ships the Conformer
237
+ * (`audio_tower.*`) tensors. Both legs matter — the 12B OptiQ snapshot has
238
+ * audio_config but a stub sidecar (embed_audio only), so config presence
239
+ * alone would advertise a capability that 400s at request time. */
240
+ export function audioCapable(m: Pick<ModelRecord, "hasAudioConfig" | "hasAudioTower">): boolean {
241
+ return m.hasAudioConfig && m.hasAudioTower;
242
+ }
243
+
221
244
  /** Of several cached revisions of ONE repo, pick the canonical snapshot: the
222
245
  * revision refs/main points at, else the most recently scanned (stable
223
246
  * tie-break on path). Paths are `<hub>/models--<repo>/snapshots/<hash>`. */
@@ -250,6 +273,8 @@ function rowToRecord(r: Record<string, unknown>): ModelRecord {
250
273
  quantMode: r.quant_mode as string | null,
251
274
  hasVisionSidecar: !!r.has_vision_sidecar,
252
275
  visionConfigType: (r.vision_config_type as string | null) ?? null,
276
+ hasAudioConfig: !!r.has_audio_config,
277
+ hasAudioTower: !!r.has_audio_tower,
253
278
  hasKvConfig: !!r.has_kv_config,
254
279
  hasToolTemplate: !!r.has_tool_template,
255
280
  numLayers: r.num_layers as number | null,
@@ -280,6 +305,34 @@ function readmeLicense(dir: string): string | null {
280
305
  }
281
306
  }
282
307
 
308
+ /** True when a safetensors file's header names any `audio_tower.*` tensor
309
+ * (header JSON only — never touches tensor data). Separates a real
310
+ * audio-capable sidecar (e4b: 752 audio tensors) from the 12B stub, whose
311
+ * only audio entry is embed_audio.embedding_projection.weight. Missing or
312
+ * unparseable files read as false. */
313
+ export function sidecarShipsAudioTower(path: string): boolean {
314
+ let fd: number;
315
+ try {
316
+ fd = openSync(path, "r");
317
+ } catch {
318
+ return false;
319
+ }
320
+ try {
321
+ const lenBuf = Buffer.alloc(8);
322
+ readSync(fd, lenBuf, 0, 8, 0);
323
+ const headerLen = Number(lenBuf.readBigUInt64LE(0));
324
+ if (headerLen <= 0 || headerLen > 256 * 2 ** 20) return false;
325
+ const headerBuf = Buffer.alloc(headerLen);
326
+ readSync(fd, headerBuf, 0, headerLen, 8);
327
+ const header = JSON.parse(headerBuf.toString("utf8")) as Record<string, unknown>;
328
+ return Object.keys(header).some((n) => n.startsWith("audio_tower."));
329
+ } catch {
330
+ return false;
331
+ } finally {
332
+ closeSync(fd);
333
+ }
334
+ }
335
+
283
336
  /** Sum the byte sizes of `.experts.` tensors from a safetensors header
284
337
  * (header JSON only — never touches tensor data). */
285
338
  function expertTensorBytes(path: string): number {
@@ -536,6 +589,8 @@ export async function scanSnapshot(dir: string, repoId: string): Promise<ModelRe
536
589
  quantMode: quant?.mode ?? (quant ? "affine" : null),
537
590
  hasVisionSidecar: existsSync(join(dir, "optiq_vision.safetensors")),
538
591
  visionConfigType: visionConfigTypeOf(config),
592
+ hasAudioConfig: !!config.audio_config,
593
+ hasAudioTower: sidecarShipsAudioTower(join(dir, "optiq_vision.safetensors")),
539
594
  hasKvConfig: existsSync(join(dir, "kv_config.json")),
540
595
  hasToolTemplate,
541
596
  numLayers: text.num_hidden_layers ?? null,