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
|
@@ -1,170 +0,0 @@
|
|
|
1
|
-
// Universal-tier primitives: dense (unquantized) linear/embedding modules
|
|
2
|
-
// (Phase 1.5 — mlx nn.Linear / nn.Embedding semantics), quantized-or-dense
|
|
3
|
-
// loader helpers, norm loaders, and the load-time weight audit.
|
|
4
|
-
//
|
|
5
|
-
// The quantized paths delegate to the proven gemma4-base primitives; the
|
|
6
|
-
// dense paths are verbatim ports of mlx.nn.Linear (`mx.addmm(bias, x, W.T)`
|
|
7
|
-
// / `x @ W.T`) and mlx.nn.Embedding (`weight[ids]`, `x @ weight.T`).
|
|
8
|
-
|
|
9
|
-
import type { ModelConfig } from "../../config";
|
|
10
|
-
import type { Weights } from "../../weights";
|
|
11
|
-
import { MlxArray } from "../../mlx/array";
|
|
12
|
-
import * as ops from "../../mlx/ops";
|
|
13
|
-
import { QuantizedEmbedding, QuantizedLinear, RMSNorm } from "../gemma4-base";
|
|
14
|
-
import { TrellisLinear } from "../trellis-linear";
|
|
15
|
-
|
|
16
|
-
/** Records every tensor a universal load consumes; `finish` diffs against
|
|
17
|
-
* the shard index so a descriptor mistake is a LOAD error (unconsumed /
|
|
18
|
-
* missing tensors named), never a silently-wrong model
|
|
19
|
-
* (docs/design/generic-model-support.md §3.4). */
|
|
20
|
-
export class WeightAudit {
|
|
21
|
-
readonly consumed = new Set<string>();
|
|
22
|
-
|
|
23
|
-
use(name: string): void {
|
|
24
|
-
this.consumed.add(name);
|
|
25
|
-
}
|
|
26
|
-
|
|
27
|
-
/** `drop` = the arch's sanitize rules (tensors mlx-lm discards on load,
|
|
28
|
-
* e.g. rotary_emb.inv_freq, or lm_head.weight under tied embeddings). */
|
|
29
|
-
finish(weights: Weights, drop: RegExp[]): void {
|
|
30
|
-
const unconsumed = weights.tensorNames.filter(
|
|
31
|
-
(n) => !this.consumed.has(n) && !drop.some((re) => re.test(n)),
|
|
32
|
-
);
|
|
33
|
-
if (unconsumed.length > 0)
|
|
34
|
-
throw new Error(
|
|
35
|
-
`weight audit: ${unconsumed.length} tensor(s) in the checkpoint were not consumed ` +
|
|
36
|
-
`by the universal module (descriptor mismatch?): ${unconsumed.slice(0, 12).join(", ")}` +
|
|
37
|
-
(unconsumed.length > 12 ? ", …" : ""),
|
|
38
|
-
);
|
|
39
|
-
}
|
|
40
|
-
}
|
|
41
|
-
|
|
42
|
-
/** Load a tensor and record it in the audit. */
|
|
43
|
-
function tensorUsed(weights: Weights, audit: WeightAudit, name: string): MlxArray {
|
|
44
|
-
audit.use(name);
|
|
45
|
-
return weights.tensor(name);
|
|
46
|
-
}
|
|
47
|
-
|
|
48
|
-
/** Plain unquantized linear — mlx nn.Linear:
|
|
49
|
-
* `mx.addmm(bias, x, weight.T)` when bias is present, else `x @ weight.T`. */
|
|
50
|
-
export class DenseLinear {
|
|
51
|
-
readonly wT: MlxArray;
|
|
52
|
-
|
|
53
|
-
constructor(readonly w: MlxArray, readonly bias: MlxArray | null) {
|
|
54
|
-
this.wT = ops.transposeAxes(w, [1, 0]); // lazy view, shared across calls
|
|
55
|
-
}
|
|
56
|
-
|
|
57
|
-
get inFeatures(): number {
|
|
58
|
-
return this.w.shape[1]!;
|
|
59
|
-
}
|
|
60
|
-
get outFeatures(): number {
|
|
61
|
-
return this.w.shape[0]!;
|
|
62
|
-
}
|
|
63
|
-
|
|
64
|
-
forward(x: MlxArray): MlxArray {
|
|
65
|
-
return this.bias ? ops.addmm(this.bias, x, this.wT) : ops.matmul(x, this.wT);
|
|
66
|
-
}
|
|
67
|
-
}
|
|
68
|
-
|
|
69
|
-
/** Either linear flavor behind one call surface. */
|
|
70
|
-
export type AnyLinear = QuantizedLinear | DenseLinear | TrellisLinear;
|
|
71
|
-
|
|
72
|
-
/** Quantized when `.scales` exists (MLX-quantized checkpoints), else dense
|
|
73
|
-
* bf16/f16/f32 (Phase 1.5). Both flavors carry the optional ADDITIVE
|
|
74
|
-
* `.bias` term (qwen2 qkv, starcoder2, …). */
|
|
75
|
-
export function loadLinear(
|
|
76
|
-
weights: Weights, path: string, config: ModelConfig, audit: WeightAudit,
|
|
77
|
-
): AnyLinear {
|
|
78
|
-
if (TrellisLinear.isTrellis(config, path)) {
|
|
79
|
-
const lin = TrellisLinear.load(weights, path, config);
|
|
80
|
-
audit.use(`${path}.weight`);
|
|
81
|
-
audit.use(`${path}.scales`);
|
|
82
|
-
return lin;
|
|
83
|
-
}
|
|
84
|
-
if (weights.has(`${path}.scales`)) {
|
|
85
|
-
const lin = QuantizedLinear.load(weights, path, config);
|
|
86
|
-
audit.use(`${path}.weight`);
|
|
87
|
-
audit.use(`${path}.scales`);
|
|
88
|
-
if (weights.has(`${path}.biases`)) audit.use(`${path}.biases`);
|
|
89
|
-
if (weights.has(`${path}.bias`)) audit.use(`${path}.bias`);
|
|
90
|
-
return lin;
|
|
91
|
-
}
|
|
92
|
-
if (!weights.has(`${path}.weight`))
|
|
93
|
-
throw new Error(`${path}: no .weight tensor (nor .scales) in checkpoint`);
|
|
94
|
-
return new DenseLinear(
|
|
95
|
-
tensorUsed(weights, audit, `${path}.weight`),
|
|
96
|
-
weights.has(`${path}.bias`) ? tensorUsed(weights, audit, `${path}.bias`) : null,
|
|
97
|
-
);
|
|
98
|
-
}
|
|
99
|
-
|
|
100
|
-
/** Plain unquantized embedding — mlx nn.Embedding (+ as_linear tied head). */
|
|
101
|
-
export class DenseEmbedding {
|
|
102
|
-
readonly wT: MlxArray;
|
|
103
|
-
|
|
104
|
-
constructor(readonly w: MlxArray) {
|
|
105
|
-
this.wT = ops.transposeAxes(w, [1, 0]);
|
|
106
|
-
}
|
|
107
|
-
|
|
108
|
-
encode(ids: MlxArray): MlxArray {
|
|
109
|
-
return ops.takeAxis(this.w, ids, 0);
|
|
110
|
-
}
|
|
111
|
-
|
|
112
|
-
/** nn.Embedding.as_linear: `x @ weight.T`. */
|
|
113
|
-
asLinear(h: MlxArray): MlxArray {
|
|
114
|
-
return ops.matmul(h, this.wT);
|
|
115
|
-
}
|
|
116
|
-
}
|
|
117
|
-
|
|
118
|
-
export type AnyEmbedding = QuantizedEmbedding | DenseEmbedding;
|
|
119
|
-
|
|
120
|
-
export function loadEmbedding(
|
|
121
|
-
weights: Weights, path: string, config: ModelConfig, audit: WeightAudit,
|
|
122
|
-
): AnyEmbedding {
|
|
123
|
-
if (weights.has(`${path}.scales`)) {
|
|
124
|
-
const emb = QuantizedEmbedding.load(weights, path, config);
|
|
125
|
-
audit.use(`${path}.weight`);
|
|
126
|
-
audit.use(`${path}.scales`);
|
|
127
|
-
if (weights.has(`${path}.biases`)) audit.use(`${path}.biases`);
|
|
128
|
-
return emb;
|
|
129
|
-
}
|
|
130
|
-
return new DenseEmbedding(tensorUsed(weights, audit, `${path}.weight`));
|
|
131
|
-
}
|
|
132
|
-
|
|
133
|
-
/** RMSNorm loader; `plusOne` = the gemma-family `1.0 + weight` variant
|
|
134
|
-
* (the add is folded once at load — identical value to mlx-lm's
|
|
135
|
-
* per-call `1.0 + self.weight`). */
|
|
136
|
-
export function loadRmsNorm(
|
|
137
|
-
weights: Weights, path: string, eps: number, plusOne: boolean, audit: WeightAudit,
|
|
138
|
-
): RMSNorm {
|
|
139
|
-
const w = tensorUsed(weights, audit, `${path}.weight`);
|
|
140
|
-
if (!plusOne) return new RMSNorm(w, eps);
|
|
141
|
-
const one = ops.scalarLike(1, w);
|
|
142
|
-
const w1 = ops.add(one, w);
|
|
143
|
-
one.dispose();
|
|
144
|
-
return new RMSNorm(w1, eps);
|
|
145
|
-
}
|
|
146
|
-
|
|
147
|
-
/** mlx nn.LayerNorm (weight + bias affine — starcoder2's norm). */
|
|
148
|
-
export class LayerNorm {
|
|
149
|
-
constructor(
|
|
150
|
-
readonly weight: MlxArray,
|
|
151
|
-
readonly bias: MlxArray | null,
|
|
152
|
-
readonly eps: number,
|
|
153
|
-
) {}
|
|
154
|
-
|
|
155
|
-
forward(x: MlxArray): MlxArray {
|
|
156
|
-
return ops.layerNorm(x, this.weight, this.bias, this.eps);
|
|
157
|
-
}
|
|
158
|
-
}
|
|
159
|
-
|
|
160
|
-
export type AnyNorm = RMSNorm | LayerNorm;
|
|
161
|
-
|
|
162
|
-
export function loadLayerNorm(
|
|
163
|
-
weights: Weights, path: string, eps: number, audit: WeightAudit,
|
|
164
|
-
): LayerNorm {
|
|
165
|
-
return new LayerNorm(
|
|
166
|
-
tensorUsed(weights, audit, `${path}.weight`),
|
|
167
|
-
weights.has(`${path}.bias`) ? tensorUsed(weights, audit, `${path}.bias`) : null,
|
|
168
|
-
eps,
|
|
169
|
-
);
|
|
170
|
-
}
|
|
@@ -1,381 +0,0 @@
|
|
|
1
|
-
// Universal rope factory — faithful port of mlx_lm/models/rope_utils.py
|
|
2
|
-
// (oracle venv, mlx-lm 0.31.3): initialize_rope with the default / linear /
|
|
3
|
-
// llama3 (Llama3RoPE) / yarn (YarnRoPE) / longrope (SuScaledRoPE) /
|
|
4
|
-
// proportional (ProportionalRoPE) variants.
|
|
5
|
-
//
|
|
6
|
-
// Porting discipline (three-level fidelity tree, L1): every frequency
|
|
7
|
-
// table is computed with the SAME mlx float32 ops the oracle uses at
|
|
8
|
-
// module init — not host-side float64 math — so the freqs bytes match the
|
|
9
|
-
// oracle's bit-for-bit. Host math (Math.log/floor/ceil) is used only where
|
|
10
|
-
// the oracle uses python-float math (yarn correction range, mscale, su
|
|
11
|
-
// scale), which is IEEE double on both sides.
|
|
12
|
-
//
|
|
13
|
-
// The existing runtime paths (ops.rope default, gemma4's proportional
|
|
14
|
-
// freqs) are untouched; this factory serves the universal (Tier-0) models.
|
|
15
|
-
|
|
16
|
-
import { MlxArray } from "../../mlx/array";
|
|
17
|
-
import { Dtype } from "../../mlx/ffi";
|
|
18
|
-
import * as ops from "../../mlx/ops";
|
|
19
|
-
|
|
20
|
-
/** Raw HF `rope_scaling` dict (values read exactly like the oracle). */
|
|
21
|
-
export type RopeScalingConfig = Record<string, unknown>;
|
|
22
|
-
|
|
23
|
-
const num = (v: unknown, fallback: number): number =>
|
|
24
|
-
typeof v === "number" ? v : fallback;
|
|
25
|
-
|
|
26
|
-
const requireNum = (sc: RopeScalingConfig, key: string): number => {
|
|
27
|
-
const v = sc[key];
|
|
28
|
-
if (typeof v !== "number") throw new Error(`rope_scaling.${key} missing or not a number`);
|
|
29
|
-
return v;
|
|
30
|
-
};
|
|
31
|
-
|
|
32
|
-
/** One rope instance (shared across layers — identical params ⇒ identical
|
|
33
|
-
* values). `apply` returns a NEW array; the caller owns disposal of both. */
|
|
34
|
-
export class UniversalRope {
|
|
35
|
-
constructor(
|
|
36
|
-
readonly dims: number,
|
|
37
|
-
readonly traditional: boolean,
|
|
38
|
-
/** Base for native fast_rope; null when `freqs` supplies the table. */
|
|
39
|
-
readonly base: number | null,
|
|
40
|
-
/** fast_rope scale (1/factor for "linear", else 1). */
|
|
41
|
-
readonly scale: number,
|
|
42
|
-
/** Precomputed frequency table (llama3/yarn/longrope/proportional). */
|
|
43
|
-
readonly freqs: MlxArray | null,
|
|
44
|
-
/** Pre-rope input scale on x[..., :dims] (yarn mscale / su scale). */
|
|
45
|
-
readonly preScale: number,
|
|
46
|
-
/** Apply preScale even when it is exactly 1.0 (SuScaledRoPE does; Yarn
|
|
47
|
-
* skips when mscale == 1.0). */
|
|
48
|
-
readonly preScaleAlways: boolean,
|
|
49
|
-
) {}
|
|
50
|
-
|
|
51
|
-
apply(x: MlxArray, offset: number): MlxArray {
|
|
52
|
-
const input = this.#preScaled(x);
|
|
53
|
-
const out = ops.ropeScaled(
|
|
54
|
-
input ?? x, this.dims, this.traditional, this.base, this.scale, offset, this.freqs,
|
|
55
|
-
);
|
|
56
|
-
input?.dispose();
|
|
57
|
-
return out;
|
|
58
|
-
}
|
|
59
|
-
|
|
60
|
-
/** Per-row array-offset variant (batched decode under left-padding: each
|
|
61
|
-
* row's REAL position rides an int32 [B] array — the scalar offset would
|
|
62
|
-
* mis-position every padded row by its pad amount). Same math as
|
|
63
|
-
* `apply`; only the offset spelling differs. */
|
|
64
|
-
applyDynamic(x: MlxArray, offsetArr: MlxArray): MlxArray {
|
|
65
|
-
const input = this.#preScaled(x);
|
|
66
|
-
const out = ops.ropeScaledDynamic(
|
|
67
|
-
input ?? x, this.dims, this.traditional, this.base, this.scale, offsetArr, this.freqs,
|
|
68
|
-
);
|
|
69
|
-
input?.dispose();
|
|
70
|
-
return out;
|
|
71
|
-
}
|
|
72
|
-
|
|
73
|
-
/** x[..., :dims] = preScale * x[..., :dims] (weak-scalar mul); null when
|
|
74
|
-
* no pre-scale applies (caller uses x unchanged). */
|
|
75
|
-
#preScaled(x: MlxArray): MlxArray | null {
|
|
76
|
-
if (!(this.preScaleAlways || this.preScale !== 1.0)) return null;
|
|
77
|
-
const D = x.shape[x.shape.length - 1]!;
|
|
78
|
-
if (this.dims === D) return ops.mulScalar(x, this.preScale);
|
|
79
|
-
const start = x.shape.map(() => 0);
|
|
80
|
-
const stop = [...x.shape];
|
|
81
|
-
stop[stop.length - 1] = this.dims;
|
|
82
|
-
const head = x.slice(start, stop);
|
|
83
|
-
const headScaled = ops.mulScalar(head, this.preScale);
|
|
84
|
-
head.dispose();
|
|
85
|
-
const scaled = ops.sliceUpdate(x, headScaled, start, stop);
|
|
86
|
-
headScaled.dispose();
|
|
87
|
-
return scaled;
|
|
88
|
-
}
|
|
89
|
-
|
|
90
|
-
dispose(): void {
|
|
91
|
-
this.freqs?.dispose();
|
|
92
|
-
}
|
|
93
|
-
}
|
|
94
|
-
|
|
95
|
-
/** SmolLM3 NoPE: identity stand-in for disabled-rope layers. */
|
|
96
|
-
export class NoRope extends UniversalRope {
|
|
97
|
-
constructor() {
|
|
98
|
-
super(0, false, null, 1, null, 1, false);
|
|
99
|
-
}
|
|
100
|
-
override apply(x: MlxArray, _offset: number): MlxArray {
|
|
101
|
-
// Return an owned copy-view so callers can dispose uniformly.
|
|
102
|
-
return x.slice(x.shape.map(() => 0), [...x.shape]);
|
|
103
|
-
}
|
|
104
|
-
override applyDynamic(x: MlxArray, _offsetArr: MlxArray): MlxArray {
|
|
105
|
-
return x.slice(x.shape.map(() => 0), [...x.shape]);
|
|
106
|
-
}
|
|
107
|
-
}
|
|
108
|
-
|
|
109
|
-
/** `base ** (arange(0, dims, 2) / dims)` with the oracle's dtypes:
|
|
110
|
-
* int32 arange for llama3 (no explicit dtype in the source), float32
|
|
111
|
-
* arange for yarn/longrope. Integer÷integer divides to float32 in mlx,
|
|
112
|
-
* so both spellings compute the pow in f32. */
|
|
113
|
-
function baseFreqs(dims: number, base: number, arangeDtype: Dtype): MlxArray {
|
|
114
|
-
const ar = ops.arange(0, dims, 2, arangeDtype);
|
|
115
|
-
const dimsScalar = arangeDtype === Dtype.int32
|
|
116
|
-
? ops.fromInt32([dims], [])
|
|
117
|
-
: MlxArray.fromFloat32(new Float32Array([dims]), []);
|
|
118
|
-
const exps = ops.div(ar, dimsScalar); // → float32
|
|
119
|
-
ar.dispose();
|
|
120
|
-
dimsScalar.dispose();
|
|
121
|
-
const baseScalar = ops.scalarLike(base, exps);
|
|
122
|
-
const out = ops.pow(baseScalar, exps);
|
|
123
|
-
baseScalar.dispose();
|
|
124
|
-
exps.dispose();
|
|
125
|
-
return out; // float32 [dims/2]
|
|
126
|
-
}
|
|
127
|
-
|
|
128
|
-
/** Llama3RoPE frequency table (rope_utils.py:74-108, verbatim). */
|
|
129
|
-
function llama3Freqs(dims: number, base: number, sc: RopeScalingConfig): MlxArray {
|
|
130
|
-
const factor = requireNum(sc, "factor");
|
|
131
|
-
const lowFreqFactor = num(sc.low_freq_factor, 1.0);
|
|
132
|
-
const highFreqFactor = num(sc.high_freq_factor, 4.0);
|
|
133
|
-
const oldContextLen = num(sc.original_max_position_embeddings, 8192);
|
|
134
|
-
|
|
135
|
-
const lowFreqWavelen = oldContextLen / lowFreqFactor;
|
|
136
|
-
const highFreqWavelen = oldContextLen / highFreqFactor;
|
|
137
|
-
|
|
138
|
-
let freqs = baseFreqs(dims, base, Dtype.int32);
|
|
139
|
-
const wavelens = ops.mulScalar(freqs, 2 * Math.PI);
|
|
140
|
-
|
|
141
|
-
// freqs = where(wavelens > low_freq_wavelen, freqs * factor, freqs)
|
|
142
|
-
const lfw = ops.scalarLike(lowFreqWavelen, wavelens);
|
|
143
|
-
const gtLow = ops.less(lfw, wavelens);
|
|
144
|
-
const freqsTimesFactor = ops.mulScalar(freqs, factor);
|
|
145
|
-
const freqs1 = ops.where(gtLow, freqsTimesFactor, freqs);
|
|
146
|
-
for (const a of [gtLow, freqsTimesFactor, freqs]) a.dispose();
|
|
147
|
-
freqs = freqs1;
|
|
148
|
-
|
|
149
|
-
// is_medium = (wavelens > high_freq_wavelen) & (wavelens < low_freq_wavelen)
|
|
150
|
-
const hfw = ops.scalarLike(highFreqWavelen, wavelens);
|
|
151
|
-
const gtHigh = ops.less(hfw, wavelens);
|
|
152
|
-
const ltLow = ops.less(wavelens, lfw);
|
|
153
|
-
const isMedium = ops.logicalAnd(gtHigh, ltLow);
|
|
154
|
-
for (const a of [hfw, lfw, gtHigh, ltLow]) a.dispose();
|
|
155
|
-
|
|
156
|
-
// smooth_factors = (old_context_len / wavelens - low) / (high - low)
|
|
157
|
-
const oldScalar = ops.scalarLike(oldContextLen, wavelens);
|
|
158
|
-
const oldOverWav = ops.div(oldScalar, wavelens);
|
|
159
|
-
oldScalar.dispose();
|
|
160
|
-
wavelens.dispose();
|
|
161
|
-
const lowScalar = ops.scalarLike(lowFreqFactor, oldOverWav);
|
|
162
|
-
const numArr = ops.sub(oldOverWav, lowScalar);
|
|
163
|
-
oldOverWav.dispose();
|
|
164
|
-
lowScalar.dispose();
|
|
165
|
-
const denScalar = ops.scalarLike(highFreqFactor - lowFreqFactor, numArr);
|
|
166
|
-
const smooth = ops.div(numArr, denScalar);
|
|
167
|
-
numArr.dispose();
|
|
168
|
-
denScalar.dispose();
|
|
169
|
-
|
|
170
|
-
// smooth_freqs = freqs / ((1 - smooth) / factor + smooth)
|
|
171
|
-
const one = ops.scalarLike(1, smooth);
|
|
172
|
-
const oneMinus = ops.sub(one, smooth);
|
|
173
|
-
one.dispose();
|
|
174
|
-
const factorScalar = ops.scalarLike(factor, oneMinus);
|
|
175
|
-
const scaledDown = ops.div(oneMinus, factorScalar);
|
|
176
|
-
oneMinus.dispose();
|
|
177
|
-
factorScalar.dispose();
|
|
178
|
-
const denom = ops.add(scaledDown, smooth);
|
|
179
|
-
scaledDown.dispose();
|
|
180
|
-
smooth.dispose();
|
|
181
|
-
const smoothFreqs = ops.div(freqs, denom);
|
|
182
|
-
denom.dispose();
|
|
183
|
-
|
|
184
|
-
const out = ops.where(isMedium, smoothFreqs, freqs);
|
|
185
|
-
for (const a of [isMedium, smoothFreqs, freqs]) a.dispose();
|
|
186
|
-
return out;
|
|
187
|
-
}
|
|
188
|
-
|
|
189
|
-
/** YarnRoPE (rope_utils.py:128-181, verbatim): returns [freqs, mscale]. */
|
|
190
|
-
function yarnFreqs(
|
|
191
|
-
dims: number,
|
|
192
|
-
base: number,
|
|
193
|
-
scalingFactor: number,
|
|
194
|
-
originalMaxPositionEmbeddings: number,
|
|
195
|
-
betaFast: number,
|
|
196
|
-
betaSlow: number,
|
|
197
|
-
mscale: number,
|
|
198
|
-
mscaleAllDim: number,
|
|
199
|
-
): [MlxArray, number] {
|
|
200
|
-
const findCorrectionDim = (numRotations: number): number =>
|
|
201
|
-
(dims * Math.log(originalMaxPositionEmbeddings / (numRotations * 2 * Math.PI))) /
|
|
202
|
-
(2 * Math.log(base));
|
|
203
|
-
let low = Math.floor(findCorrectionDim(betaFast));
|
|
204
|
-
let high = Math.ceil(findCorrectionDim(betaSlow));
|
|
205
|
-
low = Math.max(low, 0);
|
|
206
|
-
high = Math.min(high, dims - 1);
|
|
207
|
-
const getMscale = (scale: number, m: number): number =>
|
|
208
|
-
scale <= 1 ? 1.0 : 0.1 * m * Math.log(scale) + 1.0;
|
|
209
|
-
const outMscale =
|
|
210
|
-
getMscale(scalingFactor, mscale) / getMscale(scalingFactor, mscaleAllDim);
|
|
211
|
-
|
|
212
|
-
const freqExtra = baseFreqs(dims, base, Dtype.float32);
|
|
213
|
-
const freqInter = ops.mulScalar(freqExtra, scalingFactor);
|
|
214
|
-
|
|
215
|
-
// freq_mask = 1 - linear_ramp_mask(low, high, dims // 2)
|
|
216
|
-
let maxVal = high;
|
|
217
|
-
if (low === maxVal) maxVal += 0.001; // prevent singularity (oracle comment)
|
|
218
|
-
const arangeHalf = ops.arange(0, Math.floor(dims / 2), 1, Dtype.float32);
|
|
219
|
-
const minScalar = ops.scalarLike(low, arangeHalf);
|
|
220
|
-
const shifted = ops.sub(arangeHalf, minScalar);
|
|
221
|
-
arangeHalf.dispose();
|
|
222
|
-
minScalar.dispose();
|
|
223
|
-
const rangeScalar = ops.scalarLike(maxVal - low, shifted);
|
|
224
|
-
const linear = ops.div(shifted, rangeScalar);
|
|
225
|
-
shifted.dispose();
|
|
226
|
-
rangeScalar.dispose();
|
|
227
|
-
const lo = ops.scalarLike(0, linear);
|
|
228
|
-
const hi = ops.scalarLike(1, linear);
|
|
229
|
-
const ramp = ops.clip(linear, lo, hi);
|
|
230
|
-
for (const a of [linear, lo, hi]) a.dispose();
|
|
231
|
-
const one = ops.scalarLike(1, ramp);
|
|
232
|
-
const freqMask = ops.sub(one, ramp);
|
|
233
|
-
ramp.dispose();
|
|
234
|
-
|
|
235
|
-
// freqs = (inter * extra) / (inter * mask + extra * (1 - mask))
|
|
236
|
-
const numArr = ops.mul(freqInter, freqExtra);
|
|
237
|
-
const interMasked = ops.mul(freqInter, freqMask);
|
|
238
|
-
const oneMinusMask = ops.sub(one, freqMask);
|
|
239
|
-
one.dispose();
|
|
240
|
-
const extraMasked = ops.mul(freqExtra, oneMinusMask);
|
|
241
|
-
oneMinusMask.dispose();
|
|
242
|
-
const den = ops.add(interMasked, extraMasked);
|
|
243
|
-
const freqs = ops.div(numArr, den);
|
|
244
|
-
for (const a of [freqExtra, freqInter, freqMask, numArr, interMasked, extraMasked, den])
|
|
245
|
-
a.dispose();
|
|
246
|
-
return [freqs, outMscale];
|
|
247
|
-
}
|
|
248
|
-
|
|
249
|
-
/** SuScaledRoPE (rope_utils.py:10-71): [freqs, scale]. NOTE the oracle uses
|
|
250
|
-
* long_factor for ALL positions (short_factor is accepted and ignored). */
|
|
251
|
-
function suFreqsAndScale(
|
|
252
|
-
dims: number,
|
|
253
|
-
base: number,
|
|
254
|
-
maxPositionEmbeddings: number,
|
|
255
|
-
originalMaxPositionEmbeddings: number,
|
|
256
|
-
longFactor: number[] | number,
|
|
257
|
-
longMscale?: number,
|
|
258
|
-
): [MlxArray, number] {
|
|
259
|
-
const freqs0 = baseFreqs(dims, base, Dtype.float32);
|
|
260
|
-
const lf = Array.isArray(longFactor)
|
|
261
|
-
? MlxArray.fromFloat32(new Float32Array(longFactor), [longFactor.length])
|
|
262
|
-
: MlxArray.fromFloat32(new Float32Array([longFactor]), []);
|
|
263
|
-
const freqs = ops.mul(lf, freqs0);
|
|
264
|
-
lf.dispose();
|
|
265
|
-
freqs0.dispose();
|
|
266
|
-
|
|
267
|
-
const factor = maxPositionEmbeddings / originalMaxPositionEmbeddings;
|
|
268
|
-
const defaultScale =
|
|
269
|
-
factor <= 1.0
|
|
270
|
-
? 1.0
|
|
271
|
-
: Math.sqrt(1 + Math.log(factor) / Math.log(originalMaxPositionEmbeddings));
|
|
272
|
-
// python: `long_mscale or default` — 0 is falsy there too.
|
|
273
|
-
const scale = longMscale ? longMscale : defaultScale;
|
|
274
|
-
return [freqs, scale];
|
|
275
|
-
}
|
|
276
|
-
|
|
277
|
-
/** ProportionalRoPE (rope_utils.py:199-232): rotate the first rotated_dims,
|
|
278
|
-
* ∞-freq (identity) padding for the rest. */
|
|
279
|
-
function proportionalFreqs(
|
|
280
|
-
dims: number, rotatedDims: number, base: number, factor: number,
|
|
281
|
-
): MlxArray {
|
|
282
|
-
if (rotatedDims > dims) throw new Error("rotated_dims should be smaller than dims");
|
|
283
|
-
const ar = ops.arange(0, rotatedDims, 2, Dtype.float32);
|
|
284
|
-
const dimsScalar = MlxArray.fromFloat32(new Float32Array([dims]), []);
|
|
285
|
-
const exps = ops.div(ar, dimsScalar);
|
|
286
|
-
ar.dispose();
|
|
287
|
-
dimsScalar.dispose();
|
|
288
|
-
const baseScalar = ops.scalarLike(base, exps);
|
|
289
|
-
const powed = ops.pow(baseScalar, exps);
|
|
290
|
-
baseScalar.dispose();
|
|
291
|
-
exps.dispose();
|
|
292
|
-
const scaled = ops.mulScalar(powed, factor);
|
|
293
|
-
powed.dispose();
|
|
294
|
-
const padLen = Math.floor((dims - rotatedDims) / 2);
|
|
295
|
-
if (padLen === 0) return scaled;
|
|
296
|
-
const pad = MlxArray.fromFloat32(
|
|
297
|
-
new Float32Array(padLen).fill(Number.POSITIVE_INFINITY), [padLen],
|
|
298
|
-
);
|
|
299
|
-
const out = ops.concatAxis([scaled, pad], 0);
|
|
300
|
-
scaled.dispose();
|
|
301
|
-
pad.dispose();
|
|
302
|
-
return out;
|
|
303
|
-
}
|
|
304
|
-
|
|
305
|
-
/** Port of rope_utils.initialize_rope — the shared entry every
|
|
306
|
-
* initialize_rope-style arch (llama, qwen2/3, olmo2, granite, …) uses. */
|
|
307
|
-
export function initializeRope(
|
|
308
|
-
dims: number,
|
|
309
|
-
base: number,
|
|
310
|
-
traditional: boolean,
|
|
311
|
-
scalingConfig: RopeScalingConfig | null,
|
|
312
|
-
maxPositionEmbeddings: number | null,
|
|
313
|
-
): UniversalRope {
|
|
314
|
-
const ropeType = scalingConfig
|
|
315
|
-
? String(scalingConfig.type ?? scalingConfig.rope_type ?? "default")
|
|
316
|
-
: "default";
|
|
317
|
-
|
|
318
|
-
if (ropeType === "default" || ropeType === "linear") {
|
|
319
|
-
const scale = ropeType === "linear" ? 1 / requireNum(scalingConfig!, "factor") : 1.0;
|
|
320
|
-
return new UniversalRope(dims, traditional, base, scale, null, 1, false);
|
|
321
|
-
}
|
|
322
|
-
|
|
323
|
-
if (ropeType === "llama3") {
|
|
324
|
-
const freqs = llama3Freqs(dims, base, scalingConfig!);
|
|
325
|
-
return new UniversalRope(dims, traditional, null, 1.0, freqs, 1, false);
|
|
326
|
-
}
|
|
327
|
-
|
|
328
|
-
if (ropeType === "yarn" || ropeType === "deepseek_yarn" || ropeType === "telechat3-yarn") {
|
|
329
|
-
const sc = scalingConfig!;
|
|
330
|
-
const [freqs, mscale] = yarnFreqs(
|
|
331
|
-
dims,
|
|
332
|
-
base,
|
|
333
|
-
requireNum(sc, "factor"),
|
|
334
|
-
num(sc.original_max_position_embeddings, 4096),
|
|
335
|
-
num(sc.beta_fast, 32),
|
|
336
|
-
num(sc.beta_slow, 1),
|
|
337
|
-
num(sc.mscale, 1),
|
|
338
|
-
num(sc.mscale_all_dim, 0),
|
|
339
|
-
);
|
|
340
|
-
// YarnRoPE scales x only when mscale != 1.0 (preScaleAlways=false).
|
|
341
|
-
return new UniversalRope(dims, traditional, null, 1.0, freqs, mscale, false);
|
|
342
|
-
}
|
|
343
|
-
|
|
344
|
-
if (ropeType === "longrope") {
|
|
345
|
-
const sc = scalingConfig!;
|
|
346
|
-
const [freqs, scale] = suFreqsAndScale(
|
|
347
|
-
dims,
|
|
348
|
-
base,
|
|
349
|
-
maxPositionEmbeddings ?? 131072,
|
|
350
|
-
requireNum(sc, "original_max_position_embeddings"),
|
|
351
|
-
(sc.long_factor as number[] | number | undefined) ?? 1.0,
|
|
352
|
-
);
|
|
353
|
-
// SuScaledRoPE hardcodes traditional=False and ALWAYS pre-scales.
|
|
354
|
-
return new UniversalRope(dims, false, null, 1.0, freqs, scale, true);
|
|
355
|
-
}
|
|
356
|
-
|
|
357
|
-
if (ropeType === "proportional") {
|
|
358
|
-
const sc = scalingConfig!;
|
|
359
|
-
const rotatedDims = Math.floor(dims * num(sc.partial_rotary_factor, 1.0));
|
|
360
|
-
const freqs = proportionalFreqs(dims, rotatedDims, base, num(sc.factor, 1.0));
|
|
361
|
-
return new UniversalRope(dims, traditional, null, 1.0, freqs, 1, false);
|
|
362
|
-
}
|
|
363
|
-
|
|
364
|
-
throw new Error(`Unsupported RoPE type ${ropeType}`);
|
|
365
|
-
}
|
|
366
|
-
|
|
367
|
-
/** SuScaledRoPE constructed the phi3.py way (NOT via initialize_rope):
|
|
368
|
-
* original_max_position_embeddings comes from the TOP-LEVEL config, and
|
|
369
|
-
* rope_dim = head_dim * partial_rotary_factor. */
|
|
370
|
-
export function phi3SuRope(
|
|
371
|
-
ropeDim: number,
|
|
372
|
-
base: number,
|
|
373
|
-
maxPositionEmbeddings: number,
|
|
374
|
-
originalMaxPositionEmbeddings: number,
|
|
375
|
-
longFactor: number[] | number,
|
|
376
|
-
): UniversalRope {
|
|
377
|
-
const [freqs, scale] = suFreqsAndScale(
|
|
378
|
-
ropeDim, base, maxPositionEmbeddings, originalMaxPositionEmbeddings, longFactor,
|
|
379
|
-
);
|
|
380
|
-
return new UniversalRope(ropeDim, false, null, 1.0, freqs, scale, true);
|
|
381
|
-
}
|