mlx-bun 0.4.0 → 1.0.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 +1453 -88
- package/THIRD_PARTY_NOTICES.md +882 -0
- package/bin/mlx-bun.convert.mjs +4 -0
- package/bin/mlx-bun.fuse.mjs +4 -0
- package/bin/mlx-bun.generate.mjs +4 -0
- package/bin/mlx-bun.lora.mjs +4 -0
- package/bin/mlx-bun.mjs +13 -21
- package/bin/mlx-bun.server.mjs +4 -0
- package/bin/mlx-bun.upload.mjs +4 -0
- package/dist/web/app.js +9584 -0
- package/package.json +50 -53
- package/src/cli/ane-install.ts +76 -0
- package/src/cli/args.ts +179 -0
- package/src/cli/engine-entry.ts +68 -0
- package/src/cli/inference.ts +214 -0
- package/src/cli/job-entry.ts +97 -0
- package/src/cli/library-host.ts +117 -0
- package/src/cli/main.ts +57 -0
- package/src/cli/memory-engine.ts +163 -0
- package/src/cli/memory.ts +36 -0
- package/src/cli/mlx-lm-aliases.ts +309 -0
- package/src/cli/model-choice.ts +33 -0
- package/src/cli/model-selection.ts +84 -0
- package/src/cli/module-host.ts +11 -0
- package/src/cli/module-verbs.ts +50 -0
- package/src/cli/numerical-policy.ts +12 -0
- package/src/cli/perplexity.ts +89 -0
- package/src/cli/serve-host.ts +423 -0
- package/src/cli/serve-isolated.ts +217 -0
- package/src/cli/serve-options.ts +106 -0
- package/src/cli/serve-state.ts +205 -0
- package/src/cli/serve.ts +500 -0
- package/src/cli/served-model-host.ts +102 -0
- package/src/cli/server-entry.ts +148 -0
- package/src/cli/serving-unit.ts +177 -0
- package/src/{tui.ts → cli/terminal.ts} +13 -8
- package/src/cli/worker-adapters.ts +34 -0
- package/src/cli/worker-entry.ts +272 -0
- package/src/cli/worker-unit.ts +179 -0
- package/src/engine/adapter-operation.ts +26 -0
- package/src/engine/cache-services.ts +114 -0
- package/src/engine/completion.ts +39 -0
- package/src/engine/generation-gateway.ts +358 -0
- package/src/engine/index.ts +59 -0
- package/src/engine/kv-budget.ts +120 -0
- package/src/engine/media-preparation.ts +232 -0
- package/src/engine/model-binding.ts +24 -0
- package/src/engine/model-host.ts +513 -0
- package/src/engine/model-serving.ts +89 -0
- package/src/{serve → engine}/preparation.ts +10 -8
- package/src/{serve → engine}/session-completion-engine.ts +6 -6
- package/src/engine/telemetry.ts +109 -0
- package/src/hub/downloads.ts +116 -0
- package/src/jobs/db.ts +16 -7
- package/src/jobs/events.ts +29 -0
- package/src/jobs/executable.ts +5 -0
- package/src/jobs/host.ts +51 -0
- package/src/jobs/in-process.ts +55 -0
- package/src/{contracts/jobs.ts → jobs/protocol.ts} +7 -9
- package/src/jobs/runner.ts +150 -231
- package/src/jobs/service.ts +75 -0
- package/src/jobs/sse.ts +59 -23
- package/src/jobs/worker-memory.ts +46 -0
- package/src/jobs/worker-process.ts +163 -0
- package/src/jobs/worker-supervisor.ts +258 -0
- package/src/modules.ts +58 -0
- package/src/publishing/catalog-hub.ts +18 -0
- package/src/publishing/credentials.ts +31 -0
- package/src/publishing/upload.ts +37 -0
- package/src/residency/leased-adapters.ts +21 -0
- package/src/residency/local-records.ts +16 -0
- package/src/residency/model-residency.ts +405 -0
- package/src/residency/record-index.ts +28 -0
- package/src/residency/resident-estimate.ts +50 -0
- package/src/{anthropic.ts → server/anthropic.ts} +3 -81
- package/src/server/cache-routes.ts +36 -0
- package/src/{serve → server}/chat-request.ts +12 -10
- package/src/{serve → server}/chat-stage.ts +27 -22
- package/src/{client.ts → server/client.ts} +28 -2
- package/src/{serve → server}/completion-executor.ts +42 -42
- package/src/server/decision-body.ts +55 -0
- package/src/server/decision-routes.ts +134 -0
- package/src/server/decision-wire.ts +228 -0
- package/src/{serve → server}/discovery-routes.ts +42 -83
- package/src/server/generated-token-history.ts +57 -0
- package/src/{serve → server}/http.ts +25 -13
- package/src/{serve → server}/inference-request.ts +5 -5
- package/src/server/job-routes.ts +20 -0
- package/src/server/media-prompt.ts +63 -0
- package/src/server/memory-completion-client.ts +46 -0
- package/src/server/model-routes.ts +192 -0
- package/src/{serve → server}/openai-wire.ts +3 -2
- package/src/{serve → server}/pipeline.ts +1 -1
- package/src/server/prompt-builder.ts +13 -0
- package/src/server/prompt-contracts.ts +17 -0
- package/src/server/proxy-routes.ts +187 -0
- package/src/server/publishing-routes.ts +41 -0
- package/src/{serve → server}/request-plan.ts +8 -8
- package/src/{serve → server}/request-prep.ts +82 -68
- package/src/{serve → server}/responses-client.ts +30 -13
- package/src/{responses.ts → server/responses.ts} +12 -93
- package/src/server/routes.ts +205 -0
- package/src/server/start.ts +72 -0
- package/src/server/status-routes.ts +108 -0
- package/src/{serve → server}/text-completion-stage.ts +25 -4
- package/src/server/text-prompt.ts +12 -0
- package/src/{serve → server}/token-streams.ts +36 -38
- package/src/server/worker-routes.ts +311 -0
- package/src/storage/paths.ts +61 -0
- package/src/web/assets.ts +37 -0
- package/src/web/{src → browser}/api.ts +5 -23
- package/src/web/browser/chat-host.ts +17 -0
- package/src/web/browser/installed-panels.ts +6 -0
- package/src/web/browser/main.ts +55 -0
- package/src/web/browser/markdown.ts +279 -0
- package/src/web/browser/memory-panel.ts +9 -0
- package/src/web/browser/palette.ts +64 -0
- package/src/web/browser/protocol.ts +13 -0
- package/src/web/browser/shell.ts +467 -0
- package/src/web/{src → browser}/status.ts +32 -36
- package/src/web/build.ts +51 -0
- package/src/web/public/app.html +1054 -0
- package/src/web/{sw.js → public/sw.js} +4 -3
- package/src/web/{vendor → public/vendor}/hljs-theme.css +2 -2
- package/src/web/{vendor → public/vendor}/hljs.js +2 -2
- package/THIRD_PARTY_LICENSES.md +0 -51
- package/docs/README.md +0 -94
- package/docs/archive/hlg-sampling.md +0 -350
- package/docs/archive/investigations/audio-input-plan.md +0 -356
- package/docs/archive/investigations/chunk-finetune-experiment.md +0 -84
- package/docs/archive/investigations/colibri-glm52-port.md +0 -1486
- package/docs/archive/investigations/colibri-oracle-pin.md +0 -1284
- package/docs/archive/investigations/curve-bisector-routing.md +0 -123
- package/docs/archive/investigations/decode-roofline-lookagain.md +0 -348
- package/docs/archive/investigations/diffusion-gemma-port.md +0 -602
- package/docs/archive/investigations/dspark-handoff.md +0 -224
- package/docs/archive/investigations/expert-offload-single-user-moe.md +0 -292
- package/docs/archive/investigations/faithful-l1-consolidation.md +0 -310
- package/docs/archive/investigations/hlg-sampling-investigation.md +0 -970
- package/docs/archive/investigations/indexshare-performance-spike.md +0 -357
- package/docs/archive/investigations/kernel-perf-review-2026-07.md +0 -171
- package/docs/archive/investigations/lab-build-journal.md +0 -157
- package/docs/archive/investigations/minicpm5-decode-megakernel.md +0 -489
- package/docs/archive/investigations/openjarvis-fit.md +0 -198
- package/docs/archive/investigations/openwiki-evaluation.md +0 -249
- package/docs/archive/investigations/optimization_plan.md +0 -346
- package/docs/archive/investigations/orpo-base-uf-experiment-and-directions.md +0 -120
- package/docs/archive/investigations/orpo-flash-cce-pin-leak.md +0 -167
- package/docs/archive/investigations/orpo-uf-testing-handoff.md +0 -126
- package/docs/archive/investigations/pareto-specialized-runtime-findings.md +0 -92
- package/docs/archive/investigations/pi-builtin-investigation.md +0 -308
- package/docs/archive/investigations/pre-colibri-stabilization.md +0 -329
- package/docs/archive/investigations/qwen38-27b-campaign-log.md +0 -1092
- package/docs/archive/investigations/segmented-backward-handoff.md +0 -148
- package/docs/archive/investigations/starter-model-port-handoff.md +0 -342
- package/docs/archive/investigations/steel-flash-cce-handoff.md +0 -489
- package/docs/archive/investigations/the-dreaming-handoff.md +0 -141
- package/docs/archive/investigations/trainer-validation-experiment.md +0 -173
- package/docs/archive/planning/chat-ui-vision-memory-first.md +0 -199
- package/docs/archive/planning/chat-ui-vision.md +0 -50
- package/docs/archive/planning/chunk-task-roadmap.md +0 -77
- package/docs/archive/planning/curve-sampler-research-plan.md +0 -331
- package/docs/archive/planning/journal.md +0 -316
- package/docs/archive/planning/memory-docs-and-dag-plan.md +0 -194
- package/docs/archive/planning/pi-integration-review.md +0 -178
- package/docs/archive/planning/release-notes-v0.0.10.md +0 -57
- package/docs/archive/planning/release-notes-v0.0.11.md +0 -239
- package/docs/archive/planning/release-notes-v0.0.12.md +0 -80
- package/docs/archive/planning/release-notes-v0.0.13.md +0 -77
- package/docs/archive/planning/release-notes-v0.0.9.md +0 -122
- package/docs/archive/planning/release-notes-v0.2.0.md +0 -49
- package/docs/archive/planning/release-notes-v0.3.0.md +0 -35
- package/docs/archive/planning/web-ui-pass-plan.md +0 -213
- package/docs/archive/planning/website-readme-pass-plan.md +0 -224
- package/docs/design/batching.md +0 -1841
- package/docs/design/decode-speed-program.md +0 -2349
- package/docs/design/dreaming-nightly-pipeline.md +0 -684
- package/docs/design/generic-model-support.md +0 -705
- package/docs/design/kv-cache.md +0 -824
- package/docs/design/orpo-training.md +0 -680
- package/docs/design/speculative-decoding.md +0 -1990
- package/docs/design/turboquant.md +0 -2212
- package/docs/design/unified-engine-frontier-plan.md +0 -1880
- package/docs/design/web-chat-redesign.md +0 -695
- package/docs/planning/PRODUCT_ROADMAP.md +0 -594
- package/docs/planning/ResearchTopics.md +0 -119
- package/docs/planning/release-notes-v0.4.0.md +0 -45
- package/docs/planning/video-series.md +0 -421
- package/docs/reference/benchmarks.md +0 -3613
- package/docs/reference/cli.md +0 -585
- package/docs/reference/distribution.md +0 -355
- package/docs/reference/environment.md +0 -398
- package/docs/reference/glossary.md +0 -39
- package/docs/reference/library-api.md +0 -674
- package/docs/reference/memory.md +0 -260
- package/docs/reference/models.md +0 -334
- package/docs/reference/server-api.md +0 -1775
- package/docs/reference/server-config.md +0 -846
- package/docs/reference/training.md +0 -462
- package/docs/reference/troubleshooting.md +0 -158
- package/src/admission-plan.ts +0 -18
- package/src/audio/conformer.ts +0 -647
- package/src/audio/decode.ts +0 -133
- package/src/audio/features.ts +0 -202
- package/src/audio/transcode.ts +0 -67
- package/src/backends/mlx/assistant-target.ts +0 -46
- package/src/backends/mlx/autoregressive-draft-rows.ts +0 -120
- package/src/backends/mlx/autoregressive.ts +0 -108
- package/src/backends/mlx/batch-group.ts +0 -1790
- package/src/backends/mlx/cache-attachment.ts +0 -67
- package/src/backends/mlx/cache-layout.ts +0 -31
- package/src/backends/mlx/checkpoint-state.ts +0 -53
- package/src/backends/mlx/continuation-persistence.ts +0 -106
- package/src/backends/mlx/continuation-request.ts +0 -43
- package/src/backends/mlx/continuation.ts +0 -77
- package/src/backends/mlx/diffusion.ts +0 -39
- package/src/backends/mlx/draft-target.ts +0 -32
- package/src/backends/mlx/evaluation.ts +0 -105
- package/src/backends/mlx/fill-append.ts +0 -45
- package/src/backends/mlx/full-kv-row-append.ts +0 -32
- package/src/backends/mlx/full-kv-row-donor.ts +0 -54
- package/src/backends/mlx/full-transitioning-kv-rows.ts +0 -65
- package/src/backends/mlx/gateway-binding.ts +0 -130
- package/src/backends/mlx/graph.ts +0 -38
- package/src/backends/mlx/kv-codec.ts +0 -9
- package/src/backends/mlx/kv-maintenance.ts +0 -133
- package/src/backends/mlx/kv-tensor-rows.ts +0 -212
- package/src/backends/mlx/memory-client.ts +0 -189
- package/src/backends/mlx/methods.ts +0 -90
- package/src/backends/mlx/model-prompt.ts +0 -164
- package/src/backends/mlx/model-serving.ts +0 -79
- package/src/backends/mlx/prefill-cohort.ts +0 -89
- package/src/backends/mlx/prefill-rows.ts +0 -200
- package/src/backends/mlx/prefill.ts +0 -33
- package/src/backends/mlx/projected-context-rows.ts +0 -72
- package/src/backends/mlx/request-policy.ts +0 -20
- package/src/backends/mlx/request-state-policy.ts +0 -41
- package/src/backends/mlx/rollback.ts +0 -89
- package/src/backends/mlx/row-sampling.ts +0 -49
- package/src/backends/mlx/serial-executor.ts +0 -242
- package/src/backends/mlx/speculative-group.ts +0 -359
- package/src/backends/mlx/speculative-round.ts +0 -161
- package/src/backends/mlx/speculative.ts +0 -114
- package/src/backends/mlx/state-rows.ts +0 -104
- package/src/backends/mlx/state-views.ts +0 -46
- package/src/backends/mlx/target-layout-capability.ts +0 -35
- package/src/backends/mlx/text-engine.ts +0 -73
- package/src/backends/mlx/transitioning-kv-rows.ts +0 -165
- package/src/chat-template.ts +0 -382
- package/src/cli.ts +0 -3168
- package/src/config.ts +0 -399
- package/src/contracts/completion.ts +0 -16
- package/src/contracts/execution.ts +0 -60
- package/src/contracts/generation.ts +0 -103
- package/src/contracts/host.ts +0 -7
- package/src/contracts/pi-provider.ts +0 -30
- package/src/contracts/pi.ts +0 -259
- package/src/contracts/prefix-cache.ts +0 -23
- package/src/contracts/resources.ts +0 -25
- package/src/contracts/sampling.ts +0 -23
- package/src/contracts/scheduling.ts +0 -45
- package/src/contracts/tsconfig.json +0 -14
- package/src/dataset/generators.ts +0 -1087
- package/src/dataset/index.ts +0 -12
- package/src/dataset/job.ts +0 -81
- package/src/dataset/llm.ts +0 -71
- package/src/dataset/registry.ts +0 -528
- package/src/diffusion/diffusion-generate.ts +0 -507
- package/src/download.ts +0 -505
- package/src/embed.ts +0 -77
- package/src/engine/admission.ts +0 -59
- package/src/engine/cancellation.ts +0 -38
- package/src/engine/engine.ts +0 -56
- package/src/engine/execution-coordinator.ts +0 -70
- package/src/engine/execution-plan.ts +0 -53
- package/src/engine/resources.ts +0 -53
- package/src/engine/scheduler.ts +0 -72
- package/src/engine/session.ts +0 -247
- package/src/eval/README.md +0 -130
- package/src/eval/ifeval.ts +0 -81
- package/src/eval/kl-prompts.ts +0 -29
- package/src/eval/kl.ts +0 -590
- package/src/eval/perplexity.ts +0 -164
- package/src/eval/quality-db.ts +0 -92
- package/src/eval/reference-logits.ts +0 -207
- package/src/eval/runner.ts +0 -124
- package/src/eval/score.ts +0 -49
- package/src/eval/tasks/bfcl.ts +0 -535
- package/src/eval/tasks/gsm8k.ts +0 -98
- package/src/eval/tasks/hashhop.ts +0 -139
- package/src/eval/tasks/humaneval.ts +0 -275
- package/src/eval/tasks/ifeval.ts +0 -525
- package/src/eval/tasks/mmlu.ts +0 -157
- package/src/eval/ultrafeedback.ts +0 -60
- package/src/evaldb.ts +0 -199
- package/src/expert-io.ts +0 -393
- package/src/expert-offload-build.ts +0 -77
- package/src/expert-offload.ts +0 -76
- package/src/expert-residency.ts +0 -1134
- package/src/expert-usage.ts +0 -415
- package/src/fill/echo-index.ts +0 -188
- package/src/fill/fill-session.ts +0 -452
- package/src/fill/proposal.ts +0 -68
- package/src/fill/schema-rows.ts +0 -506
- package/src/fill/tool-boundary.ts +0 -103
- package/src/fit.ts +0 -249
- package/src/generate.ts +0 -1346
- package/src/grammar.ts +0 -591
- package/src/harness-pi.ts +0 -184
- package/src/hf-push.ts +0 -628
- package/src/hub-rest.ts +0 -242
- package/src/index.ts +0 -63
- package/src/inference/denoising.ts +0 -17
- package/src/inference/draft-acceptance.ts +0 -41
- package/src/inference/draft-output.ts +0 -49
- package/src/inference/graph.ts +0 -33
- package/src/inference/prefill.ts +0 -39
- package/src/inference/rollback.ts +0 -7
- package/src/jobs/index.ts +0 -31
- package/src/jobs/job-entry.ts +0 -107
- package/src/jobs/task-client.ts +0 -19
- package/src/jobs/types.ts +0 -2
- package/src/kv-scheme.ts +0 -257
- package/src/kv-store.ts +0 -1099
- package/src/lab/README.md +0 -18
- package/src/lab/curve/curve-designer.html +0 -303
- package/src/lab/curve/curve-sampler.ts +0 -122
- package/src/lab/expert-trace/expert-trace.ts +0 -69
- package/src/lab/paged-kv/paged-kv-rows.ts +0 -83
- package/src/lab/paged-kv/paged-kv.ts +0 -270
- package/src/library.ts +0 -40
- package/src/lora.ts +0 -408
- package/src/md-modules.d.ts +0 -49
- package/src/media-fetch.ts +0 -442
- package/src/memory/article.ts +0 -558
- package/src/memory/chunk-validation.ts +0 -150
- package/src/memory/chunk.ts +0 -318
- package/src/memory/cluster.ts +0 -351
- package/src/memory/coverage.ts +0 -134
- package/src/memory/crosslink.ts +0 -0
- package/src/memory/db.ts +0 -253
- package/src/memory/entity.ts +0 -203
- package/src/memory/gate.ts +0 -187
- package/src/memory/history.ts +0 -152
- package/src/memory/ingest.ts +0 -335
- package/src/memory/model.ts +0 -183
- package/src/memory/normalize.ts +0 -671
- package/src/memory/parse.ts +0 -22
- package/src/memory/pipeline.ts +0 -275
- package/src/memory/prompts.ts +0 -65
- package/src/memory/query.ts +0 -334
- package/src/memory/reconcile.ts +0 -734
- package/src/memory/redirect.ts +0 -108
- package/src/memory/reindex.ts +0 -325
- package/src/memory/resolve.ts +0 -307
- package/src/memory/rest.ts +0 -245
- package/src/memory/route.ts +0 -272
- package/src/memory/schedule.ts +0 -165
- package/src/memory/stages.ts +0 -777
- package/src/memory/synthesize.ts +0 -1638
- package/src/memory/tools.ts +0 -749
- package/src/memory/vault.ts +0 -1194
- package/src/memory/wikify.ts +0 -1064
- package/src/mlx/array.ts +0 -388
- package/src/mlx/autograd.ts +0 -304
- package/src/mlx/checkpoint.ts +0 -149
- package/src/mlx/compile.ts +0 -168
- package/src/mlx/custom-vjp.ts +0 -188
- package/src/mlx/ffi.ts +0 -512
- package/src/mlx/materialize.ts +0 -23
- package/src/mlx/metal-kernel.ts +0 -159
- package/src/mlx/ops.ts +0 -1021
- package/src/mlx/quantized-rows.ts +0 -36
- package/src/mlx/token-bitmask.ts +0 -52
- package/src/mlx/turboquant-kv-decode.ts +0 -82
- package/src/mlx/turboquant-kv-inverse.ts +0 -122
- package/src/mlx/turboquant-ops.ts +0 -861
- package/src/mlx/turboquant-tables.ts +0 -237
- package/src/mmap.ts +0 -90
- package/src/model/artifact-identity.ts +0 -73
- package/src/model/batched-kv.ts +0 -235
- package/src/model/batched-mask.ts +0 -279
- package/src/model/batched-quant.ts +0 -212
- package/src/model/batched-quantized-kv.ts +0 -144
- package/src/model/batched-rotating-quant.ts +0 -353
- package/src/model/batched-rotating-state.ts +0 -149
- package/src/model/batched-rotating.ts +0 -338
- package/src/model/batched-row-storage.ts +0 -179
- package/src/model/batched-ssm.ts +0 -52
- package/src/model/batched-turboquant-kv.ts +0 -93
- package/src/model/compiled-decode.ts +0 -712
- package/src/model/decoded-kv-donor.ts +0 -12
- package/src/model/delayed-quantized-kv.ts +0 -72
- package/src/model/delayed-rotating-quantized-kv.ts +0 -93
- package/src/model/delayed-turboquant-kv.ts +0 -55
- package/src/model/diffusion-gemma.ts +0 -1026
- package/src/model/factory.ts +0 -189
- package/src/model/fingerprint.ts +0 -33
- package/src/model/flash-attention.ts +0 -694
- package/src/model/full-prefill-padding.ts +0 -56
- package/src/model/full-prefill-row.ts +0 -105
- package/src/model/gemma4-base.ts +0 -2351
- package/src/model/gemma4.ts +0 -1315
- package/src/model/generated/gemma4-12b.ts +0 -429
- package/src/model/generated/gemma4-26b.ts +0 -448
- package/src/model/generated/gemma4-e4b.ts +0 -717
- package/src/model/generated/index.ts +0 -18
- package/src/model/glm52-atlas.ts +0 -511
- package/src/model/glm52-cache.ts +0 -792
- package/src/model/glm52-config.ts +0 -262
- package/src/model/glm52-container.ts +0 -305
- package/src/model/glm52-coupling.ts +0 -388
- package/src/model/glm52-dsa.ts +0 -606
- package/src/model/glm52-expert-layout.ts +0 -253
- package/src/model/glm52-layout.ts +0 -152
- package/src/model/glm52-memory.ts +0 -500
- package/src/model/glm52-mla.ts +0 -792
- package/src/model/glm52-moe.ts +0 -223
- package/src/model/glm52-pilot.ts +0 -537
- package/src/model/glm52-reference.ts +0 -445
- package/src/model/glm52-residency.ts +0 -557
- package/src/model/glm52-resident-weights.ts +0 -135
- package/src/model/glm52-streamed-experts.ts +0 -453
- package/src/model/glm52-streamed-metal.ts +0 -610
- package/src/model/glm52-weights.ts +0 -503
- package/src/model/glm52.ts +0 -1132
- package/src/model/implementation.ts +0 -69
- package/src/model/kv-attention-view.ts +0 -67
- package/src/model/minicpm5.ts +0 -376
- package/src/model/profile.ts +0 -443
- package/src/model/qwen-append.ts +0 -12
- package/src/model/qwen-conv.ts +0 -52
- package/src/model/qwen3-delta-state.ts +0 -74
- package/src/model/qwen3-delta.ts +0 -451
- package/src/model/qwen3-moe.ts +0 -404
- package/src/model/qwen3-mrope.ts +0 -252
- package/src/model/qwen3.ts +0 -277
- package/src/model/qwen3_5-checkpoint.ts +0 -164
- package/src/model/qwen3_5.ts +0 -819
- package/src/model/rotating-kv-donor.ts +0 -35
- package/src/model/rotating-kv-layout.ts +0 -262
- package/src/model/rotating-row-transaction.ts +0 -35
- package/src/model/speculative-rotating-kv.ts +0 -101
- package/src/model/ssm-prefill-padding.ts +0 -39
- package/src/model/support.ts +0 -102
- package/src/model/trellis-balanced-scatter.ts +0 -51
- package/src/model/trellis-linear.ts +0 -672
- package/src/model/trellis-shared-m.ts +0 -151
- package/src/model/trellis-shared-scatter.ts +0 -144
- package/src/model/trellis-splitk-prefill.ts +0 -158
- package/src/model/trellis-tiled-prefill.ts +0 -136
- package/src/model/trellis-vector-expand.ts +0 -62
- package/src/model/trellis-wide-prefill.ts +0 -94
- package/src/model/turboquant-codec.ts +0 -83
- package/src/model/universal/archs.ts +0 -332
- package/src/model/universal/dense.ts +0 -612
- package/src/model/universal/modules.ts +0 -170
- package/src/model/universal/rope.ts +0 -381
- package/src/native/expert_io.c +0 -523
- package/src/native/frame_extract.swift +0 -140
- package/src/native-pack.ts +0 -162
- package/src/pi-provider.ts +0 -53
- package/src/pi-session.ts +0 -81
- package/src/pi-terminal.ts +0 -274
- package/src/pi-web.ts +0 -1932
- package/src/preflight.ts +0 -158
- package/src/prompt-cache.ts +0 -460
- package/src/quantize/allocator.ts +0 -627
- package/src/quantize/atomic-output.ts +0 -38
- package/src/quantize/calibration.ts +0 -181
- package/src/quantize/config-writer.ts +0 -145
- package/src/quantize/index.ts +0 -84
- package/src/quantize/job.ts +0 -162
- package/src/quantize/quantizer.ts +0 -666
- package/src/quantize/rotate.ts +0 -603
- package/src/quantize/safetensors-writer.ts +0 -249
- package/src/quantize/sensitivity.ts +0 -216
- package/src/quantize/trellis.ts +0 -586
- package/src/quantize/weight-transform.ts +0 -292
- package/src/registry.ts +0 -602
- package/src/runtime-config.ts +0 -81
- package/src/safetensors.ts +0 -158
- package/src/sampler.ts +0 -1011
- package/src/serve/admin-routes.ts +0 -184
- package/src/serve/assets/curve-terrain.html +0 -105
- package/src/serve/assets/training-inference-map.html +0 -575
- package/src/serve/aux-routes.ts +0 -132
- package/src/serve/batch-scheduler.ts +0 -8
- package/src/serve/checkpoint-identity.ts +0 -64
- package/src/serve/generation-gateway.ts +0 -496
- package/src/serve/isolate.ts +0 -609
- package/src/serve/kv-budget.ts +0 -14
- package/src/serve/lab-routes.ts +0 -341
- package/src/serve/lane-registry.ts +0 -45
- package/src/serve/model-admin-routes.ts +0 -155
- package/src/serve/model-binding.ts +0 -43
- package/src/serve/model-host.ts +0 -511
- package/src/serve/parent-application.ts +0 -56
- package/src/serve/prompt-response-trace.ts +0 -147
- package/src/serve/response-lease.ts +0 -24
- package/src/serve/session-search.ts +0 -326
- package/src/serve/static-routes.ts +0 -101
- package/src/server.ts +0 -1314
- package/src/spec/assistant-rows.ts +0 -92
- package/src/spec/assistant-source.ts +0 -126
- package/src/spec/cache-identity.ts +0 -22
- package/src/spec/deepspec-rows.ts +0 -29
- package/src/spec/deepspec-source.ts +0 -145
- package/src/spec/dflash-source.ts +0 -149
- package/src/spec/drafter.ts +0 -304
- package/src/spec/dspark/ab-stats.ts +0 -136
- package/src/spec/dspark/calibration.ts +0 -112
- package/src/spec/dspark/data-dflash.ts +0 -156
- package/src/spec/dspark/data.ts +0 -201
- package/src/spec/dspark/deepspec-module.ts +0 -776
- package/src/spec/dspark/generate-dflash.ts +0 -245
- package/src/spec/dspark/generate.ts +0 -290
- package/src/spec/dspark/loader.ts +0 -23
- package/src/spec/dspark/loss.ts +0 -213
- package/src/spec/dspark/module-dflash.ts +0 -560
- package/src/spec/dspark/module.ts +0 -540
- package/src/spec/dspark/quantize-drafter.ts +0 -91
- package/src/spec/dspark/sample.ts +0 -112
- package/src/spec/generate.ts +0 -276
- package/src/spec/glm52-mtp-graph.ts +0 -29
- package/src/spec/glm52-mtp-rows.ts +0 -132
- package/src/spec/glm52-mtp-source.ts +0 -317
- package/src/spec/glm52-mtp-state.ts +0 -33
- package/src/spec/ngram-source.ts +0 -204
- package/src/spec/projected-draft-rows.ts +0 -73
- package/src/spec/qwen-mtp-module.ts +0 -171
- package/src/spec/qwen-mtp-rows.ts +0 -219
- package/src/spec/qwen-mtp-source.ts +0 -383
- package/src/spec/qwen-mtp-state.ts +0 -30
- package/src/spec/serve-loop.ts +0 -585
- package/src/spec/source.ts +0 -215
- package/src/spec/two-model-rows.ts +0 -61
- package/src/spec/two-model.ts +0 -113
- package/src/ssd-cache.ts +0 -372
- package/src/ssd-durability.ts +0 -211
- package/src/tokenizer.ts +0 -335
- package/src/tool-approvals.ts +0 -131
- package/src/tool-call.ts +0 -695
- package/src/train/dataset.ts +0 -484
- package/src/train/diffusion-lora.ts +0 -216
- package/src/train/export.ts +0 -32
- package/src/train/flash-cce.ts +0 -1234
- package/src/train/forward.ts +0 -266
- package/src/train/fuse.ts +0 -243
- package/src/train/index.ts +0 -46
- package/src/train/job.ts +0 -205
- package/src/train/lora-params.ts +0 -254
- package/src/train/loss.ts +0 -1591
- package/src/train/merge.ts +0 -195
- package/src/train/optimizer.ts +0 -183
- package/src/train/prefix-shared.ts +0 -490
- package/src/train/rank.ts +0 -137
- package/src/train/segmented.ts +0 -1828
- package/src/train/steel-qmm-header.ts +0 -1461
- package/src/train/trainer.ts +0 -1221
- package/src/train/watch.ts +0 -542
- package/src/vision/diffusion-vision.ts +0 -620
- package/src/vision/embedder.ts +0 -168
- package/src/vision/preprocess.ts +0 -247
- package/src/vision/prompt.ts +0 -373
- package/src/vision/qwen3vl-preprocess.ts +0 -409
- package/src/vision/qwen3vl-prompt.ts +0 -196
- package/src/vision/qwen3vl-tower.ts +0 -555
- package/src/vision/siglip.ts +0 -599
- package/src/vision/video-frames.ts +0 -148
- package/src/web/app.html +0 -2143
- package/src/web/app.js +0 -5703
- package/src/web/skills/memory/SKILL.md +0 -120
- package/src/web/skills/web-research/SKILL.md +0 -46
- package/src/web/skills.ts +0 -97
- package/src/web/src/adapters-panel.ts +0 -288
- package/src/web/src/assistant.ts +0 -451
- package/src/web/src/chat.ts +0 -1057
- package/src/web/src/composer.ts +0 -959
- package/src/web/src/dataset.ts +0 -132
- package/src/web/src/finetune.ts +0 -267
- package/src/web/src/hub.ts +0 -428
- package/src/web/src/main.ts +0 -98
- package/src/web/src/markdown.ts +0 -0
- package/src/web/src/memory-panel.ts +0 -572
- package/src/web/src/model-picker.ts +0 -159
- package/src/web/src/palette.ts +0 -346
- package/src/web/src/protocol.ts +0 -26
- package/src/web/src/quantize.ts +0 -167
- package/src/web/src/rag.ts +0 -274
- package/src/web/src/sessions.ts +0 -469
- package/src/web/src/shell.ts +0 -804
- package/src/web/src/ui-catalog.ts +0 -78
- package/src/web/vendor/README.md +0 -61
- package/src/web-assets.ts +0 -47
- package/src/web-tools.ts +0 -478
- package/src/weights.ts +0 -255
- /package/src/{serve → server}/completion-sink.ts +0 -0
- /package/src/web/{icon.svg → public/icon.svg} +0 -0
- /package/src/web/{manifest.webmanifest → public/manifest.webmanifest} +0 -0
- /package/src/web/{vendor → public/vendor}/hljs-LICENSE +0 -0
package/src/kv-store.ts
DELETED
|
@@ -1,1099 +0,0 @@
|
|
|
1
|
-
import { type CheckpointAttachment, attachmentBytes, disposeAttachments } from "./backends/mlx/checkpoint-state";
|
|
2
|
-
import { cleanupFailure, disposeResources } from "./engine/resources";
|
|
3
|
-
// KV-cache persistence: save prompt caches to disk, reload by streamed copy.
|
|
4
|
-
// The serialization core of the SSD cold tier (docs/design/kv-cache.md).
|
|
5
|
-
//
|
|
6
|
-
// File layout (every tensor PAGE-ALIGNED — the Phase 1 corollary: files
|
|
7
|
-
// we write can be mmap'd and handed to the GPU without copies):
|
|
8
|
-
// [magic "MLXBUNKV2\n"][u32 LE header length][u32 LE dataStart]
|
|
9
|
-
// [u64 LE header hash][JSON header][padding]
|
|
10
|
-
// [tensor data at 16 KiB-aligned offsets]
|
|
11
|
-
// Header: { formatVersion, modelId, configFingerprint, ns, tokenizerHash,
|
|
12
|
-
// createdAt, tokens, caches: [{ kind, offset, idx?, maxSize?,
|
|
13
|
-
// groupSize?, bits?, kBits?, vBits?, headDim?,
|
|
14
|
-
// tensors: [{ off, bytes, shape, dtype, hash }] }] }
|
|
15
|
-
//
|
|
16
|
-
// v2 over v1: quantized cache kinds (kv_config quantization — v1 could not
|
|
17
|
-
// persist it; NOTE bf16 is the serving default since the 2026-07-05 L1
|
|
18
|
-
// decision, quantized KV is opt-in), SSM kind (Qwen3.5 hybrid),
|
|
19
|
-
// invalidation metadata (configFingerprint covers the kv-quant scheme;
|
|
20
|
-
// tokenizerHash guards vocab drift; ns = adapter spec), per-tensor hashes
|
|
21
|
-
// (verified opt-in — the hash pass roughly doubles restore reads),
|
|
22
|
-
// a header hash (always verified — cheap), and atomic tmp+fsync+rename
|
|
23
|
-
// writes. v1 files are not migrated: nothing shipped wrote them (the
|
|
24
|
-
// format was test/experiment-only) — they read as "not a v2 file" and the
|
|
25
|
-
// SSD tier deletes-and-regenerates.
|
|
26
|
-
//
|
|
27
|
-
// v3 over v2: fixed-width (zero-padded 16-hex) tensor hashes, which let the
|
|
28
|
-
// writer STREAM — header sized up front, tensors materialized/hashed/written
|
|
29
|
-
// one at a time, header patched in last (peak host transient = one tensor,
|
|
30
|
-
// not the whole entry). Same magic; v2 files read as unsupported and the
|
|
31
|
-
// SSD tier deletes-and-regenerates (machine-local, cheap).
|
|
32
|
-
//
|
|
33
|
-
// Reload (2026-07-07, A7-restore) is a STREAMED COPY — the read-side twin
|
|
34
|
-
// of the v3 streaming writer: each tensor is copied out of a read-only mmap
|
|
35
|
-
// into an mlx-owned leaf (MlxArray.fromBytesCopy), its clean file pages
|
|
36
|
-
// dropped (MADV_DONTNEED) right after, and the mapping unmapped before
|
|
37
|
-
// loadKvCache returns. Peak host transient = the live entry + one tensor,
|
|
38
|
-
// and NOTHING outlives the call: the old zero-copy wrap (COW mmap +
|
|
39
|
-
// fromPointer) required pinning every restore mapping for the PROCESS
|
|
40
|
-
// lifetime once the FFI-dtor fix removed the unmap signal (2026-07-06) —
|
|
41
|
-
// one leaked full-entry mapping per restore — and the exactly-offset-sized
|
|
42
|
-
// restored buffers made the first post-restore decode step concat-copy the
|
|
43
|
-
// whole entry again.
|
|
44
|
-
|
|
45
|
-
import { openSync, writeSync, readSync, closeSync, fsyncSync, renameSync, rmSync } from "node:fs";
|
|
46
|
-
import { MmapFile, MADV_DONTNEED } from "./mmap";
|
|
47
|
-
import { MlxArray } from "./mlx/array";
|
|
48
|
-
import type { Dtype } from "./mlx/ffi";
|
|
49
|
-
import * as ops from "./mlx/ops";
|
|
50
|
-
import {
|
|
51
|
-
type Cache, cacheSignature, KVCache, RotatingKVCache,
|
|
52
|
-
QuantizedKVCache, RotatingQuantizedKVCache,
|
|
53
|
-
TurboQuantKVCache, type TurboQuantTensor,
|
|
54
|
-
} from "./model/gemma4-base";
|
|
55
|
-
import { Glm52Cache, MLACache } from "./model/glm52-cache";
|
|
56
|
-
import { SSMCache } from "./model/qwen3-delta";
|
|
57
|
-
|
|
58
|
-
const MAGIC = "MLXBUNKV2\n";
|
|
59
|
-
const ALIGN = 16384;
|
|
60
|
-
/** magic + u32 headerLen + u32 dataStart + u64 headerHash */
|
|
61
|
-
const PREFIX_LEN = MAGIC.length + 4 + 4 + 8;
|
|
62
|
-
|
|
63
|
-
// Fixed-width (16 hex chars): the v3 streaming writer sizes the header
|
|
64
|
-
// BEFORE materializing any tensor bytes, so hash strings must not change
|
|
65
|
-
// the header's byte length when the real values replace the placeholders.
|
|
66
|
-
const hash64 = (bytes: Uint8Array): string => Bun.hash(bytes).toString(16).padStart(16, "0");
|
|
67
|
-
|
|
68
|
-
export type CacheKind =
|
|
69
|
-
| "kv"
|
|
70
|
-
| "rotating"
|
|
71
|
-
| "qkv"
|
|
72
|
-
| "rotating-qkv"
|
|
73
|
-
| "ssm"
|
|
74
|
-
| "turboquant"
|
|
75
|
-
| "mla"
|
|
76
|
-
| "mla-dsa"
|
|
77
|
-
| "mtp-mla";
|
|
78
|
-
|
|
79
|
-
interface TensorSlot {
|
|
80
|
-
off: number;
|
|
81
|
-
bytes: number;
|
|
82
|
-
shape: number[];
|
|
83
|
-
dtype: number;
|
|
84
|
-
hash: string;
|
|
85
|
-
}
|
|
86
|
-
|
|
87
|
-
export interface CacheHeaderEntry {
|
|
88
|
-
kind: CacheKind;
|
|
89
|
-
offset: number;
|
|
90
|
-
minimumReusableOffset?: number;
|
|
91
|
-
/** rotating variants: ring write index */
|
|
92
|
-
idx?: number;
|
|
93
|
-
/** rotating variants: window size */
|
|
94
|
-
maxSize?: number;
|
|
95
|
-
/** quantized variants (qkv/rotating-qkv: mlx affine scheme) */
|
|
96
|
-
groupSize?: number;
|
|
97
|
-
bits?: number;
|
|
98
|
-
/** turboquant: per-side bit widths (reuses the `bits` field's slot class
|
|
99
|
-
* but needs both — asymmetric key/value bit widths, unlike qkv's single
|
|
100
|
-
* `bits`). `groupSize` is unused for turboquant (fixed at 32, the
|
|
101
|
-
* BLOCK_SIZE constant — not a configurable field like qkv's). */
|
|
102
|
-
kBits?: number;
|
|
103
|
-
vBits?: number;
|
|
104
|
-
/** turboquant: head_dim, needed to unpack kIdx/vPacked on restore
|
|
105
|
-
* (packed byte width alone doesn't recover the original element count). */
|
|
106
|
-
headDim?: number;
|
|
107
|
-
/** GLM checkpoint-native compressed-cache geometry. */
|
|
108
|
-
kvLoraRank?: number;
|
|
109
|
-
ropeHeadDim?: number;
|
|
110
|
-
dsaHeadDim?: number;
|
|
111
|
-
maxTokens?: number;
|
|
112
|
-
/** kv/rotating: [k, v] · qkv/rotating-qkv: [kPacked, kScales, kBiases,
|
|
113
|
-
* vPacked, vScales, vBiases] · ssm: [conv, recurrent] · turboquant:
|
|
114
|
-
* [kIdx, kScales, kZeros, vPacked, vScales] */
|
|
115
|
-
tensors: TensorSlot[];
|
|
116
|
-
}
|
|
117
|
-
|
|
118
|
-
export interface KvSaveMeta {
|
|
119
|
-
attachments?: CheckpointAttachment[];
|
|
120
|
-
modelId?: string;
|
|
121
|
-
/** configFingerprint(config) — covers every graph-shaping field incl.
|
|
122
|
-
* the kv-quant scheme, so a scheme flip invalidates naturally. */
|
|
123
|
-
configFingerprint?: string;
|
|
124
|
-
/** Adapter namespace (PromptCache ns — adapters joined with "+"). */
|
|
125
|
-
ns?: string;
|
|
126
|
-
/** sha256 of tokenizer.json: identical ids must mean identical text. */
|
|
127
|
-
tokenizerHash?: string;
|
|
128
|
-
/** In-flight generation state. The serialized caches cover `tokens`;
|
|
129
|
-
* pendingToken was sampled from that state but has not been emitted. */
|
|
130
|
-
generationCheckpoint?: {
|
|
131
|
-
key: string;
|
|
132
|
-
cacheNs: string;
|
|
133
|
-
originalPromptTokens: number;
|
|
134
|
-
generatedTokens: number;
|
|
135
|
-
pendingToken: number;
|
|
136
|
-
/** Actual sampler seed, including a server-generated default. Required to
|
|
137
|
-
* continue the random stream exactly after a restart. */
|
|
138
|
-
seed: number;
|
|
139
|
-
/** True when `seed` was part of the client request rather than generated
|
|
140
|
-
* by the server. */
|
|
141
|
-
seedWasExplicit: boolean;
|
|
142
|
-
};
|
|
143
|
-
}
|
|
144
|
-
|
|
145
|
-
export interface KvFileHeader extends Omit<KvSaveMeta, "attachments"> {
|
|
146
|
-
attachments?: Array<Omit<CheckpointAttachment, "tensors"> & { tensors: TensorSlot[] }>;
|
|
147
|
-
/** Absent in earlier v3 files, which used mlx-cache-v3. */
|
|
148
|
-
codecProvider?: string;
|
|
149
|
-
formatVersion: 3 | 4;
|
|
150
|
-
createdAt: number;
|
|
151
|
-
tokens: number[];
|
|
152
|
-
caches: CacheHeaderEntry[];
|
|
153
|
-
}
|
|
154
|
-
|
|
155
|
-
const alignUp = (n: number) => Math.ceil(n / ALIGN) * ALIGN;
|
|
156
|
-
|
|
157
|
-
/** Snapshot one cache into header entry + LAZY tensor sources (v3 streaming
|
|
158
|
-
* writer: sizes/shapes are metadata, bytes materialize ONE tensor at a
|
|
159
|
-
* time at write — the old all-blobs-up-front path doubled RSS by the whole
|
|
160
|
-
* entry, ~390 MB for a 16k cpm5 entry, on every write-behind). Quantized
|
|
161
|
-
* full caches are sliced to the live [.., :offset, :] region (token axis —
|
|
162
|
-
* quantization groups run along the FEATURE dim, so a sequence-axis cut
|
|
163
|
-
* never splits a group); rotating caches persist the whole ring
|
|
164
|
-
* as-laid-out with ringIdx. Slice handles are lazy graph nodes (no GPU
|
|
165
|
-
* materialization until rawBytes at write time). */
|
|
166
|
-
interface TensorSource { arr: MlxArray; disposeAfter: boolean }
|
|
167
|
-
|
|
168
|
-
export interface SnapshotContext {
|
|
169
|
-
slots: TensorSlot[];
|
|
170
|
-
push(a: MlxArray, disposeAfter: boolean): void;
|
|
171
|
-
liveSlice(a: MlxArray, upTo: number): MlxArray;
|
|
172
|
-
liveMlaSlice(a: MlxArray, upTo: number): MlxArray;
|
|
173
|
-
pushTriple(t: ops.QuantizedTensor, upTo: number | null): void;
|
|
174
|
-
}
|
|
175
|
-
|
|
176
|
-
export interface CloneContext {
|
|
177
|
-
view(a: MlxArray): MlxArray;
|
|
178
|
-
liveView(a: MlxArray, upTo: number): MlxArray;
|
|
179
|
-
mlaView(a: MlxArray, upTo: number): MlxArray;
|
|
180
|
-
tripleView(t: ops.QuantizedTensor, upTo: number | null): ops.QuantizedTensor;
|
|
181
|
-
}
|
|
182
|
-
|
|
183
|
-
export interface LoadContext {
|
|
184
|
-
path: string;
|
|
185
|
-
arr(slot: TensorSlot): MlxArray;
|
|
186
|
-
grownArr(slot: TensorSlot, offset: number): MlxArray;
|
|
187
|
-
triple(slots: TensorSlot[], at: number): ops.QuantizedTensor;
|
|
188
|
-
}
|
|
189
|
-
|
|
190
|
-
export interface CacheCodec {
|
|
191
|
-
matches(cache: Cache): boolean;
|
|
192
|
-
snapshot(cache: Cache, context: SnapshotContext): CacheHeaderEntry;
|
|
193
|
-
clone(cache: Cache, context: CloneContext): Cache;
|
|
194
|
-
load(entry: CacheHeaderEntry, context: LoadContext): Cache;
|
|
195
|
-
headerTrimmable(entry: CacheHeaderEntry): boolean;
|
|
196
|
-
}
|
|
197
|
-
|
|
198
|
-
const requireMlaState = (cache: Glm52Cache, operation: string): void => {
|
|
199
|
-
if (!cache.latent || !cache.rope || cache.batchSize === null || cache.offset === 0)
|
|
200
|
-
throw new Error(`cannot ${operation} an empty GLM compressed cache`);
|
|
201
|
-
if (cache.dsa && (!cache.dsa.data || cache.dsa.offset !== cache.offset))
|
|
202
|
-
throw new Error(`cannot ${operation} misaligned GLM DSA state`);
|
|
203
|
-
};
|
|
204
|
-
|
|
205
|
-
const snapshotMla = (
|
|
206
|
-
kind: Extract<CacheKind, "mla" | "mla-dsa" | "mtp-mla">,
|
|
207
|
-
cache: Cache,
|
|
208
|
-
context: SnapshotContext,
|
|
209
|
-
): CacheHeaderEntry => {
|
|
210
|
-
const c = cache as Glm52Cache;
|
|
211
|
-
requireMlaState(c, "persist");
|
|
212
|
-
context.push(context.liveMlaSlice(c.latent!, c.offset), true);
|
|
213
|
-
context.push(context.liveMlaSlice(c.rope!, c.offset), true);
|
|
214
|
-
if (c.dsa) context.push(context.liveMlaSlice(c.dsa.data!, c.offset), true);
|
|
215
|
-
return {
|
|
216
|
-
kind,
|
|
217
|
-
offset: c.offset,
|
|
218
|
-
kvLoraRank: c.kvLoraRank,
|
|
219
|
-
ropeHeadDim: c.ropeHeadDim,
|
|
220
|
-
...(c.dsa ? { dsaHeadDim: c.dsa.headDim } : {}),
|
|
221
|
-
maxTokens: c.maxTokens,
|
|
222
|
-
tensors: context.slots,
|
|
223
|
-
};
|
|
224
|
-
};
|
|
225
|
-
|
|
226
|
-
const cloneMla = (cache: Cache, context: CloneContext): Cache => {
|
|
227
|
-
const c = cache as Glm52Cache;
|
|
228
|
-
requireMlaState(c, "clone");
|
|
229
|
-
const clone = new MLACache({
|
|
230
|
-
kvLoraRank: c.kvLoraRank,
|
|
231
|
-
ropeHeadDim: c.ropeHeadDim,
|
|
232
|
-
...(c.dsa ? { dsa: { headDim: c.dsa.headDim } } : {}),
|
|
233
|
-
maxTokens: c.maxTokens,
|
|
234
|
-
role: c.role,
|
|
235
|
-
});
|
|
236
|
-
clone.restoreCompressedState(
|
|
237
|
-
context.mlaView(c.latent!, c.offset),
|
|
238
|
-
context.mlaView(c.rope!, c.offset),
|
|
239
|
-
c.dsa ? context.mlaView(c.dsa.data!, c.offset) : null,
|
|
240
|
-
c.offset,
|
|
241
|
-
);
|
|
242
|
-
return clone;
|
|
243
|
-
};
|
|
244
|
-
|
|
245
|
-
const loadMla = (entry: CacheHeaderEntry, context: LoadContext): Cache => {
|
|
246
|
-
const hasDsa = entry.kind === "mla-dsa";
|
|
247
|
-
if (entry.kind === "mtp-mla" && entry.dsaHeadDim !== undefined)
|
|
248
|
-
throw new Error(`${context.path}: native MTP cache cannot contain DSA state`);
|
|
249
|
-
const cache = new MLACache({
|
|
250
|
-
kvLoraRank: entry.kvLoraRank!,
|
|
251
|
-
ropeHeadDim: entry.ropeHeadDim!,
|
|
252
|
-
...(hasDsa ? { dsa: { headDim: entry.dsaHeadDim! } } : {}),
|
|
253
|
-
maxTokens: entry.maxTokens,
|
|
254
|
-
role: entry.kind === "mtp-mla" ? "mtp" : "target",
|
|
255
|
-
});
|
|
256
|
-
cache.restoreCompressedState(
|
|
257
|
-
context.arr(entry.tensors[0]!),
|
|
258
|
-
context.arr(entry.tensors[1]!),
|
|
259
|
-
hasDsa ? context.arr(entry.tensors[2]!) : null,
|
|
260
|
-
entry.offset,
|
|
261
|
-
);
|
|
262
|
-
return cache;
|
|
263
|
-
};
|
|
264
|
-
|
|
265
|
-
/** One authoritative persistence algebra. Adding a cache storage kind must
|
|
266
|
-
* define all four operations here; save/clone/load and SSD eligibility
|
|
267
|
-
* therefore cannot grow independent instanceof/switch chains. */
|
|
268
|
-
const CACHE_CODECS = {
|
|
269
|
-
kv: {
|
|
270
|
-
matches: (cache) => cacheSignature(cache) === "kv:plain",
|
|
271
|
-
snapshot: (cache, context) => {
|
|
272
|
-
const c = cache as KVCache;
|
|
273
|
-
if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
|
|
274
|
-
context.push(context.liveSlice(c.keys, c.offset), true);
|
|
275
|
-
context.push(context.liveSlice(c.values, c.offset), true);
|
|
276
|
-
return { kind: "kv", offset: c.offset, tensors: context.slots };
|
|
277
|
-
},
|
|
278
|
-
clone: (cache, context) => {
|
|
279
|
-
const c = cache as KVCache;
|
|
280
|
-
const clone = new KVCache();
|
|
281
|
-
if (!c.keys || !c.values) return clone;
|
|
282
|
-
clone.restoreState(
|
|
283
|
-
context.liveView(c.keys, c.offset),
|
|
284
|
-
context.liveView(c.values, c.offset),
|
|
285
|
-
c.offset,
|
|
286
|
-
);
|
|
287
|
-
return clone;
|
|
288
|
-
},
|
|
289
|
-
load: (entry, context) => {
|
|
290
|
-
const cache = new KVCache();
|
|
291
|
-
cache.restoreState(
|
|
292
|
-
context.grownArr(entry.tensors[0]!, entry.offset),
|
|
293
|
-
context.grownArr(entry.tensors[1]!, entry.offset),
|
|
294
|
-
entry.offset,
|
|
295
|
-
);
|
|
296
|
-
return cache;
|
|
297
|
-
},
|
|
298
|
-
headerTrimmable: () => true,
|
|
299
|
-
},
|
|
300
|
-
rotating: {
|
|
301
|
-
matches: (cache) => cacheSignature(cache) === "kv:rotating-plain",
|
|
302
|
-
snapshot: (cache, context) => {
|
|
303
|
-
const c = cache as RotatingKVCache;
|
|
304
|
-
if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
|
|
305
|
-
context.push(c.keys, false);
|
|
306
|
-
context.push(c.values, false);
|
|
307
|
-
return { kind: "rotating", offset: c.offset, idx: c.ringIdx,
|
|
308
|
-
maxSize: c.maxSize, tensors: context.slots };
|
|
309
|
-
},
|
|
310
|
-
clone: (cache, context) => {
|
|
311
|
-
const c = cache as RotatingKVCache;
|
|
312
|
-
const clone = new RotatingKVCache(c.maxSize);
|
|
313
|
-
if (!c.keys || !c.values) return clone;
|
|
314
|
-
clone.restoreState(context.view(c.keys), context.view(c.values), c.offset, c.ringIdx);
|
|
315
|
-
return clone;
|
|
316
|
-
},
|
|
317
|
-
load: (entry, context) => {
|
|
318
|
-
const cache = new RotatingKVCache(entry.maxSize!);
|
|
319
|
-
cache.restoreState(
|
|
320
|
-
context.arr(entry.tensors[0]!), context.arr(entry.tensors[1]!),
|
|
321
|
-
entry.offset, entry.idx!,
|
|
322
|
-
);
|
|
323
|
-
return cache;
|
|
324
|
-
},
|
|
325
|
-
headerTrimmable: (entry) => entry.offset < (entry.maxSize ?? 0),
|
|
326
|
-
},
|
|
327
|
-
qkv: {
|
|
328
|
-
matches: (cache) => cacheSignature(cache).startsWith("kv:quant:"),
|
|
329
|
-
snapshot: (cache, context) => {
|
|
330
|
-
const c = cache as QuantizedKVCache;
|
|
331
|
-
if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
|
|
332
|
-
context.pushTriple(c.keys, c.offset);
|
|
333
|
-
context.pushTriple(c.values, c.offset);
|
|
334
|
-
return { kind: "qkv", offset: c.offset, groupSize: c.groupSize,
|
|
335
|
-
bits: c.bits, tensors: context.slots };
|
|
336
|
-
},
|
|
337
|
-
clone: (cache, context) => {
|
|
338
|
-
const c = cache as QuantizedKVCache;
|
|
339
|
-
const clone = new QuantizedKVCache(c.groupSize, c.bits);
|
|
340
|
-
if (!c.keys || !c.values) return clone;
|
|
341
|
-
clone.restoreState(
|
|
342
|
-
context.tripleView(c.keys, c.offset),
|
|
343
|
-
context.tripleView(c.values, c.offset),
|
|
344
|
-
c.offset,
|
|
345
|
-
);
|
|
346
|
-
return clone;
|
|
347
|
-
},
|
|
348
|
-
load: (entry, context) => {
|
|
349
|
-
const cache = new QuantizedKVCache(entry.groupSize!, entry.bits!);
|
|
350
|
-
cache.restoreState(
|
|
351
|
-
context.triple(entry.tensors, 0), context.triple(entry.tensors, 3), entry.offset,
|
|
352
|
-
);
|
|
353
|
-
return cache;
|
|
354
|
-
},
|
|
355
|
-
headerTrimmable: () => true,
|
|
356
|
-
},
|
|
357
|
-
"rotating-qkv": {
|
|
358
|
-
matches: (cache) => cacheSignature(cache).startsWith("kv:rotating-quant:"),
|
|
359
|
-
snapshot: (cache, context) => {
|
|
360
|
-
const c = cache as RotatingQuantizedKVCache;
|
|
361
|
-
if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
|
|
362
|
-
context.pushTriple(c.keys, null);
|
|
363
|
-
context.pushTriple(c.values, null);
|
|
364
|
-
return { kind: "rotating-qkv", offset: c.offset, idx: c.ringIdx,
|
|
365
|
-
maxSize: c.maxSize, groupSize: c.groupSize, bits: c.bits,
|
|
366
|
-
tensors: context.slots };
|
|
367
|
-
},
|
|
368
|
-
clone: (cache, context) => {
|
|
369
|
-
const c = cache as RotatingQuantizedKVCache;
|
|
370
|
-
const clone = new RotatingQuantizedKVCache(c.maxSize, c.groupSize, c.bits);
|
|
371
|
-
if (!c.keys || !c.values) return clone;
|
|
372
|
-
clone.restoreState(
|
|
373
|
-
context.tripleView(c.keys, null), context.tripleView(c.values, null),
|
|
374
|
-
c.offset, c.ringIdx,
|
|
375
|
-
);
|
|
376
|
-
return clone;
|
|
377
|
-
},
|
|
378
|
-
load: (entry, context) => {
|
|
379
|
-
const cache = new RotatingQuantizedKVCache(
|
|
380
|
-
entry.maxSize!, entry.groupSize!, entry.bits!,
|
|
381
|
-
);
|
|
382
|
-
cache.restoreState(
|
|
383
|
-
context.triple(entry.tensors, 0), context.triple(entry.tensors, 3),
|
|
384
|
-
entry.offset, entry.idx!,
|
|
385
|
-
);
|
|
386
|
-
return cache;
|
|
387
|
-
},
|
|
388
|
-
headerTrimmable: (entry) => entry.offset < (entry.maxSize ?? 0),
|
|
389
|
-
},
|
|
390
|
-
ssm: {
|
|
391
|
-
matches: (cache) => cacheSignature(cache) === "ssm",
|
|
392
|
-
snapshot: (cache, context) => {
|
|
393
|
-
const c = cache as SSMCache;
|
|
394
|
-
if (!c.conv || !c.recurrent) throw new Error("cannot persist an empty cache");
|
|
395
|
-
context.push(c.conv, false);
|
|
396
|
-
context.push(c.recurrent, false);
|
|
397
|
-
return { kind: "ssm", offset: c.offset, tensors: context.slots };
|
|
398
|
-
},
|
|
399
|
-
clone: (cache, context) => {
|
|
400
|
-
const c = cache as SSMCache;
|
|
401
|
-
if (!c.conv || !c.recurrent) throw new Error("cannot clone an empty cache");
|
|
402
|
-
const clone = new SSMCache();
|
|
403
|
-
clone.conv = context.view(c.conv);
|
|
404
|
-
clone.recurrent = context.view(c.recurrent);
|
|
405
|
-
clone.offset = c.offset;
|
|
406
|
-
return clone;
|
|
407
|
-
},
|
|
408
|
-
load: (entry, context) => {
|
|
409
|
-
const cache = new SSMCache();
|
|
410
|
-
cache.conv = context.arr(entry.tensors[0]!);
|
|
411
|
-
cache.recurrent = context.arr(entry.tensors[1]!);
|
|
412
|
-
cache.offset = entry.offset;
|
|
413
|
-
return cache;
|
|
414
|
-
},
|
|
415
|
-
headerTrimmable: () => false,
|
|
416
|
-
},
|
|
417
|
-
turboquant: {
|
|
418
|
-
matches: (cache) => cacheSignature(cache).startsWith("kv:turboquant:"),
|
|
419
|
-
snapshot: (cache, context) => {
|
|
420
|
-
const c = cache as TurboQuantKVCache;
|
|
421
|
-
const state = c.state();
|
|
422
|
-
if (state.length === 0) throw new Error("cannot persist an empty cache");
|
|
423
|
-
for (const array of state) context.push(array, true);
|
|
424
|
-
return { kind: "turboquant", offset: c.offset, kBits: c.kBits,
|
|
425
|
-
vBits: c.vBits, headDim: c.headDim ?? 0, tensors: context.slots };
|
|
426
|
-
},
|
|
427
|
-
clone: (cache, context) => {
|
|
428
|
-
const c = cache as TurboQuantKVCache;
|
|
429
|
-
const state = c.state();
|
|
430
|
-
if (state.length === 0) throw new Error("cannot clone an empty cache");
|
|
431
|
-
try {
|
|
432
|
-
const [kIdx, kScales, kZeros, vPacked, vScales] = state;
|
|
433
|
-
const clone = new TurboQuantKVCache(c.kBits, c.vBits);
|
|
434
|
-
clone.restoreState({
|
|
435
|
-
kIdx: context.view(kIdx!),
|
|
436
|
-
kScales: context.view(kScales!),
|
|
437
|
-
kZeros: context.view(kZeros!),
|
|
438
|
-
vPacked: context.view(vPacked!),
|
|
439
|
-
vScales: context.view(vScales!),
|
|
440
|
-
}, c.offset, c.headDim!);
|
|
441
|
-
return clone;
|
|
442
|
-
} finally {
|
|
443
|
-
for (const array of state) array.dispose();
|
|
444
|
-
}
|
|
445
|
-
},
|
|
446
|
-
load: (entry, context) => {
|
|
447
|
-
const cache = new TurboQuantKVCache(entry.kBits!, entry.vBits!);
|
|
448
|
-
const tensors: TurboQuantTensor = {
|
|
449
|
-
kIdx: context.arr(entry.tensors[0]!),
|
|
450
|
-
kScales: context.arr(entry.tensors[1]!),
|
|
451
|
-
kZeros: context.arr(entry.tensors[2]!),
|
|
452
|
-
vPacked: context.arr(entry.tensors[3]!),
|
|
453
|
-
vScales: context.arr(entry.tensors[4]!),
|
|
454
|
-
};
|
|
455
|
-
cache.restoreState(tensors, entry.offset, entry.headDim!);
|
|
456
|
-
return cache;
|
|
457
|
-
},
|
|
458
|
-
headerTrimmable: () => true,
|
|
459
|
-
},
|
|
460
|
-
mla: {
|
|
461
|
-
matches: (cache) => cacheSignature(cache) === "kv:mla:target",
|
|
462
|
-
snapshot: (cache, context) => snapshotMla("mla", cache, context),
|
|
463
|
-
clone: cloneMla,
|
|
464
|
-
load: loadMla,
|
|
465
|
-
headerTrimmable: () => true,
|
|
466
|
-
},
|
|
467
|
-
"mla-dsa": {
|
|
468
|
-
matches: (cache) => cacheSignature(cache) === "kv:mla:target:dsa",
|
|
469
|
-
snapshot: (cache, context) => snapshotMla("mla-dsa", cache, context),
|
|
470
|
-
clone: cloneMla,
|
|
471
|
-
load: loadMla,
|
|
472
|
-
headerTrimmable: () => true,
|
|
473
|
-
},
|
|
474
|
-
"mtp-mla": {
|
|
475
|
-
matches: (cache) => cacheSignature(cache) === "kv:mla:mtp",
|
|
476
|
-
snapshot: (cache, context) => snapshotMla("mtp-mla", cache, context),
|
|
477
|
-
clone: cloneMla,
|
|
478
|
-
load: loadMla,
|
|
479
|
-
headerTrimmable: () => true,
|
|
480
|
-
},
|
|
481
|
-
} satisfies Record<CacheKind, CacheCodec>;
|
|
482
|
-
|
|
483
|
-
/** Codecs belong to the loaded backend binding. The persisted ID is checked
|
|
484
|
-
* before allocating state; codecs are engine code, never loaded from a file. */
|
|
485
|
-
export interface CacheCodecProvider {
|
|
486
|
-
readonly id: string;
|
|
487
|
-
forCache(cache: Cache): CacheCodec;
|
|
488
|
-
forHeader(entry: CacheHeaderEntry): CacheCodec;
|
|
489
|
-
}
|
|
490
|
-
|
|
491
|
-
export function createCacheCodecProvider(
|
|
492
|
-
id: string, entries: Readonly<Partial<Record<CacheKind, CacheCodec>>>,
|
|
493
|
-
): CacheCodecProvider {
|
|
494
|
-
if (!id.trim()) throw new Error("cache codec provider needs an identity");
|
|
495
|
-
const table = new Map(Object.entries(entries).map(([kind, codec]) => [kind, Object.freeze({
|
|
496
|
-
matches: codec.matches.bind(codec), snapshot: codec.snapshot.bind(codec), clone: codec.clone.bind(codec),
|
|
497
|
-
load: codec.load.bind(codec), headerTrimmable: codec.headerTrimmable.bind(codec),
|
|
498
|
-
})]));
|
|
499
|
-
const candidates = [...table.values()];
|
|
500
|
-
const selected = new WeakMap<Cache, CacheCodec>();
|
|
501
|
-
return Object.freeze({
|
|
502
|
-
id,
|
|
503
|
-
forCache(cache: Cache) {
|
|
504
|
-
const bound = selected.get(cache);
|
|
505
|
-
if (bound) return bound;
|
|
506
|
-
const matches = candidates.filter((codec) => codec.matches(cache));
|
|
507
|
-
if (matches.length !== 1)
|
|
508
|
-
throw new Error(`cache codec provider ${id}: ${matches.length} matches for ${cacheSignature(cache)}`);
|
|
509
|
-
selected.set(cache, matches[0]!);
|
|
510
|
-
return matches[0]!;
|
|
511
|
-
},
|
|
512
|
-
forHeader(entry: CacheHeaderEntry) {
|
|
513
|
-
const codec = table.get(entry.kind);
|
|
514
|
-
if (!codec) throw new Error(`cache codec provider ${id}: unsupported cache kind ${entry.kind}`);
|
|
515
|
-
return codec;
|
|
516
|
-
},
|
|
517
|
-
});
|
|
518
|
-
}
|
|
519
|
-
|
|
520
|
-
export const legacyCacheCodecs = createCacheCodecProvider("mlx-cache-v3", CACHE_CODECS);
|
|
521
|
-
|
|
522
|
-
/** Old TQ files lack the conversion boundary. Keep them reusable at their
|
|
523
|
-
* stored offset, but never infer that an earlier precision state is valid. */
|
|
524
|
-
export function cacheHeaderMinimumReusableOffset(entry: CacheHeaderEntry): number {
|
|
525
|
-
const offset = entry.minimumReusableOffset ?? (entry.kind === "turboquant" ? entry.offset : 0);
|
|
526
|
-
if (!Number.isSafeInteger(offset) || offset < 0 || offset > entry.offset)
|
|
527
|
-
throw new Error("invalid cache minimum reusable offset");
|
|
528
|
-
return offset;
|
|
529
|
-
}
|
|
530
|
-
export const cacheHeadersMinimumReusableOffset = (caches: readonly CacheHeaderEntry[]): number =>
|
|
531
|
-
Math.max(0, ...caches.map(cacheHeaderMinimumReusableOffset));
|
|
532
|
-
|
|
533
|
-
export const cacheHeadersTrimmable = (
|
|
534
|
-
caches: readonly CacheHeaderEntry[], codecs: CacheCodecProvider = legacyCacheCodecs,
|
|
535
|
-
): boolean => caches.every((entry) => codecs.forHeader(entry).headerTrimmable(entry));
|
|
536
|
-
|
|
537
|
-
function snapshotCache(c: Cache, dataOffset: number, codecs: CacheCodecProvider): { entry: CacheHeaderEntry; sources: TensorSource[]; next: number } {
|
|
538
|
-
const slots: TensorSlot[] = [];
|
|
539
|
-
const sources: TensorSource[] = [];
|
|
540
|
-
let cursor = dataOffset;
|
|
541
|
-
const HASH_PLACEHOLDER = "0".repeat(16); // fixed-width; patched at write
|
|
542
|
-
const push = (a: MlxArray, disposeAfter: boolean): void => {
|
|
543
|
-
const off = alignUp(cursor);
|
|
544
|
-
slots.push({ off, bytes: a.nbytes, shape: a.shape, dtype: a.dtype, hash: HASH_PLACEHOLDER });
|
|
545
|
-
sources.push({ arr: a, disposeAfter });
|
|
546
|
-
cursor = off + a.nbytes;
|
|
547
|
-
};
|
|
548
|
-
const liveSlice = (a: MlxArray, upTo: number): MlxArray => {
|
|
549
|
-
const [B, H, , D] = a.shape as [number, number, number, number];
|
|
550
|
-
return a.slice([0, 0, 0, 0], [B, H, upTo, D]);
|
|
551
|
-
};
|
|
552
|
-
const liveMlaSlice = (a: MlxArray, upTo: number): MlxArray => {
|
|
553
|
-
const [B, , D] = a.shape as [number, number, number];
|
|
554
|
-
return a.slice([0, 0, 0], [B, upTo, D]);
|
|
555
|
-
};
|
|
556
|
-
const pushTriple = (t: ops.QuantizedTensor, upTo: number | null): void => {
|
|
557
|
-
for (const a of [t.packed, t.scales, t.biases]) {
|
|
558
|
-
if (upTo === null) push(a, false);
|
|
559
|
-
else push(liveSlice(a, upTo), true);
|
|
560
|
-
}
|
|
561
|
-
};
|
|
562
|
-
|
|
563
|
-
try {
|
|
564
|
-
const entry = codecs.forCache(c).snapshot(c, { slots, push, liveSlice, liveMlaSlice, pushTriple });
|
|
565
|
-
if (c.minimumReusableOffset !== undefined) entry.minimumReusableOffset = c.minimumReusableOffset;
|
|
566
|
-
return { entry, sources, next: cursor };
|
|
567
|
-
} catch (error) {
|
|
568
|
-
for (const source of sources) if (source.disposeAfter) source.arr.dispose();
|
|
569
|
-
throw error;
|
|
570
|
-
}
|
|
571
|
-
}
|
|
572
|
-
|
|
573
|
-
/** Persist `caches` (+ the exact token prefix they encode) to `path`.
|
|
574
|
-
* ATOMIC: written to `<path>.tmp`, fsync'd, renamed — a crash mid-write
|
|
575
|
-
* leaves only a .tmp orphan (ignored + reaped by SsdCacheStore.scan). */
|
|
576
|
-
/** Zero-copy in-memory clone of a cache list at its CURRENT state — new
|
|
577
|
-
* cache instances holding slice VIEWS of the live tensors (mlx arrays are
|
|
578
|
-
* immutable: continued generation replaces the source caches' internal
|
|
579
|
-
* arrays, it never mutates the shared buffers, so views are stable).
|
|
580
|
-
* Built for the prompt-boundary prompt-cache snapshot (multi-turn agent
|
|
581
|
-
* traffic: the reply's decode→encode roundtrip drift makes prompt+gen
|
|
582
|
-
* entries untrimmable at context > sliding window — a prompt-only entry
|
|
583
|
-
* is ALWAYS an exact prefix of the next turn, immune to drift). Must be
|
|
584
|
-
* called while the caches hold EXACTLY the state to snapshot (post-wrap
|
|
585
|
-
* rings cannot be rewound later). Same per-kind dispatch as
|
|
586
|
-
* snapshotCache/loadKvCache. */
|
|
587
|
-
export function cloneKvCaches(caches: Cache[], codecs: CacheCodecProvider = legacyCacheCodecs): Cache[] {
|
|
588
|
-
const pending: MlxArray[] = [];
|
|
589
|
-
const hold = (array: MlxArray): MlxArray => { pending.push(array); return array; };
|
|
590
|
-
const view = (a: MlxArray): MlxArray => hold(a.slice(a.shape.map(() => 0), [...a.shape]));
|
|
591
|
-
const liveView = (a: MlxArray, upTo: number): MlxArray => {
|
|
592
|
-
const [B, H, , D] = a.shape as [number, number, number, number];
|
|
593
|
-
return hold(a.slice([0, 0, 0, 0], [B, H, upTo, D]));
|
|
594
|
-
};
|
|
595
|
-
const mlaView = (a: MlxArray, upTo: number): MlxArray => {
|
|
596
|
-
const [B, , D] = a.shape as [number, number, number];
|
|
597
|
-
return hold(a.slice([0, 0, 0], [B, upTo, D]));
|
|
598
|
-
};
|
|
599
|
-
const tripleView = (t: ops.QuantizedTensor, upTo: number | null): ops.QuantizedTensor => ({
|
|
600
|
-
packed: upTo === null ? view(t.packed) : liveView(t.packed, upTo),
|
|
601
|
-
scales: upTo === null ? view(t.scales) : liveView(t.scales, upTo),
|
|
602
|
-
biases: upTo === null ? view(t.biases) : liveView(t.biases, upTo),
|
|
603
|
-
});
|
|
604
|
-
const out: Cache[] = [];
|
|
605
|
-
try {
|
|
606
|
-
for (const c of caches) {
|
|
607
|
-
const clone = codecs.forCache(c).clone(c, { view, liveView, mlaView, tripleView });
|
|
608
|
-
if (c.minimumReusableOffset !== undefined) clone.minimumReusableOffset = c.minimumReusableOffset;
|
|
609
|
-
out.push(clone);
|
|
610
|
-
pending.length = 0; // transferred to the returned cache
|
|
611
|
-
}
|
|
612
|
-
} catch (err) {
|
|
613
|
-
cleanupFailure(err, () => disposeResources([...pending, ...out]));
|
|
614
|
-
}
|
|
615
|
-
return out;
|
|
616
|
-
}
|
|
617
|
-
|
|
618
|
-
/** Shared v3 writer core, one `yield` after each tensor write. The sync
|
|
619
|
-
* wrapper drains it in a tight loop; the async wrapper awaits a macrotask
|
|
620
|
-
* between yields so the EVENT LOOP KEEPS SERVING while a multi-hundred-MB
|
|
621
|
-
* entry flushes (the write-behind persistence contract: durability never
|
|
622
|
-
* blocks a request). try/finally still runs on early generator close. */
|
|
623
|
-
function* saveKvCacheSteps(path: string, tokens: number[], caches: Cache[], meta: KvSaveMeta, codecs: CacheCodecProvider): Generator<void, void, void> {
|
|
624
|
-
// Plan pass: header entries + lazy tensor sources, NO bytes materialized.
|
|
625
|
-
const entries: CacheHeaderEntry[] = [];
|
|
626
|
-
const sources: TensorSource[] = [];
|
|
627
|
-
let dataOffset = 0; // relative; rebased after header is sized
|
|
628
|
-
try {
|
|
629
|
-
for (const c of caches) {
|
|
630
|
-
const s = snapshotCache(c, dataOffset, codecs);
|
|
631
|
-
entries.push(s.entry);
|
|
632
|
-
sources.push(...s.sources);
|
|
633
|
-
dataOffset = s.next;
|
|
634
|
-
}
|
|
635
|
-
|
|
636
|
-
const { attachments: liveAttachments, ...identity } = meta;
|
|
637
|
-
const attachments = liveAttachments?.map((attachment) => ({
|
|
638
|
-
schema: attachment.schema,
|
|
639
|
-
metadata: { ...attachment.metadata },
|
|
640
|
-
tensors: attachment.tensors.map((arr) => {
|
|
641
|
-
const slot: TensorSlot = { off: dataOffset, bytes: arr.nbytes,
|
|
642
|
-
shape: [...arr.shape], dtype: arr.dtype, hash: "0000000000000000" };
|
|
643
|
-
dataOffset = alignUp(dataOffset + slot.bytes);
|
|
644
|
-
sources.push({ arr, disposeAfter: false });
|
|
645
|
-
return slot;
|
|
646
|
-
}),
|
|
647
|
-
}));
|
|
648
|
-
|
|
649
|
-
// Header size is final NOW: hashes are fixed-width placeholders that get
|
|
650
|
-
// patched in place (same byte length) after the data pass computes them.
|
|
651
|
-
const header: KvFileHeader = {
|
|
652
|
-
formatVersion: attachments?.length ? 4 : 3, createdAt: Date.now(), ...identity,
|
|
653
|
-
...(attachments?.length ? { attachments } : {}), codecProvider: codecs.id, tokens, caches: entries,
|
|
654
|
-
};
|
|
655
|
-
const headerLen = new TextEncoder().encode(JSON.stringify(header)).length;
|
|
656
|
-
const dataStart = alignUp(PREFIX_LEN + headerLen);
|
|
657
|
-
|
|
658
|
-
const tmp = `${path}.tmp`;
|
|
659
|
-
const fd = openSync(tmp, "w");
|
|
660
|
-
try {
|
|
661
|
-
// Data pass: materialize → hash → write → drop, ONE tensor at a time.
|
|
662
|
-
// The hash+write read a ZERO-COPY view of the contiguous mlx buffer
|
|
663
|
-
// (rawBytesView) — never a JS-heap copy. rawBytes() here looked
|
|
664
|
-
// "streamed" but each call allocated a dead per-tensor JS-heap copy
|
|
665
|
-
// whose reclamation is GC-timing-dependent — up to a whole extra
|
|
666
|
-
// entry of heap on a busy server (the A7 residual). Now zero JS
|
|
667
|
-
// allocations; the only real transient is the one mlx-side
|
|
668
|
-
// contiguous copy (allocator-pooled; none at all for contiguous
|
|
669
|
-
// sources — mlx save-transient measured 0 bytes, 2026-07-07). NOTE:
|
|
670
|
-
// the hash+write still CPU-touch the live entry's unified-memory
|
|
671
|
-
// pages, which makes them VISIBLE to ps RSS — accounting, not an
|
|
672
|
-
// allocation (see bench-serve.ts' per-leg RSS note).
|
|
673
|
-
let srcIdx = 0;
|
|
674
|
-
for (const e of [...entries, ...(attachments ?? [])]) {
|
|
675
|
-
for (const slot of e.tensors) {
|
|
676
|
-
const src = sources[srcIdx++]!;
|
|
677
|
-
const c = ops.contiguous(src.arr);
|
|
678
|
-
try {
|
|
679
|
-
const bytes = c.rawBytesView(); // evals; aliases the mlx buffer
|
|
680
|
-
if (bytes.length !== slot.bytes)
|
|
681
|
-
throw new Error(`tensor byte-length drift: planned ${slot.bytes}, got ${bytes.length}`);
|
|
682
|
-
slot.hash = hash64(bytes);
|
|
683
|
-
writeSync(fd, bytes, 0, bytes.length, dataStart + slot.off);
|
|
684
|
-
} finally {
|
|
685
|
-
c.dispose(); // the view dies with the buffer — nothing retains it
|
|
686
|
-
}
|
|
687
|
-
if (src.disposeAfter) { src.arr.dispose(); src.disposeAfter = false; }
|
|
688
|
-
yield;
|
|
689
|
-
}
|
|
690
|
-
}
|
|
691
|
-
|
|
692
|
-
// Header pass: real hashes in, byte length unchanged by construction.
|
|
693
|
-
const headerJson = new TextEncoder().encode(JSON.stringify(header));
|
|
694
|
-
if (headerJson.length !== headerLen)
|
|
695
|
-
throw new Error(`header length drift: planned ${headerLen}, got ${headerJson.length}`);
|
|
696
|
-
const pre = new Uint8Array(dataStart);
|
|
697
|
-
pre.set(new TextEncoder().encode(MAGIC), 0);
|
|
698
|
-
const dv = new DataView(pre.buffer);
|
|
699
|
-
dv.setUint32(MAGIC.length, headerJson.length, true);
|
|
700
|
-
dv.setUint32(MAGIC.length + 4, dataStart, true);
|
|
701
|
-
dv.setBigUint64(MAGIC.length + 8, BigInt(Bun.hash(headerJson)), true);
|
|
702
|
-
pre.set(headerJson, PREFIX_LEN);
|
|
703
|
-
writeSync(fd, pre, 0, pre.length, 0);
|
|
704
|
-
fsyncSync(fd);
|
|
705
|
-
} catch (err) {
|
|
706
|
-
closeSync(fd);
|
|
707
|
-
try { rmSync(tmp, { force: true }); } catch {}
|
|
708
|
-
throw err;
|
|
709
|
-
}
|
|
710
|
-
closeSync(fd);
|
|
711
|
-
renameSync(tmp, path);
|
|
712
|
-
} finally {
|
|
713
|
-
// Slices created in the plan pass are ours to free on EVERY path.
|
|
714
|
-
for (const s of sources) if (s.disposeAfter) s.arr.dispose();
|
|
715
|
-
}
|
|
716
|
-
}
|
|
717
|
-
|
|
718
|
-
export function saveKvCache(path: string, tokens: number[], caches: Cache[], meta: KvSaveMeta = {}, codecs: CacheCodecProvider = legacyCacheCodecs): void {
|
|
719
|
-
for (const _ of saveKvCacheSteps(path, tokens, caches, meta, codecs)) { /* drain */ }
|
|
720
|
-
}
|
|
721
|
-
|
|
722
|
-
/** Non-blocking variant: yields the event loop after every tensor write so
|
|
723
|
-
* serving interleaves with the flush. Caller owns `caches` lifetime for
|
|
724
|
-
* the duration (pass zero-copy clones, dispose after).
|
|
725
|
-
*
|
|
726
|
-
* `runStep` (optional) runs each blocking tensor step inside the caller's
|
|
727
|
-
* exclusion domain. Waiting for an idle observation is not sufficient: a
|
|
728
|
-
* request can acquire the engine between that check and `steps.next()`,
|
|
729
|
-
* putting rawBytesView's GPU sync on the decode stream during prefill. */
|
|
730
|
-
export async function saveKvCacheAsync(
|
|
731
|
-
path: string, tokens: number[], caches: Cache[], meta: KvSaveMeta = {},
|
|
732
|
-
runStep?: <T>(step: () => T) => Promise<T>, codecs: CacheCodecProvider = legacyCacheCodecs,
|
|
733
|
-
): Promise<void> {
|
|
734
|
-
const steps = saveKvCacheSteps(path, tokens, caches, meta, codecs);
|
|
735
|
-
try {
|
|
736
|
-
while (true) {
|
|
737
|
-
const next = runStep ? await runStep(() => steps.next()) : steps.next();
|
|
738
|
-
if (next.done) break;
|
|
739
|
-
await new Promise<void>((r) => setImmediate(r));
|
|
740
|
-
}
|
|
741
|
-
} catch (error) {
|
|
742
|
-
// A rejected execution lease can leave the writer suspended with an
|
|
743
|
-
// open fd. Throw through it so its existing failure cleanup runs.
|
|
744
|
-
steps.throw(error);
|
|
745
|
-
throw error;
|
|
746
|
-
}
|
|
747
|
-
}
|
|
748
|
-
|
|
749
|
-
/** One pending write-behind item: an entry's copied tokens + OWNED
|
|
750
|
-
* zero-copy cache clones (the queue disposes them on every exit path). */
|
|
751
|
-
export interface SpillItem {
|
|
752
|
-
attachments?: CheckpointAttachment[];
|
|
753
|
-
tokens: number[];
|
|
754
|
-
caches: Cache[];
|
|
755
|
-
ns: string;
|
|
756
|
-
}
|
|
757
|
-
|
|
758
|
-
/** Bounded write-behind queue (2026-07-07 post-merge review fix).
|
|
759
|
-
*
|
|
760
|
-
* Pending spill/snapshot clones pin their entries' GPU buffers until the
|
|
761
|
-
* generation-locked flush gets a turn; the old bare promise chain queued them
|
|
762
|
-
* WITHOUT BOUND, so under sustained traffic (gate starved, evictions
|
|
763
|
-
* ongoing) resident memory = prompt-cache cap + every queued clone —
|
|
764
|
-
* allocator pressure exactly under the load that caused the evictions.
|
|
765
|
-
*
|
|
766
|
-
* Policy: queued bytes are capped. Over cap, the OLDEST not-in-flight
|
|
767
|
-
* item is dropped and its clones disposed immediately — cache semantics
|
|
768
|
-
* (a dropped spill is a future cache miss, never a wrong result; the
|
|
769
|
-
* contention-free alternative to letting the flush cut into decode).
|
|
770
|
-
* The item being enqueued is never its own victim (soft cap: one
|
|
771
|
-
* oversized entry may exceed the cap alone rather than never spilling).
|
|
772
|
-
* Items run strictly serially in enqueue order via an internal chain;
|
|
773
|
-
* store failures are swallowed (cold tier is best-effort) but clones
|
|
774
|
-
* are disposed on every settle path. */
|
|
775
|
-
type SpillRec = SpillItem & {
|
|
776
|
-
bytes: number;
|
|
777
|
-
dropped: boolean;
|
|
778
|
-
settle: (stored: boolean) => void;
|
|
779
|
-
};
|
|
780
|
-
|
|
781
|
-
export class SpillQueue {
|
|
782
|
-
#queue: SpillRec[] = [];
|
|
783
|
-
#bytes = 0;
|
|
784
|
-
#dropped = 0;
|
|
785
|
-
#failed = 0;
|
|
786
|
-
#inFlight: object | null = null;
|
|
787
|
-
#chain: Promise<void> = Promise.resolve();
|
|
788
|
-
|
|
789
|
-
constructor(
|
|
790
|
-
readonly capBytes: number,
|
|
791
|
-
/** Byte size of a clone set (prompt-cache's cacheBytes). */
|
|
792
|
-
readonly bytesOf: (caches: Cache[]) => number,
|
|
793
|
-
/** The actual write (server passes storeAsync + the idle gate). */
|
|
794
|
-
readonly store: (item: SpillItem) => Promise<unknown>,
|
|
795
|
-
/** Clone disposal — frees the pinned GPU memory. */
|
|
796
|
-
readonly disposeClones: (caches: Cache[]) => void,
|
|
797
|
-
/** Zero-byte control barriers may need to survive capacity shedding. */
|
|
798
|
-
readonly canDrop: (item: SpillItem) => boolean = () => true,
|
|
799
|
-
) {}
|
|
800
|
-
|
|
801
|
-
/** Bytes pinned by queued (not yet flushed) clones. */
|
|
802
|
-
get pendingBytes(): number { return this.#bytes; }
|
|
803
|
-
get pendingCount(): number { return this.#queue.length; }
|
|
804
|
-
/** Spills dropped by the cap since start (each = one future cache miss). */
|
|
805
|
-
get droppedCount(): number { return this.#dropped; }
|
|
806
|
-
/** Store calls that threw or returned false. */
|
|
807
|
-
get failedCount(): number { return this.#failed; }
|
|
808
|
-
|
|
809
|
-
enqueue(item: SpillItem): Promise<boolean> {
|
|
810
|
-
let settle!: (stored: boolean) => void;
|
|
811
|
-
const result = new Promise<boolean>((resolve) => { settle = resolve; });
|
|
812
|
-
const rec: SpillRec = {
|
|
813
|
-
...item,
|
|
814
|
-
bytes: this.bytesOf(item.caches) + attachmentBytes(item.attachments),
|
|
815
|
-
dropped: false,
|
|
816
|
-
settle,
|
|
817
|
-
};
|
|
818
|
-
this.#queue.push(rec);
|
|
819
|
-
this.#bytes += rec.bytes;
|
|
820
|
-
while (this.#bytes > this.capBytes) {
|
|
821
|
-
const victim = this.#queue.find(
|
|
822
|
-
(p) => !p.dropped && p !== this.#inFlight && p !== rec && this.canDrop(p),
|
|
823
|
-
);
|
|
824
|
-
if (!victim) break; // only the new item (or in-flight) left — soft cap
|
|
825
|
-
this.#drop(victim);
|
|
826
|
-
}
|
|
827
|
-
this.#chain = this.#chain.then(async () => {
|
|
828
|
-
if (rec.dropped) return; // clones already disposed at drop time
|
|
829
|
-
this.#inFlight = rec;
|
|
830
|
-
try {
|
|
831
|
-
const stored = await this.store(rec);
|
|
832
|
-
if (stored === false) this.#failed++;
|
|
833
|
-
rec.settle(stored !== false);
|
|
834
|
-
} catch {
|
|
835
|
-
// best-effort tier — the entry simply won't restore
|
|
836
|
-
this.#failed++;
|
|
837
|
-
rec.settle(false);
|
|
838
|
-
} finally {
|
|
839
|
-
this.#inFlight = null;
|
|
840
|
-
this.#remove(rec);
|
|
841
|
-
disposeResources([{ dispose: () => this.disposeClones(rec.caches) },
|
|
842
|
-
{ dispose: () => disposeAttachments(rec.attachments) }]);
|
|
843
|
-
}
|
|
844
|
-
});
|
|
845
|
-
return result;
|
|
846
|
-
}
|
|
847
|
-
|
|
848
|
-
/** Cancel queued ownership, excluding the write already in flight. */
|
|
849
|
-
cancelWhere(matches: (item: SpillItem) => boolean): void {
|
|
850
|
-
for (const rec of [...this.#queue])
|
|
851
|
-
if (!rec.dropped && rec !== this.#inFlight && matches(rec)) this.#drop(rec);
|
|
852
|
-
}
|
|
853
|
-
|
|
854
|
-
/** Settles when everything enqueued so far has flushed or dropped
|
|
855
|
-
* (tests and graceful shutdown). */
|
|
856
|
-
drain(): Promise<void> { return this.#chain; }
|
|
857
|
-
|
|
858
|
-
#drop(rec: SpillRec): void {
|
|
859
|
-
rec.dropped = true;
|
|
860
|
-
this.#dropped++;
|
|
861
|
-
rec.settle(false);
|
|
862
|
-
this.#remove(rec);
|
|
863
|
-
disposeResources([{ dispose: () => this.disposeClones(rec.caches) },
|
|
864
|
-
{ dispose: () => disposeAttachments(rec.attachments) }]);
|
|
865
|
-
}
|
|
866
|
-
|
|
867
|
-
#remove(rec: SpillRec): void {
|
|
868
|
-
const i = this.#queue.indexOf(rec);
|
|
869
|
-
if (i >= 0) {
|
|
870
|
-
this.#queue.splice(i, 1);
|
|
871
|
-
this.#bytes -= rec.bytes;
|
|
872
|
-
}
|
|
873
|
-
}
|
|
874
|
-
}
|
|
875
|
-
|
|
876
|
-
/** Read only the header (cheap — for prefix matching across many files).
|
|
877
|
-
* Always verifies the header hash; throws on any structural mismatch. */
|
|
878
|
-
export function readKvHeader(path: string): KvFileHeader & { dataStart: number } {
|
|
879
|
-
const fd = openSync(path, "r");
|
|
880
|
-
try {
|
|
881
|
-
const head = new Uint8Array(PREFIX_LEN);
|
|
882
|
-
readSync(fd, head, 0, head.length, 0);
|
|
883
|
-
if (new TextDecoder().decode(head.subarray(0, MAGIC.length)) !== MAGIC)
|
|
884
|
-
throw new Error(`${path}: not an mlx-bun KV v2 cache file`);
|
|
885
|
-
const dv = new DataView(head.buffer);
|
|
886
|
-
const len = dv.getUint32(MAGIC.length, true);
|
|
887
|
-
const dataStart = dv.getUint32(MAGIC.length + 4, true);
|
|
888
|
-
const expectHash = dv.getBigUint64(MAGIC.length + 8, true);
|
|
889
|
-
const body = new Uint8Array(len);
|
|
890
|
-
readSync(fd, body, 0, len, PREFIX_LEN);
|
|
891
|
-
if (BigInt(Bun.hash(body)) !== expectHash)
|
|
892
|
-
throw new Error(`${path}: header hash mismatch (truncated or corrupt)`);
|
|
893
|
-
const header = JSON.parse(new TextDecoder().decode(body)) as KvFileHeader;
|
|
894
|
-
if (header.formatVersion !== 3 && header.formatVersion !== 4)
|
|
895
|
-
throw new Error(`${path}: unsupported formatVersion ${header.formatVersion}`);
|
|
896
|
-
return { ...header, dataStart };
|
|
897
|
-
} finally {
|
|
898
|
-
closeSync(fd);
|
|
899
|
-
}
|
|
900
|
-
}
|
|
901
|
-
|
|
902
|
-
export interface LoadedKvCache {
|
|
903
|
-
attachments?: CheckpointAttachment[];
|
|
904
|
-
tokens: number[];
|
|
905
|
-
header: KvFileHeader;
|
|
906
|
-
caches: Cache[];
|
|
907
|
-
}
|
|
908
|
-
|
|
909
|
-
export interface KvLoadExpect {
|
|
910
|
-
/** Reject on metadata mismatch (pass what the server is running). */
|
|
911
|
-
modelId?: string;
|
|
912
|
-
configFingerprint?: string;
|
|
913
|
-
tokenizerHash?: string;
|
|
914
|
-
ns?: string;
|
|
915
|
-
/** Verify every tensor hash before copying it in (off by default,
|
|
916
|
-
* `--ssd-cache-verify` — the hash pass roughly doubles restore reads). */
|
|
917
|
-
verify?: boolean;
|
|
918
|
-
}
|
|
919
|
-
|
|
920
|
-
/** Reject GLM cache/layout drift before opening the tensor mmap. Generic KV
|
|
921
|
-
* kinds may legitimately differ from makeCache() after runtime quantization;
|
|
922
|
-
* GLM's checkpoint-native cache geometry does not. */
|
|
923
|
-
function validateGlm52Prototype(
|
|
924
|
-
path: string,
|
|
925
|
-
entry: CacheHeaderEntry,
|
|
926
|
-
prototype: Cache,
|
|
927
|
-
): void {
|
|
928
|
-
const glmKind = entry.kind === "mla" || entry.kind === "mla-dsa" ||
|
|
929
|
-
entry.kind === "mtp-mla";
|
|
930
|
-
const prototypeSignature = cacheSignature(prototype);
|
|
931
|
-
if (!prototypeSignature.startsWith("kv:mla:")) {
|
|
932
|
-
if (glmKind)
|
|
933
|
-
throw new Error(`${path}: GLM cache kind ${entry.kind} does not match model cache`);
|
|
934
|
-
return;
|
|
935
|
-
}
|
|
936
|
-
const glmPrototype = prototype as Glm52Cache;
|
|
937
|
-
const expectedKind: CacheKind = glmPrototype.role === "mtp"
|
|
938
|
-
? "mtp-mla"
|
|
939
|
-
: glmPrototype.dsa ? "mla-dsa" : "mla";
|
|
940
|
-
if (entry.kind !== expectedKind) {
|
|
941
|
-
throw new Error(
|
|
942
|
-
`${path}: cache kind ${entry.kind} != model cache kind ${expectedKind}`,
|
|
943
|
-
);
|
|
944
|
-
}
|
|
945
|
-
for (const [field, actual, expected] of [
|
|
946
|
-
["kvLoraRank", entry.kvLoraRank, glmPrototype.kvLoraRank],
|
|
947
|
-
["ropeHeadDim", entry.ropeHeadDim, glmPrototype.ropeHeadDim],
|
|
948
|
-
["dsaHeadDim", entry.dsaHeadDim, glmPrototype.dsa?.headDim],
|
|
949
|
-
["maxTokens", entry.maxTokens, glmPrototype.maxTokens],
|
|
950
|
-
] as const) {
|
|
951
|
-
if (actual !== expected) {
|
|
952
|
-
throw new Error(
|
|
953
|
-
`${path}: ${field} ${String(actual)} != model ${String(expected)}`,
|
|
954
|
-
);
|
|
955
|
-
}
|
|
956
|
-
}
|
|
957
|
-
const expectedTensors = glmPrototype.dsa ? 3 : 2;
|
|
958
|
-
if (entry.tensors.length !== expectedTensors) {
|
|
959
|
-
throw new Error(
|
|
960
|
-
`${path}: ${entry.kind} has ${entry.tensors.length} tensors, ` +
|
|
961
|
-
`expected ${expectedTensors}`,
|
|
962
|
-
);
|
|
963
|
-
}
|
|
964
|
-
}
|
|
965
|
-
|
|
966
|
-
/** Copy-restore a full-attention KV tensor into STEP-rounded capacity with
|
|
967
|
-
* ≥1 token of slack (what a live mid-generation cache looks like) so the
|
|
968
|
-
* first post-restore updateAndFetch never takes the grow path — the old
|
|
969
|
-
* exactly-offset-sized restore made that first step concat-copy the ENTIRE
|
|
970
|
-
* entry into a fresh buffer. Zero padding is bit-safe: writes land at
|
|
971
|
-
* [offset..) before any read, and attention only reads [:offset+L).
|
|
972
|
-
* Evaluated EAGERLY so the host-copy leaf frees before the next tensor
|
|
973
|
-
* streams in (bounded transient). Plain KVCache only — rotating rings and
|
|
974
|
-
* quantized triples are position-exact layouts. */
|
|
975
|
-
function withStepCapacity(a: MlxArray, offset: number): MlxArray {
|
|
976
|
-
const [B, H, S, D] = a.shape as [number, number, number, number];
|
|
977
|
-
const cap = Math.ceil((offset + 1) / KVCache.STEP) * KVCache.STEP;
|
|
978
|
-
if (S >= cap) return a;
|
|
979
|
-
// Internals drained on throw (2026-07-07 review: a grown.eval() failure
|
|
980
|
-
// orphaned z/grown; `a` stays the CALLER's to free — loadKvCache's
|
|
981
|
-
// pending[] holds it).
|
|
982
|
-
let z: MlxArray | null = null;
|
|
983
|
-
let grown: MlxArray | null = null;
|
|
984
|
-
try {
|
|
985
|
-
z = ops.zeros([B, H, cap, D], a.dtype);
|
|
986
|
-
grown = ops.sliceUpdate(z, a, [0, 0, 0, 0], [B, H, S, D]);
|
|
987
|
-
z.dispose();
|
|
988
|
-
z = null;
|
|
989
|
-
grown.eval();
|
|
990
|
-
a.dispose();
|
|
991
|
-
const out = grown;
|
|
992
|
-
grown = null;
|
|
993
|
-
return out;
|
|
994
|
-
} finally {
|
|
995
|
-
z?.dispose();
|
|
996
|
-
grown?.dispose();
|
|
997
|
-
}
|
|
998
|
-
}
|
|
999
|
-
|
|
1000
|
-
/** Reload by STREAMED COPY (see the header note): every tensor is copied
|
|
1001
|
-
* into an mlx-owned leaf; the mapping is read-only, its pages dropped
|
|
1002
|
-
* per-tensor, and unmapped before returning — the caches own their bytes
|
|
1003
|
-
* outright (no dtor contract, no pinned mapping, nothing to retain).
|
|
1004
|
-
* `model` is anything with makeCache() — the entry count is validated
|
|
1005
|
-
* against the DONOR cache list (model.layers.length was wrong for
|
|
1006
|
-
* KV-shared models like e4b, whose makeCache() returns donors only). */
|
|
1007
|
-
export function loadKvCache(
|
|
1008
|
-
path: string,
|
|
1009
|
-
model: { makeCache(): Cache[] },
|
|
1010
|
-
expect: KvLoadExpect = {}, codecs: CacheCodecProvider = legacyCacheCodecs,
|
|
1011
|
-
): LoadedKvCache {
|
|
1012
|
-
const header = readKvHeader(path);
|
|
1013
|
-
if ((header.codecProvider ?? legacyCacheCodecs.id) !== codecs.id)
|
|
1014
|
-
throw new Error(`${path}: cache codec provider mismatch`);
|
|
1015
|
-
for (const entry of header.caches) codecs.forHeader(entry);
|
|
1016
|
-
for (const key of [
|
|
1017
|
-
"modelId",
|
|
1018
|
-
"configFingerprint",
|
|
1019
|
-
"tokenizerHash",
|
|
1020
|
-
"ns",
|
|
1021
|
-
] as const) {
|
|
1022
|
-
if (expect[key] !== undefined && header[key] !== expect[key])
|
|
1023
|
-
throw new Error(`${path}: ${key} mismatch (file ${header[key]}, expected ${expect[key]})`);
|
|
1024
|
-
}
|
|
1025
|
-
const proto = model.makeCache();
|
|
1026
|
-
try {
|
|
1027
|
-
if (header.caches.length !== proto.length) {
|
|
1028
|
-
throw new Error(
|
|
1029
|
-
`${path}: ${header.caches.length} cached layers but model has ${proto.length}`,
|
|
1030
|
-
);
|
|
1031
|
-
}
|
|
1032
|
-
for (let index = 0; index < proto.length; index++) {
|
|
1033
|
-
validateGlm52Prototype(path, header.caches[index]!, proto[index]!);
|
|
1034
|
-
}
|
|
1035
|
-
} finally {
|
|
1036
|
-
for (const c of proto) c.dispose();
|
|
1037
|
-
}
|
|
1038
|
-
|
|
1039
|
-
const mmap = MmapFile.open(path, "ro");
|
|
1040
|
-
const dataStart = header.dataStart;
|
|
1041
|
-
// Every tensor materialized for the CURRENT entry, drained on any throw
|
|
1042
|
-
// (2026-07-07 review: a hash mismatch under --ssd-cache-verify — or any
|
|
1043
|
-
// per-tensor failure — orphaned the entry's already-built tensors: the
|
|
1044
|
-
// catch below only saw completed caches). Cleared once the entry's cache
|
|
1045
|
-
// takes ownership; disposing an array withStepCapacity already consumed
|
|
1046
|
-
// is safe (dispose is idempotent).
|
|
1047
|
-
const pending: MlxArray[] = [];
|
|
1048
|
-
const arr = (slot: TensorSlot): MlxArray => {
|
|
1049
|
-
const view = mmap.view(dataStart + slot.off, slot.bytes);
|
|
1050
|
-
if (expect.verify && hash64(view) !== slot.hash)
|
|
1051
|
-
throw new Error(`${path}: tensor hash mismatch at offset ${slot.off}`);
|
|
1052
|
-
// mlx_array_new_data COPIES synchronously — the view only has to
|
|
1053
|
-
// outlive this call, so the mapping can be advised/unmapped freely.
|
|
1054
|
-
const a = MlxArray.fromBytesCopy(view, slot.shape, slot.dtype as Dtype);
|
|
1055
|
-
pending.push(a);
|
|
1056
|
-
// Drop the just-read clean pages so the load's transient stays at
|
|
1057
|
-
// live entry + one tensor (ALIGN = 16 KiB == the arm64 page size, so
|
|
1058
|
-
// every tensor offset is page-aligned by construction).
|
|
1059
|
-
const len = Math.min(alignUp(slot.bytes), mmap.size - (dataStart + slot.off));
|
|
1060
|
-
mmap.advise(dataStart + slot.off, len, MADV_DONTNEED);
|
|
1061
|
-
return a;
|
|
1062
|
-
};
|
|
1063
|
-
/** withStepCapacity whose (possibly fresh) result is pending-tracked. */
|
|
1064
|
-
const grownArr = (slot: TensorSlot, offset: number): MlxArray => {
|
|
1065
|
-
const g = withStepCapacity(arr(slot), offset);
|
|
1066
|
-
pending.push(g);
|
|
1067
|
-
return g;
|
|
1068
|
-
};
|
|
1069
|
-
const triple = (slots: TensorSlot[], at: number): ops.QuantizedTensor =>
|
|
1070
|
-
({ packed: arr(slots[at]!), scales: arr(slots[at + 1]!), biases: arr(slots[at + 2]!) });
|
|
1071
|
-
|
|
1072
|
-
const caches: Cache[] = [];
|
|
1073
|
-
const attachments: CheckpointAttachment[] = [];
|
|
1074
|
-
try {
|
|
1075
|
-
for (const e of header.caches) {
|
|
1076
|
-
const codec = codecs.forHeader(e);
|
|
1077
|
-
const minimumOffset = cacheHeaderMinimumReusableOffset(e);
|
|
1078
|
-
const cache = codec.load(e, { path, arr, grownArr, triple });
|
|
1079
|
-
cache.minimumReusableOffset = minimumOffset;
|
|
1080
|
-
caches.push(cache);
|
|
1081
|
-
// This entry's tensors are now owned by its cache — stop tracking them
|
|
1082
|
-
// (the catch must not double-free through both pending AND caches).
|
|
1083
|
-
pending.length = 0;
|
|
1084
|
-
}
|
|
1085
|
-
for (const attachment of header.attachments ?? []) {
|
|
1086
|
-
attachments.push({ schema: attachment.schema, metadata: { ...attachment.metadata },
|
|
1087
|
-
tensors: attachment.tensors.map(arr) });
|
|
1088
|
-
pending.length = 0;
|
|
1089
|
-
}
|
|
1090
|
-
} catch (err) {
|
|
1091
|
-
disposeAttachments(attachments);
|
|
1092
|
-
for (const a of pending) a.dispose(); // the mid-entry orphans
|
|
1093
|
-
for (const c of caches) c.dispose();
|
|
1094
|
-
throw err;
|
|
1095
|
-
} finally {
|
|
1096
|
-
mmap.unmap(); // every tensor was copied out — nothing aliases the file
|
|
1097
|
-
}
|
|
1098
|
-
return { tokens: header.tokens, header, caches, ...(attachments.length ? { attachments } : {}) };
|
|
1099
|
-
}
|