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/server.ts CHANGED
@@ -16,22 +16,56 @@ import type { Server } from "bun";
16
16
  // the runtime value a string — hence the double cast.
17
17
  import appHtml from "./web/app.html" with { type: "text" };
18
18
  import curveDesignerHtml from "./assets/curve-designer.html" with { type: "text" };
19
+ // Vendored, no-CDN static assets referenced by app.html (convention: any
20
+ // self-contained JS/CSS too big to inline gets `with { type: "text" }`
21
+ // imported here and served under /assets/<name>; see src/web/vendor/README
22
+ // for how hljs.js was built). Add new vendored assets the same way.
23
+ import hljsJs from "./web/vendor/hljs.js" with { type: "text" };
24
+ import hljsCss from "./web/vendor/hljs-theme.css" with { type: "text" };
25
+ // The frontend bundle (plan §7/§9 Phase 2 module split): GENERATED from
26
+ // src/web/src/*.ts by `bun scripts/build-web.ts` — see that file's header
27
+ // and tests/web-build.test.ts (the freshness gate). Same
28
+ // with { type: "text" } + /assets/<name> pattern as the vendored assets
29
+ // above; app.html's <script defer src="/assets/app.js"> loads it.
30
+ import appJs from "./web/app.js" with { type: "text" };
31
+ // PWA installability (plan §9 Phase 3, beat-matrix Axis 10): a manifest +
32
+ // a single inline SVG icon (no binary PNGs — the hygiene gate forbids
33
+ // tracked binary files; some browsers won't show an SVG app icon, which is
34
+ // an accepted tradeoff, see docs/reference/features-matrix.md) + a shell-
35
+ // only service worker. Same with { type: "text" } + /assets/<name>-shaped
36
+ // pattern as the vendored assets above; see src/web/sw.js's header for why
37
+ // it deliberately does NOT cache API/WS traffic.
38
+ import manifestWebmanifest from "./web/manifest.webmanifest" with { type: "text" };
39
+ import iconSvg from "./web/icon.svg" with { type: "text" };
40
+ import swJs from "./web/sw.js" with { type: "text" };
19
41
  import pkgJson from "../package.json" with { type: "json" };
20
42
  import { readFileSync } from "node:fs";
21
43
  const APP_PAGE = appHtml as unknown as string;
44
+ const HLJS_JS = hljsJs as unknown as string;
45
+ const HLJS_CSS = hljsCss as unknown as string;
46
+ const APP_JS = appJs as unknown as string;
47
+ const MANIFEST_WEBMANIFEST = manifestWebmanifest as unknown as string;
48
+ const ICON_SVG = iconSvg as unknown as string;
49
+ const SW_JS = swJs as unknown as string;
22
50
  const pkgVersion = (pkgJson as { version: string }).version;
23
- import { loadModelConfig, type KvQuantSpec, type ModelConfig } from "./config";
51
+ import { loadModelConfig, type KvQuantSpec, type ModelConfig, type TurboQuantScheme } from "./config";
24
52
  import { Weights } from "./weights";
25
53
  import { Gemma4Model } from "./model/gemma4";
26
54
  import { DiffusionGemmaModel } from "./model/diffusion-gemma";
27
55
  import { spliceImageTokens } from "./vision/diffusion-vision";
28
56
  import { createModel, type RuntimeModel } from "./model/factory";
29
- import { isMiniCPM5Config, isSupportedModelRecord } from "./model/support";
57
+ import { isMiniCPM5Config } from "./model/support";
30
58
  import { generate, type GenerateOptions, type TokenLogprobs } from "./generate";
59
+ import { cloneKvCaches, SpillQueue } from "./kv-store";
60
+ import { TURBOQUANT_HEAD_DIMS } from "./mlx/turboquant-tables";
61
+ import {
62
+ compileGrammarRequest, grammarEnabled, type GrammarRequest,
63
+ } from "./grammar";
31
64
  import type { HlgConfig } from "./sampler";
32
65
  import { isMonotone, CURVE_UMIN, type CurveParams } from "./curve-sampler";
33
66
  const CURVE_PAGE = curveDesignerHtml as unknown as string;
34
67
  import { GenerationGateway } from "./serve/generation-gateway";
68
+ import { recordLane, type Lane } from "./serve/lane-registry";
35
69
  import {
36
70
  ChatTemplate, type ChatMessage, type ToolDefinition,
37
71
  } from "./chat-template";
@@ -40,7 +74,7 @@ import {
40
74
  CHANNEL_END, CHANNEL_START, parseGeneratedToolCalls, parseToolCalls,
41
75
  TOOL_CALL_END, TOOL_CALL_START,
42
76
  } from "./tool-call";
43
- import { PromptCache } from "./prompt-cache";
77
+ import { PromptCache, cacheBytes } from "./prompt-cache";
44
78
  import { SsdCacheStore } from "./ssd-cache";
45
79
  import { configFingerprint } from "./model/fingerprint";
