mlx-bun 0.0.13 → 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.
- package/README.md +71 -578
- package/bin/mlx-bun.mjs +1 -1
- package/docs/README.md +91 -0
- package/docs/{design → archive}/hlg-sampling.md +3 -3
- package/docs/{design → archive/investigations}/audio-input-plan.md +4 -4
- package/docs/{design → archive/investigations}/colibri-glm52-port.md +1 -1
- package/docs/{design → archive/investigations}/diffusion-gemma-port.md +3 -3
- package/docs/{design → archive/investigations}/faithful-l1-consolidation.md +4 -4
- package/docs/{investigations → archive/investigations}/hlg-sampling-investigation.md +1 -1
- package/docs/{design → archive/investigations}/minicpm5-decode-megakernel.md +11 -11
- package/docs/{design → archive/investigations}/optimization_plan.md +1 -1
- package/docs/archive/investigations/pareto-specialized-runtime-findings.md +92 -0
- package/docs/{design → archive/investigations}/pre-colibri-stabilization.md +5 -5
- package/docs/{design → archive/investigations}/the-dreaming-handoff.md +2 -2
- package/docs/{planning → archive/planning}/chat-ui-vision.md +1 -1
- package/docs/{planning → archive/planning}/chunk-task-roadmap.md +1 -1
- package/docs/{planning → archive/planning}/curve-sampler-research-plan.md +4 -4
- package/docs/{planning → archive/planning}/memory-docs-and-dag-plan.md +1 -1
- package/docs/{planning → archive/planning}/release-notes-v0.0.9.md +1 -1
- package/docs/archive/planning/release-notes-v0.2.0.md +49 -0
- package/docs/{planning → archive/planning}/website-readme-pass-plan.md +1 -1
- package/docs/design/batching.md +569 -0
- package/docs/design/decode-speed-program.md +3178 -169
- package/docs/design/dreaming-nightly-pipeline.md +675 -282
- package/docs/design/generic-model-support.md +579 -334
- package/docs/design/kv-cache.md +680 -0
- package/docs/design/orpo-training.md +673 -1239
- package/docs/design/speculative-decoding.md +1627 -0
- package/docs/design/turboquant.md +2102 -0
- package/docs/design/unified-engine-frontier-plan.md +1831 -727
- package/docs/design/web-chat-redesign.md +666 -1083
- package/docs/planning/PRODUCT_ROADMAP.md +34 -1
- package/docs/planning/release-notes-v0.3.0.md +35 -0
- package/docs/planning/video-series.md +421 -0
- package/docs/reference/benchmarks.md +1162 -0
- package/docs/reference/cli.md +395 -149
- package/docs/reference/distribution.md +321 -181
- package/docs/reference/environment.md +348 -0
- package/docs/reference/glossary.md +39 -0
- package/docs/reference/library-api.md +580 -142
- package/docs/reference/memory.md +227 -64
- package/docs/reference/models.md +137 -6
- package/docs/reference/server-api.md +938 -422
- package/docs/reference/server-config.md +652 -320
- package/docs/reference/training.md +393 -254
- package/docs/reference/troubleshooting.md +166 -0
- package/package.json +10 -5
- package/src/admission-plan.ts +18 -0
- package/src/anthropic.ts +92 -3
- package/src/audio/decode.ts +2 -2
- package/src/audio/features.ts +2 -2
- package/src/backends/mlx/autoregressive.ts +108 -0
- package/src/backends/mlx/batch-group.ts +1646 -0
- package/src/backends/mlx/diffusion.ts +39 -0
- package/src/backends/mlx/draft-target.ts +47 -0
- package/src/backends/mlx/evaluation.ts +105 -0
- package/src/backends/mlx/fill-append.ts +45 -0
- package/src/backends/mlx/gateway-binding.ts +71 -0
- package/src/backends/mlx/graph.ts +38 -0
- package/src/backends/mlx/memory-client.ts +189 -0
- package/src/backends/mlx/memory-guard.ts +40 -0
- package/src/backends/mlx/methods.ts +90 -0
- package/src/backends/mlx/model-prompt.ts +164 -0
- package/src/backends/mlx/model-serving.ts +79 -0
- package/src/backends/mlx/prefill.ts +33 -0
- package/src/backends/mlx/request-policy.ts +20 -0
- package/src/backends/mlx/rollback.ts +51 -0
- package/src/backends/mlx/serial-executor.ts +273 -0
- package/src/backends/mlx/speculative.ts +114 -0
- package/src/backends/mlx/state-views.ts +26 -0
- package/src/backends/mlx/text-engine.ts +73 -0
- package/src/cli.ts +189 -61
- package/src/client.ts +66 -0
- package/src/config.ts +21 -4
- package/src/contracts/completion.ts +16 -0
- package/src/contracts/execution.ts +49 -0
- package/src/contracts/generation.ts +103 -0
- package/src/contracts/host.ts +7 -0
- package/src/contracts/jobs.ts +66 -0
- package/src/contracts/pi-provider.ts +30 -0
- package/src/contracts/pi.ts +259 -0
- package/src/contracts/resources.ts +18 -0
- package/src/contracts/scheduling.ts +27 -0
- package/src/contracts/tsconfig.json +14 -0
- package/src/diffusion/diffusion-generate.ts +218 -153
- package/src/embed.ts +1 -1
- package/src/engine/admission.ts +59 -0
- package/src/engine/cancellation.ts +38 -0
- package/src/engine/engine.ts +56 -0
- package/src/engine/execution-coordinator.ts +70 -0
- package/src/engine/execution-plan.ts +51 -0
- package/src/engine/resources.ts +42 -0
- package/src/engine/scheduler.ts +67 -0
- package/src/engine/session.ts +247 -0
- package/src/eval/README.md +44 -1
- package/src/eval/kl.ts +263 -23
- package/src/eval/reference-logits.ts +207 -0
- package/src/eval/runner.ts +10 -79
- package/src/eval/tasks/bfcl.ts +2 -1
- package/src/eval/tasks/gsm8k.ts +2 -1
- package/src/eval/tasks/hashhop.ts +2 -1
- package/src/eval/tasks/humaneval.ts +5 -4
- package/src/eval/tasks/ifeval.ts +2 -1
- package/src/eval/tasks/mmlu.ts +2 -1
- package/src/expert-io.ts +3 -1
- package/src/expert-offload.ts +3 -1
- package/src/fill/echo-index.ts +188 -0
- package/src/fill/fill-session.ts +452 -0
- package/src/fill/proposal.ts +68 -0
- package/src/fill/schema-rows.ts +506 -0
- package/src/fill/tool-boundary.ts +103 -0
- package/src/fit.ts +24 -43
- package/src/generate.ts +740 -337
- package/src/grammar.ts +14 -8
- package/src/harness-pi.ts +15 -34
- package/src/hub-rest.ts +4 -4
- package/src/index.ts +27 -0
- package/src/inference/denoising.ts +17 -0
- package/src/inference/graph.ts +33 -0
- package/src/inference/prefill.ts +39 -0
- package/src/inference/rollback.ts +7 -0
- package/src/jobs/db.ts +4 -0
- package/src/jobs/index.ts +2 -0
- package/src/jobs/job-entry.ts +15 -14
- package/src/jobs/runner.ts +98 -13
- package/src/jobs/task-client.ts +19 -0
- package/src/jobs/types.ts +2 -58
- package/src/kv-scheme.ts +199 -0
- package/src/kv-store.ts +459 -234
- package/src/lab/README.md +18 -0
- package/src/{assets → lab/curve}/curve-designer.html +2 -2
- package/src/{curve-sampler.ts → lab/curve/curve-sampler.ts} +2 -2
- package/src/{expert-trace.ts → lab/expert-trace/expert-trace.ts} +7 -5
- package/src/{model → lab/paged-kv}/paged-kv.ts +9 -6
- package/src/library.ts +40 -0
- package/src/lora.ts +18 -0
- package/src/media-fetch.ts +20 -4
- package/src/memory/article.ts +1 -1
- package/src/memory/chunk.ts +1 -1
- package/src/memory/model.ts +32 -170
- package/src/memory/prompts.ts +1 -1
- package/src/memory/rest.ts +1 -1
- package/src/memory/schedule.ts +1 -1
- package/src/memory/stages.ts +4 -1
- package/src/memory/vault.ts +3 -4
- package/src/mlx/array.ts +9 -10
- package/src/mlx/compile.ts +12 -2
- package/src/mlx/ffi.ts +66 -22
- package/src/mlx/materialize.ts +23 -0
- package/src/mlx/metal-kernel.ts +10 -11
- package/src/mlx/ops.ts +64 -19
- package/src/mlx/quantized-rows.ts +36 -0
- package/src/mlx/token-bitmask.ts +52 -0
- package/src/mlx/turboquant-kv-decode.ts +82 -0
- package/src/mlx/turboquant-kv-inverse.ts +122 -0
- package/src/mlx/turboquant-ops.ts +33 -3
- package/src/mlx/turboquant-tables.ts +2 -2
- package/src/model/batched-mask.ts +9 -3
- package/src/model/batched-quant.ts +2 -2
- package/src/model/batched-rotating-quant.ts +45 -99
- package/src/model/batched-rotating-state.ts +73 -0
- package/src/model/batched-rotating.ts +57 -106
- package/src/model/batched-row-storage.ts +127 -0
- package/src/model/compiled-decode.ts +4 -3
- package/src/model/diffusion-gemma.ts +8 -8
- package/src/model/factory.ts +87 -49
- package/src/model/fingerprint.ts +1 -1
- package/src/model/flash-attention.ts +3 -3
- package/src/model/gemma4-base.ts +187 -17
- package/src/model/gemma4.ts +6 -6
- package/src/model/generated/gemma4-12b.ts +52 -148
- package/src/model/generated/gemma4-26b.ts +34 -94
- package/src/model/generated/gemma4-e4b.ts +28 -76
- package/src/model/glm52-cache.ts +19 -9
- package/src/model/glm52-dsa.ts +13 -9
- package/src/model/glm52-memory.ts +28 -1
- package/src/model/implementation.ts +69 -0
- package/src/model/minicpm5.ts +7 -8
- package/src/model/profile.ts +443 -0
- package/src/model/qwen-append.ts +12 -0
- package/src/model/qwen-conv.ts +52 -0
- package/src/model/qwen3-delta.ts +98 -2
- package/src/model/qwen3-mrope.ts +252 -0
- package/src/model/qwen3_5-checkpoint.ts +164 -0
- package/src/model/qwen3_5.ts +336 -60
- package/src/model/support.ts +1 -1
- package/src/model/trellis-balanced-scatter.ts +51 -0
- package/src/model/trellis-linear.ts +672 -0
- package/src/model/trellis-shared-m.ts +151 -0
- package/src/model/trellis-shared-scatter.ts +144 -0
- package/src/model/trellis-splitk-prefill.ts +158 -0
- package/src/model/trellis-tiled-prefill.ts +136 -0
- package/src/model/trellis-vector-expand.ts +62 -0
- package/src/model/trellis-wide-prefill.ts +94 -0
- package/src/model/universal/dense.ts +1 -1
- package/src/model/universal/modules.ts +8 -1
- package/src/native/frame_extract.swift +140 -0
- package/src/native-pack.ts +15 -7
- package/src/pi-provider.ts +6 -64
- package/src/pi-web.ts +14 -320
- package/src/prompt-cache.ts +97 -54
- package/src/quantize/allocator.ts +1 -1
- package/src/quantize/atomic-output.ts +38 -0
- package/src/quantize/calibration.ts +2 -1
- package/src/quantize/config-writer.ts +16 -3
- package/src/quantize/index.ts +28 -1
- package/src/quantize/job.ts +11 -0
- package/src/quantize/quantizer.ts +265 -100
- package/src/quantize/rotate.ts +603 -0
- package/src/quantize/safetensors-writer.ts +79 -1
- package/src/quantize/sensitivity.ts +41 -59
- package/src/quantize/trellis.ts +586 -0
- package/src/quantize/weight-transform.ts +292 -0
- package/src/responses.ts +130 -3
- package/src/runtime-config.ts +81 -0
- package/src/sampler.ts +276 -4
- package/src/serve/admin-routes.ts +184 -0
- package/src/serve/aux-routes.ts +132 -0
- package/src/serve/batch-scheduler.ts +8 -1558
- package/src/serve/chat-request.ts +391 -0
- package/src/serve/chat-stage.ts +195 -0
- package/src/serve/checkpoint-identity.ts +64 -0
- package/src/serve/completion-executor.ts +354 -0
- package/src/serve/completion-sink.ts +157 -0
- package/src/serve/discovery-routes.ts +204 -0
- package/src/serve/generation-gateway.ts +211 -243
- package/src/serve/http.ts +159 -0
- package/src/serve/inference-request.ts +79 -0
- package/src/serve/isolate.ts +297 -105
- package/src/serve/kv-budget.ts +14 -0
- package/src/serve/lab-routes.ts +341 -0
- package/src/serve/lane-registry.ts +2 -1
- package/src/serve/model-admin-routes.ts +155 -0
- package/src/serve/model-binding.ts +43 -0
- package/src/serve/model-host.ts +512 -0
- package/src/serve/openai-wire.ts +103 -0
- package/src/serve/parent-application.ts +56 -0
- package/src/serve/pipeline.ts +33 -0
- package/src/serve/preparation.ts +33 -0
- package/src/serve/prompt-response-trace.ts +147 -0
- package/src/serve/request-plan.ts +162 -0
- package/src/serve/request-prep.ts +346 -0
- package/src/serve/response-lease.ts +24 -0
- package/src/serve/responses-client.ts +69 -0
- package/src/serve/session-completion-engine.ts +110 -0
- package/src/serve/session-search.ts +1 -1
- package/src/serve/static-routes.ts +101 -0
- package/src/serve/text-completion-stage.ts +104 -0
- package/src/serve/token-streams.ts +519 -0
- package/src/server.ts +523 -3551
- package/src/spec/assistant-source.ts +10 -44
- package/src/spec/deepspec-source.ts +6 -6
- package/src/spec/dflash-source.ts +9 -8
- package/src/spec/dspark/ab-stats.ts +3 -3
- package/src/spec/dspark/calibration.ts +2 -2
- package/src/spec/dspark/deepspec-module.ts +2 -2
- package/src/spec/dspark/generate-dflash.ts +2 -2
- package/src/spec/dspark/generate.ts +2 -2
- package/src/spec/dspark/loader.ts +1 -1
- package/src/spec/dspark/module-dflash.ts +5 -4
- package/src/spec/dspark/module.ts +3 -3
- package/src/spec/dspark/quantize-drafter.ts +1 -1
- package/src/spec/generate.ts +1 -1
- package/src/spec/glm52-mtp-source.ts +9 -1
- package/src/spec/ngram-source.ts +1 -1
- package/src/spec/prefix-state.ts +45 -0
- package/src/spec/qwen-mtp-source.ts +528 -0
- package/src/spec/serve-loop.ts +179 -213
- package/src/spec/source.ts +42 -10
- package/src/spec/two-model.ts +1 -1
- package/src/ssd-cache.ts +95 -24
- package/src/ssd-durability.ts +207 -0
- package/src/tool-call.ts +51 -14
- package/src/train/flash-cce.ts +14 -13
- package/src/train/forward.ts +38 -1
- package/src/train/loss.ts +3 -2
- package/src/train/prefix-shared.ts +5 -1
- package/src/train/segmented.ts +9 -6
- package/src/train/trainer.ts +9 -8
- package/src/vision/preprocess.ts +1 -1
- package/src/vision/prompt.ts +48 -2
- package/src/vision/qwen3vl-preprocess.ts +409 -0
- package/src/vision/qwen3vl-prompt.ts +196 -0
- package/src/vision/qwen3vl-tower.ts +555 -0
- package/src/vision/video-frames.ts +148 -0
- package/src/web/app.html +1 -1
- package/src/web/app.js +15 -6
- package/src/web/src/adapters-panel.ts +3 -3
- package/src/web/src/assistant.ts +2 -36
- package/src/web/src/chat.ts +16 -9
- package/src/web/src/composer.ts +4 -4
- package/src/web/src/hub.ts +1 -1
- package/src/web/src/markdown.ts +0 -0
- package/src/web/src/memory-panel.ts +1 -1
- package/src/web/src/model-picker.ts +1 -1
- package/src/web/src/palette.ts +1 -1
- package/src/web/src/protocol.ts +3 -39
- package/src/web/src/sessions.ts +1 -1
- package/src/web/src/ui-catalog.ts +4 -18
- package/src/web-assets.ts +47 -0
- package/src/weights.ts +134 -33
- package/docs/archive/mlx-bun-lab-report.html +0 -275
- package/docs/archive/wwdc-mlx-bun.html +0 -848
- package/docs/design/adapters-end-to-end.md +0 -205
- package/docs/design/batching-perf-path.md +0 -255
- package/docs/design/batching-v2-plan.md +0 -529
- package/docs/design/bucketing-stage.md +0 -197
- package/docs/design/compat-cli-surface-design.md +0 -230
- package/docs/design/docs-reorg-plan.md +0 -165
- package/docs/design/dspark-serving-program.md +0 -371
- package/docs/design/dspark-speculative-decoding.md +0 -268
- package/docs/design/grammar-spec-batching-integration.md +0 -301
- package/docs/design/memory-inference-path.md +0 -422
- package/docs/design/memory-synthesis.md +0 -287
- package/docs/design/memory-system.md +0 -237
- package/docs/design/mlx-lm-tool-parity-plan.md +0 -545
- package/docs/design/omlx-adoption-map.md +0 -98
- package/docs/design/orpo-dynamic-lambda.md +0 -220
- package/docs/design/orpo-future-enhancements.md +0 -340
- package/docs/design/paged-kv-cache.md +0 -214
- package/docs/design/parallel-slots.md +0 -567
- package/docs/design/parity-tier-dag.md +0 -161
- package/docs/design/repo-cleanup-plan.md +0 -189
- package/docs/design/runtime-isolation.md +0 -263
- package/docs/design/segmented-backward-training.md +0 -543
- package/docs/design/spec-decode-larger-targets.md +0 -229
- package/docs/design/ssd-kv-cold-tier.md +0 -288
- package/docs/design/structured-output.md +0 -484
- package/docs/design/tauri-desktop-app.md +0 -127
- package/docs/design/the-dreaming-master-plan.md +0 -445
- package/docs/design/turboquant-kv.md +0 -235
- package/docs/design/w4a16-compute-precision-spike.md +0 -232
- package/docs/design/web-chat-beat-matrix.md +0 -315
- package/docs/design/write-pipeline-entity.md +0 -468
- package/docs/investigations/curve-runs/bisector-route.json +0 -98
- package/docs/investigations/curve-runs/bisector-trace.json +0 -1393
- package/docs/investigations/curve-runs/curve-terrain.json +0 -1179
- package/docs/investigations/curve-runs/curve-terrain.svg +0 -42
- package/docs/investigations/hlg-ab.html +0 -34
- package/docs/investigations/hlg-explorer.html +0 -238
- package/docs/investigations/hlg-figs/cliff.svg +0 -1
- package/docs/investigations/hlg-figs/cross.svg +0 -1
- package/docs/investigations/hlg-figs/edges.svg +0 -1
- package/docs/investigations/hlg-figs/k20-grids.svg +0 -1
- package/docs/investigations/hlg-figs/k20-peaks.svg +0 -1
- package/docs/investigations/hlg-figs/lw-curve.svg +0 -1
- package/docs/investigations/hlg-figs/terrain.svg +0 -1
- package/docs/investigations/hlg-report.html +0 -262
- package/docs/investigations/hlg-runs/hlg-ab.json +0 -149
- package/docs/investigations/hlg-runs/hlg-blind-prompts.md +0 -87
- package/docs/investigations/hlg-runs/hlg-blind.json +0 -107
- package/docs/investigations/hlg-runs/hlg-cliff.json +0 -197
- package/docs/investigations/hlg-runs/hlg-cross.json +0 -505
- package/docs/investigations/hlg-runs/hlg-edge.json +0 -388
- package/docs/investigations/hlg-runs/hlg-eval1.json +0 -25
- package/docs/investigations/hlg-runs/hlg-frontier.json +0 -86
- package/docs/investigations/hlg-runs/hlg-grid.json +0 -242
- package/docs/investigations/hlg-runs/hlg-grid2.json +0 -343
- package/docs/investigations/hlg-runs/hlg-lw-coarse.json +0 -110
- package/docs/investigations/hlg-runs/hlg-map.json +0 -383
- package/docs/investigations/hlg-runs/run0-original.log +0 -159
- package/docs/investigations/hlg-runs/runA-recommended.log +0 -131
- package/docs/investigations/hlg-runs/runAA-cliff-cpm5.log +0 -26
- package/docs/investigations/hlg-runs/runAB-grid2-e4b.log +0 -18
- package/docs/investigations/hlg-runs/runAC-peak-e4b-partial.log +0 -25
- package/docs/investigations/hlg-runs/runAD-ab-samples.log +0 -8
- package/docs/investigations/hlg-runs/runAE-edge-e4b.log +0 -51
- package/docs/investigations/hlg-runs/runAF-logit-calibration.log +0 -10
- package/docs/investigations/hlg-runs/runAG-eval-config2.log +0 -16
- package/docs/investigations/hlg-runs/runB-gain-variants.log +0 -192
- package/docs/investigations/hlg-runs/runC-pivots.log +0 -162
- package/docs/investigations/hlg-runs/runD-oetf.log +0 -132
- package/docs/investigations/hlg-runs/runE-eotf.log +0 -162
- package/docs/investigations/hlg-runs/runF-shaper-working.log +0 -86
- package/docs/investigations/hlg-runs/runG-shaper-sweep.log +0 -101
- package/docs/investigations/hlg-runs/runH-wsweep.log +0 -176
- package/docs/investigations/hlg-runs/runI-asweep.log +0 -176
- package/docs/investigations/hlg-runs/runJ-knobmap.log +0 -49
- package/docs/investigations/hlg-runs/runK-frontier.log +0 -19
- package/docs/investigations/hlg-runs/runL-widemap.log +0 -65
- package/docs/investigations/hlg-runs/runM-frontier-emb.log +0 -15
- package/docs/investigations/hlg-runs/runN-verify.log +0 -116
- package/docs/investigations/hlg-runs/runO-frontier-seed5000.log +0 -15
- package/docs/investigations/hlg-runs/runP-grid-AxSm.log +0 -15
- package/docs/investigations/hlg-runs/runQ-grid-tgW.log +0 -15
- package/docs/investigations/hlg-runs/runR-blind-transcript.log +0 -79
- package/docs/investigations/hlg-runs/runS-cross-coherence.log +0 -33
- package/docs/investigations/hlg-runs/runT-cross-cpm5.log +0 -33
- package/docs/investigations/hlg-runs/runU-cross-12b.log +0 -33
- package/docs/investigations/hlg-runs/runV-lw-e4b.log +0 -22
- package/docs/investigations/hlg-runs/runW-lw-12b.log +0 -22
- package/docs/investigations/hlg-runs/runX-lw-cpm5.log +0 -22
- package/docs/investigations/hlg-runs/runY-cliff-e4b.log +0 -26
- package/docs/investigations/hlg-runs/runZ-cliff-12b.log +0 -26
- package/docs/planning/IDEAS.md +0 -69
- package/docs/reference/embedding.md +0 -197
- package/docs/reference/features-matrix.md +0 -100
- package/docs/reference/orpo-quickstart.md +0 -135
- package/src/flags.ts +0 -13
- /package/docs/{investigations → archive/investigations}/chunk-finetune-experiment.md +0 -0
- /package/docs/{investigations → archive/investigations}/colibri-oracle-pin.md +0 -0
- /package/docs/{investigations → archive/investigations}/curve-bisector-routing.md +0 -0
- /package/docs/{investigations → archive/investigations}/decode-roofline-lookagain.md +0 -0
- /package/docs/{investigations → archive/investigations}/dspark-handoff.md +0 -0
- /package/docs/{investigations → archive/investigations}/expert-offload-single-user-moe.md +0 -0
- /package/docs/{investigations → archive/investigations}/indexshare-performance-spike.md +0 -0
- /package/docs/{investigations → archive/investigations}/kernel-perf-review-2026-07.md +0 -0
- /package/docs/{investigations → archive/investigations}/lab-build-journal.md +0 -0
- /package/docs/{investigations → archive/investigations}/openwiki-evaluation.md +0 -0
- /package/docs/{investigations → archive/investigations}/orpo-base-uf-experiment-and-directions.md +0 -0
- /package/docs/{investigations → archive/investigations}/orpo-flash-cce-pin-leak.md +0 -0
- /package/docs/{investigations → archive/investigations}/orpo-uf-testing-handoff.md +0 -0
- /package/docs/{investigations → archive/investigations}/pi-builtin-investigation.md +0 -0
- /package/docs/{investigations → archive/investigations}/segmented-backward-handoff.md +0 -0
- /package/docs/{investigations → archive/investigations}/starter-model-port-handoff.md +0 -0
- /package/docs/{investigations → archive/investigations}/steel-flash-cce-handoff.md +0 -0
- /package/docs/{investigations → archive/investigations}/trainer-validation-experiment.md +0 -0
- /package/docs/{planning → archive/planning}/chat-ui-vision-memory-first.md +0 -0
- /package/docs/{planning → archive/planning}/journal.md +0 -0
- /package/docs/{planning → archive/planning}/pi-integration-review.md +0 -0
- /package/docs/{planning → archive/planning}/release-notes-v0.0.10.md +0 -0
- /package/docs/{planning → archive/planning}/release-notes-v0.0.11.md +0 -0
- /package/docs/{planning → archive/planning}/release-notes-v0.0.12.md +0 -0
- /package/docs/{planning → archive/planning}/release-notes-v0.0.13.md +0 -0
- /package/docs/{planning → archive/planning}/web-ui-pass-plan.md +0 -0
- /package/{docs/investigations → src/serve/assets}/curve-terrain.html +0 -0
- /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
|
-
|
|
188
|
-
|
|
189
|
-
promptIds
|
|
190
|
-
|
|
191
|
-
|
|
192
|
-
|
|
193
|
-
const
|
|
194
|
-
|
|
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
|
-
|
|
218
|
-
|
|
219
|
-
|
|
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
|
-
|
|
242
|
-
prevArr.dispose();
|
|
281
|
+
try { graph.extendPrefill(prevArr, cache); }
|
|
282
|
+
finally { prevArr.dispose(); }
|
|
243
283
|
}
|
|
244
284
|
isFirstBlock = false;
|
|
245
285
|
|
|
246
|
-
let currentCanvas
|
|
247
|
-
let draftReveal
|
|
248
|
-
let draftCanvas
|
|
249
|
-
let argmaxCanvas
|
|
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
|
-
|
|
255
|
-
|
|
256
|
-
|
|
257
|
-
|
|
258
|
-
|
|
259
|
-
|
|
260
|
-
|
|
261
|
-
|
|
262
|
-
|
|
263
|
-
|
|
264
|
-
|
|
265
|
-
|
|
266
|
-
|
|
267
|
-
|
|
268
|
-
|
|
269
|
-
|
|
270
|
-
|
|
271
|
-
|
|
272
|
-
|
|
273
|
-
|
|
274
|
-
|
|
275
|
-
|
|
276
|
-
|
|
277
|
-
|
|
278
|
-
|
|
279
|
-
|
|
280
|
-
|
|
281
|
-
|
|
282
|
-
|
|
283
|
-
|
|
284
|
-
|
|
285
|
-
|
|
286
|
-
|
|
287
|
-
|
|
288
|
-
|
|
289
|
-
|
|
290
|
-
|
|
291
|
-
|
|
292
|
-
|
|
293
|
-
|
|
294
|
-
|
|
295
|
-
|
|
296
|
-
|
|
297
|
-
|
|
298
|
-
|
|
299
|
-
|
|
300
|
-
|
|
301
|
-
|
|
302
|
-
|
|
303
|
-
|
|
304
|
-
|
|
305
|
-
|
|
306
|
-
|
|
307
|
-
|
|
308
|
-
|
|
309
|
-
|
|
310
|
-
|
|
311
|
-
|
|
312
|
-
|
|
313
|
-
|
|
314
|
-
|
|
315
|
-
|
|
316
|
-
|
|
317
|
-
|
|
318
|
-
|
|
319
|
-
|
|
320
|
-
|
|
321
|
-
|
|
322
|
-
|
|
323
|
-
|
|
324
|
-
|
|
325
|
-
|
|
326
|
-
|
|
327
|
-
|
|
328
|
-
|
|
329
|
-
|
|
330
|
-
|
|
331
|
-
|
|
332
|
-
|
|
333
|
-
|
|
334
|
-
|
|
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
|
-
|
|
411
|
+
} finally {
|
|
412
|
+
disposeResources([logits, processed, acceptance, nextSc]
|
|
413
|
+
.filter((array): array is MlxArray => array !== null));
|
|
337
414
|
}
|
|
338
415
|
}
|
|
339
416
|
|
|
340
|
-
|
|
341
|
-
|
|
342
|
-
|
|
343
|
-
|
|
344
|
-
|
|
345
|
-
|
|
346
|
-
|
|
347
|
-
|
|
348
|
-
|
|
349
|
-
|
|
350
|
-
|
|
351
|
-
|
|
352
|
-
|
|
353
|
-
|
|
354
|
-
|
|
355
|
-
|
|
356
|
-
|
|
357
|
-
if (
|
|
358
|
-
|
|
359
|
-
|
|
360
|
-
|
|
361
|
-
|
|
362
|
-
|
|
363
|
-
|
|
364
|
-
|
|
365
|
-
|
|
366
|
-
|
|
367
|
-
|
|
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
|
-
|
|
395
|
-
|
|
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
|
+
}
|