mlx-bun 0.2.0 → 0.4.0

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 (433) hide show
  1. package/README.md +71 -588
  2. package/docs/README.md +94 -0
  3. package/docs/{design → archive/investigations}/audio-input-plan.md +4 -4
  4. package/docs/{design → archive/investigations}/diffusion-gemma-port.md +3 -3
  5. package/docs/{design → archive/investigations}/faithful-l1-consolidation.md +4 -4
  6. package/docs/{design → archive/investigations}/minicpm5-decode-megakernel.md +11 -11
  7. package/docs/archive/investigations/openjarvis-fit.md +198 -0
  8. package/docs/{design → archive/investigations}/optimization_plan.md +1 -1
  9. package/docs/{design → archive/investigations}/pre-colibri-stabilization.md +5 -5
  10. package/docs/archive/investigations/qwen38-27b-campaign-log.md +1092 -0
  11. package/docs/{design → archive/investigations}/the-dreaming-handoff.md +2 -2
  12. package/docs/archive/planning/release-notes-v0.3.0.md +35 -0
  13. package/docs/design/batching.md +1841 -0
  14. package/docs/design/decode-speed-program.md +2346 -169
  15. package/docs/design/dreaming-nightly-pipeline.md +675 -282
  16. package/docs/design/generic-model-support.md +655 -334
  17. package/docs/design/kv-cache.md +824 -0
  18. package/docs/design/orpo-training.md +673 -1239
  19. package/docs/design/speculative-decoding.md +1990 -0
  20. package/docs/design/turboquant.md +2212 -0
  21. package/docs/design/unified-engine-frontier-plan.md +1833 -727
  22. package/docs/design/web-chat-redesign.md +666 -1083
  23. package/docs/planning/PRODUCT_ROADMAP.md +50 -1
  24. package/docs/planning/release-notes-v0.4.0.md +45 -0
  25. package/docs/planning/video-series.md +2 -2
  26. package/docs/reference/benchmarks.md +3613 -0
  27. package/docs/reference/cli.md +434 -160
  28. package/docs/reference/distribution.md +328 -182
  29. package/docs/reference/environment.md +398 -0
  30. package/docs/reference/glossary.md +39 -0
  31. package/docs/reference/library-api.md +566 -168
  32. package/docs/reference/memory.md +227 -64
  33. package/docs/reference/models.md +113 -6
  34. package/docs/reference/server-api.md +944 -429
  35. package/docs/reference/server-config.md +740 -338
  36. package/docs/reference/training.md +393 -254
  37. package/docs/reference/troubleshooting.md +158 -0
  38. package/package.json +9 -4
  39. package/src/audio/decode.ts +2 -2
  40. package/src/audio/features.ts +2 -2
  41. package/src/backends/mlx/assistant-target.ts +46 -0
  42. package/src/backends/mlx/autoregressive-draft-rows.ts +120 -0
  43. package/src/backends/mlx/autoregressive.ts +108 -0
  44. package/src/backends/mlx/batch-group.ts +1790 -0
  45. package/src/backends/mlx/cache-attachment.ts +67 -0
  46. package/src/backends/mlx/cache-layout.ts +31 -0
  47. package/src/backends/mlx/checkpoint-state.ts +53 -0
  48. package/src/backends/mlx/continuation-persistence.ts +106 -0
  49. package/src/backends/mlx/continuation-request.ts +43 -0
  50. package/src/backends/mlx/continuation.ts +77 -0
  51. package/src/backends/mlx/diffusion.ts +39 -0
  52. package/src/backends/mlx/draft-target.ts +32 -0
  53. package/src/backends/mlx/evaluation.ts +105 -0
  54. package/src/backends/mlx/fill-append.ts +45 -0
  55. package/src/backends/mlx/full-kv-row-append.ts +32 -0
  56. package/src/backends/mlx/full-kv-row-donor.ts +54 -0
  57. package/src/backends/mlx/full-transitioning-kv-rows.ts +65 -0
  58. package/src/backends/mlx/gateway-binding.ts +130 -0
  59. package/src/backends/mlx/graph.ts +38 -0
  60. package/src/backends/mlx/kv-codec.ts +9 -0
  61. package/src/backends/mlx/kv-maintenance.ts +133 -0
  62. package/src/backends/mlx/kv-tensor-rows.ts +212 -0
  63. package/src/backends/mlx/memory-client.ts +189 -0
  64. package/src/backends/mlx/methods.ts +90 -0
  65. package/src/backends/mlx/model-prompt.ts +164 -0
  66. package/src/backends/mlx/model-serving.ts +79 -0
  67. package/src/backends/mlx/prefill-cohort.ts +89 -0
  68. package/src/backends/mlx/prefill-rows.ts +200 -0
  69. package/src/backends/mlx/prefill.ts +33 -0
  70. package/src/backends/mlx/projected-context-rows.ts +72 -0
  71. package/src/backends/mlx/request-policy.ts +20 -0
  72. package/src/backends/mlx/request-state-policy.ts +41 -0
  73. package/src/backends/mlx/rollback.ts +89 -0
  74. package/src/backends/mlx/row-sampling.ts +49 -0
  75. package/src/backends/mlx/serial-executor.ts +242 -0
  76. package/src/backends/mlx/speculative-group.ts +359 -0
  77. package/src/backends/mlx/speculative-round.ts +161 -0
  78. package/src/backends/mlx/speculative.ts +114 -0
  79. package/src/backends/mlx/state-rows.ts +104 -0
  80. package/src/backends/mlx/state-views.ts +46 -0
  81. package/src/backends/mlx/target-layout-capability.ts +35 -0
  82. package/src/backends/mlx/text-engine.ts +73 -0
  83. package/src/backends/mlx/transitioning-kv-rows.ts +165 -0
  84. package/src/cli.ts +66 -58
  85. package/src/client.ts +66 -0
  86. package/src/config.ts +21 -4
  87. package/src/contracts/completion.ts +16 -0
  88. package/src/contracts/execution.ts +60 -0
  89. package/src/contracts/generation.ts +103 -0
  90. package/src/contracts/host.ts +7 -0
  91. package/src/contracts/jobs.ts +66 -0
  92. package/src/contracts/pi-provider.ts +30 -0
  93. package/src/contracts/pi.ts +259 -0
  94. package/src/contracts/prefix-cache.ts +23 -0
  95. package/src/contracts/resources.ts +25 -0
  96. package/src/contracts/sampling.ts +23 -0
  97. package/src/contracts/scheduling.ts +45 -0
  98. package/src/contracts/tsconfig.json +14 -0
  99. package/src/diffusion/diffusion-generate.ts +218 -153
  100. package/src/embed.ts +1 -1
  101. package/src/engine/admission.ts +59 -0
  102. package/src/engine/cancellation.ts +38 -0
  103. package/src/engine/engine.ts +56 -0
  104. package/src/engine/execution-coordinator.ts +70 -0
  105. package/src/engine/execution-plan.ts +53 -0
  106. package/src/engine/resources.ts +53 -0
  107. package/src/engine/scheduler.ts +72 -0
  108. package/src/engine/session.ts +247 -0
  109. package/src/eval/README.md +44 -1
  110. package/src/eval/kl.ts +236 -3
  111. package/src/eval/reference-logits.ts +207 -0
  112. package/src/eval/runner.ts +8 -96
  113. package/src/fill/echo-index.ts +188 -0
  114. package/src/fill/fill-session.ts +452 -0
  115. package/src/fill/proposal.ts +68 -0
  116. package/src/fill/schema-rows.ts +506 -0
  117. package/src/fill/tool-boundary.ts +103 -0
  118. package/src/fit.ts +30 -16
  119. package/src/generate.ts +709 -417
  120. package/src/grammar.ts +12 -7
  121. package/src/harness-pi.ts +14 -33
  122. package/src/hub-rest.ts +4 -4
  123. package/src/index.ts +7 -0
  124. package/src/inference/denoising.ts +17 -0
  125. package/src/inference/draft-acceptance.ts +41 -0
  126. package/src/inference/draft-output.ts +49 -0
  127. package/src/inference/graph.ts +33 -0
  128. package/src/inference/prefill.ts +39 -0
  129. package/src/inference/rollback.ts +7 -0
  130. package/src/jobs/db.ts +4 -0
  131. package/src/jobs/index.ts +2 -0
  132. package/src/jobs/job-entry.ts +12 -12
  133. package/src/jobs/runner.ts +98 -13
  134. package/src/jobs/task-client.ts +19 -0
  135. package/src/jobs/types.ts +2 -58
  136. package/src/kv-scheme.ts +70 -12
  137. package/src/kv-store.ts +177 -63
  138. package/src/lab/README.md +18 -0
  139. package/src/{assets → lab/curve}/curve-designer.html +2 -2
  140. package/src/{curve-sampler.ts → lab/curve/curve-sampler.ts} +2 -2
  141. package/src/{expert-trace.ts → lab/expert-trace/expert-trace.ts} +5 -5
  142. package/src/lab/paged-kv/paged-kv-rows.ts +83 -0
  143. package/src/{model → lab/paged-kv}/paged-kv.ts +32 -15
  144. package/src/library.ts +40 -0
  145. package/src/lora.ts +18 -0
  146. package/src/media-fetch.ts +1 -1
  147. package/src/memory/article.ts +1 -1
  148. package/src/memory/chunk.ts +1 -1
  149. package/src/memory/model.ts +30 -169
  150. package/src/memory/prompts.ts +1 -1
  151. package/src/memory/rest.ts +1 -1
  152. package/src/memory/schedule.ts +1 -1
  153. package/src/memory/stages.ts +4 -1
  154. package/src/memory/vault.ts +1 -3
  155. package/src/mlx/array.ts +11 -11
  156. package/src/mlx/compile.ts +12 -2
  157. package/src/mlx/ffi.ts +58 -11
  158. package/src/mlx/materialize.ts +23 -0
  159. package/src/mlx/metal-kernel.ts +10 -11
  160. package/src/mlx/ops.ts +33 -13
  161. package/src/mlx/quantized-rows.ts +36 -0
  162. package/src/mlx/token-bitmask.ts +52 -0
  163. package/src/mlx/turboquant-kv-decode.ts +82 -0
  164. package/src/mlx/turboquant-kv-inverse.ts +122 -0
  165. package/src/mlx/turboquant-ops.ts +33 -3
  166. package/src/mlx/turboquant-tables.ts +2 -2
  167. package/src/model/artifact-identity.ts +73 -0
  168. package/src/model/batched-kv.ts +235 -0
  169. package/src/model/batched-mask.ts +20 -6
  170. package/src/model/batched-quant.ts +9 -3
  171. package/src/model/batched-quantized-kv.ts +144 -0
  172. package/src/model/batched-rotating-quant.ts +78 -17
  173. package/src/model/batched-rotating-state.ts +76 -0
  174. package/src/model/batched-rotating.ts +65 -23
  175. package/src/model/batched-row-storage.ts +52 -0
  176. package/src/model/batched-ssm.ts +52 -0
  177. package/src/model/batched-turboquant-kv.ts +93 -0
  178. package/src/model/compiled-decode.ts +8 -2
  179. package/src/model/decoded-kv-donor.ts +12 -0
  180. package/src/model/delayed-quantized-kv.ts +72 -0
  181. package/src/model/delayed-rotating-quantized-kv.ts +93 -0
  182. package/src/model/delayed-turboquant-kv.ts +55 -0
  183. package/src/model/diffusion-gemma.ts +8 -8
  184. package/src/model/factory.ts +76 -33
  185. package/src/model/fingerprint.ts +1 -1
  186. package/src/model/flash-attention.ts +3 -3
  187. package/src/model/full-prefill-padding.ts +56 -0
  188. package/src/model/full-prefill-row.ts +105 -0
  189. package/src/model/gemma4-base.ts +179 -80
  190. package/src/model/gemma4.ts +40 -32
  191. package/src/model/generated/gemma4-12b.ts +6 -3
  192. package/src/model/generated/gemma4-26b.ts +6 -3
  193. package/src/model/generated/gemma4-e4b.ts +6 -3
  194. package/src/model/glm52-cache.ts +80 -4
  195. package/src/model/glm52-dsa.ts +13 -9
  196. package/src/model/glm52-memory.ts +1 -1
  197. package/src/model/glm52-mla.ts +12 -6
  198. package/src/model/glm52.ts +19 -40
  199. package/src/model/implementation.ts +69 -0
  200. package/src/model/kv-attention-view.ts +67 -0
  201. package/src/model/minicpm5.ts +17 -15
  202. package/src/model/profile.ts +12 -2
  203. package/src/model/qwen-append.ts +12 -0
  204. package/src/model/qwen-conv.ts +52 -0
  205. package/src/model/qwen3-delta-state.ts +74 -0
  206. package/src/model/qwen3-delta.ts +59 -36
  207. package/src/model/qwen3_5-checkpoint.ts +164 -0
  208. package/src/model/qwen3_5.ts +202 -78
  209. package/src/model/rotating-kv-donor.ts +35 -0
  210. package/src/model/rotating-kv-layout.ts +262 -0
  211. package/src/model/rotating-row-transaction.ts +35 -0
  212. package/src/model/speculative-rotating-kv.ts +101 -0
  213. package/src/model/ssm-prefill-padding.ts +39 -0
  214. package/src/model/support.ts +1 -1
  215. package/src/model/trellis-balanced-scatter.ts +51 -0
  216. package/src/model/trellis-linear.ts +672 -0
  217. package/src/model/trellis-shared-m.ts +151 -0
  218. package/src/model/trellis-shared-scatter.ts +144 -0
  219. package/src/model/trellis-splitk-prefill.ts +158 -0
  220. package/src/model/trellis-tiled-prefill.ts +136 -0
  221. package/src/model/trellis-vector-expand.ts +62 -0
  222. package/src/model/trellis-wide-prefill.ts +94 -0
  223. package/src/model/turboquant-codec.ts +83 -0
  224. package/src/model/universal/dense.ts +15 -14
  225. package/src/model/universal/modules.ts +8 -1
  226. package/src/native-pack.ts +6 -5
  227. package/src/pi-provider.ts +6 -64
  228. package/src/pi-web.ts +8 -315
  229. package/src/preflight.ts +31 -15
  230. package/src/prompt-cache.ts +110 -69
  231. package/src/quantize/allocator.ts +1 -1
  232. package/src/quantize/config-writer.ts +10 -3
  233. package/src/quantize/rotate.ts +4 -4
  234. package/src/quantize/trellis.ts +586 -0
  235. package/src/responses.ts +28 -1
  236. package/src/runtime-config.ts +18 -4
  237. package/src/sampler.ts +40 -29
  238. package/src/serve/batch-scheduler.ts +8 -1671
  239. package/src/serve/chat-request.ts +391 -0
  240. package/src/serve/chat-stage.ts +195 -0
  241. package/src/serve/checkpoint-identity.ts +64 -0
  242. package/src/serve/completion-executor.ts +15 -7
  243. package/src/serve/discovery-routes.ts +12 -13
  244. package/src/serve/generation-gateway.ts +183 -213
  245. package/src/serve/http.ts +159 -0
  246. package/src/serve/inference-request.ts +79 -0
  247. package/src/serve/isolate.ts +297 -105
  248. package/src/serve/lab-routes.ts +3 -0
  249. package/src/serve/lane-registry.ts +2 -1
  250. package/src/serve/model-admin-routes.ts +7 -5
  251. package/src/serve/model-binding.ts +43 -0
  252. package/src/serve/model-host.ts +511 -0
  253. package/src/serve/openai-wire.ts +103 -0
  254. package/src/serve/parent-application.ts +56 -0
  255. package/src/serve/pipeline.ts +33 -0
  256. package/src/serve/preparation.ts +33 -0
  257. package/src/serve/request-plan.ts +23 -18
  258. package/src/serve/request-prep.ts +344 -0
  259. package/src/serve/response-lease.ts +24 -0
  260. package/src/serve/responses-client.ts +69 -0
  261. package/src/serve/session-completion-engine.ts +110 -0
  262. package/src/serve/session-search.ts +1 -1
  263. package/src/serve/static-routes.ts +5 -5
  264. package/src/serve/text-completion-stage.ts +104 -0
  265. package/src/serve/token-streams.ts +519 -0
  266. package/src/server.ts +335 -2711
  267. package/src/spec/assistant-rows.ts +92 -0
  268. package/src/spec/assistant-source.ts +31 -55
  269. package/src/spec/cache-identity.ts +22 -0
  270. package/src/spec/deepspec-rows.ts +29 -0
  271. package/src/spec/deepspec-source.ts +16 -14
  272. package/src/spec/dflash-source.ts +36 -12
  273. package/src/spec/drafter.ts +153 -142
  274. package/src/spec/dspark/ab-stats.ts +3 -3
  275. package/src/spec/dspark/calibration.ts +2 -2
  276. package/src/spec/dspark/deepspec-module.ts +97 -122
  277. package/src/spec/dspark/generate-dflash.ts +2 -2
  278. package/src/spec/dspark/generate.ts +2 -2
  279. package/src/spec/dspark/loader.ts +1 -1
  280. package/src/spec/dspark/module-dflash.ts +129 -172
  281. package/src/spec/dspark/module.ts +3 -3
  282. package/src/spec/dspark/quantize-drafter.ts +1 -1
  283. package/src/spec/generate.ts +1 -1
  284. package/src/spec/glm52-mtp-graph.ts +29 -0
  285. package/src/spec/glm52-mtp-rows.ts +132 -0
  286. package/src/spec/glm52-mtp-source.ts +54 -137
  287. package/src/spec/glm52-mtp-state.ts +33 -0
  288. package/src/spec/ngram-source.ts +93 -21
  289. package/src/spec/projected-draft-rows.ts +73 -0
  290. package/src/spec/qwen-mtp-module.ts +171 -0
  291. package/src/spec/qwen-mtp-rows.ts +219 -0
  292. package/src/spec/qwen-mtp-source.ts +147 -195
  293. package/src/spec/qwen-mtp-state.ts +30 -0
  294. package/src/spec/serve-loop.ts +201 -201
  295. package/src/spec/source.ts +110 -13
  296. package/src/spec/two-model-rows.ts +61 -0
  297. package/src/spec/two-model.ts +44 -102
  298. package/src/ssd-cache.ts +99 -19
  299. package/src/ssd-durability.ts +36 -20
  300. package/src/tool-call.ts +51 -14
  301. package/src/train/flash-cce.ts +2 -2
  302. package/src/train/forward.ts +6 -0
  303. package/src/train/loss.ts +1 -1
  304. package/src/train/prefix-shared.ts +5 -1
  305. package/src/train/segmented.ts +5 -3
  306. package/src/train/trainer.ts +6 -6
  307. package/src/vision/prompt.ts +1 -1
  308. package/src/vision/qwen3vl-tower.ts +52 -10
  309. package/src/web/app.html +1 -1
  310. package/src/web/app.js +62 -60
  311. package/src/web/src/adapters-panel.ts +3 -3
  312. package/src/web/src/assistant.ts +2 -36
  313. package/src/web/src/chat.ts +4 -4
  314. package/src/web/src/composer.ts +4 -4
  315. package/src/web/src/hub.ts +1 -1
  316. package/src/web/src/markdown.ts +0 -0
  317. package/src/web/src/memory-panel.ts +1 -1
  318. package/src/web/src/model-picker.ts +1 -1
  319. package/src/web/src/palette.ts +1 -1
  320. package/src/web/src/protocol.ts +3 -39
  321. package/src/web/src/sessions.ts +1 -1
  322. package/src/web/src/ui-catalog.ts +4 -18
  323. package/src/web-assets.ts +47 -0
  324. package/src/weights.ts +55 -9
  325. package/docs/archive/investigations/curve-runs/bisector-route.json +0 -98
  326. package/docs/archive/investigations/curve-runs/bisector-trace.json +0 -1393
  327. package/docs/archive/investigations/curve-runs/curve-terrain.json +0 -1179
  328. package/docs/archive/investigations/curve-runs/curve-terrain.svg +0 -42
  329. package/docs/archive/investigations/hlg-ab.html +0 -34
  330. package/docs/archive/investigations/hlg-explorer.html +0 -238
  331. package/docs/archive/investigations/hlg-figs/cliff.svg +0 -1
  332. package/docs/archive/investigations/hlg-figs/cross.svg +0 -1
  333. package/docs/archive/investigations/hlg-figs/edges.svg +0 -1
  334. package/docs/archive/investigations/hlg-figs/k20-grids.svg +0 -1
  335. package/docs/archive/investigations/hlg-figs/k20-peaks.svg +0 -1
  336. package/docs/archive/investigations/hlg-figs/lw-curve.svg +0 -1
  337. package/docs/archive/investigations/hlg-figs/terrain.svg +0 -1
  338. package/docs/archive/investigations/hlg-report.html +0 -262
  339. package/docs/archive/investigations/hlg-runs/hlg-ab.json +0 -149
  340. package/docs/archive/investigations/hlg-runs/hlg-blind-prompts.md +0 -87
  341. package/docs/archive/investigations/hlg-runs/hlg-blind.json +0 -107
  342. package/docs/archive/investigations/hlg-runs/hlg-cliff.json +0 -197
  343. package/docs/archive/investigations/hlg-runs/hlg-cross.json +0 -505
  344. package/docs/archive/investigations/hlg-runs/hlg-edge.json +0 -388
  345. package/docs/archive/investigations/hlg-runs/hlg-eval1.json +0 -25
  346. package/docs/archive/investigations/hlg-runs/hlg-frontier.json +0 -86
  347. package/docs/archive/investigations/hlg-runs/hlg-grid.json +0 -242
  348. package/docs/archive/investigations/hlg-runs/hlg-grid2.json +0 -343
  349. package/docs/archive/investigations/hlg-runs/hlg-lw-coarse.json +0 -110
  350. package/docs/archive/investigations/hlg-runs/hlg-map.json +0 -383
  351. package/docs/archive/investigations/hlg-runs/run0-original.log +0 -159
  352. package/docs/archive/investigations/hlg-runs/runA-recommended.log +0 -131
  353. package/docs/archive/investigations/hlg-runs/runAA-cliff-cpm5.log +0 -26
  354. package/docs/archive/investigations/hlg-runs/runAB-grid2-e4b.log +0 -18
  355. package/docs/archive/investigations/hlg-runs/runAC-peak-e4b-partial.log +0 -25
  356. package/docs/archive/investigations/hlg-runs/runAD-ab-samples.log +0 -8
  357. package/docs/archive/investigations/hlg-runs/runAE-edge-e4b.log +0 -51
  358. package/docs/archive/investigations/hlg-runs/runAF-logit-calibration.log +0 -10
  359. package/docs/archive/investigations/hlg-runs/runAG-eval-config2.log +0 -16
  360. package/docs/archive/investigations/hlg-runs/runB-gain-variants.log +0 -192
  361. package/docs/archive/investigations/hlg-runs/runC-pivots.log +0 -162
  362. package/docs/archive/investigations/hlg-runs/runD-oetf.log +0 -132
  363. package/docs/archive/investigations/hlg-runs/runE-eotf.log +0 -162
  364. package/docs/archive/investigations/hlg-runs/runF-shaper-working.log +0 -86
  365. package/docs/archive/investigations/hlg-runs/runG-shaper-sweep.log +0 -101
  366. package/docs/archive/investigations/hlg-runs/runH-wsweep.log +0 -176
  367. package/docs/archive/investigations/hlg-runs/runI-asweep.log +0 -176
  368. package/docs/archive/investigations/hlg-runs/runJ-knobmap.log +0 -49
  369. package/docs/archive/investigations/hlg-runs/runK-frontier.log +0 -19
  370. package/docs/archive/investigations/hlg-runs/runL-widemap.log +0 -65
  371. package/docs/archive/investigations/hlg-runs/runM-frontier-emb.log +0 -15
  372. package/docs/archive/investigations/hlg-runs/runN-verify.log +0 -116
  373. package/docs/archive/investigations/hlg-runs/runO-frontier-seed5000.log +0 -15
  374. package/docs/archive/investigations/hlg-runs/runP-grid-AxSm.log +0 -15
  375. package/docs/archive/investigations/hlg-runs/runQ-grid-tgW.log +0 -15
  376. package/docs/archive/investigations/hlg-runs/runR-blind-transcript.log +0 -79
  377. package/docs/archive/investigations/hlg-runs/runS-cross-coherence.log +0 -33
  378. package/docs/archive/investigations/hlg-runs/runT-cross-cpm5.log +0 -33
  379. package/docs/archive/investigations/hlg-runs/runU-cross-12b.log +0 -33
  380. package/docs/archive/investigations/hlg-runs/runV-lw-e4b.log +0 -22
  381. package/docs/archive/investigations/hlg-runs/runW-lw-12b.log +0 -22
  382. package/docs/archive/investigations/hlg-runs/runX-lw-cpm5.log +0 -22
  383. package/docs/archive/investigations/hlg-runs/runY-cliff-e4b.log +0 -26
  384. package/docs/archive/investigations/hlg-runs/runZ-cliff-12b.log +0 -26
  385. package/docs/archive/mlx-bun-lab-report.html +0 -275
  386. package/docs/archive/wwdc-mlx-bun.html +0 -848
  387. package/docs/design/adapters-end-to-end.md +0 -205
  388. package/docs/design/batching-perf-path.md +0 -255
  389. package/docs/design/batching-v2-plan.md +0 -529
  390. package/docs/design/bucketing-stage.md +0 -197
  391. package/docs/design/compat-cli-surface-design.md +0 -230
  392. package/docs/design/docs-reorg-plan.md +0 -165
  393. package/docs/design/dspark-serving-program.md +0 -371
  394. package/docs/design/dspark-speculative-decoding.md +0 -318
  395. package/docs/design/grammar-spec-batching-integration.md +0 -301
  396. package/docs/design/memory-inference-path.md +0 -422
  397. package/docs/design/memory-synthesis.md +0 -287
  398. package/docs/design/memory-system.md +0 -237
  399. package/docs/design/mlx-lm-serving-execution-seam.md +0 -81
  400. package/docs/design/mlx-lm-tool-parity-plan.md +0 -545
  401. package/docs/design/omlx-adoption-map.md +0 -98
  402. package/docs/design/orpo-dynamic-lambda.md +0 -220
  403. package/docs/design/orpo-future-enhancements.md +0 -340
  404. package/docs/design/paged-kv-cache.md +0 -214
  405. package/docs/design/parallel-slots.md +0 -567
  406. package/docs/design/parity-tier-dag.md +0 -161
  407. package/docs/design/repo-cleanup-plan.md +0 -231
  408. package/docs/design/runtime-isolation.md +0 -263
  409. package/docs/design/segmented-backward-training.md +0 -543
  410. package/docs/design/spec-decode-larger-targets.md +0 -229
  411. package/docs/design/ssd-kv-cold-tier.md +0 -415
  412. package/docs/design/structured-output.md +0 -484
  413. package/docs/design/tauri-desktop-app.md +0 -127
  414. package/docs/design/the-dreaming-master-plan.md +0 -445
  415. package/docs/design/turboquant-kv.md +0 -235
  416. package/docs/design/turboquant-weights.md +0 -467
  417. package/docs/design/w4a16-compute-precision-spike.md +0 -232
  418. package/docs/design/web-chat-beat-matrix.md +0 -315
  419. package/docs/design/write-pipeline-entity.md +0 -468
  420. package/docs/planning/IDEAS.md +0 -69
  421. package/docs/reference/embedding.md +0 -197
  422. package/docs/reference/features-matrix.md +0 -101
  423. package/docs/reference/orpo-quickstart.md +0 -135
  424. package/src/flags.ts +0 -11
  425. /package/docs/{design → archive/investigations}/colibri-glm52-port.md +0 -0
  426. /package/docs/{planning → archive/planning}/chunk-task-roadmap.md +0 -0
  427. /package/docs/{planning → archive/planning}/curve-sampler-research-plan.md +0 -0
  428. /package/docs/{planning → archive/planning}/memory-docs-and-dag-plan.md +0 -0
  429. /package/docs/{planning → archive/planning}/pi-integration-review.md +0 -0
  430. /package/docs/{planning → archive/planning}/release-notes-v0.2.0.md +0 -0
  431. /package/src/{memory-plan.ts → admission-plan.ts} +0 -0
  432. /package/{docs/archive/investigations → src/serve/assets}/curve-terrain.html +0 -0
  433. /package/{docs/dag → src/serve/assets}/training-inference-map.html +0 -0