46
80
  import {
@@ -59,14 +93,23 @@ import { setMemoryLimit } from "./mlx/ffi";
59
93
  import { VisionTower } from "./vision/embedder";
60
94
  import { SiglipVisionTower, parseSiglipConfig } from "./vision/siglip";
61
95
  import {
62
- buildVisionPrompt, extractImages,
63
- type VisionTokenIds, type VisionEncoder,
96
+ buildMultimodalPrompt, buildVisionPrompt, extractAudio, extractImages,
97
+ type AudioTokenIds, type VisionTokenIds, type VisionEncoder,
64
98
  } from "./vision/prompt";
99
+ import { AudioTower, parseAudioConfig } from "./audio/conformer";
100
+ import { ensureWav } from "./audio/transcode";
101
+ import { sidecarShipsAudioTower } from "./registry";
65
102
  import { makePiWsHandler, type PiWsData } from "./pi-web";
66
103
 
67
104
  export interface ServerOptions {
68
- /** Byte cap for the prompt (KV) cache. Default 2 GB. */
105
+ /** Byte cap for the prompt (KV) cache. Default 8 GB. */
69
106
  promptCacheBytes?: number;
107
+ /** Aggregate KV-byte budget across concurrently-admitted batch rows
108
+ * (`--kv-budget`, batching-perf-path P3). Joiners whose projected KV
109
+ * (prompt + max_tokens, window-capped) would exceed it QUEUE until rows
110
+ * evict; a request over the budget alone is rejected. Unset = no
111
+ * aggregate cap (per-request admission via memoryBudget still applies). */
112
+ kvBudgetBytes?: number;
70
113
  /** KV quantization override. When unset: apply ctx.kvConfig (mixed
71
114
  * per-layer) for serial serving, but bf16 under `--batch N` (the batched
72
115
  * engine is bf16-only — a mode switch, see `batch` below). "config" forces
@@ -74,6 +117,19 @@ export interface ServerOptions {
74
117
  * the serial path); "off" forces bf16; a number forces uniform bits
75
118
  * (group size 64, start 0) ignoring the config file. */
76
119
  kvQuant?: "off" | "config" | number;
120
+ /** TurboQuant scheme (docs/design/turboquant-kv.md): a separate axis from
121
+ * kvQuant above, mutually exclusive with it (`--kv-quant turbo[:k<bits>v
122
+ * <bits>]` sets this instead of kvQuant). Solo-only in v1 — see
123
+ * GenerationGateway's explicit refusal. */
124
+ turboQuant?: TurboQuantScheme;
125
+ /** OPTIONAL paged KV cache (`--paged-kv`, docs/design/paged-kv-cache.md):
126
+ * vLLM-style block-pool storage for full-attention layers,
127
+ * gather-to-contiguous before the stock SDPA. Default off (unset = the
128
+ * plain KVCache path, byte-identical). v1 scope: serial batch=1,
129
+ * Gemma4-family, bf16 — startup refuses `--batch N>1`, `--kv-quant`,
130
+ * turbo, and `--draft-model` combinations; paged requests bypass the
131
+ * prompt cache and run uncompiled. Gated bit-exact vs the plain path. */
132
+ pagedKv?: { blockSize?: number };
77
133
  /** Memory budget for the serving process (admission control — Phase 5).
78
134
  * Requests whose prompt + max_tokens exceed the budget's max safe
79
135
  * context are rejected with 400 instead of crashing the GPU: the OOM
@@ -88,6 +144,11 @@ export interface ServerOptions {
88
144
  * it). Exposed at /stats.server so other mlx-bun processes can
89
145
  * warn before attaching to a server that may vanish. */
90
146
  owner?: "serve" | "pi-session" | "embedded";
147
+ /** Listen on a UNIX DOMAIN SOCKET instead of a TCP port — the engine-
148
+ * child mode of the isolation architecture (runtime-isolation.md): the
149
+ * parent process reverse-proxies HTTP to this socket. Stale socket
150
+ * files are unlinked before bind. When set, `hostname`/port are ignored. */
151
+ unixSocket?: string;
91
152
  /** Interface to bind (Bun.serve hostname). Unset ⇒ Bun's default
92
153
  * (all interfaces) — embedded/library use. The CLI always passes one:
93
154
  * "127.0.0.1" (loopback, mlx_lm.server parity) unless --host says
@@ -111,9 +172,9 @@ export interface ServerOptions {
111
172
  * thinking traces never truncate. `--max-tokens 512` = mlx-lm behavior. */
112
173
  defaultMaxTokens?: number;
113
174
  /** Max concurrent requests batched through the mlx-lm-parity engine
114
- * (`--batch N`). Default 1 = today's serialized single-queue path. >1 opts
115
- * the WHOLE server into the batched engine (continuous batching, B floats
116
- * 1..N, bit-parity with mlx-lm B=N) — a mode switch, not a load-dependent
175
+ * (`--batch N`). Default 8 (continuous batching, B floats 1..N,
176
+ * bit-parity with mlx-lm B=N); `--batch 1` pins the serialized
177
+ * single-queue path — a mode switch, not a load-dependent
117
178
  * fallback. See docs/design/parallel-slots.md. NOTE: the batched executor
118
179
  * is mid-build; until it lands, >1 warns and runs serially. */
119
180
  batch?: number;
@@ -132,6 +193,11 @@ export interface ServerOptions {
132
193
  ssdCacheDir?: string;
133
194
  /** Byte cap for the SSD tier (default 32 GiB). */
134
195
  ssdCacheMaxBytes?: number;
196
+ /** Idle-demotion threshold in seconds (`--ssd-demote-idle`): prompt-cache
197
+ * entries unused this long spill to the SSD tier and free their GPU
198
+ * memory (Layer 0 — RAM drains between bursts, prefixes stay reachable
199
+ * via zero-copy restore). Default 300 when the tier is on; 0 disables. */
200
+ ssdDemoteIdleSec?: number;
135
201
  /** Verify every tensor hash on restore (`--ssd-cache-verify`) — reads all
136
202
  * bytes eagerly, defeating lazy fault-in; integrity paranoia only. */
137
203
  ssdCacheVerify?: boolean;
@@ -151,6 +217,17 @@ export interface ServerContext {
151
217
  * `vision`. */
152
218
  loadVision: (() => VisionEncoder) | null;
153
219
  visionTokenIds: VisionTokenIds;
220
+ /** Lazily-loaded Conformer audio tower — null until the first audio
221
+ * request (see `getAudioTower`). Same sidecar file as vision
222
+ * (optiq_vision.safetensors), separate tower; text-only sessions never
223
+ * pay for it. */
224
+ audio: AudioTower | null;
225
+ /** Loads the audio tower on demand; null when the model can't do audio
226
+ * (no `audio_config` in config.json or no sidecar). No flags — audio
227
+ * auto-enables exactly like vision. */
228
+ loadAudio: (() => AudioTower) | null;
229
+ /** null when the model has no `audio_config` (audio-incapable). */
230
+ audioTokenIds: AudioTokenIds | null;
154
231
  adapters: AdapterManager;
155
232
  /** Per-layer KV quantization from the repo's kv_config.json (null if
156
233
  * absent). Applied by default — optiq serve's headline behavior;
@@ -160,6 +237,14 @@ export interface ServerContext {
160
237
  * optiq serve injects these as server defaults (gen_config.py);
161
238
  * explicit request fields always win. */
162
239
  genDefaults: GenSamplingDefaults;
240
+ /** Speculative decoding (`serve --draft-model`, mlx_lm.server parity).
241
+ * Server-level: when set, EVERY request routes to the serial lane
242
+ * (upstream: is_batchable = draft is None) and spec-eligible ones decode
243
+ * through src/spec/serve-loop.ts. null = no draft configured. */
244
+ draft?: {
245
+ provider: import("./spec/source").DraftProvider;
246
+ numDraftTokens: number;
247
+ } | null;
163
248
  }
164
249
 
165
250
  export interface GenSamplingDefaults {
@@ -169,6 +254,30 @@ export interface GenSamplingDefaults {
169
254
  repetitionPenalty?: number;
170
255
  }
171
256
 
257
+ export type DraftKind = "dspark" | "deepspec" | "assistant" | "two-model" | "ngram";
258
+
259
+ /** Detect the draft artifact's kind so the right provider is loaded. All
260
+ * providers share ONE serve loop (src/spec/serve-loop.ts). Exported for the
261
+ * bench harness (scripts/bench-feature-matrix.ts) — one detection, no drift.
262
+ * "ngram" is never detected — it has no artifact (model-free prompt lookup,
263
+ * src/spec/ngram-source.ts) and mounts via an explicit `--draft-kind ngram`. */
264
+ export async function detectDraftKind(dir: string): Promise<DraftKind> {
265
+ if (await Bun.file(`${dir}/dspark.json`).exists()) return "dspark"; // our trained module
266
+ try {
267
+ const cfg = (await Bun.file(`${dir}/config.json`).json()) as {
268
+ model_type?: string;
269
+ architectures?: string[];
270
+ };
271
+ // DeepSeek's released DSpark drafters (DeepSpec reference): no
272
+ // dspark.json, plain HF config stamped Gemma4DSparkModel.
273
+ if (cfg.architectures?.[0] === "Gemma4DSparkModel") return "deepspec";
274
+ if (String(cfg.model_type ?? "").includes("assistant")) return "assistant";
275
+ } catch {
276
+ // no/unreadable config → fall through to a full second model
277
+ }
278
+ return "two-model";
279
+ }
280
+
172
281
  async function loadGenSamplingDefaults(modelDir: string): Promise<GenSamplingDefaults> {
173
282
  const file = Bun.file(`${modelDir}/generation_config.json`);
174
283
  if (!(await file.exists())) return {};
@@ -188,7 +297,24 @@ async function loadGenSamplingDefaults(modelDir: string): Promise<GenSamplingDef
188
297
 
189
298
  export async function loadContext(
190
299
  modelDir: string, modelId?: string,
191
- opts: { memoryBudgetBytes?: number } = {},
300
+ opts: {
301
+ memoryBudgetBytes?: number;
302
+ /** Snapshot dir of a draft model for speculative decoding
303
+ * (`--draft-model`). Loaded alongside the target; the pair must share
304
+ * a tokenizer family (hard check below — upstream's silent-garbage
305
+ * mode on mismatched tokenizers isn't worth inheriting). */
306
+ draftModelDir?: string;
307
+ /** Drafts per round (`--num-draft-tokens`, mlx_lm.server default 3). */
308
+ numDraftTokens?: number;
309
+ /** Draft-provider kind override (`--draft-kind`); auto-detected from the
310
+ * draft artifact when absent (dspark.json → dspark; *_assistant config →
311
+ * assistant; otherwise two-model). "ngram" mounts WITHOUT a draft dir. */
312
+ draftKind?: DraftKind;
313
+ /** Prompt-lookup window bounds (`--ngram-max` / `--ngram-min`), ngram
314
+ * kind only. Defaults 3 / 1 (Saxena's reference values). */
315
+ ngramMax?: number;
316
+ ngramMin?: number;
317
+ } = {},
192
318
  ): Promise<ServerContext> {
193
319
  const config = await loadModelConfig(modelDir);
194
320
  const weights = await Weights.open(modelDir);
@@ -219,7 +345,116 @@ export async function loadContext(
219
345
  // the tokenizer eos; union it in. No-op when already present (Gemma, 27B).
220
346
  if (tokenizer.eosTokenId != null && !config.eosTokenIds.includes(tokenizer.eosTokenId))
221
347
  config.eosTokenIds = [...config.eosTokenIds, tokenizer.eosTokenId];
348
+
349
+ // Speculative decoding: load the draft (mlx_lm.server --draft-model). The
350
+ // draft artifact's KIND selects the provider — all three share ONE serve
351
+ // loop (src/spec/serve-loop.ts): dspark.json → DSpark (KV-injected), a
352
+ // *_assistant config → the optiq KV-borrowing Gemma drafter, otherwise a
353
+ // full second model (mlx-lm parity). `--draft-kind` overrides the detect.
354
+ let draft: ServerContext["draft"] = null;
355
+ if (opts.draftKind === "ngram") {
356
+ // Model-free prompt lookup: no artifact, no dir, no probe/budget concerns
357
+ // (weightsBytes 0, open() never throws). Default γ=10 per the reference
358
+ // implementation — drafting is free, so wide blocks cost only verify-window
359
+ // width when wrong.
360
+ if (opts.draftModelDir)
361
+ throw new Error(
362
+ "--draft-kind ngram is model-free — drop --draft-model (it would be ignored)",
363
+ );
364
+ const { NgramProvider } = await import("./spec/ngram-source");
365
+ draft = {
366
+ provider: new NgramProvider({ max: opts.ngramMax, min: opts.ngramMin }),
367
+ numDraftTokens: Math.max(1, opts.numDraftTokens ?? 10),
368
+ };
369
+ } else if (opts.draftModelDir) {
370
+ const dir = opts.draftModelDir;
371
+ const kind = opts.draftKind ?? (await detectDraftKind(dir));
372
+ let provider: import("./spec/source").DraftProvider;
373
+ let numDraftTokens = Math.max(1, opts.numDraftTokens ?? 3);
374
+ if (kind === "dspark") {
375
+ const { DflashProvider } = await import("./spec/dflash-source");
376
+ const p = await DflashProvider.load(dir);
377
+ provider = p;
378
+ // Pin to the trained block width — the serve loop must never ask for
379
+ // more positions than the DSpark block was trained for (n ≤ cfg.gamma).
380
+ numDraftTokens = Math.max(1, Math.min(opts.numDraftTokens ?? p.gamma, p.gamma));
381
+ } else if (kind === "deepspec") {
382
+ const { DeepspecProvider } = await import("./spec/deepspec-source");
383
+ const p = await DeepspecProvider.load(dir);
384
+ provider = p;
385
+ // Same pin, from their config's block_size (e.g. 7 for the released
386
+ // dspark_gemma4_12b_block7).
387
+ numDraftTokens = Math.max(1, Math.min(opts.numDraftTokens ?? p.gamma, p.gamma));
388
+ } else if (kind === "assistant") {
389
+ const { AssistantProvider } = await import("./spec/assistant-source");
390
+ provider = await AssistantProvider.load(dir);
391
+ } else {
392
+ const { TwoModelProvider } = await import("./spec/two-model");
393
+ provider = await TwoModelProvider.load(dir, config.text.vocabSize);
394
+ // Tokenizer-family hard check — two-model ONLY (it ships its own
395
+ // tokenizer). Exact-token-match acceptance is meaningful only when both
396
+ // models tokenize identically; a probe that ENCODES differently means
397
+ // different families — refuse instead of silently accepting ~0% of
398
+ // drafts. The KV-borrowing drafters share the target's tokenization by
399
+ // construction, so this check does not apply to them.
400
+ const draftTok = await loadTokenizer(dir);
401
+ const probe = "The 3 quick brown foxes jumped över the lazy dog?! 🦊";
402
+ if (JSON.stringify(tokenizer.encode(probe)) !== JSON.stringify(draftTok.encode(probe))) {
403
+ provider.dispose();
404
+ throw new Error(
405
+ `--draft-model tokenizer differs from the target's (probe string encodes ` +
406
+ `differently) — speculation needs the same tokenizer family`,
407
+ );
408
+ }
409
+ }
410
+ // Fail-fast pairing validation (2026-07-07 review): the KV-borrowing
411
+ // sources validate the (target, drafter) pairing in open() — non-Gemma4
412
+ // target, DeepSpec target-layer-count mismatch — which used to surface
413
+ // as a 500 from inside specServeRun on EVERY text request. Probe-open
414
+ // once with throwaway caches here so a mismatch refuses at load; open()
415
+ // allocates no per-request tensors before prefill/draft, so this is
416
+ // free. The probe sampler is never called during open().
417
+ {
418
+ const probeCaches = model.makeCache();
419
+ try {
420
+ provider
421
+ .open({
422
+ sampler: () => { throw new Error("probe sampler never samples"); },
423
+ target: { model, caches: probeCaches },
424
+ })
425
+ .dispose();
426
+ } catch (err) {
427
+ provider.dispose();
428
+ throw new Error(
429
+ `--draft-model is incompatible with this target: ${err instanceof Error ? err.message : String(err)}`,
430
+ );
431
+ } finally {
432
+ for (const c of probeCaches) c.dispose();
433
+ }
434
+ }
435
+ if (opts.memoryBudgetBytes) {
436
+ const targetBytes = [...weights.shards.files.values()].reduce((a, f) => a + f.mmap.size, 0);
437
+ // Draft weights shrink the target's envelope. Draft KV is not modeled
438
+ // (small relative to its weights at serve contexts); admission stays
439
+ // approximately conservative via the combined-weights term.
440
+ const report = fit(config, targetBytes + provider.weightsBytes, 1, undefined, undefined, 0, opts.memoryBudgetBytes);
441
+ if (report.maxSafeContext < 1) {
442
+ provider.dispose();
443
+ throw new Error(
444
+ `target + draft do not fit the memory budget (draft adds ` +
445
+ `${(provider.weightsBytes / 1e9).toFixed(2)} GB)`,
446
+ );
447
+ }
448
+ }
449
+ draft = { provider, numDraftTokens };
450
+ } else if (opts.draftKind) {
451
+ // Every other kind names an artifact to load — refuse instead of silently
452
+ // serving without speculation.
453
+ throw new Error(`--draft-kind ${opts.draftKind} requires --draft-model`);
454
+ }
455
+
222
456
  return {
457
+ draft,
223
458
  model,
224
459
  adapters: new AdapterManager(model),
225
460
  kvConfig: config.kvQuant,
@@ -239,6 +474,17 @@ export async function loadContext(
239
474
  boiTokenId: (config.raw.boi_token_id as number) ?? 255999,
240
475
  eoiTokenId: (config.raw.eoi_token_id as number) ?? 258882,
241
476
  },
477
+ // Audio mirrors vision: lazy tower from the same sidecar, loaded on the
478
+ // first audio request only (docs/design/audio-input-plan.md A4).
479
+ audio: null,
480
+ loadAudio: makeAudioLoader(modelDir, model, config),
481
+ audioTokenIds: config.raw.audio_config
482
+ ? {
483
+ audioTokenId: (config.raw.audio_token_id as number) ?? 258881,
484
+ boaTokenId: (config.raw.boa_token_id as number) ?? 256000,
485
+ eoaTokenId: (config.raw.eoa_token_id as number) ?? 258883,
486
+ }
487
+ : null,
242
488
  };
243
489
  }
244
490
 
@@ -276,6 +522,46 @@ function getVisionTower(ctx: ServerContext): VisionEncoder | null {
276
522
  }
277
523
  }
278
524
 
525
+ /** Build the on-demand audio-tower loader (gemma-4 Conformer, A4 of
526
+ * docs/design/audio-input-plan.md). Auto-enables — no flags — when
527
+ * config.json carries an `audio_config` AND the optiq_vision.safetensors
528
+ * sidecar exists (the audio tensors ship in the same sidecar as vision).
529
+ * Returns null when the model can't do audio (no audio_config: 26B-A4B,
530
+ * DiffusionGemma, bf16 assistants — architectural, not a porting gap). */
531
+ function makeAudioLoader(
532
+ modelDir: string, model: RuntimeModel, config: ModelConfig,
533
+ ): (() => AudioTower) | null {
534
+ if (!(config.hasVisionSidecar && config.raw.audio_config && model instanceof Gemma4Model))
535
+ return null;
536
+ // The sidecar header must actually name the Conformer tensors: the local
537
+ // 12B snapshot pairs audio_config with a STUB sidecar (embed_audio only),
538
+ // and a non-null loader here is what advertises `audio: true` on every
539
+ // capability surface (ws handshake, /v1/models). Header-only read.
540
+ if (!sidecarShipsAudioTower(`${modelDir}/optiq_vision.safetensors`)) return null;
541
+ const audioCfg = parseAudioConfig(config.raw.audio_config as Record<string, any>);
542
+ return () => AudioTower.load(modelDir, audioCfg, model.embedScale);
543
+ }
544
+
545
+ /** Lazily load + cache the audio tower on first use. Unlike vision's
546
+ * warn-and-continue (text-only degrade is correct for requests WITHOUT
547
+ * images), this is only ever consulted for requests WITH audio — the
548
+ * caller turns null into an explicit 400, never a silent text-only
549
+ * degrade. A failed load is not retried every request. (The stub-sidecar
550
+ * case — the local 12B state — never gets here: makeAudioLoader checks the
551
+ * sidecar header and returns a null loader.) */
552
+ function getAudioTower(ctx: ServerContext): AudioTower | null {
553
+ if (ctx.audio) return ctx.audio;
554
+ if (!ctx.loadAudio) return null;
555
+ try {
556
+ ctx.audio = ctx.loadAudio();
557
+ return ctx.audio;
558
+ } catch (e) {
559
+ console.warn(`audio sidecar not loadable (${(e as Error).message})`);
560
+ ctx.loadAudio = null;
561
+ return null;
562
+ }
563
+ }
564
+
279
565
  interface ChatRequest {
280
566
  messages: ChatMessage[];
281
567
  stream?: boolean;
@@ -340,6 +626,20 @@ interface ChatRequest {
340
626
  toe?: number;
341
627
  pivot_offset?: number;
342
628
  };
629
+ /** OpenAI structured output: {type:"json_object"} | {type:"json_schema",
630
+ * json_schema:{name,schema,strict?}} | {type:"text"}. Enforced at the
631
+ * sampler via xgrammar token-bitmasks (src/grammar.ts). L2-class (oMLX
632
+ * oracle). On compile failure, degrades to a system-prompt injection +
633
+ * Warning header (oMLX parity), never 500. */
634
+ response_format?: unknown;
635
+ /** vLLM/oMLX grammar aliases (all compiled via xgrammar): raw EBNF/LARK
636
+ * grammar, regex, enum choice, and a bare JSON-schema object. Precedence
637
+ * (guided_grammar > json_schema > json_object > structured_outputs >
638
+ * guided_regex > guided_choice) mirrors oMLX _effective_guided_grammar. */
639
+ guided_grammar?: string;
640
+ guided_regex?: string;
641
+ guided_choice?: string[];
642
+ structured_outputs?: unknown;
343
643
  }
344
644
 
345
645
  /** POST /v1/completions body (mlx_lm.server's raw text completion — no chat
@@ -419,9 +719,9 @@ export class ToolAwareStream {
419
719
  * chain-of-thought as `<|channel>thought\n…<channel|>` using special tokens
420
720
  * 100/101 that the content decoder strips, so reasoning is captured here at
421
721
  * the token level. A SEPARATE decoder keeps the reasoning byte-stream's
422
- * incremental state independent of content's. The `thought\n` channel-name
423
- * line is stripped before the reasoning text (per the model's parser spec,
424
- * it sits outside the captured thought). */
722
+ * incremental state independent of content's. The `thought` channel-name
723
+ * word is stripped before the reasoning text (mlx-lm's think-start marker
724
+ * is `<|channel>thought`; the "\n" after it is reasoning content). */
425
725
  readonly #channelDecoder: StreamDecoder;
426
726
  #inChannel = false;
427
727
  #channelNamePending = "";
@@ -442,15 +742,20 @@ export class ToolAwareStream {
442
742
  this.#channelDecoder = new StreamDecoder(tokenizer, true);
443
743
  }
444
744
 
445
- /** Feed decoded channel text, stripping the leading `thought\n` name line
446
- * (which may arrive across tokens). Returns the reasoning delta. */
745
+ /** Feed decoded channel text, stripping the leading `thought` channel-name
746
+ * word (which may arrive across tokens). Returns the reasoning delta.
747
+ * mlx-lm parity: its think-start marker is exactly `<|channel>thought`
748
+ * (tokenizer_utils.py `_infer_thinking`) and only the MARKER tokens get
749
+ * their text blanked (server.py `_process_control_tokens`) — the "\n"
750
+ * after the name is an ordinary generated token, so it is the FIRST byte
751
+ * of the reasoning stream and must be kept, not swallowed with the name. */
447
752
  #feedChannel(text: string): string {
448
753
  if (this.#channelNameDone) return text;
449
754
  this.#channelNamePending += text;
450
755
  const nl = this.#channelNamePending.indexOf("\n");
451
- if (nl === -1) return ""; // still inside the channel-name line
756
+ if (nl === -1) return ""; // still inside the channel-name word
452
757
  this.#channelNameDone = true;
453
- const rest = this.#channelNamePending.slice(nl + 1);
758
+ const rest = this.#channelNamePending.slice(nl); // keep the "\n" (mlx-lm does)
454
759
  this.#channelNamePending = "";
455
760
  return rest;
456
761
  }
@@ -515,7 +820,8 @@ export class ToolAwareStream {
515
820
  // Reasoning channel: tokens between <|channel> and <channel|> are thought,
516
821
  // captured to #reasoning (drained via takeReasoning), never content. An
517
822
  // empty block (<|channel>thought\n<channel|>, emitted by larger Gemmas even
518
- // with thinking off) yields no reasoning and leaks nothing.
823
+ // with thinking off) yields only the "\n" as reasoning (mlx-lm parity) and
824
+ // leaks nothing into content.
519
825
  if (this.#inChannel) {
520
826
  if (token === CHANNEL_END) {
521
827
  this.#inChannel = false;
@@ -756,10 +1062,15 @@ function contentPartsToText(parts: Array<Record<string, unknown>>): string {
756
1062
  .join("");
757
1063
  }
758
1064
 
759
- /** True if any content part is an image (so the vision path must keep the
760
- * array form for extractImages). */
761
- function hasImagePart(parts: Array<Record<string, unknown>>): boolean {
762
- return parts.some((p) => p && (p.type === "image" || p.type === "image_url"));
1065
+ /** True if any content part is media — image OR audio — so the multimodal
1066
+ * path must keep the array form for extractImages/extractAudio. */
1067
+ function hasMediaPart(parts: Array<Record<string, unknown>>): boolean {
1068
+ return parts.some(
1069
+ (p) =>
1070
+ p &&
1071
+ (p.type === "image" || p.type === "image_url" ||
1072
+ p.type === "audio" || p.type === "input_audio" || p.type === "audio_url"),
1073
+ );
763
1074
  }
764
1075
 
765
1076
  /** OpenAI sends assistant tool_call arguments as JSON strings; the
@@ -780,13 +1091,14 @@ function hasImagePart(parts: Array<Record<string, unknown>>): boolean {
780
1091
  * OpenAI multimodal client) sends user content as `[{type:"text",text}]`,
781
1092
  * but non-vision chat templates expect `content` to be a STRING and render
782
1093
  * nothing for an array — so the user's turn silently vanishes and the model
783
- * replies "I don't see any message." Arrays that carry an image part are
784
- * left intact for the vision path (extractImages); only text-only arrays
785
- * are collapsed here, which is a no-op for the vision templates too. */
1094
+ * replies "I don't see any message." Arrays that carry a media part
1095
+ * (image or audio) are left intact for the multimodal path
1096
+ * (extractImages/extractAudio); only text-only arrays are collapsed
1097
+ * here, which is a no-op for the vision templates too. */
786
1098
  export function normalizeMessages(messages: ChatMessage[]): ChatMessage[] {
787
1099
  return messages.map((raw) => {
788
1100
  let m = raw.role === "developer" ? { ...raw, role: "system" } : raw;
789
- if (Array.isArray(m.content) && !hasImagePart(m.content)) {
1101
+ if (Array.isArray(m.content) && !hasMediaPart(m.content)) {
790
1102
  m = { ...m, content: contentPartsToText(m.content) };
791
1103
  }
792
1104
  if (!m.tool_calls) return m;
@@ -850,6 +1162,33 @@ export function nextDefaultSeed(): number {
850
1162
  * with defaults logprobs=False / top_logprobs=-1). Returns the reference's
851
1163
  * error message (→ 400) or null when valid. Note mlx-lm caps top_logprobs
852
1164
  * at 11, not OpenAI's 20 — we mirror the reference. */
1165
+ /** Build a degrade-path system-prompt instruction for JSON output, mirroring
1166
+ * oMLX's api.tool_calling.build_json_system_prompt (used when xgrammar
1167
+ * compile fails — the response_format degrades to prompt injection rather
1168
+ * than a 500, oMLX parity). Returns null for {type:"text"} / unset. */
1169
+ export function degradeJsonSystemPrompt(body: ChatRequest): string | null {
1170
+ const rf = body.response_format as { type?: string; json_schema?: { name?: string; description?: string; schema?: unknown } } | undefined;
1171
+ if (!rf || typeof rf !== "object") return null;
1172
+ const type = rf.type ?? "text";
1173
+ if (type === "text") return null;
1174
+ if (type === "json_object") {
1175
+ return "You must respond with valid JSON only. " +
1176
+ "Do not include any explanation or text outside the JSON object.";
1177
+ }
1178
+ if (type === "json_schema") {
1179
+ const spec = rf.json_schema ?? {};
1180
+ const schema = spec.schema ?? body.structured_outputs ?? {};
1181
+ const name = spec.name ?? "response";
1182
+ const description = spec.description ?? "";
1183
+ let prompt = `You must respond with valid JSON matching the '${name}' schema.`;
1184
+ if (description) prompt += ` ${description}`;
1185
+ prompt += `\n\nJSON Schema:\n\`\`\`json\n${JSON.stringify(schema, null, 2)}\n\`\`\`\n\n` +
1186
+ "Respond with only the JSON object, no additional text or explanation.";
1187
+ return prompt;
1188
+ }
1189
+ return null;
1190
+ }
1191
+
853
1192
  export function validateLogprobsParams(body: {
854
1193
  logprobs?: unknown;
855
1194
  top_logprobs?: unknown;
@@ -929,24 +1268,79 @@ class LogprobsCollector {
929
1268
  }
930
1269
  }
931
1270
 
932
- /** Incremental detokenizer: emits the longest stable decoded prefix. */
933
- class StreamDecoder {
1271
+ /** Incremental detokenizer: emits the longest stable decoded prefix.
1272
+ *
1273
+ * Byte parity with mlx-lm's streaming detokenizers (the drop-in contract is
1274
+ * rendered BYTES, not just token ids). For BPE/ByteLevel tokenizers, two
1275
+ * mlx-lm 0.31.3 BPEStreamingDetokenizer behaviors our full-sequence decode
1276
+ * lacks (tokenizer_utils.py:195-226):
1277
+ *
1278
+ * 1. `trimsLeadingSpace` — mlx-lm drops ONE leading " " at the start of the
1279
+ * generated sequence (`_maybe_trim_space`); trim it here. SPM decode
1280
+ * already matches (see LoadedTokenizer.trimsLeadingSpace).
1281
+ * 2. `bareSpaceTokenId` — add_token WITHHOLDS a single-char byte-32 token
1282
+ * ("Ġ") in `_unflushed` ("For single spaces wait until the next token"),
1283
+ * flushing it together with the NEXT token — and mlx_lm.server NEVER
1284
+ * calls detokenizer.finalize() (zero hits in server.py 0.31.3), so a
1285
+ * generation ENDING on bare-space token(s) silently drops their spaces
1286
+ * from the served bytes. push() withholds those spaces; flush() drops a
1287
+ * trailing bare-space run (the held text dies with the request, exactly
1288
+ * like mlx-lm serve). Re-check both if upstream ever adds a finalize()
1289
+ * call. LATENT HAZARD (deliberately not emulated): models with
1290
+ * clean_up_tokenization_spaces=true get an ADDITIONAL mid-stream rule
1291
+ * (`_space_matches`: held space dropped before "." "," "'s" …) — both
1292
+ * current BPE targets have it false (MiniCPM5), so it never fires here.
1293
+ *
1294
+ * Exported for unit tests (serve-detok mlx-lm byte parity). */
1295
+ export class StreamDecoder {
934
1296
  #ids: number[] = [];
935
1297
  #emitted = "";
1298
+ #warnedRevision = false;
1299
+ readonly #trimLeadingSpace: boolean;
1300
+ readonly #bareSpaceId: number | undefined;
936
1301
 
937
1302
  constructor(
938
1303
  readonly tokenizer: LoadedTokenizer,
939
1304
  readonly skipSpecialTokens = true,
940
- ) {}
1305
+ ) {
1306
+ this.#trimLeadingSpace = tokenizer.trimsLeadingSpace === true;
1307
+ this.#bareSpaceId = tokenizer.bareSpaceTokenId;
1308
+ }
1309
+
1310
+ #decode(ids: number[]): string {
1311
+ const full = this.tokenizer.decode(ids, this.skipSpecialTokens);
1312
+ return this.#trimLeadingSpace && full.startsWith(" ") ? full.slice(1) : full;
1313
+ }
941
1314
 
942
1315
  push(token: number): string {
943
1316
  this.#ids.push(token);
944
- const full = this.tokenizer.decode(this.#ids, this.skipSpecialTokens);
1317
+ // Bare-space hold-back (mlx-lm add_token keeps "Ġ" in _unflushed): don't
1318
+ // advance #emitted; the held space(s) flush as part of the next
1319
+ // non-bare-space token's delta — consecutive bare spaces accumulate.
1320
+ if (token === this.#bareSpaceId) return "";
1321
+ const full = this.#decode(this.#ids);
945
1322
  // hold back a trailing replacement char (partial multi-byte sequence)
946
1323
  const stable = full.endsWith("�") ? full.slice(0, -1) : full;
947
1324
  if (!stable.startsWith(this.#emitted)) {
948
- // decoder revised earlier text (rare); re-emit from scratch
949
- const out = stable;
1325
+ // The decoder revised already-streamed text (cleanup rules — never
1326
+ // fires for the shipped tokenizers, which all have
1327
+ // clean_up_tokenization_spaces=false; see the LATENT HAZARD note in
1328
+ // the class doc). An SSE client cannot un-receive bytes, so the old
1329
+ // "re-emit from scratch" answer DUPLICATED the whole stream
1330
+ // (2026-07-07 review). Truncate-safe resync instead: emit only the
1331
+ // length-extension and keep the emitted watermark monotone — byte
1332
+ // drift is confined to the revised span, and the once-per-stream
1333
+ // warning makes the tokenizer that needs real _space_matches
1334
+ // emulation loud instead of silently wrong.
1335
+ if (!this.#warnedRevision) {
1336
+ this.#warnedRevision = true;
1337
+ console.warn(
1338
+ "[detok] decoder revised already-streamed text (clean_up_tokenization_spaces?) — " +
1339
+ "resyncing without re-emit; add mlx-lm's cleanup-rule emulation for this tokenizer",
1340
+ );
1341
+ }
1342
+ if (stable.length <= this.#emitted.length) return "";
1343
+ const out = stable.slice(this.#emitted.length);
950
1344
  this.#emitted = stable;
951
1345
  return out;
952
1346
  }
@@ -956,7 +1350,15 @@ class StreamDecoder {
956
1350
  }
957
1351
 
958
1352
  flush(): string {
959
- const full = this.tokenizer.decode(this.#ids, this.skipSpecialTokens);
1353
+ // mlx_lm.server never finalize()s: a trailing bare-space run stays
1354
+ // withheld forever, so its text is dropped from the served bytes.
1355
+ let ids = this.#ids;
1356
+ if (this.#bareSpaceId !== undefined) {
1357
+ let n = ids.length;
1358
+ while (n > 0 && ids[n - 1] === this.#bareSpaceId) n--;
1359
+ if (n < ids.length) ids = ids.slice(0, n);
1360
+ }
1361
+ const full = this.#decode(ids);
960
1362
  const delta = full.slice(this.#emitted.length);
961
1363
  this.#emitted = full;
962
1364
  return delta;
@@ -986,41 +1388,139 @@ export function createServer(
986
1388
  // cache by attention type. Non-batchable requests (vision / adapters /
987
1389
  // repetition penalty / user seed / explicit kv-quant) drain to the serial
988
1390
  // lane (see GenerationGateway.willBatch).
989
- const batch = Math.max(1, Math.floor(serverOptions.batch ?? 1));
990
-
991
- // KV-quant scheme, resolved once: kv_config.json by default (optiq
992
- // serve's headline behavior), overridable to uniform bits or off.
993
- // KV quant scheme. The serial default is the model's mixed-precision config
994
- // (optiq parity). But `--batch N` is a bf16 continuous-batching MODE (= mlx-lm
995
- // B=N parity), so when KV quant is left UNSET under `--batch N` it defaults to
996
- // bf16 — the batch path engages out of the box. An EXPLICIT --kv-quant
997
- // (config / bits) is still honored, but those requests then route to the
998
- // serial path (batched quantized KV is the L2 follow-up); warned below.
1391
+ // DEFAULT 8 (flipped 2026-07-05, Josh's call, after GATE-B1-SPEED): a
1392
+ // lone request through the batch lane IS the serial engine (adopted
1393
+ // serial-class caches, compiled decode, prompt cache + SSD restore;
1394
+ // 0.992-0.996 paired decode ratios, byte-identical output), so the cap
1395
+ // only changes behavior when concurrent requests actually arrive — the
1396
+ // agentic sub-agent workload. --batch 1 pins strict serial for
1397
+ // arrival-independent numerics. 8 = optiq's Mac-safe concurrency.
1398
+ const batch = Math.max(1, Math.floor(serverOptions.batch ?? 8));
1399
+
1400
+ // KV-quant scheme, resolved once. UNSET now means bf16 (flipped
1401
+ // 2026-07-05 with the naked-=-L1 default): quantized KV measured 5–20%
1402
+ // SLOWER decode than bf16 at ≤16k on every model — on mlx-lm too (its
1403
+ // kv8 oracle trails its own bf16) — so it pays only in memory headroom
1404
+ // and must be an explicit opt-in (--kv-quant config|4|8, or --l2
1405
+ // whose presets pass it explicitly). The CLI always passes kvQuant now;
1406
+ // this fallback is the library-user default and matches the CLI's.
999
1407
  const configScheme = ctx.kvConfig?.length ? { kvConfig: ctx.kvConfig } : {};
1000
- const kvScheme: Pick<GenerateOptions, "kvBits" | "kvConfig" | "quantizedKvStart"> =
1001
- serverOptions.kvQuant === "off" ? {}
1408
+ // Mutually exclusive by contract (GenerateOptions.turboQuant doc): a
1409
+ // programmatic caller setting both gets turboQuant — say so, like the
1410
+ // other risky-combination warnings below.
1411
+ if (serverOptions.turboQuant && serverOptions.kvQuant && serverOptions.kvQuant !== "off")
1412
+ console.warn(
1413
+ `[kv-quant] both turboQuant and --kv-quant ${serverOptions.kvQuant} are set; ` +
1414
+ `turboQuant wins (they are mutually exclusive).`,
1415
+ );
1416
+ const kvScheme: Pick<GenerateOptions, "kvBits" | "kvConfig" | "quantizedKvStart" | "turboQuant"> =
1417
+ serverOptions.turboQuant ? { turboQuant: serverOptions.turboQuant, quantizedKvStart: 0 }
1418
+ : serverOptions.kvQuant === "off" ? {}
1002
1419
  : serverOptions.kvQuant === "config" ? configScheme
1003
1420
  : typeof serverOptions.kvQuant === "number"
1004
1421
  ? { kvBits: serverOptions.kvQuant, quantizedKvStart: 0 }
1005
- : batch > 1 ? {} // unset + batching → bf16 mode (Option B)
1006
- : configScheme; // unset + serial → optiq-parity mixed-precision default
1007
- if (batch > 1 && (kvScheme.kvConfig?.length || kvScheme.kvBits))
1422
+ : {}; // unset → bf16 (L1 default; quantized KV is opt-in)
1423
+ // Phase 3.1: kvConfig whose layers are all full-attention BATCHES (the
1424
+ // scheduler applies the mixed scheme per row); uniform kvBits and configs
1425
+ // touching rotating layers still route those requests serial. The warning
1426
+ // only fires for the still-serial compositions.
1427
+ if (batch > 1 && kvScheme.kvBits)
1008
1428
  console.warn(
1009
- `[batch] --batch ${batch} with explicit --kv-quant: batched serving (v1) is ` +
1010
- `bf16-only, so kv-quant requests route to the serial path (no batching for them). ` +
1011
- `Omit --kv-quant to batch in bf16. (docs/design/parallel-slots.md)`,
1429
+ `[batch] --batch ${batch} with uniform --kv-quant ${kvScheme.kvBits}: uniform ` +
1430
+ `quantized KV is serial-only (it touches rotating layers) — those requests ` +
1431
+ `won't batch. Use --kv-quant config for batched mixed-precision KV, or omit ` +
1432
+ `--kv-quant to batch in bf16. (docs/design/unified-engine-frontier-plan.md)`,
1012
1433
  );
1434
+ // TurboQuant is solo-only in v1 (novel cache class, not batchable by
1435
+ // construction — see GenerationGateway#modelCachesBatchable/willBatch).
1436
+ if (batch > 1 && kvScheme.turboQuant)
1437
+ console.warn(
1438
+ `[batch] --batch ${batch} with --kv-quant turbo: TurboQuant is serial-only in v1 ` +
1439
+ `— those requests won't batch. Omit --kv-quant to batch in bf16. ` +
1440
+ `(docs/design/turboquant-kv.md)`,
1441
+ );
1442
+ // Quantized KV (any axis) excludes the spec lane: drafted requests fall to
1443
+ // the normal serial path WITH the KV scheme applied rather than losing it
1444
+ // silently (mirrors the affine kvBits/kvConfig exclusion; the spec loop is
1445
+ // bf16-KV-only in v1).
1446
+ if (ctx.draft && (kvScheme.turboQuant || kvScheme.kvBits || kvScheme.kvConfig?.length))
1447
+ console.warn(
1448
+ `[spec] --draft-model with quantized KV (--kv-quant ${serverOptions.turboQuant ? "turbo" : String(serverOptions.kvQuant)}): ` +
1449
+ `the speculative lane is bf16-KV-only in v1 — requests keep the KV scheme and ` +
1450
+ `decode serially WITHOUT speculation. Omit --kv-quant to speculate. ` +
1451
+ `(docs/design/dspark-serving-program.md Phase 4)`,
1452
+ );
1453
+ // TurboQuant head-dim fail-fast (2026-07-07 review): the cache class only
1454
+ // supports {64,128,256,512} (sign-vector + Lloyd-Max table coverage) and
1455
+ // used to validate LAZILY on the first append — an unsupported model
1456
+ // (e.g. a 72/80/96 head dim) accepted --kv-quant turbo at startup and
1457
+ // then 500'd EVERY request from inside prefill. The config knows the
1458
+ // full-attention head dim (the only kind that converts; sliding layers
1459
+ // stay bf16) — refuse at createServer instead.
1460
+ if (kvScheme.turboQuant) {
1461
+ const dim = ctx.model.config.text.globalHeadDim;
1462
+ if (!(TURBOQUANT_HEAD_DIMS as readonly number[]).includes(dim))
1463
+ throw new Error(
1464
+ `--kv-quant turbo: this model's full-attention head_dim is ${dim}; ` +
1465
+ `TurboQuant supports {${TURBOQUANT_HEAD_DIMS.join(",")}} ` +
1466
+ `(docs/design/turboquant-kv.md) — use --kv-quant config|4|8 or omit it`,
1467
+ );
1468
+ }
1469
+ // Paged KV v1 (docs/design/paged-kv-cache.md): explicit refusals, not
1470
+ // silent downgrades — the incompatible combos would otherwise degrade
1471
+ // quietly (batch: paged caches can't merge; kv-quant: the swap would
1472
+ // drop the scheme, the exact composition bug the kvQuant gate above
1473
+ // exists to prevent; draft: the spec lane assumes serial-class caches).
1474
+ if (serverOptions.pagedKv) {
1475
+ if (batch > 1)
1476
+ throw new Error(
1477
+ `--paged-kv is serial-only in v1 — add --batch 1 (got --batch ${batch}). ` +
1478
+ `Batched paging is the follow-up PR (docs/design/paged-kv-cache.md).`,
1479
+ );
1480
+ if (kvScheme.kvBits || kvScheme.kvConfig?.length || kvScheme.turboQuant)
1481
+ throw new Error(
1482
+ `--paged-kv is bf16-only in v1 — omit --kv-quant (quantized paged ` +
1483
+ `blocks are a documented follow-up, docs/design/paged-kv-cache.md).`,
1484
+ );
1485
+ if (ctx.draft)
1486
+ throw new Error(
1487
+ `--paged-kv cannot combine with --draft-model in v1 ` +
1488
+ `(docs/design/paged-kv-cache.md non-goals).`,
1489
+ );
1490
+ if (!ctx.model.config.modelType.startsWith("gemma4"))
1491
+ throw new Error(
1492
+ `--paged-kv v1 supports Gemma4-family models only ` +
1493
+ `(this model: ${ctx.model.config.modelType}) — docs/design/paged-kv-cache.md.`,
1494
+ );
1495
+ const bs = serverOptions.pagedKv.blockSize;
1496
+ if (bs !== undefined && (!Number.isInteger(bs) || bs <= 0))
1497
+ throw new Error(`--paged-kv-block-size must be a positive integer (got ${bs})`);
1498
+ if (serverOptions.ssdCacheDir)
1499
+ console.warn(
1500
+ "[paged-kv] --ssd-cache has no effect: paged requests bypass the prompt " +
1501
+ "cache (v1 non-goal), so nothing reaches the SSD tier.",
1502
+ );
1503
+ }
1013
1504
 
1014
1505
  // SSD cold tier (docs/design/ssd-kv-cold-tier.md): prefix KV survives RAM
1015
1506
  // eviction and restarts. Compatibility key = configFingerprint (graph
1016
1507
  // shape) + the EFFECTIVE kv scheme (flags pick bf16 vs config vs uniform
1017
1508
  // on the same model — restored caches must match what serving produces) +
1018
1509
  // tokenizer hash (ids must keep meaning the same text).
1510
+ // RAM prompt-cache cap: 8 GB default (Josh's call, 2026-07-06). The old
1511
+ // flat 2e9 was an anti-OOM reflex sized for a 1B model in 24 GB — a
1512
+ // single full-context 12B entry is ~0.6 GB, so it silently shrank to
1513
+ // "three entries" on big models. 8 GB holds a dozen 12B contexts; idle
1514
+ // entries still demote to SSD and LRU eviction bounds pressure.
1515
+ // --prompt-cache <GB> overrides; 0 disables.
1516
+ const promptCacheCap = serverOptions.promptCacheBytes ?? 8e9;
1019
1517
  let ssdStore: SsdCacheStore | null = null;
1020
1518
  if (serverOptions.ssdCacheDir) {
1021
- if ((serverOptions.promptCacheBytes ?? 2e9) <= 0)
1519
+ if (promptCacheCap <= 0)
1022
1520
  throw new Error("--ssd-cache requires the RAM prompt cache (--prompt-cache 0 disables it)");
1023
- const schemeKey = kvScheme.kvBits
1521
+ const schemeKey = kvScheme.turboQuant
1522
+ ? `turbo-k${kvScheme.turboQuant.kBits}v${kvScheme.turboQuant.vBits}`
1523
+ : kvScheme.kvBits
1024
1524
  ? `kv${kvScheme.kvBits}`
1025
1525
  : kvScheme.kvConfig?.length ? "config" : "bf16";
1026
1526
  const tokJson = readFileSync(`${ctx.model.config.modelDir}/tokenizer.json`);
@@ -1039,11 +1539,51 @@ export function createServer(
1039
1539
  );
1040
1540
  }
1041
1541
 
1542
+ // Layer 0 (unified-engine plan): the cold tier is bound INTO the prompt
1543
+ // cache — take() itself tiers over it, so the batch scheduler and every
1544
+ // other consumer get SSD restores through the same API; eviction AND
1545
+ // idle demotion spill into it.
1546
+ const coldTier = ssdStore
1547
+ ? {
1548
+ find: (prompt: number[], ns: string) => {
1549
+ const h = ssdStore!.find(prompt, ns);
1550
+ return h ? { prefixLen: h.prefixLen, handle: h.entry } : null;
1551
+ },
1552
+ restore: (handle: unknown) => {
1553
+ const loaded = ssdStore!.restore(handle as import("./ssd-cache").SsdIndexEntry, ctx.model);
1554
+ if (!loaded) return null;
1555
+ // Restore is a STREAMED COPY (2026-07-07 A7-restore): the caches
1556
+ // own their bytes and no mapping outlives loadKvCache — nothing
1557
+ // to pin, nothing to unmap. retain stays in the entry contract
1558
+ // as a no-op so callers' dispose ordering is unchanged.
1559
+ return { tokens: loaded.tokens, caches: loaded.caches, retain: () => {} };
1560
+ },
1561
+ store: (tokens: number[], caches: import("./model/gemma4").Cache[], ns: string) => {
1562
+ ssdStore!.store(tokens, caches, ns);
1563
+ },
1564
+ }
1565
+ : null;
1042
1566
  const promptCache = new PromptCache(
1043
- serverOptions.promptCacheBytes ?? 2e9,
1044
- // RAM eviction spills to the cold tier (whole-prefix entries, no work
1045
- // on the token loop — the spill runs at eviction time only).
1046
- ssdStore ? (entry) => { ssdStore!.store(entry.tokens, entry.caches, entry.ns); } : null,
1567
+ promptCacheCap,
1568
+ // RAM eviction AND idle demotion spill to the cold tier NON-BLOCKING
1569
+ // (2026-07-06, same contract as the write-behind snapshot below): the
1570
+ // cache hands us OWNED zero-copy clones + copied tokens (made under
1571
+ // the generation lock — microseconds; entries are immutable so the
1572
+ // clones stay consistent), and the flush goes through the BOUNDED
1573
+ // SpillQueue -> storeAsync (idle-gated per tensor, see below) ->
1574
+ // clone disposal on every settle/drop path (that dispose is what
1575
+ // actually frees the demoted GPU memory; the queue's byte cap keeps
1576
+ // starved-gate retention bounded — 2026-07-07 review fix). spillQueue
1577
+ // is declared with the snapshot scheduler below; safe to close over
1578
+ // here because spills only fire at put()/demoteIdle time, long after
1579
+ // init.
1580
+ ssdStore
1581
+ ? {
1582
+ spillOwned: (entry) =>
1583
+ spillQueue!.enqueue({ tokens: entry.tokens, caches: entry.caches, ns: entry.ns }),
1584
+ }
1585
+ : null,
1586
+ coldTier,
1047
1587
  );
1048
1588
  // Responses-API store for previous_response_id resumption (Phase 11):
1049
1589
  // TTL + byte-capped LRU, port of optiq/response_store.py. Pairs with
@@ -1057,68 +1597,143 @@ export function createServer(
1057
1597
  * `enqueue` promise queue is gone; the gateway's AsyncMutex is THE lock.)
1058
1598
  * onToken returning `false` halts generation early (stop
1059
1599
  * sequence fired); the cache snapshot stays valid and is still kept.
1060
- * Vision requests bypass the prompt cache: image tokens are
1061
- * identical placeholder ids, so prefix matching across different
1062
- * images would false-hit. */
1600
+ * Media (vision/audio embeddings-prefill) requests bypass the prompt
1601
+ * cache: soft tokens are identical placeholder ids, so prefix matching
1602
+ * across different images/clips would false-hit. */
1063
1603
  const runGeneration = async (
1064
1604
  promptIds: number[],
1065
1605
  options: GenerateOptions,
1066
1606
  onToken: (token: number, logprobs?: TokenLogprobs) => void | boolean | Promise<void | boolean>,
1067
- vision?: { embeddings: import("./mlx/array").MlxArray; imageMask: import("./mlx/array").MlxArray },
1607
+ vision?: {
1608
+ embeddings: import("./mlx/array").MlxArray;
1609
+ imageMask?: import("./mlx/array").MlxArray;
1610
+ multimodalMask?: import("./mlx/array").MlxArray;
1611
+ },
1068
1612
  ) => {
1613
+ // Speculative decoding (serve --draft-model): spec-ELIGIBLE requests
1614
+ // decode through the verify loop; the rest fall through to the normal
1615
+ // serial path (never wrong results, just no speedup — logged once per
1616
+ // combination class would be noise, so silent). Eligibility v1:
1617
+ // text-only, base weights, no logprobs capture, bf16 KV — ALL quantized
1618
+ // KV axes excluded (affine kvBits/kvConfig AND turboQuant; the spec loop
1619
+ // builds fresh bf16 caches and never calls maybeQuantizeKv, so routing a
1620
+ // turbo request here would silently drop the operator's KV scheme —
1621
+ // 2026-07-07 review). Grammar COMPOSES (Phase C constrained verify walk
1622
+ // — see the serve-loop.ts header). Prompt-cache reuse is bypassed on
1623
+ // the spec path v1.
1624
+ if (
1625
+ ctx.draft &&
1626
+ !vision &&
1627
+ !options.adapters?.length &&
1628
+ !options.logprobs &&
1629
+ !options.kvBits &&
1630
+ !options.kvConfig &&
1631
+ !options.turboQuant &&
1632
+ // Belt: --paged-kv + --draft-model is refused at createServer; this
1633
+ // keeps the spec lane paged-free even for programmatic callers.
1634
+ !options.pagedKv
1635
+ ) {
1636
+ const { specServeRun } = await import("./spec/serve-loop");
1637
+ return specServeRun(
1638
+ ctx.model, ctx.draft.provider, ctx.draft.numDraftTokens,
1639
+ promptIds, options, onToken,
1640
+ );
1641
+ }
1069
1642
  // Cache entries are adapter-specific: KV computed under one adapter
1070
1643
  // must never seed another's (or the base's) prefill.
1071
1644
  const cacheNs = options.adapters?.join("+") ?? "";
1072
- let entry = vision ? null : promptCache.take(promptIds, cacheNs);
1073
- // Cold tier: restore from SSD when it has a STRICTLY longer usable
1074
- // prefix than RAM offered (zero-copy COW mmap; pages fault in lazily
1075
- // during the continuation prefill). The restored entry carries an
1076
- // unmap thunk that PromptCache runs after its caches are disposed.
1077
- if (!vision && ssdStore) {
1078
- const ssdHit = ssdStore.find(promptIds, cacheNs);
1079
- if (ssdHit && ssdHit.prefixLen > (entry?.tokens.length ?? 0)) {
1080
- const loaded = ssdStore.restore(ssdHit.entry, ctx.model);
1081
- if (loaded) {
1082
- const trimNeeded = loaded.tokens.length - ssdHit.prefixLen;
1083
- if (trimNeeded > 0 && !loaded.caches.every((c) => c.isTrimmable())) {
1084
- // Divergent tail on an untrimmable kind (ring post-wrap, SSM):
1085
- // this file can't seed the prompt — fall back to the RAM hit.
1086
- for (const c of loaded.caches) c.dispose();
1087
- loaded.mmap.unmap();
1088
- } else {
1089
- if (trimNeeded > 0) for (const c of loaded.caches) c.trim(trimNeeded);
1090
- if (entry) promptCache.put(entry.tokens, entry.caches, cacheNs, entry.retain); // return the RAM hit untouched
1091
- entry = { tokens: loaded.tokens.slice(0, ssdHit.prefixLen), caches: loaded.caches,
1092
- ns: cacheNs, retain: () => loaded.mmap.unmap() };
1093
- }
1094
- }
1095
- }
1096
- }
1645
+ // Paged-KV request scope (docs/design/paged-kv-cache.md): media
1646
+ // prompts (bidir overlay) and LoRA-adapter requests are v1 non-goals —
1647
+ // they run the PLAIN cache path even under --paged-kv (scope the flag
1648
+ // per request, never 400). Effective value computed ONCE so the
1649
+ // prompt-cache bypass below and the generate() options can't disagree.
1650
+ const pagedKv = vision || options.adapters?.length ? undefined : options.pagedKv;
1651
+ // Paged requests bypass the prompt cache entirely (v1 non-goal:
1652
+ // PagedKVCache has no cloneKvCaches/restore path — the vision
1653
+ // precedent). Fresh caches per request, disposed on completion.
1654
+ const skipPromptCache = Boolean(vision || pagedKv);
1655
+ // Both tiers in one call (Layer 0): take() prefers a strictly-longer
1656
+ // SSD prefix, restores it zero-copy, and trims — see PromptCache.take.
1657
+ const entry = skipPromptCache ? null : promptCache.take(promptIds, cacheNs);
1097
1658
  const caches = entry?.caches ?? ctx.model.makeCache();
1659
+ // Prompt-boundary snapshot (the multi-turn agent fix, 2026-07-04): the
1660
+ // prompt+gen entry put() below is UNTRIMMABLE at context > sliding
1661
+ // window (wrapped rings) and under quantized KV (mid-group), so any
1662
+ // decode→encode roundtrip drift in the reply the client sends back
1663
+ // turns the next turn into a total miss (measured: 12B turn-2 TTFT
1664
+ // 8.9 s instead of ~0.2 s). A prompt-ONLY entry is always an exact
1665
+ // prefix of the next turn's rendering regardless of reply drift.
1666
+ // Zero-copy (cloneKvCaches = slice views); only for substantial cold
1667
+ // prefills, where the re-prefill it saves is worth an extra entry.
1668
+ // The oracle invariant (mlx-lm insert_segments): a trim-free STRICT
1669
+ // prefix of the prompt exists for EVERY substantial request — cap the
1670
+ // boundary at len-1 so even a stableLen == len prompt (e4b: the
1671
+ // template tail survives the probe render) snapshots prompt[:-1]. An
1672
+ // exact repeat then matches with trimNeeded == 0, bypassing
1673
+ // isTrimmable() entirely — the only reuse path a wrapped ring has.
1674
+ const boundary = Math.min(options.snapshotAt ?? promptIds.length, promptIds.length - 1);
1675
+ // Re-snapshot on EVERY substantial request whose stable boundary extends
1676
+ // past the cached prefix; the clone is zero-copy views, so re-putting
1677
+ // is ~free.
1678
+ const snapshotBoundary =
1679
+ !skipPromptCache && boundary >= 256 && boundary > (entry?.tokens.length ?? 0);
1098
1680
  try {
1099
1681
  const gen = generate(ctx.model, promptIds, {
1100
1682
  ...options,
1683
+ pagedKv, // request-scoped (undefined strips the server-wide flag)
1101
1684
  cache: caches,
1102
- ...(vision ? { promptEmbeddings: vision.embeddings, imageMask: vision.imageMask } : {}),
1685
+ ...(snapshotBoundary
1686
+ ? {
1687
+ // snapshotAt MUST travel with the hook: generate() splits the
1688
+ // prefill at exactly this many tokens and fires the hook while
1689
+ // the caches hold exactly that prefix — putting boundary
1690
+ // tokens against caches at any other offset is silent KV
1691
+ // corruption.
1692
+ snapshotAt: boundary,
1693
+ onPrefillDone: () => {
1694
+ try {
1695
+ promptCache.put(promptIds.slice(0, boundary), cloneKvCaches(caches), cacheNs);
1696
+ } catch (e) {
1697
+ console.warn(`prompt-boundary snapshot skipped: ${(e as Error).message}`);
1698
+ }
1699
+ },
1700
+ }
1701
+ : {}),
1702
+ ...(vision
1703
+ ? {
1704
+ promptEmbeddings: vision.embeddings,
1705
+ ...(vision.imageMask ? { imageMask: vision.imageMask } : {}),
1706
+ ...(vision.multimodalMask ? { multimodalMask: vision.multimodalMask } : {}),
1707
+ }
1708
+ : {}),
1103
1709
  });
1104
1710
  for await (const t of gen) {
1105
1711
  if ((await onToken(t.token, t.logprobs)) === false) break;
1106
1712
  }
1107
1713
  const s = gen.stats!; // set on completion AND on early break
1108
- if (vision) {
1714
+ if (skipPromptCache) {
1715
+ // Vision and paged-KV requests own their caches for exactly one
1716
+ // generation (paged: v1 non-goal — no PromptCache integration).
1109
1717
  for (const c of caches) c.dispose();
1110
1718
  } else {
1719
+ // put() fires onPut → the debounced write-behind SSD snapshot
1720
+ // (wired below), covering the batch lane's puts too.
1111
1721
  promptCache.put(s.cacheTokens, caches, cacheNs, entry?.retain);
1112
- scheduleSsdSnapshot(s.cacheTokens, cacheNs);
1113
1722
  }
1114
1723
  return s;
1115
1724
  } catch (e) {
1116
1725
  for (const c of caches) c.dispose();
1117
1726
  entry?.retain?.();
1727
+ // A throw BEFORE generate()/the gateway took ownership (promptCache
1728
+ // take / makeCache) would leak the grammar's WASM matcher; dispose()
1729
+ // is idempotent, so this is safe when the throw came from inside run
1730
+ // (whose finally already disposed it).
1731
+ options.grammar?.dispose();
1118
1732
  throw e;
1119
1733
  } finally {
1120
1734
  vision?.embeddings.dispose();
1121
- vision?.imageMask.dispose();
1735
+ vision?.imageMask?.dispose();
1736
+ vision?.multimodalMask?.dispose();
1122
1737
  options.visionPixels?.dispose();
1123
1738
  }
1124
1739
  };
@@ -1132,24 +1747,134 @@ export function createServer(
1132
1747
  // extraction must not race a generation mutating the entry); if the entry
1133
1748
  // was evicted (already spilled) or extended (a newer schedule pending)
1134
1749
  // meanwhile, findExact misses and nothing is written.
1750
+ // Keyed by ns AND entry length: a sub-second final [prompt+gen] put must
1751
+ // not cancel the pending boundary-snapshot write (they are DIFFERENT
1752
+ // entries; for wrapped-ring models the boundary file is the only one a
1753
+ // restart can use). Same-length reschedules still coalesce; stale keys
1754
+ // self-clean at fire time (findExact misses once the RAM tier superseded
1755
+ // the entry, so nothing extra is written).
1756
+ //
1757
+ // NON-BLOCKING (2026-07-06, the write-behind persistence contract): the
1758
+ // gateway lock is held only for a zero-copy SNAPSHOT (findExact +
1759
+ // cloneKvCaches — microseconds; entries are immutable so the clones are
1760
+ // consistent forever). The flush itself runs OFF the lock via
1761
+ // storeAsync, and writes chain serially so two multi-hundred-MB flushes
1762
+ // never overlap.
1763
+ //
1764
+ // IDLE-GATED (2026-07-07, the decode@ctx fix): "off the lock" was not
1765
+ // enough — every per-tensor flush step is a blocking GPU sync on the
1766
+ // SAME stream decode uses (ops.contiguous enqueues a kernel, rawBytesView
1767
+ // evals) plus a synchronous multi-MB writeSync, and the setImmediate
1768
+ // pacing interleaved those slices exactly between decode tokens. A ~16k
1769
+ // entry's flush overlapping the bench's ctx repeats depressed decode@ctx
1770
+ // ~9% on e4b (mlx-lm runs no equivalent background work). Now every step
1771
+ // — including the first — awaits gateway.onIdle() (ssdFlushGate), so the
1772
+ // flush only progresses while NOTHING is generating and pauses when a
1773
+ // request arrives mid-flush. Tradeoffs, accepted: durability waits for a
1774
+ // quiet moment (single-user serving quiesces constantly; sustained
1775
+ // hammering defers the flush AND the spill clones' GPU-memory release —
1776
+ // bounded by the chain), and one in-flight tensor step (~10-15 MB) can
1777
+ // still land ahead of a just-arrived request. MLX_BUN_SSD_WRITEBEHIND=0
1778
+ // disables write-behind snapshots entirely — the paired-A/B lever + kill
1779
+ // switch (restart survival then degrades to spill-on-evict only).
1780
+ const writeBehindOn = process.env.MLX_BUN_SSD_WRITEBEHIND !== "0";
1781
+ const ssdFlushGate = (): Promise<void> => gateway.onIdle();
1135
1782
  const ssdPending = new Map<string, ReturnType<typeof setTimeout>>();
1783
+ // Bounded write-behind queue (2026-07-07 review fix — see SpillQueue in
1784
+ // kv-store.ts): pending clones pin their entries' GPU buffers while the
1785
+ // idle gate starves under sustained traffic, so QUEUED bytes are capped —
1786
+ // over cap the oldest queued spill drops (clone disposed immediately, a
1787
+ // future cache miss, never a wrong result) instead of accumulating past
1788
+ // the prompt-cache cap. Default 2 GB (a quarter of the 8 GB RAM-cache
1789
+ // default); MLX_BUN_SSD_SPILL_QUEUE_GB overrides. There is deliberately
1790
+ // NO shutdown flush: exiting under traffic loses queued write-behind
1791
+ // entries (restart survival degrades to whatever flushed) — accepted for
1792
+ // a best-effort cache tier; drops/pending are visible in /stats.
1793
+ // `|| 2` would coerce an explicit "0" back to 2 GB — parse so 0 works
1794
+ // (cap 0 = keep only the newest + in-flight clone pinned; the soft cap
1795
+ // never drops the item just enqueued).
1796
+ const spillQueueGbRaw = Number(process.env.MLX_BUN_SSD_SPILL_QUEUE_GB);
1797
+ const spillQueueCapBytes =
1798
+ (Number.isFinite(spillQueueGbRaw) && spillQueueGbRaw >= 0 ? spillQueueGbRaw : 2) * 1024 ** 3;
1799
+ const spillQueue = ssdStore
1800
+ ? new SpillQueue(
1801
+ spillQueueCapBytes,
1802
+ cacheBytes,
1803
+ (item) => ssdStore!.storeAsync(item.tokens, item.caches, item.ns, ssdFlushGate),
1804
+ (caches) => { for (const c of caches) c.dispose(); },
1805
+ )
1806
+ : null;
1136
1807
  const scheduleSsdSnapshot = (tokens: number[], ns: string): void => {
1137
- if (!ssdStore || tokens.length === 0) return;
1138
- const prev = ssdPending.get(ns);
1808
+ if (!ssdStore || !writeBehindOn || tokens.length === 0) return;
1809
+ const key = `${ns}|${tokens.length}`;
1810
+ const prev = ssdPending.get(key);
1139
1811
  if (prev) clearTimeout(prev);
1140
- ssdPending.set(ns, setTimeout(() => {
1141
- ssdPending.delete(ns);
1812
+ const fire = (): void => {
1813
+ ssdPending.delete(key);
1814
+ // Same guard as demoteIdle below (2026-07-07 review fix): the
1815
+ // runExclusive registers a serial waiter, which HOLDS BATCH ADMISSION
1816
+ // until every running row drains — a background durability timer must
1817
+ // never do that. Busy → re-arm and try again in 5 s (entries are
1818
+ // immutable; a late snapshot is just as valid).
1819
+ if (gateway.activeRows > 0 || gateway.pendingRows > 0) {
1820
+ const t = setTimeout(fire, 5000);
1821
+ t.unref?.();
1822
+ ssdPending.set(key, t);
1823
+ return;
1824
+ }
1142
1825
  void gateway.runExclusive(async () => {
1143
1826
  const e = promptCache.findExact(tokens, ns);
1144
- if (e) ssdStore!.store(e.tokens, e.caches, ns);
1827
+ if (!e) return null;
1828
+ return { tokens: e.tokens, caches: cloneKvCaches(e.caches) };
1829
+ }).then((snap) => {
1830
+ if (!snap) return;
1831
+ spillQueue!.enqueue({ tokens: snap.tokens, caches: snap.caches, ns });
1145
1832
  }).catch(() => { /* cold tier is best-effort */ });
1146
- }, 1000));
1833
+ };
1834
+ ssdPending.set(key, setTimeout(fire, 1000));
1147
1835
  };
1836
+ // Every put() — serial lane AND batch scheduler — schedules the snapshot
1837
+ // (Layer 0: batch-lane entries survive restarts too, not just evictions).
1838
+ promptCache.onPut = (tokens, ns) => scheduleSsdSnapshot(tokens, ns);
1839
+
1840
+ // Idle demotion (Layer 0): entries unused for --ssd-demote-idle seconds
1841
+ // spill to SSD and free their GPU memory — the RAM tier drains between
1842
+ // bursts while every prefix stays reachable (take() restores via
1843
+ // zero-copy mmap). The exclusive section below covers only the sweep's
1844
+ // zero-copy clone + dispose (spillOwned above); the SSD write itself
1845
+ // runs off-lock on ssdWriteChain. Swept only when the engine is TRULY idle: the
1846
+ // runExclusive below registers as a serial waiter, which would DRAIN a
1847
+ // running batch, so the activity check guards it (a momentary exclusive
1848
+ // grab of an idle engine is free). 0 disables.
1849
+ const demoteIdleMs = (serverOptions.ssdDemoteIdleSec ?? (ssdStore ? 300 : 0)) * 1000;
1850
+ let demoteTimer: ReturnType<typeof setInterval> | null = null;
1851
+ if (ssdStore && demoteIdleMs > 0) {
1852
+ demoteTimer = setInterval(() => {
1853
+ if (gateway.activeRows > 0 || gateway.pendingRows > 0) return;
1854
+ void gateway.runExclusive(async () => {
1855
+ const n = promptCache.demoteIdle(demoteIdleMs);
1856
+ if (n > 0) console.log(`[ssd-cache] demoted ${n} idle entr${n === 1 ? "y" : "ies"} to disk`);
1857
+ }).catch(() => {});
1858
+ }, Math.max(30_000, demoteIdleMs / 4));
1859
+ demoteTimer.unref?.();
1860
+ }
1148
1861
 
1149
1862
  // The lane picker: routes each request to the serial path (runGeneration,
1150
1863
  // above) or the continuous-batching scheduler, keeping the two off the GPU
1151
1864
  // (and shared loraState) at the same time. See src/serve/generation-gateway.ts.
1152
- const gateway = new GenerationGateway(ctx.model, batch, runGeneration);
1865
+ const gateway = new GenerationGateway(ctx.model, batch, runGeneration, {
1866
+ kvBudgetBytes: serverOptions.kvBudgetBytes,
1867
+ // Phase 3.1: the server-wide scheme travels to the gateway once; the
1868
+ // gateway batches kv-quant requests when the scheme is the batchable
1869
+ // kvConfig/full-attention composition (see #kvBatchable) and the
1870
+ // scheduler applies it — otherwise those requests route serial as before.
1871
+ kvScheme,
1872
+ // Phase 3.2: batch-lane prompt-cache reuse — joiners take() the longest
1873
+ // usable prefix (multi-turn chat TTFT under --batch N); never-merged
1874
+ // lone rows put() back on finish. Vision/adapter requests never batch,
1875
+ // so the serial lane's bypass rules are preserved by routing.
1876
+ promptCache,
1877
+ });
1153
1878
 
1154
1879
  // Admission ceiling, resolved once (Phase 5 memoryBudget enforcement).
1155
1880
  // fit() solves max safe context from weights + KV growth + prefill
@@ -1265,9 +1990,30 @@ export function createServer(
1265
1990
  stopSequences: (typeof req.stop === "string" ? [req.stop] : req.stop ?? [])
1266
1991
  .filter((s) => typeof s === "string" && s.length > 0),
1267
1992
  ...kvScheme,
1993
+ // Paged KV is a server-wide mode like the kv scheme above (v1:
1994
+ // serial bf16 gemma4 — the createServer refusals hold the invariants).
1995
+ ...(serverOptions.pagedKv ? { pagedKv: serverOptions.pagedKv } : {}),
1268
1996
  };
1269
1997
  };
1270
1998
 
1999
+ /** Compile a grammar controller from the request's structured-output fields
2000
+ * (response_format / guided_* / structured_outputs). Returns null when no
2001
+ * constraint is requested OR when compile failed (degrade path). On degrade,
2002
+ * `degradeHint` carries a human description for the system-prompt injection
2003
+ * + Warning header (oMLX parity — never 500). Honors MLX_BUN_GRAMMAR=0. */
2004
+ const compileGrammarForRequest = async (
2005
+ req: ChatRequest,
2006
+ ): Promise<{ controller: import("./grammar").GrammarController | null; degradeHint: string | null }> => {
2007
+ if (!grammarEnabled()) return { controller: null, degradeHint: null };
2008
+ const r = await compileGrammarRequest(
2009
+ req as GrammarRequest,
2010
+ ctx.tokenizer,
2011
+ ctx.model.config.text.vocabSize,
2012
+ );
2013
+ if (!r) return { controller: null, degradeHint: null };
2014
+ return { controller: r.controller, degradeHint: r.degradeHint };
2015
+ };
2016
+
1271
2017
  const templateOptionsFor = (req: ChatRequest, tools: ToolDefinition[] | null) => {
1272
2018
  // enableThinking resolution (and its precedence) lives in
1273
2019
  // resolveEnableThinking so template rendering and sampling stay in sync.
@@ -1278,13 +2024,60 @@ export function createServer(
1278
2024
  req: ChatRequest,
1279
2025
  tools: ToolDefinition[] | null,
1280
2026
  ): { ids: number[]; startInThinking: boolean } => {
1281
- const rendered = ctx.template.render(normalizeMessages(req.messages), templateOptionsFor(req, tools));
2027
+ const opts = templateOptionsFor(req, tools);
2028
+ const rendered = ctx.template.render(normalizeMessages(req.messages), opts);
1282
2029
  const ids = ctx.tokenizer.encode(rendered);
1283
2030
  // template includes <bos>; tokenizer post-processor also prepends one
1284
2031
  const trimmed = ids[0] === ids[1] && ids[0] === ctx.tokenizer.bosTokenId ? ids.slice(1) : ids;
1285
2032
  return { ids: trimmed, startInThinking: promptEndsInOpenThink(rendered) };
1286
2033
  };
1287
2034
 
2035
+ /** STABLE cache boundary: the prompt's tail is a generation primer (e.g.
2036
+ * 12B's `<|channel>thought` tokens) that the NEXT turn's re-render does
2037
+ * not contain — a cache entry ending there can never prefix-match again
2038
+ * once the rings wrap. Probe by rendering this conversation as if a
2039
+ * reply existed; the common prefix is what every future turn preserves.
2040
+ * The probe is EXPENSIVE (a second full render + encode — ~150 ms at
2041
+ * 16k with our tokenizer; 30-50% of the 12B ctx-repeat TTFT,
2042
+ * 2026-07-06b), but the PRIMER LENGTH it measures is constant per
2043
+ * template mode: the divergence sits after a special-token turn
2044
+ * delimiter, and special tokens break BPE merges, so message content
2045
+ * can't shift it. So the full probe runs ONCE per mode and later
2046
+ * requests pay a subtraction. (A wrong boundary is a QUALITY knob, not
2047
+ * a correctness one — any prefix ≤ len-1 is a valid snapshot point.) */
2048
+ const primerLenByMode = new Map<string, number>();
2049
+ const stableLenFor = (
2050
+ req: ChatRequest,
2051
+ tools: ToolDefinition[] | null,
2052
+ trimmed: number[],
2053
+ ): number => {
2054
+ const opts = templateOptionsFor(req, tools);
2055
+ const mode = `${opts.enableThinking}|${!!tools?.length}`;
2056
+ const primer = primerLenByMode.get(mode);
2057
+ if (primer !== undefined) return Math.max(0, trimmed.length - primer);
2058
+ let stableLen = trimmed.length;
2059
+ try {
2060
+ const probe = ctx.tokenizer.encode(
2061
+ ctx.template.render(
2062
+ [...normalizeMessages(req.messages), { role: "assistant", content: "x" }],
2063
+ opts,
2064
+ ),
2065
+ );
2066
+ const probeTrimmed =
2067
+ probe[0] === probe[1] && probe[0] === ctx.tokenizer.bosTokenId ? probe.slice(1) : probe;
2068
+ let i = 0;
2069
+ const n = Math.min(trimmed.length, probeTrimmed.length);
2070
+ while (i < n && trimmed[i] === probeTrimmed[i]) i++;
2071
+ stableLen = i;
2072
+ // Memoize the mode's primer length (sanity-capped: a "primer" longer
2073
+ // than 64 tokens means the probe diverged for content reasons —
2074
+ // don't generalize that).
2075
+ const p = trimmed.length - stableLen;
2076
+ if (p >= 0 && p <= 64) primerLenByMode.set(mode, p);
2077
+ } catch { /* probe is best-effort; full boundary is the fallback */ }
2078
+ return stableLen;
2079
+ };
2080
+
1288
2081
  const toolStreamMode = (tools: ToolDefinition[] | null): ToolStreamMode =>
1289
2082
  // Gemma-4 keeps the token-level sentinel path; every other family
1290
2083
  // (MiniCPM5, Qwen3/3.5, Tier-0 generics) emits tool calls as DECODED TEXT
@@ -1294,9 +2087,18 @@ export function createServer(
1294
2087
  const toolRouter = (tools: ToolDefinition[] | null): ToolAwareStream =>
1295
2088
  new ToolAwareStream(ctx.tokenizer, toolStreamMode(tools), tools);
1296
2089
 
2090
+ if (serverOptions.unixSocket) {
2091
+ try { require("node:fs").unlinkSync(serverOptions.unixSocket); } catch {}
2092
+ }
2093
+ // Bun's types make unix/port mutually exclusive variants that a spread
2094
+ // union can't prove — runtime accepts either; cast once at the call.
1297
2095
  serverRef = Bun.serve({
1298
- port,
1299
- ...(serverOptions.hostname ? { hostname: serverOptions.hostname } : {}),
2096
+ ...(serverOptions.unixSocket
2097
+ ? ({ unix: serverOptions.unixSocket } as unknown as Record<string, never>)
2098
+ : {
2099
+ port,
2100
+ ...(serverOptions.hostname ? { hostname: serverOptions.hostname } : {}),
2101
+ }),
1300
2102
  idleTimeout: 0,
1301
2103
  // Web chat rides pi's AgentSession events over a WebSocket; the embedded
1302
2104
  // pi provider points back at THIS server's own loopback /v1 (port
@@ -1305,9 +2107,21 @@ export function createServer(
1305
2107
  port: () => serverRef.port ?? port,
1306
2108
  modelId: ctx.modelId,
1307
2109
  contextWindow: ctx.model.config.text.maxPositionEmbeddings,
1308
- // capability flag — true if a tower is loaded or loadable (lazy)
2110
+ // capability flags — true if a tower is loaded or loadable (lazy)
1309
2111
  vision: !!(ctx.vision || ctx.loadVision),
2112
+ audio: !!(ctx.audio || ctx.loadAudio),
1310
2113
  thinking: ctx.template.supportsThinking,
2114
+ // Model-author sampling defaults (generation_config.json, server-CLI
2115
+ // overrides applied) for the sampling popover's per-model recommended
2116
+ // values (web-ui-pass-plan.md #14 groundwork) — same resolution
2117
+ // /v1/models' gen_defaults uses, so the two surfaces agree. The
2118
+ // think-mode-vs-not distinction lives in toOptions (per-request); this
2119
+ // is the single "model's own defaults" snapshot sent once at ready.
2120
+ genDefaults: {
2121
+ temperature: serverOptions.defaultTemperature ?? ctx.genDefaults.temperature ?? null,
2122
+ topP: serverOptions.defaultTopP ?? ctx.genDefaults.topP ?? null,
2123
+ topK: serverOptions.defaultTopK ?? ctx.genDefaults.topK ?? null,
2124
+ },
1311
2125
  }),
1312
2126
  async fetch(request, server) {
1313
2127
  const url = new URL(request.url);
@@ -1318,6 +2132,81 @@ export function createServer(
1318
2132
  return new Response("expected websocket", { status: 426 });
1319
2133
  }
1320
2134
 
2135
+ // Memory REST wrappers (web-chat-redesign.md §2.3/§9 Phase 2): thin
2136
+ // loopback JSON routes over src/memory/vault.ts for the web chat's
2137
+ // Memory panel. Handlers live in src/memory/rest.ts (pure functions,
2138
+ // no `ctx` dependency) so they're unit-testable without a loaded
2139
+ // model; dispatch here just matches path+method. The agent-tool
2140
+ // surface (src/memory/tools.ts) is untouched by this block.
2141
+ if (url.pathname === "/api/memory/status" && request.method === "GET") {
2142
+ const { handleMemoryStatus } = await import("./memory/rest");
2143
+ return handleMemoryStatus();
2144
+ }
2145
+ if (url.pathname === "/api/memory/list" && request.method === "GET") {
2146
+ const { handleMemoryList } = await import("./memory/rest");
2147
+ return handleMemoryList();
2148
+ }
2149
+ if (url.pathname === "/api/memory/search" && request.method === "GET") {
2150
+ const { handleMemorySearch } = await import("./memory/rest");
2151
+ return handleMemorySearch(url);
2152
+ }
2153
+ if (url.pathname === "/api/memory/article" && request.method === "GET") {
2154
+ const { handleMemoryArticle } = await import("./memory/rest");
2155
+ return handleMemoryArticle(url);
2156
+ }
2157
+ if (url.pathname === "/api/memory/links" && request.method === "GET") {
2158
+ const { handleMemoryLinks } = await import("./memory/rest");
2159
+ return handleMemoryLinks(url);
2160
+ }
2161
+ if (url.pathname === "/api/memory/history" && request.method === "GET") {
2162
+ const { handleMemoryHistory } = await import("./memory/rest");
2163
+ return handleMemoryHistory(url);
2164
+ }
2165
+ if (url.pathname === "/api/memory/diff" && request.method === "GET") {
2166
+ const { handleMemoryDiff } = await import("./memory/rest");
2167
+ return handleMemoryDiff(url);
2168
+ }
2169
+ if (url.pathname === "/api/memory/init" && request.method === "POST") {
2170
+ const { handleMemoryInit } = await import("./memory/rest");
2171
+ return handleMemoryInit(request);
2172
+ }
2173
+
2174
+ // Model Hub (web-chat-redesign.md §9 Phase 3, beat-matrix Axis 3):
2175
+ // thin loopback JSON routes over the registry/download/fit machinery
2176
+ // for the web chat's Hub panel. Handlers live in src/hub-rest.ts
2177
+ // (pure functions, no `ctx` dependency); dispatch here just matches
2178
+ // path+method, same convention as the memory wrappers above.
2179
+ if (url.pathname === "/api/hub/local" && request.method === "GET") {
2180
+ const { handleHubLocal } = await import("./hub-rest");
2181
+ return handleHubLocal();
2182
+ }
2183
+ if (url.pathname === "/api/hub/search" && request.method === "GET") {
2184
+ const { handleHubSearch } = await import("./hub-rest");
2185
+ return handleHubSearch(url);
2186
+ }
2187
+ if (url.pathname === "/api/hub/download" && request.method === "POST") {
2188
+ const { handleHubDownload } = await import("./hub-rest");
2189
+ return handleHubDownload(request);
2190
+ }
2191
+ if (url.pathname === "/api/hub/serve" && request.method === "POST") {
2192
+ const { handleHubServe } = await import("./hub-rest");
2193
+ return handleHubServe(request);
2194
+ }
2195
+
2196
+ // Session full-text search + export (docs/design/web-chat-redesign.md
2197
+ // §9 Phase 3, beat-matrix Axis 10/11 — the full-text-search BEAT and
2198
+ // Markdown/JSON chat export). Handlers live in src/serve/session-search.ts
2199
+ // (pure, read-only, no `ctx` dependency), dispatched here by
2200
+ // path+method, same convention as the memory/hub blocks above.
2201
+ if (url.pathname === "/api/sessions/search" && request.method === "GET") {
2202
+ const { handleSessionsSearch } = await import("./serve/session-search");
2203
+ return handleSessionsSearch(url);
2204
+ }
2205
+ if (url.pathname === "/api/sessions/export" && request.method === "GET") {
2206
+ const { handleSessionsExport } = await import("./serve/session-search");
2207
+ return handleSessionsExport(url);
2208
+ }
2209
+
1321
2210
  // The unified SPA is served at "/"; legacy deep links redirect into
1322
2211
  // the hash router so old bookmarks still land on the right section.
1323
2212
  if (url.pathname === "/" && request.method === "GET") {
@@ -1325,6 +2214,50 @@ export function createServer(
1325
2214
  headers: { "content-type": "text/html; charset=utf-8" },
1326
2215
  });
1327
2216
  }
2217
+ // Vendored syntax-highlighting assets (src/web/vendor/*, no CDN — see
2218
+ // that dir's README). Long cache: content is versioned by the vendor
2219
+ // rebuild step, not by URL, but a hard reload always wins during dev.
2220
+ if (url.pathname === "/assets/hljs.js" && request.method === "GET") {
2221
+ return new Response(HLJS_JS, {
2222
+ headers: { "content-type": "text/javascript; charset=utf-8", "cache-control": "public, max-age=3600" },
2223
+ });
2224
+ }
2225
+ if (url.pathname === "/assets/hljs.css" && request.method === "GET") {
2226
+ return new Response(HLJS_CSS, {
2227
+ headers: { "content-type": "text/css; charset=utf-8", "cache-control": "public, max-age=3600" },
2228
+ });
2229
+ }
2230
+ // The frontend bundle (see the import comment above) — cache header
2231
+ // matches the other /assets/* entries; content is versioned by the
2232
+ // build step, not the URL, and a hard reload always wins during dev.
2233
+ if (url.pathname === "/assets/app.js" && request.method === "GET") {
2234
+ return new Response(APP_JS, {
2235
+ headers: { "content-type": "text/javascript; charset=utf-8", "cache-control": "public, max-age=3600" },
2236
+ });
2237
+ }
2238
+ // PWA manifest + icon (plan §9 Phase 3): linked from app.html's
2239
+ // <link rel="manifest">. Same versioned-by-content cache posture as
2240
+ // the other /assets/* entries above.
2241
+ if (url.pathname === "/manifest.webmanifest" && request.method === "GET") {
2242
+ return new Response(MANIFEST_WEBMANIFEST, {
2243
+ headers: { "content-type": "application/manifest+json; charset=utf-8", "cache-control": "public, max-age=3600" },
2244
+ });
2245
+ }
2246
+ if (url.pathname === "/assets/icon.svg" && request.method === "GET") {
2247
+ return new Response(ICON_SVG, {
2248
+ headers: { "content-type": "image/svg+xml; charset=utf-8", "cache-control": "public, max-age=3600" },
2249
+ });
2250
+ }
2251
+ // Service worker MUST be served at the root scope (/sw.js, not
2252
+ // /assets/sw.js) — a worker's default scope is the directory it's
2253
+ // served from, and the shell it manages includes "/" itself.
2254
+ // no-store: the browser's own update check needs a fresh byte
2255
+ // comparison every time to notice a new worker, not a cached 304.
2256
+ if (url.pathname === "/sw.js" && request.method === "GET") {
2257
+ return new Response(SW_JS, {
2258
+ headers: { "content-type": "text/javascript; charset=utf-8", "cache-control": "no-store" },
2259
+ });
2260
+ }
1328
2261
  // v2 HLG Curve Designer — served same-origin so /generate + /signal need no CORS.
1329
2262
  // Read fresh from disk in dev (edits show on reload, no restart); fall back to the
1330
2263
  // embedded copy when running as the compiled single binary.
@@ -1368,7 +2301,7 @@ export function createServer(
1368
2301
  // surface on their own. (scan() is INSERT-OR-REPLACE + prunes deleted,
1369
2302
  // so it's idempotent and cheap for a local cache.)
1370
2303
  await reg.scan();
1371
- const { visionCapable } = await import("./registry");
2304
+ const { visionCapable, audioCapable } = await import("./registry");
1372
2305
  const { supportTier } = await import("./model/support");
1373
2306
  const rows = [];
1374
2307
  // listCanonical: one row per repo (refs/main) — duplicate snapshots
@@ -1389,7 +2322,8 @@ export function createServer(
1389
2322
  rows.push({
1390
2323
  repo_id: m.repoId, model_type: m.modelType,
1391
2324
  size_bytes: m.sizeBytes, quant_bits: m.quantBits,
1392
- vision: visionCapable(m), supported, support_tier: tier,
2325
+ vision: visionCapable(m), audio: audioCapable(m),
2326
+ supported, support_tier: tier,
1393
2327
  serving: m.repoId === ctx.modelId,
1394
2328
  assessment,
1395
2329
  });
@@ -1516,7 +2450,13 @@ export function createServer(
1516
2450
  const layerTypes = ctx.model.config.text.layerTypes;
1517
2451
  const kvLayers: Record<string, number> = {};
1518
2452
  let kvMode = "bf16";
1519
- if (kvScheme.kvBits) {
2453
+ if (kvScheme.turboQuant) {
2454
+ // v1: full-attention layers only (sliding-window stays bf16 —
2455
+ // docs/design/turboquant-kv.md non-goal).
2456
+ const fullAttn = layerTypes.filter((l) => l !== "sliding_attention").length;
2457
+ kvMode = `turbo k${kvScheme.turboQuant.kBits}v${kvScheme.turboQuant.vBits}`;
2458
+ kvLayers[`turbo-k${kvScheme.turboQuant.kBits}v${kvScheme.turboQuant.vBits}`] = fullAttn;
2459
+ } else if (kvScheme.kvBits) {
1520
2460
  kvMode = `uniform-kv${kvScheme.kvBits}`;
1521
2461
  kvLayers[`kv${kvScheme.kvBits}`] = layerTypes.length;
1522
2462
  } else if (kvScheme.kvConfig) {
@@ -1548,6 +2488,13 @@ export function createServer(
1548
2488
  restores: ssdStore.stats.restores,
1549
2489
  spills: ssdStore.stats.spills,
1550
2490
  restore_ms_last: Math.round(ssdStore.stats.restoreMsLast),
2491
+ demotions: promptCache.demotions,
2492
+ // Bounded write-behind queue (2026-07-07): clones waiting for
2493
+ // the idle-gated flush, and spills the byte cap dropped
2494
+ // (each = one future cache miss, never a wrong result).
2495
+ pending_spills: spillQueue?.pendingCount ?? 0,
2496
+ pending_spill_bytes: spillQueue?.pendingBytes ?? 0,
2497
+ dropped_spills: spillQueue?.droppedCount ?? 0,
1551
2498
  },
1552
2499
  } : {}),
1553
2500
  response_store: {
@@ -1576,12 +2523,31 @@ export function createServer(
1576
2523
  configured: batch,
1577
2524
  batched: gateway.batchingEnabled,
1578
2525
  active_rows: gateway.activeRows,
2526
+ pending_rows: gateway.pendingRows,
2527
+ submitted_rows: gateway.submittedRows,
2528
+ kv_bytes: gateway.kvBytes.projected,
2529
+ kv_budget_bytes: gateway.kvBytes.budget,
1579
2530
  },
1580
2531
  });
1581
2532
  }
1582
2533
 
1583
2534
  // mlx_lm.server parity: GET /health → the exact body it writes
1584
2535
  // ('{"status": "ok"}', note the space) so byte-for-byte health checks pass.
2536
+ // Engine-mode admin (unix-socket children only — never exposed on
2537
+ // TCP): drain = quiesce the gateway + demote the whole prompt cache
2538
+ // to the SSD tier. The pool calls this before evicting a model
2539
+ // child, so its state survives the eviction losslessly.
2540
+ if (
2541
+ serverOptions.unixSocket &&
2542
+ url.pathname === "/admin/drain" &&
2543
+ request.method === "POST"
2544
+ ) {
2545
+ await gateway.runExclusive(async () => {
2546
+ promptCache.demoteIdle(0);
2547
+ });
2548
+ return Response.json({ drained: true, demotions: promptCache.demotions });
2549
+ }
2550
+
1585
2551
  if (url.pathname === "/health" && request.method === "GET") {
1586
2552
  return new Response('{"status": "ok"}', {
1587
2553
  headers: { "content-type": "application/json" },
@@ -1601,26 +2567,43 @@ export function createServer(
1601
2567
  ? decodeURIComponent(url.pathname.slice("/v1/models/".length))
1602
2568
  : null;
1603
2569
  const created = Math.floor(startedAt / 1000);
2570
+ // Model-author sampling defaults (generation_config.json), resolved once
2571
+ // at load (web-ui-pass-plan.md #14 groundwork): lets a client (the web
2572
+ // chat's sampling popover, external pi) show the SERVED model's actual
2573
+ // recommended values instead of one hardcoded shape.
2574
+ const genDefaultsOut = {
2575
+ temperature: ctx.genDefaults.temperature ?? null,
2576
+ top_p: ctx.genDefaults.topP ?? null,
2577
+ top_k: ctx.genDefaults.topK ?? null,
2578
+ };
1604
2579
  const data: Array<Record<string, unknown>> = [{
1605
2580
  id: ctx.modelId, object: "model", created, owned_by: "mlx-bun",
1606
2581
  context_window: ctx.model.config.text.maxPositionEmbeddings,
1607
2582
  // Capability flags for clients (CLI/external pi) that build a
1608
2583
  // provider from discovery — `reasoning` gates the thinking toggle,
1609
- // `vision` the image input declaration (tower loaded or lazily
1610
- // loadable; same signal the web embed uses).
2584
+ // `vision` the image input declaration, `audio` whether
2585
+ // `input_audio` content parts are accepted (tower loaded or lazily
2586
+ // loadable; same signals the web embed uses).
1611
2587
  reasoning: ctx.template.supportsThinking,
1612
2588
  vision: !!(ctx.vision || ctx.loadVision),
2589
+ audio: !!(ctx.audio || ctx.loadAudio),
2590
+ gen_defaults: genDefaultsOut,
1613
2591
  }];
1614
2592
  try {
1615
- const { Registry } = await import("./registry");
2593
+ const { Registry, visionCapable } = await import("./registry");
2594
+ const { supportTier } = await import("./model/support");
1616
2595
  const reg = new Registry();
1617
2596
  try {
1618
2597
  if (reg.list().length === 0) await reg.scan();
1619
2598
  // Canonical rows only — a repo with N snapshots is ONE model id.
1620
2599
  for (const m of reg.listCanonical()) {
1621
2600
  if (m.repoId === ctx.modelId) continue;
1622
- if (!isSupportedModelRecord(m.modelType, m.repoId)) continue;
1623
- data.push({ id: m.repoId, object: "model", created });
2601
+ const tier = supportTier(m.modelType, m.repoId);
2602
+ if (tier === null) continue;
2603
+ data.push({
2604
+ id: m.repoId, object: "model", created,
2605
+ vision: visionCapable(m), tier,
2606
+ });
1624
2607
  }
1625
2608
  } finally {
1626
2609
  reg.close();
@@ -1676,13 +2659,16 @@ export function createServer(
1676
2659
  // unmount. Mount and unmount go through the generation queue so
1677
2660
  // they never race an in-flight forward pass.
1678
2661
  if (url.pathname === "/v1/adapters/available" && request.method === "GET") {
1679
- // On-disk adapters that can be mounted (the chat selector's source),
1680
- // FILTERED to ones compatible with the currently-served model — a
1681
- // MiniCPM5 adapter can't mount on Gemma, so it must not appear in the
1682
- // picker. Compatibility = the adapter's recorded base repo id matches the
1683
- // served model (compared on the bare name, lenient about the org); an
1684
- // adapter with no recorded base is kept (mount validates it). `mounted`
1685
- // flags ones already loaded so the UI auto-loads on select only if needed.
2662
+ // On-disk adapters that can be mounted (the chat selector's source).
2663
+ // Every on-disk adapter is returned — the web chat routing chip (§5.2)
2664
+ // needs incompatible entries visible-but-grayed, not silently dropped,
2665
+ // so a user understands "10 adapters exist, 2 apply here" rather than
2666
+ // seeing a shorter list with no explanation. `compatible` = the
2667
+ // adapter's recorded base repo id matches the served model (compared
2668
+ // on the bare name, lenient about the org); an adapter with no
2669
+ // recorded base is treated as compatible (mount validates it for
2670
+ // real). `mounted` flags ones already loaded so the UI auto-loads on
2671
+ // select only if needed.
1686
2672
  const { homedir } = await import("node:os");
1687
2673
  const stores = [
1688
2674
  `${homedir()}/.cache/mlx-bun-finetunes`,
@@ -1692,10 +2678,10 @@ export function createServer(
1692
2678
  const bareName = (s: string) => s.split("/").pop()!.toLowerCase();
1693
2679
  const servedName = bareName(ctx.modelId);
1694
2680
  const adapters = (await listAvailableAdapters(stores))
1695
- .filter((a) => a.baseModel == null || bareName(a.baseModel) === servedName)
1696
2681
  .map((a) => ({
1697
2682
  id: a.id, path: a.path, rank: a.rank, scale: a.scale,
1698
2683
  base_model: a.baseModel, mounted: mounted.has(a.id),
2684
+ compatible: a.baseModel == null || bareName(a.baseModel) === servedName,
1699
2685
  }));
1700
2686
  return Response.json({ adapters });
1701
2687
  }
@@ -1704,6 +2690,7 @@ export function createServer(
1704
2690
  adapters: ctx.adapters.list().map((a) => ({
1705
2691
  id: a.id, path: a.path, rank: a.rank, scale: a.scale,
1706
2692
  size_bytes: a.sizeBytes, mounted_layers: a.mountedLayers,
2693
+ ram_bytes: a.ramBytes,
1707
2694
  })),
1708
2695
  });
1709
2696
  }
@@ -1723,7 +2710,7 @@ export function createServer(
1723
2710
  const info = await gateway.runExclusive(() => ctx.adapters.mount(body.id!, body.path!));
1724
2711
  return Response.json({
1725
2712
  id: info.id, mounted_layers: info.mountedLayers,
1726
- rank: info.rank, scale: info.scale,
2713
+ rank: info.rank, scale: info.scale, ram_bytes: info.ramBytes,
1727
2714
  });
1728
2715
  } catch (e) {
1729
2716
  return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
@@ -1841,29 +2828,140 @@ export function createServer(
1841
2828
  (m) => Array.isArray(m.content) &&
1842
2829
  m.content.some((p: any) => p.type === "image_url" || p.type === "image"),
1843
2830
  );
2831
+ // Same shapes extractAudio accepts: OpenAI-canonical input_audio plus
2832
+ // optiq's audio / audio_url aliases (audio-input-plan.md §3.2).
2833
+ const hasAudio = body.messages.some(
2834
+ (m) => Array.isArray(m.content) &&
2835
+ m.content.some(
2836
+ (p: any) => p.type === "input_audio" || p.type === "audio" || p.type === "audio_url",
2837
+ ),
2838
+ );
1844
2839
  let promptIds: number[];
2840
+ let stableLen: number | null = null;
1845
2841
  // Whether the prompt primed an open <think> (Qwen3.5/MiniCPM5 thinking
1846
2842
  // on) so the model's output starts mid-reasoning — seeds the splitter.
1847
2843
  // Vision is Gemma4-only (no switchable thinking channel), so it's false.
1848
2844
  let startInThinking = false;
1849
2845
  let vision: Parameters<typeof runGeneration>[3];
1850
2846
  let diffusionPixels: import("./mlx/array").MlxArray | null = null;
2847
+ // Grammar-constrained decoding (src/grammar.ts). Compile BEFORE prompt
2848
+ // rendering: on the degrade path (compile failed but a constraint was
2849
+ // requested) inject a system message instructing valid JSON so the
2850
+ // model still best-efforts schema-conformant output (oMLX parity —
2851
+ // _compile_grammar_for_request returning None + build_json_system_prompt).
2852
+ let grammarCtrl: import("./grammar").GrammarController | null = null;
2853
+ let grammarWarning: string | null = null;
2854
+ const grammarReq =
2855
+ body.response_format != null || !!body.guided_grammar ||
2856
+ !!body.guided_regex || !!body.guided_choice?.length ||
2857
+ body.structured_outputs != null;
2858
+ if (grammarReq) {
2859
+ const g = await compileGrammarForRequest(body);
2860
+ grammarCtrl = g.controller;
2861
+ if (!g.controller && g.degradeHint) {
2862
+ grammarWarning = `grammar not enforced: ${g.degradeHint} — falling back to prompt injection`;
2863
+ body = {
2864
+ ...body,
2865
+ messages: [
2866
+ { role: "system", content: degradeJsonSystemPrompt(body) },
2867
+ ...body.messages,
2868
+ ],
2869
+ };
2870
+ }
2871
+ }
2872
+ // Free what a rejected request has already allocated (grammar WASM
2873
+ // matcher/compiled-grammar, vision embeddings, diffusion pixels).
2874
+ // Every early return between here and gateway.run() must call this —
2875
+ // generate()/the gateway only take disposal ownership once running
2876
+ // (leaked one matcher per SSRF-blocked-media + response_format
2877
+ // request before the 2026-07-07 sweep).
2878
+ const disposeRejected = () => {
2879
+ grammarCtrl?.dispose();
2880
+ vision?.embeddings.dispose();
2881
+ vision?.imageMask?.dispose();
2882
+ vision?.multimodalMask?.dispose();
2883
+ diffusionPixels?.dispose();
2884
+ };
1851
2885
  try {
1852
- if (hasImages && ctx.model instanceof DiffusionGemmaModel) {
2886
+ if (hasAudio) {
2887
+ // Audio (and MIXED image+audio) input — A4 of
2888
+ // docs/design/audio-input-plan.md. One buildMultimodalPrompt call
2889
+ // splices both media kinds in document order. A request WITH
2890
+ // audio on a model whose tower is unavailable is an explicit 400
2891
+ // — never a silent text-only degrade (that leniency is only for
2892
+ // requests without the media, getVisionTower's contract).
2893
+ const audioTower = getAudioTower(ctx);
2894
+ if (!audioTower || !ctx.audioTokenIds) {
2895
+ disposeRejected();
2896
+ return Response.json(
2897
+ {
2898
+ error: {
2899
+ message:
2900
+ `model ${ctx.modelId} has no audio tower — audio input needs ` +
2901
+ `a model whose config.json carries audio_config and whose ` +
2902
+ `sidecar ships the audio tensors (e.g. gemma-4 e4b OptiQ)`,
2903
+ },
2904
+ },
2905
+ { status: 400 },
2906
+ );
2907
+ }
2908
+ let visionSide:
2909
+ | { tower: VisionEncoder; tokenIds: VisionTokenIds }
2910
+ | undefined;
2911
+ if (hasImages) {
2912
+ const tower = getVisionTower(ctx);
2913
+ if (!tower) {
2914
+ disposeRejected();
2915
+ return Response.json(
2916
+ { error: { message: "model has no vision sidecar" } }, { status: 400 },
2917
+ );
2918
+ }
2919
+ visionSide = { tower, tokenIds: ctx.visionTokenIds };
2920
+ }
2921
+ const { messages: withAudioParts, images } =
2922
+ await extractImages(normalizeMessages(body.messages));
2923
+ const { messages, audio } = await extractAudio(withAudioParts);
2924
+ // Non-WAV containers (mp3/m4a/flac/ogg/aiff/…) transcode to
2925
+ // 16 kHz WAV via CoreAudio; RIFF bytes pass through untouched.
2926
+ // Failures throw into the prompt-build 400 below.
2927
+ const wavs = await Promise.all(audio.map(ensureWav));
2928
+ // The towers are only ever non-null for Gemma4 (loader gates).
2929
+ const mp = await buildMultimodalPrompt(
2930
+ ctx.model as Gemma4Model,
2931
+ {
2932
+ ...(visionSide ? { vision: visionSide } : {}),
2933
+ audio: { tower: audioTower, tokenIds: ctx.audioTokenIds },
2934
+ },
2935
+ ctx.tokenizer, ctx.template, messages, images, wavs, tools,
2936
+ );
2937
+ promptIds = mp.ids;
2938
+ // bidirMask is null whenever audio is present (§3.3 Q1: mixed
2939
+ // prompts run fully causal); the union mask does the per-layer
2940
+ // id zeroing either way.
2941
+ vision = {
2942
+ embeddings: mp.embeddings,
2943
+ ...(mp.bidirMask ? { imageMask: mp.bidirMask } : {}),
2944
+ multimodalMask: mp.multimodalMask,
2945
+ };
2946
+ } else if (hasImages && ctx.model instanceof DiffusionGemmaModel) {
1853
2947
  // DiffusionGemma image-text-to-text: its OWN dedicated SigLIP tower +
1854
2948
  // encoder vision merge feed the denoising engine (NOT the AR
1855
2949
  // forwardEmbeddings path). v1 supports a single image.
1856
2950
  const dm = ctx.model;
1857
- if (!dm.visionTower)
2951
+ if (!dm.visionTower) {
2952
+ disposeRejected();
1858
2953
  return Response.json(
1859
2954
  { error: { message: "this checkpoint has no vision tower" } }, { status: 400 },
1860
2955
  );
2956
+ }
1861
2957
  const { messages, images } = await extractImages(normalizeMessages(body.messages));
1862
- if (images.length !== 1)
2958
+ if (images.length !== 1) {
2959
+ disposeRejected();
1863
2960
  return Response.json(
1864
2961
  { error: { message: "DiffusionGemma image input supports exactly one image" } },
1865
2962
  { status: 400 },
1866
2963
  );
2964
+ }
1867
2965
  const rendered = ctx.template.render(messages, { tools, addGenerationPrompt: true });
1868
2966
  const rawIds = ctx.tokenizer.encode(rendered, /* addSpecialTokens */ false);
1869
2967
  const { pixels, softTokens } = await dm.visionTower.preprocess(images[0]!);
@@ -1877,10 +2975,12 @@ export function createServer(
1877
2975
  // Loads (and caches) the tower on first image request — text-only
1878
2976
  // sessions never pay for it.
1879
2977
  const tower = getVisionTower(ctx);
1880
- if (!tower)
2978
+ if (!tower) {
2979
+ disposeRejected();
1881
2980
  return Response.json(
1882
2981
  { error: { message: "model has no vision sidecar" } }, { status: 400 },
1883
2982
  );
2983
+ }
1884
2984
  const { messages, images } = await extractImages(normalizeMessages(body.messages));
1885
2985
  // The tower is only ever non-null for Gemma4 (sidecar gate in
1886
2986
  // makeVisionLoader), so the model narrow is safe here.
@@ -1894,8 +2994,10 @@ export function createServer(
1894
2994
  const built = promptIdsFor(body, tools);
1895
2995
  promptIds = built.ids;
1896
2996
  startInThinking = built.startInThinking;
2997
+ stableLen = -1; // marker: chat path, probe lazily below
1897
2998
  }
1898
2999
  } catch (e) {
3000
+ disposeRejected();
1899
3001
  return Response.json(
1900
3002
  { error: { message: `prompt build failed: ${(e as Error).message}` } },
1901
3003
  { status: 400 },
@@ -1904,22 +3006,33 @@ export function createServer(
1904
3006
  let options: ReturnType<typeof toOptions>;
1905
3007
  try {
1906
3008
  options = toOptions(body);
3009
+ // Stable cache boundary for the prompt-boundary snapshot (chat
3010
+ // prompts end in a generation primer the next turn won't contain).
3011
+ // The probe costs a second full render+encode (~150 ms at 16k), so
3012
+ // run it ONLY when a snapshot can actually be taken: the cache
3013
+ // must not already hold the strict prefix (warm repeats peek at
3014
+ // len-1 → skip; the snapshot gate needs boundary > cached prefix).
3015
+ if (stableLen !== null) {
3016
+ const cachePeek = promptCache.peekPrefixLen(
3017
+ promptIds, options.adapters?.join("+") ?? "");
3018
+ if (cachePeek < promptIds.length - 1)
3019
+ options.snapshotAt = stableLenFor(body, tools, promptIds);
3020
+ }
1907
3021
  } catch (e) {
1908
3022
  // bad logit_bias (non-numeric keys/values) — mlx-lm's coercion error
1909
- vision?.embeddings.dispose();
1910
- vision?.imageMask.dispose();
1911
- diffusionPixels?.dispose();
3023
+ disposeRejected();
1912
3024
  return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
1913
3025
  }
1914
3026
  if (diffusionPixels) options.visionPixels = diffusionPixels;
3027
+ // Attach the compiled grammar controller (null when no constraint /
3028
+ // degrade — generate() runs the unmasked fast pipelined loop).
3029
+ if (grammarCtrl) options.grammar = grammarCtrl;
1915
3030
  // Admission: reject what cannot finish within the memory budget
1916
3031
  // (the GPU OOM it would otherwise hit is uncatchable and kills
1917
3032
  // the process — Phase 6 finding).
1918
3033
  const requiredCtx = promptIds.length + (options.maxTokens ?? 1024);
1919
3034
  if (requiredCtx > admission.maxSafeContext) {
1920
- vision?.embeddings.dispose();
1921
- vision?.imageMask.dispose();
1922
- diffusionPixels?.dispose();
3035
+ disposeRejected();
1923
3036
  return Response.json(
1924
3037
  {
1925
3038
  error: {
@@ -1941,6 +3054,7 @@ export function createServer(
1941
3054
  const adapterIds = ctx.adapters.resolveSpec(body.adapter ?? serverOptions.defaultAdapter);
1942
3055
  if (adapterIds.length) options.adapters = adapterIds;
1943
3056
  } catch (e) {
3057
+ disposeRejected();
1944
3058
  return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
1945
3059
  }
1946
3060
 
@@ -1958,14 +3072,18 @@ export function createServer(
1958
3072
  options.topLogprobs = topLogprobs;
1959
3073
  }
1960
3074
 
1961
- // What lane this request takes (vision / adapters / logits processors /
1962
- // a user-fixed seed → serial; everything else batches when --batch N).
3075
+ // What lane this request takes (vision/audio media / adapters /
3076
+ // logprobs / seed / explicit kv-quant / a mounted draft → serial;
3077
+ // sampler extras, repetition penalty, and grammar all BATCH — see
3078
+ // willBatch). `vision` is the embeddings-prefill payload for BOTH
3079
+ // media kinds, so hasVision routes audio serial too — batching is a
3080
+ // mode, not a fallback, and batched audio has no oracle (§3.2).
1963
3081
  const shape = {
1964
3082
  hasVision: !!vision,
1965
3083
  hasAdapters: !!options.adapters?.length,
1966
3084
  hasRepetitionPenalty: !!options.repetitionPenalty,
1967
- // min_p / XTC / logit_bias / presence+frequency penalties: serial-only
1968
- // in v1 (see GenerationGateway.RequestShape.hasLogitsExtras).
3085
+ // Informational since 2026-07-02: per-row logits processors batch;
3086
+ // willBatch no longer gates on these fields.
1969
3087
  hasLogitsExtras: !!(
1970
3088
  options.minP || options.xtcProbability || options.logitBias ||
1971
3089
  options.presencePenalty || options.frequencyPenalty
@@ -1973,10 +3091,26 @@ export function createServer(
1973
3091
  wantsLogprobs: captureLogprobs,
1974
3092
  userSeed: body.seed !== undefined,
1975
3093
  kvQuant: !!(options.kvConfig?.length || options.kvBits),
3094
+ turboQuant: !!options.turboQuant,
3095
+ // grammarCtrl is null on the degrade path (prompt injection) —
3096
+ // those stay batchable. A real controller batches via per-row
3097
+ // matchers (MLX_BUN_GRAMMAR_BATCH=0 forces it serial).
3098
+ hasGrammar: !!grammarCtrl,
3099
+ hasDraft: !!ctx.draft,
1976
3100
  };
1977
3101
  const batched = gateway.willBatch(shape);
1978
3102
  if (process.env.MLX_BUN_LANE_DEBUG === "1")
1979
3103
  console.error(`[lane] batched=${batched} shape=${JSON.stringify(shape)} t=${Date.now() % 100000}`);
3104
+ // Per-turn lane (docs/design/web-chat-redesign.md §2.3 caveat / risk #5):
3105
+ // reported on usage AND recorded in the in-process lane registry (keyed
3106
+ // by this response's `id`) so pi-web's WS bridge — which the pi SDK's
3107
+ // usage parsing can't carry a custom field through — can correlate it
3108
+ // via AssistantMessage.responseId. batched wins over hasDraft (mutually
3109
+ // exclusive by construction: willBatch already excludes hasDraft), so
3110
+ // the only ambiguity is serial vs serial+spec, resolved once s.spec is
3111
+ // known post-generation.
3112
+ const lane: Lane = batched ? "batched" : shape.hasDraft ? "serial+spec" : "serial";
3113
+ recordLane(id, lane);
1980
3114
 
1981
3115
  if (body.stream) {
1982
3116
  const stream = new ReadableStream<Uint8Array>({
@@ -2054,6 +3188,11 @@ export function createServer(
2054
3188
  ? "tool_calls"
2055
3189
  : stopper.stopped ? "stop"
2056
3190
  : s.generatedTokens >= (options.maxTokens ?? 1024) ? "length" : "stop";
3191
+ // Refine serial vs serial+spec now that s.spec is known (a mounted
3192
+ // draft can still decode zero speculative tokens on a very short
3193
+ // reply) and re-record for pi-web's WS correlation.
3194
+ const finalLane: Lane = batched ? "batched" : s.spec ? "serial+spec" : lane;
3195
+ if (finalLane !== lane) recordLane(id, finalLane);
2057
3196
  send({
2058
3197
  ...chunk({}, finish),
2059
3198
  usage: {
@@ -2061,6 +3200,8 @@ export function createServer(
2061
3200
  completion_tokens: s.generatedTokens,
2062
3201
  total_tokens: s.promptTokens + s.generatedTokens,
2063
3202
  prompt_tokens_details: { cached_tokens: s.cachedTokens },
3203
+ lane: finalLane,
3204
+ ...(s.spec ? { speculation: s.spec } : {}),
2064
3205
  },
2065
3206
  });
2066
3207
  // bare sentinel per the OpenAI spec — JSON.stringify would
@@ -2078,6 +3219,7 @@ export function createServer(
2078
3219
  "content-type": "text/event-stream",
2079
3220
  "cache-control": "no-cache",
2080
3221
  connection: "keep-alive",
3222
+ ...(grammarWarning ? { Warning: grammarWarning } : {}),
2081
3223
  },
2082
3224
  });
2083
3225
  }
@@ -2123,6 +3265,10 @@ export function createServer(
2123
3265
  : stopper.stopped ? "stop"
2124
3266
  : s.generatedTokens >= (options.maxTokens ?? 1024) ? "length" : "stop";
2125
3267
  const logprobsBlock = lpc?.payload() ?? null;
3268
+ // Refine serial vs serial+spec now that s.spec is known; re-record
3269
+ // for pi-web's WS correlation (see the streaming branch above).
3270
+ const finalLane: Lane = batched ? "batched" : s.spec ? "serial+spec" : lane;
3271
+ if (finalLane !== lane) recordLane(id, finalLane);
2126
3272
  return Response.json({
2127
3273
  id, object: "chat.completion", created, model: ctx.modelId,
2128
3274
  choices: [{
@@ -2141,8 +3287,10 @@ export function createServer(
2141
3287
  completion_tokens: s.generatedTokens,
2142
3288
  total_tokens: s.promptTokens + s.generatedTokens,
2143
3289
  prompt_tokens_details: { cached_tokens: s.cachedTokens },
3290
+ lane: finalLane,
3291
+ ...(s.spec ? { speculation: s.spec } : {}),
2144
3292
  },
2145
- });
3293
+ }, grammarWarning ? { headers: { Warning: grammarWarning } } : undefined);
2146
3294
  }
2147
3295
  } catch (e) {
2148
3296
  return Response.json({ error: { message: (e as Error).message } }, { status: 500 });
@@ -2196,12 +3344,34 @@ export function createServer(
2196
3344
  // bad logit_bias (non-numeric keys/values) — mlx-lm's coercion error
2197
3345
  return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
2198
3346
  }
3347
+ // Grammar-constrained decoding on raw completions too (response_format /
3348
+ // guided_* / structured_outputs). The degrade path for /v1/completions
3349
+ // has no chat template to inject a system message into, so a degrade
3350
+ // only emits the Warning header (no prompt injection) — documented gap
3351
+ // vs the chat lane, which mirrors oMLX's text-completions behavior.
3352
+ let textGrammarWarning: string | null = null;
3353
+ const textGrammarReq =
3354
+ body.response_format != null || !!body.guided_grammar ||
3355
+ !!body.guided_regex || !!body.guided_choice?.length ||
3356
+ body.structured_outputs != null;
3357
+ if (textGrammarReq) {
3358
+ const g = await compileGrammarForRequest(body as unknown as ChatRequest);
3359
+ if (g.controller) options.grammar = g.controller;
3360
+ else if (g.degradeHint)
3361
+ textGrammarWarning = `grammar not enforced: ${g.degradeHint} — no prompt injection on /v1/completions`;
3362
+ }
3363
+ // Mirrors the chat lane: a real controller (not the degrade path)
3364
+ // shapes the request for per-row grammar batching.
3365
+ const textGrammarCtrl = options.grammar ?? null;
2199
3366
  // mlx_lm.server's default max_tokens is 512 (its --max-tokens CLI
2200
3367
  // default). The chat lane's very generous default is wrong for raw
2201
3368
  // completion: with no template an EOS may never come.
2202
3369
  options.maxTokens = body.max_completion_tokens ?? body.max_tokens ?? serverOptions.defaultMaxTokens ?? 512;
2203
3370
  const requiredCtx = promptIds.length + options.maxTokens;
2204
- if (requiredCtx > admission.maxSafeContext)
3371
+ if (requiredCtx > admission.maxSafeContext) {
3372
+ // Pre-run reject: the gateway only takes grammar disposal ownership
3373
+ // once run — free the WASM matcher here (chat lane: disposeRejected).
3374
+ textGrammarCtrl?.dispose();
2205
3375
  return Response.json(
2206
3376
  {
2207
3377
  error: {
@@ -2216,10 +3386,12 @@ export function createServer(
2216
3386
  },
2217
3387
  { status: 400 },
2218
3388
  );
3389
+ }
2219
3390
  try {
2220
3391
  const adapterIds = ctx.adapters.resolveSpec(body.adapter ?? serverOptions.defaultAdapter);
2221
3392
  if (adapterIds.length) options.adapters = adapterIds;
2222
3393
  } catch (e) {
3394
+ textGrammarCtrl?.dispose();
2223
3395
  return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
2224
3396
  }
2225
3397
  // logprobs capture — same mlx-lm block as chat (generate_response is
@@ -2245,6 +3417,11 @@ export function createServer(
2245
3417
  wantsLogprobs: captureLogprobs,
2246
3418
  userSeed: body.seed !== undefined,
2247
3419
  kvQuant: !!(options.kvConfig?.length || options.kvBits),
3420
+ turboQuant: !!options.turboQuant,
3421
+ // /v1/completions grammar (textGrammarCtrl). Same null-on-degrade
3422
+ // contract as the chat lane.
3423
+ hasGrammar: !!textGrammarCtrl,
3424
+ hasDraft: !!ctx.draft,
2248
3425
  };
2249
3426
  const finishReason = (stopped: boolean, generated: number): "stop" | "length" =>
2250
3427
  stopped ? "stop" : generated >= (options.maxTokens ?? 512) ? "length" : "stop";
@@ -2292,6 +3469,7 @@ export function createServer(
2292
3469
  completion_tokens: s.generatedTokens,
2293
3470
  total_tokens: s.promptTokens + s.generatedTokens,
2294
3471
  prompt_tokens_details: { cached_tokens: s.cachedTokens },
3472
+ ...(s.spec ? { speculation: s.spec } : {}),
2295
3473
  },
2296
3474
  });
2297
3475
  controller.enqueue(enc.encode("data: [DONE]\n\n"));
@@ -2307,6 +3485,7 @@ export function createServer(
2307
3485
  "content-type": "text/event-stream",
2308
3486
  "cache-control": "no-cache",
2309
3487
  connection: "keep-alive",
3488
+ ...(textGrammarWarning ? { Warning: textGrammarWarning } : {}),
2310
3489
  },
2311
3490
  });
2312
3491
  }
@@ -2341,8 +3520,9 @@ export function createServer(
2341
3520
  completion_tokens: s.generatedTokens,
2342
3521
  total_tokens: s.promptTokens + s.generatedTokens,
2343
3522
  prompt_tokens_details: { cached_tokens: s.cachedTokens },
3523
+ ...(s.spec ? { speculation: s.spec } : {}),
2344
3524
  },
2345
- });
3525
+ }, textGrammarWarning ? { headers: { Warning: textGrammarWarning } } : undefined);
2346
3526
  } catch (e) {
2347
3527
  return Response.json({ error: { message: (e as Error).message } }, { status: 500 });
2348
3528
  }
@@ -2708,6 +3888,27 @@ export function createServer(
2708
3888
  saveHfToken(body.token);
2709
3889
  return Response.json({ ok: true });
2710
3890
  }
3891
+
3892
+ // --- Durable tool-approvals settings (plan §5.4/§6.5/§9 Phase 2) ---
3893
+ // The web chat's own "always allow this tool" list
3894
+ // (~/.mlx-bun/tool-approvals.json, src/tool-approvals.ts). The
3895
+ // approval card itself grants an always-allow via the `approval` WS
3896
+ // frame (alwaysAllow:true) — these REST routes are read/forget only,
3897
+ // for the settings panel's list view (matches the hf-token GET/POST
3898
+ // pattern's separation of "the card that grants" vs "settings that
3899
+ // manage" — there is no POST here on purpose, granting always goes
3900
+ // through the approval card, never a bare settings form).
3901
+ if (url.pathname === "/api/settings/tool-approvals" && request.method === "GET") {
3902
+ const { listAlwaysAllowedTools } = await import("./tool-approvals");
3903
+ return Response.json({ ok: true, alwaysAllow: listAlwaysAllowedTools() });
3904
+ }
3905
+ if (url.pathname === "/api/settings/tool-approvals" && request.method === "DELETE") {
3906
+ const body = (await request.json().catch(() => ({}))) as { tool?: string };
3907
+ if (!body.tool) return Response.json({ ok: false, error: "tool required" }, { status: 400 });
3908
+ const { revokeToolAlwaysAllowed } = await import("./tool-approvals");
3909
+ const file = revokeToolAlwaysAllowed(body.tool);
3910
+ return Response.json({ ok: true, alwaysAllow: Object.keys(file.allows).sort() });
3911
+ }
2711
3912
  {
2712
3913
  const m = url.pathname.match(/^\/api\/(quantize|finetune|dataset)\/push$/);
2713
3914
  if (m && request.method === "POST") {
@@ -2736,6 +3937,50 @@ export function createServer(
2736
3937
  }
2737
3938
  }
2738
3939
 
3940
+ // gc plan/execute (web-ui-pass-plan.md #9): thin wrappers over the CLI's
3941
+ // own `mlx-bun gc` planner/executor (src/registry.ts planGc/executeGc).
3942
+ // Loopback-served admin route — no auth beyond the server's own bind
3943
+ // (matches every other /api/* route here). GET plans (read-only, cheap:
3944
+ // config.json + safetensors index reads, no tensor bytes); POST executes
3945
+ // and requires an explicit {yes:true} body (mirrors the CLI's --yes gate)
3946
+ // so a stray call can't delete anything by accident.
3947
+ if (url.pathname === "/api/gc/plan" && request.method === "GET") {
3948
+ const { planGc } = await import("./registry");
3949
+ const plans = planGc().filter(
3950
+ (p) => p.pruneSnapshots.length || p.skippedSnapshots.length || p.deadBlobs.length,
3951
+ );
3952
+ const superseded = plans.map((p) => ({
3953
+ repo_id: p.repoId,
3954
+ prune_snapshots: p.pruneSnapshots.length,
3955
+ skipped_snapshots: p.skippedSnapshots.length,
3956
+ dead_blobs: p.deadBlobs.length,
3957
+ reclaim_bytes: p.reclaimBytes,
3958
+ }));
3959
+ const reclaim_bytes = plans.reduce((a, p) => a + p.reclaimBytes, 0);
3960
+ return Response.json({ ok: true, superseded, reclaim_bytes });
3961
+ }
3962
+ if (url.pathname === "/api/gc/execute" && request.method === "POST") {
3963
+ const body = (await request.json().catch(() => ({}))) as { yes?: boolean };
3964
+ if (body.yes !== true)
3965
+ return Response.json({ ok: false, error: "pass {\"yes\": true} to confirm deletion" }, { status: 400 });
3966
+ const { planGc, executeGc, Registry } = await import("./registry");
3967
+ const plans = planGc().filter(
3968
+ (p) => p.pruneSnapshots.length || p.skippedSnapshots.length || p.deadBlobs.length,
3969
+ );
3970
+ const res = executeGc(plans);
3971
+ const reg = new Registry();
3972
+ try {
3973
+ await reg.scan(); // reap deleted snapshots from the registry
3974
+ } finally {
3975
+ reg.close();
3976
+ }
3977
+ libraryCache = null; // force /library to reflect the deletion
3978
+ return Response.json({
3979
+ ok: true,
3980
+ snapshots: res.snapshots, blobs: res.blobs, reclaimed_bytes: res.reclaimedBytes,
3981
+ });
3982
+ }
3983
+
2739
3984
  if (url.pathname === "/api/jobs" && request.method === "GET") {
2740
3985
  const store = await ensureJobs();
2741
3986
  const limit = Number(url.searchParams.get("limit") ?? "50");