mlx-bun 0.0.10 → 0.0.12

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 (216) hide show
  1. package/README.md +107 -32
  2. package/THIRD_PARTY_LICENSES.md +10 -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/colibri-glm52-port.md +1479 -0
  10. package/docs/design/compat-cli-surface-design.md +4 -3
  11. package/docs/design/decode-speed-program.md +172 -0
  12. package/docs/design/docs-reorg-plan.md +1 -1
  13. package/docs/design/dspark-serving-program.md +371 -0
  14. package/docs/design/dspark-speculative-decoding.md +167 -5
  15. package/docs/design/faithful-l1-consolidation.md +310 -0
  16. package/docs/design/generic-model-support.md +8 -5
  17. package/docs/design/grammar-spec-batching-integration.md +301 -0
  18. package/docs/design/minicpm5-decode-megakernel.md +16 -11
  19. package/docs/design/mlx-lm-tool-parity-plan.md +11 -1
  20. package/docs/design/omlx-adoption-map.md +14 -6
  21. package/docs/design/orpo-training.md +34 -30
  22. package/docs/design/paged-kv-cache.md +214 -0
  23. package/docs/design/parallel-slots.md +82 -49
  24. package/docs/design/parity-tier-dag.md +43 -33
  25. package/docs/design/pre-colibri-stabilization.md +329 -0
  26. package/docs/design/repo-cleanup-plan.md +189 -0
  27. package/docs/design/runtime-isolation.md +19 -3
  28. package/docs/design/segmented-backward-training.md +3 -2
  29. package/docs/design/ssd-kv-cold-tier.md +130 -0
  30. package/docs/design/structured-output.md +484 -0
  31. package/docs/design/the-dreaming-master-plan.md +2 -3
  32. package/docs/design/turboquant-kv.md +235 -0
  33. package/docs/design/unified-engine-frontier-plan.md +774 -0
  34. package/docs/design/web-chat-beat-matrix.md +315 -0
  35. package/docs/design/web-chat-redesign.md +1112 -0
  36. package/docs/investigations/colibri-oracle-pin.md +1284 -0
  37. package/docs/investigations/dspark-handoff.md +182 -16
  38. package/docs/investigations/indexshare-performance-spike.md +357 -0
  39. package/docs/investigations/lab-build-journal.md +2 -2
  40. package/docs/investigations/openwiki-evaluation.md +249 -0
  41. package/docs/planning/IDEAS.md +14 -0
  42. package/docs/planning/PRODUCT_ROADMAP.md +2 -3
  43. package/docs/planning/chat-ui-vision-memory-first.md +199 -0
  44. package/docs/planning/chat-ui-vision.md +50 -0
  45. package/docs/planning/memory-docs-and-dag-plan.md +15 -6
  46. package/docs/planning/release-notes-v0.0.11.md +239 -0
  47. package/docs/planning/release-notes-v0.0.12.md +80 -0
  48. package/docs/reference/cli.md +63 -18
  49. package/docs/reference/distribution.md +37 -6
  50. package/docs/reference/embedding.md +40 -16
  51. package/docs/reference/features-matrix.md +99 -0
  52. package/docs/reference/library-api.md +69 -20
  53. package/docs/reference/memory.md +26 -6
  54. package/docs/reference/models.md +60 -3
  55. package/docs/reference/orpo-quickstart.md +3 -4
  56. package/docs/reference/server-api.md +717 -24
  57. package/docs/reference/server-config.md +265 -129
  58. package/docs/reference/training.md +24 -16
  59. package/package.json +5 -3
  60. package/src/anthropic.ts +67 -44
  61. package/src/audio/conformer.ts +647 -0
  62. package/src/audio/decode.ts +133 -0
  63. package/src/audio/features.ts +202 -0
  64. package/src/audio/transcode.ts +67 -0
  65. package/src/chat-template.ts +155 -10
  66. package/src/cli.ts +483 -113
  67. package/src/config.ts +116 -24
  68. package/src/diffusion/diffusion-generate.ts +14 -8
  69. package/src/download.ts +143 -12
  70. package/src/eval/README.md +7 -6
  71. package/src/eval/ifeval.ts +63 -120
  72. package/src/eval/kl.ts +59 -3
  73. package/src/eval/quality-db.ts +2 -2
  74. package/src/eval/runner.ts +17 -8
  75. package/src/eval/tasks/ifeval.ts +167 -31
  76. package/src/expert-io.ts +391 -0
  77. package/src/expert-offload.ts +15 -2
  78. package/src/expert-residency.ts +1134 -0
  79. package/src/expert-usage.ts +415 -0
  80. package/src/flags.ts +13 -0
  81. package/src/generate.ts +438 -35
  82. package/src/grammar.ts +585 -0
  83. package/src/hf-push.ts +1 -1
  84. package/src/hub-rest.ts +242 -0
  85. package/src/index.ts +14 -2
  86. package/src/jobs/runner.ts +37 -18
  87. package/src/kv-store.ts +586 -93
  88. package/src/lora.ts +85 -14
  89. package/src/md-modules.d.ts +42 -0
  90. package/src/media-fetch.ts +426 -0
  91. package/src/memory/rest.ts +245 -0
  92. package/src/memory/vault.ts +74 -0
  93. package/src/mlx/array.ts +149 -19
  94. package/src/mlx/autograd.ts +53 -16
  95. package/src/mlx/checkpoint.ts +46 -13
  96. package/src/mlx/custom-vjp.ts +83 -35
  97. package/src/mlx/ffi.ts +27 -3
  98. package/src/mlx/metal-kernel.ts +7 -4
  99. package/src/mlx/ops.ts +153 -23
  100. package/src/mlx/turboquant-ops.ts +831 -0
  101. package/src/mlx/turboquant-tables.ts +237 -0
  102. package/src/model/batched-mask.ts +62 -1
  103. package/src/model/batched-quant.ts +206 -0
  104. package/src/model/batched-rotating-quant.ts +346 -0
  105. package/src/model/batched-rotating.ts +46 -1
  106. package/src/model/compiled-decode.ts +4 -15
  107. package/src/model/factory.ts +103 -6
  108. package/src/model/gemma4-base.ts +395 -39
  109. package/src/model/gemma4.ts +122 -33
  110. package/src/model/generated/gemma4-12b.ts +9 -13
  111. package/src/model/generated/gemma4-26b.ts +10 -5
  112. package/src/model/generated/gemma4-e4b.ts +14 -9
  113. package/src/model/glm52-atlas.ts +511 -0
  114. package/src/model/glm52-cache.ts +706 -0
  115. package/src/model/glm52-config.ts +262 -0
  116. package/src/model/glm52-container.ts +305 -0
  117. package/src/model/glm52-coupling.ts +388 -0
  118. package/src/model/glm52-dsa.ts +602 -0
  119. package/src/model/glm52-expert-layout.ts +253 -0
  120. package/src/model/glm52-layout.ts +152 -0
  121. package/src/model/glm52-memory.ts +473 -0
  122. package/src/model/glm52-mla.ts +786 -0
  123. package/src/model/glm52-moe.ts +223 -0
  124. package/src/model/glm52-pilot.ts +537 -0
  125. package/src/model/glm52-reference.ts +445 -0
  126. package/src/model/glm52-residency.ts +557 -0
  127. package/src/model/glm52-resident-weights.ts +135 -0
  128. package/src/model/glm52-streamed-experts.ts +453 -0
  129. package/src/model/glm52-streamed-metal.ts +610 -0
  130. package/src/model/glm52-weights.ts +503 -0
  131. package/src/model/glm52.ts +1153 -0
  132. package/src/model/minicpm5.ts +58 -102
  133. package/src/model/paged-kv.ts +253 -0
  134. package/src/model/qwen3-delta.ts +76 -15
  135. package/src/model/qwen3-moe.ts +404 -0
  136. package/src/model/qwen3.ts +9 -8
  137. package/src/model/qwen3_5.ts +88 -23
  138. package/src/model/support.ts +19 -1
  139. package/src/model/universal/dense.ts +40 -5
  140. package/src/model/universal/rope.ts +36 -21
  141. package/src/native/expert_io.c +523 -0
  142. package/src/native-pack.ts +6 -4
  143. package/src/pi-web.ts +1299 -82
  144. package/src/preflight.ts +27 -8
  145. package/src/prompt-cache.ts +293 -33
  146. package/src/quantize/quantizer.ts +16 -1
  147. package/src/registry.ts +57 -2
  148. package/src/responses.ts +62 -49
  149. package/src/serve/batch-scheduler.ts +1050 -78
  150. package/src/serve/generation-gateway.ts +252 -24
  151. package/src/serve/isolate.ts +417 -0
  152. package/src/serve/lane-registry.ts +44 -0
  153. package/src/serve/session-search.ts +326 -0
  154. package/src/server.ts +1701 -175
  155. package/src/spec/assistant-source.ts +150 -0
  156. package/src/spec/deepspec-source.ts +143 -0
  157. package/src/spec/dflash-source.ts +124 -0
  158. package/src/spec/dspark/ab-stats.ts +136 -0
  159. package/src/spec/dspark/calibration.ts +112 -0
  160. package/src/spec/dspark/deepspec-module.ts +801 -0
  161. package/src/spec/dspark/generate-dflash.ts +72 -20
  162. package/src/spec/dspark/loader.ts +23 -0
  163. package/src/spec/dspark/module-dflash.ts +278 -36
  164. package/src/spec/dspark/quantize-drafter.ts +91 -0
  165. package/src/spec/generate.ts +59 -25
  166. package/src/spec/glm52-mtp-source.ts +392 -0
  167. package/src/spec/ngram-source.ts +132 -0
  168. package/src/spec/serve-loop.ts +617 -0
  169. package/src/spec/source.ts +118 -0
  170. package/src/spec/two-model.ts +171 -0
  171. package/src/ssd-cache.ts +111 -35
  172. package/src/tokenizer.ts +288 -2
  173. package/src/tool-approvals.ts +131 -0
  174. package/src/tool-call.ts +438 -22
  175. package/src/train/flash-cce.ts +2 -2
  176. package/src/train/job.ts +0 -12
  177. package/src/train/lora-params.ts +17 -6
  178. package/src/train/trainer.ts +60 -30
  179. package/src/vision/prompt.ts +217 -48
  180. package/src/web/app.html +1197 -1791
  181. package/src/web/app.js +5694 -0
  182. package/src/web/icon.svg +13 -0
  183. package/src/web/manifest.webmanifest +24 -0
  184. package/src/web/src/adapters-panel.ts +288 -0
  185. package/src/web/src/api.ts +57 -0
  186. package/src/web/src/assistant.ts +485 -0
  187. package/src/web/src/chat.ts +1050 -0
  188. package/src/web/src/composer.ts +959 -0
  189. package/src/web/src/dataset.ts +132 -0
  190. package/src/web/src/finetune.ts +267 -0
  191. package/src/web/src/hub.ts +428 -0
  192. package/src/web/src/main.ts +98 -0
  193. package/src/web/src/markdown.ts +0 -0
  194. package/src/web/src/memory-panel.ts +572 -0
  195. package/src/web/src/model-picker.ts +159 -0
  196. package/src/web/src/palette.ts +346 -0
  197. package/src/web/src/protocol.ts +62 -0
  198. package/src/web/src/quantize.ts +167 -0
  199. package/src/web/src/rag.ts +274 -0
  200. package/src/web/src/sessions.ts +469 -0
  201. package/src/web/src/shell.ts +804 -0
  202. package/src/web/src/status.ts +300 -0
  203. package/src/web/src/ui-catalog.ts +92 -0
  204. package/src/web/sw.js +71 -0
  205. package/src/web/vendor/README.md +61 -0
  206. package/src/web/vendor/hljs-LICENSE +29 -0
  207. package/src/web/vendor/hljs-theme.css +59 -0
  208. package/src/web/vendor/hljs.js +26 -0
  209. package/src/web-tools.ts +54 -2
  210. package/src/weights.ts +82 -19
  211. package/src/model/fused-decode-kernel.ts +0 -312
  212. package/src/model/fused-geglu-kernel.ts +0 -198
  213. package/src/model/fused-mlp-kernel.ts +0 -169
  214. package/src/model/fused-swiglu-kernel.ts +0 -240
  215. package/src/model/steel-linear-kernel.ts +0 -96
  216. package/src/status-page.html +0 -497
package/src/server.ts CHANGED
@@ -16,22 +16,76 @@ 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
- import { createModel, type RuntimeModel } from "./model/factory";
29
- import { isMiniCPM5Config, isSupportedModelRecord } from "./model/support";
30
- import { generate, type GenerateOptions, type TokenLogprobs } from "./generate";
56
+ import {
57
+ createModel,
58
+ openGlm52RuntimeModel,
59
+ type Glm52RuntimeOpenOptions,
60
+ type RuntimeModel,
61
+ } from "./model/factory";
62
+ import { Glm52Model } from "./model/glm52";
63
+ import {
64
+ GLM52_G5_ASPIRATIONAL_DECODE_TPS,
65
+ GLM52_G5_DIRECT_ORACLE_WARM_DECODE_TPS,
66
+ GLM52_G5_MEASURED_AT,
67
+ GLM52_G5_MEASURED_WARM_DECODE_TPS,
68
+ type Glm52MemoryPlan,
69
+ } from "./model/glm52-memory";
70
+ import { isGlm52Config, isMiniCPM5Config } from "./model/support";
71
+ import { Glm52NativeMtpProvider } from "./spec/glm52-mtp-source";
72
+ import {
73
+ generate,
74
+ type GenerateOptions,
75
+ type TokenLogprobs,
76
+ withModelUsageFlush,
77
+ withModelWiredLimit,
78
+ } from "./generate";
79
+ import { cloneKvCaches, SpillQueue } from "./kv-store";
80
+ import { TURBOQUANT_HEAD_DIMS } from "./mlx/turboquant-tables";
81
+ import {
82
+ compileGrammarRequest, grammarEnabled, type GrammarRequest,
83
+ } from "./grammar";
31
84
  import type { HlgConfig } from "./sampler";
32
85
  import { isMonotone, CURVE_UMIN, type CurveParams } from "./curve-sampler";
33
86
  const CURVE_PAGE = curveDesignerHtml as unknown as string;
34
87
  import { GenerationGateway } from "./serve/generation-gateway";
88
+ import { recordLane, type Lane } from "./serve/lane-registry";
35
89
  import {
36
90
  ChatTemplate, type ChatMessage, type ToolDefinition,
37
91
  } from "./chat-template";
@@ -40,7 +94,7 @@ import {
40
94
  CHANNEL_END, CHANNEL_START, parseGeneratedToolCalls, parseToolCalls,
41
95
  TOOL_CALL_END, TOOL_CALL_START,
42
96
  } from "./tool-call";
43
- import { PromptCache } from "./prompt-cache";
97
+ import { PromptCache, cacheBytes } from "./prompt-cache";
44
98
  import { SsdCacheStore } from "./ssd-cache";
45
99
  import { configFingerprint } from "./model/fingerprint";