package/src/server.ts CHANGED
@@ -1,3 +1,7 @@
1
+ import { ContinuationPersistence } from "./backends/mlx/continuation-persistence";
2
+ import { modelServingBinding } from "./backends/mlx/model-serving";
3
+ import type { ServingContext } from "./serve/model-host";
4
+ export { generationCheckpointKey } from "./serve/checkpoint-identity";
1
5
  // OpenAI-compatible HTTP server: /v1/chat/completions (+ SSE streaming)
2
6
  // and /v1/models. Phase 4 core — tool calling, vision, and the
3
7
  // byte-capped prompt cache land on top of this.
@@ -9,89 +13,30 @@
9
13
  // Generation is serialized through a single queue (one GPU, batch=1).
10
14
 
11
15
  import type { Server } from "bun";
12
- // Embedded web app — the unified SPA (chat / quantize / finetune /
13
- // dataset / status). `with { type: "text" }` inlines the file in both
14
- // `bun run` and the compiled single binary. bun-types types *.html
15
- // imports as HTMLBundle (the html loader), but the text attribute makes
16
- // the runtime value a string — hence the double cast.
17
- import appHtml from "./web/app.html" with { type: "text" };
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" };
41
- import { existsSync, readFileSync } from "node:fs";
42
- const APP_PAGE = appHtml as unknown as string;
43
- const HLJS_JS = hljsJs as unknown as string;
44
- const HLJS_CSS = hljsCss as unknown as string;
45
- const APP_JS = appJs as unknown as string;
46
- const MANIFEST_WEBMANIFEST = manifestWebmanifest as unknown as string;
47
- const ICON_SVG = iconSvg as unknown as string;
48
- const SW_JS = swJs as unknown as string;
49
- import { loadModelConfig, type KvQuantSpec, type ModelConfig, type TurboQuantScheme } from "./config";
50
- import { Weights } from "./weights";
51
- import { Gemma4Model } from "./model/gemma4";
52
- import { DiffusionGemmaModel } from "./model/diffusion-gemma";
53
- import { spliceImageTokens } from "./vision/diffusion-vision";
54
- import {
55
- createModel,
56
- openGlm52RuntimeModel,
57
- type Glm52RuntimeOpenOptions,
58
- type RuntimeModel,
59
- } from "./model/factory";
60
- import { Glm52Model } from "./model/glm52";
16
+ import { STATIC_ROUTE_ASSETS } from "./web-assets";
17
+ import { readFileSync } from "node:fs";
18
+ import { type TurboQuantScheme } from "./config";
61
19
  import {
62
20
  GLM52_G5_ASPIRATIONAL_DECODE_TPS,
63
21
  GLM52_G5_DIRECT_ORACLE_WARM_DECODE_TPS,
64
22
  GLM52_G5_MEASURED_AT,
65
23
  GLM52_G5_MEASURED_WARM_DECODE_TPS,
66
- type Glm52MemoryPlan,
67
24
  } from "./model/glm52-memory";
68
- import { isMiniCPM5Config } from "./model/support";
69
- import { resolveModelProfile, type ResolvedModelProfile } from "./model/profile";
70
- import { Glm52NativeMtpProvider } from "./spec/glm52-mtp-source";
71
25
  import {
72
- generate,
73
26
  type GenerateOptions,
74
- type TokenLogprobs,
75
- withModelUsageFlush,
76
- withModelWiredLimit,
77
27
  } from "./generate";
78
- import { cloneKvCaches, SpillQueue } from "./kv-store";
28
+ import { createPreparationExecutor } from "./serve/preparation";
29
+ import { cloneKvCaches, legacyCacheCodecs, SpillQueue } from "./kv-store";
30
+ import type { Cache } from "./model/gemma4";
79
31
  import { resolveKvScheme } from "./kv-scheme";
80
32
  import { runtimeValue } from "./runtime-config";
81
33
  import { TURBOQUANT_HEAD_DIMS } from "./mlx/turboquant-tables";
82
- import {
83
- compileGrammarRequest, grammarEnabled, type GrammarRequest,
84
- } from "./grammar";
85
34
  import type { HlgConfig } from "./sampler";
86
- import { isMonotone, CURVE_UMIN, type CurveParams } from "./curve-sampler";
87
- const CURVE_PAGE = curveDesignerHtml as unknown as string;
88
- import { GenerationGateway } from "./serve/generation-gateway";
35
+ import { isMonotone, CURVE_UMIN, type CurveParams } from "./lab/curve/curve-sampler";
36
+ import { GenerationGateway, disposeUnstartedRequest } from "./serve/generation-gateway";
37
+ import { createSessionCompletionEngine } from "./serve/session-completion-engine";
89
38
  import {
90
39
  CompletionExecutor,
91
- CompletionRejected,
92
- prepareCompletion,
93
- type CompletionSummary,
94
- type PreparedCompletion,
95
40
  } from "./serve/completion-executor";
