mlx-bun 0.2.0 → 0.3.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 (364) hide show
  1. package/README.md +71 -588
  2. package/docs/README.md +91 -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/{design → archive/investigations}/optimization_plan.md +1 -1
  8. package/docs/{design → archive/investigations}/pre-colibri-stabilization.md +5 -5
  9. package/docs/{design → archive/investigations}/the-dreaming-handoff.md +2 -2
  10. package/docs/design/batching.md +569 -0
  11. package/docs/design/decode-speed-program.md +3178 -169
  12. package/docs/design/dreaming-nightly-pipeline.md +675 -282
  13. package/docs/design/generic-model-support.md +579 -334
  14. package/docs/design/kv-cache.md +680 -0
  15. package/docs/design/orpo-training.md +673 -1239
  16. package/docs/design/speculative-decoding.md +1627 -0
  17. package/docs/design/turboquant.md +2102 -0
  18. package/docs/design/unified-engine-frontier-plan.md +1831 -727
  19. package/docs/design/web-chat-redesign.md +666 -1083
  20. package/docs/planning/PRODUCT_ROADMAP.md +34 -1
  21. package/docs/planning/release-notes-v0.3.0.md +35 -0
  22. package/docs/planning/video-series.md +2 -2
  23. package/docs/reference/benchmarks.md +1162 -0
  24. package/docs/reference/cli.md +394 -159
  25. package/docs/reference/distribution.md +321 -182
  26. package/docs/reference/environment.md +348 -0
  27. package/docs/reference/glossary.md +39 -0
  28. package/docs/reference/library-api.md +566 -168
  29. package/docs/reference/memory.md +227 -64
  30. package/docs/reference/models.md +104 -6
  31. package/docs/reference/server-api.md +886 -425
  32. package/docs/reference/server-config.md +638 -330
  33. package/docs/reference/training.md +393 -254
  34. package/docs/reference/troubleshooting.md +166 -0
  35. package/package.json +9 -4
  36. package/src/audio/decode.ts +2 -2
  37. package/src/audio/features.ts +2 -2
  38. package/src/backends/mlx/autoregressive.ts +108 -0
  39. package/src/backends/mlx/batch-group.ts +1646 -0
  40. package/src/backends/mlx/diffusion.ts +39 -0
  41. package/src/backends/mlx/draft-target.ts +47 -0
  42. package/src/backends/mlx/evaluation.ts +105 -0
  43. package/src/backends/mlx/fill-append.ts +45 -0
  44. package/src/backends/mlx/gateway-binding.ts +71 -0
  45. package/src/backends/mlx/graph.ts +38 -0
  46. package/src/backends/mlx/memory-client.ts +189 -0
  47. package/src/backends/mlx/memory-guard.ts +40 -0
  48. package/src/backends/mlx/methods.ts +90 -0
  49. package/src/backends/mlx/model-prompt.ts +164 -0
  50. package/src/backends/mlx/model-serving.ts +79 -0
  51. package/src/backends/mlx/prefill.ts +33 -0
  52. package/src/backends/mlx/request-policy.ts +20 -0
  53. package/src/backends/mlx/rollback.ts +51 -0
  54. package/src/backends/mlx/serial-executor.ts +273 -0
  55. package/src/backends/mlx/speculative.ts +114 -0
  56. package/src/backends/mlx/state-views.ts +26 -0
  57. package/src/backends/mlx/text-engine.ts +73 -0
  58. package/src/cli.ts +58 -25
  59. package/src/client.ts +66 -0
  60. package/src/config.ts +21 -4
  61. package/src/contracts/completion.ts +16 -0
  62. package/src/contracts/execution.ts +49 -0
  63. package/src/contracts/generation.ts +103 -0
  64. package/src/contracts/host.ts +7 -0
  65. package/src/contracts/jobs.ts +66 -0
  66. package/src/contracts/pi-provider.ts +30 -0
  67. package/src/contracts/pi.ts +259 -0
  68. package/src/contracts/resources.ts +18 -0
  69. package/src/contracts/scheduling.ts +27 -0
  70. package/src/contracts/tsconfig.json +14 -0
  71. package/src/diffusion/diffusion-generate.ts +218 -153
  72. package/src/embed.ts +1 -1
  73. package/src/engine/admission.ts +59 -0
  74. package/src/engine/cancellation.ts +38 -0
  75. package/src/engine/engine.ts +56 -0
  76. package/src/engine/execution-coordinator.ts +70 -0
  77. package/src/engine/execution-plan.ts +51 -0
  78. package/src/engine/resources.ts +42 -0
  79. package/src/engine/scheduler.ts +67 -0
  80. package/src/engine/session.ts +247 -0
  81. package/src/eval/README.md +44 -1
  82. package/src/eval/kl.ts +236 -3
  83. package/src/eval/reference-logits.ts +207 -0
  84. package/src/eval/runner.ts +8 -96
  85. package/src/fill/echo-index.ts +188 -0
  86. package/src/fill/fill-session.ts +452 -0
  87. package/src/fill/proposal.ts +68 -0
  88. package/src/fill/schema-rows.ts +506 -0
  89. package/src/fill/tool-boundary.ts +103 -0
  90. package/src/fit.ts +1 -1
  91. package/src/generate.ts +656 -268
  92. package/src/grammar.ts +12 -7
  93. package/src/harness-pi.ts +14 -33
  94. package/src/hub-rest.ts +4 -4
  95. package/src/index.ts +7 -0
  96. package/src/inference/denoising.ts +17 -0
  97. package/src/inference/graph.ts +33 -0
  98. package/src/inference/prefill.ts +39 -0
  99. package/src/inference/rollback.ts +7 -0
  100. package/src/jobs/db.ts +4 -0
  101. package/src/jobs/index.ts +2 -0
  102. package/src/jobs/job-entry.ts +12 -12
  103. package/src/jobs/runner.ts +98 -13
  104. package/src/jobs/task-client.ts +19 -0
  105. package/src/jobs/types.ts +2 -58
  106. package/src/kv-store.ts +109 -49
  107. package/src/lab/README.md +18 -0
  108. package/src/{assets → lab/curve}/curve-designer.html +2 -2
  109. package/src/{curve-sampler.ts → lab/curve/curve-sampler.ts} +2 -2
  110. package/src/{expert-trace.ts → lab/expert-trace/expert-trace.ts} +5 -5
  111. package/src/{model → lab/paged-kv}/paged-kv.ts +9 -6
  112. package/src/library.ts +40 -0
  113. package/src/lora.ts +18 -0
  114. package/src/media-fetch.ts +1 -1
  115. package/src/memory/article.ts +1 -1
  116. package/src/memory/chunk.ts +1 -1
  117. package/src/memory/model.ts +30 -169
  118. package/src/memory/prompts.ts +1 -1
  119. package/src/memory/rest.ts +1 -1
  120. package/src/memory/schedule.ts +1 -1
  121. package/src/memory/stages.ts +4 -1
  122. package/src/memory/vault.ts +1 -3
  123. package/src/mlx/array.ts +9 -10
  124. package/src/mlx/compile.ts +12 -2
  125. package/src/mlx/ffi.ts +58 -11
  126. package/src/mlx/materialize.ts +23 -0
  127. package/src/mlx/metal-kernel.ts +10 -11
  128. package/src/mlx/ops.ts +33 -13
  129. package/src/mlx/quantized-rows.ts +36 -0
  130. package/src/mlx/token-bitmask.ts +52 -0
  131. package/src/mlx/turboquant-kv-decode.ts +82 -0
  132. package/src/mlx/turboquant-kv-inverse.ts +122 -0
  133. package/src/mlx/turboquant-ops.ts +33 -3
  134. package/src/mlx/turboquant-tables.ts +2 -2
  135. package/src/model/batched-mask.ts +8 -2
  136. package/src/model/batched-quant.ts +1 -1
  137. package/src/model/batched-rotating-quant.ts +2 -2
  138. package/src/model/batched-rotating.ts +1 -1
  139. package/src/model/compiled-decode.ts +1 -1
  140. package/src/model/diffusion-gemma.ts +8 -8
  141. package/src/model/factory.ts +76 -33
  142. package/src/model/fingerprint.ts +1 -1
  143. package/src/model/flash-attention.ts +3 -3
  144. package/src/model/gemma4-base.ts +26 -9
  145. package/src/model/gemma4.ts +6 -6
  146. package/src/model/generated/gemma4-12b.ts +2 -2
  147. package/src/model/generated/gemma4-26b.ts +2 -2
  148. package/src/model/generated/gemma4-e4b.ts +2 -2
  149. package/src/model/glm52-dsa.ts +13 -9
  150. package/src/model/glm52-memory.ts +1 -1
  151. package/src/model/implementation.ts +69 -0
  152. package/src/model/minicpm5.ts +6 -8
  153. package/src/model/profile.ts +12 -2
  154. package/src/model/qwen-append.ts +12 -0
  155. package/src/model/qwen-conv.ts +52 -0
  156. package/src/model/qwen3_5-checkpoint.ts +164 -0
  157. package/src/model/qwen3_5.ts +140 -56
  158. package/src/model/support.ts +1 -1
  159. package/src/model/trellis-balanced-scatter.ts +51 -0
  160. package/src/model/trellis-linear.ts +672 -0
  161. package/src/model/trellis-shared-m.ts +151 -0
  162. package/src/model/trellis-shared-scatter.ts +144 -0
  163. package/src/model/trellis-splitk-prefill.ts +158 -0
  164. package/src/model/trellis-tiled-prefill.ts +136 -0
  165. package/src/model/trellis-vector-expand.ts +62 -0
  166. package/src/model/trellis-wide-prefill.ts +94 -0
  167. package/src/model/universal/dense.ts +1 -1
  168. package/src/model/universal/modules.ts +8 -1
  169. package/src/native-pack.ts +6 -5
  170. package/src/pi-provider.ts +6 -64
  171. package/src/pi-web.ts +8 -315
  172. package/src/prompt-cache.ts +96 -53
  173. package/src/quantize/allocator.ts +1 -1
  174. package/src/quantize/config-writer.ts +10 -3
  175. package/src/quantize/rotate.ts +4 -4
  176. package/src/quantize/trellis.ts +586 -0
  177. package/src/responses.ts +28 -1
  178. package/src/runtime-config.ts +18 -4
  179. package/src/sampler.ts +4 -4
  180. package/src/serve/batch-scheduler.ts +8 -1671
  181. package/src/serve/chat-request.ts +391 -0
  182. package/src/serve/chat-stage.ts +195 -0
  183. package/src/serve/checkpoint-identity.ts +64 -0
  184. package/src/serve/completion-executor.ts +15 -7
  185. package/src/serve/discovery-routes.ts +12 -13
  186. package/src/serve/generation-gateway.ts +154 -193
  187. package/src/serve/http.ts +159 -0
  188. package/src/serve/inference-request.ts +79 -0
  189. package/src/serve/isolate.ts +297 -105
  190. package/src/serve/lab-routes.ts +3 -0
  191. package/src/serve/lane-registry.ts +2 -1
  192. package/src/serve/model-admin-routes.ts +7 -5
  193. package/src/serve/model-binding.ts +43 -0
  194. package/src/serve/model-host.ts +512 -0
  195. package/src/serve/openai-wire.ts +103 -0
  196. package/src/serve/parent-application.ts +56 -0
  197. package/src/serve/pipeline.ts +33 -0
  198. package/src/serve/preparation.ts +33 -0
  199. package/src/serve/request-plan.ts +19 -15
  200. package/src/serve/request-prep.ts +346 -0
  201. package/src/serve/response-lease.ts +24 -0
  202. package/src/serve/responses-client.ts +69 -0
  203. package/src/serve/session-completion-engine.ts +110 -0
  204. package/src/serve/session-search.ts +1 -1
  205. package/src/serve/static-routes.ts +5 -5
  206. package/src/serve/text-completion-stage.ts +104 -0
  207. package/src/serve/token-streams.ts +519 -0
  208. package/src/server.ts +317 -2649
  209. package/src/spec/assistant-source.ts +10 -44
  210. package/src/spec/deepspec-source.ts +6 -6
  211. package/src/spec/dflash-source.ts +8 -8
  212. package/src/spec/dspark/ab-stats.ts +3 -3
  213. package/src/spec/dspark/calibration.ts +2 -2
  214. package/src/spec/dspark/deepspec-module.ts +2 -2
  215. package/src/spec/dspark/generate-dflash.ts +2 -2
  216. package/src/spec/dspark/generate.ts +2 -2
  217. package/src/spec/dspark/loader.ts +1 -1
  218. package/src/spec/dspark/module-dflash.ts +4 -3
  219. package/src/spec/dspark/module.ts +3 -3
  220. package/src/spec/dspark/quantize-drafter.ts +1 -1
  221. package/src/spec/generate.ts +1 -1
  222. package/src/spec/glm52-mtp-source.ts +1 -1
  223. package/src/spec/ngram-source.ts +1 -1
  224. package/src/spec/prefix-state.ts +45 -0
  225. package/src/spec/qwen-mtp-source.ts +166 -69
  226. package/src/spec/serve-loop.ts +154 -156
  227. package/src/spec/source.ts +42 -10
  228. package/src/spec/two-model.ts +1 -1
  229. package/src/ssd-cache.ts +80 -12
  230. package/src/ssd-durability.ts +25 -13
  231. package/src/tool-call.ts +51 -14
  232. package/src/train/flash-cce.ts +2 -2
  233. package/src/train/forward.ts +6 -0
  234. package/src/train/loss.ts +1 -1
  235. package/src/train/prefix-shared.ts +5 -1
  236. package/src/train/segmented.ts +5 -3
  237. package/src/train/trainer.ts +6 -6
  238. package/src/vision/prompt.ts +1 -1
  239. package/src/vision/qwen3vl-tower.ts +52 -10
  240. package/src/web/app.html +1 -1
  241. package/src/web/app.js +5 -3
  242. package/src/web/src/adapters-panel.ts +3 -3
  243. package/src/web/src/assistant.ts +2 -36
  244. package/src/web/src/chat.ts +4 -4
  245. package/src/web/src/composer.ts +4 -4
  246. package/src/web/src/hub.ts +1 -1
  247. package/src/web/src/markdown.ts +0 -0
  248. package/src/web/src/memory-panel.ts +1 -1
  249. package/src/web/src/model-picker.ts +1 -1
  250. package/src/web/src/palette.ts +1 -1
  251. package/src/web/src/protocol.ts +3 -39
  252. package/src/web/src/sessions.ts +1 -1
  253. package/src/web/src/ui-catalog.ts +4 -18
  254. package/src/web-assets.ts +47 -0
  255. package/src/weights.ts +55 -9
  256. package/docs/archive/investigations/curve-runs/bisector-route.json +0 -98
  257. package/docs/archive/investigations/curve-runs/bisector-trace.json +0 -1393
  258. package/docs/archive/investigations/curve-runs/curve-terrain.json +0 -1179
  259. package/docs/archive/investigations/curve-runs/curve-terrain.svg +0 -42
  260. package/docs/archive/investigations/hlg-ab.html +0 -34
  261. package/docs/archive/investigations/hlg-explorer.html +0 -238
  262. package/docs/archive/investigations/hlg-figs/cliff.svg +0 -1
  263. package/docs/archive/investigations/hlg-figs/cross.svg +0 -1
  264. package/docs/archive/investigations/hlg-figs/edges.svg +0 -1
  265. package/docs/archive/investigations/hlg-figs/k20-grids.svg +0 -1
  266. package/docs/archive/investigations/hlg-figs/k20-peaks.svg +0 -1
  267. package/docs/archive/investigations/hlg-figs/lw-curve.svg +0 -1
  268. package/docs/archive/investigations/hlg-figs/terrain.svg +0 -1
  269. package/docs/archive/investigations/hlg-report.html +0 -262
  270. package/docs/archive/investigations/hlg-runs/hlg-ab.json +0 -149
  271. package/docs/archive/investigations/hlg-runs/hlg-blind-prompts.md +0 -87
  272. package/docs/archive/investigations/hlg-runs/hlg-blind.json +0 -107
  273. package/docs/archive/investigations/hlg-runs/hlg-cliff.json +0 -197
  274. package/docs/archive/investigations/hlg-runs/hlg-cross.json +0 -505
  275. package/docs/archive/investigations/hlg-runs/hlg-edge.json +0 -388
  276. package/docs/archive/investigations/hlg-runs/hlg-eval1.json +0 -25
  277. package/docs/archive/investigations/hlg-runs/hlg-frontier.json +0 -86
  278. package/docs/archive/investigations/hlg-runs/hlg-grid.json +0 -242
  279. package/docs/archive/investigations/hlg-runs/hlg-grid2.json +0 -343
  280. package/docs/archive/investigations/hlg-runs/hlg-lw-coarse.json +0 -110
  281. package/docs/archive/investigations/hlg-runs/hlg-map.json +0 -383
  282. package/docs/archive/investigations/hlg-runs/run0-original.log +0 -159
  283. package/docs/archive/investigations/hlg-runs/runA-recommended.log +0 -131
  284. package/docs/archive/investigations/hlg-runs/runAA-cliff-cpm5.log +0 -26
  285. package/docs/archive/investigations/hlg-runs/runAB-grid2-e4b.log +0 -18
  286. package/docs/archive/investigations/hlg-runs/runAC-peak-e4b-partial.log +0 -25
  287. package/docs/archive/investigations/hlg-runs/runAD-ab-samples.log +0 -8
  288. package/docs/archive/investigations/hlg-runs/runAE-edge-e4b.log +0 -51
  289. package/docs/archive/investigations/hlg-runs/runAF-logit-calibration.log +0 -10
  290. package/docs/archive/investigations/hlg-runs/runAG-eval-config2.log +0 -16
  291. package/docs/archive/investigations/hlg-runs/runB-gain-variants.log +0 -192
  292. package/docs/archive/investigations/hlg-runs/runC-pivots.log +0 -162
  293. package/docs/archive/investigations/hlg-runs/runD-oetf.log +0 -132
  294. package/docs/archive/investigations/hlg-runs/runE-eotf.log +0 -162
  295. package/docs/archive/investigations/hlg-runs/runF-shaper-working.log +0 -86
  296. package/docs/archive/investigations/hlg-runs/runG-shaper-sweep.log +0 -101
  297. package/docs/archive/investigations/hlg-runs/runH-wsweep.log +0 -176
  298. package/docs/archive/investigations/hlg-runs/runI-asweep.log +0 -176
  299. package/docs/archive/investigations/hlg-runs/runJ-knobmap.log +0 -49
  300. package/docs/archive/investigations/hlg-runs/runK-frontier.log +0 -19
  301. package/docs/archive/investigations/hlg-runs/runL-widemap.log +0 -65
  302. package/docs/archive/investigations/hlg-runs/runM-frontier-emb.log +0 -15
  303. package/docs/archive/investigations/hlg-runs/runN-verify.log +0 -116
  304. package/docs/archive/investigations/hlg-runs/runO-frontier-seed5000.log +0 -15
  305. package/docs/archive/investigations/hlg-runs/runP-grid-AxSm.log +0 -15
  306. package/docs/archive/investigations/hlg-runs/runQ-grid-tgW.log +0 -15
  307. package/docs/archive/investigations/hlg-runs/runR-blind-transcript.log +0 -79
  308. package/docs/archive/investigations/hlg-runs/runS-cross-coherence.log +0 -33
  309. package/docs/archive/investigations/hlg-runs/runT-cross-cpm5.log +0 -33
  310. package/docs/archive/investigations/hlg-runs/runU-cross-12b.log +0 -33
  311. package/docs/archive/investigations/hlg-runs/runV-lw-e4b.log +0 -22
  312. package/docs/archive/investigations/hlg-runs/runW-lw-12b.log +0 -22
  313. package/docs/archive/investigations/hlg-runs/runX-lw-cpm5.log +0 -22
  314. package/docs/archive/investigations/hlg-runs/runY-cliff-e4b.log +0 -26
  315. package/docs/archive/investigations/hlg-runs/runZ-cliff-12b.log +0 -26
  316. package/docs/archive/mlx-bun-lab-report.html +0 -275
  317. package/docs/archive/wwdc-mlx-bun.html +0 -848
  318. package/docs/design/adapters-end-to-end.md +0 -205
  319. package/docs/design/batching-perf-path.md +0 -255
  320. package/docs/design/batching-v2-plan.md +0 -529
  321. package/docs/design/bucketing-stage.md +0 -197
  322. package/docs/design/compat-cli-surface-design.md +0 -230
  323. package/docs/design/docs-reorg-plan.md +0 -165
  324. package/docs/design/dspark-serving-program.md +0 -371
  325. package/docs/design/dspark-speculative-decoding.md +0 -318
  326. package/docs/design/grammar-spec-batching-integration.md +0 -301
  327. package/docs/design/memory-inference-path.md +0 -422
  328. package/docs/design/memory-synthesis.md +0 -287
  329. package/docs/design/memory-system.md +0 -237
  330. package/docs/design/mlx-lm-serving-execution-seam.md +0 -81
  331. package/docs/design/mlx-lm-tool-parity-plan.md +0 -545
  332. package/docs/design/omlx-adoption-map.md +0 -98
  333. package/docs/design/orpo-dynamic-lambda.md +0 -220
  334. package/docs/design/orpo-future-enhancements.md +0 -340
  335. package/docs/design/paged-kv-cache.md +0 -214
  336. package/docs/design/parallel-slots.md +0 -567
  337. package/docs/design/parity-tier-dag.md +0 -161
  338. package/docs/design/repo-cleanup-plan.md +0 -231
  339. package/docs/design/runtime-isolation.md +0 -263
  340. package/docs/design/segmented-backward-training.md +0 -543
  341. package/docs/design/spec-decode-larger-targets.md +0 -229
  342. package/docs/design/ssd-kv-cold-tier.md +0 -415
  343. package/docs/design/structured-output.md +0 -484
  344. package/docs/design/tauri-desktop-app.md +0 -127
  345. package/docs/design/the-dreaming-master-plan.md +0 -445
  346. package/docs/design/turboquant-kv.md +0 -235
  347. package/docs/design/turboquant-weights.md +0 -467
  348. package/docs/design/w4a16-compute-precision-spike.md +0 -232
  349. package/docs/design/web-chat-beat-matrix.md +0 -315
  350. package/docs/design/write-pipeline-entity.md +0 -468
  351. package/docs/planning/IDEAS.md +0 -69
  352. package/docs/reference/embedding.md +0 -197
  353. package/docs/reference/features-matrix.md +0 -101
  354. package/docs/reference/orpo-quickstart.md +0 -135
  355. package/src/flags.ts +0 -11
  356. /package/docs/{design → archive/investigations}/colibri-glm52-port.md +0 -0
  357. /package/docs/{planning → archive/planning}/chunk-task-roadmap.md +0 -0
  358. /package/docs/{planning → archive/planning}/curve-sampler-research-plan.md +0 -0
  359. /package/docs/{planning → archive/planning}/memory-docs-and-dag-plan.md +0 -0
  360. /package/docs/{planning → archive/planning}/pi-integration-review.md +0 -0
  361. /package/docs/{planning → archive/planning}/release-notes-v0.2.0.md +0 -0
  362. /package/src/{memory-plan.ts → admission-plan.ts} +0 -0
  363. /package/{docs/archive/investigations → src/serve/assets}/curve-terrain.html +0 -0
  364. /package/{docs/dag → src/serve/assets}/training-inference-map.html +0 -0