46
100
  import {
@@ -59,14 +113,23 @@ import { setMemoryLimit } from "./mlx/ffi";
59
113
  import { VisionTower } from "./vision/embedder";
60
114
  import { SiglipVisionTower, parseSiglipConfig } from "./vision/siglip";
61
115
  import {
62
- buildVisionPrompt, extractImages,
63
- type VisionTokenIds, type VisionEncoder,
116
+ buildMultimodalPrompt, buildVisionPrompt, extractAudio, extractImages,
117
+ type AudioTokenIds, type VisionTokenIds, type VisionEncoder,
64
118
  } from "./vision/prompt";
119
+ import { AudioTower, parseAudioConfig } from "./audio/conformer";
120
+ import { ensureWav } from "./audio/transcode";
121
+ import { sidecarShipsAudioTower } from "./registry";
65
122
  import { makePiWsHandler, type PiWsData } from "./pi-web";
66
123
 
67
124
  export interface ServerOptions {
68
- /** Byte cap for the prompt (KV) cache. Default 2 GB. */
125
+ /** Byte cap for the prompt (KV) cache. Default 8 GB. */
69
126
  promptCacheBytes?: number;
127
+ /** Aggregate KV-byte budget across concurrently-admitted batch rows
128
+ * (`--kv-budget`, batching-perf-path P3). Joiners whose projected KV
129
+ * (prompt + max_tokens, window-capped) would exceed it QUEUE until rows
130
+ * evict; a request over the budget alone is rejected. Unset = no
131
+ * aggregate cap (per-request admission via memoryBudget still applies). */
132
+ kvBudgetBytes?: number;
70
133
  /** KV quantization override. When unset: apply ctx.kvConfig (mixed
71
134
  * per-layer) for serial serving, but bf16 under `--batch N` (the batched
72
135
  * engine is bf16-only — a mode switch, see `batch` below). "config" forces
@@ -74,6 +137,19 @@ export interface ServerOptions {
74
137
  * the serial path); "off" forces bf16; a number forces uniform bits
75
138
  * (group size 64, start 0) ignoring the config file. */
76
139
  kvQuant?: "off" | "config" | number;
140
+ /** TurboQuant scheme (docs/design/turboquant-kv.md): a separate axis from
141
+ * kvQuant above, mutually exclusive with it (`--kv-quant turbo[:k<bits>v
142
+ * <bits>]` sets this instead of kvQuant). Solo-only in v1 — see
143
+ * GenerationGateway's explicit refusal. */
144
+ turboQuant?: TurboQuantScheme;
145
+ /** OPTIONAL paged KV cache (`--paged-kv`, docs/design/paged-kv-cache.md):
146
+ * vLLM-style block-pool storage for full-attention layers,
147
+ * gather-to-contiguous before the stock SDPA. Default off (unset = the
148
+ * plain KVCache path, byte-identical). v1 scope: serial batch=1,
149
+ * Gemma4-family, bf16 — startup refuses `--batch N>1`, `--kv-quant`,
150
+ * turbo, and `--draft-model` combinations; paged requests bypass the
151
+ * prompt cache and run uncompiled. Gated bit-exact vs the plain path. */
152
+ pagedKv?: { blockSize?: number };
77
153
  /** Memory budget for the serving process (admission control — Phase 5).
78
154
  * Requests whose prompt + max_tokens exceed the budget's max safe
79
155
  * context are rejected with 400 instead of crashing the GPU: the OOM
@@ -88,6 +164,11 @@ export interface ServerOptions {
88
164
  * it). Exposed at /stats.server so other mlx-bun processes can
89
165
  * warn before attaching to a server that may vanish. */
90
166
  owner?: "serve" | "pi-session" | "embedded";
167
+ /** Listen on a UNIX DOMAIN SOCKET instead of a TCP port — the engine-
168
+ * child mode of the isolation architecture (runtime-isolation.md): the
169
+ * parent process reverse-proxies HTTP to this socket. Stale socket
170
+ * files are unlinked before bind. When set, `hostname`/port are ignored. */
171
+ unixSocket?: string;
91
172
  /** Interface to bind (Bun.serve hostname). Unset ⇒ Bun's default
92
173
  * (all interfaces) — embedded/library use. The CLI always passes one:
93
174
  * "127.0.0.1" (loopback, mlx_lm.server parity) unless --host says
@@ -111,9 +192,9 @@ export interface ServerOptions {
111
192
  * thinking traces never truncate. `--max-tokens 512` = mlx-lm behavior. */
112
193
  defaultMaxTokens?: number;
113
194
  /** 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
195
+ * (`--batch N`). Default 8 (continuous batching, B floats 1..N,
196
+ * bit-parity with mlx-lm B=N); `--batch 1` pins the serialized
197
+ * single-queue path — a mode switch, not a load-dependent
117
198
  * fallback. See docs/design/parallel-slots.md. NOTE: the batched executor
118
199
  * is mid-build; until it lands, >1 warns and runs serially. */
119
200
  batch?: number;
@@ -132,6 +213,11 @@ export interface ServerOptions {
132
213
  ssdCacheDir?: string;
133
214
  /** Byte cap for the SSD tier (default 32 GiB). */
134
215
  ssdCacheMaxBytes?: number;
216
+ /** Idle-demotion threshold in seconds (`--ssd-demote-idle`): prompt-cache
217
+ * entries unused this long spill to the SSD tier and free their GPU
218
+ * memory (Layer 0 — RAM drains between bursts, prefixes stay reachable
219
+ * via zero-copy restore). Default 300 when the tier is on; 0 disables. */
220
+ ssdDemoteIdleSec?: number;
135
221
  /** Verify every tensor hash on restore (`--ssd-cache-verify`) — reads all
136
222
  * bytes eagerly, defeating lazy fault-in; integrity paranoia only. */
137
223
  ssdCacheVerify?: boolean;
@@ -151,6 +237,17 @@ export interface ServerContext {
151
237
  * `vision`. */
152
238
  loadVision: (() => VisionEncoder) | null;
153
239
  visionTokenIds: VisionTokenIds;
240
+ /** Lazily-loaded Conformer audio tower — null until the first audio
241
+ * request (see `getAudioTower`). Same sidecar file as vision
242
+ * (optiq_vision.safetensors), separate tower; text-only sessions never
243
+ * pay for it. */
244
+ audio: AudioTower | null;
245
+ /** Loads the audio tower on demand; null when the model can't do audio
246
+ * (no `audio_config` in config.json or no sidecar). No flags — audio
247
+ * auto-enables exactly like vision. */
248
+ loadAudio: (() => AudioTower) | null;
249
+ /** null when the model has no `audio_config` (audio-incapable). */
250
+ audioTokenIds: AudioTokenIds | null;
154
251
  adapters: AdapterManager;
155
252
  /** Per-layer KV quantization from the repo's kv_config.json (null if
156
253
  * absent). Applied by default — optiq serve's headline behavior;
@@ -160,6 +257,33 @@ export interface ServerContext {
160
257
  * optiq serve injects these as server defaults (gen_config.py);
161
258
  * explicit request fields always win. */
162
259
  genDefaults: GenSamplingDefaults;
260
+ /** Speculative decoding (`serve --draft-model`, mlx_lm.server parity).
261
+ * Server-level: when set, EVERY request routes to the serial lane
262
+ * (upstream: is_batchable = draft is None) and spec-eligible ones decode
263
+ * through src/spec/serve-loop.ts. null = no draft configured. */
264
+ draft?: {
265
+ provider: import("./spec/source").DraftProvider;
266
+ numDraftTokens: number;
267
+ } | null;
268
+ /** Present only for the direct Colibri GLM runtime. This is the exact
269
+ * header-derived process equation used before opening resident state. */
270
+ glmMemoryPlan?: Glm52MemoryPlan | null;
271
+ }
272
+
273
+ export interface LoadContextOptions {
274
+ memoryBudgetBytes?: number;
275
+ /** Direct Colibri runtime/resource overrides. Ignored by other models. */
276
+ glm?: Glm52RuntimeOpenOptions;
277
+ /** Snapshot dir of a draft model for speculative decoding
278
+ * (`--draft-model`). Loaded alongside the target; the pair must share
279
+ * a tokenizer family. */
280
+ draftModelDir?: string;
281
+ /** Drafts per round (`--num-draft-tokens`, mlx_lm.server default 3). */
282
+ numDraftTokens?: number;
283
+ /** Draft-provider kind override (`--draft-kind`). */
284
+ draftKind?: DraftKind;
285
+ ngramMax?: number;
286
+ ngramMin?: number;
163
287
  }
164
288
 
165
289
  export interface GenSamplingDefaults {
@@ -169,6 +293,30 @@ export interface GenSamplingDefaults {
169
293
  repetitionPenalty?: number;
170
294
  }
171
295
 
296
+ export type DraftKind = "dspark" | "deepspec" | "assistant" | "two-model" | "ngram";
297
+
298
+ /** Detect the draft artifact's kind so the right provider is loaded. All
299
+ * providers share ONE serve loop (src/spec/serve-loop.ts). Exported for the
300
+ * bench harness (scripts/bench-feature-matrix.ts) — one detection, no drift.
301
+ * "ngram" is never detected — it has no artifact (model-free prompt lookup,
302
+ * src/spec/ngram-source.ts) and mounts via an explicit `--draft-kind ngram`. */
303
+ export async function detectDraftKind(dir: string): Promise<DraftKind> {
304
+ if (await Bun.file(`${dir}/dspark.json`).exists()) return "dspark"; // our trained module
305
+ try {
306
+ const cfg = (await Bun.file(`${dir}/config.json`).json()) as {
307
+ model_type?: string;
308
+ architectures?: string[];
309
+ };
310
+ // DeepSeek's released DSpark drafters (DeepSpec reference): no
311
+ // dspark.json, plain HF config stamped Gemma4DSparkModel.
312
+ if (cfg.architectures?.[0] === "Gemma4DSparkModel") return "deepspec";
313
+ if (String(cfg.model_type ?? "").includes("assistant")) return "assistant";
314
+ } catch {
315
+ // no/unreadable config → fall through to a full second model
316
+ }
317
+ return "two-model";
318
+ }
319
+
172
320
  async function loadGenSamplingDefaults(modelDir: string): Promise<GenSamplingDefaults> {
173
321
  const file = Bun.file(`${modelDir}/generation_config.json`);
174
322
  if (!(await file.exists())) return {};
@@ -188,16 +336,38 @@ async function loadGenSamplingDefaults(modelDir: string): Promise<GenSamplingDef
188
336
 
189
337
  export async function loadContext(
190
338
  modelDir: string, modelId?: string,
191
- opts: { memoryBudgetBytes?: number } = {},
339
+ opts: LoadContextOptions = {},
192
340
  ): Promise<ServerContext> {
193
341
  const config = await loadModelConfig(modelDir);
194
- const weights = await Weights.open(modelDir);
342
+ const glm = isGlm52Config(config);
343
+ const externalDraft = opts.draftModelDir !== undefined || opts.draftKind !== undefined;
344
+ if (glm && externalDraft && opts.glm?.enableMtp === true)
345
+ throw new Error("native GLM MTP and --draft-model/--draft-kind are mutually exclusive");
346
+ let weights: Weights | null = null;
347
+ let model!: RuntimeModel;
348
+ let glmMemoryPlan: Glm52MemoryPlan | null = null;
349
+ if (glm) {
350
+ const enableNativeMtp = !externalDraft && opts.glm?.enableMtp !== false;
351
+ const opened = await openGlm52RuntimeModel(modelDir, {
352
+ ...opts.glm,
353
+ memoryBudgetBytes: opts.memoryBudgetBytes ?? opts.glm?.memoryBudgetBytes,
354
+ // Every drafter owns the serial speculative lane. Do not reserve a
355
+ // simultaneous ordinary batch that cannot be admitted while it is live.
356
+ batchSize: enableNativeMtp || externalDraft ? 1 : opts.glm?.batchSize,
357
+ // An explicit alternate drafter owns the one speculative lane.
358
+ enableMtp: enableNativeMtp,
359
+ });
360
+ model = opened.model;
361
+ glmMemoryPlan = opened.plan;
362
+ } else {
363
+ weights = await Weights.open(modelDir);
364
+ }
195
365
  // memoryBudget enforcement at load (Phase 5): Weights.open only mmaps
196
366
  // (no GPU allocation yet), so a model whose weights can never serve
197
367
  // within the budget is refused HERE — before any unified-memory
198
368
  // commitment — with an actionable error instead of a Metal OOM later.
199
- if (opts.memoryBudgetBytes) {
200
- const weightsBytes = [...weights.shards.files.values()]
369
+ if (!glm && opts.memoryBudgetBytes) {
370
+ const weightsBytes = [...weights!.shards.files.values()]
201
371
  .reduce((a, f) => a + f.mmap.size, 0);
202
372
  const report = fit(config, weightsBytes, 1, undefined, undefined, 0, opts.memoryBudgetBytes);
203
373
  if (report.maxSafeContext < 1)
@@ -209,7 +379,7 @@ export async function loadContext(
209
379
  }
210
380
  // generated-specialization dispatch by config fingerprint (Phase C);
211
381
  // unmatched configs run the monolith — slow, never broken
212
- const model = createModel(weights, config);
382
+ if (!glm) model = createModel(weights!, config);
213
383
  const tokenizer = await loadTokenizer(modelDir);
214
384
  // Generation must stop on the tokenizer's eos_token — the chat turn
215
385
  // terminator (e.g. Qwen <|im_end|> = 248046). Some configs (Qwen3.5-4B)
@@ -219,8 +389,130 @@ export async function loadContext(
219
389
  // the tokenizer eos; union it in. No-op when already present (Gemma, 27B).
220
390
  if (tokenizer.eosTokenId != null && !config.eosTokenIds.includes(tokenizer.eosTokenId))
221
391
  config.eosTokenIds = [...config.eosTokenIds, tokenizer.eosTokenId];
392
+
393
+ // Speculative decoding: load the draft (mlx_lm.server --draft-model). The
394
+ // draft artifact's KIND selects the provider — all three share ONE serve
395
+ // loop (src/spec/serve-loop.ts): dspark.json → DSpark (KV-injected), a
396
+ // *_assistant config → the optiq KV-borrowing Gemma drafter, otherwise a
397
+ // full second model (mlx-lm parity). `--draft-kind` overrides the detect.
398
+ let draft: ServerContext["draft"] = null;
399
+ if (opts.draftKind === "ngram") {
400
+ // Model-free prompt lookup: no artifact, no dir, no probe/budget concerns
401
+ // (weightsBytes 0, open() never throws). Default γ=10 per the reference
402
+ // implementation — drafting is free, so wide blocks cost only verify-window
403
+ // width when wrong.
404
+ if (opts.draftModelDir)
405
+ throw new Error(
406
+ "--draft-kind ngram is model-free — drop --draft-model (it would be ignored)",
407
+ );
408
+ const { NgramProvider } = await import("./spec/ngram-source");
409
+ draft = {
410
+ provider: new NgramProvider({ max: opts.ngramMax, min: opts.ngramMin }),
411
+ numDraftTokens: Math.max(1, opts.numDraftTokens ?? 10),
412
+ };
413
+ } else if (opts.draftModelDir) {
414
+ const dir = opts.draftModelDir;
415
+ const kind = opts.draftKind ?? (await detectDraftKind(dir));
416
+ let provider: import("./spec/source").DraftProvider;
417
+ let numDraftTokens = Math.max(1, opts.numDraftTokens ?? 3);
418
+ if (kind === "dspark") {
419
+ const { DflashProvider } = await import("./spec/dflash-source");
420
+ const p = await DflashProvider.load(dir);
421
+ provider = p;
422
+ // Pin to the trained block width — the serve loop must never ask for
423
+ // more positions than the DSpark block was trained for (n ≤ cfg.gamma).
424
+ numDraftTokens = Math.max(1, Math.min(opts.numDraftTokens ?? p.gamma, p.gamma));
425
+ } else if (kind === "deepspec") {
426
+ const { DeepspecProvider } = await import("./spec/deepspec-source");
427
+ const p = await DeepspecProvider.load(dir);
428
+ provider = p;
429
+ // Same pin, from their config's block_size (e.g. 7 for the released
430
+ // dspark_gemma4_12b_block7).
431
+ numDraftTokens = Math.max(1, Math.min(opts.numDraftTokens ?? p.gamma, p.gamma));
432
+ } else if (kind === "assistant") {
433
+ const { AssistantProvider } = await import("./spec/assistant-source");
434
+ provider = await AssistantProvider.load(dir);
435
+ } else {
436
+ const { TwoModelProvider } = await import("./spec/two-model");
437
+ provider = await TwoModelProvider.load(dir, config.text.vocabSize);
438
+ // Tokenizer-family hard check — two-model ONLY (it ships its own
439
+ // tokenizer). Exact-token-match acceptance is meaningful only when both
440
+ // models tokenize identically; a probe that ENCODES differently means
441
+ // different families — refuse instead of silently accepting ~0% of
442
+ // drafts. The KV-borrowing drafters share the target's tokenization by
443
+ // construction, so this check does not apply to them.
444
+ const draftTok = await loadTokenizer(dir);
445
+ const probe = "The 3 quick brown foxes jumped över the lazy dog?! 🦊";
446
+ if (JSON.stringify(tokenizer.encode(probe)) !== JSON.stringify(draftTok.encode(probe))) {
447
+ provider.dispose();
448
+ throw new Error(
449
+ `--draft-model tokenizer differs from the target's (probe string encodes ` +
450
+ `differently) — speculation needs the same tokenizer family`,
451
+ );
452
+ }
453
+ }
454
+ // Fail-fast pairing validation (2026-07-07 review): the KV-borrowing
455
+ // sources validate the (target, drafter) pairing in open() — non-Gemma4
456
+ // target, DeepSpec target-layer-count mismatch — which used to surface
457
+ // as a 500 from inside specServeRun on EVERY text request. Probe-open
458
+ // once with throwaway caches here so a mismatch refuses at load; open()
459
+ // allocates no per-request tensors before prefill/draft, so this is
460
+ // free. The probe sampler is never called during open().
461
+ {
462
+ const probeCaches = model.makeCache();
463
+ try {
464
+ provider
465
+ .open({
466
+ sampler: () => { throw new Error("probe sampler never samples"); },
467
+ target: { model, caches: probeCaches },
468
+ })
469
+ .dispose();
470
+ } catch (err) {
471
+ provider.dispose();
472
+ throw new Error(
473
+ `--draft-model is incompatible with this target: ${err instanceof Error ? err.message : String(err)}`,
474
+ );
475
+ } finally {
476
+ for (const c of probeCaches) c.dispose();
477
+ }
478
+ }
479
+ if (opts.memoryBudgetBytes) {
480
+ const targetBytes = weights
481
+ ? [...weights.shards.files.values()].reduce((a, f) => a + f.mmap.size, 0)
482
+ : model.weightsBytes;
483
+ // Draft weights shrink the target's envelope. Draft KV is not modeled
484
+ // (small relative to its weights at serve contexts); admission stays
485
+ // approximately conservative via the combined-weights term.
486
+ const report = fit(config, targetBytes + provider.weightsBytes, 1, undefined, undefined, 0, opts.memoryBudgetBytes);
487
+ if (report.maxSafeContext < 1) {
488
+ provider.dispose();
489
+ throw new Error(
490
+ `target + draft do not fit the memory budget (draft adds ` +
491
+ `${(provider.weightsBytes / 1e9).toFixed(2)} GB)`,
492
+ );
493
+ }
494
+ }
495
+ draft = { provider, numDraftTokens };
496
+ } else if (opts.draftKind) {
497
+ // Every other kind names an artifact to load — refuse instead of silently
498
+ // serving without speculation.
499
+ throw new Error(`--draft-kind ${opts.draftKind} requires --draft-model`);
500
+ }
501
+
502
+ // GLM's checkpoint-native MTP row is the production default. It uses the
503
+ // already-planned bounded auxiliary expert tier and the same tokenizer, so
504
+ // there is no second artifact or compatibility probe to load.
505
+ if (!draft && model instanceof Glm52Model && glmMemoryPlan?.enableMtp) {
506
+ draft = {
507
+ provider: new Glm52NativeMtpProvider(model),
508
+ numDraftTokens: glmMemoryPlan.mtpDraftTokens,
509
+ };
510
+ }
511
+
222
512
  return {
513
+ draft,
223
514
  model,
515
+ glmMemoryPlan,
224
516
  adapters: new AdapterManager(model),
225
517
  kvConfig: config.kvQuant,
226
518
  genDefaults: await loadGenSamplingDefaults(modelDir),
@@ -239,6 +531,17 @@ export async function loadContext(
239
531
  boiTokenId: (config.raw.boi_token_id as number) ?? 255999,
240
532
  eoiTokenId: (config.raw.eoi_token_id as number) ?? 258882,
241
533
  },
534
+ // Audio mirrors vision: lazy tower from the same sidecar, loaded on the
535
+ // first audio request only (docs/design/audio-input-plan.md A4).
536
+ audio: null,
537
+ loadAudio: makeAudioLoader(modelDir, model, config),
538
+ audioTokenIds: config.raw.audio_config
539
+ ? {
540
+ audioTokenId: (config.raw.audio_token_id as number) ?? 258881,
541
+ boaTokenId: (config.raw.boa_token_id as number) ?? 256000,
542
+ eoaTokenId: (config.raw.eoa_token_id as number) ?? 258883,
543
+ }
544
+ : null,
242
545
  };
243
546
  }
244
547
 
@@ -276,6 +579,46 @@ function getVisionTower(ctx: ServerContext): VisionEncoder | null {
276
579
  }
277
580
  }
278
581
 
582
+ /** Build the on-demand audio-tower loader (gemma-4 Conformer, A4 of
583
+ * docs/design/audio-input-plan.md). Auto-enables — no flags — when
584
+ * config.json carries an `audio_config` AND the optiq_vision.safetensors
585
+ * sidecar exists (the audio tensors ship in the same sidecar as vision).
586
+ * Returns null when the model can't do audio (no audio_config: 26B-A4B,
587
+ * DiffusionGemma, bf16 assistants — architectural, not a porting gap). */
588
+ function makeAudioLoader(
589
+ modelDir: string, model: RuntimeModel, config: ModelConfig,
590
+ ): (() => AudioTower) | null {
591
+ if (!(config.hasVisionSidecar && config.raw.audio_config && model instanceof Gemma4Model))
592
+ return null;
593
+ // The sidecar header must actually name the Conformer tensors: the local
594
+ // 12B snapshot pairs audio_config with a STUB sidecar (embed_audio only),
595
+ // and a non-null loader here is what advertises `audio: true` on every
596
+ // capability surface (ws handshake, /v1/models). Header-only read.
597
+ if (!sidecarShipsAudioTower(`${modelDir}/optiq_vision.safetensors`)) return null;
598
+ const audioCfg = parseAudioConfig(config.raw.audio_config as Record<string, any>);
599
+ return () => AudioTower.load(modelDir, audioCfg, model.embedScale);
600
+ }
601
+
602
+ /** Lazily load + cache the audio tower on first use. Unlike vision's
603
+ * warn-and-continue (text-only degrade is correct for requests WITHOUT
604
+ * images), this is only ever consulted for requests WITH audio — the
605
+ * caller turns null into an explicit 400, never a silent text-only
606
+ * degrade. A failed load is not retried every request. (The stub-sidecar
607
+ * case — the local 12B state — never gets here: makeAudioLoader checks the
608
+ * sidecar header and returns a null loader.) */
609
+ function getAudioTower(ctx: ServerContext): AudioTower | null {
610
+ if (ctx.audio) return ctx.audio;
611
+ if (!ctx.loadAudio) return null;
612
+ try {
613
+ ctx.audio = ctx.loadAudio();
614
+ return ctx.audio;
615
+ } catch (e) {
616
+ console.warn(`audio sidecar not loadable (${(e as Error).message})`);
617
+ ctx.loadAudio = null;
618
+ return null;
619
+ }
620
+ }
621
+
279
622
  interface ChatRequest {
280
623
  messages: ChatMessage[];
281
624
  stream?: boolean;
@@ -340,6 +683,20 @@ interface ChatRequest {
340
683
  toe?: number;
341
684
  pivot_offset?: number;
342
685
  };
686
+ /** OpenAI structured output: {type:"json_object"} | {type:"json_schema",
687
+ * json_schema:{name,schema,strict?}} | {type:"text"}. Enforced at the
688
+ * sampler via xgrammar token-bitmasks (src/grammar.ts). L2-class (oMLX
689
+ * oracle). On compile failure, degrades to a system-prompt injection +
690
+ * Warning header (oMLX parity), never 500. */
691
+ response_format?: unknown;
692
+ /** vLLM/oMLX grammar aliases (all compiled via xgrammar): raw EBNF/LARK
693
+ * grammar, regex, enum choice, and a bare JSON-schema object. Precedence
694
+ * (guided_grammar > json_schema > json_object > structured_outputs >
695
+ * guided_regex > guided_choice) mirrors oMLX _effective_guided_grammar. */
696
+ guided_grammar?: string;
697
+ guided_regex?: string;
698
+ guided_choice?: string[];
699
+ structured_outputs?: unknown;
343
700
  }
344
701
 
345
702
  /** POST /v1/completions body (mlx_lm.server's raw text completion — no chat
@@ -419,9 +776,9 @@ export class ToolAwareStream {
419
776
  * chain-of-thought as `<|channel>thought\n…<channel|>` using special tokens
420
777
  * 100/101 that the content decoder strips, so reasoning is captured here at
421
778
  * 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). */
779
+ * incremental state independent of content's. The `thought` channel-name
780
+ * word is stripped before the reasoning text (mlx-lm's think-start marker
781
+ * is `<|channel>thought`; the "\n" after it is reasoning content). */
425
782
  readonly #channelDecoder: StreamDecoder;
426
783
  #inChannel = false;
427
784
  #channelNamePending = "";
@@ -442,15 +799,20 @@ export class ToolAwareStream {
442
799
  this.#channelDecoder = new StreamDecoder(tokenizer, true);
443
800
  }
444
801
 
445
- /** Feed decoded channel text, stripping the leading `thought\n` name line
446
- * (which may arrive across tokens). Returns the reasoning delta. */
802
+ /** Feed decoded channel text, stripping the leading `thought` channel-name
803
+ * word (which may arrive across tokens). Returns the reasoning delta.
804
+ * mlx-lm parity: its think-start marker is exactly `<|channel>thought`
805
+ * (tokenizer_utils.py `_infer_thinking`) and only the MARKER tokens get
806
+ * their text blanked (server.py `_process_control_tokens`) — the "\n"
807
+ * after the name is an ordinary generated token, so it is the FIRST byte
808
+ * of the reasoning stream and must be kept, not swallowed with the name. */
447
809
  #feedChannel(text: string): string {
448
810
  if (this.#channelNameDone) return text;
449
811
  this.#channelNamePending += text;
450
812
  const nl = this.#channelNamePending.indexOf("\n");
451
- if (nl === -1) return ""; // still inside the channel-name line
813
+ if (nl === -1) return ""; // still inside the channel-name word
452
814
  this.#channelNameDone = true;
453
- const rest = this.#channelNamePending.slice(nl + 1);
815
+ const rest = this.#channelNamePending.slice(nl); // keep the "\n" (mlx-lm does)
454
816
  this.#channelNamePending = "";
455
817
  return rest;
456
818
  }
@@ -515,7 +877,8 @@ export class ToolAwareStream {
515
877
  // Reasoning channel: tokens between <|channel> and <channel|> are thought,
516
878
  // captured to #reasoning (drained via takeReasoning), never content. An
517
879
  // empty block (<|channel>thought\n<channel|>, emitted by larger Gemmas even
518
- // with thinking off) yields no reasoning and leaks nothing.
880
+ // with thinking off) yields only the "\n" as reasoning (mlx-lm parity) and
881
+ // leaks nothing into content.
519
882
  if (this.#inChannel) {
520
883
  if (token === CHANNEL_END) {
521
884
  this.#inChannel = false;
@@ -756,10 +1119,15 @@ function contentPartsToText(parts: Array<Record<string, unknown>>): string {
756
1119
  .join("");
757
1120
  }
758
1121
 
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"));
1122
+ /** True if any content part is media — image OR audio — so the multimodal
1123
+ * path must keep the array form for extractImages/extractAudio. */
1124
+ function hasMediaPart(parts: Array<Record<string, unknown>>): boolean {
1125
+ return parts.some(
1126
+ (p) =>
1127
+ p &&
1128
+ (p.type === "image" || p.type === "image_url" ||
1129
+ p.type === "audio" || p.type === "input_audio" || p.type === "audio_url"),
1130
+ );
763
1131
  }
764
1132
 
765
1133
  /** OpenAI sends assistant tool_call arguments as JSON strings; the
@@ -780,13 +1148,14 @@ function hasImagePart(parts: Array<Record<string, unknown>>): boolean {
780
1148
  * OpenAI multimodal client) sends user content as `[{type:"text",text}]`,
781
1149
  * but non-vision chat templates expect `content` to be a STRING and render
782
1150
  * 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. */
1151
+ * replies "I don't see any message." Arrays that carry a media part
1152
+ * (image or audio) are left intact for the multimodal path
1153
+ * (extractImages/extractAudio); only text-only arrays are collapsed
1154
+ * here, which is a no-op for the vision templates too. */
786
1155
  export function normalizeMessages(messages: ChatMessage[]): ChatMessage[] {
787
1156
  return messages.map((raw) => {
788
1157
  let m = raw.role === "developer" ? { ...raw, role: "system" } : raw;
789
- if (Array.isArray(m.content) && !hasImagePart(m.content)) {
1158
+ if (Array.isArray(m.content) && !hasMediaPart(m.content)) {
790
1159
  m = { ...m, content: contentPartsToText(m.content) };
791
1160
  }
792
1161
  if (!m.tool_calls) return m;
@@ -850,6 +1219,62 @@ export function nextDefaultSeed(): number {
850
1219
  * with defaults logprobs=False / top_logprobs=-1). Returns the reference's
851
1220
  * error message (→ 400) or null when valid. Note mlx-lm caps top_logprobs
852
1221
  * at 11, not OpenAI's 20 — we mirror the reference. */
1222
+ /** Build a degrade-path system-prompt instruction for JSON output, mirroring
1223
+ * oMLX's api.tool_calling.build_json_system_prompt (used when xgrammar
1224
+ * compile fails — the response_format degrades to prompt injection rather
1225
+ * than a 500, oMLX parity). Returns null for {type:"text"} / unset. */
1226
+ export function degradeJsonSystemPrompt(body: ChatRequest): string | null {
1227
+ if (body.guided_grammar) {
1228
+ return "You must respond with text matching this grammar:\n\n" + body.guided_grammar;
1229
+ }
1230
+ const rf = body.response_format as { type?: string; json_schema?: { name?: string; description?: string; schema?: unknown } } | undefined;
1231
+ if (rf && typeof rf === "object" && rf.type === "json_object") {
1232
+ return "You must respond with valid JSON only. " +
1233
+ "Do not include any explanation or text outside the JSON object.";
1234
+ }
1235
+ const schemaSpec = rf && typeof rf === "object" && rf.type === "json_schema"
1236
+ ? rf.json_schema
1237
+ : body.structured_outputs && typeof body.structured_outputs === "object"
1238
+ ? { schema: body.structured_outputs }
1239
+ : null;
1240
+ if (schemaSpec) {
1241
+ const spec = schemaSpec;
1242
+ const schema = spec.schema ?? body.structured_outputs ?? {};
1243
+ const name = spec.name ?? "response";
1244
+ const description = spec.description ?? "";
1245
+ let prompt = `You must respond with valid JSON matching the '${name}' schema.`;
1246
+ if (description) prompt += ` ${description}`;
1247
+ prompt += `\n\nJSON Schema:\n\`\`\`json\n${JSON.stringify(schema, null, 2)}\n\`\`\`\n\n` +
1248
+ "Respond with only the JSON object, no additional text or explanation.";
1249
+ return prompt;
1250
+ }
1251
+ if (body.guided_regex) {
1252
+ return "You must respond with text matching this regular expression:\n\n" +
1253
+ body.guided_regex;
1254
+ }
1255
+ if (body.guided_choice?.length) {
1256
+ return "You must respond with exactly one of these choices:\n\n" +
1257
+ body.guided_choice.map((choice) => JSON.stringify(choice)).join("\n");
1258
+ }
1259
+ return null;
1260
+ }
1261
+
1262
+ export function applyGrammarDegrade(
1263
+ body: ChatRequest,
1264
+ degradeHint: string,
1265
+ ): { body: ChatRequest; warning: string } {
1266
+ const content = degradeJsonSystemPrompt(body) ??
1267
+ "Follow the requested output constraint exactly.";
1268
+ return {
1269
+ warning:
1270
+ `grammar not enforced: ${degradeHint} - falling back to prompt injection`,
1271
+ body: {
1272
+ ...body,
1273
+ messages: [{ role: "system", content }, ...body.messages],
1274
+ },
1275
+ };
1276
+ }
1277
+
853
1278
  export function validateLogprobsParams(body: {
854
1279
  logprobs?: unknown;
855
1280
  top_logprobs?: unknown;
@@ -929,24 +1354,117 @@ class LogprobsCollector {
929
1354
  }
930
1355
  }
931
1356
 
932
- /** Incremental detokenizer: emits the longest stable decoded prefix. */
933
- class StreamDecoder {
1357
+ /** Incremental detokenizer: emits the longest stable decoded prefix.
1358
+ *
1359
+ * Byte parity with mlx-lm's streaming detokenizers (the drop-in contract is
1360
+ * rendered BYTES, not just token ids). For BPE/ByteLevel tokenizers, two
1361
+ * mlx-lm 0.31.3 BPEStreamingDetokenizer behaviors our full-sequence decode
1362
+ * lacks (tokenizer_utils.py:195-226):
1363
+ *
1364
+ * 1. `trimsLeadingSpace` — mlx-lm drops ONE leading " " at the start of the
1365
+ * generated sequence (`_maybe_trim_space`); trim it here. SPM decode
1366
+ * already matches (see LoadedTokenizer.trimsLeadingSpace).
1367
+ * 2. `bareSpaceTokenId` — add_token WITHHOLDS a single-char byte-32 token
1368
+ * ("Ġ") in `_unflushed` ("For single spaces wait until the next token"),
1369
+ * flushing it together with the NEXT token — and mlx_lm.server NEVER
1370
+ * calls detokenizer.finalize() (zero hits in server.py 0.31.3), so a
1371
+ * generation ENDING on bare-space token(s) silently drops their spaces
1372
+ * from the served bytes. push() withholds those spaces; flush() drops a
1373
+ * trailing bare-space run (the held text dies with the request, exactly
1374
+ * like mlx-lm serve). Re-check both if upstream ever adds a finalize()
1375
+ * call. LATENT HAZARD (deliberately not emulated): models with
1376
+ * clean_up_tokenization_spaces=true get an ADDITIONAL mid-stream rule
1377
+ * (`_space_matches`: held space dropped before "." "," "'s" …) — both
1378
+ * current BPE targets have it false (MiniCPM5), so it never fires here.
1379
+ *
1380
+ * Exported for unit tests (serve-detok mlx-lm byte parity). */
1381
+ export class StreamDecoder {
1382
+ static readonly #WINDOW_TOKENS = 32;
934
1383
  #ids: number[] = [];
935
1384
  #emitted = "";
1385
+ #warnedRevision = false;
1386
+ /** Token index where the bounded decode suffix begins. */
1387
+ #windowStart = 0;
1388
+ /** Exact decoded text before #windowStart, established by suffix matching. */
1389
+ #windowPrefix = "";
1390
+ readonly #trimLeadingSpace: boolean;
1391
+ readonly #bareSpaceId: number | undefined;
936
1392
 
937
1393
  constructor(
938
1394
  readonly tokenizer: LoadedTokenizer,
939
1395
  readonly skipSpecialTokens = true,
940
- ) {}
1396
+ ) {
1397
+ this.#trimLeadingSpace = tokenizer.trimsLeadingSpace === true;
1398
+ this.#bareSpaceId = tokenizer.bareSpaceTokenId;
1399
+ }
1400
+
1401
+ #decode(ids: number[]): string {
1402
+ const full = this.#decodeRaw(ids);
1403
+ return this.#trimLeadingSpace && full.startsWith(" ") ? full.slice(1) : full;
1404
+ }
1405
+
1406
+ #decodeRaw(ids: number[]): string {
1407
+ return this.tokenizer.decode(ids, this.skipSpecialTokens);
1408
+ }
1409
+
1410
+ /** Decode only a bounded suffix once an exact text anchor can be proven.
1411
+ *
1412
+ * At each rebase, decode a candidate suffix independently and accept it only
1413
+ * when it is literally the suffix of the exact text decoded so far. This
1414
+ * preserves ByteLevel/SPM boundary behavior while reducing steady-state
1415
+ * decode work from the whole generation to 32–64 tokens. If no suffix can be
1416
+ * anchored (an unusual/global decoder), leave #windowStart at zero and keep
1417
+ * the full-history correctness path. */
1418
+ #decodeIncremental(): string {
1419
+ const full = this.#windowStart === 0
1420
+ ? this.#decode(this.#ids)
1421
+ : this.#windowPrefix + this.#decodeRaw(this.#ids.slice(this.#windowStart));
1422
+
1423
+ if (this.#ids.length - this.#windowStart < StreamDecoder.#WINDOW_TOKENS * 2)
1424
+ return full;
1425
+
1426
+ const preferred = this.#ids.length - StreamDecoder.#WINDOW_TOKENS;
1427
+ const oldest = Math.max(this.#windowStart, preferred - StreamDecoder.#WINDOW_TOKENS);
1428
+ for (let start = preferred; start >= oldest; start--) {
1429
+ const tail = this.#decodeRaw(this.#ids.slice(start));
1430
+ if ((tail.length > 0 || full.length === 0) && full.endsWith(tail)) {
1431
+ this.#windowStart = start;
1432
+ this.#windowPrefix = full.slice(0, full.length - tail.length);
1433
+ break;
1434
+ }
1435
+ }
1436
+ return full;
1437
+ }
941
1438
 
942
1439
  push(token: number): string {
943
1440
  this.#ids.push(token);
944
- const full = this.tokenizer.decode(this.#ids, this.skipSpecialTokens);
1441
+ // Bare-space hold-back (mlx-lm add_token keeps "Ġ" in _unflushed): don't
1442
+ // advance #emitted; the held space(s) flush as part of the next
1443
+ // non-bare-space token's delta — consecutive bare spaces accumulate.
1444
+ if (token === this.#bareSpaceId) return "";
1445
+ const full = this.#decodeIncremental();
945
1446
  // hold back a trailing replacement char (partial multi-byte sequence)
946
1447
  const stable = full.endsWith("�") ? full.slice(0, -1) : full;
947
1448
  if (!stable.startsWith(this.#emitted)) {
948
- // decoder revised earlier text (rare); re-emit from scratch
949
- const out = stable;
1449
+ // The decoder revised already-streamed text (cleanup rules — never
1450
+ // fires for the shipped tokenizers, which all have
1451
+ // clean_up_tokenization_spaces=false; see the LATENT HAZARD note in
1452
+ // the class doc). An SSE client cannot un-receive bytes, so the old
1453
+ // "re-emit from scratch" answer DUPLICATED the whole stream
1454
+ // (2026-07-07 review). Truncate-safe resync instead: emit only the
1455
+ // length-extension and keep the emitted watermark monotone — byte
1456
+ // drift is confined to the revised span, and the once-per-stream
1457
+ // warning makes the tokenizer that needs real _space_matches
1458
+ // emulation loud instead of silently wrong.
1459
+ if (!this.#warnedRevision) {
1460
+ this.#warnedRevision = true;
1461
+ console.warn(
1462
+ "[detok] decoder revised already-streamed text (clean_up_tokenization_spaces?) — " +
1463
+ "resyncing without re-emit; add mlx-lm's cleanup-rule emulation for this tokenizer",
1464
+ );
1465
+ }
1466
+ if (stable.length <= this.#emitted.length) return "";
1467
+ const out = stable.slice(this.#emitted.length);
950
1468
  this.#emitted = stable;
951
1469
  return out;
952
1470
  }
@@ -956,7 +1474,15 @@ class StreamDecoder {
956
1474
  }
957
1475
 
958
1476
  flush(): string {
959
- const full = this.tokenizer.decode(this.#ids, this.skipSpecialTokens);
1477
+ // mlx_lm.server never finalize()s: a trailing bare-space run stays
1478
+ // withheld forever, so its text is dropped from the served bytes.
1479
+ let ids = this.#ids;
1480
+ if (this.#bareSpaceId !== undefined) {
1481
+ let n = ids.length;
1482
+ while (n > 0 && ids[n - 1] === this.#bareSpaceId) n--;
1483
+ if (n < ids.length) ids = ids.slice(0, n);
1484
+ }
1485
+ const full = this.#decode(ids);
960
1486
  const delta = full.slice(this.#emitted.length);
961
1487
  this.#emitted = full;
962
1488
  return delta;
@@ -986,41 +1512,141 @@ export function createServer(
986
1512
  // cache by attention type. Non-batchable requests (vision / adapters /
987
1513
  // repetition penalty / user seed / explicit kv-quant) drain to the serial
988
1514
  // 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.
1515
+ // DEFAULT 8 (flipped 2026-07-05, Josh's call, after GATE-B1-SPEED): a
1516
+ // lone request through the batch lane IS the serial engine (adopted
1517
+ // serial-class caches, compiled decode, prompt cache + SSD restore;
1518
+ // 0.992-0.996 paired decode ratios, byte-identical output), so the cap
1519
+ // only changes behavior when concurrent requests actually arrive — the
1520
+ // agentic sub-agent workload. --batch 1 pins strict serial for
1521
+ // arrival-independent numerics. 8 = optiq's Mac-safe concurrency.
1522
+ const batch = Math.max(1, Math.floor(serverOptions.batch ?? 8));
1523
+ const defaultGeneratedTokens =
1524
+ serverOptions.defaultMaxTokens ?? ctx.glmMemoryPlan?.maxGenerationTokens;
1525
+
1526
+ // KV-quant scheme, resolved once. UNSET now means bf16 (flipped
1527
+ // 2026-07-05 with the naked-=-L1 default): quantized KV measured 5–20%
1528
+ // SLOWER decode than bf16 at ≤16k on every model — on mlx-lm too (its
1529
+ // kv8 oracle trails its own bf16) — so it pays only in memory headroom
1530
+ // and must be an explicit opt-in (--kv-quant config|4|8, or --l2
1531
+ // whose presets pass it explicitly). The CLI always passes kvQuant now;
1532
+ // this fallback is the library-user default and matches the CLI's.
999
1533
  const configScheme = ctx.kvConfig?.length ? { kvConfig: ctx.kvConfig } : {};
1000
- const kvScheme: Pick<GenerateOptions, "kvBits" | "kvConfig" | "quantizedKvStart"> =
1001
- serverOptions.kvQuant === "off" ? {}
1534
+ // Mutually exclusive by contract (GenerateOptions.turboQuant doc): a
1535
+ // programmatic caller setting both gets turboQuant — say so, like the
1536
+ // other risky-combination warnings below.
1537
+ if (serverOptions.turboQuant && serverOptions.kvQuant && serverOptions.kvQuant !== "off")
1538
+ console.warn(
1539
+ `[kv-quant] both turboQuant and --kv-quant ${serverOptions.kvQuant} are set; ` +
1540
+ `turboQuant wins (they are mutually exclusive).`,
1541
+ );
1542
+ const kvScheme: Pick<GenerateOptions, "kvBits" | "kvConfig" | "quantizedKvStart" | "turboQuant"> =
1543
+ serverOptions.turboQuant ? { turboQuant: serverOptions.turboQuant, quantizedKvStart: 0 }
1544
+ : serverOptions.kvQuant === "off" ? {}
1002
1545
  : serverOptions.kvQuant === "config" ? configScheme
1003
1546
  : typeof serverOptions.kvQuant === "number"
1004
1547
  ? { 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))
1548
+ : {}; // unset → bf16 (L1 default; quantized KV is opt-in)
1549
+ // Phase 3.1: kvConfig whose layers are all full-attention BATCHES (the
1550
+ // scheduler applies the mixed scheme per row); uniform kvBits and configs
1551
+ // touching rotating layers still route those requests serial. The warning
1552
+ // only fires for the still-serial compositions.
1553
+ if (batch > 1 && kvScheme.kvBits)
1554
+ console.warn(
1555
+ `[batch] --batch ${batch} with uniform --kv-quant ${kvScheme.kvBits}: uniform ` +
1556
+ `quantized KV is serial-only (it touches rotating layers) — those requests ` +
1557
+ `won't batch. Use --kv-quant config for batched mixed-precision KV, or omit ` +
1558
+ `--kv-quant to batch in bf16. (docs/design/unified-engine-frontier-plan.md)`,
1559
+ );
1560
+ // TurboQuant is solo-only in v1 (novel cache class, not batchable by
1561
+ // construction — see GenerationGateway#modelCachesBatchable/willBatch).
1562
+ if (batch > 1 && kvScheme.turboQuant)
1563
+ console.warn(
1564
+ `[batch] --batch ${batch} with --kv-quant turbo: TurboQuant is serial-only in v1 ` +
1565
+ `— those requests won't batch. Omit --kv-quant to batch in bf16. ` +
1566
+ `(docs/design/turboquant-kv.md)`,
1567
+ );
1568
+ // Quantized KV (any axis) excludes the spec lane: drafted requests fall to
1569
+ // the normal serial path WITH the KV scheme applied rather than losing it
1570
+ // silently (mirrors the affine kvBits/kvConfig exclusion; the spec loop is
1571
+ // bf16-KV-only in v1).
1572
+ if (ctx.draft && (kvScheme.turboQuant || kvScheme.kvBits || kvScheme.kvConfig?.length))
1008
1573
  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)`,
1574
+ `[spec] --draft-model with quantized KV (--kv-quant ${serverOptions.turboQuant ? "turbo" : String(serverOptions.kvQuant)}): ` +
1575
+ `the speculative lane is bf16-KV-only in v1 — requests keep the KV scheme and ` +
1576
+ `decode serially WITHOUT speculation. Omit --kv-quant to speculate. ` +
1577
+ `(docs/design/dspark-serving-program.md Phase 4)`,
1012
1578
  );
1579
+ // TurboQuant head-dim fail-fast (2026-07-07 review): the cache class only
1580
+ // supports {64,128,256,512} (sign-vector + Lloyd-Max table coverage) and
1581
+ // used to validate LAZILY on the first append — an unsupported model
1582
+ // (e.g. a 72/80/96 head dim) accepted --kv-quant turbo at startup and
1583
+ // then 500'd EVERY request from inside prefill. The config knows the
1584
+ // full-attention head dim (the only kind that converts; sliding layers
1585
+ // stay bf16) — refuse at createServer instead.
1586
+ if (kvScheme.turboQuant) {
1587
+ const dim = ctx.model.config.text.globalHeadDim;
1588
+ if (!(TURBOQUANT_HEAD_DIMS as readonly number[]).includes(dim))
1589
+ throw new Error(
1590
+ `--kv-quant turbo: this model's full-attention head_dim is ${dim}; ` +
1591
+ `TurboQuant supports {${TURBOQUANT_HEAD_DIMS.join(",")}} ` +
1592
+ `(docs/design/turboquant-kv.md) — use --kv-quant config|4|8 or omit it`,
1593
+ );
1594
+ }
1595
+ // Paged KV v1 (docs/design/paged-kv-cache.md): explicit refusals, not
1596
+ // silent downgrades — the incompatible combos would otherwise degrade
1597
+ // quietly (batch: paged caches can't merge; kv-quant: the swap would
1598
+ // drop the scheme, the exact composition bug the kvQuant gate above
1599
+ // exists to prevent; draft: the spec lane assumes serial-class caches).
1600
+ if (serverOptions.pagedKv) {
1601
+ if (batch > 1)
1602
+ throw new Error(
1603
+ `--paged-kv is serial-only in v1 — add --batch 1 (got --batch ${batch}). ` +
1604
+ `Batched paging is the follow-up PR (docs/design/paged-kv-cache.md).`,
1605
+ );
1606
+ if (kvScheme.kvBits || kvScheme.kvConfig?.length || kvScheme.turboQuant)
1607
+ throw new Error(
1608
+ `--paged-kv is bf16-only in v1 — omit --kv-quant (quantized paged ` +
1609
+ `blocks are a documented follow-up, docs/design/paged-kv-cache.md).`,
1610
+ );
1611
+ if (ctx.draft)
1612
+ throw new Error(
1613
+ `--paged-kv cannot combine with --draft-model in v1 ` +
1614
+ `(docs/design/paged-kv-cache.md non-goals).`,
1615
+ );
1616
+ if (!ctx.model.config.modelType.startsWith("gemma4"))
1617
+ throw new Error(
1618
+ `--paged-kv v1 supports Gemma4-family models only ` +
1619
+ `(this model: ${ctx.model.config.modelType}) — docs/design/paged-kv-cache.md.`,
1620
+ );
1621
+ const bs = serverOptions.pagedKv.blockSize;
1622
+ if (bs !== undefined && (!Number.isInteger(bs) || bs <= 0))
1623
+ throw new Error(`--paged-kv-block-size must be a positive integer (got ${bs})`);
1624
+ if (serverOptions.ssdCacheDir)
1625
+ console.warn(
1626
+ "[paged-kv] --ssd-cache has no effect: paged requests bypass the prompt " +
1627
+ "cache (v1 non-goal), so nothing reaches the SSD tier.",
1628
+ );
1629
+ }
1013
1630
 
1014
1631
  // SSD cold tier (docs/design/ssd-kv-cold-tier.md): prefix KV survives RAM
1015
1632
  // eviction and restarts. Compatibility key = configFingerprint (graph
1016
1633
  // shape) + the EFFECTIVE kv scheme (flags pick bf16 vs config vs uniform
1017
1634
  // on the same model — restored caches must match what serving produces) +
1018
1635
  // tokenizer hash (ids must keep meaning the same text).
1636
+ // RAM prompt-cache cap: 8 GB default (Josh's call, 2026-07-06). The old
1637
+ // flat 2e9 was an anti-OOM reflex sized for a 1B model in 24 GB — a
1638
+ // single full-context 12B entry is ~0.6 GB, so it silently shrank to
1639
+ // "three entries" on big models. 8 GB holds a dozen 12B contexts; idle
1640
+ // entries still demote to SSD and LRU eviction bounds pressure.
1641
+ // --prompt-cache <GB> overrides; 0 disables.
1642
+ const promptCacheCap = serverOptions.promptCacheBytes ?? 8e9;
1019
1643
  let ssdStore: SsdCacheStore | null = null;
1020
1644
  if (serverOptions.ssdCacheDir) {
1021
- if ((serverOptions.promptCacheBytes ?? 2e9) <= 0)
1645
+ if (promptCacheCap <= 0)
1022
1646
  throw new Error("--ssd-cache requires the RAM prompt cache (--prompt-cache 0 disables it)");
1023
- const schemeKey = kvScheme.kvBits
1647
+ const schemeKey = kvScheme.turboQuant
1648
+ ? `turbo-k${kvScheme.turboQuant.kBits}v${kvScheme.turboQuant.vBits}`
1649
+ : kvScheme.kvBits
1024
1650
  ? `kv${kvScheme.kvBits}`
1025
1651
  : kvScheme.kvConfig?.length ? "config" : "bf16";
1026
1652
  const tokJson = readFileSync(`${ctx.model.config.modelDir}/tokenizer.json`);
@@ -1039,11 +1665,51 @@ export function createServer(
1039
1665
  );
1040
1666
  }
1041
1667
 
1668
+ // Layer 0 (unified-engine plan): the cold tier is bound INTO the prompt
1669
+ // cache — take() itself tiers over it, so the batch scheduler and every
1670
+ // other consumer get SSD restores through the same API; eviction AND
1671
+ // idle demotion spill into it.
1672
+ const coldTier = ssdStore
1673
+ ? {
1674
+ find: (prompt: number[], ns: string) => {
1675
+ const h = ssdStore!.find(prompt, ns);
1676
+ return h ? { prefixLen: h.prefixLen, handle: h.entry } : null;
1677
+ },
1678
+ restore: (handle: unknown) => {
1679
+ const loaded = ssdStore!.restore(handle as import("./ssd-cache").SsdIndexEntry, ctx.model);
1680
+ if (!loaded) return null;
1681
+ // Restore is a STREAMED COPY (2026-07-07 A7-restore): the caches
1682
+ // own their bytes and no mapping outlives loadKvCache — nothing
1683
+ // to pin, nothing to unmap. retain stays in the entry contract
1684
+ // as a no-op so callers' dispose ordering is unchanged.
1685
+ return { tokens: loaded.tokens, caches: loaded.caches, retain: () => {} };
1686
+ },
1687
+ store: (tokens: number[], caches: import("./model/gemma4").Cache[], ns: string) => {
1688
+ ssdStore!.store(tokens, caches, ns);
1689
+ },
1690
+ }
1691
+ : null;
1042
1692
  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,
1693
+ promptCacheCap,
1694
+ // RAM eviction AND idle demotion spill to the cold tier NON-BLOCKING
1695
+ // (2026-07-06, same contract as the write-behind snapshot below): the
1696
+ // cache hands us OWNED zero-copy clones + copied tokens (made under
1697
+ // the generation lock — microseconds; entries are immutable so the
1698
+ // clones stay consistent), and the flush goes through the BOUNDED
1699
+ // SpillQueue -> storeAsync (idle-gated per tensor, see below) ->
1700
+ // clone disposal on every settle/drop path (that dispose is what
1701
+ // actually frees the demoted GPU memory; the queue's byte cap keeps
1702
+ // starved-gate retention bounded — 2026-07-07 review fix). spillQueue
1703
+ // is declared with the snapshot scheduler below; safe to close over
1704
+ // here because spills only fire at put()/demoteIdle time, long after
1705
+ // init.
1706
+ ssdStore
1707
+ ? {
1708
+ spillOwned: (entry) =>
1709
+ spillQueue!.enqueue({ tokens: entry.tokens, caches: entry.caches, ns: entry.ns }),
1710
+ }
1711
+ : null,
1712
+ coldTier,
1047
1713
  );
1048
1714
  // Responses-API store for previous_response_id resumption (Phase 11):
1049
1715
  // TTL + byte-capped LRU, port of optiq/response_store.py. Pairs with
@@ -1057,68 +1723,149 @@ export function createServer(
1057
1723
  * `enqueue` promise queue is gone; the gateway's AsyncMutex is THE lock.)
1058
1724
  * onToken returning `false` halts generation early (stop
1059
1725
  * 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. */
1726
+ * Media (vision/audio embeddings-prefill) requests bypass the prompt
1727
+ * cache: soft tokens are identical placeholder ids, so prefix matching
1728
+ * across different images/clips would false-hit. */
1063
1729
  const runGeneration = async (
1064
1730
  promptIds: number[],
1065
1731
  options: GenerateOptions,
1066
1732
  onToken: (token: number, logprobs?: TokenLogprobs) => void | boolean | Promise<void | boolean>,
1067
- vision?: { embeddings: import("./mlx/array").MlxArray; imageMask: import("./mlx/array").MlxArray },
1733
+ vision?: {
1734
+ embeddings: import("./mlx/array").MlxArray;
1735
+ imageMask?: import("./mlx/array").MlxArray;
1736
+ multimodalMask?: import("./mlx/array").MlxArray;
1737
+ },
1068
1738
  ) => {
1739
+ // Speculative decoding (serve --draft-model): spec-ELIGIBLE requests
1740
+ // decode through the verify loop; the rest fall through to the normal
1741
+ // serial path (never wrong results, just no speedup — logged once per
1742
+ // combination class would be noise, so silent). Eligibility v1:
1743
+ // text-only, base weights, no logprobs capture, bf16 KV — ALL quantized
1744
+ // KV axes excluded (affine kvBits/kvConfig AND turboQuant; the spec loop
1745
+ // builds fresh bf16 caches and never calls maybeQuantizeKv, so routing a
1746
+ // turbo request here would silently drop the operator's KV scheme —
1747
+ // 2026-07-07 review). Grammar COMPOSES (Phase C constrained verify walk
1748
+ // — see the serve-loop.ts header). Prompt-cache reuse is bypassed on
1749
+ // the spec path v1.
1750
+ if (
1751
+ ctx.draft &&
1752
+ !vision &&
1753
+ !options.adapters?.length &&
1754
+ !options.logprobs &&
1755
+ !options.kvBits &&
1756
+ !options.kvConfig &&
1757
+ !options.turboQuant &&
1758
+ // Belt: --paged-kv + --draft-model is refused at createServer; this
1759
+ // keeps the spec lane paged-free even for programmatic callers.
1760
+ !options.pagedKv
1761
+ ) {
1762
+ const { specServeRun } = await import("./spec/serve-loop");
1763
+ return withModelWiredLimit(
1764
+ ctx.model,
1765
+ () => withModelUsageFlush(
1766
+ ctx.model,
1767
+ () => specServeRun(
1768
+ ctx.model, ctx.draft!.provider, ctx.draft!.numDraftTokens,
1769
+ promptIds, options, onToken,
1770
+ ),
1771
+ ),
1772
+ );
1773
+ }
1069
1774
  // Cache entries are adapter-specific: KV computed under one adapter
1070
1775
  // must never seed another's (or the base's) prefill.
1071
1776
  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
- }
1777
+ // Paged-KV request scope (docs/design/paged-kv-cache.md): media
1778
+ // prompts (bidir overlay) and LoRA-adapter requests are v1 non-goals —
1779
+ // they run the PLAIN cache path even under --paged-kv (scope the flag
1780
+ // per request, never 400). Effective value computed ONCE so the
1781
+ // prompt-cache bypass below and the generate() options can't disagree.
1782
+ const pagedKv = vision || options.adapters?.length ? undefined : options.pagedKv;
1783
+ // Paged requests bypass the prompt cache entirely (v1 non-goal:
1784
+ // PagedKVCache has no cloneKvCaches/restore path — the vision
1785
+ // precedent). Fresh caches per request, disposed on completion.
1786
+ const skipPromptCache = Boolean(vision || pagedKv);
1787
+ // Both tiers in one call (Layer 0): take() prefers a strictly-longer
1788
+ // SSD prefix, restores it zero-copy, and trims — see PromptCache.take.
1789
+ const entry = skipPromptCache ? null : promptCache.take(promptIds, cacheNs);
1097
1790
  const caches = entry?.caches ?? ctx.model.makeCache();
1791
+ // Prompt-boundary snapshot (the multi-turn agent fix, 2026-07-04): the
1792
+ // prompt+gen entry put() below is UNTRIMMABLE at context > sliding
1793
+ // window (wrapped rings) and under quantized KV (mid-group), so any
1794
+ // decode→encode roundtrip drift in the reply the client sends back
1795
+ // turns the next turn into a total miss (measured: 12B turn-2 TTFT
1796
+ // 8.9 s instead of ~0.2 s). A prompt-ONLY entry is always an exact
1797
+ // prefix of the next turn's rendering regardless of reply drift.
1798
+ // Zero-copy (cloneKvCaches = slice views); only for substantial cold
1799
+ // prefills, where the re-prefill it saves is worth an extra entry.
1800
+ // The oracle invariant (mlx-lm insert_segments): a trim-free STRICT
1801
+ // prefix of the prompt exists for EVERY substantial request — cap the
1802
+ // boundary at len-1 so even a stableLen == len prompt (e4b: the
1803
+ // template tail survives the probe render) snapshots prompt[:-1]. An
1804
+ // exact repeat then matches with trimNeeded == 0, bypassing
1805
+ // isTrimmable() entirely — the only reuse path a wrapped ring has.
1806
+ const boundary = Math.min(options.snapshotAt ?? promptIds.length, promptIds.length - 1);
1807
+ // Re-snapshot on EVERY substantial request whose stable boundary extends
1808
+ // past the cached prefix; the clone is zero-copy views, so re-putting
1809
+ // is ~free.
1810
+ const snapshotBoundary =
1811
+ !skipPromptCache && boundary >= 256 && boundary > (entry?.tokens.length ?? 0);
1098
1812
  try {
1099
1813
  const gen = generate(ctx.model, promptIds, {
1100
1814
  ...options,
1815
+ pagedKv, // request-scoped (undefined strips the server-wide flag)
1101
1816
  cache: caches,
1102
- ...(vision ? { promptEmbeddings: vision.embeddings, imageMask: vision.imageMask } : {}),
1817
+ ...(snapshotBoundary
1818
+ ? {
1819
+ // snapshotAt MUST travel with the hook: generate() splits the
1820
+ // prefill at exactly this many tokens and fires the hook while
1821
+ // the caches hold exactly that prefix — putting boundary
1822
+ // tokens against caches at any other offset is silent KV
1823
+ // corruption.
1824
+ snapshotAt: boundary,
1825
+ onPrefillDone: () => {
1826
+ try {
1827
+ promptCache.put(promptIds.slice(0, boundary), cloneKvCaches(caches), cacheNs);
1828
+ } catch (e) {
1829
+ console.warn(`prompt-boundary snapshot skipped: ${(e as Error).message}`);
1830
+ }
1831
+ },
1832
+ }
1833
+ : {}),
1834
+ ...(vision
1835
+ ? {
1836
+ promptEmbeddings: vision.embeddings,
1837
+ ...(vision.imageMask ? { imageMask: vision.imageMask } : {}),
1838
+ ...(vision.multimodalMask ? { multimodalMask: vision.multimodalMask } : {}),
1839
+ }
1840
+ : {}),
1103
1841
  });
1104
1842
  for await (const t of gen) {
1105
1843
  if ((await onToken(t.token, t.logprobs)) === false) break;
1106
1844
  }
1107
1845
  const s = gen.stats!; // set on completion AND on early break
1108
- if (vision) {
1846
+ if (skipPromptCache) {
1847
+ // Vision and paged-KV requests own their caches for exactly one
1848
+ // generation (paged: v1 non-goal — no PromptCache integration).
1109
1849
  for (const c of caches) c.dispose();
1110
1850
  } else {
1851
+ // put() fires onPut → the debounced write-behind SSD snapshot
1852
+ // (wired below), covering the batch lane's puts too.
1111
1853
  promptCache.put(s.cacheTokens, caches, cacheNs, entry?.retain);
1112
- scheduleSsdSnapshot(s.cacheTokens, cacheNs);
1113
1854
  }
1114
1855
  return s;
1115
1856
  } catch (e) {
1116
1857
  for (const c of caches) c.dispose();
1117
1858
  entry?.retain?.();
1859
+ // A throw BEFORE generate()/the gateway took ownership (promptCache
1860
+ // take / makeCache) would leak the grammar's WASM matcher; dispose()
1861
+ // is idempotent, so this is safe when the throw came from inside run
1862
+ // (whose finally already disposed it).
1863
+ options.grammar?.dispose();
1118
1864
  throw e;
1119
1865
  } finally {
1120
1866
  vision?.embeddings.dispose();
1121
- vision?.imageMask.dispose();
1867
+ vision?.imageMask?.dispose();
1868
+ vision?.multimodalMask?.dispose();
1122
1869
  options.visionPixels?.dispose();
1123
1870
  }
1124
1871
  };
@@ -1132,40 +1879,160 @@ export function createServer(
1132
1879
  // extraction must not race a generation mutating the entry); if the entry
1133
1880
  // was evicted (already spilled) or extended (a newer schedule pending)
1134
1881
  // meanwhile, findExact misses and nothing is written.
1882
+ // Keyed by ns AND entry length: a sub-second final [prompt+gen] put must
1883
+ // not cancel the pending boundary-snapshot write (they are DIFFERENT
1884
+ // entries; for wrapped-ring models the boundary file is the only one a
1885
+ // restart can use). Same-length reschedules still coalesce; stale keys
1886
+ // self-clean at fire time (findExact misses once the RAM tier superseded
1887
+ // the entry, so nothing extra is written).
1888
+ //
1889
+ // NON-BLOCKING (2026-07-06, the write-behind persistence contract): the
1890
+ // gateway lock is held only for a zero-copy SNAPSHOT (findExact +
1891
+ // cloneKvCaches — microseconds; entries are immutable so the clones are
1892
+ // consistent forever). The flush itself runs OFF the lock via
1893
+ // storeAsync, and writes chain serially so two multi-hundred-MB flushes
1894
+ // never overlap.
1895
+ //
1896
+ // IDLE-GATED (2026-07-07, the decode@ctx fix): "off the lock" was not
1897
+ // enough — every per-tensor flush step is a blocking GPU sync on the
1898
+ // SAME stream decode uses (ops.contiguous enqueues a kernel, rawBytesView
1899
+ // evals) plus a synchronous multi-MB writeSync, and the setImmediate
1900
+ // pacing interleaved those slices exactly between decode tokens. A ~16k
1901
+ // entry's flush overlapping the bench's ctx repeats depressed decode@ctx
1902
+ // ~9% on e4b (mlx-lm runs no equivalent background work). Now every step
1903
+ // — including the first — awaits gateway.onIdle() (ssdFlushGate), so the
1904
+ // flush only progresses while NOTHING is generating and pauses when a
1905
+ // request arrives mid-flush. Tradeoffs, accepted: durability waits for a
1906
+ // quiet moment (single-user serving quiesces constantly; sustained
1907
+ // hammering defers the flush AND the spill clones' GPU-memory release —
1908
+ // bounded by the chain), and one in-flight tensor step (~10-15 MB) can
1909
+ // still land ahead of a just-arrived request. MLX_BUN_SSD_WRITEBEHIND=0
1910
+ // disables write-behind snapshots entirely — the paired-A/B lever + kill
1911
+ // switch (restart survival then degrades to spill-on-evict only).
1912
+ const writeBehindOn = process.env.MLX_BUN_SSD_WRITEBEHIND !== "0";
1913
+ const ssdFlushGate = (): Promise<void> => gateway.onIdle();
1135
1914
  const ssdPending = new Map<string, ReturnType<typeof setTimeout>>();
1915
+ // Bounded write-behind queue (2026-07-07 review fix — see SpillQueue in
1916
+ // kv-store.ts): pending clones pin their entries' GPU buffers while the
1917
+ // idle gate starves under sustained traffic, so QUEUED bytes are capped —
1918
+ // over cap the oldest queued spill drops (clone disposed immediately, a
1919
+ // future cache miss, never a wrong result) instead of accumulating past
1920
+ // the prompt-cache cap. Default 2 GB (a quarter of the 8 GB RAM-cache
1921
+ // default); MLX_BUN_SSD_SPILL_QUEUE_GB overrides. There is deliberately
1922
+ // NO shutdown flush: exiting under traffic loses queued write-behind
1923
+ // entries (restart survival degrades to whatever flushed) — accepted for
1924
+ // a best-effort cache tier; drops/pending are visible in /stats.
1925
+ // `|| 2` would coerce an explicit "0" back to 2 GB — parse so 0 works
1926
+ // (cap 0 = keep only the newest + in-flight clone pinned; the soft cap
1927
+ // never drops the item just enqueued).
1928
+ const spillQueueGbRaw = Number(process.env.MLX_BUN_SSD_SPILL_QUEUE_GB);
1929
+ const spillQueueCapBytes =
1930
+ (Number.isFinite(spillQueueGbRaw) && spillQueueGbRaw >= 0 ? spillQueueGbRaw : 2) * 1024 ** 3;
1931
+ const spillQueue = ssdStore
1932
+ ? new SpillQueue(
1933
+ spillQueueCapBytes,
1934
+ cacheBytes,
1935
+ (item) => ssdStore!.storeAsync(item.tokens, item.caches, item.ns, ssdFlushGate),
1936
+ (caches) => { for (const c of caches) c.dispose(); },
1937
+ )
1938
+ : null;
1136
1939
  const scheduleSsdSnapshot = (tokens: number[], ns: string): void => {
1137
- if (!ssdStore || tokens.length === 0) return;
1138
- const prev = ssdPending.get(ns);
1940
+ if (!ssdStore || !writeBehindOn || tokens.length === 0) return;
1941
+ const key = `${ns}|${tokens.length}`;
1942
+ const prev = ssdPending.get(key);
1139
1943
  if (prev) clearTimeout(prev);
1140
- ssdPending.set(ns, setTimeout(() => {
1141
- ssdPending.delete(ns);
1944
+ const fire = (): void => {
1945
+ ssdPending.delete(key);
1946
+ // Same guard as demoteIdle below (2026-07-07 review fix): the
1947
+ // runExclusive registers a serial waiter, which HOLDS BATCH ADMISSION
1948
+ // until every running row drains — a background durability timer must
1949
+ // never do that. Busy → re-arm and try again in 5 s (entries are
1950
+ // immutable; a late snapshot is just as valid).
1951
+ if (gateway.activeRows > 0 || gateway.pendingRows > 0) {
1952
+ const t = setTimeout(fire, 5000);
1953
+ t.unref?.();
1954
+ ssdPending.set(key, t);
1955
+ return;
1956
+ }
1142
1957
  void gateway.runExclusive(async () => {
1143
1958
  const e = promptCache.findExact(tokens, ns);
1144
- if (e) ssdStore!.store(e.tokens, e.caches, ns);
1959
+ if (!e) return null;
1960
+ return { tokens: e.tokens, caches: cloneKvCaches(e.caches) };
1961
+ }).then((snap) => {
1962
+ if (!snap) return;
1963
+ spillQueue!.enqueue({ tokens: snap.tokens, caches: snap.caches, ns });
1145
1964
  }).catch(() => { /* cold tier is best-effort */ });
1146
- }, 1000));
1965
+ };
1966
+ ssdPending.set(key, setTimeout(fire, 1000));
1147
1967
  };
1968
+ // Every put() — serial lane AND batch scheduler — schedules the snapshot
1969
+ // (Layer 0: batch-lane entries survive restarts too, not just evictions).
1970
+ promptCache.onPut = (tokens, ns) => scheduleSsdSnapshot(tokens, ns);
1971
+
1972
+ // Idle demotion (Layer 0): entries unused for --ssd-demote-idle seconds
1973
+ // spill to SSD and free their GPU memory — the RAM tier drains between
1974
+ // bursts while every prefix stays reachable (take() restores via
1975
+ // zero-copy mmap). The exclusive section below covers only the sweep's
1976
+ // zero-copy clone + dispose (spillOwned above); the SSD write itself
1977
+ // runs off-lock on ssdWriteChain. Swept only when the engine is TRULY idle: the
1978
+ // runExclusive below registers as a serial waiter, which would DRAIN a
1979
+ // running batch, so the activity check guards it (a momentary exclusive
1980
+ // grab of an idle engine is free). 0 disables.
1981
+ const demoteIdleMs = (serverOptions.ssdDemoteIdleSec ?? (ssdStore ? 300 : 0)) * 1000;
1982
+ let demoteTimer: ReturnType<typeof setInterval> | null = null;
1983
+ if (ssdStore && demoteIdleMs > 0) {
1984
+ demoteTimer = setInterval(() => {
1985
+ if (gateway.activeRows > 0 || gateway.pendingRows > 0) return;
1986
+ void gateway.runExclusive(async () => {
1987
+ const n = promptCache.demoteIdle(demoteIdleMs);
1988
+ if (n > 0) console.log(`[ssd-cache] demoted ${n} idle entr${n === 1 ? "y" : "ies"} to disk`);
1989
+ }).catch(() => {});
1990
+ }, Math.max(30_000, demoteIdleMs / 4));
1991
+ demoteTimer.unref?.();
1992
+ }
1148
1993
 
1149
1994
  // The lane picker: routes each request to the serial path (runGeneration,
1150
1995
  // above) or the continuous-batching scheduler, keeping the two off the GPU
1151
1996
  // (and shared loraState) at the same time. See src/serve/generation-gateway.ts.
1152
- const gateway = new GenerationGateway(ctx.model, batch, runGeneration);
1997
+ const gateway = new GenerationGateway(ctx.model, batch, runGeneration, {
1998
+ kvBudgetBytes: serverOptions.kvBudgetBytes,
1999
+ // Phase 3.1: the server-wide scheme travels to the gateway once; the
2000
+ // gateway batches kv-quant requests when the scheme is the batchable
2001
+ // kvConfig/full-attention composition (see #kvBatchable) and the
2002
+ // scheduler applies it — otherwise those requests route serial as before.
2003
+ kvScheme,
2004
+ // Phase 3.2: batch-lane prompt-cache reuse — joiners take() the longest
2005
+ // usable prefix (multi-turn chat TTFT under --batch N); never-merged
2006
+ // lone rows put() back on finish. Vision/adapter requests never batch,
2007
+ // so the serial lane's bypass rules are preserved by routing.
2008
+ promptCache,
2009
+ });
1153
2010
 
1154
2011
  // Admission ceiling, resolved once (Phase 5 memoryBudget enforcement).
1155
2012
  // fit() solves max safe context from weights + KV growth + prefill
1156
2013
  // transient; the KV term assumes bf16 (a kv-quant scheme stretches the
1157
2014
  // real ceiling, never shrinks it — admission stays conservative).
1158
- const admission = fit(
2015
+ const genericAdmission = fit(
1159
2016
  ctx.model.config, ctx.model.weightsBytes, 1,
1160
2017
  undefined, undefined, 0, serverOptions.memoryBudgetBytes,
1161
2018
  );
2019
+ const admission = ctx.glmMemoryPlan
2020
+ ? {
2021
+ ...genericAdmission,
2022
+ maxSafeContext: ctx.glmMemoryPlan.contextTokens,
2023
+ usableBytes: ctx.glmMemoryPlan.processLimitBytes,
2024
+ }
2025
+ : genericAdmission;
1162
2026
  if (admission.maxSafeContext < 1)
1163
2027
  throw new Error(
1164
2028
  `memory budget ${(admission.usableBytes / 1e9).toFixed(2)} GB cannot serve ` +
1165
2029
  `${ctx.modelId} (weights ${(ctx.model.weightsBytes / 1e9).toFixed(2)} GB): ` +
1166
2030
  `no context fits — raise the budget or pick a smaller model`,
1167
2031
  );
1168
- if (serverOptions.memoryBudgetBytes) setMemoryLimit(serverOptions.memoryBudgetBytes);
2032
+ if (ctx.glmMemoryPlan)
2033
+ setMemoryLimit(ctx.glmMemoryPlan.lineItems.allocatorReserveBytes);
2034
+ else if (serverOptions.memoryBudgetBytes)
2035
+ setMemoryLimit(serverOptions.memoryBudgetBytes);
1169
2036
 
1170
2037
  // /library response cache (30 s) — registry + config reads only.
1171
2038
  let libraryCache: { at: number; rows: unknown[] } | null = null;
@@ -1239,7 +2106,8 @@ export function createServer(
1239
2106
  // truncates the visible reply. Default very generously (the model's
1240
2107
  // context is far larger); only an explicit max_tokens narrows it. The
1241
2108
  // model still stops at its eos_token well before this in normal replies.
1242
- maxTokens: req.max_completion_tokens ?? req.max_tokens ?? serverOptions.defaultMaxTokens ?? 65_536,
2109
+ maxTokens: req.max_completion_tokens ?? req.max_tokens ??
2110
+ defaultGeneratedTokens ?? 65_536,
1243
2111
  temperature: req.temperature ?? serverOptions.defaultTemperature ?? defaultTemp,
1244
2112
  topP: req.top_p ?? serverOptions.defaultTopP ?? ctx.genDefaults.topP ?? 0,
1245
2113
  topK: req.top_k ?? serverOptions.defaultTopK ?? ctx.genDefaults.topK ?? 0,
@@ -1265,9 +2133,35 @@ export function createServer(
1265
2133
  stopSequences: (typeof req.stop === "string" ? [req.stop] : req.stop ?? [])
1266
2134
  .filter((s) => typeof s === "string" && s.length > 0),
1267
2135
  ...kvScheme,
2136
+ // Paged KV is a server-wide mode like the kv scheme above (v1:
2137
+ // serial bf16 gemma4 — the createServer refusals hold the invariants).
2138
+ ...(serverOptions.pagedKv ? { pagedKv: serverOptions.pagedKv } : {}),
1268
2139
  };
1269
2140
  };
1270
2141
 
2142
+ /** Compile a grammar controller from the request's structured-output fields
2143
+ * (response_format / guided_* / structured_outputs). Returns null when no
2144
+ * constraint is requested OR when compile failed (degrade path). On degrade,
2145
+ * `degradeHint` carries a human description for the system-prompt injection
2146
+ * + Warning header (oMLX parity — never 500). Honors MLX_BUN_GRAMMAR=0. */
2147
+ const compileGrammarForRequest = async (
2148
+ req: ChatRequest,
2149
+ ): Promise<{ controller: import("./grammar").GrammarController | null; degradeHint: string | null }> => {
2150
+ if (!grammarEnabled()) {
2151
+ return {
2152
+ controller: null,
2153
+ degradeHint: "grammar compilation disabled by MLX_BUN_GRAMMAR=0",
2154
+ };
2155
+ }
2156
+ const r = await compileGrammarRequest(
2157
+ req as GrammarRequest,
2158
+ ctx.tokenizer,
2159
+ ctx.model.config.text.vocabSize,
2160
+ );
2161
+ if (!r) return { controller: null, degradeHint: null };
2162
+ return { controller: r.controller, degradeHint: r.degradeHint };
2163
+ };
2164
+
1271
2165
  const templateOptionsFor = (req: ChatRequest, tools: ToolDefinition[] | null) => {
1272
2166
  // enableThinking resolution (and its precedence) lives in
1273
2167
  // resolveEnableThinking so template rendering and sampling stay in sync.
@@ -1278,13 +2172,60 @@ export function createServer(
1278
2172
  req: ChatRequest,
1279
2173
  tools: ToolDefinition[] | null,
1280
2174
  ): { ids: number[]; startInThinking: boolean } => {
1281
- const rendered = ctx.template.render(normalizeMessages(req.messages), templateOptionsFor(req, tools));
2175
+ const opts = templateOptionsFor(req, tools);
2176
+ const rendered = ctx.template.render(normalizeMessages(req.messages), opts);
1282
2177
  const ids = ctx.tokenizer.encode(rendered);
1283
2178
  // template includes <bos>; tokenizer post-processor also prepends one
1284
2179
  const trimmed = ids[0] === ids[1] && ids[0] === ctx.tokenizer.bosTokenId ? ids.slice(1) : ids;
1285
2180
  return { ids: trimmed, startInThinking: promptEndsInOpenThink(rendered) };
1286
2181
  };
1287
2182
 
2183
+ /** STABLE cache boundary: the prompt's tail is a generation primer (e.g.
2184
+ * 12B's `<|channel>thought` tokens) that the NEXT turn's re-render does
2185
+ * not contain — a cache entry ending there can never prefix-match again
2186
+ * once the rings wrap. Probe by rendering this conversation as if a
2187
+ * reply existed; the common prefix is what every future turn preserves.
2188
+ * The probe is EXPENSIVE (a second full render + encode — ~150 ms at
2189
+ * 16k with our tokenizer; 30-50% of the 12B ctx-repeat TTFT,
2190
+ * 2026-07-06b), but the PRIMER LENGTH it measures is constant per
2191
+ * template mode: the divergence sits after a special-token turn
2192
+ * delimiter, and special tokens break BPE merges, so message content
2193
+ * can't shift it. So the full probe runs ONCE per mode and later
2194
+ * requests pay a subtraction. (A wrong boundary is a QUALITY knob, not
2195
+ * a correctness one — any prefix ≤ len-1 is a valid snapshot point.) */
2196
+ const primerLenByMode = new Map<string, number>();
2197
+ const stableLenFor = (
2198
+ req: ChatRequest,
2199
+ tools: ToolDefinition[] | null,
2200
+ trimmed: number[],
2201
+ ): number => {
2202
+ const opts = templateOptionsFor(req, tools);
2203
+ const mode = `${opts.enableThinking}|${!!tools?.length}`;
2204
+ const primer = primerLenByMode.get(mode);
2205
+ if (primer !== undefined) return Math.max(0, trimmed.length - primer);
2206
+ let stableLen = trimmed.length;
2207
+ try {
2208
+ const probe = ctx.tokenizer.encode(
2209
+ ctx.template.render(
2210
+ [...normalizeMessages(req.messages), { role: "assistant", content: "x" }],
2211
+ opts,
2212
+ ),
2213
+ );
2214
+ const probeTrimmed =
2215
+ probe[0] === probe[1] && probe[0] === ctx.tokenizer.bosTokenId ? probe.slice(1) : probe;
2216
+ let i = 0;
2217
+ const n = Math.min(trimmed.length, probeTrimmed.length);
2218
+ while (i < n && trimmed[i] === probeTrimmed[i]) i++;
2219
+ stableLen = i;
2220
+ // Memoize the mode's primer length (sanity-capped: a "primer" longer
2221
+ // than 64 tokens means the probe diverged for content reasons —
2222
+ // don't generalize that).
2223
+ const p = trimmed.length - stableLen;
2224
+ if (p >= 0 && p <= 64) primerLenByMode.set(mode, p);
2225
+ } catch { /* probe is best-effort; full boundary is the fallback */ }
2226
+ return stableLen;
2227
+ };
2228
+
1288
2229
  const toolStreamMode = (tools: ToolDefinition[] | null): ToolStreamMode =>
1289
2230
  // Gemma-4 keeps the token-level sentinel path; every other family
1290
2231
  // (MiniCPM5, Qwen3/3.5, Tier-0 generics) emits tool calls as DECODED TEXT
@@ -1294,9 +2235,18 @@ export function createServer(
1294
2235
  const toolRouter = (tools: ToolDefinition[] | null): ToolAwareStream =>
1295
2236
  new ToolAwareStream(ctx.tokenizer, toolStreamMode(tools), tools);
1296
2237
 
2238
+ if (serverOptions.unixSocket) {
2239
+ try { require("node:fs").unlinkSync(serverOptions.unixSocket); } catch {}
2240
+ }
2241
+ // Bun's types make unix/port mutually exclusive variants that a spread
2242
+ // union can't prove — runtime accepts either; cast once at the call.
1297
2243
  serverRef = Bun.serve({
1298
- port,
1299
- ...(serverOptions.hostname ? { hostname: serverOptions.hostname } : {}),
2244
+ ...(serverOptions.unixSocket
2245
+ ? ({ unix: serverOptions.unixSocket } as unknown as Record<string, never>)
2246
+ : {
2247
+ port,
2248
+ ...(serverOptions.hostname ? { hostname: serverOptions.hostname } : {}),
2249
+ }),
1300
2250
  idleTimeout: 0,
1301
2251
  // Web chat rides pi's AgentSession events over a WebSocket; the embedded
1302
2252
  // pi provider points back at THIS server's own loopback /v1 (port
@@ -1305,9 +2255,21 @@ export function createServer(
1305
2255
  port: () => serverRef.port ?? port,
1306
2256
  modelId: ctx.modelId,
1307
2257
  contextWindow: ctx.model.config.text.maxPositionEmbeddings,
1308
- // capability flag — true if a tower is loaded or loadable (lazy)
2258
+ // capability flags — true if a tower is loaded or loadable (lazy)
1309
2259
  vision: !!(ctx.vision || ctx.loadVision),
2260
+ audio: !!(ctx.audio || ctx.loadAudio),
1310
2261
  thinking: ctx.template.supportsThinking,
2262
+ // Model-author sampling defaults (generation_config.json, server-CLI
2263
+ // overrides applied) for the sampling popover's per-model recommended
2264
+ // values (web-ui-pass-plan.md #14 groundwork) — same resolution
2265
+ // /v1/models' gen_defaults uses, so the two surfaces agree. The
2266
+ // think-mode-vs-not distinction lives in toOptions (per-request); this
2267
+ // is the single "model's own defaults" snapshot sent once at ready.
2268
+ genDefaults: {
2269
+ temperature: serverOptions.defaultTemperature ?? ctx.genDefaults.temperature ?? null,
2270
+ topP: serverOptions.defaultTopP ?? ctx.genDefaults.topP ?? null,
2271
+ topK: serverOptions.defaultTopK ?? ctx.genDefaults.topK ?? null,
2272
+ },
1311
2273
  }),
1312
2274
  async fetch(request, server) {
1313
2275
  const url = new URL(request.url);
@@ -1318,6 +2280,81 @@ export function createServer(
1318
2280
  return new Response("expected websocket", { status: 426 });
1319
2281
  }
1320
2282
 
2283
+ // Memory REST wrappers (web-chat-redesign.md §2.3/§9 Phase 2): thin
2284
+ // loopback JSON routes over src/memory/vault.ts for the web chat's
2285
+ // Memory panel. Handlers live in src/memory/rest.ts (pure functions,
2286
+ // no `ctx` dependency) so they're unit-testable without a loaded
2287
+ // model; dispatch here just matches path+method. The agent-tool
2288
+ // surface (src/memory/tools.ts) is untouched by this block.
2289
+ if (url.pathname === "/api/memory/status" && request.method === "GET") {
2290
+ const { handleMemoryStatus } = await import("./memory/rest");
2291
+ return handleMemoryStatus();
2292
+ }
2293
+ if (url.pathname === "/api/memory/list" && request.method === "GET") {
2294
+ const { handleMemoryList } = await import("./memory/rest");
2295
+ return handleMemoryList();
2296
+ }
2297
+ if (url.pathname === "/api/memory/search" && request.method === "GET") {
2298
+ const { handleMemorySearch } = await import("./memory/rest");
2299
+ return handleMemorySearch(url);
2300
+ }
2301
+ if (url.pathname === "/api/memory/article" && request.method === "GET") {
2302
+ const { handleMemoryArticle } = await import("./memory/rest");
2303
+ return handleMemoryArticle(url);
2304
+ }
2305
+ if (url.pathname === "/api/memory/links" && request.method === "GET") {
2306
+ const { handleMemoryLinks } = await import("./memory/rest");
2307
+ return handleMemoryLinks(url);
2308
+ }
2309
+ if (url.pathname === "/api/memory/history" && request.method === "GET") {
2310
+ const { handleMemoryHistory } = await import("./memory/rest");
2311
+ return handleMemoryHistory(url);
2312
+ }
2313
+ if (url.pathname === "/api/memory/diff" && request.method === "GET") {
2314
+ const { handleMemoryDiff } = await import("./memory/rest");
2315
+ return handleMemoryDiff(url);
2316
+ }
2317
+ if (url.pathname === "/api/memory/init" && request.method === "POST") {
2318
+ const { handleMemoryInit } = await import("./memory/rest");
2319
+ return handleMemoryInit(request);
2320
+ }
2321
+
2322
+ // Model Hub (web-chat-redesign.md §9 Phase 3, beat-matrix Axis 3):
2323
+ // thin loopback JSON routes over the registry/download/fit machinery
2324
+ // for the web chat's Hub panel. Handlers live in src/hub-rest.ts
2325
+ // (pure functions, no `ctx` dependency); dispatch here just matches
2326
+ // path+method, same convention as the memory wrappers above.
2327
+ if (url.pathname === "/api/hub/local" && request.method === "GET") {
2328
+ const { handleHubLocal } = await import("./hub-rest");
2329
+ return handleHubLocal();
2330
+ }
2331
+ if (url.pathname === "/api/hub/search" && request.method === "GET") {
2332
+ const { handleHubSearch } = await import("./hub-rest");
2333
+ return handleHubSearch(url);
2334
+ }
2335
+ if (url.pathname === "/api/hub/download" && request.method === "POST") {
2336
+ const { handleHubDownload } = await import("./hub-rest");
2337
+ return handleHubDownload(request);
2338
+ }
2339
+ if (url.pathname === "/api/hub/serve" && request.method === "POST") {
2340
+ const { handleHubServe } = await import("./hub-rest");
2341
+ return handleHubServe(request);
2342
+ }
2343
+
2344
+ // Session full-text search + export (docs/design/web-chat-redesign.md
2345
+ // §9 Phase 3, beat-matrix Axis 10/11 — the full-text-search BEAT and
2346
+ // Markdown/JSON chat export). Handlers live in src/serve/session-search.ts
2347
+ // (pure, read-only, no `ctx` dependency), dispatched here by
2348
+ // path+method, same convention as the memory/hub blocks above.
2349
+ if (url.pathname === "/api/sessions/search" && request.method === "GET") {
2350
+ const { handleSessionsSearch } = await import("./serve/session-search");
2351
+ return handleSessionsSearch(url);
2352
+ }
2353
+ if (url.pathname === "/api/sessions/export" && request.method === "GET") {
2354
+ const { handleSessionsExport } = await import("./serve/session-search");
2355
+ return handleSessionsExport(url);
2356
+ }
2357
+
1321
2358
  // The unified SPA is served at "/"; legacy deep links redirect into
1322
2359
  // the hash router so old bookmarks still land on the right section.
1323
2360
  if (url.pathname === "/" && request.method === "GET") {
@@ -1325,6 +2362,50 @@ export function createServer(
1325
2362
  headers: { "content-type": "text/html; charset=utf-8" },
1326
2363
  });
1327
2364
  }
2365
+ // Vendored syntax-highlighting assets (src/web/vendor/*, no CDN — see
2366
+ // that dir's README). Long cache: content is versioned by the vendor
2367
+ // rebuild step, not by URL, but a hard reload always wins during dev.
2368
+ if (url.pathname === "/assets/hljs.js" && request.method === "GET") {
2369
+ return new Response(HLJS_JS, {
2370
+ headers: { "content-type": "text/javascript; charset=utf-8", "cache-control": "public, max-age=3600" },
2371
+ });
2372
+ }
2373
+ if (url.pathname === "/assets/hljs.css" && request.method === "GET") {
2374
+ return new Response(HLJS_CSS, {
2375
+ headers: { "content-type": "text/css; charset=utf-8", "cache-control": "public, max-age=3600" },
2376
+ });
2377
+ }
2378
+ // The frontend bundle (see the import comment above) — cache header
2379
+ // matches the other /assets/* entries; content is versioned by the
2380
+ // build step, not the URL, and a hard reload always wins during dev.
2381
+ if (url.pathname === "/assets/app.js" && request.method === "GET") {
2382
+ return new Response(APP_JS, {
2383
+ headers: { "content-type": "text/javascript; charset=utf-8", "cache-control": "public, max-age=3600" },
2384
+ });
2385
+ }
2386
+ // PWA manifest + icon (plan §9 Phase 3): linked from app.html's
2387
+ // <link rel="manifest">. Same versioned-by-content cache posture as
2388
+ // the other /assets/* entries above.
2389
+ if (url.pathname === "/manifest.webmanifest" && request.method === "GET") {
2390
+ return new Response(MANIFEST_WEBMANIFEST, {
2391
+ headers: { "content-type": "application/manifest+json; charset=utf-8", "cache-control": "public, max-age=3600" },
2392
+ });
2393
+ }
2394
+ if (url.pathname === "/assets/icon.svg" && request.method === "GET") {
2395
+ return new Response(ICON_SVG, {
2396
+ headers: { "content-type": "image/svg+xml; charset=utf-8", "cache-control": "public, max-age=3600" },
2397
+ });
2398
+ }
2399
+ // Service worker MUST be served at the root scope (/sw.js, not
2400
+ // /assets/sw.js) — a worker's default scope is the directory it's
2401
+ // served from, and the shell it manages includes "/" itself.
2402
+ // no-store: the browser's own update check needs a fresh byte
2403
+ // comparison every time to notice a new worker, not a cached 304.
2404
+ if (url.pathname === "/sw.js" && request.method === "GET") {
2405
+ return new Response(SW_JS, {
2406
+ headers: { "content-type": "text/javascript; charset=utf-8", "cache-control": "no-store" },
2407
+ });
2408
+ }
1328
2409
  // v2 HLG Curve Designer — served same-origin so /generate + /signal need no CORS.
1329
2410
  // Read fresh from disk in dev (edits show on reload, no restart); fall back to the
1330
2411
  // embedded copy when running as the compiled single binary.
@@ -1368,7 +2449,7 @@ export function createServer(
1368
2449
  // surface on their own. (scan() is INSERT-OR-REPLACE + prunes deleted,
1369
2450
  // so it's idempotent and cheap for a local cache.)
1370
2451
  await reg.scan();
1371
- const { visionCapable } = await import("./registry");
2452
+ const { visionCapable, audioCapable } = await import("./registry");
1372
2453
  const { supportTier } = await import("./model/support");
1373
2454
  const rows = [];
1374
2455
  // listCanonical: one row per repo (refs/main) — duplicate snapshots
@@ -1389,7 +2470,8 @@ export function createServer(
1389
2470
  rows.push({
1390
2471
  repo_id: m.repoId, model_type: m.modelType,
1391
2472
  size_bytes: m.sizeBytes, quant_bits: m.quantBits,
1392
- vision: visionCapable(m), supported, support_tier: tier,
2473
+ vision: visionCapable(m), audio: audioCapable(m),
2474
+ supported, support_tier: tier,
1393
2475
  serving: m.repoId === ctx.modelId,
1394
2476
  assessment,
1395
2477
  });
@@ -1449,16 +2531,65 @@ export function createServer(
1449
2531
  const machine = thisMachine();
1450
2532
  const chip = detectChip();
1451
2533
  let expertsBytes = 0;
2534
+ let artifactDiskBytes: number | null = null;
1452
2535
  let measured: { decodeTps: number; ts: number } | null = null;
1453
2536
  try {
1454
2537
  const { Registry } = await import("./registry");
1455
2538
  const rec = new Registry().list().find((r) => r.repoId === ctx.modelId);
1456
2539
  if (rec) {
1457
2540
  expertsBytes = rec.expertsBytes;
2541
+ artifactDiskBytes = rec.sizeBytes;
1458
2542
  const { EvalDB } = await import("./evaldb");
1459
2543
  measured = new EvalDB().latestFor(rec.path);
1460
2544
  }
1461
2545
  } catch {}
2546
+ if (ctx.glmMemoryPlan) {
2547
+ const plan = ctx.glmMemoryPlan;
2548
+ const li = plan.lineItems;
2549
+ const kvBytes = li.targetKvBytes + li.mtpKvBytes;
2550
+ const transientBytes = plan.plannedProcessBytes -
2551
+ li.residentWeightsBytes - li.mainExpertSlabBytes -
2552
+ li.mtpExpertSlabBytes - kvBytes;
2553
+ return Response.json({
2554
+ machine: {
2555
+ chip: chip.name,
2556
+ ram_bytes: machine.ramBytes,
2557
+ bandwidth_gbs: machine.bandwidthGBs,
2558
+ },
2559
+ context_tokens: plan.contextTokens,
2560
+ typical_context_tokens: plan.contextTokens,
2561
+ typical_decode_tps: GLM52_G5_MEASURED_WARM_DECODE_TPS,
2562
+ measured_decode_tps: GLM52_G5_MEASURED_WARM_DECODE_TPS,
2563
+ measured_at: GLM52_G5_MEASURED_AT,
2564
+ report: {
2565
+ fits: true,
2566
+ weights_bytes: li.residentWeightsBytes,
2567
+ kv_bytes: kvBytes,
2568
+ transient_bytes: transientBytes,
2569
+ total_bytes: plan.plannedProcessBytes,
2570
+ usable_bytes: plan.processLimitBytes,
2571
+ max_safe_context: plan.contextTokens,
2572
+ predicted_decode_tps: null,
2573
+ },
2574
+ glm52: {
2575
+ artifact_disk_bytes: artifactDiskBytes,
2576
+ main_expert_slab_bytes: li.mainExpertSlabBytes,
2577
+ mtp_expert_slab_bytes: li.mtpExpertSlabBytes,
2578
+ max_generation_tokens: plan.maxGenerationTokens,
2579
+ direct_oracle_warm_decode_tps:
2580
+ GLM52_G5_DIRECT_ORACLE_WARM_DECODE_TPS,
2581
+ aspirational_decode_tps: GLM52_G5_ASPIRATIONAL_DECODE_TPS,
2582
+ },
2583
+ sku_matrix_ctx: plan.contextTokens,
2584
+ sku_matrix: [{
2585
+ sku: chip.name,
2586
+ ram_gb: Math.round(machine.ramBytes / 2 ** 30),
2587
+ fits: true,
2588
+ max_context: plan.contextTokens,
2589
+ decode_tps: GLM52_G5_MEASURED_WARM_DECODE_TPS,
2590
+ }],
2591
+ });
2592
+ }
1462
2593
  const report = fit(
1463
2594
  ctx.model.config, ctx.model.weightsBytes, admission.maxSafeContext,
1464
2595
  machine, undefined, expertsBytes, serverOptions.memoryBudgetBytes,
@@ -1516,7 +2647,13 @@ export function createServer(
1516
2647
  const layerTypes = ctx.model.config.text.layerTypes;
1517
2648
  const kvLayers: Record<string, number> = {};
1518
2649
  let kvMode = "bf16";
1519
- if (kvScheme.kvBits) {
2650
+ if (kvScheme.turboQuant) {
2651
+ // v1: full-attention layers only (sliding-window stays bf16 —
2652
+ // docs/design/turboquant-kv.md non-goal).
2653
+ const fullAttn = layerTypes.filter((l) => l !== "sliding_attention").length;
2654
+ kvMode = `turbo k${kvScheme.turboQuant.kBits}v${kvScheme.turboQuant.vBits}`;
2655
+ kvLayers[`turbo-k${kvScheme.turboQuant.kBits}v${kvScheme.turboQuant.vBits}`] = fullAttn;
2656
+ } else if (kvScheme.kvBits) {
1520
2657
  kvMode = `uniform-kv${kvScheme.kvBits}`;
1521
2658
  kvLayers[`kv${kvScheme.kvBits}`] = layerTypes.length;
1522
2659
  } else if (kvScheme.kvConfig) {
@@ -1548,6 +2685,13 @@ export function createServer(
1548
2685
  restores: ssdStore.stats.restores,
1549
2686
  spills: ssdStore.stats.spills,
1550
2687
  restore_ms_last: Math.round(ssdStore.stats.restoreMsLast),
2688
+ demotions: promptCache.demotions,
2689
+ // Bounded write-behind queue (2026-07-07): clones waiting for
2690
+ // the idle-gated flush, and spills the byte cap dropped
2691
+ // (each = one future cache miss, never a wrong result).
2692
+ pending_spills: spillQueue?.pendingCount ?? 0,
2693
+ pending_spill_bytes: spillQueue?.pendingBytes ?? 0,
2694
+ dropped_spills: spillQueue?.droppedCount ?? 0,
1551
2695
  },
1552
2696
  } : {}),
1553
2697
  response_store: {
@@ -1566,22 +2710,71 @@ export function createServer(
1566
2710
  },
1567
2711
  admission: {
1568
2712
  max_safe_context: admission.maxSafeContext,
1569
- memory_budget_bytes: serverOptions.memoryBudgetBytes ?? null,
2713
+ memory_budget_bytes:
2714
+ ctx.glmMemoryPlan?.processLimitBytes ??
2715
+ serverOptions.memoryBudgetBytes ?? null,
1570
2716
  usable_bytes: admission.usableBytes,
1571
2717
  weights_bytes: ctx.model.weightsBytes,
1572
2718
  },
2719
+ ...(ctx.glmMemoryPlan ? {
2720
+ glm52: {
2721
+ preset: ctx.glmMemoryPlan.preset,
2722
+ planned_process_bytes: ctx.glmMemoryPlan.plannedProcessBytes,
2723
+ process_limit_bytes: ctx.glmMemoryPlan.processLimitBytes,
2724
+ context_tokens: ctx.glmMemoryPlan.contextTokens,
2725
+ max_generation_tokens: ctx.glmMemoryPlan.maxGenerationTokens,
2726
+ batch_size: ctx.glmMemoryPlan.batchSize,
2727
+ dsa: ctx.model instanceof Glm52Model && ctx.model.capabilities.dsa,
2728
+ mtp: ctx.draft?.provider.id === "glm52-native-mtp",
2729
+ mtp_draft_tokens: ctx.glmMemoryPlan.mtpDraftTokens,
2730
+ resident_weight_bytes: ctx.glmMemoryPlan.lineItems.residentWeightsBytes,
2731
+ main_expert_slab_bytes: ctx.glmMemoryPlan.lineItems.mainExpertSlabBytes,
2732
+ mtp_expert_slab_bytes: ctx.glmMemoryPlan.lineItems.mtpExpertSlabBytes,
2733
+ expert_runtime: ctx.model instanceof Glm52Model &&
2734
+ ctx.model.expertRuntime
2735
+ ? {
2736
+ main_residency:
2737
+ ctx.model.expertRuntime.manager.snapshot(),
2738
+ mtp_residency:
2739
+ ctx.model.expertRuntime.mtp?.manager.snapshot() ?? null,
2740
+ last_turn: ctx.model.expertRuntime.lastTelemetry,
2741
+ last_repin: ctx.model.expertRuntime.lastRepin,
2742
+ }
2743
+ : null,
2744
+ },
2745
+ } : {}),
1573
2746
  // --batch: configured cap, whether batching is live for this model,
1574
2747
  // and rows currently decoding in the batch.
1575
2748
  batch: {
1576
2749
  configured: batch,
2750
+ mode: gateway.batchMode,
1577
2751
  batched: gateway.batchingEnabled,
1578
2752
  active_rows: gateway.activeRows,
2753
+ pending_rows: gateway.pendingRows,
2754
+ submitted_rows: gateway.submittedRows,
2755
+ kv_bytes: gateway.kvBytes.projected,
2756
+ kv_budget_bytes: gateway.kvBytes.budget,
1579
2757
  },
1580
2758
  });
1581
2759
  }
1582
2760
 
1583
2761
  // mlx_lm.server parity: GET /health → the exact body it writes
1584
2762
  // ('{"status": "ok"}', note the space) so byte-for-byte health checks pass.
2763
+ // Engine-mode admin (unix-socket children only — never exposed on
2764
+ // TCP): drain = quiesce the gateway + demote the whole prompt cache
2765
+ // to the SSD tier. The pool calls this before evicting a model
2766
+ // child, so its state survives the eviction losslessly.
2767
+ if (
2768
+ serverOptions.unixSocket &&
2769
+ url.pathname === "/admin/drain" &&
2770
+ request.method === "POST"
2771
+ ) {
2772
+ await gateway.runExclusive(async () => {
2773
+ promptCache.demoteIdle(0);
2774
+ });
2775
+ return Response.json({ drained: true, demotions: promptCache.demotions });
2776
+ }
2777
+
1585
2778
  if (url.pathname === "/health" && request.method === "GET") {
1586
2779
  return new Response('{"status": "ok"}', {
1587
2780
  headers: { "content-type": "application/json" },
@@ -1601,26 +2794,68 @@ export function createServer(
1601
2794
  ? decodeURIComponent(url.pathname.slice("/v1/models/".length))
1602
2795
  : null;
1603
2796
  const created = Math.floor(startedAt / 1000);
2797
+ // Model-author sampling defaults (generation_config.json), resolved once
2798
+ // at load (web-ui-pass-plan.md #14 groundwork): lets a client (the web
2799
+ // chat's sampling popover, external pi) show the SERVED model's actual
2800
+ // recommended values instead of one hardcoded shape.
2801
+ const genDefaultsOut = {
2802
+ temperature: ctx.genDefaults.temperature ?? null,
2803
+ top_p: ctx.genDefaults.topP ?? null,
2804
+ top_k: ctx.genDefaults.topK ?? null,
2805
+ };
1604
2806
  const data: Array<Record<string, unknown>> = [{
1605
2807
  id: ctx.modelId, object: "model", created, owned_by: "mlx-bun",
1606
- context_window: ctx.model.config.text.maxPositionEmbeddings,
2808
+ context_window:
2809
+ ctx.glmMemoryPlan?.contextTokens ??
2810
+ ctx.model.config.text.maxPositionEmbeddings,
1607
2811
  // Capability flags for clients (CLI/external pi) that build a
1608
2812
  // 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).
2813
+ // `vision` the image input declaration, `audio` whether
2814
+ // `input_audio` content parts are accepted (tower loaded or lazily
2815
+ // loadable; same signals the web embed uses).
1611
2816
  reasoning: ctx.template.supportsThinking,
1612
2817
  vision: !!(ctx.vision || ctx.loadVision),
2818
+ audio: !!(ctx.audio || ctx.loadAudio),
2819
+ batch_mode: gateway.batchMode,
2820
+ tools: true,
2821
+ structured_output: true,
2822
+ embeddings: isEmbeddingModel(ctx.model),
2823
+ adapters: !(ctx.model instanceof Glm52Model),
2824
+ training: !(ctx.model instanceof Glm52Model),
2825
+ dsa: ctx.model instanceof Glm52Model && ctx.model.capabilities.dsa,
2826
+ mtp: ctx.draft?.provider.id === "glm52-native-mtp",
2827
+ capabilities: {
2828
+ chat_completions: true,
2829
+ text_completions: true,
2830
+ anthropic_messages: true,
2831
+ responses: true,
2832
+ streaming: true,
2833
+ tools: true,
2834
+ structured_output: true,
2835
+ logprobs: true,
2836
+ embeddings: isEmbeddingModel(ctx.model),
2837
+ vision: !!(ctx.vision || ctx.loadVision),
2838
+ audio: !!(ctx.audio || ctx.loadAudio),
2839
+ adapters: !(ctx.model instanceof Glm52Model),
2840
+ training: !(ctx.model instanceof Glm52Model),
2841
+ },
2842
+ gen_defaults: genDefaultsOut,
1613
2843
  }];
1614
2844
  try {
1615
- const { Registry } = await import("./registry");
2845
+ const { Registry, visionCapable } = await import("./registry");
2846
+ const { supportTier } = await import("./model/support");
1616
2847
  const reg = new Registry();
1617
2848
  try {
1618
2849
  if (reg.list().length === 0) await reg.scan();
1619
2850
  // Canonical rows only — a repo with N snapshots is ONE model id.
1620
2851
  for (const m of reg.listCanonical()) {
1621
2852
  if (m.repoId === ctx.modelId) continue;
1622
- if (!isSupportedModelRecord(m.modelType, m.repoId)) continue;
1623
- data.push({ id: m.repoId, object: "model", created });
2853
+ const tier = supportTier(m.modelType, m.repoId);
2854
+ if (tier === null) continue;
2855
+ data.push({
2856
+ id: m.repoId, object: "model", created,
2857
+ vision: visionCapable(m), tier,
2858
+ });
1624
2859
  }
1625
2860
  } finally {
1626
2861
  reg.close();
@@ -1676,13 +2911,16 @@ export function createServer(
1676
2911
  // unmount. Mount and unmount go through the generation queue so
1677
2912
  // they never race an in-flight forward pass.
1678
2913
  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.
2914
+ // On-disk adapters that can be mounted (the chat selector's source).
2915
+ // Every on-disk adapter is returned — the web chat routing chip (§5.2)
2916
+ // needs incompatible entries visible-but-grayed, not silently dropped,
2917
+ // so a user understands "10 adapters exist, 2 apply here" rather than
2918
+ // seeing a shorter list with no explanation. `compatible` = the
2919
+ // adapter's recorded base repo id matches the served model (compared
2920
+ // on the bare name, lenient about the org); an adapter with no
2921
+ // recorded base is treated as compatible (mount validates it for
2922
+ // real). `mounted` flags ones already loaded so the UI auto-loads on
2923
+ // select only if needed.
1686
2924
  const { homedir } = await import("node:os");
1687
2925
  const stores = [
1688
2926
  `${homedir()}/.cache/mlx-bun-finetunes`,
@@ -1692,10 +2930,10 @@ export function createServer(
1692
2930
  const bareName = (s: string) => s.split("/").pop()!.toLowerCase();
1693
2931
  const servedName = bareName(ctx.modelId);
1694
2932
  const adapters = (await listAvailableAdapters(stores))
1695
- .filter((a) => a.baseModel == null || bareName(a.baseModel) === servedName)
1696
2933
  .map((a) => ({
1697
2934
  id: a.id, path: a.path, rank: a.rank, scale: a.scale,
1698
2935
  base_model: a.baseModel, mounted: mounted.has(a.id),
2936
+ compatible: a.baseModel == null || bareName(a.baseModel) === servedName,
1699
2937
  }));
1700
2938
  return Response.json({ adapters });
1701
2939
  }
@@ -1704,6 +2942,7 @@ export function createServer(
1704
2942
  adapters: ctx.adapters.list().map((a) => ({
1705
2943
  id: a.id, path: a.path, rank: a.rank, scale: a.scale,
1706
2944
  size_bytes: a.sizeBytes, mounted_layers: a.mountedLayers,
2945
+ ram_bytes: a.ramBytes,
1707
2946
  })),
1708
2947
  });
1709
2948
  }
@@ -1723,7 +2962,7 @@ export function createServer(
1723
2962
  const info = await gateway.runExclusive(() => ctx.adapters.mount(body.id!, body.path!));
1724
2963
  return Response.json({
1725
2964
  id: info.id, mounted_layers: info.mountedLayers,
1726
- rank: info.rank, scale: info.scale,
2965
+ rank: info.rank, scale: info.scale, ram_bytes: info.ramBytes,
1727
2966
  });
1728
2967
  } catch (e) {
1729
2968
  return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
@@ -1825,7 +3064,7 @@ export function createServer(
1825
3064
  // translates its body into this shape and the Response back —
1826
3065
  // generation, tools, vision, stop sequences, prompt cache, and
1827
3066
  // admission control all live here exactly once.
1828
- const handleChat = async (body: ChatRequest): Promise<Response> => {
3067
+ const handleChat = async (body: ChatRequest, signal: AbortSignal): Promise<Response> => {
1829
3068
  if (!Array.isArray(body.messages) || body.messages.length === 0)
1830
3069
  return Response.json({ error: { message: "messages required" } }, { status: 400 });
1831
3070
  // mlx-lm validates logprobs params up front (ValueError → 400)
@@ -1841,29 +3080,135 @@ export function createServer(
1841
3080
  (m) => Array.isArray(m.content) &&
1842
3081
  m.content.some((p: any) => p.type === "image_url" || p.type === "image"),
1843
3082
  );
3083
+ // Same shapes extractAudio accepts: OpenAI-canonical input_audio plus
3084
+ // optiq's audio / audio_url aliases (audio-input-plan.md §3.2).
3085
+ const hasAudio = body.messages.some(
3086
+ (m) => Array.isArray(m.content) &&
3087
+ m.content.some(
3088
+ (p: any) => p.type === "input_audio" || p.type === "audio" || p.type === "audio_url",
3089
+ ),
3090
+ );
1844
3091
  let promptIds: number[];
3092
+ let stableLen: number | null = null;
1845
3093
  // Whether the prompt primed an open <think> (Qwen3.5/MiniCPM5 thinking
1846
3094
  // on) so the model's output starts mid-reasoning — seeds the splitter.
1847
3095
  // Vision is Gemma4-only (no switchable thinking channel), so it's false.
1848
3096
  let startInThinking = false;
1849
3097
  let vision: Parameters<typeof runGeneration>[3];
1850
3098
  let diffusionPixels: import("./mlx/array").MlxArray | null = null;
3099
+ // Grammar-constrained decoding (src/grammar.ts). Compile BEFORE prompt
3100
+ // rendering: on the degrade path (compile failed but a constraint was
3101
+ // requested) inject a system message instructing valid JSON so the
3102
+ // model still best-efforts schema-conformant output (oMLX parity —
3103
+ // _compile_grammar_for_request returning None + build_json_system_prompt).
3104
+ let grammarCtrl: import("./grammar").GrammarController | null = null;
3105
+ let grammarWarning: string | null = null;
3106
+ const grammarReq =
3107
+ body.response_format != null || !!body.guided_grammar ||
3108
+ !!body.guided_regex || !!body.guided_choice?.length ||
3109
+ body.structured_outputs != null;
3110
+ if (grammarReq) {
3111
+ const g = await compileGrammarForRequest(body);
3112
+ grammarCtrl = g.controller;
3113
+ if (!g.controller && g.degradeHint) {
3114
+ const degraded = applyGrammarDegrade(body, g.degradeHint);
3115
+ grammarWarning = degraded.warning;
3116
+ body = degraded.body;
3117
+ }
3118
+ }
3119
+ // Free what a rejected request has already allocated (grammar WASM
3120
+ // matcher/compiled-grammar, vision embeddings, diffusion pixels).
3121
+ // Every early return between here and gateway.run() must call this —
3122
+ // generate()/the gateway only take disposal ownership once running
3123
+ // (leaked one matcher per SSRF-blocked-media + response_format
3124
+ // request before the 2026-07-07 sweep).
3125
+ const disposeRejected = () => {
3126
+ grammarCtrl?.dispose();
3127
+ vision?.embeddings.dispose();
3128
+ vision?.imageMask?.dispose();
3129
+ vision?.multimodalMask?.dispose();
3130
+ diffusionPixels?.dispose();
3131
+ };
1851
3132
  try {
1852
- if (hasImages && ctx.model instanceof DiffusionGemmaModel) {
3133
+ if (hasAudio) {
3134
+ // Audio (and MIXED image+audio) input — A4 of
3135
+ // docs/design/audio-input-plan.md. One buildMultimodalPrompt call
3136
+ // splices both media kinds in document order. A request WITH
3137
+ // audio on a model whose tower is unavailable is an explicit 400
3138
+ // — never a silent text-only degrade (that leniency is only for
3139
+ // requests without the media, getVisionTower's contract).
3140
+ const audioTower = getAudioTower(ctx);
3141
+ if (!audioTower || !ctx.audioTokenIds) {
3142
+ disposeRejected();
3143
+ return Response.json(
3144
+ {
3145
+ error: {
3146
+ message:
3147
+ `model ${ctx.modelId} has no audio tower — audio input needs ` +
3148
+ `a model whose config.json carries audio_config and whose ` +
3149
+ `sidecar ships the audio tensors (e.g. gemma-4 e4b OptiQ)`,
3150
+ },
3151
+ },
3152
+ { status: 400 },
3153
+ );
3154
+ }
3155
+ let visionSide:
3156
+ | { tower: VisionEncoder; tokenIds: VisionTokenIds }
3157
+ | undefined;
3158
+ if (hasImages) {
3159
+ const tower = getVisionTower(ctx);
3160
+ if (!tower) {
3161
+ disposeRejected();
3162
+ return Response.json(
3163
+ { error: { message: "model has no vision sidecar" } }, { status: 400 },
3164
+ );
3165
+ }
3166
+ visionSide = { tower, tokenIds: ctx.visionTokenIds };
3167
+ }
3168
+ const { messages: withAudioParts, images } =
3169
+ await extractImages(normalizeMessages(body.messages));
3170
+ const { messages, audio } = await extractAudio(withAudioParts);
3171
+ // Non-WAV containers (mp3/m4a/flac/ogg/aiff/…) transcode to
3172
+ // 16 kHz WAV via CoreAudio; RIFF bytes pass through untouched.
3173
+ // Failures throw into the prompt-build 400 below.
3174
+ const wavs = await Promise.all(audio.map(ensureWav));
3175
+ // The towers are only ever non-null for Gemma4 (loader gates).
3176
+ const mp = await buildMultimodalPrompt(
3177
+ ctx.model as Gemma4Model,
3178
+ {
3179
+ ...(visionSide ? { vision: visionSide } : {}),
3180
+ audio: { tower: audioTower, tokenIds: ctx.audioTokenIds },
3181
+ },
3182
+ ctx.tokenizer, ctx.template, messages, images, wavs, tools,
3183
+ );
3184
+ promptIds = mp.ids;
3185
+ // bidirMask is null whenever audio is present (§3.3 Q1: mixed
3186
+ // prompts run fully causal); the union mask does the per-layer
3187
+ // id zeroing either way.
3188
+ vision = {
3189
+ embeddings: mp.embeddings,
3190
+ ...(mp.bidirMask ? { imageMask: mp.bidirMask } : {}),
3191
+ multimodalMask: mp.multimodalMask,
3192
+ };
3193
+ } else if (hasImages && ctx.model instanceof DiffusionGemmaModel) {
1853
3194
  // DiffusionGemma image-text-to-text: its OWN dedicated SigLIP tower +
1854
3195
  // encoder vision merge feed the denoising engine (NOT the AR
1855
3196
  // forwardEmbeddings path). v1 supports a single image.
1856
3197
  const dm = ctx.model;
1857
- if (!dm.visionTower)
3198
+ if (!dm.visionTower) {
3199
+ disposeRejected();
1858
3200
  return Response.json(
1859
3201
  { error: { message: "this checkpoint has no vision tower" } }, { status: 400 },
1860
3202
  );
3203
+ }
1861
3204
  const { messages, images } = await extractImages(normalizeMessages(body.messages));
1862
- if (images.length !== 1)
3205
+ if (images.length !== 1) {
3206
+ disposeRejected();
1863
3207
  return Response.json(
1864
3208
  { error: { message: "DiffusionGemma image input supports exactly one image" } },
1865
3209
  { status: 400 },
1866
3210
  );
3211
+ }
1867
3212
  const rendered = ctx.template.render(messages, { tools, addGenerationPrompt: true });
1868
3213
  const rawIds = ctx.tokenizer.encode(rendered, /* addSpecialTokens */ false);
1869
3214
  const { pixels, softTokens } = await dm.visionTower.preprocess(images[0]!);
@@ -1877,10 +3222,12 @@ export function createServer(
1877
3222
  // Loads (and caches) the tower on first image request — text-only
1878
3223
  // sessions never pay for it.
1879
3224
  const tower = getVisionTower(ctx);
1880
- if (!tower)
3225
+ if (!tower) {
3226
+ disposeRejected();
1881
3227
  return Response.json(
1882
3228
  { error: { message: "model has no vision sidecar" } }, { status: 400 },
1883
3229
  );
3230
+ }
1884
3231
  const { messages, images } = await extractImages(normalizeMessages(body.messages));
1885
3232
  // The tower is only ever non-null for Gemma4 (sidecar gate in
1886
3233
  // makeVisionLoader), so the model narrow is safe here.
@@ -1894,8 +3241,10 @@ export function createServer(
1894
3241
  const built = promptIdsFor(body, tools);
1895
3242
  promptIds = built.ids;
1896
3243
  startInThinking = built.startInThinking;
3244
+ stableLen = -1; // marker: chat path, probe lazily below
1897
3245
  }
1898
3246
  } catch (e) {
3247
+ disposeRejected();
1899
3248
  return Response.json(
1900
3249
  { error: { message: `prompt build failed: ${(e as Error).message}` } },
1901
3250
  { status: 400 },
@@ -1904,22 +3253,33 @@ export function createServer(
1904
3253
  let options: ReturnType<typeof toOptions>;
1905
3254
  try {
1906
3255
  options = toOptions(body);
3256
+ // Stable cache boundary for the prompt-boundary snapshot (chat
3257
+ // prompts end in a generation primer the next turn won't contain).
3258
+ // The probe costs a second full render+encode (~150 ms at 16k), so
3259
+ // run it ONLY when a snapshot can actually be taken: the cache
3260
+ // must not already hold the strict prefix (warm repeats peek at
3261
+ // len-1 → skip; the snapshot gate needs boundary > cached prefix).
3262
+ if (stableLen !== null) {
3263
+ const cachePeek = promptCache.peekPrefixLen(
3264
+ promptIds, options.adapters?.join("+") ?? "");
3265
+ if (cachePeek < promptIds.length - 1)
3266
+ options.snapshotAt = stableLenFor(body, tools, promptIds);
3267
+ }
1907
3268
  } catch (e) {
1908
3269
  // bad logit_bias (non-numeric keys/values) — mlx-lm's coercion error
1909
- vision?.embeddings.dispose();
1910
- vision?.imageMask.dispose();
1911
- diffusionPixels?.dispose();
3270
+ disposeRejected();
1912
3271
  return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
1913
3272
  }
1914
3273
  if (diffusionPixels) options.visionPixels = diffusionPixels;
3274
+ // Attach the compiled grammar controller (null when no constraint /
3275
+ // degrade — generate() runs the unmasked fast pipelined loop).
3276
+ if (grammarCtrl) options.grammar = grammarCtrl;
1915
3277
  // Admission: reject what cannot finish within the memory budget
1916
3278
  // (the GPU OOM it would otherwise hit is uncatchable and kills
1917
3279
  // the process — Phase 6 finding).
1918
3280
  const requiredCtx = promptIds.length + (options.maxTokens ?? 1024);
1919
3281
  if (requiredCtx > admission.maxSafeContext) {
1920
- vision?.embeddings.dispose();
1921
- vision?.imageMask.dispose();
1922
- diffusionPixels?.dispose();
3282
+ disposeRejected();
1923
3283
  return Response.json(
1924
3284
  {
1925
3285
  error: {
@@ -1941,6 +3301,7 @@ export function createServer(
1941
3301
  const adapterIds = ctx.adapters.resolveSpec(body.adapter ?? serverOptions.defaultAdapter);
1942
3302
  if (adapterIds.length) options.adapters = adapterIds;
1943
3303
  } catch (e) {
3304
+ disposeRejected();
1944
3305
  return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
1945
3306
  }
1946
3307
 
@@ -1958,14 +3319,18 @@ export function createServer(
1958
3319
  options.topLogprobs = topLogprobs;
1959
3320
  }
1960
3321
 
1961
- // What lane this request takes (vision / adapters / logits processors /
1962
- // a user-fixed seed → serial; everything else batches when --batch N).
3322
+ // What lane this request takes (vision/audio media / adapters /
3323
+ // logprobs / seed / explicit kv-quant / a mounted draft → serial;
3324
+ // sampler extras, repetition penalty, and grammar all BATCH — see
3325
+ // willBatch). `vision` is the embeddings-prefill payload for BOTH
3326
+ // media kinds, so hasVision routes audio serial too — batching is a
3327
+ // mode, not a fallback, and batched audio has no oracle (§3.2).
1963
3328
  const shape = {
1964
3329
  hasVision: !!vision,
1965
3330
  hasAdapters: !!options.adapters?.length,
1966
3331
  hasRepetitionPenalty: !!options.repetitionPenalty,
1967
- // min_p / XTC / logit_bias / presence+frequency penalties: serial-only
1968
- // in v1 (see GenerationGateway.RequestShape.hasLogitsExtras).
3332
+ // Informational since 2026-07-02: per-row logits processors batch;
3333
+ // willBatch no longer gates on these fields.
1969
3334
  hasLogitsExtras: !!(
1970
3335
  options.minP || options.xtcProbability || options.logitBias ||
1971
3336
  options.presencePenalty || options.frequencyPenalty
@@ -1973,22 +3338,44 @@ export function createServer(
1973
3338
  wantsLogprobs: captureLogprobs,
1974
3339
  userSeed: body.seed !== undefined,
1975
3340
  kvQuant: !!(options.kvConfig?.length || options.kvBits),
3341
+ turboQuant: !!options.turboQuant,
3342
+ // grammarCtrl is null on the degrade path (prompt injection) —
3343
+ // those stay batchable. A real controller batches via per-row
3344
+ // matchers (MLX_BUN_GRAMMAR_BATCH=0 forces it serial).
3345
+ hasGrammar: !!grammarCtrl,
3346
+ hasDraft: !!ctx.draft,
1976
3347
  };
1977
3348
  const batched = gateway.willBatch(shape);
1978
3349
  if (process.env.MLX_BUN_LANE_DEBUG === "1")
1979
3350
  console.error(`[lane] batched=${batched} shape=${JSON.stringify(shape)} t=${Date.now() % 100000}`);
3351
+ // Per-turn lane (docs/design/web-chat-redesign.md §2.3 caveat / risk #5):
3352
+ // reported on usage AND recorded in the in-process lane registry (keyed
3353
+ // by this response's `id`) so pi-web's WS bridge — which the pi SDK's
3354
+ // usage parsing can't carry a custom field through — can correlate it
3355
+ // via AssistantMessage.responseId. batched wins over hasDraft (mutually
3356
+ // exclusive by construction: willBatch already excludes hasDraft), so
3357
+ // the only ambiguity is serial vs serial+spec, resolved once s.spec is
3358
+ // known post-generation.
3359
+ const lane: Lane = batched ? "batched" : shape.hasDraft ? "serial+spec" : "serial";
3360
+ recordLane(id, lane);
1980
3361
 
1981
3362
  if (body.stream) {
3363
+ const streamAbort = new AbortController();
3364
+ const generationSignal = AbortSignal.any([signal, streamAbort.signal]);
3365
+ let cancelled = false;
1982
3366
  const stream = new ReadableStream<Uint8Array>({
1983
- async start(controller) {
3367
+ start(controller) {
1984
3368
  const enc = new TextEncoder();
1985
- const send = (obj: unknown) =>
1986
- controller.enqueue(enc.encode(`data: ${JSON.stringify(obj)}\n\n`));
3369
+ const send = (obj: unknown) => {
3370
+ if (!generationSignal.aborted)
3371
+ controller.enqueue(enc.encode(`data: ${JSON.stringify(obj)}\n\n`));
3372
+ };
1987
3373
  const chunk = (delta: Record<string, unknown>, finish: string | null) => ({
1988
3374
  id, object: "chat.completion.chunk", created, model: ctx.modelId,
1989
3375
  choices: [{ index: 0, delta, finish_reason: finish }],
1990
3376
  });
1991
- try {
3377
+ void (async () => {
3378
+ try {
1992
3379
  // The gateway owns lane selection + GPU exclusivity; this body
1993
3380
  // runs per-request (concurrently in batched mode, each writing
1994
3381
  // its own SSE stream — the per-row fan-out).
@@ -2023,7 +3410,8 @@ export function createServer(
2023
3410
  return new Promise<void>((r) => setImmediate(r));
2024
3411
  }
2025
3412
  }
2026
- }, vision, shape);
3413
+ }, vision, shape, generationSignal);
3414
+ if (generationSignal.aborted) return;
2027
3415
  // a stop match discards everything from the match on,
2028
3416
  // including text still held by the decoders
2029
3417
  let tail = "";
@@ -2054,6 +3442,11 @@ export function createServer(
2054
3442
  ? "tool_calls"
2055
3443
  : stopper.stopped ? "stop"
2056
3444
  : s.generatedTokens >= (options.maxTokens ?? 1024) ? "length" : "stop";
3445
+ // Refine serial vs serial+spec now that s.spec is known (a mounted
3446
+ // draft can still decode zero speculative tokens on a very short
3447
+ // reply) and re-record for pi-web's WS correlation.
3448
+ const finalLane: Lane = batched ? "batched" : s.spec ? "serial+spec" : lane;
3449
+ if (finalLane !== lane) recordLane(id, finalLane);
2057
3450
  send({
2058
3451
  ...chunk({}, finish),
2059
3452
  usage: {
@@ -2061,16 +3454,27 @@ export function createServer(
2061
3454
  completion_tokens: s.generatedTokens,
2062
3455
  total_tokens: s.promptTokens + s.generatedTokens,
2063
3456
  prompt_tokens_details: { cached_tokens: s.cachedTokens },
3457
+ lane: finalLane,
3458
+ ...(s.spec ? { speculation: s.spec } : {}),
2064
3459
  },
2065
3460
  });
2066
3461
  // bare sentinel per the OpenAI spec — JSON.stringify would
2067
3462
  // quote it and strict SDK clients never see the terminator
2068
- controller.enqueue(enc.encode("data: [DONE]\n\n"));
2069
- } catch (e) {
2070
- send({ error: { message: (e as Error).message } });
2071
- } finally {
2072
- controller.close();
2073
- }
3463
+ controller.enqueue(enc.encode("data: [DONE]\n\n"));
3464
+ } catch (e) {
3465
+ if (!generationSignal.aborted)
3466
+ send({ error: { message: (e as Error).message } });
3467
+ } finally {
3468
+ if (!cancelled) {
3469
+ if (generationSignal.aborted) controller.error(generationSignal.reason);
3470
+ else controller.close();
3471
+ }
3472
+ }
3473
+ })();
3474
+ },
3475
+ cancel(reason) {
3476
+ cancelled = true;
3477
+ streamAbort.abort(reason);
2074
3478
  },
2075
3479
  });
2076
3480
  return new Response(stream, {
@@ -2078,6 +3482,7 @@ export function createServer(
2078
3482
  "content-type": "text/event-stream",
2079
3483
  "cache-control": "no-cache",
2080
3484
  connection: "keep-alive",
3485
+ ...(grammarWarning ? { Warning: grammarWarning } : {}),
2081
3486
  },
2082
3487
  });
2083
3488
  }
@@ -2102,7 +3507,7 @@ export function createServer(
2102
3507
  content += parts.content;
2103
3508
  reasoning += parts.reasoning;
2104
3509
  if (stopper.stopped) return false; // halt generation
2105
- }, vision, shape);
3510
+ }, vision, shape, signal);
2106
3511
  if (!stopper.stopped) {
2107
3512
  const flushed = router.flush();
2108
3513
  reasoning += router.takeReasoning();
@@ -2123,6 +3528,10 @@ export function createServer(
2123
3528
  : stopper.stopped ? "stop"
2124
3529
  : s.generatedTokens >= (options.maxTokens ?? 1024) ? "length" : "stop";
2125
3530
  const logprobsBlock = lpc?.payload() ?? null;
3531
+ // Refine serial vs serial+spec now that s.spec is known; re-record
3532
+ // for pi-web's WS correlation (see the streaming branch above).
3533
+ const finalLane: Lane = batched ? "batched" : s.spec ? "serial+spec" : lane;
3534
+ if (finalLane !== lane) recordLane(id, finalLane);
2126
3535
  return Response.json({
2127
3536
  id, object: "chat.completion", created, model: ctx.modelId,
2128
3537
  choices: [{
@@ -2141,8 +3550,10 @@ export function createServer(
2141
3550
  completion_tokens: s.generatedTokens,
2142
3551
  total_tokens: s.promptTokens + s.generatedTokens,
2143
3552
  prompt_tokens_details: { cached_tokens: s.cachedTokens },
3553
+ lane: finalLane,
3554
+ ...(s.spec ? { speculation: s.spec } : {}),
2144
3555
  },
2145
- });
3556
+ }, grammarWarning ? { headers: { Warning: grammarWarning } } : undefined);
2146
3557
  }
2147
3558
  } catch (e) {
2148
3559
  return Response.json({ error: { message: (e as Error).message } }, { status: 500 });
@@ -2156,7 +3567,7 @@ export function createServer(
2156
3567
  } catch {
2157
3568
  return Response.json({ error: { message: "invalid JSON body" } }, { status: 400 });
2158
3569
  }
2159
- return handleChat(body);
3570
+ return handleChat(body, request.signal);
2160
3571
  }
2161
3572
 
2162
3573
  // Raw text completion (mlx_lm.server's /v1/completions, request_type
@@ -2196,12 +3607,36 @@ export function createServer(
2196
3607
  // bad logit_bias (non-numeric keys/values) — mlx-lm's coercion error
2197
3608
  return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
2198
3609
  }
3610
+ // Grammar-constrained decoding on raw completions too (response_format /
3611
+ // guided_* / structured_outputs). The degrade path for /v1/completions
3612
+ // has no chat template to inject a system message into, so a degrade
3613
+ // only emits the Warning header (no prompt injection) — documented gap
3614
+ // vs the chat lane, which mirrors oMLX's text-completions behavior.
3615
+ let textGrammarWarning: string | null = null;
3616
+ const textGrammarReq =
3617
+ body.response_format != null || !!body.guided_grammar ||
3618
+ !!body.guided_regex || !!body.guided_choice?.length ||
3619
+ body.structured_outputs != null;
3620
+ if (textGrammarReq) {
3621
+ const g = await compileGrammarForRequest(body as unknown as ChatRequest);
3622
+ if (g.controller) options.grammar = g.controller;
3623
+ else if (g.degradeHint)
3624
+ textGrammarWarning =
3625
+ `grammar not enforced: ${g.degradeHint} - no prompt injection on /v1/completions`;
3626
+ }
3627
+ // Mirrors the chat lane: a real controller (not the degrade path)
3628
+ // shapes the request for per-row grammar batching.
3629
+ const textGrammarCtrl = options.grammar ?? null;
2199
3630
  // mlx_lm.server's default max_tokens is 512 (its --max-tokens CLI
2200
3631
  // default). The chat lane's very generous default is wrong for raw
2201
3632
  // completion: with no template an EOS may never come.
2202
- options.maxTokens = body.max_completion_tokens ?? body.max_tokens ?? serverOptions.defaultMaxTokens ?? 512;
3633
+ options.maxTokens = body.max_completion_tokens ?? body.max_tokens ??
3634
+ defaultGeneratedTokens ?? 512;
2203
3635
  const requiredCtx = promptIds.length + options.maxTokens;
2204
- if (requiredCtx > admission.maxSafeContext)
3636
+ if (requiredCtx > admission.maxSafeContext) {
3637
+ // Pre-run reject: the gateway only takes grammar disposal ownership
3638
+ // once run — free the WASM matcher here (chat lane: disposeRejected).
3639
+ textGrammarCtrl?.dispose();
2205
3640
  return Response.json(
2206
3641
  {
2207
3642
  error: {
@@ -2216,10 +3651,12 @@ export function createServer(
2216
3651
  },
2217
3652
  { status: 400 },
2218
3653
  );
3654
+ }
2219
3655
  try {
2220
3656
  const adapterIds = ctx.adapters.resolveSpec(body.adapter ?? serverOptions.defaultAdapter);
2221
3657
  if (adapterIds.length) options.adapters = adapterIds;
2222
3658
  } catch (e) {
3659
+ textGrammarCtrl?.dispose();
2223
3660
  return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
2224
3661
  }
2225
3662
  // logprobs capture — same mlx-lm block as chat (generate_response is
@@ -2245,21 +3682,32 @@ export function createServer(
2245
3682
  wantsLogprobs: captureLogprobs,
2246
3683
  userSeed: body.seed !== undefined,
2247
3684
  kvQuant: !!(options.kvConfig?.length || options.kvBits),
3685
+ turboQuant: !!options.turboQuant,
3686
+ // /v1/completions grammar (textGrammarCtrl). Same null-on-degrade
3687
+ // contract as the chat lane.
3688
+ hasGrammar: !!textGrammarCtrl,
3689
+ hasDraft: !!ctx.draft,
2248
3690
  };
2249
3691
  const finishReason = (stopped: boolean, generated: number): "stop" | "length" =>
2250
3692
  stopped ? "stop" : generated >= (options.maxTokens ?? 512) ? "length" : "stop";
2251
3693
 
2252
3694
  if (body.stream) {
3695
+ const streamAbort = new AbortController();
3696
+ const generationSignal = AbortSignal.any([request.signal, streamAbort.signal]);
3697
+ let cancelled = false;
2253
3698
  const stream = new ReadableStream<Uint8Array>({
2254
- async start(controller) {
3699
+ start(controller) {
2255
3700
  const enc = new TextEncoder();
2256
- const send = (obj: unknown) =>
2257
- controller.enqueue(enc.encode(`data: ${JSON.stringify(obj)}\n\n`));
3701
+ const send = (obj: unknown) => {
3702
+ if (!generationSignal.aborted)
3703
+ controller.enqueue(enc.encode(`data: ${JSON.stringify(obj)}\n\n`));
3704
+ };
2258
3705
  const chunk = (text: string, finish: string | null) => ({
2259
3706
  id, object: "text_completion", created, model: ctx.modelId,
2260
3707
  choices: [{ index: 0, text, finish_reason: finish }],
2261
3708
  });
2262
- try {
3709
+ void (async () => {
3710
+ try {
2263
3711
  const decoder = new StreamDecoder(ctx.tokenizer);
2264
3712
  const stopper = new StopMatcher(options.stopSequences);
2265
3713
  // ≥25 ms macrotask hop, same reason as the chat lane: keep the
@@ -2276,7 +3724,8 @@ export function createServer(
2276
3724
  return new Promise<void>((r) => setImmediate(r));
2277
3725
  }
2278
3726
  }
2279
- }, undefined, shape);
3727
+ }, undefined, shape, generationSignal);
3728
+ if (generationSignal.aborted) return;
2280
3729
  if (!stopper.stopped) {
2281
3730
  let tail = stopper.push(decoder.flush());
2282
3731
  if (!stopper.stopped) tail += stopper.flush();
@@ -2292,14 +3741,24 @@ export function createServer(
2292
3741
  completion_tokens: s.generatedTokens,
2293
3742
  total_tokens: s.promptTokens + s.generatedTokens,
2294
3743
  prompt_tokens_details: { cached_tokens: s.cachedTokens },
3744
+ ...(s.spec ? { speculation: s.spec } : {}),
2295
3745
  },
2296
3746
  });
2297
- controller.enqueue(enc.encode("data: [DONE]\n\n"));
2298
- } catch (e) {
2299
- send({ error: { message: (e as Error).message } });
2300
- } finally {
2301
- controller.close();
2302
- }
3747
+ controller.enqueue(enc.encode("data: [DONE]\n\n"));
3748
+ } catch (e) {
3749
+ if (!generationSignal.aborted)
3750
+ send({ error: { message: (e as Error).message } });
3751
+ } finally {
3752
+ if (!cancelled) {
3753
+ if (generationSignal.aborted) controller.error(generationSignal.reason);
3754
+ else controller.close();
3755
+ }
3756
+ }
3757
+ })();
3758
+ },
3759
+ cancel(reason) {
3760
+ cancelled = true;
3761
+ streamAbort.abort(reason);
2303
3762
  },
2304
3763
  });
2305
3764
  return new Response(stream, {
@@ -2307,6 +3766,7 @@ export function createServer(
2307
3766
  "content-type": "text/event-stream",
2308
3767
  "cache-control": "no-cache",
2309
3768
  connection: "keep-alive",
3769
+ ...(textGrammarWarning ? { Warning: textGrammarWarning } : {}),
2310
3770
  },
2311
3771
  });
2312
3772
  }
@@ -2322,7 +3782,7 @@ export function createServer(
2322
3782
  lpc?.push(token, lpInfo);
2323
3783
  text += stopper.push(decoder.push(token));
2324
3784
  if (stopper.stopped) return false; // halt generation
2325
- }, undefined, shape);
3785
+ }, undefined, shape, request.signal);
2326
3786
  if (!stopper.stopped) {
2327
3787
  let tail = stopper.push(decoder.flush());
2328
3788
  if (!stopper.stopped) tail += stopper.flush();
@@ -2341,8 +3801,9 @@ export function createServer(
2341
3801
  completion_tokens: s.generatedTokens,
2342
3802
  total_tokens: s.promptTokens + s.generatedTokens,
2343
3803
  prompt_tokens_details: { cached_tokens: s.cachedTokens },
3804
+ ...(s.spec ? { speculation: s.spec } : {}),
2344
3805
  },
2345
- });
3806
+ }, textGrammarWarning ? { headers: { Warning: textGrammarWarning } } : undefined);
2346
3807
  } catch (e) {
2347
3808
  return Response.json({ error: { message: (e as Error).message } }, { status: 500 });
2348
3809
  }
@@ -2368,7 +3829,7 @@ export function createServer(
2368
3829
  } catch (e) {
2369
3830
  return anthropicError(400, "invalid_request_error", (e as Error).message);
2370
3831
  }
2371
- const resp = await handleChat(chatBody);
3832
+ const resp = await handleChat(chatBody, request.signal);
2372
3833
  if (!resp.ok) {
2373
3834
  const err = (await resp.json().catch(() => null)) as
2374
3835
  | { error?: { message?: string } }
@@ -2449,7 +3910,7 @@ export function createServer(
2449
3910
  } catch (e) {
2450
3911
  return responsesError(400, (e as Error).message);
2451
3912
  }
2452
- const resp = await handleChat(chatBody);
3913
+ const resp = await handleChat(chatBody, request.signal);
2453
3914
  if (!resp.ok) {
2454
3915
  const err = (await resp.json().catch(() => null)) as
2455
3916
  | { error?: { message?: string } }
@@ -2708,6 +4169,27 @@ export function createServer(
2708
4169
  saveHfToken(body.token);
2709
4170
  return Response.json({ ok: true });
2710
4171
  }
4172
+
4173
+ // --- Durable tool-approvals settings (plan §5.4/§6.5/§9 Phase 2) ---
4174
+ // The web chat's own "always allow this tool" list
4175
+ // (~/.mlx-bun/tool-approvals.json, src/tool-approvals.ts). The
4176
+ // approval card itself grants an always-allow via the `approval` WS
4177
+ // frame (alwaysAllow:true) — these REST routes are read/forget only,
4178
+ // for the settings panel's list view (matches the hf-token GET/POST
4179
+ // pattern's separation of "the card that grants" vs "settings that
4180
+ // manage" — there is no POST here on purpose, granting always goes
4181
+ // through the approval card, never a bare settings form).
4182
+ if (url.pathname === "/api/settings/tool-approvals" && request.method === "GET") {
4183
+ const { listAlwaysAllowedTools } = await import("./tool-approvals");
4184
+ return Response.json({ ok: true, alwaysAllow: listAlwaysAllowedTools() });
4185
+ }
4186
+ if (url.pathname === "/api/settings/tool-approvals" && request.method === "DELETE") {
4187
+ const body = (await request.json().catch(() => ({}))) as { tool?: string };
4188
+ if (!body.tool) return Response.json({ ok: false, error: "tool required" }, { status: 400 });
4189
+ const { revokeToolAlwaysAllowed } = await import("./tool-approvals");
4190
+ const file = revokeToolAlwaysAllowed(body.tool);
4191
+ return Response.json({ ok: true, alwaysAllow: Object.keys(file.allows).sort() });
4192
+ }
2711
4193
  {
2712
4194
  const m = url.pathname.match(/^\/api\/(quantize|finetune|dataset)\/push$/);
2713
4195
  if (m && request.method === "POST") {
@@ -2736,6 +4218,50 @@ export function createServer(
2736
4218
  }
2737
4219
  }
2738
4220
 
4221
+ // gc plan/execute (web-ui-pass-plan.md #9): thin wrappers over the CLI's
4222
+ // own `mlx-bun gc` planner/executor (src/registry.ts planGc/executeGc).
4223
+ // Loopback-served admin route — no auth beyond the server's own bind
4224
+ // (matches every other /api/* route here). GET plans (read-only, cheap:
4225
+ // config.json + safetensors index reads, no tensor bytes); POST executes
4226
+ // and requires an explicit {yes:true} body (mirrors the CLI's --yes gate)
4227
+ // so a stray call can't delete anything by accident.
4228
+ if (url.pathname === "/api/gc/plan" && request.method === "GET") {
4229
+ const { planGc } = await import("./registry");
4230
+ const plans = planGc().filter(
4231
+ (p) => p.pruneSnapshots.length || p.skippedSnapshots.length || p.deadBlobs.length,
4232
+ );
4233
+ const superseded = plans.map((p) => ({
4234
+ repo_id: p.repoId,
4235
+ prune_snapshots: p.pruneSnapshots.length,
4236
+ skipped_snapshots: p.skippedSnapshots.length,
4237
+ dead_blobs: p.deadBlobs.length,
4238
+ reclaim_bytes: p.reclaimBytes,
4239
+ }));
4240
+ const reclaim_bytes = plans.reduce((a, p) => a + p.reclaimBytes, 0);
4241
+ return Response.json({ ok: true, superseded, reclaim_bytes });
4242
+ }
4243
+ if (url.pathname === "/api/gc/execute" && request.method === "POST") {
4244
+ const body = (await request.json().catch(() => ({}))) as { yes?: boolean };
4245
+ if (body.yes !== true)
4246
+ return Response.json({ ok: false, error: "pass {\"yes\": true} to confirm deletion" }, { status: 400 });
4247
+ const { planGc, executeGc, Registry } = await import("./registry");
4248
+ const plans = planGc().filter(
4249
+ (p) => p.pruneSnapshots.length || p.skippedSnapshots.length || p.deadBlobs.length,
4250
+ );
4251
+ const res = executeGc(plans);
4252
+ const reg = new Registry();
4253
+ try {
4254
+ await reg.scan(); // reap deleted snapshots from the registry
4255
+ } finally {
4256
+ reg.close();
4257
+ }
4258
+ libraryCache = null; // force /library to reflect the deletion
4259
+ return Response.json({
4260
+ ok: true,
4261
+ snapshots: res.snapshots, blobs: res.blobs, reclaimed_bytes: res.reclaimedBytes,
4262
+ });
4263
+ }
4264
+
2739
4265
  if (url.pathname === "/api/jobs" && request.method === "GET") {
2740
4266
  const store = await ensureJobs();
2741
4267
  const limit = Number(url.searchParams.get("limit") ?? "50");