96
41
  import {
97
42
  createPromptResponseTrace,
@@ -99,34 +44,11 @@ import {
99
44
  } from "./serve/prompt-response-trace";
100
45
  import { handleAdminRoute } from "./serve/admin-routes";
101
46
  import { handleAuxiliaryRoute } from "./serve/aux-routes";
102
- import {
103
- createTimedFlowControl,
104
- type CompletionEvent,
105
- type CompletionStreamProtocol,
106
- } from "./serve/completion-sink";
107
47
  import { createDiscoveryRoutes } from "./serve/discovery-routes";
108
48
  import { handleLabRoute } from "./serve/lab-routes";
109
49
  import { handleModelAdminRoute } from "./serve/model-admin-routes";
110
- import { RequestOwnership } from "./serve/request-plan";
111
50
  import { handleStaticRoute } from "./serve/static-routes";
112
- const STATIC_ROUTE_ASSETS = {
113
- appPage: APP_PAGE,
114
- appJs: APP_JS,
115
- hljsJs: HLJS_JS,
116
- hljsCss: HLJS_CSS,
117
- manifest: MANIFEST_WEBMANIFEST,
118
- iconSvg: ICON_SVG,
119
- serviceWorker: SW_JS,
120
- curvePage: CURVE_PAGE,
121
- };
122
- import {
123
- ChatTemplate, type ChatMessage, type ToolDefinition,
124
- } from "./chat-template";
125
- import { loadTokenizer, type LoadedTokenizer } from "./tokenizer";
126
- import {
127
- CHANNEL_END, CHANNEL_START, parseGeneratedToolCalls, parseToolCalls,
128
- TOOL_CALL_END, TOOL_CALL_START,
129
- } from "./tool-call";
51
+
130
52
  import { PromptCache, cacheBytes } from "./prompt-cache";
131
53
  import { SsdCacheStore } from "./ssd-cache";
132
54
  import {
@@ -140,27 +62,46 @@ import {
140
62
  type AnthropicRequest,
141
63
  } from "./anthropic";
142
64
  import {
143
- ResponseStore, chatJsonToResponses, outputItemsToInputItems,
65
+ ResponseStore, chatJsonToResponses, resolveResponsesConversation,
144
66
  responsesToChatBody, createResponsesStreamProtocol,
145
67
  type ResponsesRequest,
146
68
  } from "./responses";
147
- import { AdapterManager } from "./lora";
148
69
  import { fit } from "./fit";
149
70
  import { setMemoryLimit } from "./mlx/ffi";
150
- import { VisionTower } from "./vision/embedder";
151
- import { SiglipVisionTower, parseSiglipConfig } from "./vision/siglip";
152
- import { Qwen3VLVisionTower } from "./vision/qwen3vl-tower";
153
- import { buildQwen3VLVisionPrompt } from "./vision/qwen3vl-prompt";
154
- import { Qwen35Model } from "./model/qwen3_5";
155
- import type { MropeRequestState } from "./model/qwen3-mrope";
156
- import {
157
- buildMultimodalPrompt, buildVisionPrompt, extractAudio, extractImages, extractVideos,
158
- type AudioTokenIds, type VisionTokenIds, type VisionEncoder,
159
- } from "./vision/prompt";
160
- import { AudioTower, parseAudioConfig } from "./audio/conformer";
161
- import { ensureWav } from "./audio/transcode";
162
- import { sidecarShipsAudioTower } from "./registry";
163
71
  import { makePiWsHandler, type PiWsData } from "./pi-web";
72
+ import { ChatStage } from "./serve/chat-stage";
73
+ import { TextCompletionStage } from "./serve/text-completion-stage";
74
+ import { InferenceStage } from "./serve/inference-request";
75
+ import { admit, respondJson, respondStream, type ErrorFormatter } from "./serve/http";
76
+ import {
77
+ chatCompletionJson, chatCompletionStream, textCompletionJson, textCompletionStream,
78
+ } from "./serve/openai-wire";
79
+ import { createRequestPrep } from "./serve/request-prep";
80
+ import {
81
+ detectDraftKind,
82
+ loadContext,
83
+ type DraftKind,
84
+ type GenSamplingDefaults,
85
+ type LoadContextOptions,
86
+ type ServerContext,
87
+ } from "./serve/model-host";
88
+ import {
89
+ ChatRequest,
90
+ TextCompletionRequest,
91
+ type ChatRequestParams,
92
+ type TextCompletionParams,
93
+ } from "./serve/chat-request";
94
+
95
+ // The serving facade: createServer's companion loaders stay importable from
96
+ // here (public library API via src/index.ts; bench scripts; parity tests).
97
+ export {
98
+ detectDraftKind,
99
+ loadContext,
100
+ type DraftKind,
101
+ type GenSamplingDefaults,
102
+ type LoadContextOptions,
103
+ type ServerContext,
104
+ };
164
105
 
165
106
  export interface ServerOptions {
166
107
  /** Byte cap for the prompt (KV) cache. Default 8 GB. */
@@ -169,34 +110,30 @@ export interface ServerOptions {
169
110
  * (`--kv-budget`, batching-perf-path P3). Joiners whose projected KV
170
111
  * (prompt + max_tokens, window-capped) would exceed it QUEUE until rows
171
112
  * evict; a request over the budget alone is rejected. Unset = no
172
- * aggregate cap (per-request admission via memoryBudget still applies). */
113
+ * aggregate cap (an explicit memoryBudget also limits each request). */
173
114
  kvBudgetBytes?: number;
174
115
  /** KV quantization override. Unset/"off" is bf16. "config" applies the
175
116
  * model's declared mixed-precision kv_config; supported per-layer schemes
176
117
  * compose with continuous scheduling. A number forces uniform bits
177
118
  * (group size 64, start 0) and uses the preserved serial executor. */
178
119
  kvQuant?: "off" | "config" | number;
179
- /** TurboQuant scheme (docs/design/turboquant-kv.md): a separate axis from
120
+ /** Optional absolute row offset for KV conversion; omitted preserves start-zero server policy. */
121
+ quantizedKvStart?: number;
122
+ /** TurboQuant scheme (docs/design/turboquant.md): a separate axis from
180
123
  * kvQuant above, mutually exclusive with it (`--kv-quant turbo[:k<bits>v
181
124
  * <bits>]` sets this instead of kvQuant). Solo-only in v1 — see
182
125
  * GenerationGateway's explicit refusal. */
183
126
  turboQuant?: TurboQuantScheme;
184
- /** OPTIONAL paged KV cache (`--paged-kv`, docs/design/paged-kv-cache.md):
127
+ /** OPTIONAL paged KV cache (`--paged-kv`, docs/design/kv-cache.md):
185
128
  * vLLM-style block-pool storage for full-attention layers,
186
129
  * gather-to-contiguous before the stock SDPA. Default off (unset = the
187
- * plain KVCache path, byte-identical). v1 scope: serial batch=1,
188
- * Gemma4-family, bf16 — startup refuses `--batch N>1`, `--kv-quant`,
130
+ * plain KVCache path, byte-identical). Gemma4-family, bf16, B1/B>1;
131
+ * startup refuses `--kv-quant`,
189
132
  * turbo, and `--draft-model` combinations; paged requests bypass the
190
133
  * prompt cache and run uncompiled. Gated bit-exact vs the plain path. */
191
134
  pagedKv?: { blockSize?: number };
192
- /** Memory budget for the serving process (admission control — Phase 5).
193
- * Requests whose prompt + max_tokens exceed the budget's max safe
194
- * context are rejected with 400 instead of crashing the GPU: the OOM
195
- * crash class is UNCATCHABLE (Phase 6 — mlx throws from a Metal
196
- * completion handler ⇒ std::terminate; optiq serve died exactly this
197
- * way loading the 26B). Also caps the mlx allocator
198
- * (mlx_set_memory_limit) as defense in depth. Default: machine RAM ×
199
- * WIRED_FRACTION, admission check only, allocator untouched. */
135
+ /** Explicit opt-in process budget. Enforces the fit estimate's context
136
+ * limit and sets the MLX allocator limit. Unset keeps estimates advisory. */
200
137
  memoryBudgetBytes?: number;
201
138
  /** Who owns this server's lifetime: "serve" (persistent, mlx-bun
202
139
  * serve) or "pi-session" (dies with the pi session that started
@@ -204,7 +141,7 @@ export interface ServerOptions {
204
141
  * warn before attaching to a server that may vanish. */
205
142
  owner?: "serve" | "pi-session" | "embedded";
206
143
  /** Listen on a UNIX DOMAIN SOCKET instead of a TCP port — the engine-
207
- * child mode of the isolation architecture (runtime-isolation.md): the
144
+ * child mode of the isolation architecture (docs/reference/server-config.md): the
208
145
  * parent process reverse-proxies HTTP to this socket. Stale socket
209
146
  * files are unlinked before bind. When set, `hostname`/port are ignored. */
210
147
  unixSocket?: string;
@@ -244,10 +181,10 @@ export interface ServerOptions {
244
181
  defaultAdapter?: string;
245
182
  /** SSD cold tier for the prompt/KV cache (`--ssd-cache <dir>`): prefix
246
183
  * KV survives RAM eviction AND server restarts — the coding-agent
247
- * long-context TTFT win (docs/design/ssd-kv-cold-tier.md). Off unless
184
+ * long-context TTFT win (docs/design/kv-cache.md). Off unless
248
185
  * a directory is given. Serial lane only (where the prompt cache lives). */
249
186
  ssdCacheDir?: string;
250
- /** Byte cap for the SSD tier (default 32 GiB). */
187
+ /** Optional byte cap for the SSD tier (default unbounded). */
251
188
  ssdCacheMaxBytes?: number;
252
189
  /** Idle-demotion threshold in seconds (`--ssd-demote-idle`): prompt-cache
253
190
  * entries unused this long spill to the SSD tier and free their GPU
@@ -257,9 +194,16 @@ export interface ServerOptions {
257
194
  /** Verify every tensor hash on restore (`--ssd-cache-verify`) — reads all
258
195
  * bytes eagerly, defeating lazy fault-in; integrity paranoia only. */
259
196
  ssdCacheVerify?: boolean;
197
+ /** Periodically persist an in-flight serial generation every N emitted
198
+ * tokens. An identical request after restart replays the saved assistant
199
+ * prefix and continues from the already-sampled next token. Requires the
200
+ * SSD cache and --batch 1. */
201
+ generationCheckpointTokens?: number;
260
202
  }
261
203
 
262
204
  interface ServerLifecycle {
205
+ stopJobs: () => Promise<void>;
206
+ close: () => Promise<void>;
263
207
  flush: () => Promise<DurabilityFlushResult>;
264
208
  stats: () => DurabilitySnapshotStats;
265
209
  stopTimers: () => void;
@@ -302,8 +246,10 @@ export async function shutdownServer(
302
246
  lifecycle?.stopTimers();
303
247
  const stopped = Promise.resolve(server.stop(false));
304
248
  const work = (async (): Promise<ServerShutdownResult> => {
249
+ await lifecycle?.stopJobs();
305
250
  await stopped;
306
251
  const durability = lifecycle ? await lifecycle.flush() : emptyDurabilityResult();
252
+ await lifecycle?.close();
307
253
  return { stopped: true, timedOut: false, durability };
308
254
  })();
309
255
  let timeoutHandle: ReturnType<typeof setTimeout> | null = null;
@@ -330,548 +276,6 @@ export async function shutdownServer(
330
276
  }
331
277
  }
332
278
 
333
- export interface ServerContext {
334
- model: RuntimeModel;
335
- /** Declared external artifact/family profile that selected model
336
- * construction. Request-level methods are resolved separately. */
337
- profile: ResolvedModelProfile;
338
- tokenizer: LoadedTokenizer;
339
- template: ChatTemplate;
340
- modelId: string;
341
- /** Lazily-loaded vision tower cache — null until the first image request
342
- * (see `getVisionTower`). The tower (SigLIP ~hundreds of MB, encoder-free
343
- * smaller) is not loaded for text-only sessions. */
344
- vision: VisionEncoder | null;
345
- /** Loads + selects the vision tower on demand; null when the model has no
346
- * (supported) vision sidecar. Invoked at most once, then cached in
347
- * `vision`. */
348
- loadVision: (() => VisionEncoder) | null;
349
- visionTokenIds: VisionTokenIds;
350
- /** Lazily-loaded Conformer audio tower — null until the first audio
351
- * request (see `getAudioTower`). Same sidecar file as vision
352
- * (optiq_vision.safetensors), separate tower; text-only sessions never
353
- * pay for it. */
354
- audio: AudioTower | null;
355
- /** Loads the audio tower on demand; null when the model can't do audio
356
- * (no `audio_config` in config.json or no sidecar). No flags — audio
357
- * auto-enables exactly like vision. */
358
- loadAudio: (() => AudioTower) | null;
359
- /** null when the model has no `audio_config` (audio-incapable). */
360
- audioTokenIds: AudioTokenIds | null;
361
- adapters: AdapterManager;
362
- /** Per-layer KV quantization from the repo's kv_config.json (null if
363
- * absent). Applied by default — optiq serve's headline behavior;
364
- * ServerOptions.kvQuant overrides ("off" | uniform bits). */
365
- kvConfig: KvQuantSpec[] | null;
366
- /** Model-author recommended sampling from generation_config.json —
367
- * optiq serve injects these as server defaults (gen_config.py);
368
- * explicit request fields always win. */
369
- genDefaults: GenSamplingDefaults;
370
- /** Speculative decoding (`serve --draft-model`, mlx_lm.server parity).
371
- * Server-level: when set, EVERY request routes to the serial lane
372
- * (upstream: is_batchable = draft is None) and spec-eligible ones decode
373
- * through src/spec/serve-loop.ts. null = no draft configured. */
374
- draft?: {
375
- provider: import("./spec/source").DraftProvider;
376
- numDraftTokens: number;
377
- } | null;
378
- /** Present only for the direct Colibri GLM runtime. This is the exact
379
- * header-derived process equation used before opening resident state. */
380
- glmMemoryPlan?: Glm52MemoryPlan | null;
381
- }
382
-
383
- export interface LoadContextOptions {
384
- memoryBudgetBytes?: number;
385
- /** Direct Colibri runtime/resource overrides. Ignored by other models. */
386
- glm?: Glm52RuntimeOpenOptions;
387
- /** Snapshot dir of a draft model for speculative decoding
388
- * (`--draft-model`). Loaded alongside the target; the pair must share
389
- * a tokenizer family. */
390
- draftModelDir?: string;
391
- /** Drafts per round (`--num-draft-tokens`, mlx_lm.server default 3). */
392
- numDraftTokens?: number;
393
- /** Draft-provider kind override (`--draft-kind`). */
394
- draftKind?: DraftKind;
395
- ngramMax?: number;
396
- ngramMin?: number;
397
- }
398
-
399
- export interface GenSamplingDefaults {
400
- temperature?: number;
401
- topP?: number;
402
- topK?: number;
403
- repetitionPenalty?: number;
404
- }
405
-
406
- export type DraftKind = "dspark" | "deepspec" | "assistant" | "two-model" | "ngram" | "mtp";
407
-
408
- /** Detect the draft artifact's kind so the right provider is loaded. All
409
- * providers share ONE serve loop (src/spec/serve-loop.ts). Exported for the
410
- * bench harness (scripts/bench-feature-matrix.ts) — one detection, no drift.
411
- * "ngram" is never detected — it has no artifact (model-free prompt lookup,
412
- * src/spec/ngram-source.ts) and mounts via an explicit `--draft-kind ngram`. */
413
- export async function detectDraftKind(dir: string): Promise<DraftKind> {
414
- if (await Bun.file(`${dir}/dspark.json`).exists()) return "dspark"; // our trained module
415
- try {
416
- const cfg = (await Bun.file(`${dir}/config.json`).json()) as {
417
- model_type?: string;
418
- architectures?: string[];
419
- };
420
- // DeepSeek's released DSpark drafters (DeepSpec reference): no
421
- // dspark.json, plain HF config stamped Gemma4DSparkModel.
422
- if (cfg.architectures?.[0] === "Gemma4DSparkModel") return "deepspec";
423
- if (String(cfg.model_type ?? "").includes("assistant")) return "assistant";
424
- // Native MTP heads split from a qwen3_5-family release
425
- // (mlx-community/Qwen3.8-27B-MTP-*): model_type "qwen3_5_mtp". The
426
- // target's recurrent DeltaNet caches roll back via the serve loop's
427
- // spec-round snapshot/replay contract (SSMCache.specRound*).
428
- if (String(cfg.model_type ?? "").endsWith("_mtp")) return "mtp";
429
- } catch {
430
- // no/unreadable config → fall through to a full second model
431
- }
432
- return "two-model";
433
- }
434
-
435
- async function loadGenSamplingDefaults(modelDir: string): Promise<GenSamplingDefaults> {
436
- const file = Bun.file(`${modelDir}/generation_config.json`);
437
- if (!(await file.exists())) return {};
438
- try {
439
- const raw = (await file.json()) as Record<string, unknown>;
440
- const num = (v: unknown) => (typeof v === "number" ? v : undefined);
441
- return {
442
- temperature: num(raw.temperature),
443
- topP: num(raw.top_p),
444
- topK: num(raw.top_k),
445
- repetitionPenalty: num(raw.repetition_penalty),
446
- };
447
- } catch {
448
- return {};
449
- }
450
- }
451
-
452
- export async function loadContext(
453
- modelDir: string, modelId?: string,
454
- opts: LoadContextOptions = {},
455
- ): Promise<ServerContext> {
456
- const config = await loadModelConfig(modelDir);
457
- const profile = resolveModelProfile(config);
458
- const glm = profile.profile.execution.loader === "colibri";
459
- // Bundled MTP companion: `--draft-kind mtp` with no --draft-model resolves
460
- // to the artifact's own mtp/ subfolder (single-repo packaging — the
461
- // companion is a complete model dir the provider already loads). Explicit
462
- // --draft-model still wins; a missing bundle is a clear refusal below.
463
- if (opts.draftKind === "mtp" && !opts.draftModelDir) {
464
- const bundled = `${modelDir}/mtp`;
465
- if (await Bun.file(`${bundled}/config.json`).exists()) {
466
- opts = { ...opts, draftModelDir: bundled };
467
- } else {
468
- throw new Error(
469
- `--draft-kind mtp needs a companion: pass --draft-model <dir> or use ` +
470
- `an artifact that bundles one at <model>/mtp/ (none at ${bundled})`,
471
- );
472
- }
473
- }
474
- const externalDraft = opts.draftModelDir !== undefined || opts.draftKind !== undefined;
475
- const resolvedDraftKind = opts.draftModelDir
476
- ? opts.draftKind ?? await detectDraftKind(opts.draftModelDir)
477
- : opts.draftKind;
478
- if (glm && externalDraft && opts.glm?.enableMtp === true)
479
- throw new Error("native GLM MTP and --draft-model/--draft-kind are mutually exclusive");
480
- let weights: Weights | null = null;
481
- let model!: RuntimeModel;
482
- let glmMemoryPlan: Glm52MemoryPlan | null = null;
483
- if (glm) {
484
- const enableNativeMtp = !externalDraft && opts.glm?.enableMtp !== false;
485
- const opened = await openGlm52RuntimeModel(modelDir, {
486
- ...opts.glm,
487
- memoryBudgetBytes: opts.memoryBudgetBytes ?? opts.glm?.memoryBudgetBytes,
488
- // Every drafter owns the serial speculative lane. Do not reserve a
489
- // simultaneous ordinary batch that cannot be admitted while it is live.
490
- batchSize: enableNativeMtp || externalDraft ? 1 : opts.glm?.batchSize,
491
- // An explicit alternate drafter owns the one speculative lane.
492
- enableMtp: enableNativeMtp,
493
- });
494
- model = opened.model;
495
- glmMemoryPlan = opened.plan;
496
- } else {
497
- weights = await Weights.open(modelDir);
498
- }
499
- // memoryBudget enforcement at load (Phase 5): Weights.open only mmaps
500
- // (no GPU allocation yet), so a model whose weights can never serve
501
- // within the budget is refused HERE — before any unified-memory
502
- // commitment — with an actionable error instead of a Metal OOM later.
503
- if (!glm && opts.memoryBudgetBytes) {
504
- const weightsBytes = [...weights!.shards.files.values()]
505
- .reduce((a, f) => a + f.mmap.size, 0);
506
- const report = fit(config, weightsBytes, 1, undefined, undefined, 0, opts.memoryBudgetBytes);
507
- if (report.maxSafeContext < 1)
508
- throw new Error(
509
- `model does not fit the memory budget: weights ${(weightsBytes / 1e9).toFixed(2)} GB ` +
510
- `+ prefill transient leave no room for any context within ` +
511
- `${(opts.memoryBudgetBytes / 1e9).toFixed(2)} GB`,
512
- );
513
- }
514
- if (!glm) model = createModel(weights!, config, profile);
515
- const tokenizer = await loadTokenizer(modelDir);
516
- // Generation must stop on the tokenizer's eos_token — the chat turn
517
- // terminator (e.g. Qwen <|im_end|> = 248046). Some configs (Qwen3.5-4B)
518
- // declare a different eos_token_id in config.json than the chat format
519
- // emits, so without this a turn never ends and generation runs away,
520
- // hallucinating both sides of the dialogue until max_tokens. mlx-lm stops on
521
- // the tokenizer eos; union it in. No-op when already present (Gemma, 27B).
522
- if (tokenizer.eosTokenId != null && !config.eosTokenIds.includes(tokenizer.eosTokenId))
523
- config.eosTokenIds = [...config.eosTokenIds, tokenizer.eosTokenId];
524
-
525
- // Speculative decoding: load the draft (mlx_lm.server --draft-model). The
526
- // draft artifact's KIND selects the provider — all three share ONE serve
527
- // loop (src/spec/serve-loop.ts): dspark.json → DSpark (KV-injected), a
528
- // *_assistant config → the optiq KV-borrowing Gemma drafter, otherwise a
529
- // full second model (mlx-lm parity). `--draft-kind` overrides the detect.
530
- let draft: ServerContext["draft"] = null;
531
- if (opts.draftKind === "ngram") {
532
- // Model-free prompt lookup: no artifact, no dir, no probe/budget concerns
533
- // (weightsBytes 0, open() never throws). Default γ=10 per the reference
534
- // implementation — drafting is free, so wide blocks cost only verify-window
535
- // width when wrong.
536
- if (opts.draftModelDir)
537
- throw new Error(
538
- "--draft-kind ngram is model-free — drop --draft-model (it would be ignored)",
539
- );
540
- const { NgramProvider } = await import("./spec/ngram-source");
541
- draft = {
542
- provider: new NgramProvider({ max: opts.ngramMax, min: opts.ngramMin }),
543
- numDraftTokens: Math.max(1, opts.numDraftTokens ?? 10),
544
- };
545
- } else if (opts.draftModelDir) {
546
- const dir = opts.draftModelDir;
547
- const kind = resolvedDraftKind!;
548
- let provider: import("./spec/source").DraftProvider;
549
- let numDraftTokens = Math.max(1, opts.numDraftTokens ?? 3);
550
- if (kind === "dspark") {
551
- const { DflashProvider } = await import("./spec/dflash-source");
552
- const p = await DflashProvider.load(dir);
553
- provider = p;
554
- // Pin to the trained block width — the serve loop must never ask for
555
- // more positions than the DSpark block was trained for (n ≤ cfg.gamma).
556
- numDraftTokens = Math.max(1, Math.min(opts.numDraftTokens ?? p.gamma, p.gamma));
557
- } else if (kind === "deepspec") {
558
- const { DeepspecProvider } = await import("./spec/deepspec-source");
559
- const p = await DeepspecProvider.load(dir);
560
- provider = p;
561
- // Same pin, from their config's block_size (e.g. 7 for the released
562
- // dspark_gemma4_12b_block7).
563
- numDraftTokens = Math.max(1, Math.min(opts.numDraftTokens ?? p.gamma, p.gamma));
564
- } else if (kind === "assistant") {
565
- const { AssistantProvider } = await import("./spec/assistant-source");
566
- provider = await AssistantProvider.load(dir);
567
- } else if (kind === "mtp") {
568
- const { QwenMtpProvider } = await import("./spec/qwen-mtp-source");
569
- const p = await QwenMtpProvider.load(dir);
570
- provider = p;
571
- // Default the round width to the head's trained block (block_size 3 →
572
- // 2 recursive drafts + the pending row per round; the head was trained
573
- // multi-step, so an explicit larger --num-draft-tokens is allowed but
574
- // acceptance decides whether it pays).
575
- const block = (await Bun.file(`${dir}/config.json`).json() as { block_size?: number }).block_size;
576
- if (opts.numDraftTokens === undefined && typeof block === "number")
577
- numDraftTokens = Math.max(1, block - 1);
578
- } else {
579
- const { TwoModelProvider } = await import("./spec/two-model");
580
- provider = await TwoModelProvider.load(dir, config.text.vocabSize);
581
- // Tokenizer-family hard check — two-model ONLY (it ships its own
582
- // tokenizer). Exact-token-match acceptance is meaningful only when both
583
- // models tokenize identically; a probe that ENCODES differently means
584
- // different families — refuse instead of silently accepting ~0% of
585
- // drafts. The KV-borrowing drafters share the target's tokenization by
586
- // construction, so this check does not apply to them.
587
- const draftTok = await loadTokenizer(dir);
588
- const probe = "The 3 quick brown foxes jumped över the lazy dog?! 🦊";
589
- if (JSON.stringify(tokenizer.encode(probe)) !== JSON.stringify(draftTok.encode(probe))) {
590
- provider.dispose();
591
- throw new Error(
592
- `--draft-model tokenizer differs from the target's (probe string encodes ` +
593
- `differently) — speculation needs the same tokenizer family`,
594
- );
595
- }
596
- }
597
- // Fail-fast pairing validation (2026-07-07 review): the KV-borrowing
598
- // sources validate the (target, drafter) pairing in open() — non-Gemma4
599
- // target, DeepSpec target-layer-count mismatch — which used to surface
600
- // as a 500 from inside specServeRun on EVERY text request. Probe-open
601
- // once with throwaway caches here so a mismatch refuses at load; open()
602
- // allocates no per-request tensors before prefill/draft, so this is
603
- // free. The probe sampler is never called during open().
604
- {
605
- const probeCaches = model.makeCache();
606
- try {
607
- provider
608
- .open({
609
- sampler: () => { throw new Error("probe sampler never samples"); },
610
- target: { model, caches: probeCaches },
611
- })
612
- .dispose();
613
- } catch (err) {
614
- provider.dispose();
615
- throw new Error(
616
- `--draft-model is incompatible with this target: ${err instanceof Error ? err.message : String(err)}`,
617
- );
618
- } finally {
619
- for (const c of probeCaches) c.dispose();
620
- }
621
- }
622
- if (opts.memoryBudgetBytes) {
623
- const targetBytes = weights
624
- ? [...weights.shards.files.values()].reduce((a, f) => a + f.mmap.size, 0)
625
- : model.weightsBytes;
626
- // Draft weights shrink the target's envelope. Draft KV is not modeled
627
- // (small relative to its weights at serve contexts); admission stays
628
- // approximately conservative via the combined-weights term.
629
- const report = fit(config, targetBytes + provider.weightsBytes, 1, undefined, undefined, 0, opts.memoryBudgetBytes);
630
- if (report.maxSafeContext < 1) {
631
- provider.dispose();
632
- throw new Error(
633
- `target + draft do not fit the memory budget (draft adds ` +
634
- `${(provider.weightsBytes / 1e9).toFixed(2)} GB)`,
635
- );
636
- }
637
- }
638
- draft = { provider, numDraftTokens };
639
- } else if (opts.draftKind) {
640
- // Every other kind names an artifact to load — refuse instead of silently
641
- // serving without speculation.
642
- throw new Error(`--draft-kind ${opts.draftKind} requires --draft-model`);
643
- }
644
-
645
- // GLM's checkpoint-native MTP row is the production default. It uses the
646
- // already-planned bounded auxiliary expert tier and the same tokenizer, so
647
- // there is no second artifact or compatibility probe to load.
648
- if (!draft && model instanceof Glm52Model && glmMemoryPlan?.enableMtp) {
649
- draft = {
650
- provider: new Glm52NativeMtpProvider(model),
651
- numDraftTokens: glmMemoryPlan.mtpDraftTokens,
652
- };
653
- }
654
-
655
- return {
656
- draft,
657
- model,
658
- profile,
659
- glmMemoryPlan,
660
- adapters: new AdapterManager(model),
661
- kvConfig: config.kvQuant,
662
- genDefaults: await loadGenSamplingDefaults(modelDir),
663
- tokenizer,
664
- template: await ChatTemplate.load(modelDir),
665
- modelId: modelId ?? modelDir.split("/").filter(Boolean).at(-1)!,
666
- // Vision is loaded lazily (getVisionTower) — text-only sessions never
667
- // pay for the tower. The loader picks the encoder-free gemma4_unified
668
- // (12B) tower vs the SigLIP encoder (e2b/e4b/26B/31B) by the sidecar's
669
- // vision_config.model_type. Vision sidecars are a Gemma4 feature;
670
- // MiniCPM5 never ships one.
671
- vision: null,
672
- loadVision: makeVisionLoader(modelDir, model, config),
673
- visionTokenIds: {
674
- imageTokenId: (config.raw.image_token_id as number) ?? 258880,
675
- boiTokenId: (config.raw.boi_token_id as number) ?? 255999,
676
- eoiTokenId: (config.raw.eoi_token_id as number) ?? 258882,
677
- },
678
- // Audio mirrors vision: lazy tower from the same sidecar, loaded on the
679
- // first audio request only (docs/design/audio-input-plan.md A4).
680
- audio: null,
681
- loadAudio: makeAudioLoader(modelDir, model, config),
682
- audioTokenIds: config.raw.audio_config
683
- ? {
684
- audioTokenId: (config.raw.audio_token_id as number) ?? 258881,
685
- boaTokenId: (config.raw.boa_token_id as number) ?? 256000,
686
- eoaTokenId: (config.raw.eoa_token_id as number) ?? 258883,
687
- }
688
- : null,
689
- };
690
- }
691
-
692
- /** Build the on-demand vision-tower loader, selecting the encoder-free
693
- * (gemma4_unified, 12B) tower vs the SigLIP encoder (gemma4_vision:
694
- * e2b/e4b/26B/31B) by the sidecar's vision_config.model_type. Returns null
695
- * when the model has no usable vision sidecar. */
696
- function makeVisionLoader(
697
- modelDir: string, model: RuntimeModel, config: ModelConfig,
698
- ): (() => VisionEncoder) | null {
699
- // Qwen3.5/3.8: the tower ships as optiq/optiq_vision.safetensors (the
700
- // artifact's bf16 sidecar). The returned tower is a Qwen3VLVisionTower —
701
- // it rides the same lazy slot/capability flags; the qwen chat branch is
702
- // the only consumer and casts it back (the gemma branches are gated on
703
- // `instanceof Gemma4Model`, so the union never crosses).
704
- if (model instanceof Qwen35Model) {
705
- // Vision weights arrive either as the OptiQ-convention sidecar OR in-main
706
- // (mlx-vlm convention; our artifacts ship one copy in-main since
707
- // 2026-08-18 — the tower loader handles both).
708
- const hasSidecar = existsSync(`${modelDir}/optiq/optiq_vision.safetensors`);
709
- const hasInMain = config.raw.vision_config !== undefined;
710
- if (!hasSidecar && !hasInMain) return null;
711
- return () =>
712
- Qwen3VLVisionTower.load(modelDir) as unknown as VisionEncoder;
713
- }
714
- if (!(config.hasVisionSidecar && model instanceof Gemma4Model)) return null;
715
- const vc = config.raw.vision_config as Record<string, any> | undefined;
716
- if (vc?.model_type === "gemma4_vision") {
717
- const sigCfg = parseSiglipConfig(vc);
718
- return () => SiglipVisionTower.load(modelDir, sigCfg, model.embedScale);
719
- }
720
- // gemma4_unified_vision (or unlabelled): the encoder-free patch embedder.
721
- return () => VisionTower.load(modelDir, model.embedScale, config.text.rmsNormEps);
722
- }
723
-
724
- /** Lazily load + cache the vision tower on first use. A sidecar that fails
725
- * to load is a capability gap, not a fatal error: returns null and the
726
- * request is answered with a 400 (the loader is cleared so we don't retry
727
- * a known-bad load every request). */
728
- function getVisionTower(ctx: ServerContext): VisionEncoder | null {
729
- if (ctx.vision) return ctx.vision;
730
- if (!ctx.loadVision) return null;
731
- try {
732
- ctx.vision = ctx.loadVision();
733
- return ctx.vision;
734
- } catch (e) {
735
- console.warn(`vision sidecar not loadable (${(e as Error).message}) — serving text-only`);
736
- ctx.loadVision = null;
737
- return null;
738
- }
739
- }
740
-
741
- /** Build the on-demand audio-tower loader (gemma-4 Conformer, A4 of
742
- * docs/design/audio-input-plan.md). Auto-enables — no flags — when
743
- * config.json carries an `audio_config` AND the optiq_vision.safetensors
744
- * sidecar exists (the audio tensors ship in the same sidecar as vision).
745
- * Returns null when the model can't do audio (no audio_config: 26B-A4B,
746
- * DiffusionGemma, bf16 assistants — architectural, not a porting gap). */
747
- function makeAudioLoader(
748
- modelDir: string, model: RuntimeModel, config: ModelConfig,
749
- ): (() => AudioTower) | null {
750
- if (!(config.hasVisionSidecar && config.raw.audio_config && model instanceof Gemma4Model))
751
- return null;
752
- // The sidecar header must actually name the Conformer tensors: the local
753
- // 12B snapshot pairs audio_config with a STUB sidecar (embed_audio only),
754
- // and a non-null loader here is what advertises `audio: true` on every
755
- // capability surface (ws handshake, /v1/models). Header-only read.
756
- if (!sidecarShipsAudioTower(`${modelDir}/optiq_vision.safetensors`)) return null;
757
- const audioCfg = parseAudioConfig(config.raw.audio_config as Record<string, any>);
758
- return () => AudioTower.load(modelDir, audioCfg, model.embedScale);
759
- }
760
-
761
- /** Lazily load + cache the audio tower on first use. Unlike vision's
762
- * warn-and-continue (text-only degrade is correct for requests WITHOUT
763
- * images), this is only ever consulted for requests WITH audio — the
764
- * caller turns null into an explicit 400, never a silent text-only
765
- * degrade. A failed load is not retried every request. (The stub-sidecar
766
- * case — the local 12B state — never gets here: makeAudioLoader checks the
767
- * sidecar header and returns a null loader.) */
768
- function getAudioTower(ctx: ServerContext): AudioTower | null {
769
- if (ctx.audio) return ctx.audio;
770
- if (!ctx.loadAudio) return null;
771
- try {
772
- ctx.audio = ctx.loadAudio();
773
- return ctx.audio;
774
- } catch (e) {
775
- console.warn(`audio sidecar not loadable (${(e as Error).message})`);
776
- ctx.loadAudio = null;
777
- return null;
778
- }
779
- }
780
-
781
- interface ChatRequest {
782
- messages: ChatMessage[];
783
- stream?: boolean;
784
- max_tokens?: number;
785
- max_completion_tokens?: number;
786
- temperature?: number;
787
- top_p?: number;
788
- top_k?: number;
789
- seed?: number;
790
- repetition_penalty?: number;
791
- /** mlx-lm extension: recent-token window for repetition_penalty
792
- * (default 20; 0 = whole history, Python `[-0:]` semantics). */
793
- repetition_context_size?: number;
794
- /** min-p sampling (mlx_lm.server's `min_p`): keep tokens whose probability
795
- * is ≥ min_p · p(top token). 0 = off. */
796
- min_p?: number;
797
- /** XTC sampling (mlx_lm.server names): with probability `xtc_probability`
798
- * per step, remove every token above `xtc_threshold` except the least
799
- * likely of them. EOS + the newline token are always exempt (the server
800
- * injects them as xtc special tokens, matching mlx_lm.server). */
801
- xtc_probability?: number;
802
- xtc_threshold?: number;
803
- /** OpenAI logit_bias: {tokenId: additive bias}. JSON object keys arrive as
804
- * strings; coerced to int keys / float values like mlx-lm (400 on failure). */
805
- logit_bias?: Record<string, number>;
806
- /** OpenAI presence/frequency penalties + mlx-lm's context-size extensions
807
- * (window of recent tokens the penalty looks at; default 20). */
808
- presence_penalty?: number;
809
- presence_context_size?: number;
810
- frequency_penalty?: number;
811
- frequency_context_size?: number;
812
- /** mlx_lm.server logprobs: `logprobs` is a BOOL (even on /v1/completions —
813
- * not OpenAI's legacy int), `top_logprobs` an int in [0, 11] or the -1
814
- * "unset" sentinel (server.py validates exactly that; OpenAI's cap is 20,
815
- * mlx-lm's is 11 — we copy the reference). Non-stream responses carry
816
- * mlx-lm's logprobs block; stream chunks never do (reference behavior). */
817
- logprobs?: boolean;
818
- top_logprobs?: number;
819
- /** OpenAI stop sequences: plain string or array (spec allows up to 4).
820
- * Matched on DECODED text, not token ids — see StopMatcher. */
821
- stop?: string | string[];
822
- tools?: ToolDefinition[];
823
- tool_choice?: "auto" | "none" | { type: string; function?: { name: string } };
824
- /** Forwarded to HF chat templates, matching optiq serve. MiniCPM5 uses
825
- * enable_thinking to select direct answers vs the <think> channel; Qwen3.8
826
- * adds preserve_thinking (keep think blocks from history; template default
827
- * true — better prompt-cache reuse in agent loops). */
828
- chat_template_kwargs?: {
829
- enable_thinking?: boolean;
830
- preserve_thinking?: boolean;
831
- [key: string]: unknown;
832
- };
833
- /** OpenAI reasoning control. For models with a switchable <think> channel
834
- * (Qwen3.5/MiniCPM5) it gates enable_thinking: "none" → off, any level → on.
835
- * This is what Pi sends when the provider advertises reasoning. For
836
- * templates that read a reasoning-depth variable (Qwen3.8: xhigh/medium/
837
- * low), the level ALSO maps into the template via qwenReasoningEffort —
838
- * "xhigh" accepted as a first-class value (Qwen3.8's default/deepest). */
839
- reasoning_effort?: "none" | "minimal" | "low" | "medium" | "high" | "xhigh";
840
- /** Mounted LoRA adapter selection: "id", "a+b" (stacked), or "none". */
841
- adapter?: string;
842
- /** HLG tone-curve sampling override (per request). Snake_case wire fields,
843
- * merged over the server's --hlg-sampling config. docs/design/hlg-sampling.md. */
844
- hlg?: {
845
- enabled?: boolean;
846
- width?: number;
847
- shoulder?: number;
848
- toe?: number;
849
- pivot_offset?: number;
850
- };
851
- /** OpenAI structured output: {type:"json_object"} | {type:"json_schema",
852
- * json_schema:{name,schema,strict?}} | {type:"text"}. Enforced at the
853
- * sampler via xgrammar token-bitmasks (src/grammar.ts). L2-class (oMLX
854
- * oracle). On compile failure, degrades to a system-prompt injection +
855
- * Warning header (oMLX parity), never 500. */
856
- response_format?: unknown;
857
- /** vLLM/oMLX grammar aliases (all compiled via xgrammar): raw EBNF/LARK
858
- * grammar, regex, enum choice, and a bare JSON-schema object. Precedence
859
- * (guided_grammar > json_schema > json_object > structured_outputs >
860
- * guided_regex > guided_choice) mirrors oMLX _effective_guided_grammar. */
861
- guided_grammar?: string;
862
- guided_regex?: string;
863
- guided_choice?: string[];
864
- structured_outputs?: unknown;
865
- }
866
-
867
- /** POST /v1/completions body (mlx_lm.server's raw text completion — no chat
868
- * template). Sampling/penalty/stop fields are the same names as ChatRequest;
869
- * `prompt` replaces `messages`. mlx_lm.server accepts only a STRING prompt
870
- * (it calls `tokenizer.encode(request.prompt)` directly) — token-array
871
- * prompts are rejected there too, so we match. No `echo` (mlx-lm has none). */
872
- type TextCompletionRequest = Omit<ChatRequest, "messages" | "tools" | "tool_choice"> & {
873
- prompt?: unknown;
874
- };
875
279
 
876
280
  export interface ContextAdmissionDecision {
877
281
  /** Effective completion ceiling after admission. */
@@ -880,16 +284,8 @@ export interface ContextAdmissionDecision {
880
284
  clamped: boolean;
881
285
  }
882
286
 
883
- /** Resolve a request's completion upper bound against the admitted context.
884
- *
885
- * `max_tokens` is a ceiling, not a promise that every token will be emitted:
886
- * when a prompt fits but a broad client-wide `max_tokens` would push the
887
- * reservation past the safe context, the bound is CLAMPED to the remaining
888
- * room instead of rejecting an otherwise valid request (the pre-v0.0.13
889
- * behavior 400'd a prompt with thousands of tokens of generation room over a
890
- * 17-token overshoot). This never weakens the OOM guard — generation stops at
891
- * the clamped ceiling, inside the admitted context. Only a prompt that leaves
892
- * no generation slot at all is rejected. */
287
+ /** Resolve a completion cap against an explicitly enforced context limit.
288
+ * Serving keeps fit estimates advisory unless the operator chooses a budget. */
893
289
  export function admitRequestContext(
894
290
  promptTokens: number,
895
291
  requestedMaxTokens: number,
@@ -902,744 +298,7 @@ export function admitRequestContext(
902
298
  return { maxTokens: available, clamped: true };
903
299
  }
904
300
 
905
- /** Per-field default HLG knobs when enabling without specifying them. */
906
- const HLG_DEFAULTS = { width: 4, shoulder: 4, toe: 6, pivotOffset: 6 } as const;
907
-
908
- /** Resolve the effective HLG config: a per-request `hlg` object overrides the
909
- * server's --hlg-sampling default field-by-field. Returns undefined (HLG off)
910
- * unless enabled by the request or the server. */
911
- function resolveHlg(
912
- reqHlg: ChatRequest["hlg"],
913
- serverHlg: HlgConfig | undefined,
914
- ): HlgConfig | undefined {
915
- const enabled = reqHlg?.enabled ?? serverHlg?.enabled ?? false;
916
- if (!enabled) return undefined;
917
- const base = serverHlg ?? HLG_DEFAULTS;
918
- return {
919
- enabled: true,
920
- width: reqHlg?.width ?? base.width,
921
- shoulder: reqHlg?.shoulder ?? base.shoulder,
922
- toe: reqHlg?.toe ?? base.toe,
923
- pivotOffset: reqHlg?.pivot_offset ?? base.pivotOffset,
924
- pivot: "top",
925
- };
926
- }
927
-
928
- interface OpenAIToolCall {
929
- id: string;
930
- type: "function";
931
- function: { name: string; arguments: string };
932
- }
933
-
934
- export type ToolStreamMode = "gemma-sentinel" | "plain" | "buffered-text";
935
-
936
- /** Pick the stream router for a model family. The token-id sentinel router is
937
- * Gemma-4-ONLY: ids 48/49 (<|tool_call>/<tool_call|>) and 100/101
938
- * (<|channel>/<channel|>) are special tokens of that tokenizer family
939
- * (src/tool-call.ts). On every other tokenizer — MiniCPM5, Qwen3/3.5, and
940
- * the Tier-0 generics (llama, qwen2, phi3, …) — those ids are ordinary
941
- * low-id vocab, so the sentinel router would silently swallow output into a
942
- * phantom tool/reasoning segment. Everyone else parses tool calls from
943
- * decoded text (buffered-text; parseGeneratedToolCalls covers the
944
- * OpenAI-JSON <tool_call>, Qwen <function=…>, and MiniCPM5 <function name=…>
945
- * shapes) when tools are present, and streams plain otherwise. Models whose
946
- * markup isn't covered fail soft: the markup stays in content. */
947
- export function selectToolStreamMode(modelType: string, hasTools: boolean): ToolStreamMode {
948
- if (modelType.startsWith("gemma4")) return "gemma-sentinel";
949
- return hasTools ? "buffered-text" : "plain";
950
- }
951
-
952
- /** Routes generated tokens. Gemma uses family-specific sentinel token ids;
953
- * MiniCPM5 and other text-template models use decoded-text parsing so
954
- * ordinary tokenizer ids like "<" are never swallowed globally. Exported for
955
- * unit tests (gemma-channel reasoning split). */
956
- export class ToolAwareStream {
957
- readonly #decoder: StreamDecoder;
958
- #inTool = false;
959
- #toolTokens: number[] = [];
960
- #text = "";
961
- /** Chars of #text already returned as content. */
962
- #sent = 0;
963
- /** Index where tool markup starts; content emission stops there. */
964
- #frozen = -1;
965
- #textToolCalls: OpenAIToolCall[] | null = null;
966
- #textToolParseFailed = false;
967
- readonly toolSegments: number[][] = [];
968
-
969
- /** Gemma reasoning-channel state (gemma-sentinel mode). The model wraps
970
- * chain-of-thought as `<|channel>thought\n…<channel|>` using special tokens
971
- * 100/101 that the content decoder strips, so reasoning is captured here at
972
- * the token level. A SEPARATE decoder keeps the reasoning byte-stream's
973
- * incremental state independent of content's. The `thought` channel-name
974
- * word is stripped before the reasoning text (mlx-lm's think-start marker
975
- * is `<|channel>thought`; the "\n" after it is reasoning content). */
976
- readonly #channelDecoder: StreamDecoder;
977
- #inChannel = false;
978
- #channelNamePending = "";
979
- #channelNameDone = false;
980
- #reasoning = "";
981
-
982
- /** Decoded-text markers that open tool markup (oracle: the streaming
983
- * parser buffers from `<tool_call`/`<function` on, never the whole
984
- * response — content before a tool call still streams live). */
985
- static readonly TOOL_MARKERS = ["<tool_call>", "<function"];
986
-
987
- constructor(
988
- readonly tokenizer: LoadedTokenizer,
989
- readonly mode: ToolStreamMode,
990
- readonly tools: ToolDefinition[] | null,
991
- ) {
992
- this.#decoder = new StreamDecoder(tokenizer, mode !== "buffered-text");
993
- this.#channelDecoder = new StreamDecoder(tokenizer, true);
994
- }
995
-
996
- /** Feed decoded channel text, stripping the leading `thought` channel-name
997
- * word (which may arrive across tokens). Returns the reasoning delta.
998
- * mlx-lm parity: its think-start marker is exactly `<|channel>thought`
999
- * (tokenizer_utils.py `_infer_thinking`) and only the MARKER tokens get
1000
- * their text blanked (server.py `_process_control_tokens`) — the "\n"
1001
- * after the name is an ordinary generated token, so it is the FIRST byte
1002
- * of the reasoning stream and must be kept, not swallowed with the name. */
1003
- #feedChannel(text: string): string {
1004
- if (this.#channelNameDone) return text;
1005
- this.#channelNamePending += text;
1006
- const nl = this.#channelNamePending.indexOf("\n");
1007
- if (nl === -1) return ""; // still inside the channel-name word
1008
- this.#channelNameDone = true;
1009
- const rest = this.#channelNamePending.slice(nl); // keep the "\n" (mlx-lm does)
1010
- this.#channelNamePending = "";
1011
- return rest;
1012
- }
1013
-
1014
- /** Drain reasoning captured since the last call (gemma-channel thinking). */
1015
- takeReasoning(): string {
1016
- const r = this.#reasoning;
1017
- this.#reasoning = "";
1018
- return r;
1019
- }
1020
-
1021
- /** Emit the longest #text prefix that cannot be (the start of) tool
1022
- * markup; hold back ambiguous tails until disambiguated. */
1023
- #textDelta(): string {
1024
- if (this.#frozen >= 0) return "";
1025
- let markerAt = -1;
1026
- for (const mk of ToolAwareStream.TOOL_MARKERS) {
1027
- const i = this.#text.indexOf(mk, this.#sent);
1028
- if (i !== -1 && (markerAt === -1 || i < markerAt)) markerAt = i;
1029
- }
1030
- if (markerAt !== -1) {
1031
- this.#frozen = markerAt;
1032
- const out = this.#text.slice(this.#sent, markerAt);
1033
- this.#sent = markerAt;
1034
- return out;
1035
- }
1036
- let hold = 0;
1037
- for (const mk of ToolAwareStream.TOOL_MARKERS) {
1038
- const max = Math.min(mk.length - 1, this.#text.length - this.#sent);
1039
- for (let k = max; k > hold; k--) {
1040
- if (this.#text.endsWith(mk.slice(0, k))) { hold = k; break; }
1041
- }
1042
- }
1043
- const limit = this.#text.length - hold;
1044
- if (limit <= this.#sent) return "";
1045
- const out = this.#text.slice(this.#sent, limit);
1046
- this.#sent = limit;
1047
- return out;
1048
- }
1049
-
1050
- /** Returns the content text delta for this token ("" while capturing). */
1051
- push(token: number): string {
1052
- if (this.mode !== "gemma-sentinel") {
1053
- this.#text += this.#decoder.push(token);
1054
- if (this.mode === "plain") {
1055
- const out = this.#text.slice(this.#sent);
1056
- this.#sent = this.#text.length;
1057
- return out;
1058
- }
1059
- return this.#textDelta();
1060
- }
1061
- if (this.#inTool) {
1062
- if (token === TOOL_CALL_END) {
1063
- this.#inTool = false;
1064
- this.toolSegments.push(this.#toolTokens);
1065
- this.#toolTokens = [];
1066
- } else {
1067
- this.#toolTokens.push(token);
1068
- }
1069
- return "";
1070
- }
1071
- // Reasoning channel: tokens between <|channel> and <channel|> are thought,
1072
- // captured to #reasoning (drained via takeReasoning), never content. An
1073
- // empty block (<|channel>thought\n<channel|>, emitted by larger Gemmas even
1074
- // with thinking off) yields only the "\n" as reasoning (mlx-lm parity) and
1075
- // leaks nothing into content.
1076
- if (this.#inChannel) {
1077
- if (token === CHANNEL_END) {
1078
- this.#inChannel = false;
1079
- this.#reasoning += this.#feedChannel(this.#channelDecoder.flush());
1080
- } else {
1081
- this.#reasoning += this.#feedChannel(this.#channelDecoder.push(token));
1082
- }
1083
- return "";
1084
- }
1085
- if (token === CHANNEL_START) {
1086
- this.#inChannel = true;
1087
- this.#channelNameDone = false;
1088
- this.#channelNamePending = "";
1089
- return "";
1090
- }
1091
- if (token === TOOL_CALL_START) {
1092
- this.#inTool = true;
1093
- return "";
1094
- }
1095
- return this.#decoder.push(token);
1096
- }
1097
-
1098
- flush(): string {
1099
- if (this.mode !== "gemma-sentinel") {
1100
- this.#text += this.#decoder.flush();
1101
- if (this.mode === "buffered-text") {
1102
- const calls = this.toolCalls();
1103
- if (calls.length && !this.#textToolParseFailed && this.#frozen >= 0) {
1104
- // markup parsed into tool_calls — emit any prose still held
1105
- // before it; the markup itself never reaches content
1106
- const out = this.#text.slice(this.#sent, this.#frozen);
1107
- this.#sent = this.#text.length;
1108
- return out;
1109
- }
1110
- // no tool call (or parse fallback): release everything withheld
1111
- const out = this.#text.slice(this.#sent);
1112
- this.#sent = this.#text.length;
1113
- return out;
1114
- }
1115
- const out = this.#text.slice(this.#sent);
1116
- this.#sent = this.#text.length;
1117
- return out;
1118
- }
1119
- if (this.#inChannel) {
1120
- // truncated mid-reasoning (hit max_tokens); surface the partial thought
1121
- this.#reasoning += this.#feedChannel(this.#channelDecoder.flush());
1122
- this.#inChannel = false;
1123
- return "";
1124
- }
1125
- if (this.#inTool && this.#toolTokens.length) {
1126
- // truncated mid-tool-call (hit max_tokens); surface what we have
1127
- this.toolSegments.push(this.#toolTokens);
1128
- this.#toolTokens = [];
1129
- }
1130
- return this.#decoder.flush();
1131
- }
1132
-
1133
- toolCalls(): OpenAIToolCall[] {
1134
- if (this.mode !== "gemma-sentinel") {
1135
- if (this.#textToolCalls) return this.#textToolCalls;
1136
- try {
1137
- this.#textToolCalls = parseGeneratedToolCalls(this.#text, this.tools ?? []).map((c) => ({
1138
- id: `call_${crypto.randomUUID().slice(0, 8)}`,
1139
- type: "function",
1140
- function: { name: c.name, arguments: JSON.stringify(c.arguments) },
1141
- }));
1142
- } catch {
1143
- this.#textToolParseFailed = true;
1144
- this.#textToolCalls = [];
1145
- }
1146
- return this.#textToolCalls;
1147
- }
1148
- const out: OpenAIToolCall[] = [];
1149
- for (const seg of this.toolSegments) {
1150
- const text = this.tokenizer.decode(seg, false); // keep <|"|> markers
1151
- for (const c of parseToolCalls(text)) {
1152
- out.push({
1153
- id: `call_${crypto.randomUUID().slice(0, 8)}`,
1154
- type: "function",
1155
- function: { name: c.name, arguments: JSON.stringify(c.arguments) },
1156
- });
1157
- }
1158
- }
1159
- return out;
1160
- }
1161
- }
1162
-
1163
- /** Decoded-text stop-sequence matcher with streaming hold-back. Matching
1164
- * on text (not token ids) catches sequences that span token boundaries
1165
- * or tokenize differently in context; current mlx-lm uses token-id
1166
- * state machines and misses those. Text that could be the start of a
1167
- * stop sequence is withheld until disambiguated, so SSE clients never
1168
- * see any part of the stop sequence itself. */
1169
- export class StopMatcher {
1170
- #pending = "";
1171
- stopped = false;
1172
-
1173
- constructor(readonly sequences: string[]) {}
1174
-
1175
- /** Feed a text delta; returns the prefix that is now safe to emit.
1176
- * After a match fires (`stopped`), text before the match is returned
1177
- * and everything from the match on is discarded. */
1178
- push(text: string): string {
1179
- if (this.stopped) return "";
1180
- if (this.sequences.length === 0) return text;
1181
- this.#pending += text;
1182
- // earliest full match wins
1183
- let cut = -1;
1184
- for (const seq of this.sequences) {
1185
- const i = this.#pending.indexOf(seq);
1186
- if (i !== -1 && (cut === -1 || i < cut)) cut = i;
1187
- }
1188
- if (cut !== -1) {
1189
- this.stopped = true;
1190
- const out = this.#pending.slice(0, cut);
1191
- this.#pending = "";
1192
- return out;
1193
- }
1194
- // hold back the longest tail that is a proper prefix of any sequence
1195
- let hold = 0;
1196
- for (const seq of this.sequences) {
1197
- const max = Math.min(seq.length - 1, this.#pending.length);
1198
- for (let k = max; k > hold; k--) {
1199
- if (this.#pending.endsWith(seq.slice(0, k))) {
1200
- hold = k;
1201
- break;
1202
- }
1203
- }
1204
- }
1205
- if (hold === 0) {
1206
- const out = this.#pending;
1207
- this.#pending = "";
1208
- return out;
1209
- }
1210
- const out = this.#pending.slice(0, -hold);
1211
- this.#pending = this.#pending.slice(-hold);
1212
- return out;
1213
- }
1214
-
1215
- /** Generation ended without a match — release any held-back text. */
1216
- flush(): string {
1217
- const out = this.#pending;
1218
- this.#pending = "";
1219
- return out;
1220
- }
1221
- }
1222
-
1223
- /** Split Qwen-style inline <think>...</think> markup into OpenAI reasoning
1224
- * deltas/content. This keeps raw tags out of normal chat text while giving
1225
- * pi (TUI + web) proper thinking_delta events. It is streaming-safe: partial
1226
- * tag prefixes are held until disambiguated. */
1227
- export class ThinkingTagSplitter {
1228
- #pending = "";
1229
- #inThinking: boolean;
1230
- reasoning = "";
1231
- content = "";
1232
-
1233
- /** `startInThinking` seeds the parser INSIDE a <think> block. Needed for
1234
- * templates (Qwen3.5, MiniCPM5) that prime an OPEN `<think>` in the
1235
- * generation prompt when thinking is enabled: the model's output then
1236
- * starts mid-reasoning and emits only the closing `</think>`, never an
1237
- * opening tag. Without this seed the whole chain-of-thought leaks into
1238
- * `content` and the `reasoning` field stays empty. */
1239
- constructor(private readonly enabled: boolean, startInThinking = false) {
1240
- this.#inThinking = startInThinking;
1241
- }
1242
-
1243
- #safePrefixUntilTag(tag: string): string {
1244
- const i = this.#pending.indexOf(tag);
1245
- if (i !== -1) return this.#pending.slice(0, i);
1246
- let hold = 0;
1247
- for (let k = Math.min(tag.length - 1, this.#pending.length); k > 0; k--) {
1248
- if (this.#pending.endsWith(tag.slice(0, k))) { hold = k; break; }
1249
- }
1250
- return this.#pending.slice(0, this.#pending.length - hold);
1251
- }
1252
-
1253
- push(text: string): { content: string; reasoning: string } {
1254
- if (!this.enabled) {
1255
- this.content += text;
1256
- return { content: text, reasoning: "" };
1257
- }
1258
- this.#pending += text;
1259
- let content = "";
1260
- let reasoning = "";
1261
- while (this.#pending) {
1262
- const tag = this.#inThinking ? "</think>" : "<think>";
1263
- const i = this.#pending.indexOf(tag);
1264
- const emit = i === -1 ? this.#safePrefixUntilTag(tag) : this.#pending.slice(0, i);
1265
- if (!emit && i === -1) break;
1266
- if (emit) {
1267
- if (this.#inThinking) reasoning += emit;
1268
- else content += emit;
1269
- this.#pending = this.#pending.slice(emit.length);
1270
- }
1271
- if (i !== -1 && this.#pending.startsWith(tag)) {
1272
- this.#pending = this.#pending.slice(tag.length);
1273
- this.#inThinking = !this.#inThinking;
1274
- continue;
1275
- }
1276
- if (i === -1) break;
1277
- }
1278
- this.content += content;
1279
- this.reasoning += reasoning;
1280
- return { content, reasoning };
1281
- }
1282
-
1283
- flush(): { content: string; reasoning: string } {
1284
- if (!this.enabled) return { content: "", reasoning: "" };
1285
- const out = this.#inThinking
1286
- ? { content: "", reasoning: this.#pending }
1287
- : { content: this.#pending, reasoning: "" };
1288
- this.#pending = "";
1289
- this.content += out.content;
1290
- this.reasoning += out.reasoning;
1291
- return out;
1292
- }
1293
- }
1294
-
1295
- /** True when a rendered prompt ends INSIDE an unclosed `<think>` block — the
1296
- * generation prompt primed reasoning (Qwen3.5 / MiniCPM5 with thinking on),
1297
- * so the model continues the chain-of-thought and emits only the closing
1298
- * `</think>`. Seeds ThinkingTagSplitter so reasoning is split out correctly.
1299
- * Thinking-off primes a CLOSED empty block (`<think>\n\n</think>`), and
1300
- * no-thinking templates have no `<think>` at all — both return false. */
1301
- export function promptEndsInOpenThink(rendered: string): boolean {
1302
- const open = rendered.lastIndexOf("<think>");
1303
- return open !== -1 && open > rendered.lastIndexOf("</think>");
1304
- }
1305
-
1306
- /** Concatenate the text of an OpenAI content-part array, ignoring non-text
1307
- * parts. Tolerant of the part shapes clients actually send: `{type:"text",
1308
- * text}` (OpenAI/pi), `{type:"input_text", text}` (Responses-style), or any
1309
- * part carrying a string `text`. */
1310
- function contentPartsToText(parts: Array<Record<string, unknown>>): string {
1311
- return parts
1312
- .map((p) => (p && typeof p.text === "string" ? p.text : ""))
1313
- .join("");
1314
- }
1315
-
1316
- /** True if any content part is media — image OR audio — so the multimodal
1317
- * path must keep the array form for extractImages/extractAudio. */
1318
- function hasMediaPart(parts: Array<Record<string, unknown>>): boolean {
1319
- return parts.some(
1320
- (p) =>
1321
- p &&
1322
- (p.type === "image" || p.type === "image_url" ||
1323
- p.type === "audio" || p.type === "input_audio" || p.type === "audio_url" ||
1324
- p.type === "video" || p.type === "video_url"),
1325
- );
1326
- }
1327
-
1328
- /** OpenAI sends assistant tool_call arguments as JSON strings; the
1329
- * template renders the object form natively — normalize before render.
1330
- *
1331
- * Two more wire-format → template-format fixes, both "match the format the
1332
- * model's chat template expects":
1333
- *
1334
- * 1. Map the OpenAI reasoning-model "developer" role to "system": pi-ai (and
1335
- * OpenAI's own SDKs) rename the system prompt to `developer` whenever a
1336
- * model advertises reasoning, but our chat templates only know
1337
- * system/user/assistant/tool and raise "Unexpected message role." This is
1338
- * why Qwen3.5/MiniCPM5 chat got no messages while Gemma (non-reasoning, so
1339
- * pi keeps `system`) worked. `developer` IS the system prompt, so the remap
1340
- * is semantics-preserving.
1341
- *
1342
- * 2. Flatten text-only content-part arrays to a plain string. pi (and any
1343
- * OpenAI multimodal client) sends user content as `[{type:"text",text}]`,
1344
- * but non-vision chat templates expect `content` to be a STRING and render
1345
- * nothing for an array — so the user's turn silently vanishes and the model
1346
- * replies "I don't see any message." Arrays that carry a media part
1347
- * (image or audio) are left intact for the multimodal path
1348
- * (extractImages/extractAudio); only text-only arrays are collapsed
1349
- * here, which is a no-op for the vision templates too. */
1350
- export function normalizeMessages(messages: ChatMessage[]): ChatMessage[] {
1351
- return messages.map((raw) => {
1352
- let m = raw.role === "developer" ? { ...raw, role: "system" } : raw;
1353
- if (Array.isArray(m.content) && !hasMediaPart(m.content)) {
1354
- m = { ...m, content: contentPartsToText(m.content) };
1355
- }
1356
- if (!m.tool_calls) return m;
1357
- return {
1358
- ...m,
1359
- tool_calls: m.tool_calls.map((tc) => ({
1360
- ...tc,
1361
- function: {
1362
- ...tc.function,
1363
- arguments:
1364
- typeof tc.function.arguments === "string"
1365
- ? (JSON.parse(tc.function.arguments) as Record<string, unknown>)
1366
- : tc.function.arguments,
1367
- },
1368
- })),
1369
- };
1370
- });
1371
- }
1372
-
1373
- /** Coerce a wire `logit_bias` ({"tokenId": bias}) to numeric keys/values —
1374
- * mlx_lm.server's `{int(k): float(v)}` coercion; throws its exact error
1375
- * message on anything non-numeric (surfaced as a 400). JSON object keys
1376
- * always arrive as strings, hence the coercion. */
1377
- export function parseLogitBias(
1378
- raw: Record<string, number> | undefined | null,
1379
- ): Record<number, number> | undefined {
1380
- if (raw == null) return undefined;
1381
- if (typeof raw !== "object" || Array.isArray(raw))
1382
- throw new Error("logit_bias must be a dict of int to float");
1383
- const out: Record<number, number> = {};
1384
- for (const [k, v] of Object.entries(raw)) {
1385
- const id = Number(k);
1386
- const bias = Number(v);
1387
- if (!Number.isInteger(id) || !Number.isFinite(bias))
1388
- throw new Error("logit_bias must be a dict of int to float");
1389
- out[id] = bias;
1390
- }
1391
- return Object.keys(out).length ? out : undefined;
1392
- }
1393
-
1394
- /** Default seed for a request that didn't pin one. `Date.now()` alone is NOT
1395
- * request-unique: under `--batch N` the batch lane serves ONLY default-seed
1396
- * requests (explicit-seed requests use the serial mechanism, see
1397
- * GenerationGateway.place),
1398
- * so two identical prompts arriving in the same millisecond would share a seed
1399
- * and — with per-row RNG keyed as stepKey(seed, generatedCount) — produce
1400
- * byte-identical completions, silently collapsing best-of-N diversity. Mix a
1401
- * per-process Weyl counter (golden-ratio increment, period 2^32) into the
1402
- * timestamp so every call yields a distinct uint32 within any given ms.
1403
- * Determinism contract unchanged: reproducibility is only promised for an
1404
- * EXPLICIT request seed (`req.seed ?? nextDefaultSeed()` — explicit wins,
1405
- * byte-identical to before); a default seed is fresh entropy per request. */
1406
- let seedWeyl = 0;
1407
- export function nextDefaultSeed(): number {
1408
- seedWeyl = (seedWeyl + 0x9e3779b9) >>> 0;
1409
- return ((Date.now() & 0xffffffff) ^ seedWeyl) >>> 0;
1410
- }
1411
-
1412
- /** mlx_lm.server's logprobs request validation, copied exactly (server.py
1413
- * APIHandler.validate_model_parameters: `_validate("logprobs", bool)` and
1414
- * `_validate("top_logprobs", int, min_val=0, max_val=11, whitelist=[-1])`
1415
- * with defaults logprobs=False / top_logprobs=-1). Returns the reference's
1416
- * error message (→ 400) or null when valid. Note mlx-lm caps top_logprobs
1417
- * at 11, not OpenAI's 20 — we mirror the reference. */
1418
- /** Build a degrade-path system-prompt instruction for JSON output, mirroring
1419
- * oMLX's api.tool_calling.build_json_system_prompt (used when xgrammar
1420
- * compile fails — the response_format degrades to prompt injection rather
1421
- * than a 500, oMLX parity). Returns null for {type:"text"} / unset. */
1422
- export function degradeJsonSystemPrompt(body: ChatRequest): string | null {
1423
- if (body.guided_grammar) {
1424
- return "You must respond with text matching this grammar:\n\n" + body.guided_grammar;
1425
- }
1426
- const rf = body.response_format as { type?: string; json_schema?: { name?: string; description?: string; schema?: unknown } } | undefined;
1427
- if (rf && typeof rf === "object" && rf.type === "json_object") {
1428
- return "You must respond with valid JSON only. " +
1429
- "Do not include any explanation or text outside the JSON object.";
1430
- }
1431
- const schemaSpec = rf && typeof rf === "object" && rf.type === "json_schema"
1432
- ? rf.json_schema
1433
- : body.structured_outputs && typeof body.structured_outputs === "object"
1434
- ? { schema: body.structured_outputs }
1435
- : null;
1436
- if (schemaSpec) {
1437
- const spec = schemaSpec;
1438
- const schema = spec.schema ?? body.structured_outputs ?? {};
1439
- const name = spec.name ?? "response";
1440
- const description = spec.description ?? "";
1441
- let prompt = `You must respond with valid JSON matching the '${name}' schema.`;
1442
- if (description) prompt += ` ${description}`;
1443
- prompt += `\n\nJSON Schema:\n\`\`\`json\n${JSON.stringify(schema, null, 2)}\n\`\`\`\n\n` +
1444
- "Respond with only the JSON object, no additional text or explanation.";
1445
- return prompt;
1446
- }
1447
- if (body.guided_regex) {
1448
- return "You must respond with text matching this regular expression:\n\n" +
1449
- body.guided_regex;
1450
- }
1451
- if (body.guided_choice?.length) {
1452
- return "You must respond with exactly one of these choices:\n\n" +
1453
- body.guided_choice.map((choice) => JSON.stringify(choice)).join("\n");
1454
- }
1455
- return null;
1456
- }
1457
-
1458
- export function applyGrammarDegrade(
1459
- body: ChatRequest,
1460
- degradeHint: string,
1461
- ): { body: ChatRequest; warning: string } {
1462
- const content = degradeJsonSystemPrompt(body) ??
1463
- "Follow the requested output constraint exactly.";
1464
- return {
1465
- warning:
1466
- `grammar not enforced: ${degradeHint} - falling back to prompt injection`,
1467
- body: {
1468
- ...body,
1469
- messages: [{ role: "system", content }, ...body.messages],
1470
- },
1471
- };
1472
- }
1473
-
1474
- export function validateLogprobsParams(body: {
1475
- logprobs?: unknown;
1476
- top_logprobs?: unknown;
1477
- }): string | null {
1478
- const lp = body.logprobs ?? false;
1479
- if (typeof lp !== "boolean") return "logprobs must be of type bool";
1480
- const tl = body.top_logprobs ?? -1;
1481
- if (typeof tl !== "number" || !Number.isInteger(tl))
1482
- return "top_logprobs must be of type int";
1483
- if (tl === -1) return null; // the "unset" whitelist sentinel
1484
- if (tl < 0) return "top_logprobs must be at least 0";
1485
- if (tl > 11) return "top_logprobs must be at most 11";
1486
- return null;
1487
- }
1488
-
1489
- const REASONING_EFFORT_LEVELS =
1490
- ["none", "minimal", "low", "medium", "high", "xhigh"] as const;
1491
-
1492
- /** Validate `reasoning_effort` against the accepted level names. An invalid
1493
- * string must NOT pass through: resolveEnableThinking treats any defined
1494
- * value ≠ "none" as thinking-ON, so a typo like "hihg" silently flipped the
1495
- * thinking channel (and its hotter temperature) on thinking-off models and
1496
- * burned a full default-depth reasoning budget (2026-08-18 review). Same
1497
- * contract style as validateLogprobsParams above. */
1498
- export function validateReasoningEffort(body: {
1499
- reasoning_effort?: unknown;
1500
- }): string | null {
1501
- const v = body.reasoning_effort;
1502
- if (v === undefined || v === null) return null;
1503
- if (typeof v !== "string" || !(REASONING_EFFORT_LEVELS as readonly string[]).includes(v))
1504
- return "reasoning_effort must be one of " +
1505
- "'none', 'minimal', 'low', 'medium', 'high', 'xhigh'";
1506
- return null;
1507
- }
1508
-
1509
- /** Incremental detokenizer: emits the longest stable decoded prefix.
1510
- *
1511
- * Byte parity with mlx-lm's streaming detokenizers (the drop-in contract is
1512
- * rendered BYTES, not just token ids). For BPE/ByteLevel tokenizers, two
1513
- * mlx-lm 0.31.3 BPEStreamingDetokenizer behaviors our full-sequence decode
1514
- * lacks (tokenizer_utils.py:195-226):
1515
- *
1516
- * 1. `trimsLeadingSpace` — mlx-lm drops ONE leading " " at the start of the
1517
- * generated sequence (`_maybe_trim_space`); trim it here. SPM decode
1518
- * already matches (see LoadedTokenizer.trimsLeadingSpace).
1519
- * 2. `bareSpaceTokenId` — add_token WITHHOLDS a single-char byte-32 token
1520
- * ("Ġ") in `_unflushed` ("For single spaces wait until the next token"),
1521
- * flushing it together with the NEXT token — and mlx_lm.server NEVER
1522
- * calls detokenizer.finalize() (zero hits in server.py 0.31.3), so a
1523
- * generation ENDING on bare-space token(s) silently drops their spaces
1524
- * from the served bytes. push() withholds those spaces; flush() drops a
1525
- * trailing bare-space run (the held text dies with the request, exactly
1526
- * like mlx-lm serve). Re-check both if upstream ever adds a finalize()
1527
- * call. LATENT HAZARD (deliberately not emulated): models with
1528
- * clean_up_tokenization_spaces=true get an ADDITIONAL mid-stream rule
1529
- * (`_space_matches`: held space dropped before "." "," "'s" …) — both
1530
- * current BPE targets have it false (MiniCPM5), so it never fires here.
1531
- *
1532
- * Exported for unit tests (serve-detok mlx-lm byte parity). */
1533
- export class StreamDecoder {
1534
- static readonly #WINDOW_TOKENS = 32;
1535
- #ids: number[] = [];
1536
- #emitted = "";
1537
- #warnedRevision = false;
1538
- /** Token index where the bounded decode suffix begins. */
1539
- #windowStart = 0;
1540
- /** Exact decoded text before #windowStart, established by suffix matching. */
1541
- #windowPrefix = "";
1542
- readonly #trimLeadingSpace: boolean;
1543
- readonly #bareSpaceId: number | undefined;
1544
301
 
1545
- constructor(
1546
- readonly tokenizer: LoadedTokenizer,
1547
- readonly skipSpecialTokens = true,
1548
- ) {
1549
- this.#trimLeadingSpace = tokenizer.trimsLeadingSpace === true;
1550
- this.#bareSpaceId = tokenizer.bareSpaceTokenId;
1551
- }
1552
-
1553
- #decode(ids: number[]): string {
1554
- const full = this.#decodeRaw(ids);
1555
- return this.#trimLeadingSpace && full.startsWith(" ") ? full.slice(1) : full;
1556
- }
1557
-
1558
- #decodeRaw(ids: number[]): string {
1559
- return this.tokenizer.decode(ids, this.skipSpecialTokens);
1560
- }
1561
-
1562
- /** Decode only a bounded suffix once an exact text anchor can be proven.
1563
- *
1564
- * At each rebase, decode a candidate suffix independently and accept it only
1565
- * when it is literally the suffix of the exact text decoded so far. This
1566
- * preserves ByteLevel/SPM boundary behavior while reducing steady-state
1567
- * decode work from the whole generation to 32–64 tokens. If no suffix can be
1568
- * anchored (an unusual/global decoder), leave #windowStart at zero and keep
1569
- * the full-history correctness path. */
1570
- #decodeIncremental(): string {
1571
- const full = this.#windowStart === 0
1572
- ? this.#decode(this.#ids)
1573
- : this.#windowPrefix + this.#decodeRaw(this.#ids.slice(this.#windowStart));
1574
-
1575
- if (this.#ids.length - this.#windowStart < StreamDecoder.#WINDOW_TOKENS * 2)
1576
- return full;
1577
-
1578
- const preferred = this.#ids.length - StreamDecoder.#WINDOW_TOKENS;
1579
- const oldest = Math.max(this.#windowStart, preferred - StreamDecoder.#WINDOW_TOKENS);
1580
- for (let start = preferred; start >= oldest; start--) {
1581
- const tail = this.#decodeRaw(this.#ids.slice(start));
1582
- if ((tail.length > 0 || full.length === 0) && full.endsWith(tail)) {
1583
- this.#windowStart = start;
1584
- this.#windowPrefix = full.slice(0, full.length - tail.length);
1585
- break;
1586
- }
1587
- }
1588
- return full;
1589
- }
1590
-
1591
- push(token: number): string {
1592
- this.#ids.push(token);
1593
- // Bare-space hold-back (mlx-lm add_token keeps "Ġ" in _unflushed): don't
1594
- // advance #emitted; the held space(s) flush as part of the next
1595
- // non-bare-space token's delta — consecutive bare spaces accumulate.
1596
- if (token === this.#bareSpaceId) return "";
1597
- const full = this.#decodeIncremental();
1598
- // hold back a trailing replacement char (partial multi-byte sequence)
1599
- const stable = full.endsWith("�") ? full.slice(0, -1) : full;
1600
- if (!stable.startsWith(this.#emitted)) {
1601
- // The decoder revised already-streamed text (cleanup rules — never
1602
- // fires for the shipped tokenizers, which all have
1603
- // clean_up_tokenization_spaces=false; see the LATENT HAZARD note in
1604
- // the class doc). An SSE client cannot un-receive bytes, so the old
1605
- // "re-emit from scratch" answer DUPLICATED the whole stream
1606
- // (2026-07-07 review). Truncate-safe resync instead: emit only the
1607
- // length-extension and keep the emitted watermark monotone — byte
1608
- // drift is confined to the revised span, and the once-per-stream
1609
- // warning makes the tokenizer that needs real _space_matches
1610
- // emulation loud instead of silently wrong.
1611
- if (!this.#warnedRevision) {
1612
- this.#warnedRevision = true;
1613
- console.warn(
1614
- "[detok] decoder revised already-streamed text (clean_up_tokenization_spaces?) — " +
1615
- "resyncing without re-emit; add mlx-lm's cleanup-rule emulation for this tokenizer",
1616
- );
1617
- }
1618
- if (stable.length <= this.#emitted.length) return "";
1619
- const out = stable.slice(this.#emitted.length);
1620
- this.#emitted = stable;
1621
- return out;
1622
- }
1623
- const delta = stable.slice(this.#emitted.length);
1624
- this.#emitted = stable;
1625
- return delta;
1626
- }
1627
-
1628
- flush(): string {
1629
- // mlx_lm.server never finalize()s: a trailing bare-space run stays
1630
- // withheld forever, so its text is dropped from the served bytes.
1631
- let ids = this.#ids;
1632
- if (this.#bareSpaceId !== undefined) {
1633
- let n = ids.length;
1634
- while (n > 0 && ids[n - 1] === this.#bareSpaceId) n--;
1635
- if (n < ids.length) ids = ids.slice(0, n);
1636
- }
1637
- const full = this.#decode(ids);
1638
- const delta = full.slice(this.#emitted.length);
1639
- this.#emitted = full;
1640
- return delta;
1641
- }
1642
- }
1643
302
 
1644
303
  // v2 curve designer: CORS-open so a file:// editor can call a localhost engine.
1645
304
  const CURVE_CORS: Record<string, string> = {
@@ -1654,8 +313,9 @@ function curveJunk(s: string): boolean {
1654
313
  return nonLatin >= 0.02 || /�/.test(s) || /<unused\d+>/.test(s);
1655
314
  }
1656
315
 
316
+
1657
317
  export function createServer(
1658
- ctx: ServerContext, port = 0, serverOptions: ServerOptions = {},
318
+ ctx: ServingContext, port = 0, serverOptions: ServerOptions = {},
1659
319
  ): Server<unknown> {
1660
320
  // --batch N is a concurrency cap: N===1 pins the strict serial executor;
1661
321
  // N>1 admits supported execution compositions to the continuous scheduler.
@@ -1663,7 +323,7 @@ export function createServer(
1663
323
  // Both full-attention (CPM) and
1664
324
  // sliding-window (Gemma) models batch — the scheduler assembles each layer's
1665
325
  // cache by attention type. Non-batchable requests (vision / adapters /
1666
- // user seed / unsupported explicit kv-quant) drain to the serial executor
326
+ // unsupported explicit kv-quant) drain to the serial executor
1667
327
  // (see GenerationGateway.place). No inference setting is rewritten.
1668
328
  // DEFAULT 8 (flipped 2026-07-05, Josh's call, after GATE-B1-SPEED): a
1669
329
  // lone request through the batch lane IS the serial engine (adopted
@@ -1672,7 +332,17 @@ export function createServer(
1672
332
  // only changes behavior when concurrent requests actually arrive — the
1673
333
  // agentic sub-agent workload. --batch 1 pins strict serial for
1674
334
  // arrival-independent numerics. 8 = optiq's Mac-safe concurrency.
335
+ const serving = modelServingBinding(ctx);
336
+ ctx = { ...ctx, serving };
1675
337
  const batch = Math.max(1, Math.floor(serverOptions.batch ?? 8));
338
+ if (serverOptions.generationCheckpointTokens !== undefined) {
339
+ if (!Number.isInteger(serverOptions.generationCheckpointTokens) ||
340
+ serverOptions.generationCheckpointTokens < 1)
341
+ throw new Error("--generation-checkpoint expects a positive integer token interval");
342
+ if (!serverOptions.ssdCacheDir)
343
+ throw new Error("--generation-checkpoint requires --ssd-cache <dir>");
344
+
345
+ }
1676
346
  const defaultGeneratedTokens =
1677
347
  serverOptions.defaultMaxTokens ?? ctx.glmMemoryPlan?.maxGenerationTokens;
1678
348
 
@@ -1694,38 +364,17 @@ export function createServer(
1694
364
  const resolvedKvScheme = resolveKvScheme({
1695
365
  override: serverOptions.kvQuant,
1696
366
  turboQuant: serverOptions.turboQuant,
367
+ quantizedKvStart: serverOptions.quantizedKvStart,
1697
368
  config: ctx.kvConfig,
1698
369
  });
1699
370
  const kvScheme = resolvedKvScheme.generationOptions;
1700
- // Phase 3.1: kvConfig whose layers are all full-attention BATCHES (the
1701
- // scheduler applies the mixed scheme per row); uniform kvBits and configs
1702
- // touching rotating layers still route those requests serial. The warning
1703
- // only fires for the still-serial compositions.
1704
- if (batch > 1 && kvScheme.kvBits)
1705
- console.warn(
1706
- `[batch] --batch ${batch} with uniform --kv-quant ${kvScheme.kvBits}: uniform ` +
1707
- `quantized KV is serial-only (it touches rotating layers) — those requests ` +
1708
- `won't batch. Use --kv-quant config for batched mixed-precision KV, or omit ` +
1709
- `--kv-quant to batch in bf16. (docs/design/unified-engine-frontier-plan.md)`,
1710
- );
1711
- // TurboQuant is solo-only in v1 (novel cache class, not batchable by
1712
- // construction — see GenerationGateway placement's cache-capability gate).
1713
- if (batch > 1 && kvScheme.turboQuant)
1714
- console.warn(
1715
- `[batch] --batch ${batch} with --kv-quant turbo: TurboQuant is serial-only in v1 ` +
1716
- `— those requests won't batch. Omit --kv-quant to batch in bf16. ` +
1717
- `(docs/design/turboquant-kv.md)`,
1718
- );
1719
- // Quantized KV (any axis) excludes the spec lane: drafted requests fall to
1720
- // the normal serial path WITH the KV scheme applied rather than losing it
1721
- // silently (mirrors the affine kvBits/kvConfig exclusion; the spec loop is
1722
- // bf16-KV-only in v1).
371
+ // The model execution plan qualifies the method while retaining KV policy.
1723
372
  if (ctx.draft && (kvScheme.turboQuant || kvScheme.kvBits || kvScheme.kvConfig?.length))
1724
373
  console.warn(
1725
374
  `[spec] --draft-model with quantized KV (--kv-quant ${serverOptions.turboQuant ? "turbo" : String(serverOptions.kvQuant)}): ` +
1726
- `the speculative lane is bf16-KV-only in v1 — requests keep the KV scheme and ` +
1727
- `decode serially WITHOUT speculation. Omit --kv-quant to speculate. ` +
1728
- `(docs/design/dspark-serving-program.md Phase 4)`,
375
+ `requests keep the KV scheme; the model execution plan enables speculation ` +
376
+ `only for a qualified combination. Check response usage for the active method. ` +
377
+ `(docs/design/speculative-decoding.md Phase 4)`,
1729
378
  );
1730
379
  // TurboQuant head-dim fail-fast (2026-07-07 review): the cache class only
1731
380
  // supports {64,128,256,512} (sign-vector + Lloyd-Max table coverage) and
@@ -1740,34 +389,28 @@ export function createServer(
1740
389
  throw new Error(
1741
390
  `--kv-quant turbo: this model's full-attention head_dim is ${dim}; ` +
1742
391
  `TurboQuant supports {${TURBOQUANT_HEAD_DIMS.join(",")}} ` +
1743
- `(docs/design/turboquant-kv.md) — use --kv-quant config|4|8 or omit it`,
392
+ `(docs/design/turboquant.md) — use --kv-quant config|4|8 or omit it`,
1744
393
  );
1745
394
  }
1746
- // Paged KV v1 (docs/design/paged-kv-cache.md): explicit refusals, not
395
+ // Paged KV v1 (docs/design/kv-cache.md): explicit refusals, not
1747
396
  // silent downgrades — the incompatible combos would otherwise degrade
1748
- // quietly (batch: paged caches can't merge; kv-quant: the swap would
1749
- // drop the scheme, the exact composition bug the kvQuant gate above
1750
- // exists to prevent; draft: the spec lane assumes serial-class caches).
397
+ // quietly (kv-quant: the swap would drop the scheme; speculative
398
+ // provider composition with paged storage is not implemented).
1751
399
  if (serverOptions.pagedKv) {
1752
- if (batch > 1)
1753
- throw new Error(
1754
- `--paged-kv is serial-only in v1 — add --batch 1 (got --batch ${batch}). ` +
1755
- `Batched paging is the follow-up PR (docs/design/paged-kv-cache.md).`,
1756
- );
1757
400
  if (kvScheme.kvBits || kvScheme.kvConfig?.length || kvScheme.turboQuant)
1758
401
  throw new Error(
1759
402
  `--paged-kv is bf16-only in v1 — omit --kv-quant (quantized paged ` +
1760
- `blocks are a documented follow-up, docs/design/paged-kv-cache.md).`,
403
+ `blocks are a documented follow-up, docs/design/kv-cache.md).`,
1761
404
  );
1762
405
  if (ctx.draft)
1763
406
  throw new Error(
1764
407
  `--paged-kv cannot combine with --draft-model in v1 ` +
1765
- `(docs/design/paged-kv-cache.md non-goals).`,
408
+ `(docs/design/kv-cache.md non-goals).`,
1766
409
  );
1767
410
  if (!ctx.model.config.modelType.startsWith("gemma4"))
1768
411
  throw new Error(
1769
412
  `--paged-kv v1 supports Gemma4-family models only ` +
1770
- `(this model: ${ctx.model.config.modelType}) — docs/design/paged-kv-cache.md.`,
413
+ `(this model: ${ctx.model.config.modelType}) — docs/design/kv-cache.md.`,
1771
414
  );
1772
415
  const bs = serverOptions.pagedKv.blockSize;
1773
416
  if (bs !== undefined && (!Number.isInteger(bs) || bs <= 0))
@@ -1779,11 +422,12 @@ export function createServer(
1779
422
  );
1780
423
  }
1781
424
 
1782
- // SSD cold tier (docs/design/ssd-kv-cold-tier.md): prefix KV survives RAM
425
+ // SSD cold tier (docs/design/kv-cache.md): prefix KV survives RAM
1783
426
  // eviction and restarts. Compatibility key = configFingerprint (graph
1784
427
  // shape) + the EFFECTIVE kv scheme (flags pick bf16 vs config vs uniform
1785
428
  // on the same model — restored caches must match what serving produces) +
1786
- // tokenizer hash (ids must keep meaning the same text).
429
+ // tokenizer hash (ids must keep meaning the same text). The model binding
430
+ // also identifies the runtime numerics that produced the persisted state.
1787
431
  // RAM prompt-cache cap: 8 GB default (Josh's call, 2026-07-06). The old
1788
432
  // flat 2e9 was an anti-OOM reflex sized for a 1B model in 24 GB — a
1789
433
  // single full-context 12B entry is ~0.6 GB, so it silently shrank to
@@ -1791,24 +435,33 @@ export function createServer(
1791
435
  // entries still demote to SSD and LRU eviction bounds pressure.
1792
436
  // --prompt-cache <GB> overrides; 0 disables.
1793
437
  const promptCacheCap = serverOptions.promptCacheBytes ?? 8e9;
438
+ const stateCodecs = ctx.stateCodecs ?? legacyCacheCodecs;
439
+ const cloneState = (caches: Cache[]) => cloneKvCaches(caches, stateCodecs);
1794
440
  let ssdStore: SsdCacheStore | null = null;
1795
441
  if (serverOptions.ssdCacheDir) {
1796
442
  if (promptCacheCap <= 0)
1797
443
  throw new Error("--ssd-cache requires the RAM prompt cache (--prompt-cache 0 disables it)");
444
+ if (typeof serving.stateCompatibility !== "string" || !serving.stateCompatibility.length)
445
+ throw new Error("--ssd-cache requires the model binding's stateCompatibility identity");
1798
446
  const schemeKey = resolvedKvScheme.cacheKey;
447
+ const stateKey = Bun.hash(serving.stateCompatibility).toString(16);
1799
448
  const tokJson = readFileSync(`${ctx.model.config.modelDir}/tokenizer.json`);
1800
449
  ssdStore = new SsdCacheStore({
450
+ codecs: stateCodecs,
1801
451
  dir: serverOptions.ssdCacheDir,
1802
- maxBytes: serverOptions.ssdCacheMaxBytes ?? 32 * 2 ** 30,
1803
- configFingerprint: `${configFingerprint(ctx.model.config)}-${schemeKey}`,
452
+ maxBytes: serverOptions.ssdCacheMaxBytes ?? Number.POSITIVE_INFINITY,
453
+ configFingerprint: `${configFingerprint(ctx.model.config)}-${schemeKey}-${stateKey}`,
1804
454
  tokenizerHash: Bun.hash(tokJson).toString(16),
1805
455
  modelId: ctx.modelId,
1806
456
  verify: serverOptions.ssdCacheVerify,
1807
457
  });
1808
458
  const recovered = ssdStore.scan();
459
+ const capacity = Number.isFinite(ssdStore.maxBytes)
460
+ ? `${(ssdStore.maxBytes / 2 ** 30).toFixed(0)} GiB cap`
461
+ : "unlimited";
1809
462
  console.log(
1810
463
  `[ssd-cache] ${serverOptions.ssdCacheDir} — ${recovered} entr${recovered === 1 ? "y" : "ies"} recovered, ` +
1811
- `${(ssdStore.totalBytes / 2 ** 30).toFixed(2)} GiB of ${(ssdStore.maxBytes / 2 ** 30).toFixed(0)} GiB cap`,
464
+ `${(ssdStore.totalBytes / 2 ** 30).toFixed(2)} GiB, ${capacity}`,
1812
465
  );
1813
466
  }
1814
467
 
@@ -1823,16 +476,18 @@ export function createServer(
1823
476
  return h ? { prefixLen: h.prefixLen, handle: h.entry } : null;
1824
477
  },
1825
478
  restore: (handle: unknown) => {
1826
- const loaded = ssdStore!.restore(handle as import("./ssd-cache").SsdIndexEntry, ctx.model);
479
+ const loaded = serving.restore(ssdStore!, handle as import("./ssd-cache").SsdIndexEntry);
1827
480
  if (!loaded) return null;
1828
481
  // Restore is a STREAMED COPY (2026-07-07 A7-restore): the caches
1829
482
  // own their bytes and no mapping outlives loadKvCache — nothing
1830
483
  // to pin, nothing to unmap. retain stays in the entry contract
1831
484
  // as a no-op so callers' dispose ordering is unchanged.
1832
- return { tokens: loaded.tokens, caches: loaded.caches, retain: () => {} };
485
+ return { tokens: loaded.tokens, caches: loaded.caches, attachments: loaded.attachments, retain: () => {} };
1833
486
  },
1834
- store: (tokens: number[], caches: import("./model/gemma4").Cache[], ns: string) => {
1835
- ssdStore!.store(tokens, caches, ns);
487
+ store: (tokens: number[], caches: import("./model/gemma4").Cache[], ns: string,
488
+ attachments?: import("./backends/mlx/checkpoint-state").CheckpointAttachment[]) => {
489
+ if (ssdStore!.hasDurablePrefix(tokens, ns)) return true;
490
+ return ssdStore!.store(tokens, caches, ns, attachments);
1836
491
  },
1837
492
  }
1838
493
  : null;
@@ -1843,7 +498,7 @@ export function createServer(
1843
498
  // cache hands us OWNED zero-copy clones + copied tokens (made under
1844
499
  // the generation lock — microseconds; entries are immutable so the
1845
500
  // clones stay consistent), and the flush goes through the BOUNDED
1846
- // SpillQueue -> storeAsync (idle-gated per tensor, see below) ->
501
+ // SpillQueue -> storeAsync (generation-locked per tensor, see below) ->
1847
502
  // clone disposal on every settle/drop path (that dispose is what
1848
503
  // actually frees the demoted GPU memory; the queue's byte cap keeps
1849
504
  // starved-gate retention bounded — 2026-07-07 review fix). spillQueue
@@ -1853,10 +508,11 @@ export function createServer(
1853
508
  ssdStore
1854
509
  ? {
1855
510
  spillOwned: (entry) =>
1856
- spillQueue!.enqueue({ tokens: entry.tokens, caches: entry.caches, ns: entry.ns }),
511
+ spillQueue!.enqueue({ tokens: entry.tokens, caches: entry.caches, ns: entry.ns, attachments: entry.attachments }),
1857
512
  }
1858
513
  : null,
1859
514
  coldTier,
515
+ cloneState,
1860
516
  );
1861
517
  // Responses-API store for previous_response_id resumption (Phase 11):
1862
518
  // TTL + byte-capped LRU, port of optiq/response_store.py. Pairs with
@@ -1864,175 +520,38 @@ export function createServer(
1864
520
  // so its KV prefill is already cached.
1865
521
  const responseStore = new ResponseStore();
1866
522
 
1867
- /** Run one generation with prompt-cache reuse. Must be called under the
1868
- * gateway's exclusive lock — the gateway invokes it on the serial lane;
1869
- * the curve /generate endpoint wraps it in gateway.runExclusive. (The old
1870
- * `enqueue` promise queue is gone; the gateway's AsyncMutex is THE lock.)
1871
- * onToken returning `false` halts generation early (stop
1872
- * sequence fired); the cache snapshot stays valid and is still kept.
1873
- * Media (vision/audio embeddings-prefill) requests bypass the prompt
1874
- * cache: soft tokens are identical placeholder ids, so prefix matching
1875
- * across different images/clips would false-hit. */
1876
- const runGeneration = async (
1877
- promptIds: number[],
1878
- options: GenerateOptions,
1879
- onToken: (token: number, logprobs?: TokenLogprobs) => void | boolean | Promise<void | boolean>,
1880
- vision?: {
1881
- embeddings: import("./mlx/array").MlxArray;
1882
- imageMask?: import("./mlx/array").MlxArray;
1883
- multimodalMask?: import("./mlx/array").MlxArray;
1884
- mrope?: MropeRequestState;
1885
- },
1886
- trace?: PromptResponseTrace,
1887
- ) => {
1888
- // Qwen vision: install the request's mRoPE state for every forward of
1889
- // this serial run (prefill AND decode use the 3D interleaved positions +
1890
- // delta). Scoped inside the try below — under the serial mutex — so
1891
- // queued text requests never see it and a throw between here and the
1892
- // try (makeCache under memory pressure, spec-path early return) can't
1893
- // leave a DEAD request's positions installed for the next generation
1894
- // (2026-08-18 review).
1895
- // Speculative decoding (serve --draft-model): spec-ELIGIBLE requests
1896
- // decode through the verify loop; the rest fall through to the normal
1897
- // serial path (never wrong results, just no speedup — logged once per
1898
- // combination class would be noise, so silent). Eligibility v1:
1899
- // text-only, base weights, no logprobs capture, bf16 KV — ALL quantized
1900
- // KV axes excluded (affine kvBits/kvConfig AND turboQuant; the spec loop
1901
- // builds fresh bf16 caches and never calls maybeQuantizeKv, so routing a
1902
- // turbo request here would silently drop the operator's KV scheme —
1903
- // 2026-07-07 review). Grammar COMPOSES (Phase C constrained verify walk
1904
- // — see the serve-loop.ts header). Prompt-cache reuse is bypassed on
1905
- // the spec path v1.
1906
- if (
1907
- ctx.draft &&
1908
- !vision &&
1909
- !options.adapters?.length &&
1910
- !options.logprobs &&
1911
- !options.kvBits &&
1912
- !options.kvConfig &&
1913
- !options.turboQuant &&
1914
- // Belt: --paged-kv + --draft-model is refused at createServer; this
1915
- // keeps the spec lane paged-free even for programmatic callers.
1916
- !options.pagedKv
1917
- ) {
1918
- const { specServeRun } = await import("./spec/serve-loop");
1919
- return withModelWiredLimit(
1920
- ctx.model,
1921
- () => withModelUsageFlush(
1922
- ctx.model,
1923
- () => specServeRun(
1924
- ctx.model, ctx.draft!.provider, ctx.draft!.numDraftTokens,
1925
- promptIds, options, onToken,
1926
- ),
1927
- ),
1928
- );
1929
- }
1930
- // Cache entries are adapter-specific: KV computed under one adapter
1931
- // must never seed another's (or the base's) prefill.
1932
- const cacheNs = options.adapters?.join("+") ?? "";
1933
- // Paged-KV request scope (docs/design/paged-kv-cache.md): media
1934
- // prompts (bidir overlay) and LoRA-adapter requests are v1 non-goals —
1935
- // they run the PLAIN cache path even under --paged-kv (scope the flag
1936
- // per request, never 400). Effective value computed ONCE so the
1937
- // prompt-cache bypass below and the generate() options can't disagree.
1938
- const pagedKv = vision || options.adapters?.length ? undefined : options.pagedKv;
1939
- // Paged requests bypass the prompt cache entirely (v1 non-goal:
1940
- // PagedKVCache has no cloneKvCaches/restore path — the vision
1941
- // precedent). Fresh caches per request, disposed on completion.
1942
- const skipPromptCache = Boolean(vision || pagedKv);
1943
- // Both tiers in one call (Layer 0): take() prefers a strictly-longer
1944
- // SSD prefix, restores it zero-copy, and trims — see PromptCache.take.
1945
- const closeCacheLookup = trace?.begin("cache.lookup_restore", {
1946
- mechanism: "serial",
1947
- bypassed: skipPromptCache,
1948
- });
1949
- const entry = skipPromptCache ? null : promptCache.take(promptIds, cacheNs);
1950
- closeCacheLookup?.();
1951
- const caches = entry?.caches ?? ctx.model.makeCache();
1952
- // Prompt-boundary snapshot (the multi-turn agent fix, 2026-07-04): the
1953
- // prompt+gen entry put() below is UNTRIMMABLE at context > sliding
1954
- // window (wrapped rings) and under quantized KV (mid-group), so any
1955
- // decode→encode roundtrip drift in the reply the client sends back
1956
- // turns the next turn into a total miss (measured: 12B turn-2 TTFT
1957
- // 8.9 s instead of ~0.2 s). A prompt-ONLY entry is always an exact
1958
- // prefix of the next turn's rendering regardless of reply drift.
1959
- // Zero-copy (cloneKvCaches = slice views); only for substantial cold
1960
- // prefills, where the re-prefill it saves is worth an extra entry.
1961
- // The oracle invariant (mlx-lm insert_segments): a trim-free STRICT
1962
- // prefix of the prompt exists for EVERY substantial request — cap the
1963
- // boundary at len-1 so even a stableLen == len prompt (e4b: the
1964
- // template tail survives the probe render) snapshots prompt[:-1]. An
1965
- // exact repeat then matches with trimNeeded == 0, bypassing
1966
- // isTrimmable() entirely — the only reuse path a wrapped ring has.
1967
- const boundary = Math.min(options.snapshotAt ?? promptIds.length, promptIds.length - 1);
1968
- // Re-snapshot on EVERY substantial request whose stable boundary extends
1969
- // past the cached prefix; the clone is zero-copy views, so re-putting
1970
- // is ~free.
1971
- const snapshotBoundary =
1972
- !skipPromptCache && boundary >= 256 && boundary > (entry?.tokens.length ?? 0);
1973
- try {
1974
- if (vision?.mrope && ctx.model instanceof Qwen35Model)
1975
- ctx.model.mrope = vision.mrope;
1976
- const gen = generate(ctx.model, promptIds, {
1977
- ...options,
1978
- pagedKv, // request-scoped (undefined strips the server-wide flag)
1979
- cache: caches,
1980
- ...(snapshotBoundary
1981
- ? {
1982
- // snapshotAt MUST travel with the hook: generate() splits the
1983
- // prefill at exactly this many tokens and fires the hook while
1984
- // the caches hold exactly that prefix — putting boundary
1985
- // tokens against caches at any other offset is silent KV
1986
- // corruption.
1987
- snapshotAt: boundary,
1988
- onPrefillDone: () => {
1989
- try {
1990
- promptCache.put(promptIds.slice(0, boundary), cloneKvCaches(caches), cacheNs);
1991
- } catch (e) {
1992
- console.warn(`prompt-boundary snapshot skipped: ${(e as Error).message}`);
1993
- }
1994
- },
1995
- }
1996
- : {}),
1997
- ...(vision
1998
- ? {
1999
- promptEmbeddings: vision.embeddings,
2000
- ...(vision.imageMask ? { imageMask: vision.imageMask } : {}),
2001
- ...(vision.multimodalMask ? { multimodalMask: vision.multimodalMask } : {}),
2002
- }
2003
- : {}),
2004
- }, { trace, mechanism: "serial" });
2005
- for await (const t of gen) {
2006
- if ((await onToken(t.token, t.logprobs)) === false) break;
2007
- }
2008
- const s = gen.stats!; // set on completion AND on early break
2009
- if (skipPromptCache) {
2010
- // Vision and paged-KV requests own their caches for exactly one
2011
- // generation (paged: v1 non-goal — no PromptCache integration).
2012
- for (const c of caches) c.dispose();
2013
- } else {
2014
- // put() fires onPut → the debounced write-behind SSD snapshot
2015
- // (wired below), covering the batch lane's puts too.
2016
- promptCache.put(s.cacheTokens, caches, cacheNs, entry?.retain);
2017
- }
2018
- return s;
2019
- } catch (e) {
2020
- for (const c of caches) c.dispose();
2021
- entry?.retain?.();
2022
- // A throw BEFORE generate()/the gateway took ownership (promptCache
2023
- // take / makeCache) would leak the grammar's WASM matcher; dispose()
2024
- // is idempotent, so this is safe when the throw came from inside run
2025
- // (whose finally already disposed it).
2026
- options.grammar?.dispose();
2027
- throw e;
2028
- } finally {
2029
- if (ctx.model instanceof Qwen35Model) ctx.model.mrope = null;
2030
- vision?.embeddings.dispose();
2031
- vision?.imageMask?.dispose();
2032
- vision?.multimodalMask?.dispose();
2033
- options.visionPixels?.dispose();
2034
- }
523
+ // Memory estimates remain available in diagnostics. Only an explicit
524
+ // operator budget or context cap constrains request planning.
525
+ const memoryAdmission = ctx.glmMemoryPlan ?? fit(
526
+ ctx.model.config, ctx.model.weightsBytes, 1,
527
+ undefined, undefined, 0, serverOptions.memoryBudgetBytes,
528
+ resolvedKvScheme.fitOptions,
529
+ );
530
+ const profileContext = runtimeValue("MLX_BUN_RD_CONTEXT_LIMIT");
531
+ const profileLimit = profileContext === undefined ? null : Number(profileContext);
532
+ if (profileLimit !== null && (!Number.isSafeInteger(profileLimit) || profileLimit < 1))
533
+ throw new Error("MLX_BUN_RD_CONTEXT_LIMIT must be a positive integer");
534
+ const budgetLimit = serverOptions.memoryBudgetBytes !== undefined
535
+ ? memoryAdmission.maxSafeContext : ctx.glmMemoryPlan?.contextTokens ?? null;
536
+ const contextLimit = profileLimit === null ? budgetLimit
537
+ : Math.min(profileLimit, budgetLimit ?? Infinity);
538
+ const admission = memoryAdmission;
539
+ const allocatorLimit =
540
+ admission.allocatorLimitBytes ?? serverOptions.memoryBudgetBytes;
541
+ if (allocatorLimit) setMemoryLimit(allocatorLimit);
542
+
543
+ let checkpointPersistence: ContinuationPersistence | undefined;
544
+ const executionServices: import("./backends/mlx/serial-executor").MlxSerialServices = {
545
+ get checkpointPersistence() { return checkpointPersistence; },
546
+ promptCache, checkpoints: ssdStore,
547
+ checkpointEveryTokens: serverOptions.generationCheckpointTokens,
548
+ identity: { artifact: ctx.profile.artifact, implementation: ctx.profile.profile.execution,
549
+ stateAbi: "legacy-cache-array-v1", codecs: stateCodecs.id },
550
+ adapterNamespace: (adapters) => ctx.adapters.cacheNamespace(adapters),
551
+ cloneState,
2035
552
  };
553
+ serving.gateway.configureContinuation?.(executionServices);
554
+ const runGeneration = serving.createSerial(executionServices);
2036
555
 
2037
556
  // Write-behind persistence (restart survival — the oMLX boundary-snapshot
2038
557
  // idea at whole-entry granularity, spill-on-evict alone can't survive a
@@ -2062,17 +581,17 @@ export function createServer(
2062
581
  // pacing interleaved those slices exactly between decode tokens. A ~16k
2063
582
  // entry's flush overlapping the bench's ctx repeats depressed decode@ctx
2064
583
  // ~9% on e4b (mlx-lm runs no equivalent background work). Now every step
2065
- // — including the first — awaits gateway.onIdle() (ssdFlushGate), so the
584
+ // — including the first — runs through gateway.runWhenIdle(), so the
2066
585
  // flush only progresses while NOTHING is generating and pauses when a
2067
586
  // request arrives mid-flush. Tradeoffs, accepted: durability waits for a
2068
587
  // quiet moment (single-user serving quiesces constantly; sustained
2069
588
  // hammering defers the flush AND the spill clones' GPU-memory release —
2070
- // bounded by the chain), and one in-flight tensor step (~10-15 MB) can
2071
- // still land ahead of a just-arrived request. MLX_BUN_SSD_WRITEBEHIND=0
589
+ // bounded by the chain). MLX_BUN_SSD_WRITEBEHIND=0
2072
590
  // disables write-behind snapshots entirely — the paired-A/B lever + kill
2073
591
  // switch (restart survival then degrades to spill-on-evict only).
2074
592
  const writeBehindOn = runtimeValue("MLX_BUN_SSD_WRITEBEHIND") !== "0";
2075
- const ssdFlushGate = (): Promise<void> => gateway.onIdle();
593
+ const ssdFlushStep = <T>(step: () => T): Promise<T> =>
594
+ gateway.runWhenIdle(async () => step());
2076
595
  // Bounded write-behind queue (2026-07-07 review fix — see SpillQueue in
2077
596
  // kv-store.ts): pending clones pin their entries' GPU buffers while the
2078
597
  // idle gate starves under sustained traffic, so QUEUED bytes are capped —
@@ -2087,19 +606,28 @@ export function createServer(
2087
606
  // never drops the item just enqueued).
2088
607
  // The gateway exists before the spill queue because the SSD writer uses
2089
608
  // its idle boundary to avoid competing with generation.
2090
- const gateway = new GenerationGateway(ctx.model, batch, runGeneration, {
609
+ const gateway = new GenerationGateway(serving.gateway, batch, runGeneration, {
2091
610
  kvBudgetBytes: serverOptions.kvBudgetBytes,
611
+ checkpoints: !!(serverOptions.generationCheckpointTokens && ssdStore),
612
+ stateCodecs,
2092
613
  kvScheme: resolvedKvScheme,
2093
614
  promptCache,
615
+ adapterNamespace: (adapters) => ctx.adapters.cacheNamespace(adapters),
2094
616
  });
2095
617
  const spillQueueGbRaw = Number(runtimeValue("MLX_BUN_SSD_SPILL_QUEUE_GB"));
2096
618
  const spillQueueCapBytes =
2097
619
  (Number.isFinite(spillQueueGbRaw) && spillQueueGbRaw >= 0 ? spillQueueGbRaw : 2) * 1024 ** 3;
620
+ // Both queues share the configured total: reserve half for checkpoints
621
+ // only when that service exists. Each retains the existing soft-cap rule.
622
+ const checkpointQueueBytes = ssdStore && serverOptions.generationCheckpointTokens
623
+ ? Math.floor(spillQueueCapBytes / 2) : 0;
624
+ if (ssdStore && serverOptions.generationCheckpointTokens)
625
+ checkpointPersistence = new ContinuationPersistence(ssdStore, { maxBytes: checkpointQueueBytes, runStep: ssdFlushStep });
2098
626
  const spillQueue = ssdStore
2099
627
  ? new SpillQueue(
2100
- spillQueueCapBytes,
628
+ spillQueueCapBytes - checkpointQueueBytes,
2101
629
  cacheBytes,
2102
- (item) => ssdStore!.storeAsync(item.tokens, item.caches, item.ns, ssdFlushGate),
630
+ (item) => ssdStore!.storeAsync(item.tokens, item.caches, item.ns, ssdFlushStep, item.attachments),
2103
631
  (caches) => { for (const c of caches) c.dispose(); },
2104
632
  )
2105
633
  : null;
@@ -2108,28 +636,35 @@ export function createServer(
2108
636
  gateway,
2109
637
  promptCache,
2110
638
  spillQueue,
2111
- cloneKvCaches,
639
+ cloneState,
2112
640
  (tokens, ns) => ssdStore.hasDurablePrefix(tokens, ns),
2113
641
  )
2114
642
  : null;
2115
- const durabilityStats = (): DurabilitySnapshotStats =>
2116
- durability?.stats ?? {
643
+ const durabilityStats = (): DurabilitySnapshotStats => {
644
+ const stats = durability?.stats ?? {
2117
645
  pendingSnapshots: 0,
2118
646
  pendingSpills: spillQueue?.pendingCount ?? 0,
2119
647
  pendingSpillBytes: spillQueue?.pendingBytes ?? 0,
2120
648
  droppedSpills: spillQueue?.droppedCount ?? 0,
2121
649
  failedSpills: spillQueue?.failedCount ?? 0,
2122
650
  };
651
+ const checkpoints = checkpointPersistence?.stats;
652
+ return { ...stats, pendingSpills: stats.pendingSpills + (checkpoints?.pendingCount ?? 0),
653
+ pendingSpillBytes: stats.pendingSpillBytes + (checkpoints?.pendingBytes ?? 0),
654
+ droppedSpills: stats.droppedSpills + (checkpoints?.dropped ?? 0),
655
+ failedSpills: stats.failedSpills + (checkpoints?.failed ?? 0) };
656
+ };
2123
657
  const flushDurability = async (): Promise<DurabilityFlushResult> => {
2124
- if (durability) return durability.flush();
2125
658
  const started = performance.now();
2126
- await spillQueue?.drain();
659
+ const checkpoints = await checkpointPersistence?.flush();
660
+ const result = durability ? await durability.flush() : null;
661
+ if (!result) await spillQueue?.drain();
2127
662
  const stats = durabilityStats();
2128
663
  return {
2129
664
  ...stats,
2130
- durable: stats.pendingSpills === 0,
2131
- flushedSnapshots: 0,
2132
- missingSnapshots: 0,
665
+ durable: (result?.durable ?? stats.pendingSpills === 0) && (checkpoints?.durable ?? true) && stats.pendingSpills === 0,
666
+ flushedSnapshots: result?.flushedSnapshots ?? 0,
667
+ missingSnapshots: result?.missingSnapshots ?? 0,
2133
668
  elapsedMs: performance.now() - started,
2134
669
  };
2135
670
  };
@@ -2150,7 +685,7 @@ export function createServer(
2150
685
  let demoteTimer: ReturnType<typeof setInterval> | null = null;
2151
686
  if (ssdStore && demoteIdleMs > 0) {
2152
687
  demoteTimer = setInterval(() => {
2153
- if (gateway.activeRows > 0 || gateway.pendingRows > 0) return;
688
+ if (gateway.busy) return;
2154
689
  void gateway.runExclusive(async () => {
2155
690
  const n = promptCache.demoteIdle(demoteIdleMs);
2156
691
  if (n > 0) console.log(`[ssd-cache] demoted ${n} idle entr${n === 1 ? "y" : "ies"} to disk`);
@@ -2159,45 +694,8 @@ export function createServer(
2159
694
  demoteTimer.unref?.();
2160
695
  }
2161
696
 
2162
- const completionExecutor = new CompletionExecutor(gateway);
2163
- const completionProtocolUsage = (usage: CompletionSummary["usage"]) => ({
2164
- prompt_tokens: usage.promptTokens,
2165
- completion_tokens: usage.completionTokens,
2166
- total_tokens: usage.totalTokens,
2167
- prompt_tokens_details: { cached_tokens: usage.cachedTokens },
2168
- ...(usage.speculation ? { speculation: usage.speculation } : {}),
2169
- });
2170
- const completionUsage = (summary: CompletionSummary) => ({
2171
- ...completionProtocolUsage(summary.usage),
2172
- lane: summary.lane,
2173
- });
2174
-
2175
- // Admission ceiling, resolved once (Phase 5 memoryBudget enforcement).
2176
- // fit() solves max safe context from weights + KV growth + prefill
2177
- // transient. The active kv-quant scheme (uniform kvBits / per-layer
2178
- // kvConfig) is billed at its true bytes/element so a quantized cache
2179
- // advertises and admits the larger window it actually enables; only
2180
- // TurboQuant still bills bf16 (conservative — no projector for its
2181
- // layout yet, and it is solo-only in v1).
2182
- const admission = ctx.glmMemoryPlan ?? fit(
2183
- ctx.model.config, ctx.model.weightsBytes, 1,
2184
- undefined, undefined, 0, serverOptions.memoryBudgetBytes,
2185
- resolvedKvScheme.fitOptions,
2186
- );
2187
- // A zero ceiling is a warning, never a startup refusal: killing the
2188
- // server here can only parrot the per-request admission message (which
2189
- // still fires, with this same ceiling) or be a false positive from the
2190
- // fit model itself — it can never save anything the request gate
2191
- // doesn't. Serve until physically incapable.
2192
- if (admission.maxSafeContext < 1)
2193
- console.warn(
2194
- `[admission] memory budget ${(admission.usableBytes / 1e9).toFixed(2)} GB leaves no ` +
2195
- `safe context for ${ctx.modelId} (weights ${(ctx.model.weightsBytes / 1e9).toFixed(2)} GB) ` +
2196
- `— serving anyway; generation requests will be refused until the budget is raised`,
2197
- );
2198
- const allocatorLimit =
2199
- admission.allocatorLimitBytes ?? serverOptions.memoryBudgetBytes;
2200
- if (allocatorLimit) setMemoryLimit(allocatorLimit);
697
+ const sessionEngine = createSessionCompletionEngine(gateway, disposeUnstartedRequest);
698
+ const completionExecutor = new CompletionExecutor(sessionEngine);
2201
699
 
2202
700
  // /library response cache (30 s) — registry + config reads only.
2203
701
  const startedAt = Date.now();
@@ -2224,207 +722,37 @@ export function createServer(
2224
722
  return jobStore;
2225
723
  };
2226
724
 
2227
- // (kvScheme is resolved earlier, above the prompt-cache construction —
2228
- // the SSD tier's compatibility fingerprint folds the effective scheme in.)
2229
-
2230
- // XTC never removes EOS or the newline token: mlx_lm.server passes
2231
- // [tokenizer.eos_token_id, tokenizer.encode("\n")] as xtc_special_tokens.
2232
- // We pass the flat equivalent — ALL configured EOS ids plus the bare newline
2233
- // id(s), encoded WITHOUT special tokens (mlx-lm's encode("\n") can drag a
2234
- // BOS along on some tokenizers; protecting BOS from XTC is a no-op, so the
2235
- // flat/no-BOS form is behaviorally identical and cleaner).
2236
- const xtcSpecialTokens = [
2237
- ...ctx.model.config.eosTokenIds,
2238
- ...ctx.tokenizer.encode("\n", false),
2239
- ];
2240
-
2241
- // Effective enable_thinking for a request, with the same precedence the chat
2242
- // template uses (extracted so sampling and template rendering can't disagree):
2243
- // explicit chat_template_kwargs.enable_thinking → reasoning_effort ("none" =
2244
- // off) → server --thinking default → model default (MiniCPM5 → off). undefined
2245
- // means "not a switchable-thinking model / leave the template default".
2246
- const resolveEnableThinking = (req: ChatRequest): boolean | undefined => {
2247
- const explicit = req.chat_template_kwargs?.enable_thinking;
2248
- if (typeof explicit === "boolean") return explicit;
2249
- const effort = req.reasoning_effort;
2250
- if (effort !== undefined) return effort !== "none";
2251
- if (serverOptions.defaultThinking !== undefined) return serverOptions.defaultThinking;
2252
- return isMiniCPM5Config(ctx.model.config) ? false : undefined;
2253
- };
2254
-
2255
- const toOptions = (req: ChatRequest): GenerateOptions & { stopSequences: string[] } => {
2256
- // Sampling follows the thinking state — which the web UI's thinking button
2257
- // drives via enable_thinking. Model authors publish a SINGLE
2258
- // generation_config temperature (the think-mode value) but recommend a
2259
- // cooler one for direct, no-think replies (MiniCPM5 card: 0.9 think / 0.7
2260
- // no-think, top_p 0.95 for both). So with no explicit temperature set, a
2261
- // no-think turn runs at most NO_THINK_TEMPERATURE while a think turn keeps
2262
- // the model's hotter configured default. An explicit request/CLI
2263
- // temperature always wins. top_p is unchanged by mode.
2264
- const NO_THINK_TEMPERATURE = 0.7;
2265
- const genTemp = ctx.genDefaults.temperature ?? 0.7;
2266
- const defaultTemp = resolveEnableThinking(req) === false
2267
- ? Math.min(genTemp, NO_THINK_TEMPERATURE)
2268
- : genTemp;
2269
- return {
2270
- // A thinking turn emits <think>…</think> AND the answer, so a tight cap
2271
- // truncates the visible reply. Default very generously (the model's
2272
- // context is far larger); only an explicit max_tokens narrows it. The
2273
- // model still stops at its eos_token well before this in normal replies.
2274
- maxTokens: req.max_completion_tokens ?? req.max_tokens ??
2275
- defaultGeneratedTokens ?? 65_536,
2276
- temperature: req.temperature ?? serverOptions.defaultTemperature ?? defaultTemp,
2277
- topP: req.top_p ?? serverOptions.defaultTopP ?? ctx.genDefaults.topP ?? 0,
2278
- topK: req.top_k ?? serverOptions.defaultTopK ?? ctx.genDefaults.topK ?? 0,
2279
- seed: req.seed ?? nextDefaultSeed(),
2280
- repetitionPenalty: req.repetition_penalty ?? ctx.genDefaults.repetitionPenalty,
2281
- // mlx_lm.server sampling extensions (defaults mirror server.py: min_p /
2282
- // xtc off, context windows 20). parseLogitBias throws on non-numeric
2283
- // input — callers surface it as a 400 (mlx-lm's coercion error).
2284
- minP: req.min_p ?? 0,
2285
- xtcProbability: req.xtc_probability ?? 0,
2286
- xtcThreshold: req.xtc_threshold ?? 0,
2287
- ...(req.xtc_probability ? { xtcSpecialTokens } : {}),
2288
- logitBias: parseLogitBias(req.logit_bias),
2289
- repetitionContextSize: req.repetition_context_size ?? 20,
2290
- presencePenalty: req.presence_penalty ?? 0,
2291
- presenceContextSize: req.presence_context_size ?? 20,
2292
- frequencyPenalty: req.frequency_penalty ?? 0,
2293
- frequencyContextSize: req.frequency_context_size ?? 20,
2294
- hlg: resolveHlg(req.hlg, serverOptions.hlg),
2295
- // generate() yields token ids; stop sequences match on decoded text
2296
- // (they can span token boundaries), so the StopMatcher sits at the
2297
- // decode layer below and halts the loop via onToken → false.
2298
- stopSequences: (typeof req.stop === "string" ? [req.stop] : req.stop ?? [])
2299
- .filter((s) => typeof s === "string" && s.length > 0),
2300
- ...kvScheme,
2301
- // Paged KV is a server-wide mode like the kv scheme above (v1:
2302
- // serial bf16 gemma4 — the createServer refusals hold the invariants).
2303
- ...(serverOptions.pagedKv ? { pagedKv: serverOptions.pagedKv } : {}),
2304
- };
2305
- };
2306
-
2307
- /** Compile a grammar controller from the request's structured-output fields
2308
- * (response_format / guided_* / structured_outputs). Returns null when no
2309
- * constraint is requested OR when compile failed (degrade path). On degrade,
2310
- * `degradeHint` carries a human description for the system-prompt injection
2311
- * + Warning header (oMLX parity — never 500). Honors MLX_BUN_GRAMMAR=0. */
2312
- const compileGrammarForRequest = async (
2313
- req: ChatRequest,
2314
- ): Promise<{ controller: import("./grammar").GrammarController | null; degradeHint: string | null }> => {
2315
- if (!grammarEnabled()) {
2316
- return {
2317
- controller: null,
2318
- degradeHint: "grammar compilation disabled by MLX_BUN_GRAMMAR=0",
2319
- };
2320
- }
2321
- const r = await compileGrammarRequest(
2322
- req as GrammarRequest,
2323
- ctx.tokenizer,
2324
- ctx.model.config.text.vocabSize,
2325
- );
2326
- if (!r) return { controller: null, degradeHint: null };
2327
- return { controller: r.controller, degradeHint: r.degradeHint };
2328
- };
2329
-
2330
- // Map OpenAI reasoning_effort levels onto the Qwen3.8 template's supported
2331
- // set (xhigh|medium|low — the template raises on anything else). "none"
2332
- // means thinking off (handled by resolveEnableThinking), so no depth is
2333
- // passed. Only consumed for templates with readsReasoningEffort.
2334
- const qwenReasoningEffort = (
2335
- effort: ChatRequest["reasoning_effort"],
2336
- ): "xhigh" | "medium" | "low" | undefined => {
2337
- switch (effort) {
2338
- case "minimal":
2339
- case "low":
2340
- return "low";
2341
- case "medium":
2342
- return "medium";
2343
- case "high":
2344
- case "xhigh":
2345
- return "xhigh";
2346
- default:
2347
- return undefined; // "none" / unset → template default (xhigh)
2348
- }
2349
- };
2350
-
2351
- const templateOptionsFor = (req: ChatRequest, tools: ToolDefinition[] | null) => {
2352
- // enableThinking resolution (and its precedence) lives in
2353
- // resolveEnableThinking so template rendering and sampling stay in sync.
2354
- return {
2355
- tools,
2356
- enableThinking: resolveEnableThinking(req),
2357
- reasoningEffort: qwenReasoningEffort(req.reasoning_effort),
2358
- preserveThinking: req.chat_template_kwargs?.preserve_thinking,
2359
- };
2360
- };
2361
-
2362
- const promptIdsFor = (
2363
- req: ChatRequest,
2364
- tools: ToolDefinition[] | null,
2365
- ): { ids: number[]; startInThinking: boolean } => {
2366
- const opts = templateOptionsFor(req, tools);
2367
- const rendered = ctx.template.render(normalizeMessages(req.messages), opts);
2368
- const ids = ctx.tokenizer.encode(rendered);
2369
- // template includes <bos>; tokenizer post-processor also prepends one
2370
- const trimmed = ids[0] === ids[1] && ids[0] === ctx.tokenizer.bosTokenId ? ids.slice(1) : ids;
2371
- return { ids: trimmed, startInThinking: promptEndsInOpenThink(rendered) };
2372
- };
2373
-
2374
- /** STABLE cache boundary: the prompt's tail is a generation primer (e.g.
2375
- * 12B's `<|channel>thought` tokens) that the NEXT turn's re-render does
2376
- * not contain — a cache entry ending there can never prefix-match again
2377
- * once the rings wrap. Probe by rendering this conversation as if a
2378
- * reply existed; the common prefix is what every future turn preserves.
2379
- * The probe is EXPENSIVE (a second full render + encode — ~150 ms at
2380
- * 16k with our tokenizer; 30-50% of the 12B ctx-repeat TTFT,
2381
- * 2026-07-06b), but the PRIMER LENGTH it measures is constant per
2382
- * template mode: the divergence sits after a special-token turn
2383
- * delimiter, and special tokens break BPE merges, so message content
2384
- * can't shift it. So the full probe runs ONCE per mode and later
2385
- * requests pay a subtraction. (A wrong boundary is a QUALITY knob, not
2386
- * a correctness one — any prefix ≤ len-1 is a valid snapshot point.) */
2387
- const primerLenByMode = new Map<string, number>();
2388
- const stableLenFor = (
2389
- req: ChatRequest,
2390
- tools: ToolDefinition[] | null,
2391
- trimmed: number[],
2392
- ): number => {
2393
- const opts = templateOptionsFor(req, tools);
2394
- const mode = `${opts.enableThinking}|${!!tools?.length}|${opts.reasoningEffort ?? ""}|${opts.preserveThinking ?? ""}`;
2395
- const primer = primerLenByMode.get(mode);
2396
- if (primer !== undefined) return Math.max(0, trimmed.length - primer);
2397
- let stableLen = trimmed.length;
725
+ // Per-request preparation (options/template/prompt-ids/grammar/router)
726
+ // and the chat core, each built once with their collaborators injected.
727
+ // The request pipeline, composed once:
728
+ // new ChatRequest(body) → chatStage.run → InferenceRequest
729
+ // → inferenceStage.admit → inferenceStage.run → InferenceResult → wire
730
+ // (/v1/completions substitutes textStage; Anthropic and Responses reuse
731
+ // chatStage with their own wire formats.)
732
+ const prep = createRequestPrep({ ctx, serverOptions, kvScheme, defaultGeneratedTokens });
733
+ const { templateOptionsFor } = prep;
734
+ const preparation = createPreparationExecutor((work, signal) => gateway.runExclusive(work, undefined, signal), batch);
735
+ const chatStage = new ChatStage(
736
+ ctx, prep, promptCache, contextLimit, serverOptions.defaultAdapter,
737
+ preparation, serving.buildPrompt);
738
+ const textStage = new TextCompletionStage(
739
+ ctx, prep, contextLimit, defaultGeneratedTokens, serverOptions.defaultAdapter);
740
+ const inferenceStage = new InferenceStage(completionExecutor);
741
+ const openAiMeta = (id: string) => ({ id, created: Math.floor(Date.now() / 1000), model: ctx.modelId });
742
+ /** Parse the JSON body under the request's trace; a bad body is a 400. */
743
+ const parseBody = async <T,>(request: Request, trace: PromptResponseTrace | undefined): Promise<T | Response> => {
744
+ const closeBodyParse = trace?.begin("request.body_parse");
2398
745
  try {
2399
- const probe = ctx.tokenizer.encode(
2400
- ctx.template.render(
2401
- [...normalizeMessages(req.messages), { role: "assistant", content: "x" }],
2402
- opts,
2403
- ),
2404
- );
2405
- const probeTrimmed =
2406
- probe[0] === probe[1] && probe[0] === ctx.tokenizer.bosTokenId ? probe.slice(1) : probe;
2407
- let i = 0;
2408
- const n = Math.min(trimmed.length, probeTrimmed.length);
2409
- while (i < n && trimmed[i] === probeTrimmed[i]) i++;
2410
- stableLen = i;
2411
- // Memoize the mode's primer length (sanity-capped: a "primer" longer
2412
- // than 64 tokens means the probe diverged for content reasons —
2413
- // don't generalize that).
2414
- const p = trimmed.length - stableLen;
2415
- if (p >= 0 && p <= 64) primerLenByMode.set(mode, p);
2416
- } catch { /* probe is best-effort; full boundary is the fallback */ }
2417
- return stableLen;
746
+ return (await request.json()) as T;
747
+ } catch {
748
+ trace?.finish("error", { stage: "body_parse" });
749
+ return Response.json({ error: { message: "invalid JSON body" } }, { status: 400 });
750
+ } finally {
751
+ closeBodyParse?.();
752
+ }
2418
753
  };
2419
754
 
2420
- const toolStreamMode = (tools: ToolDefinition[] | null): ToolStreamMode =>
2421
- // Gemma-4 keeps the token-level sentinel path; every other family
2422
- // (MiniCPM5, Qwen3/3.5, Tier-0 generics) emits tool calls as DECODED TEXT
2423
- // — see selectToolStreamMode for why sentinel ids must stay family-gated.
2424
- selectToolStreamMode(ctx.model.config.modelType, !!tools?.length);
2425
755
 
2426
- const toolRouter = (tools: ToolDefinition[] | null): ToolAwareStream =>
2427
- new ToolAwareStream(ctx.tokenizer, toolStreamMode(tools), tools);
2428
756
 
2429
757
  if (serverOptions.unixSocket) {
2430
758
  try { require("node:fs").unlinkSync(serverOptions.unixSocket); } catch {}
@@ -2616,7 +944,7 @@ export function createServer(
2616
944
  let kvMode = "bf16";
2617
945
  if (kvScheme.turboQuant) {
2618
946
  // v1: full-attention layers only (sliding-window stays bf16 —
2619
- // docs/design/turboquant-kv.md non-goal).
947
+ // docs/design/turboquant.md non-goal).
2620
948
  const fullAttn = layerTypes.filter((l) => l !== "sliding_attention").length;
2621
949
  kvMode = `turbo k${kvScheme.turboQuant.kBits}v${kvScheme.turboQuant.vBits}`;
2622
950
  kvLayers[`turbo-k${kvScheme.turboQuant.kBits}v${kvScheme.turboQuant.vBits}`] = fullAttn;
@@ -2648,7 +976,7 @@ export function createServer(
2648
976
  dir: serverOptions.ssdCacheDir,
2649
977
  entries: ssdStore.entries,
2650
978
  bytes: ssdStore.totalBytes,
2651
- max_bytes: ssdStore.maxBytes,
979
+ max_bytes: Number.isFinite(ssdStore.maxBytes) ? ssdStore.maxBytes : null,
2652
980
  restores: ssdStore.stats.restores,
2653
981
  spills: ssdStore.stats.spills,
2654
982
  restore_ms_last: Math.round(ssdStore.stats.restoreMsLast),
@@ -2677,39 +1005,14 @@ export function createServer(
2677
1005
  },
2678
1006
  admission: {
2679
1007
  max_safe_context: admission.maxSafeContext,
1008
+ enforced_context_tokens: contextLimit,
2680
1009
  memory_budget_bytes:
2681
1010
  ctx.glmMemoryPlan?.processLimitBytes ??
2682
1011
  serverOptions.memoryBudgetBytes ?? null,
2683
1012
  usable_bytes: admission.usableBytes,
2684
1013
  weights_bytes: ctx.model.weightsBytes,
2685
1014
  },
2686
- ...(ctx.glmMemoryPlan ? {
2687
- glm52: {
2688
- preset: ctx.glmMemoryPlan.preset,
2689
- planned_process_bytes: ctx.glmMemoryPlan.plannedProcessBytes,
2690
- process_limit_bytes: ctx.glmMemoryPlan.processLimitBytes,
2691
- context_tokens: ctx.glmMemoryPlan.contextTokens,
2692
- max_generation_tokens: ctx.glmMemoryPlan.maxGenerationTokens,
2693
- batch_size: ctx.glmMemoryPlan.batchSize,
2694
- dsa: ctx.model instanceof Glm52Model && ctx.model.capabilities.dsa,
2695
- mtp: ctx.draft?.provider.id === "glm52-native-mtp",
2696
- mtp_draft_tokens: ctx.glmMemoryPlan.mtpDraftTokens,
2697
- resident_weight_bytes: ctx.glmMemoryPlan.lineItems.residentWeightsBytes,
2698
- main_expert_slab_bytes: ctx.glmMemoryPlan.lineItems.mainExpertSlabBytes,
2699
- mtp_expert_slab_bytes: ctx.glmMemoryPlan.lineItems.mtpExpertSlabBytes,
2700
- expert_runtime: ctx.model instanceof Glm52Model &&
2701
- ctx.model.expertRuntime
2702
- ? {
2703
- main_residency:
2704
- ctx.model.expertRuntime.manager.snapshot(),
2705
- mtp_residency:
2706
- ctx.model.expertRuntime.mtp?.manager.snapshot() ?? null,
2707
- last_turn: ctx.model.expertRuntime.lastTelemetry,
2708
- last_repin: ctx.model.expertRuntime.lastRepin,
2709
- }
2710
- : null,
2711
- },
2712
- } : {}),
1015
+ ...serving.diagnostics(),
2713
1016
  // --batch: configured cap, whether batching is live for this model,
2714
1017
  // and rows currently decoding in the batch.
2715
1018
  batch: {
@@ -2725,6 +1028,25 @@ export function createServer(
2725
1028
  });
2726
1029
  }
2727
1030
 
1031
+ // A parent-managed GPU job holds this response open. The connection owns
1032
+ // the native lease, so parent death cannot strand the worker lock.
1033
+ if (serverOptions.unixSocket && url.pathname === "/admin/lease" && request.method === "POST") {
1034
+ const lease = await gateway.acquireExecutionLease(request.signal);
1035
+ let released = false;
1036
+ const release = () => {
1037
+ if (released) return;
1038
+ released = true;
1039
+ request.signal.removeEventListener("abort", release);
1040
+ lease.dispose();
1041
+ };
1042
+ request.signal.addEventListener("abort", release, { once: true });
1043
+ if (request.signal.aborted) release();
1044
+ return new Response(new ReadableStream({
1045
+ start(controller) { controller.enqueue(new TextEncoder().encode("leased\n")); },
1046
+ cancel() { release(); },
1047
+ }), { headers: { "content-type": "application/octet-stream" } });
1048
+ }
1049
+
2728
1050
  // Engine-mode admin (unix-socket children only — never exposed on
2729
1051
  // TCP): drain = quiesce the gateway + demote the whole prompt cache
2730
1052
  // to the SSD tier. The pool calls this before evicting a model
@@ -2749,27 +1071,14 @@ export function createServer(
2749
1071
  let sbody: { prompt?: string };
2750
1072
  try { sbody = (await request.json()) as typeof sbody; }
2751
1073
  catch { return Response.json({ error: "invalid JSON" }, { status: 400, headers: CURVE_CORS }); }
2752
- let sids = ctx.tokenizer.encode(ctx.template.render([{ role: "user", content: typeof sbody.prompt === "string" ? sbody.prompt : "" }], templateOptionsFor({} as ChatRequest, null)));
1074
+ let sids = ctx.tokenizer.encode(ctx.template.render([{ role: "user", content: typeof sbody.prompt === "string" ? sbody.prompt : "" }], templateOptionsFor({} as ChatRequestParams, null)));
2753
1075
  if (sids[0] === sids[1] && sids[0] === ctx.tokenizer.bosTokenId) sids = sids.slice(1);
2754
1076
  try {
2755
1077
  const NB = 80;
2756
1078
  // gateway.runExclusive: this raw forward must not run concurrently
2757
1079
  // with batched decode steps or a serial generation (GPU + shared
2758
1080
  // model state are single-owner — D3, one lock).
2759
- const result = await gateway.runExclusive(async () => {
2760
- const cache = ctx.model.makeCache();
2761
- try {
2762
- const logits = ctx.model.forward(sids, cache); // [1, L, V]
2763
- const [, Ln, V] = logits.shape as [number, number, number];
2764
- const last = logits.slice([0, Ln - 1, 0], [1, Ln, V]);
2765
- const f = last.toFloat32(); logits.dispose(); last.dispose();
2766
- let mx = -Infinity; for (const v of f) if (v > mx) mx = v;
2767
- let Z = 0; for (const v of f) Z += Math.exp(v - mx); const lse = mx + Math.log(Z);
2768
- const bins = new Array<number>(NB).fill(0);
2769
- for (const v of f) { const t = Math.max(0, Math.min(1, (v - lse - CURVE_UMIN) / (-CURVE_UMIN))); const bi = Math.min(NB - 1, Math.floor(t * NB)); bins[bi] = (bins[bi] ?? 0) + 1; }
2770
- return { bins, vocab: V };
2771
- } finally { for (const c of cache) c.dispose(); }
2772
- });
1081
+ const result = await gateway.runExclusive(() => serving.signal(sids, NB, CURVE_UMIN));
2773
1082
  return Response.json(result, { headers: CURVE_CORS });
2774
1083
  } catch (e) {
2775
1084
  return Response.json({ error: `signal failed: ${(e as Error).message}` }, { status: 500, headers: CURVE_CORS });
@@ -2778,7 +1087,7 @@ export function createServer(
2778
1087
 
2779
1088
  // ---- v2 HLG Curve Designer: POST /generate {prompt, curve, n, max_tokens, seed} ----
2780
1089
  // The drawn log-prob transfer curve REPLACES temperature+softmax entirely
2781
- // (src/curve-sampler.ts). The browser editor calls this; same curve object the
1090
+ // (src/lab/curve/curve-sampler.ts). The browser editor calls this; same curve object the
2782
1091
  // tool's "Copy values" emits is the one the sampler consumes — one contract.
2783
1092
  if (url.pathname === "/generate" && request.method === "OPTIONS")
2784
1093
  return new Response(null, { headers: CURVE_CORS });
@@ -2802,7 +1111,7 @@ export function createServer(
2802
1111
  const n = Math.max(1, Math.min(8, Math.floor(Number(body.n) || 3)));
2803
1112
  const maxTokens = Math.max(1, Math.min(256, Math.floor(Number(body.max_tokens) || 80)));
2804
1113
  const baseSeed = Number.isFinite(body.seed) ? Number(body.seed) >>> 0 : (Date.now() & 0xffffffff);
2805
- let ids = ctx.tokenizer.encode(ctx.template.render([{ role: "user", content: prompt }], templateOptionsFor({} as ChatRequest, null)));
1114
+ let ids = ctx.tokenizer.encode(ctx.template.render([{ role: "user", content: prompt }], templateOptionsFor({} as ChatRequestParams, null)));
2806
1115
  if (ids[0] === ids[1] && ids[0] === ctx.tokenizer.bosTokenId) ids = ids.slice(1);
2807
1116
  const samples: { text: string; junk: boolean }[] = [];
2808
1117
  try {
@@ -2811,9 +1120,14 @@ export function createServer(
2811
1120
  const genOpts: GenerateOptions = useCurve
2812
1121
  ? { curve, seed: baseSeed + i, maxTokens, ...kvScheme }
2813
1122
  : { temperature: recipe.temperature, topP: recipe.topP, topK: recipe.topK, seed: baseSeed + i, maxTokens, ...kvScheme };
2814
- // runExclusive: runGeneration touches the GPU + prompt cache, so
2815
- // it needs the gateway lock (this endpoint bypasses gateway.run).
2816
- await gateway.runExclusive(() => runGeneration(ids, genOpts, (t) => { toks.push(t); }));
1123
+ // The same plan and execution lease cover the Lab comparison endpoint.
1124
+ const shape = { hasVision: false, hasAdapters: false, hasRepetitionPenalty: false,
1125
+ hasLogitsExtras: false, wantsLogprobs: false, userSeed: true,
1126
+ kvQuant: !!(genOpts.kvBits || genOpts.kvConfig?.length), turboQuant: !!genOpts.turboQuant,
1127
+ hasGrammar: false, hasDraft: !!ctx.draft };
1128
+ const placement = gateway.place(shape, genOpts);
1129
+ await gateway.run(ids, genOpts, (t) => { toks.push(t); }, undefined,
1130
+ shape, placement, request.signal);
2817
1131
  const text = ctx.tokenizer.decode(toks, true).trim();
2818
1132
  samples.push({ text, junk: curveJunk(text) });
2819
1133
  }
@@ -2823,478 +1137,6 @@ export function createServer(
2823
1137
  return Response.json({ mode: useCurve ? "curve" : "default", recipe: useCurve ? undefined : recipe, n, seed: baseSeed, samples }, { headers: CURVE_CORS });
2824
1138
  }
2825
1139
 
2826
- // The chat-completions core, shared by both protocol surfaces:
2827
- // /v1/chat/completions calls it directly; /v1/messages (Anthropic)
2828
- // translates its body into this shape and the Response back —
2829
- // generation, tools, vision, stop sequences, prompt cache, and
2830
- // admission control all live here exactly once.
2831
- const handleChat = async (
2832
- body: ChatRequest,
2833
- signal: AbortSignal,
2834
- streamProtocol?: CompletionStreamProtocol,
2835
- requestId?: string,
2836
- trace?: PromptResponseTrace,
2837
- ): Promise<Response> => {
2838
- if (!Array.isArray(body.messages) || body.messages.length === 0)
2839
- return Response.json({ error: { message: "messages required" } }, { status: 400 });
2840
- // mlx-lm validates logprobs params up front (ValueError → 400)
2841
- const lpParamError = validateLogprobsParams(body);
2842
- if (lpParamError)
2843
- return Response.json({ error: { message: lpParamError } }, { status: 400 });
2844
- // Covers /v1/messages and /v1/responses too — both funnel through
2845
- // this handler after translation.
2846
- const effortError = validateReasoningEffort(body);
2847
- if (effortError)
2848
- return Response.json({ error: { message: effortError } }, { status: 400 });
2849
-
2850
- const id = requestId ?? `chatcmpl-${crypto.randomUUID()}`;
2851
- const closePromptPrepare = trace?.begin("request.prompt_prepare");
2852
- const created = Math.floor(Date.now() / 1000);
2853
- const tools =
2854
- body.tool_choice === "none" ? null : (body.tools?.length ? body.tools : null);
2855
- const hasImages = body.messages.some(
2856
- (m) => Array.isArray(m.content) &&
2857
- m.content.some((p: any) => p.type === "image_url" || p.type === "image"),
2858
- );
2859
- // Same shapes extractAudio accepts: OpenAI-canonical input_audio plus
2860
- // optiq's audio / audio_url aliases (audio-input-plan.md §3.2).
2861
- const hasAudio = body.messages.some(
2862
- (m) => Array.isArray(m.content) &&
2863
- m.content.some(
2864
- (p: any) => p.type === "input_audio" || p.type === "audio" || p.type === "audio_url",
2865
- ),
2866
- );
2867
- // Video content parts (video_url / video with base64 data) —
2868
- // Qwen3.5-family only (decoded to sampled frames via the
2869
- // AVFoundation sidecar, vision/video-frames.ts).
2870
- const hasVideos = body.messages.some(
2871
- (m) => Array.isArray(m.content) &&
2872
- m.content.some((p: any) => p.type === "video_url" || p.type === "video"),
2873
- );
2874
- let promptIds: number[];
2875
- let stableLen: number | null = null;
2876
- // Whether the prompt primed an open <think> (Qwen3.5/MiniCPM5 thinking
2877
- // on) so the model's output starts mid-reasoning — seeds the splitter.
2878
- // Vision is Gemma4-only (no switchable thinking channel), so it's false.
2879
- let startInThinking = false;
2880
- let vision: Parameters<typeof runGeneration>[3];
2881
- let diffusionPixels: import("./mlx/array").MlxArray | null = null;
2882
- const ownership = new RequestOwnership();
2883
- // Grammar-constrained decoding (src/grammar.ts). Compile BEFORE prompt
2884
- // rendering: on the degrade path (compile failed but a constraint was
2885
- // requested) inject a system message instructing valid JSON so the
2886
- // model still best-efforts schema-conformant output (oMLX parity —
2887
- // _compile_grammar_for_request returning None + build_json_system_prompt).
2888
- let grammarCtrl: import("./grammar").GrammarController | null = null;
2889
- let grammarWarning: string | null = null;
2890
- const grammarReq =
2891
- body.response_format != null || !!body.guided_grammar ||
2892
- !!body.guided_regex || !!body.guided_choice?.length ||
2893
- body.structured_outputs != null;
2894
- if (grammarReq) {
2895
- const g = await compileGrammarForRequest(body);
2896
- grammarCtrl = ownership.own(g.controller);
2897
- if (!g.controller && g.degradeHint) {
2898
- const degraded = applyGrammarDegrade(body, g.degradeHint);
2899
- grammarWarning = degraded.warning;
2900
- body = degraded.body;
2901
- }
2902
- }
2903
- const rejectBeforeRun = (response: Response): Response => {
2904
- ownership.dispose();
2905
- return response;
2906
- };
2907
- try {
2908
- // Video is Qwen3.5-family only and never composes with audio —
2909
- // one early guard so no downstream branch can silently drop a
2910
- // video part (the gemma/diffusion builders don't know the type).
2911
- if (hasVideos && (hasAudio || !(ctx.model instanceof Qwen35Model))) {
2912
- return rejectBeforeRun(Response.json(
2913
- { error: { message: hasAudio
2914
- ? "video and audio content parts cannot be combined"
2915
- : `model ${ctx.modelId} does not accept video input — video ` +
2916
- `content parts need a Qwen3.5-family model (e.g. Qwen3.8-27B)` } },
2917
- { status: 400 },
2918
- ));
2919
- }
2920
- if (hasAudio) {
2921
- // Audio (and MIXED image+audio) input — A4 of
2922
- // docs/design/audio-input-plan.md. One buildMultimodalPrompt call
2923
- // splices both media kinds in document order. A request WITH
2924
- // audio on a model whose tower is unavailable is an explicit 400
2925
- // — never a silent text-only degrade (that leniency is only for
2926
- // requests without the media, getVisionTower's contract).
2927
- const audioTower = getAudioTower(ctx);
2928
- if (!audioTower || !ctx.audioTokenIds) {
2929
- return rejectBeforeRun(Response.json(
2930
- {
2931
- error: {
2932
- message:
2933
- `model ${ctx.modelId} has no audio tower — audio input needs ` +
2934
- `a model whose config.json carries audio_config and whose ` +
2935
- `sidecar ships the audio tensors (e.g. gemma-4 e4b OptiQ)`,
2936
- },
2937
- },
2938
- { status: 400 },
2939
- ));
2940
- }
2941
- let visionSide:
2942
- | { tower: VisionEncoder; tokenIds: VisionTokenIds }
2943
- | undefined;
2944
- if (hasImages) {
2945
- const tower = getVisionTower(ctx);
2946
- if (!tower) {
2947
- return rejectBeforeRun(Response.json(
2948
- { error: { message: "model has no vision sidecar" } }, { status: 400 },
2949
- ));
2950
- }
2951
- visionSide = { tower, tokenIds: ctx.visionTokenIds };
2952
- }
2953
- const { messages: withAudioParts, images } =
2954
- await extractImages(normalizeMessages(body.messages));
2955
- const { messages, audio } = await extractAudio(withAudioParts);
2956
- // Non-WAV containers (mp3/m4a/flac/ogg/aiff/…) transcode to
2957
- // 16 kHz WAV via CoreAudio; RIFF bytes pass through untouched.
2958
- // Failures throw into the prompt-build 400 below.
2959
- const wavs = await Promise.all(audio.map(ensureWav));
2960
- // The towers are only ever non-null for Gemma4 (loader gates).
2961
- const mp = await buildMultimodalPrompt(
2962
- ctx.model as Gemma4Model,
2963
- {
2964
- ...(visionSide ? { vision: visionSide } : {}),
2965
- audio: { tower: audioTower, tokenIds: ctx.audioTokenIds },
2966
- },
2967
- ctx.tokenizer, ctx.template, messages, images, wavs, tools,
2968
- );
2969
- promptIds = mp.ids;
2970
- // bidirMask is null whenever audio is present (§3.3 Q1: mixed
2971
- // prompts run fully causal); the union mask does the per-layer
2972
- // id zeroing either way.
2973
- vision = {
2974
- embeddings: mp.embeddings,
2975
- ...(mp.bidirMask ? { imageMask: mp.bidirMask } : {}),
2976
- multimodalMask: mp.multimodalMask,
2977
- };
2978
- } else if (hasImages && ctx.model instanceof DiffusionGemmaModel) {
2979
- // DiffusionGemma image-text-to-text: its OWN dedicated SigLIP tower +
2980
- // encoder vision merge feed the denoising engine (NOT the AR
2981
- // forwardEmbeddings path). v1 supports a single image.
2982
- const dm = ctx.model;
2983
- if (!dm.visionTower) {
2984
- return rejectBeforeRun(Response.json(
2985
- { error: { message: "this checkpoint has no vision tower" } }, { status: 400 },
2986
- ));
2987
- }
2988
- const { messages, images } = await extractImages(normalizeMessages(body.messages));
2989
- if (images.length !== 1) {
2990
- return rejectBeforeRun(Response.json(
2991
- { error: { message: "DiffusionGemma image input supports exactly one image" } },
2992
- { status: 400 },
2993
- ));
2994
- }
2995
- const rendered = ctx.template.render(messages, { tools, addGenerationPrompt: true });
2996
- const rawIds = ctx.tokenizer.encode(rendered, /* addSpecialTokens */ false);
2997
- const { pixels, softTokens } = await dm.visionTower.preprocess(images[0]!);
2998
- diffusionPixels = ownership.own(pixels);
2999
- promptIds = spliceImageTokens(rawIds, [softTokens], {
3000
- image: ctx.visionTokenIds.imageTokenId,
3001
- boi: ctx.visionTokenIds.boiTokenId,
3002
- eoi: ctx.visionTokenIds.eoiTokenId,
3003
- });
3004
- } else if ((hasImages || hasVideos) && ctx.model instanceof Qwen35Model) {
3005
- // Qwen3.8 vision + video: the tower is a Qwen3VLVisionTower
3006
- // riding the shared lazy slot (makeVisionLoader's qwen branch);
3007
- // image/video spans splice into input embeddings and the request
3008
- // carries mRoPE positions + delta (PLAN 14v/14w). Videos decode
3009
- // to sampled frames via the AVFoundation sidecar.
3010
- const tower = getVisionTower(ctx) as unknown as Qwen3VLVisionTower | null;
3011
- if (!tower) {
3012
- return rejectBeforeRun(Response.json(
3013
- { error: { message: "model has no vision sidecar" } }, { status: 400 },
3014
- ));
3015
- }
3016
- const { messages: withVideos, images } =
3017
- await extractImages(normalizeMessages(body.messages));
3018
- const { messages, videos } = await extractVideos(withVideos);
3019
- const vp = await buildQwen3VLVisionPrompt(
3020
- ctx.model, tower, ctx.tokenizer, ctx.template, messages, images,
3021
- {
3022
- imageTokenId: (ctx.model.config.raw.image_token_id as number) ?? 248056,
3023
- videoTokenId: (ctx.model.config.raw.video_token_id as number) ?? 248057,
3024
- visionStartId: (ctx.model.config.raw.vision_start_token_id as number) ?? 248053,
3025
- visionEndId: (ctx.model.config.raw.vision_end_token_id as number) ?? 248054,
3026
- },
3027
- templateOptionsFor(body, tools),
3028
- videos,
3029
- );
3030
- promptIds = vp.ids;
3031
- vision = { embeddings: vp.embeddings, mrope: vp.mrope };
3032
- } else if (hasImages) {
3033
- // Loads (and caches) the tower on first image request — text-only
3034
- // sessions never pay for it.
3035
- const tower = getVisionTower(ctx);
3036
- if (!tower) {
3037
- return rejectBeforeRun(Response.json(
3038
- { error: { message: "model has no vision sidecar" } }, { status: 400 },
3039
- ));
3040
- }
3041
- const { messages, images } = await extractImages(normalizeMessages(body.messages));
3042
- // The tower is only ever non-null for Gemma4 (sidecar gate in
3043
- // makeVisionLoader), so the model narrow is safe here.
3044
- const vp = await buildVisionPrompt(
3045
- ctx.model as Gemma4Model, tower, ctx.tokenizer, ctx.template,
3046
- messages, images, ctx.visionTokenIds, tools,
3047
- );
3048
- promptIds = vp.ids;
3049
- vision = { embeddings: vp.embeddings, imageMask: vp.imageMask };
3050
- } else {
3051
- const built = promptIdsFor(body, tools);
3052
- promptIds = built.ids;
3053
- startInThinking = built.startInThinking;
3054
- stableLen = -1; // marker: chat path, probe lazily below
3055
- }
3056
- } catch (e) {
3057
- return rejectBeforeRun(Response.json(
3058
- { error: { message: `prompt build failed: ${(e as Error).message}` } },
3059
- { status: 400 },
3060
- ));
3061
- }
3062
- ownership.own(vision?.embeddings);
3063
- ownership.own(vision?.imageMask);
3064
- ownership.own(vision?.multimodalMask);
3065
- ownership.own(diffusionPixels);
3066
- let options: ReturnType<typeof toOptions>;
3067
- try {
3068
- options = toOptions(body);
3069
- // Stable cache boundary for the prompt-boundary snapshot (chat
3070
- // prompts end in a generation primer the next turn won't contain).
3071
- // The probe costs a second full render+encode (~150 ms at 16k), so
3072
- // run it ONLY when a snapshot can actually be taken: the cache
3073
- // must not already hold the strict prefix (warm repeats peek at
3074
- // len-1 → skip; the snapshot gate needs boundary > cached prefix).
3075
- if (stableLen !== null) {
3076
- const cachePeek = promptCache.peekPrefixLen(
3077
- promptIds, options.adapters?.join("+") ?? "");
3078
- if (cachePeek < promptIds.length - 1)
3079
- options.snapshotAt = stableLenFor(body, tools, promptIds);
3080
- }
3081
- } catch (e) {
3082
- // bad logit_bias (non-numeric keys/values) — mlx-lm's coercion error
3083
- return rejectBeforeRun(
3084
- Response.json({ error: { message: (e as Error).message } }, { status: 400 }),
3085
- );
3086
- }
3087
- if (diffusionPixels) options.visionPixels = diffusionPixels;
3088
- // Attach the compiled grammar controller (null when no constraint /
3089
- // degrade — generate() runs the unmasked fast pipelined loop).
3090
- if (grammarCtrl) options.grammar = grammarCtrl;
3091
- const requestedMaxTokens = options.maxTokens ?? 1024;
3092
- let adapterIds: string[];
3093
- try {
3094
- // A request's explicit `adapter` (incl. "none") wins over the
3095
- // startup default from `serve --adapter <dir>`.
3096
- adapterIds = ctx.adapters.resolveSpec(body.adapter ?? serverOptions.defaultAdapter);
3097
- } catch (e) {
3098
- return rejectBeforeRun(
3099
- Response.json({ error: { message: (e as Error).message } }, { status: 400 }),
3100
- );
3101
- }
3102
-
3103
- const wantLogprobs = body.logprobs === true;
3104
- const topLogprobs =
3105
- typeof body.top_logprobs === "number" && body.top_logprobs > 0
3106
- ? body.top_logprobs : 0;
3107
- let prepared: PreparedCompletion;
3108
- try {
3109
- prepared = prepareCompletion({
3110
- requestId: id,
3111
- plan: {
3112
- promptIds,
3113
- options,
3114
- requestedMaxTokens,
3115
- maxSafeContext: admission.maxSafeContext,
3116
- stream: body.stream === true,
3117
- wantLogprobs,
3118
- topLogprobs,
3119
- adapterIds,
3120
- hasVision: !!vision,
3121
- userSeed: body.seed !== undefined,
3122
- hasGrammar: !!grammarCtrl,
3123
- hasDraft: !!ctx.draft,
3124
- ownership,
3125
- },
3126
- vision,
3127
- pipeline: {
3128
- router: toolRouter(tools),
3129
- stopper: new StopMatcher(options.stopSequences),
3130
- thinking: new ThinkingTagSplitter(
3131
- ctx.template.thinkingFormat === "think-tag",
3132
- startInThinking,
3133
- ),
3134
- collectToolCalls: true,
3135
- },
3136
- ...(body.stream
3137
- ? {
3138
- createFlowControl: ({ mechanism }) =>
3139
- createTimedFlowControl(mechanism === "serial"),
3140
- }
3141
- : {}),
3142
- onPlacement: ({ mechanism, shape }) => {
3143
- if (runtimeValue("MLX_BUN_LANE_DEBUG") === "1")
3144
- console.error(
3145
- `[scheduling] mechanism=${mechanism} shape=${JSON.stringify(shape)} ` +
3146
- `t=${Date.now() % 100000}`,
3147
- );
3148
- },
3149
- idToToken: (tokenId) => ctx.tokenizer.idToToken(tokenId),
3150
- });
3151
- } catch (error) {
3152
- closePromptPrepare?.();
3153
- trace?.finish("error", { stage: "prepare_completion" });
3154
- if (!(error instanceof CompletionRejected)) throw error;
3155
- return Response.json({ error: error.error }, { status: error.status });
3156
- }
3157
- closePromptPrepare?.();
3158
-
3159
- if (body.stream) {
3160
- const streamAbort = new AbortController();
3161
- const generationSignal = AbortSignal.any([signal, streamAbort.signal]);
3162
- let cancelled = false;
3163
- const stream = new ReadableStream<Uint8Array>({
3164
- start(controller) {
3165
- const enc = new TextEncoder();
3166
- const send = (obj: unknown) => {
3167
- if (!generationSignal.aborted)
3168
- controller.enqueue(enc.encode(`data: ${JSON.stringify(obj)}\n\n`));
3169
- };
3170
- const emitFrames = (frames: string[]) => {
3171
- if (generationSignal.aborted) return;
3172
- for (const frame of frames) controller.enqueue(enc.encode(frame));
3173
- };
3174
- let latestUsage: Readonly<CompletionSummary["usage"]> | null = null;
3175
- let wroteFirstEvent = false;
3176
- let traceOutcome: "success" | "error" | "abort" = "success";
3177
- const chunk = (delta: Record<string, unknown>, finish: string | null) => ({
3178
- id, object: "chat.completion.chunk", created, model: ctx.modelId,
3179
- choices: [{ index: 0, delta, finish_reason: finish }],
3180
- });
3181
- void (async () => {
3182
- try {
3183
- // The gateway owns lane selection + GPU exclusivity; this body
3184
- // runs per-request (concurrently in batched mode, each writing
3185
- // its own SSE stream — the per-row fan-out).
3186
- if (streamProtocol) emitFrames(streamProtocol.start());
3187
- else send(chunk({ role: "assistant", content: "" }, null));
3188
- const sendEvents = (events: readonly CompletionEvent[]) => {
3189
- if (events.length && !wroteFirstEvent) {
3190
- wroteFirstEvent = true;
3191
- trace?.mark("response.first_write");
3192
- }
3193
- if (streamProtocol) {
3194
- emitFrames(streamProtocol.addEvents([...events]));
3195
- return;
3196
- }
3197
- for (const event of events) {
3198
- if (event.type === "reasoning") {
3199
- send(chunk({ reasoning: event.text }, null));
3200
- } else if (event.type === "content") {
3201
- send(chunk({ content: event.text }, null));
3202
- } else {
3203
- send(chunk({
3204
- tool_calls: event.calls.map((call, index) => ({ index, ...call })),
3205
- }, null));
3206
- }
3207
- }
3208
- };
3209
- const summary = await completionExecutor.execute(prepared, {
3210
- signal: generationSignal,
3211
- trace,
3212
- onEvents: sendEvents,
3213
- ...(streamProtocol
3214
- ? { onUsageProgress: (usage) => { latestUsage = usage; } }
3215
- : {}),
3216
- });
3217
- const usage = completionUsage(summary);
3218
- if (streamProtocol) {
3219
- emitFrames(streamProtocol.finish(summary.finishReason, usage));
3220
- } else {
3221
- send({ ...chunk({}, summary.finishReason), usage });
3222
- // bare sentinel per the OpenAI spec — strict SDK clients
3223
- // require the unquoted terminator.
3224
- controller.enqueue(enc.encode("data: [DONE]\n\n"));
3225
- }
3226
- trace?.mark("response.final_write");
3227
- } catch (e) {
3228
- traceOutcome = generationSignal.aborted ? "abort" : "error";
3229
- if (!generationSignal.aborted) {
3230
- const message = (e as Error).message;
3231
- if (streamProtocol) {
3232
- emitFrames([
3233
- ...streamProtocol.error(message),
3234
- ...streamProtocol.finish(
3235
- "stop",
3236
- latestUsage ? completionProtocolUsage(latestUsage) : {},
3237
- ),
3238
- ]);
3239
- } else {
3240
- send({ error: { message } });
3241
- }
3242
- }
3243
- } finally {
3244
- trace?.finish(traceOutcome);
3245
- if (!cancelled) {
3246
- if (generationSignal.aborted) controller.error(generationSignal.reason);
3247
- else controller.close();
3248
- }
3249
- }
3250
- })();
3251
- },
3252
- cancel(reason) {
3253
- cancelled = true;
3254
- streamAbort.abort(reason);
3255
- },
3256
- });
3257
- return new Response(stream, {
3258
- headers: {
3259
- "content-type": "text/event-stream",
3260
- "cache-control": "no-cache",
3261
- connection: "keep-alive",
3262
- ...(grammarWarning ? { Warning: grammarWarning } : {}),
3263
- },
3264
- });
3265
- }
3266
-
3267
- try {
3268
- {
3269
- const summary = await completionExecutor.execute(prepared, { signal, trace });
3270
- trace?.mark("response.final_write");
3271
- const response = Response.json({
3272
- id, object: "chat.completion", created, model: ctx.modelId,
3273
- choices: [{
3274
- index: 0,
3275
- message: {
3276
- role: "assistant",
3277
- content: summary.content || (summary.toolCalls.length ? null : ""),
3278
- ...(summary.reasoning ? { reasoning: summary.reasoning } : {}),
3279
- ...(summary.toolCalls.length ? { tool_calls: summary.toolCalls } : {}),
3280
- },
3281
- ...(summary.logprobs ? { logprobs: summary.logprobs } : {}),
3282
- finish_reason: summary.finishReason,
3283
- }],
3284
- usage: completionUsage(summary),
3285
- }, grammarWarning ? { headers: { Warning: grammarWarning } } : undefined);
3286
- trace?.finish("success");
3287
- return response;
3288
- }
3289
- } catch (e) {
3290
- trace?.finish(signal.aborted ? "abort" : "error");
3291
- // A 500 with no server-side trace is undebuggable in the field —
3292
- // always log the stack (the JSON body keeps only the message).
3293
- console.error(`[serve] 500 on chat request:\n${(e as Error).stack ?? e}`);
3294
- return Response.json({ error: { message: (e as Error).message } }, { status: 500 });
3295
- }
3296
- };
3297
-
3298
1140
  if (url.pathname === "/v1/chat/completions" && request.method === "POST") {
3299
1141
  const id = `chatcmpl-${crypto.randomUUID()}`;
3300
1142
  const trace = createPromptResponseTrace({
@@ -3302,25 +1144,19 @@ export function createServer(
3302
1144
  requestId: id,
3303
1145
  route: url.pathname,
3304
1146
  });
3305
- const closeBodyParse = trace?.begin("request.body_parse");
3306
- let body: ChatRequest;
3307
- try {
3308
- body = (await request.json()) as ChatRequest;
3309
- } catch {
3310
- closeBodyParse?.();
3311
- trace?.finish("error", { stage: "body_parse" });
3312
- return Response.json({ error: { message: "invalid JSON body" } }, { status: 400 });
3313
- }
3314
- closeBodyParse?.();
3315
- return handleChat(body, request.signal, undefined, id, trace);
1147
+ const body = await parseBody<ChatRequestParams>(request, trace);
1148
+ if (body instanceof Response) return body;
1149
+ const meta = openAiMeta(id);
1150
+ const a = await admit(
1151
+ inferenceStage, () => chatStage.run(new ChatRequest(body), id, request.signal), trace, "chat request");
1152
+ if ("response" in a) return a.response;
1153
+ return body.stream
1154
+ ? respondStream(inferenceStage, a.admitted, chatCompletionStream(meta), request.signal, trace)
1155
+ : respondJson(inferenceStage, a.admitted, (r) => chatCompletionJson(r, meta), request.signal, trace);
3316
1156
  }
3317
1157
 
3318
- // Raw text completion (mlx_lm.server's /v1/completions, request_type
3319
- // "text"): NO chat template — the prompt string is tokenized directly
3320
- // (tokenizer.encode with the tokenizer's own special-token handling,
3321
- // exactly mlx-lm's `tokenizer.encode(request.prompt)`). Rides the same
3322
- // GenerationGateway + admission + adapter path as chat; no tool router
3323
- // or thinking splitter (raw text in, raw text out).
1158
+ // Raw text completion (mlx_lm.server's /v1/completions): no chat
1159
+ // template — TextCompletionStage tokenizes the prompt string directly.
3324
1160
  if (url.pathname === "/v1/completions" && request.method === "POST") {
3325
1161
  const id = `cmpl-${crypto.randomUUID()}`;
3326
1162
  const trace = createPromptResponseTrace({
@@ -3328,203 +1164,15 @@ export function createServer(
3328
1164
  requestId: id,
3329
1165
  route: url.pathname,
3330
1166
  });
3331
- const closeBodyParse = trace?.begin("request.body_parse");
3332
- let body: TextCompletionRequest;
3333
- try {
3334
- body = (await request.json()) as TextCompletionRequest;
3335
- } catch {
3336
- closeBodyParse?.();
3337
- trace?.finish("error", { stage: "body_parse" });
3338
- return Response.json({ error: { message: "invalid JSON body" } }, { status: 400 });
3339
- }
3340
- closeBodyParse?.();
3341
- const closePromptPrepare = trace?.begin("request.prompt_prepare");
3342
- if (typeof body.prompt !== "string" || body.prompt.length === 0)
3343
- return Response.json(
3344
- {
3345
- error: {
3346
- message: "prompt (a non-empty string) is required " +
3347
- "(token-array prompts are not accepted, matching mlx_lm.server)",
3348
- },
3349
- },
3350
- { status: 400 },
3351
- );
3352
- // mlx-lm validates logprobs params up front (ValueError → 400)
3353
- const lpParamError = validateLogprobsParams(body);
3354
- if (lpParamError)
3355
- return Response.json({ error: { message: lpParamError } }, { status: 400 });
3356
- const created = Math.floor(Date.now() / 1000);
3357
- const promptIds = ctx.tokenizer.encode(body.prompt);
3358
- const ownership = new RequestOwnership();
3359
- let options: ReturnType<typeof toOptions>;
3360
- try {
3361
- options = toOptions(body as unknown as ChatRequest);
3362
- } catch (e) {
3363
- // bad logit_bias (non-numeric keys/values) — mlx-lm's coercion error
3364
- return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
3365
- }
3366
- // Grammar-constrained decoding on raw completions too (response_format /
3367
- // guided_* / structured_outputs). The degrade path for /v1/completions
3368
- // has no chat template to inject a system message into, so a degrade
3369
- // only emits the Warning header (no prompt injection) — documented gap
3370
- // vs the chat lane, which mirrors oMLX's text-completions behavior.
3371
- let textGrammarWarning: string | null = null;
3372
- const textGrammarReq =
3373
- body.response_format != null || !!body.guided_grammar ||
3374
- !!body.guided_regex || !!body.guided_choice?.length ||
3375
- body.structured_outputs != null;
3376
- if (textGrammarReq) {
3377
- const g = await compileGrammarForRequest(body as unknown as ChatRequest);
3378
- if (g.controller) options.grammar = ownership.own(g.controller);
3379
- else if (g.degradeHint)
3380
- textGrammarWarning =
3381
- `grammar not enforced: ${g.degradeHint} - no prompt injection on /v1/completions`;
3382
- }
3383
- // mlx_lm.server's default max_tokens is 512 (its --max-tokens CLI
3384
- // default). The chat lane's very generous default is wrong for raw
3385
- // completion: with no template an EOS may never come.
3386
- const requestedMaxTokens = body.max_completion_tokens ?? body.max_tokens ??
3387
- defaultGeneratedTokens ?? 512;
3388
- let adapterIds: string[];
3389
- try {
3390
- adapterIds = ctx.adapters.resolveSpec(body.adapter ?? serverOptions.defaultAdapter);
3391
- } catch (e) {
3392
- ownership.dispose();
3393
- return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
3394
- }
3395
- const wantLogprobs = body.logprobs === true;
3396
- const topLogprobs =
3397
- typeof body.top_logprobs === "number" && body.top_logprobs > 0
3398
- ? body.top_logprobs : 0;
3399
- let prepared: PreparedCompletion;
3400
- try {
3401
- prepared = prepareCompletion({
3402
- requestId: id,
3403
- plan: {
3404
- promptIds,
3405
- options,
3406
- requestedMaxTokens,
3407
- maxSafeContext: admission.maxSafeContext,
3408
- stream: body.stream === true,
3409
- wantLogprobs,
3410
- topLogprobs,
3411
- adapterIds,
3412
- hasVision: false,
3413
- userSeed: body.seed !== undefined,
3414
- hasGrammar: !!options.grammar,
3415
- hasDraft: !!ctx.draft,
3416
- ownership,
3417
- },
3418
- pipeline: {
3419
- router: new ToolAwareStream(ctx.tokenizer, "plain", null),
3420
- stopper: new StopMatcher(options.stopSequences),
3421
- thinking: new ThinkingTagSplitter(false),
3422
- collectToolCalls: false,
3423
- },
3424
- ...(body.stream
3425
- ? { createFlowControl: () => createTimedFlowControl(true) }
3426
- : {}),
3427
- idToToken: (tokenId) => ctx.tokenizer.idToToken(tokenId),
3428
- });
3429
- } catch (error) {
3430
- closePromptPrepare?.();
3431
- trace?.finish("error", { stage: "prepare_completion" });
3432
- if (!(error instanceof CompletionRejected)) throw error;
3433
- return Response.json({ error: error.error }, { status: error.status });
3434
- }
3435
- closePromptPrepare?.();
3436
-
3437
- if (body.stream) {
3438
- const streamAbort = new AbortController();
3439
- const generationSignal = AbortSignal.any([request.signal, streamAbort.signal]);
3440
- let cancelled = false;
3441
- const stream = new ReadableStream<Uint8Array>({
3442
- start(controller) {
3443
- const enc = new TextEncoder();
3444
- const send = (obj: unknown) => {
3445
- if (!generationSignal.aborted)
3446
- controller.enqueue(enc.encode(`data: ${JSON.stringify(obj)}\n\n`));
3447
- };
3448
- const chunk = (text: string, finish: string | null) => ({
3449
- id, object: "text_completion", created, model: ctx.modelId,
3450
- choices: [{ index: 0, text, finish_reason: finish }],
3451
- });
3452
- let wroteFirstEvent = false;
3453
- let traceOutcome: "success" | "error" | "abort" = "success";
3454
- void (async () => {
3455
- try {
3456
- const sendEvents = (events: readonly CompletionEvent[]) => {
3457
- if (events.length && !wroteFirstEvent) {
3458
- wroteFirstEvent = true;
3459
- trace?.mark("response.first_write");
3460
- }
3461
- for (const event of events) {
3462
- if (event.type === "content") send(chunk(event.text, null));
3463
- }
3464
- };
3465
- const summary = await completionExecutor.execute(prepared, {
3466
- signal: generationSignal,
3467
- trace,
3468
- onEvents: sendEvents,
3469
- });
3470
- // final chunk: finish_reason + usage (mlx-lm gates usage behind
3471
- // stream_options.include_usage; we always attach it, matching
3472
- // our chat lane — an additive superset OpenAI clients ignore)
3473
- send({
3474
- ...chunk("", summary.finishReason),
3475
- usage: completionUsage(summary),
3476
- });
3477
- controller.enqueue(enc.encode("data: [DONE]\n\n"));
3478
- trace?.mark("response.final_write");
3479
- } catch (e) {
3480
- traceOutcome = generationSignal.aborted ? "abort" : "error";
3481
- if (!generationSignal.aborted)
3482
- send({ error: { message: (e as Error).message } });
3483
- } finally {
3484
- trace?.finish(traceOutcome);
3485
- if (!cancelled) {
3486
- if (generationSignal.aborted) controller.error(generationSignal.reason);
3487
- else controller.close();
3488
- }
3489
- }
3490
- })();
3491
- },
3492
- cancel(reason) {
3493
- cancelled = true;
3494
- streamAbort.abort(reason);
3495
- },
3496
- });
3497
- return new Response(stream, {
3498
- headers: {
3499
- "content-type": "text/event-stream",
3500
- "cache-control": "no-cache",
3501
- connection: "keep-alive",
3502
- ...(textGrammarWarning ? { Warning: textGrammarWarning } : {}),
3503
- },
3504
- });
3505
- }
3506
-
3507
- try {
3508
- const summary = await completionExecutor.execute(prepared, {
3509
- signal: request.signal,
3510
- trace,
3511
- });
3512
- trace?.mark("response.final_write");
3513
- const response = Response.json({
3514
- id, object: "text_completion", created, model: ctx.modelId,
3515
- choices: [{
3516
- index: 0, text: summary.content,
3517
- ...(summary.logprobs ? { logprobs: summary.logprobs } : {}),
3518
- finish_reason: summary.finishReason,
3519
- }],
3520
- usage: completionUsage(summary),
3521
- }, textGrammarWarning ? { headers: { Warning: textGrammarWarning } } : undefined);
3522
- trace?.finish("success");
3523
- return response;
3524
- } catch (e) {
3525
- trace?.finish(request.signal.aborted ? "abort" : "error");
3526
- return Response.json({ error: { message: (e as Error).message } }, { status: 500 });
3527
- }
1167
+ const body = await parseBody<TextCompletionParams>(request, trace);
1168
+ if (body instanceof Response) return body;
1169
+ const meta = openAiMeta(id);
1170
+ const a = await admit(
1171
+ inferenceStage, () => textStage.run(new TextCompletionRequest(body), id), trace, "text completion");
1172
+ if ("response" in a) return a.response;
1173
+ return body.stream
1174
+ ? respondStream(inferenceStage, a.admitted, textCompletionStream(meta), request.signal, trace)
1175
+ : respondJson(inferenceStage, a.admitted, (r) => textCompletionJson(r, meta), request.signal, trace);
3528
1176
  }
3529
1177
 
3530
1178
  // Anthropic Messages API (Phase 11) — on by default, mirroring
@@ -3533,41 +1181,35 @@ export function createServer(
3533
1181
  // src/anthropic.ts. Point Claude Code at this port via
3534
1182
  // ANTHROPIC_BASE_URL for a fully local backend.
3535
1183
  if (url.pathname === "/v1/messages" && request.method === "POST") {
3536
- const anthropicError = (status: number, type: string, message: string) =>
3537
- Response.json({ type: "error", error: { type, message } }, { status });
1184
+ const anthropicError: ErrorFormatter = (status, message) =>
1185
+ Response.json(
1186
+ { type: "error", error: { type: status >= 500 ? "api_error" : "invalid_request_error", message } },
1187
+ { status },
1188
+ );
3538
1189
  let anthropicBody: AnthropicRequest;
3539
1190
  try {
3540
1191
  anthropicBody = (await request.json()) as AnthropicRequest;
3541
1192
  } catch {
3542
- return anthropicError(400, "invalid_request_error", "invalid JSON body");
1193
+ return anthropicError(400, "invalid JSON body", {});
3543
1194
  }
3544
- let chatBody: ChatRequest;
1195
+ let chatBody: ChatRequestParams;
3545
1196
  try {
3546
- chatBody = anthropicToChatBody(anthropicBody) as unknown as ChatRequest;
1197
+ chatBody = anthropicToChatBody(anthropicBody) as unknown as ChatRequestParams;
3547
1198
  } catch (e) {
3548
- return anthropicError(400, "invalid_request_error", (e as Error).message);
3549
- }
3550
- const resp = await handleChat(
3551
- chatBody,
3552
- request.signal,
3553
- anthropicBody.stream
3554
- ? createAnthropicStreamProtocol(ctx.modelId)
3555
- : undefined,
3556
- );
3557
- if (!resp.ok) {
3558
- const err = (await resp.json().catch(() => null)) as
3559
- | { error?: { message?: string } }
3560
- | null;
3561
- return anthropicError(
3562
- resp.status,
3563
- resp.status >= 500 ? "api_error" : "invalid_request_error",
3564
- err?.error?.message ?? "request failed",
3565
- );
3566
- }
3567
- if (anthropicBody.stream) {
3568
- return resp;
1199
+ return anthropicError(400, (e as Error).message, {});
3569
1200
  }
3570
- return Response.json(chatJsonToAnthropic(await resp.json(), ctx.modelId));
1201
+ const id = `chatcmpl-${crypto.randomUUID()}`;
1202
+ const a = await admit(
1203
+ inferenceStage, () => chatStage.run(new ChatRequest(chatBody), id, request.signal), undefined,
1204
+ "anthropic request", anthropicError);
1205
+ if ("response" in a) return a.response;
1206
+ if (anthropicBody.stream)
1207
+ return respondStream(
1208
+ inferenceStage, a.admitted, createAnthropicStreamProtocol(ctx.modelId), request.signal);
1209
+ return respondJson(
1210
+ inferenceStage, a.admitted,
1211
+ (r) => chatJsonToAnthropic(chatCompletionJson(r, openAiMeta(id)), ctx.modelId),
1212
+ request.signal, undefined, anthropicError);
3571
1213
  }
3572
1214
 
3573
1215
  // OpenAI Responses API (Phase 11) — Codex/Cursor/Continue speak
@@ -3575,7 +1217,7 @@ export function createServer(
3575
1217
  // ported in src/responses.ts. previous_response_id resumes a prior
3576
1218
  // conversation from the in-process store (TTL + byte-capped LRU).
3577
1219
  if (url.pathname === "/v1/responses" && request.method === "POST") {
3578
- const responsesError = (status: number, message: string) =>
1220
+ const responsesError: ErrorFormatter = (status, message) =>
3579
1221
  Response.json(
3580
1222
  {
3581
1223
  error: {
@@ -3590,80 +1232,55 @@ export function createServer(
3590
1232
  try {
3591
1233
  responsesBody = (await request.json()) as ResponsesRequest;
3592
1234
  } catch {
3593
- return responsesError(400, "invalid JSON body");
1235
+ return responsesError(400, "invalid JSON body", {});
3594
1236
  }
3595
1237
 
3596
- // previous_response_id: prepend the prior conversation's input
3597
- // + output (as input items); carry instructions forward only if
3598
- // the new request omits them (oracle semantics).
3599
- const prevId = responsesBody.previous_response_id ?? null;
3600
- if (prevId) {
3601
- const prior = responseStore.get(prevId);
3602
- if (!prior)
3603
- return responsesError(404, `previous_response_id '${prevId}' not found or expired`);
3604
- const prepended = [...prior.input, ...outputItemsToInputItems(prior.output)];
3605
- const newInput =
3606
- typeof responsesBody.input === "string"
3607
- ? responsesBody.input
3608
- ? [{ type: "message", role: "user", content: responsesBody.input }]
3609
- : []
3610
- : responsesBody.input ?? [];
3611
- responsesBody = {
3612
- ...responsesBody,
3613
- input: [...prepended, ...newInput] as Array<Record<string, unknown>>,
3614
- instructions: responsesBody.instructions ?? prior.instructions ?? undefined,
3615
- };
3616
- }
3617
- // Remember the effective input so a later follow-up that chains
3618
- // off THIS response sees the full history.
3619
- const capturedInput: unknown[] =
3620
- typeof responsesBody.input === "string"
3621
- ? [{ type: "message", role: "user", content: responsesBody.input }]
3622
- : [...(responsesBody.input ?? [])];
3623
- const capturedInstructions = responsesBody.instructions ?? null;
1238
+ let conversation: ReturnType<typeof resolveResponsesConversation>;
1239
+ try { conversation = resolveResponsesConversation(responsesBody, responseStore); }
1240
+ catch (error) { return responsesError(404, (error as Error).message, {}); }
1241
+ responsesBody = conversation.body;
1242
+ const prevId = conversation.previousId;
1243
+ const capturedInput = conversation.input;
1244
+ const capturedInstructions = conversation.instructions;
1245
+ const storeHere = !(serverOptions.unixSocket && request.headers.get("x-mlx-bun-response-owner") === "parent");
1246
+ const remember = (id: string, output: unknown[]) => {
1247
+ if (storeHere) responseStore.put(id, { input: capturedInput, output, instructions: capturedInstructions });
1248
+ };
3624
1249
 
3625
- let chatBody: ChatRequest;
1250
+ let chatBody: ChatRequestParams;
3626
1251
  try {
3627
- chatBody = responsesToChatBody(responsesBody) as unknown as ChatRequest;
1252
+ chatBody = responsesToChatBody(responsesBody) as unknown as ChatRequestParams;
3628
1253
  } catch (e) {
3629
- return responsesError(400, (e as Error).message);
3630
- }
3631
- const resp = await handleChat(
3632
- chatBody,
3633
- request.signal,
3634
- responsesBody.stream
3635
- ? createResponsesStreamProtocol(
3636
- ctx.modelId,
3637
- prevId,
3638
- (final) =>
3639
- responseStore.put(final.id as string, {
3640
- input: capturedInput,
3641
- output: final.output as unknown[],
3642
- instructions: capturedInstructions,
3643
- }),
3644
- )
3645
- : undefined,
3646
- );
3647
- if (!resp.ok) {
3648
- const err = (await resp.json().catch(() => null)) as
3649
- | { error?: { message?: string } }
3650
- | null;
3651
- return responsesError(resp.status, err?.error?.message ?? "request failed");
3652
- }
3653
- if (responsesBody.stream) {
3654
- return resp;
1254
+ return responsesError(400, (e as Error).message, {});
3655
1255
  }
3656
- const responses = chatJsonToResponses(await resp.json(), ctx.modelId, prevId);
3657
- responseStore.put(responses.id as string, {
3658
- input: capturedInput,
3659
- output: responses.output as unknown[],
3660
- instructions: capturedInstructions,
3661
- });
3662
- return Response.json(responses);
1256
+ const id = `chatcmpl-${crypto.randomUUID()}`;
1257
+ const a = await admit(
1258
+ inferenceStage, () => chatStage.run(new ChatRequest(chatBody), id, request.signal), undefined,
1259
+ "responses request", responsesError);
1260
+ if ("response" in a) return a.response;
1261
+ if (responsesBody.stream)
1262
+ return respondStream(
1263
+ inferenceStage, a.admitted,
1264
+ createResponsesStreamProtocol(
1265
+ ctx.modelId,
1266
+ prevId,
1267
+ (final) =>
1268
+ remember(final.id as string, final.output as unknown[]),
1269
+ ),
1270
+ request.signal);
1271
+ return respondJson(
1272
+ inferenceStage, a.admitted,
1273
+ (r) => {
1274
+ const responses = chatJsonToResponses(chatCompletionJson(r, openAiMeta(id)), ctx.modelId, prevId);
1275
+ remember(responses.id as string, responses.output as unknown[]);
1276
+ return responses;
1277
+ },
1278
+ request.signal, undefined, responsesError);
3663
1279
  }
3664
1280
 
3665
1281
  const labResponse = await handleLabRoute(url, request, {
3666
1282
  ensureJobs,
1283
+ acquireGpu: (signal) => gateway.acquireExecutionLease(signal),
3667
1284
  serverPort: () => server.port,
3668
1285
  invalidateLibrary: () => discoveryRoutes.invalidateLibrary(),
3669
1286
  });
@@ -3679,6 +1296,13 @@ export function createServer(
3679
1296
  },
3680
1297
  });
3681
1298
  serverLifecycles.set(serverRef, {
1299
+ stopJobs: async () => {
1300
+ if (jobStore) {
1301
+ const jobs = await import("./jobs");
1302
+ await Promise.all([jobs.closeSubprocessJobs(jobStore), jobs.closeInProcessJobs(jobStore)]);
1303
+ }
1304
+ },
1305
+ close: async () => { preparation.close(); await sessionEngine.close(); await gateway.close(); },
3682
1306
  flush: flushDurability,
3683
1307
  stats: durabilityStats,
3684
1308
  stopTimers: () => {