@@ -14,6 +14,10 @@
14
14
  // only RNG consumers then are the canvas draws.
15
15
 
16
16
  import type { DiffusionGemmaModel } from "../model/diffusion-gemma";
17
+ import type { DenoisingGraph } from "../inference/denoising";
18
+ import { assertMlxDenoisingGraph, bindLegacyDenoisingModel } from "../backends/mlx/diffusion";
19
+ import { cleanupFailure, disposeResources, ownResource } from "../engine/resources";
20
+ import type { DisposableResource, ResourceOwner } from "../contracts/resources";
17
21
  import { MlxArray } from "../mlx/array";
18
22
  import { Dtype } from "../mlx/ffi";
19
23
  import * as ops from "../mlx/ops";
@@ -184,14 +188,48 @@ function entropyTransferMask(entropy: MlxArray, bound: number): MlxArray {
184
188
  return out;
185
189
  }
186
190
 
187
- export function diffusionGenerate(
188
- model: DiffusionGemmaModel,
189
- promptIds: number[],
190
- opts: DiffusionGenOptions,
191
- ): DiffusionGenResult {
192
- const t = model.config.text;
193
- const vocab = t.vocabSize;
194
- const modelCanvas = model.canvasLength;
191
+ /** Compatibility entry. Both sync and asynchronous callers drain one algorithm. */
192
+ export function diffusionGenerate(model: DiffusionGemmaModel, promptIds: number[], opts: DiffusionGenOptions): DiffusionGenResult {
193
+ return denoiseSync(bindLegacyDenoisingModel(model).graph, promptIds, opts);
194
+ }
195
+
196
+ export function denoiseSync<State>(graph: DenoisingGraph<MlxArray, State>, promptIds: number[], opts: DiffusionGenOptions): DiffusionGenResult {
197
+ const steps = denoisingSteps(graph, promptIds, opts);
198
+ while (true) {
199
+ const step = steps.next();
200
+ if (step.done) {
201
+ if (!step.value) throw new Error("denoising ended without a result");
202
+ return step.value;
203
+ }
204
+ }
205
+ }
206
+
207
+ /** One macrotask between bounded denoising units allows cancellation and I/O.
208
+ * No tentative canvas tokens are published, and no extra device eval is added. */
209
+ export async function denoiseAsync<State>(
210
+ graph: DenoisingGraph<MlxArray, State>, promptIds: number[], opts: DiffusionGenOptions, signal?: AbortSignal,
211
+ ): Promise<DiffusionGenResult> {
212
+ signal?.throwIfAborted();
213
+ const steps = denoisingSteps(graph, promptIds, opts);
214
+ try {
215
+ while (true) {
216
+ signal?.throwIfAborted();
217
+ const step = steps.next();
218
+ if (step.done) {
219
+ if (!step.value) throw new Error("denoising ended without a result");
220
+ return step.value;
221
+ }
222
+ await new Promise<void>((resolve) => setImmediate(resolve));
223
+ }
224
+ } finally { steps.return(undefined); }
225
+ }
226
+
227
+ function* denoisingSteps<State>(
228
+ graph: DenoisingGraph<MlxArray, State>, promptIds: number[], opts: DiffusionGenOptions,
229
+ ): Generator<void, DiffusionGenResult | undefined> {
230
+ assertMlxDenoisingGraph(graph);
231
+ const vocab = graph.vocabSize;
232
+ const modelCanvas = graph.canvasLength;
195
233
  const maxSteps = opts.maxDenoisingSteps ?? 48;
196
234
  const fullCanvas = opts.fullCanvas ?? false;
197
235
  const maxCanvas = fullCanvas ? modelCanvas : Math.min(modelCanvas, opts.maxCanvasLength ?? modelCanvas);
@@ -214,10 +252,9 @@ export function diffusionGenerate(
214
252
 
215
253
  ops.randomSeed(opts.seed ?? 0n);
216
254
 
217
- const cache = opts.visionPixels
218
- ? model.prefillVision(promptIds, opts.visionPixels)
219
- : model.prefill(promptIds);
220
- const dequantWeight = model.dequantEmbedWeight();
255
+ let stateOwner: ResourceOwner<State> | undefined;
256
+ let dequantWeight: MlxArray | null = null;
257
+ let failure: { error: unknown } | undefined;
221
258
 
222
259
  const emitted: number[] = [];
223
260
  const blocks: number[][] = [];
@@ -228,6 +265,9 @@ export function diffusionGenerate(
228
265
  let isFirstBlock = true;
229
266
 
230
267
  try {
268
+ const cache = graph.prefill(promptIds, opts.visionPixels);
269
+ stateOwner = ownResource(cache, (state) => graph.closeState(state));
270
+ dequantWeight = graph.dequantEmbedWeight();
231
271
  while (generated < maxNewTokens) {
232
272
  const remaining = maxNewTokens - generated;
233
273
  const canvasLength = fullCanvas
@@ -238,161 +278,186 @@ export function diffusionGenerate(
238
278
  // Continuation: prefill the encoder over the previous accepted block.
239
279
  const prev = blocks[blocks.length - 1]!;
240
280
  const prevArr = MlxArray.fromInt32(Int32Array.from(prev), [1, prev.length]);
241
- model.extendPrefill(prevArr, cache);
242
- prevArr.dispose();
281
+ try { graph.extendPrefill(prevArr, cache); }
282
+ finally { prevArr.dispose(); }
243
283
  }
244
284
  isFirstBlock = false;
245
285
 
246
- let currentCanvas = ops.randint(0, vocab, [1, canvasLength], I32); // RNG: init
247
- let draftReveal = ops.zeros([1, canvasLength], Dtype.bool);
248
- let draftCanvas = ops.reshape(currentCanvas, [1, canvasLength]); // alias copy
249
- let argmaxCanvas = ops.reshape(currentCanvas, [1, canvasLength]);
286
+ let currentCanvas: MlxArray | null = null;
287
+ let draftReveal: MlxArray | null = null;
288
+ let draftCanvas: MlxArray | null = null;
289
+ let argmaxCanvas: MlxArray | null = null;
250
290
  let scEmbeddings: MlxArray | null = null;
251
291
  const history: MlxArray[] = [];
252
292
  let stepsThisCanvas = 0;
253
-
254
- for (let curStep = maxSteps; curStep >= 1; curStep--) {
255
- stepsThisCanvas++;
256
- const logits = model.decoderLogits(currentCanvas, cache, scEmbeddings);
257
- const schedT = linearTemp(curStep, maxSteps, tMin, tMax);
258
- // Match the reference EXACTLY: divide (not reciprocal-multiply) — the
259
- // f32 rounding differs and the confidence threshold is a hard cutoff
260
- // that flips acceptance on a 1-ULP difference, diverging the trajectory.
261
- const schedArr = ops.scalarLike(schedT, logits);
262
- const processed = ops.div(logits, schedArr);
263
- schedArr.dispose();
264
- logits.dispose();
265
-
266
- argmaxCanvas.dispose();
267
- const argmax = ops.argmaxAxis(processed, -1);
268
- argmaxCanvas = argmax.astype(I32); // [1, L]
269
- argmax.dispose();
270
-
271
- if (curStep === 1) {
272
- processed.dispose();
273
- break;
274
- }
275
-
276
- // temperature 0 → denoiser canvas is the argmax.
277
- const denoiser = argmaxCanvas;
278
-
279
- let acceptance: MlxArray;
280
- let nextSc: MlxArray | null = null;
281
- if (sampler === "entropy-bound") {
282
- // cur_step > 1 always here (cur_step==1 breaks before the sampler):
283
- // entropy AND the next soft embeddings come from the SAME entropy
284
- // chain (reference _diffusion_entropy_and_soft_embeddings) — NOT the
285
- // softmax-precise path the confidence sampler uses.
286
- const { probs, entropy } = entropyProbsChain(processed);
287
- acceptance = entropyTransferMask(entropy, entropyBound);
288
- entropy.dispose();
289
- nextSc = softEmbeddingsFromProbs(probs, dequantWeight, model.embedScale);
290
- probs.dispose();
291
- // accepted/current update (entropy variant)
292
- const accepted = ops.where(acceptance, denoiser, currentCanvas);
293
- const noise = ops.randint(0, vocab, [1, canvasLength], I32); // RNG: re-noise
294
- const newCurrent = ops.where(acceptance, accepted, noise);
295
- noise.dispose();
296
- currentCanvas.dispose();
297
- currentCanvas = newCurrent;
298
- draftReveal.dispose();
299
- const notAcceptance = ops.logicalNot(acceptance);
300
- draftReveal = ops.logicalNot(notAcceptance); // = acceptance
301
- notAcceptance.dispose();
302
- draftCanvas.dispose();
303
- draftCanvas = ops.reshape(argmaxCanvas, [1, canvasLength]);
304
- accepted.dispose();
305
- } else {
306
- const unrevealed = ops.logicalNot(draftReveal);
307
- const confidence = tokenProbability(processed, denoiser);
308
- acceptance = confidenceTransferMask(confidence, unrevealed, threshold, false);
309
- confidence.dispose();
310
- const accepted = ops.where(acceptance, denoiser, draftCanvas);
311
- const revealOrAccept = ops.logicalOr(draftReveal, acceptance);
312
- const noise = ops.randint(0, vocab, [1, canvasLength], I32); // RNG: re-noise
313
- const newCurrent = ops.where(revealOrAccept, accepted, noise);
314
- noise.dispose();
315
- currentCanvas.dispose();
316
- currentCanvas = newCurrent;
317
- const newReveal = ops.logicalOr(draftReveal, acceptance);
318
- draftReveal.dispose();
319
- draftReveal = newReveal;
320
- const newDraft = ops.where(acceptance, accepted, draftCanvas);
321
- draftCanvas.dispose();
322
- draftCanvas = newDraft;
323
- unrevealed.dispose();
324
- revealOrAccept.dispose();
325
- accepted.dispose();
326
- }
327
- acceptance.dispose();
328
-
329
- // Early stop 1 (confidence sampler): all positions revealed.
330
- if (sampler === "confidence-threshold") {
331
- const allRevealed = ops.allReduce(draftReveal);
332
- const done = ops.itemBool(allRevealed);
333
- allRevealed.dispose();
334
- if (done) {
293
+ try {
294
+ currentCanvas = ops.randint(0, vocab, [1, canvasLength], I32); // RNG: init
295
+ draftReveal = ops.zeros([1, canvasLength], Dtype.bool);
296
+ draftCanvas = ops.reshape(currentCanvas, [1, canvasLength]); // alias copy
297
+ argmaxCanvas = ops.reshape(currentCanvas, [1, canvasLength]);
298
+
299
+ for (let curStep = maxSteps; curStep >= 1; curStep--) {
300
+ yield; // safe point before the next denoising step
301
+ stepsThisCanvas++;
302
+ let logits: MlxArray | null = null;
303
+ let processed: MlxArray | null = null;
304
+ let acceptance: MlxArray | null = null;
305
+ let nextSc: MlxArray | null = null;
306
+ try {
307
+ logits = graph.decoderLogits(currentCanvas, cache, scEmbeddings);
308
+ const schedT = linearTemp(curStep, maxSteps, tMin, tMax);
309
+ // Match the reference EXACTLY: divide (not reciprocal-multiply) — the
310
+ // f32 rounding differs and the confidence threshold is a hard cutoff
311
+ // that flips acceptance on a 1-ULP difference, diverging the trajectory.
312
+ const schedArr = ops.scalarLike(schedT, logits);
313
+ processed = ops.div(logits, schedArr);
314
+ schedArr.dispose();
315
+ logits.dispose();
316
+
317
+ argmaxCanvas.dispose();
318
+ const argmax = ops.argmaxAxis(processed, -1);
319
+ argmaxCanvas = argmax.astype(I32); // [1, L]
320
+ argmax.dispose();
321
+
322
+ if (curStep === 1) {
323
+ processed.dispose();
324
+ break;
325
+ }
326
+
327
+ // temperature 0 → denoiser canvas is the argmax.
328
+ const denoiser = argmaxCanvas;
329
+
330
+ if (sampler === "entropy-bound") {
331
+ // cur_step > 1 always here (cur_step==1 breaks before the sampler):
332
+ // entropy AND the next soft embeddings come from the SAME entropy
333
+ // chain (reference _diffusion_entropy_and_soft_embeddings) — NOT the
334
+ // softmax-precise path the confidence sampler uses.
335
+ const { probs, entropy } = entropyProbsChain(processed);
336
+ acceptance = entropyTransferMask(entropy, entropyBound);
337
+ entropy.dispose();
338
+ nextSc = softEmbeddingsFromProbs(probs, dequantWeight, graph.embedScale);
339
+ probs.dispose();
340
+ // accepted/current update (entropy variant)
341
+ const accepted = ops.where(acceptance, denoiser, currentCanvas);
342
+ const noise = ops.randint(0, vocab, [1, canvasLength], I32); // RNG: re-noise
343
+ const newCurrent = ops.where(acceptance, accepted, noise);
344
+ noise.dispose();
345
+ currentCanvas.dispose();
346
+ currentCanvas = newCurrent;
347
+ draftReveal.dispose();
348
+ const notAcceptance = ops.logicalNot(acceptance);
349
+ draftReveal = ops.logicalNot(notAcceptance); // = acceptance
350
+ notAcceptance.dispose();
351
+ draftCanvas.dispose();
352
+ draftCanvas = ops.reshape(argmaxCanvas, [1, canvasLength]);
353
+ accepted.dispose();
354
+ } else {
355
+ const unrevealed = ops.logicalNot(draftReveal);
356
+ const confidence = tokenProbability(processed, denoiser);
357
+ acceptance = confidenceTransferMask(confidence, unrevealed, threshold, false);
358
+ confidence.dispose();
359
+ const accepted = ops.where(acceptance, denoiser, draftCanvas);
360
+ const revealOrAccept = ops.logicalOr(draftReveal, acceptance);
361
+ const noise = ops.randint(0, vocab, [1, canvasLength], I32); // RNG: re-noise
362
+ const newCurrent = ops.where(revealOrAccept, accepted, noise);
363
+ noise.dispose();
364
+ currentCanvas.dispose();
365
+ currentCanvas = newCurrent;
366
+ const newReveal = ops.logicalOr(draftReveal, acceptance);
367
+ draftReveal.dispose();
368
+ draftReveal = newReveal;
369
+ const newDraft = ops.where(acceptance, accepted, draftCanvas);
370
+ draftCanvas.dispose();
371
+ draftCanvas = newDraft;
372
+ unrevealed.dispose();
373
+ revealOrAccept.dispose();
374
+ accepted.dispose();
375
+ }
376
+ acceptance.dispose();
377
+
378
+ // Early stop 1 (confidence sampler): all positions revealed.
379
+ if (sampler === "confidence-threshold") {
380
+ const allRevealed = ops.allReduce(draftReveal);
381
+ const done = ops.itemBool(allRevealed);
382
+ allRevealed.dispose();
383
+ if (done) {
384
+ processed.dispose();
385
+ break;
386
+ }
387
+ }
388
+
389
+ // Early stop 2: stable + confident (only when explicitly configured;
390
+ // the as-loaded oracle has stopping_config None and never stops here).
391
+ if (
392
+ stableStop &&
393
+ stableAndConfident(argmaxCanvas, processed, history, stabilityThreshold, stopConfidence)
394
+ ) {
395
+ // Entropy sampling precomputes the next step's self-conditioning
396
+ // tensor before evaluating this stop. The break means ownership is
397
+ // never transferred to scEmbeddings, so release the pending value.
398
+ nextSc?.dispose();
399
+ processed.dispose();
400
+ break;
401
+ }
402
+
403
+ // Self-conditioning feedback for the next step.
404
+ if (curStep > 1) {
405
+ if (nextSc === null) nextSc = graph.softEmbeddings(processed, dequantWeight);
406
+ if (scEmbeddings) scEmbeddings.dispose();
407
+ scEmbeddings = nextSc;
408
+ nextSc = null; // ownership transferred to persistent feedback
409
+ }
335
410
  processed.dispose();
336
- break;
411
+ } finally {
412
+ disposeResources([logits, processed, acceptance, nextSc]
413
+ .filter((array): array is MlxArray => array !== null));
337
414
  }
338
415
  }
339
416
 
340
- // Early stop 2: stable + confident (only when explicitly configured;
341
- // the as-loaded oracle has stopping_config None and never stops here).
342
- if (
343
- stableStop &&
344
- stableAndConfident(argmaxCanvas, processed, history, stabilityThreshold, stopConfidence)
345
- ) {
346
- // Entropy sampling precomputes the next step's self-conditioning
347
- // tensor before evaluating this stop. The break means ownership is
348
- // never transferred to scEmbeddings, so release the pending value.
349
- nextSc?.dispose();
350
- processed.dispose();
351
- break;
352
- }
353
-
354
- // Self-conditioning feedback for the next step.
355
- if (curStep > 1) {
356
- if (nextSc === null) nextSc = model.softEmbeddings(processed, dequantWeight);
357
- if (scEmbeddings) scEmbeddings.dispose();
358
- scEmbeddings = nextSc;
359
- }
360
- processed.dispose();
361
- }
362
-
363
- if (scEmbeddings) scEmbeddings.dispose();
364
- for (const hh of history) hh.dispose();
365
- currentCanvas.dispose();
366
- draftReveal.dispose();
367
- draftCanvas.dispose();
368
-
369
- // Emit from the final argmax canvas.
370
- const argmaxFloat = argmaxCanvas.astype(Dtype.float32);
371
- const blockTokens = [...argmaxFloat.toFloat32()].map((x) => Math.round(x));
372
- argmaxFloat.dispose();
373
- argmaxCanvas.dispose();
374
- blocks.push(blockTokens);
375
- totalSteps += stepsThisCanvas;
376
-
377
- for (const tok of blockTokens) {
378
- generated++;
379
- if (eosIds.has(tok)) {
380
- stopped = true;
381
- finishReason = "stop";
382
- break;
383
- }
384
- emitted.push(tok);
385
- if (generated >= maxNewTokens) {
386
- stopped = true;
387
- finishReason = "length";
388
- break;
417
+ if (scEmbeddings) scEmbeddings.dispose();
418
+ for (const hh of history) hh.dispose();
419
+ currentCanvas.dispose();
420
+ draftReveal.dispose();
421
+ draftCanvas.dispose();
422
+
423
+ // Emit from the final argmax canvas.
424
+ const argmaxFloat = argmaxCanvas.astype(Dtype.float32);
425
+ let blockTokens: number[];
426
+ try { blockTokens = [...argmaxFloat.toFloat32()].map((x) => Math.round(x)); }
427
+ finally { argmaxFloat.dispose(); }
428
+ argmaxCanvas.dispose();
429
+ blocks.push(blockTokens);
430
+ totalSteps += stepsThisCanvas;
431
+
432
+ for (const tok of blockTokens) {
433
+ generated++;
434
+ if (eosIds.has(tok)) {
435
+ stopped = true;
436
+ finishReason = "stop";
437
+ break;
438
+ }
439
+ emitted.push(tok);
440
+ if (generated >= maxNewTokens) {
441
+ stopped = true;
442
+ finishReason = "length";
443
+ break;
444
+ }
389
445
  }
446
+ if (stopped) break;
447
+ } finally {
448
+ disposeResources([currentCanvas, draftReveal, draftCanvas, argmaxCanvas, scEmbeddings, ...history]
449
+ .filter((array): array is MlxArray => array !== null));
390
450
  }
391
- if (stopped) break;
392
451
  }
452
+ } catch (error) {
453
+ failure = { error };
454
+ throw error;
393
455
  } finally {
394
- dequantWeight.dispose();
395
- for (const c of cache) c.dispose();
456
+ const cleanup = () => disposeResources([dequantWeight,
457
+ ...(stateOwner ? [{ dispose: () => stateOwner!.close() }] : [])]
458
+ .filter((resource): resource is DisposableResource => resource !== null));
459
+ if (failure) cleanupFailure(failure.error, cleanup);
460
+ else cleanup();
396
461
  }
397
462
 
398
463
  return { tokens: emitted, blocks, steps: totalSteps, finishReason };
package/src/embed.ts CHANGED
@@ -2,7 +2,7 @@
2
2
  // L2-normalized vector) with the Qwen3-Embedding input convention so the CLI
3
3
  // (`mlx-bun embed`), the server (`/v1/embeddings`), and in-repo experiments all
4
4
  // produce the SAME vectors — the ones verified bit-exact vs mlx-lm in
5
- // tests/qwen3-embed-parity.test.ts.
5
+ // tests/parity/qwen3-embed-parity.test.ts.
6
6
 
7
7
  import { Qwen3Model } from "./model/qwen3";
8
8
  import type { RuntimeModel } from "./model/factory";
@@ -0,0 +1,59 @@
1
+ import type { Cancellation } from "../contracts/generation";
2
+ import type { DisposableResource } from "../contracts/resources";
3
+ import { throwIfCancelled } from "./cancellation";
4
+
5
+ export class AdmissionRejected extends Error {
6
+ constructor() { super("admission queue is full"); this.name = "AdmissionRejected"; }
7
+ }
8
+
9
+ /** FIFO reservations. A cancelled waiter releases no active owner's capacity. */
10
+ export class AdmissionPool {
11
+ #active = 0;
12
+ #closed = false;
13
+ #waiting: Array<{ resolve(lease: DisposableResource): void; reject(error: unknown): void; unsubscribe(): void }> = [];
14
+ constructor(readonly capacity: number, readonly maxQueued = 64) {
15
+ if (!Number.isSafeInteger(capacity) || capacity < 1 || !Number.isSafeInteger(maxQueued) || maxQueued < 0)
16
+ throw new Error("invalid admission capacity");
17
+ }
18
+ get active(): number { return this.#active; }
19
+ get queued(): number { return this.#waiting.length; }
20
+
21
+ acquire(cancellation?: Cancellation): Promise<DisposableResource> {
22
+ if (this.#closed) return Promise.reject(new Error("admission is closed"));
23
+ try { if (cancellation) throwIfCancelled(cancellation); }
24
+ catch (error) { return Promise.reject(error); }
25
+ if (this.#active < this.capacity) return Promise.resolve(this.#lease());
26
+ if (this.#waiting.length >= this.maxQueued) return Promise.reject(new AdmissionRejected());
27
+ return new Promise((resolve, reject) => {
28
+ const waiter = { resolve, reject, unsubscribe() {} };
29
+ this.#waiting.push(waiter);
30
+ waiter.unsubscribe = cancellation?.subscribe((reason) => {
31
+ const index = this.#waiting.indexOf(waiter);
32
+ if (index < 0) return;
33
+ this.#waiting.splice(index, 1);
34
+ waiter.unsubscribe();
35
+ reject(new Error(`admission cancelled: ${reason}`));
36
+ }) ?? (() => {});
37
+ if (!this.#waiting.includes(waiter)) waiter.unsubscribe();
38
+ });
39
+ }
40
+
41
+ #lease(): DisposableResource {
42
+ this.#active++;
43
+ let released = false;
44
+ return { dispose: () => {
45
+ if (released) return;
46
+ released = true;
47
+ this.#active--;
48
+ const next = this.#waiting.shift();
49
+ if (next) { next.unsubscribe(); next.resolve(this.#lease()); }
50
+ } };
51
+ }
52
+
53
+ close(): void {
54
+ this.#closed = true;
55
+ for (const waiter of this.#waiting.splice(0)) {
56
+ waiter.unsubscribe(); waiter.reject(new Error("admission is closed"));
57
+ }
58
+ }
59
+ }
@@ -0,0 +1,38 @@
1
+ import type { CancelReason, Cancellation } from "../contracts/generation";
2
+
3
+ export class GenerationCancelled extends Error {
4
+ constructor(readonly reason: CancelReason) {
5
+ super(`generation cancelled: ${reason}`);
6
+ this.name = "GenerationCancelled";
7
+ }
8
+ }
9
+
10
+ export function throwIfCancelled(cancellation: Cancellation): void {
11
+ if (cancellation.reason !== undefined) throw new GenerationCancelled(cancellation.reason);
12
+ }
13
+
14
+ export class CancellationSource implements Cancellation {
15
+ #reason: CancelReason | undefined;
16
+ #listeners = new Set<(reason: CancelReason) => void>();
17
+ get reason(): CancelReason | undefined { return this.#reason; }
18
+
19
+ subscribe(listener: (reason: CancelReason) => void): () => void {
20
+ if (this.#reason !== undefined) {
21
+ listener(this.#reason);
22
+ return () => {};
23
+ }
24
+ this.#listeners.add(listener);
25
+ return () => { this.#listeners.delete(listener); };
26
+ }
27
+
28
+ cancel(reason: CancelReason): void {
29
+ if (this.#reason !== undefined) return;
30
+ this.#reason = reason;
31
+ const listeners = [...this.#listeners];
32
+ this.#listeners.clear();
33
+ // A faulty observer cannot prevent the remaining owners from cancelling.
34
+ for (const listener of listeners) {
35
+ try { listener(reason); } catch { /* Observers do not own settlement. */ }
36
+ }
37
+ }
38
+ }
@@ -0,0 +1,56 @@
1
+ import type { ExecutionPlanner, GenerationSession, InferenceEngine, RunControl, Timer } from "../contracts/generation";
2
+ import { EngineSession, type SessionResources } from "./session";
3
+
4
+ export interface EngineOptions {
5
+ readonly timer: Timer;
6
+ /** Deadline for an unread stream or a producer blocked by an idle consumer. */
7
+ readonly idleTimeoutMs?: number;
8
+ readonly maxQueuedEvents?: number;
9
+ readonly maxQueuedTokens?: number;
10
+ readonly maxTopLogprobs?: number;
11
+ /** Aggregate output buffers for active collecting sessions, in tokens (4 bytes each).
12
+ * Once terminal, output storage belongs to the caller. */
13
+ readonly maxCollectTokens?: number;
14
+ }
15
+
16
+ export function createInferenceEngine<Request, Metrics>(
17
+ planner: ExecutionPlanner<Request, Metrics>, options: EngineOptions,
18
+ ): InferenceEngine<Request, Metrics> {
19
+ const idleTimeoutMs = options.idleTimeoutMs ?? 30_000;
20
+ const maxQueuedEvents = options.maxQueuedEvents ?? 32;
21
+ const maxQueuedTokens = options.maxQueuedTokens ?? 256;
22
+ const maxCollectTokens = options.maxCollectTokens ?? 65_536;
23
+ const maxTopLogprobs = options.maxTopLogprobs ?? 20;
24
+ for (const [name, value] of Object.entries({ idleTimeoutMs, maxQueuedEvents, maxQueuedTokens, maxCollectTokens, maxTopLogprobs }))
25
+ if (!Number.isSafeInteger(value) || value <= 0) throw new Error(`${name} must be a positive safe integer`);
26
+ const sessions = new Set<GenerationSession<Metrics>>();
27
+ let closed = false;
28
+ let closing: Promise<void> | undefined;
29
+ let nextId = 0;
30
+ let reserved = 0;
31
+ const resources: SessionResources = {
32
+ timer: options.timer, idleTimeoutMs, maxQueuedEvents, maxQueuedTokens, maxTopLogprobs,
33
+ reserveOutput(tokens) {
34
+ if (tokens > maxCollectTokens - reserved) throw new Error("collection output capacity exceeded");
35
+ const buffer = new Uint32Array(tokens);
36
+ reserved += tokens;
37
+ let released = false;
38
+ return { buffer, release() { if (!released) { released = true; reserved -= tokens; } } };
39
+ },
40
+ };
41
+ return {
42
+ async open(request: Request, control: RunControl) {
43
+ if (closed) throw new Error("inference engine is closed");
44
+ let session: EngineSession<Metrics> | undefined;
45
+ session = new EngineSession(String(++nextId), (cancellation) => planner.plan(request, cancellation),
46
+ { ...control }, resources,
47
+ () => { if (session) sessions.delete(session); });
48
+ if (session.state !== "terminal") sessions.add(session);
49
+ return session;
50
+ },
51
+ close() {
52
+ closed = true;
53
+ return closing ??= Promise.all([...sessions].map((session) => session.cancel("engine_closed"))).then(() => {});
54
+ },
55
+ };
56
+ }
@@ -0,0 +1,70 @@
1
+ import type { Cancellation } from "../contracts/generation";
2
+ import type { DisposableResource } from "../contracts/resources";
3
+ import { AdmissionRejected } from "./admission";
4
+ import { throwIfCancelled } from "./cancellation";
5
+
6
+ type Mode = "shared" | "exclusive";
7
+ interface Waiter {
8
+ mode: Mode;
9
+ resolve(lease: DisposableResource): void;
10
+ reject(error: unknown): void;
11
+ unsubscribe(): void;
12
+ }
13
+
14
+ /** Shared worker activity; exclusive managed jobs. FIFO writers cannot starve
15
+ * behind newly arriving inference. This coordinates execution, not memory size. */
16
+ export class ExecutionCoordinator {
17
+ #readers = 0;
18
+ #writer = false;
19
+ #closed = false;
20
+ #queue: Waiter[] = [];
21
+ constructor(readonly maxQueued = 64) {
22
+ if (!Number.isSafeInteger(maxQueued) || maxQueued < 0) throw new Error("invalid coordination queue bound");
23
+ }
24
+ acquire(mode: Mode, cancellation?: Cancellation): Promise<DisposableResource> {
25
+ if (this.#closed) return Promise.reject(new Error("execution coordinator is closed"));
26
+ try { if (cancellation) throwIfCancelled(cancellation); }
27
+ catch (error) { return Promise.reject(error); }
28
+ if (!this.#queue.length && !this.#writer && (mode === "shared" || this.#readers === 0))
29
+ return Promise.resolve(this.#lease(mode));
30
+ if (this.#queue.length >= this.maxQueued) return Promise.reject(new AdmissionRejected());
31
+ return new Promise((resolve, reject) => {
32
+ const waiter: Waiter = { mode, resolve, reject, unsubscribe() {} };
33
+ this.#queue.push(waiter);
34
+ waiter.unsubscribe = cancellation?.subscribe((reason) => {
35
+ const index = this.#queue.indexOf(waiter);
36
+ if (index < 0) return;
37
+ this.#queue.splice(index, 1);
38
+ waiter.unsubscribe();
39
+ reject(new Error(`execution admission cancelled: ${reason}`));
40
+ this.#drain();
41
+ }) ?? (() => {});
42
+ if (!this.#queue.includes(waiter)) waiter.unsubscribe();
43
+ });
44
+ }
45
+ #lease(mode: Mode): DisposableResource {
46
+ if (mode === "shared") this.#readers++;
47
+ else this.#writer = true;
48
+ let disposed = false;
49
+ return { dispose: () => {
50
+ if (disposed) return;
51
+ disposed = true;
52
+ if (mode === "shared") this.#readers--;
53
+ else this.#writer = false;
54
+ this.#drain();
55
+ } };
56
+ }
57
+ #drain(): void {
58
+ while (!this.#writer && this.#queue.length) {
59
+ const next = this.#queue[0]!;
60
+ if (next.mode === "exclusive" && this.#readers > 0) return;
61
+ this.#queue.shift(); next.unsubscribe(); next.resolve(this.#lease(next.mode));
62
+ }
63
+ }
64
+ close(): void {
65
+ this.#closed = true;
66
+ for (const next of this.#queue.splice(0)) {
67
+ next.unsubscribe(); next.reject(new Error("execution coordinator is closed"));
68
+ }
69
+ }
70
+ }