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.
Files changed (601) hide show
  1. package/README.md +1453 -88
  2. package/THIRD_PARTY_NOTICES.md +882 -0
  3. package/bin/mlx-bun.convert.mjs +4 -0
  4. package/bin/mlx-bun.fuse.mjs +4 -0
  5. package/bin/mlx-bun.generate.mjs +4 -0
  6. package/bin/mlx-bun.lora.mjs +4 -0
  7. package/bin/mlx-bun.mjs +13 -21
  8. package/bin/mlx-bun.server.mjs +4 -0
  9. package/bin/mlx-bun.upload.mjs +4 -0
  10. package/dist/web/app.js +9584 -0
  11. package/package.json +50 -53
  12. package/src/cli/ane-install.ts +76 -0
  13. package/src/cli/args.ts +179 -0
  14. package/src/cli/engine-entry.ts +68 -0
  15. package/src/cli/inference.ts +214 -0
  16. package/src/cli/job-entry.ts +97 -0
  17. package/src/cli/library-host.ts +117 -0
  18. package/src/cli/main.ts +57 -0
  19. package/src/cli/memory-engine.ts +163 -0
  20. package/src/cli/memory.ts +36 -0
  21. package/src/cli/mlx-lm-aliases.ts +309 -0
  22. package/src/cli/model-choice.ts +33 -0
  23. package/src/cli/model-selection.ts +84 -0
  24. package/src/cli/module-host.ts +11 -0
  25. package/src/cli/module-verbs.ts +50 -0
  26. package/src/cli/numerical-policy.ts +12 -0
  27. package/src/cli/perplexity.ts +89 -0
  28. package/src/cli/serve-host.ts +423 -0
  29. package/src/cli/serve-isolated.ts +217 -0
  30. package/src/cli/serve-options.ts +106 -0
  31. package/src/cli/serve-state.ts +205 -0
  32. package/src/cli/serve.ts +500 -0
  33. package/src/cli/served-model-host.ts +102 -0
  34. package/src/cli/server-entry.ts +148 -0
  35. package/src/cli/serving-unit.ts +177 -0
  36. package/src/{tui.ts → cli/terminal.ts} +13 -8
  37. package/src/cli/worker-adapters.ts +34 -0
  38. package/src/cli/worker-entry.ts +272 -0
  39. package/src/cli/worker-unit.ts +179 -0
  40. package/src/engine/adapter-operation.ts +26 -0
  41. package/src/engine/cache-services.ts +114 -0
  42. package/src/engine/completion.ts +39 -0
  43. package/src/engine/generation-gateway.ts +358 -0
  44. package/src/engine/index.ts +59 -0
  45. package/src/engine/kv-budget.ts +120 -0
  46. package/src/engine/media-preparation.ts +232 -0
  47. package/src/engine/model-binding.ts +24 -0
  48. package/src/engine/model-host.ts +513 -0
  49. package/src/engine/model-serving.ts +89 -0
  50. package/src/{serve → engine}/preparation.ts +10 -8
  51. package/src/{serve → engine}/session-completion-engine.ts +6 -6
  52. package/src/engine/telemetry.ts +109 -0
  53. package/src/hub/downloads.ts +116 -0
  54. package/src/jobs/db.ts +16 -7
  55. package/src/jobs/events.ts +29 -0
  56. package/src/jobs/executable.ts +5 -0
  57. package/src/jobs/host.ts +51 -0
  58. package/src/jobs/in-process.ts +55 -0
  59. package/src/{contracts/jobs.ts → jobs/protocol.ts} +7 -9
  60. package/src/jobs/runner.ts +150 -231
  61. package/src/jobs/service.ts +75 -0
  62. package/src/jobs/sse.ts +59 -23
  63. package/src/jobs/worker-memory.ts +46 -0
  64. package/src/jobs/worker-process.ts +163 -0
  65. package/src/jobs/worker-supervisor.ts +258 -0
  66. package/src/modules.ts +58 -0
  67. package/src/publishing/catalog-hub.ts +18 -0
  68. package/src/publishing/credentials.ts +31 -0
  69. package/src/publishing/upload.ts +37 -0
  70. package/src/residency/leased-adapters.ts +21 -0
  71. package/src/residency/local-records.ts +16 -0
  72. package/src/residency/model-residency.ts +405 -0
  73. package/src/residency/record-index.ts +28 -0
  74. package/src/residency/resident-estimate.ts +50 -0
  75. package/src/{anthropic.ts → server/anthropic.ts} +3 -81
  76. package/src/server/cache-routes.ts +36 -0
  77. package/src/{serve → server}/chat-request.ts +12 -10
  78. package/src/{serve → server}/chat-stage.ts +27 -22
  79. package/src/{client.ts → server/client.ts} +28 -2
  80. package/src/{serve → server}/completion-executor.ts +42 -42
  81. package/src/server/decision-body.ts +55 -0
  82. package/src/server/decision-routes.ts +134 -0
  83. package/src/server/decision-wire.ts +228 -0
  84. package/src/{serve → server}/discovery-routes.ts +42 -83
  85. package/src/server/generated-token-history.ts +57 -0
  86. package/src/{serve → server}/http.ts +25 -13
  87. package/src/{serve → server}/inference-request.ts +5 -5
  88. package/src/server/job-routes.ts +20 -0
  89. package/src/server/media-prompt.ts +63 -0
  90. package/src/server/memory-completion-client.ts +46 -0
  91. package/src/server/model-routes.ts +192 -0
  92. package/src/{serve → server}/openai-wire.ts +3 -2
  93. package/src/{serve → server}/pipeline.ts +1 -1
  94. package/src/server/prompt-builder.ts +13 -0
  95. package/src/server/prompt-contracts.ts +17 -0
  96. package/src/server/proxy-routes.ts +187 -0
  97. package/src/server/publishing-routes.ts +41 -0
  98. package/src/{serve → server}/request-plan.ts +8 -8
  99. package/src/{serve → server}/request-prep.ts +82 -68
  100. package/src/{serve → server}/responses-client.ts +30 -13
  101. package/src/{responses.ts → server/responses.ts} +12 -93
  102. package/src/server/routes.ts +205 -0
  103. package/src/server/start.ts +72 -0
  104. package/src/server/status-routes.ts +108 -0
  105. package/src/{serve → server}/text-completion-stage.ts +25 -4
  106. package/src/server/text-prompt.ts +12 -0
  107. package/src/{serve → server}/token-streams.ts +36 -38
  108. package/src/server/worker-routes.ts +311 -0
  109. package/src/storage/paths.ts +61 -0
  110. package/src/web/assets.ts +37 -0
  111. package/src/web/{src → browser}/api.ts +5 -23
  112. package/src/web/browser/chat-host.ts +17 -0
  113. package/src/web/browser/installed-panels.ts +6 -0
  114. package/src/web/browser/main.ts +55 -0
  115. package/src/web/browser/markdown.ts +279 -0
  116. package/src/web/browser/memory-panel.ts +9 -0
  117. package/src/web/browser/palette.ts +64 -0
  118. package/src/web/browser/protocol.ts +13 -0
  119. package/src/web/browser/shell.ts +467 -0
  120. package/src/web/{src → browser}/status.ts +32 -36
  121. package/src/web/build.ts +51 -0
  122. package/src/web/public/app.html +1054 -0
  123. package/src/web/{sw.js → public/sw.js} +4 -3
  124. package/src/web/{vendor → public/vendor}/hljs-theme.css +2 -2
  125. package/src/web/{vendor → public/vendor}/hljs.js +2 -2
  126. package/THIRD_PARTY_LICENSES.md +0 -51
  127. package/docs/README.md +0 -94
  128. package/docs/archive/hlg-sampling.md +0 -350
  129. package/docs/archive/investigations/audio-input-plan.md +0 -356
  130. package/docs/archive/investigations/chunk-finetune-experiment.md +0 -84
  131. package/docs/archive/investigations/colibri-glm52-port.md +0 -1486
  132. package/docs/archive/investigations/colibri-oracle-pin.md +0 -1284
  133. package/docs/archive/investigations/curve-bisector-routing.md +0 -123
  134. package/docs/archive/investigations/decode-roofline-lookagain.md +0 -348
  135. package/docs/archive/investigations/diffusion-gemma-port.md +0 -602
  136. package/docs/archive/investigations/dspark-handoff.md +0 -224
  137. package/docs/archive/investigations/expert-offload-single-user-moe.md +0 -292
  138. package/docs/archive/investigations/faithful-l1-consolidation.md +0 -310
  139. package/docs/archive/investigations/hlg-sampling-investigation.md +0 -970
  140. package/docs/archive/investigations/indexshare-performance-spike.md +0 -357
  141. package/docs/archive/investigations/kernel-perf-review-2026-07.md +0 -171
  142. package/docs/archive/investigations/lab-build-journal.md +0 -157
  143. package/docs/archive/investigations/minicpm5-decode-megakernel.md +0 -489
  144. package/docs/archive/investigations/openjarvis-fit.md +0 -198
  145. package/docs/archive/investigations/openwiki-evaluation.md +0 -249
  146. package/docs/archive/investigations/optimization_plan.md +0 -346
  147. package/docs/archive/investigations/orpo-base-uf-experiment-and-directions.md +0 -120
  148. package/docs/archive/investigations/orpo-flash-cce-pin-leak.md +0 -167
  149. package/docs/archive/investigations/orpo-uf-testing-handoff.md +0 -126
  150. package/docs/archive/investigations/pareto-specialized-runtime-findings.md +0 -92
  151. package/docs/archive/investigations/pi-builtin-investigation.md +0 -308
  152. package/docs/archive/investigations/pre-colibri-stabilization.md +0 -329
  153. package/docs/archive/investigations/qwen38-27b-campaign-log.md +0 -1092
  154. package/docs/archive/investigations/segmented-backward-handoff.md +0 -148
  155. package/docs/archive/investigations/starter-model-port-handoff.md +0 -342
  156. package/docs/archive/investigations/steel-flash-cce-handoff.md +0 -489
  157. package/docs/archive/investigations/the-dreaming-handoff.md +0 -141
  158. package/docs/archive/investigations/trainer-validation-experiment.md +0 -173
  159. package/docs/archive/planning/chat-ui-vision-memory-first.md +0 -199
  160. package/docs/archive/planning/chat-ui-vision.md +0 -50
  161. package/docs/archive/planning/chunk-task-roadmap.md +0 -77
  162. package/docs/archive/planning/curve-sampler-research-plan.md +0 -331
  163. package/docs/archive/planning/journal.md +0 -316
  164. package/docs/archive/planning/memory-docs-and-dag-plan.md +0 -194
  165. package/docs/archive/planning/pi-integration-review.md +0 -178
  166. package/docs/archive/planning/release-notes-v0.0.10.md +0 -57
  167. package/docs/archive/planning/release-notes-v0.0.11.md +0 -239
  168. package/docs/archive/planning/release-notes-v0.0.12.md +0 -80
  169. package/docs/archive/planning/release-notes-v0.0.13.md +0 -77
  170. package/docs/archive/planning/release-notes-v0.0.9.md +0 -122
  171. package/docs/archive/planning/release-notes-v0.2.0.md +0 -49
  172. package/docs/archive/planning/release-notes-v0.3.0.md +0 -35
  173. package/docs/archive/planning/web-ui-pass-plan.md +0 -213
  174. package/docs/archive/planning/website-readme-pass-plan.md +0 -224
  175. package/docs/design/batching.md +0 -1841
  176. package/docs/design/decode-speed-program.md +0 -2349
  177. package/docs/design/dreaming-nightly-pipeline.md +0 -684
  178. package/docs/design/generic-model-support.md +0 -705
  179. package/docs/design/kv-cache.md +0 -824
  180. package/docs/design/orpo-training.md +0 -680
  181. package/docs/design/speculative-decoding.md +0 -1990
  182. package/docs/design/turboquant.md +0 -2212
  183. package/docs/design/unified-engine-frontier-plan.md +0 -1880
  184. package/docs/design/web-chat-redesign.md +0 -695
  185. package/docs/planning/PRODUCT_ROADMAP.md +0 -594
  186. package/docs/planning/ResearchTopics.md +0 -119
  187. package/docs/planning/release-notes-v0.4.0.md +0 -45
  188. package/docs/planning/video-series.md +0 -421
  189. package/docs/reference/benchmarks.md +0 -3613
  190. package/docs/reference/cli.md +0 -585
  191. package/docs/reference/distribution.md +0 -355
  192. package/docs/reference/environment.md +0 -398
  193. package/docs/reference/glossary.md +0 -39
  194. package/docs/reference/library-api.md +0 -674
  195. package/docs/reference/memory.md +0 -260
  196. package/docs/reference/models.md +0 -334
  197. package/docs/reference/server-api.md +0 -1775
  198. package/docs/reference/server-config.md +0 -846
  199. package/docs/reference/training.md +0 -462
  200. package/docs/reference/troubleshooting.md +0 -158
  201. package/src/admission-plan.ts +0 -18
  202. package/src/audio/conformer.ts +0 -647
  203. package/src/audio/decode.ts +0 -133
  204. package/src/audio/features.ts +0 -202
  205. package/src/audio/transcode.ts +0 -67
  206. package/src/backends/mlx/assistant-target.ts +0 -46
  207. package/src/backends/mlx/autoregressive-draft-rows.ts +0 -120
  208. package/src/backends/mlx/autoregressive.ts +0 -108
  209. package/src/backends/mlx/batch-group.ts +0 -1790
  210. package/src/backends/mlx/cache-attachment.ts +0 -67
  211. package/src/backends/mlx/cache-layout.ts +0 -31
  212. package/src/backends/mlx/checkpoint-state.ts +0 -53
  213. package/src/backends/mlx/continuation-persistence.ts +0 -106
  214. package/src/backends/mlx/continuation-request.ts +0 -43
  215. package/src/backends/mlx/continuation.ts +0 -77
  216. package/src/backends/mlx/diffusion.ts +0 -39
  217. package/src/backends/mlx/draft-target.ts +0 -32
  218. package/src/backends/mlx/evaluation.ts +0 -105
  219. package/src/backends/mlx/fill-append.ts +0 -45
  220. package/src/backends/mlx/full-kv-row-append.ts +0 -32
  221. package/src/backends/mlx/full-kv-row-donor.ts +0 -54
  222. package/src/backends/mlx/full-transitioning-kv-rows.ts +0 -65
  223. package/src/backends/mlx/gateway-binding.ts +0 -130
  224. package/src/backends/mlx/graph.ts +0 -38
  225. package/src/backends/mlx/kv-codec.ts +0 -9
  226. package/src/backends/mlx/kv-maintenance.ts +0 -133
  227. package/src/backends/mlx/kv-tensor-rows.ts +0 -212
  228. package/src/backends/mlx/memory-client.ts +0 -189
  229. package/src/backends/mlx/methods.ts +0 -90
  230. package/src/backends/mlx/model-prompt.ts +0 -164
  231. package/src/backends/mlx/model-serving.ts +0 -79
  232. package/src/backends/mlx/prefill-cohort.ts +0 -89
  233. package/src/backends/mlx/prefill-rows.ts +0 -200
  234. package/src/backends/mlx/prefill.ts +0 -33
  235. package/src/backends/mlx/projected-context-rows.ts +0 -72
  236. package/src/backends/mlx/request-policy.ts +0 -20
  237. package/src/backends/mlx/request-state-policy.ts +0 -41
  238. package/src/backends/mlx/rollback.ts +0 -89
  239. package/src/backends/mlx/row-sampling.ts +0 -49
  240. package/src/backends/mlx/serial-executor.ts +0 -242
  241. package/src/backends/mlx/speculative-group.ts +0 -359
  242. package/src/backends/mlx/speculative-round.ts +0 -161
  243. package/src/backends/mlx/speculative.ts +0 -114
  244. package/src/backends/mlx/state-rows.ts +0 -104
  245. package/src/backends/mlx/state-views.ts +0 -46
  246. package/src/backends/mlx/target-layout-capability.ts +0 -35
  247. package/src/backends/mlx/text-engine.ts +0 -73
  248. package/src/backends/mlx/transitioning-kv-rows.ts +0 -165
  249. package/src/chat-template.ts +0 -382
  250. package/src/cli.ts +0 -3168
  251. package/src/config.ts +0 -399
  252. package/src/contracts/completion.ts +0 -16
  253. package/src/contracts/execution.ts +0 -60
  254. package/src/contracts/generation.ts +0 -103
  255. package/src/contracts/host.ts +0 -7
  256. package/src/contracts/pi-provider.ts +0 -30
  257. package/src/contracts/pi.ts +0 -259
  258. package/src/contracts/prefix-cache.ts +0 -23
  259. package/src/contracts/resources.ts +0 -25
  260. package/src/contracts/sampling.ts +0 -23
  261. package/src/contracts/scheduling.ts +0 -45
  262. package/src/contracts/tsconfig.json +0 -14
  263. package/src/dataset/generators.ts +0 -1087
  264. package/src/dataset/index.ts +0 -12
  265. package/src/dataset/job.ts +0 -81
  266. package/src/dataset/llm.ts +0 -71
  267. package/src/dataset/registry.ts +0 -528
  268. package/src/diffusion/diffusion-generate.ts +0 -507
  269. package/src/download.ts +0 -505
  270. package/src/embed.ts +0 -77
  271. package/src/engine/admission.ts +0 -59
  272. package/src/engine/cancellation.ts +0 -38
  273. package/src/engine/engine.ts +0 -56
  274. package/src/engine/execution-coordinator.ts +0 -70
  275. package/src/engine/execution-plan.ts +0 -53
  276. package/src/engine/resources.ts +0 -53
  277. package/src/engine/scheduler.ts +0 -72
  278. package/src/engine/session.ts +0 -247
  279. package/src/eval/README.md +0 -130
  280. package/src/eval/ifeval.ts +0 -81
  281. package/src/eval/kl-prompts.ts +0 -29
  282. package/src/eval/kl.ts +0 -590
  283. package/src/eval/perplexity.ts +0 -164
  284. package/src/eval/quality-db.ts +0 -92
  285. package/src/eval/reference-logits.ts +0 -207
  286. package/src/eval/runner.ts +0 -124
  287. package/src/eval/score.ts +0 -49
  288. package/src/eval/tasks/bfcl.ts +0 -535
  289. package/src/eval/tasks/gsm8k.ts +0 -98
  290. package/src/eval/tasks/hashhop.ts +0 -139
  291. package/src/eval/tasks/humaneval.ts +0 -275
  292. package/src/eval/tasks/ifeval.ts +0 -525
  293. package/src/eval/tasks/mmlu.ts +0 -157
  294. package/src/eval/ultrafeedback.ts +0 -60
  295. package/src/evaldb.ts +0 -199
  296. package/src/expert-io.ts +0 -393
  297. package/src/expert-offload-build.ts +0 -77
  298. package/src/expert-offload.ts +0 -76
  299. package/src/expert-residency.ts +0 -1134
  300. package/src/expert-usage.ts +0 -415
  301. package/src/fill/echo-index.ts +0 -188
  302. package/src/fill/fill-session.ts +0 -452
  303. package/src/fill/proposal.ts +0 -68
  304. package/src/fill/schema-rows.ts +0 -506
  305. package/src/fill/tool-boundary.ts +0 -103
  306. package/src/fit.ts +0 -249
  307. package/src/generate.ts +0 -1346
  308. package/src/grammar.ts +0 -591
  309. package/src/harness-pi.ts +0 -184
  310. package/src/hf-push.ts +0 -628
  311. package/src/hub-rest.ts +0 -242
  312. package/src/index.ts +0 -63
  313. package/src/inference/denoising.ts +0 -17
  314. package/src/inference/draft-acceptance.ts +0 -41
  315. package/src/inference/draft-output.ts +0 -49
  316. package/src/inference/graph.ts +0 -33
  317. package/src/inference/prefill.ts +0 -39
  318. package/src/inference/rollback.ts +0 -7
  319. package/src/jobs/index.ts +0 -31
  320. package/src/jobs/job-entry.ts +0 -107
  321. package/src/jobs/task-client.ts +0 -19
  322. package/src/jobs/types.ts +0 -2
  323. package/src/kv-scheme.ts +0 -257
  324. package/src/kv-store.ts +0 -1099
  325. package/src/lab/README.md +0 -18
  326. package/src/lab/curve/curve-designer.html +0 -303
  327. package/src/lab/curve/curve-sampler.ts +0 -122
  328. package/src/lab/expert-trace/expert-trace.ts +0 -69
  329. package/src/lab/paged-kv/paged-kv-rows.ts +0 -83
  330. package/src/lab/paged-kv/paged-kv.ts +0 -270
  331. package/src/library.ts +0 -40
  332. package/src/lora.ts +0 -408
  333. package/src/md-modules.d.ts +0 -49
  334. package/src/media-fetch.ts +0 -442
  335. package/src/memory/article.ts +0 -558
  336. package/src/memory/chunk-validation.ts +0 -150
  337. package/src/memory/chunk.ts +0 -318
  338. package/src/memory/cluster.ts +0 -351
  339. package/src/memory/coverage.ts +0 -134
  340. package/src/memory/crosslink.ts +0 -0
  341. package/src/memory/db.ts +0 -253
  342. package/src/memory/entity.ts +0 -203
  343. package/src/memory/gate.ts +0 -187
  344. package/src/memory/history.ts +0 -152
  345. package/src/memory/ingest.ts +0 -335
  346. package/src/memory/model.ts +0 -183
  347. package/src/memory/normalize.ts +0 -671
  348. package/src/memory/parse.ts +0 -22
  349. package/src/memory/pipeline.ts +0 -275
  350. package/src/memory/prompts.ts +0 -65
  351. package/src/memory/query.ts +0 -334
  352. package/src/memory/reconcile.ts +0 -734
  353. package/src/memory/redirect.ts +0 -108
  354. package/src/memory/reindex.ts +0 -325
  355. package/src/memory/resolve.ts +0 -307
  356. package/src/memory/rest.ts +0 -245
  357. package/src/memory/route.ts +0 -272
  358. package/src/memory/schedule.ts +0 -165
  359. package/src/memory/stages.ts +0 -777
  360. package/src/memory/synthesize.ts +0 -1638
  361. package/src/memory/tools.ts +0 -749
  362. package/src/memory/vault.ts +0 -1194
  363. package/src/memory/wikify.ts +0 -1064
  364. package/src/mlx/array.ts +0 -388
  365. package/src/mlx/autograd.ts +0 -304
  366. package/src/mlx/checkpoint.ts +0 -149
  367. package/src/mlx/compile.ts +0 -168
  368. package/src/mlx/custom-vjp.ts +0 -188
  369. package/src/mlx/ffi.ts +0 -512
  370. package/src/mlx/materialize.ts +0 -23
  371. package/src/mlx/metal-kernel.ts +0 -159
  372. package/src/mlx/ops.ts +0 -1021
  373. package/src/mlx/quantized-rows.ts +0 -36
  374. package/src/mlx/token-bitmask.ts +0 -52
  375. package/src/mlx/turboquant-kv-decode.ts +0 -82
  376. package/src/mlx/turboquant-kv-inverse.ts +0 -122
  377. package/src/mlx/turboquant-ops.ts +0 -861
  378. package/src/mlx/turboquant-tables.ts +0 -237
  379. package/src/mmap.ts +0 -90
  380. package/src/model/artifact-identity.ts +0 -73
  381. package/src/model/batched-kv.ts +0 -235
  382. package/src/model/batched-mask.ts +0 -279
  383. package/src/model/batched-quant.ts +0 -212
  384. package/src/model/batched-quantized-kv.ts +0 -144
  385. package/src/model/batched-rotating-quant.ts +0 -353
  386. package/src/model/batched-rotating-state.ts +0 -149
  387. package/src/model/batched-rotating.ts +0 -338
  388. package/src/model/batched-row-storage.ts +0 -179
  389. package/src/model/batched-ssm.ts +0 -52
  390. package/src/model/batched-turboquant-kv.ts +0 -93
  391. package/src/model/compiled-decode.ts +0 -712
  392. package/src/model/decoded-kv-donor.ts +0 -12
  393. package/src/model/delayed-quantized-kv.ts +0 -72
  394. package/src/model/delayed-rotating-quantized-kv.ts +0 -93
  395. package/src/model/delayed-turboquant-kv.ts +0 -55
  396. package/src/model/diffusion-gemma.ts +0 -1026
  397. package/src/model/factory.ts +0 -189
  398. package/src/model/fingerprint.ts +0 -33
  399. package/src/model/flash-attention.ts +0 -694
  400. package/src/model/full-prefill-padding.ts +0 -56
  401. package/src/model/full-prefill-row.ts +0 -105
  402. package/src/model/gemma4-base.ts +0 -2351
  403. package/src/model/gemma4.ts +0 -1315
  404. package/src/model/generated/gemma4-12b.ts +0 -429
  405. package/src/model/generated/gemma4-26b.ts +0 -448
  406. package/src/model/generated/gemma4-e4b.ts +0 -717
  407. package/src/model/generated/index.ts +0 -18
  408. package/src/model/glm52-atlas.ts +0 -511
  409. package/src/model/glm52-cache.ts +0 -792
  410. package/src/model/glm52-config.ts +0 -262
  411. package/src/model/glm52-container.ts +0 -305
  412. package/src/model/glm52-coupling.ts +0 -388
  413. package/src/model/glm52-dsa.ts +0 -606
  414. package/src/model/glm52-expert-layout.ts +0 -253
  415. package/src/model/glm52-layout.ts +0 -152
  416. package/src/model/glm52-memory.ts +0 -500
  417. package/src/model/glm52-mla.ts +0 -792
  418. package/src/model/glm52-moe.ts +0 -223
  419. package/src/model/glm52-pilot.ts +0 -537
  420. package/src/model/glm52-reference.ts +0 -445
  421. package/src/model/glm52-residency.ts +0 -557
  422. package/src/model/glm52-resident-weights.ts +0 -135
  423. package/src/model/glm52-streamed-experts.ts +0 -453
  424. package/src/model/glm52-streamed-metal.ts +0 -610
  425. package/src/model/glm52-weights.ts +0 -503
  426. package/src/model/glm52.ts +0 -1132
  427. package/src/model/implementation.ts +0 -69
  428. package/src/model/kv-attention-view.ts +0 -67
  429. package/src/model/minicpm5.ts +0 -376
  430. package/src/model/profile.ts +0 -443
  431. package/src/model/qwen-append.ts +0 -12
  432. package/src/model/qwen-conv.ts +0 -52
  433. package/src/model/qwen3-delta-state.ts +0 -74
  434. package/src/model/qwen3-delta.ts +0 -451
  435. package/src/model/qwen3-moe.ts +0 -404
  436. package/src/model/qwen3-mrope.ts +0 -252
  437. package/src/model/qwen3.ts +0 -277
  438. package/src/model/qwen3_5-checkpoint.ts +0 -164
  439. package/src/model/qwen3_5.ts +0 -819
  440. package/src/model/rotating-kv-donor.ts +0 -35
  441. package/src/model/rotating-kv-layout.ts +0 -262
  442. package/src/model/rotating-row-transaction.ts +0 -35
  443. package/src/model/speculative-rotating-kv.ts +0 -101
  444. package/src/model/ssm-prefill-padding.ts +0 -39
  445. package/src/model/support.ts +0 -102
  446. package/src/model/trellis-balanced-scatter.ts +0 -51
  447. package/src/model/trellis-linear.ts +0 -672
  448. package/src/model/trellis-shared-m.ts +0 -151
  449. package/src/model/trellis-shared-scatter.ts +0 -144
  450. package/src/model/trellis-splitk-prefill.ts +0 -158
  451. package/src/model/trellis-tiled-prefill.ts +0 -136
  452. package/src/model/trellis-vector-expand.ts +0 -62
  453. package/src/model/trellis-wide-prefill.ts +0 -94
  454. package/src/model/turboquant-codec.ts +0 -83
  455. package/src/model/universal/archs.ts +0 -332
  456. package/src/model/universal/dense.ts +0 -612
  457. package/src/model/universal/modules.ts +0 -170
  458. package/src/model/universal/rope.ts +0 -381
  459. package/src/native/expert_io.c +0 -523
  460. package/src/native/frame_extract.swift +0 -140
  461. package/src/native-pack.ts +0 -162
  462. package/src/pi-provider.ts +0 -53
  463. package/src/pi-session.ts +0 -81
  464. package/src/pi-terminal.ts +0 -274
  465. package/src/pi-web.ts +0 -1932
  466. package/src/preflight.ts +0 -158
  467. package/src/prompt-cache.ts +0 -460
  468. package/src/quantize/allocator.ts +0 -627
  469. package/src/quantize/atomic-output.ts +0 -38
  470. package/src/quantize/calibration.ts +0 -181
  471. package/src/quantize/config-writer.ts +0 -145
  472. package/src/quantize/index.ts +0 -84
  473. package/src/quantize/job.ts +0 -162
  474. package/src/quantize/quantizer.ts +0 -666
  475. package/src/quantize/rotate.ts +0 -603
  476. package/src/quantize/safetensors-writer.ts +0 -249
  477. package/src/quantize/sensitivity.ts +0 -216
  478. package/src/quantize/trellis.ts +0 -586
  479. package/src/quantize/weight-transform.ts +0 -292
  480. package/src/registry.ts +0 -602
  481. package/src/runtime-config.ts +0 -81
  482. package/src/safetensors.ts +0 -158
  483. package/src/sampler.ts +0 -1011
  484. package/src/serve/admin-routes.ts +0 -184
  485. package/src/serve/assets/curve-terrain.html +0 -105
  486. package/src/serve/assets/training-inference-map.html +0 -575
  487. package/src/serve/aux-routes.ts +0 -132
  488. package/src/serve/batch-scheduler.ts +0 -8
  489. package/src/serve/checkpoint-identity.ts +0 -64
  490. package/src/serve/generation-gateway.ts +0 -496
  491. package/src/serve/isolate.ts +0 -609
  492. package/src/serve/kv-budget.ts +0 -14
  493. package/src/serve/lab-routes.ts +0 -341
  494. package/src/serve/lane-registry.ts +0 -45
  495. package/src/serve/model-admin-routes.ts +0 -155
  496. package/src/serve/model-binding.ts +0 -43
  497. package/src/serve/model-host.ts +0 -511
  498. package/src/serve/parent-application.ts +0 -56
  499. package/src/serve/prompt-response-trace.ts +0 -147
  500. package/src/serve/response-lease.ts +0 -24
  501. package/src/serve/session-search.ts +0 -326
  502. package/src/serve/static-routes.ts +0 -101
  503. package/src/server.ts +0 -1314
  504. package/src/spec/assistant-rows.ts +0 -92
  505. package/src/spec/assistant-source.ts +0 -126
  506. package/src/spec/cache-identity.ts +0 -22
  507. package/src/spec/deepspec-rows.ts +0 -29
  508. package/src/spec/deepspec-source.ts +0 -145
  509. package/src/spec/dflash-source.ts +0 -149
  510. package/src/spec/drafter.ts +0 -304
  511. package/src/spec/dspark/ab-stats.ts +0 -136
  512. package/src/spec/dspark/calibration.ts +0 -112
  513. package/src/spec/dspark/data-dflash.ts +0 -156
  514. package/src/spec/dspark/data.ts +0 -201
  515. package/src/spec/dspark/deepspec-module.ts +0 -776
  516. package/src/spec/dspark/generate-dflash.ts +0 -245
  517. package/src/spec/dspark/generate.ts +0 -290
  518. package/src/spec/dspark/loader.ts +0 -23
  519. package/src/spec/dspark/loss.ts +0 -213
  520. package/src/spec/dspark/module-dflash.ts +0 -560
  521. package/src/spec/dspark/module.ts +0 -540
  522. package/src/spec/dspark/quantize-drafter.ts +0 -91
  523. package/src/spec/dspark/sample.ts +0 -112
  524. package/src/spec/generate.ts +0 -276
  525. package/src/spec/glm52-mtp-graph.ts +0 -29
  526. package/src/spec/glm52-mtp-rows.ts +0 -132
  527. package/src/spec/glm52-mtp-source.ts +0 -317
  528. package/src/spec/glm52-mtp-state.ts +0 -33
  529. package/src/spec/ngram-source.ts +0 -204
  530. package/src/spec/projected-draft-rows.ts +0 -73
  531. package/src/spec/qwen-mtp-module.ts +0 -171
  532. package/src/spec/qwen-mtp-rows.ts +0 -219
  533. package/src/spec/qwen-mtp-source.ts +0 -383
  534. package/src/spec/qwen-mtp-state.ts +0 -30
  535. package/src/spec/serve-loop.ts +0 -585
  536. package/src/spec/source.ts +0 -215
  537. package/src/spec/two-model-rows.ts +0 -61
  538. package/src/spec/two-model.ts +0 -113
  539. package/src/ssd-cache.ts +0 -372
  540. package/src/ssd-durability.ts +0 -211
  541. package/src/tokenizer.ts +0 -335
  542. package/src/tool-approvals.ts +0 -131
  543. package/src/tool-call.ts +0 -695
  544. package/src/train/dataset.ts +0 -484
  545. package/src/train/diffusion-lora.ts +0 -216
  546. package/src/train/export.ts +0 -32
  547. package/src/train/flash-cce.ts +0 -1234
  548. package/src/train/forward.ts +0 -266
  549. package/src/train/fuse.ts +0 -243
  550. package/src/train/index.ts +0 -46
  551. package/src/train/job.ts +0 -205
  552. package/src/train/lora-params.ts +0 -254
  553. package/src/train/loss.ts +0 -1591
  554. package/src/train/merge.ts +0 -195
  555. package/src/train/optimizer.ts +0 -183
  556. package/src/train/prefix-shared.ts +0 -490
  557. package/src/train/rank.ts +0 -137
  558. package/src/train/segmented.ts +0 -1828
  559. package/src/train/steel-qmm-header.ts +0 -1461
  560. package/src/train/trainer.ts +0 -1221
  561. package/src/train/watch.ts +0 -542
  562. package/src/vision/diffusion-vision.ts +0 -620
  563. package/src/vision/embedder.ts +0 -168
  564. package/src/vision/preprocess.ts +0 -247
  565. package/src/vision/prompt.ts +0 -373
  566. package/src/vision/qwen3vl-preprocess.ts +0 -409
  567. package/src/vision/qwen3vl-prompt.ts +0 -196
  568. package/src/vision/qwen3vl-tower.ts +0 -555
  569. package/src/vision/siglip.ts +0 -599
  570. package/src/vision/video-frames.ts +0 -148
  571. package/src/web/app.html +0 -2143
  572. package/src/web/app.js +0 -5703
  573. package/src/web/skills/memory/SKILL.md +0 -120
  574. package/src/web/skills/web-research/SKILL.md +0 -46
  575. package/src/web/skills.ts +0 -97
  576. package/src/web/src/adapters-panel.ts +0 -288
  577. package/src/web/src/assistant.ts +0 -451
  578. package/src/web/src/chat.ts +0 -1057
  579. package/src/web/src/composer.ts +0 -959
  580. package/src/web/src/dataset.ts +0 -132
  581. package/src/web/src/finetune.ts +0 -267
  582. package/src/web/src/hub.ts +0 -428
  583. package/src/web/src/main.ts +0 -98
  584. package/src/web/src/markdown.ts +0 -0
  585. package/src/web/src/memory-panel.ts +0 -572
  586. package/src/web/src/model-picker.ts +0 -159
  587. package/src/web/src/palette.ts +0 -346
  588. package/src/web/src/protocol.ts +0 -26
  589. package/src/web/src/quantize.ts +0 -167
  590. package/src/web/src/rag.ts +0 -274
  591. package/src/web/src/sessions.ts +0 -469
  592. package/src/web/src/shell.ts +0 -804
  593. package/src/web/src/ui-catalog.ts +0 -78
  594. package/src/web/vendor/README.md +0 -61
  595. package/src/web-assets.ts +0 -47
  596. package/src/web-tools.ts +0 -478
  597. package/src/weights.ts +0 -255
  598. /package/src/{serve → server}/completion-sink.ts +0 -0
  599. /package/src/web/{icon.svg → public/icon.svg} +0 -0
  600. /package/src/web/{manifest.webmanifest → public/manifest.webmanifest} +0 -0
  601. /package/src/web/{vendor → public/vendor}/hljs-LICENSE +0 -0
@@ -1,266 +0,0 @@
1
- // Training forward pass: full-sequence logits with no KV reuse.
2
- //
3
- // B=1 (no padding): trainForward runs the model's stock forwardHidden over a
4
- // fresh cache exactly as before — bit-identical to the single-example trainer.
5
- //
6
- // B>1 (padded batch): every row in the batch is padded to a common length L,
7
- // but the rows have different *true* lengths. A plain causal mask would let a
8
- // real query attend to padded keys (and let padded queries pollute the cache
9
- // view), so for batched forwards we build an explicit [B, 1, L, L] boolean
10
- // mask that is the causal mask AND-ed with a per-row "key is within this row's
11
- // valid length" mask, and route it through the model's attention via a cache
12
- // wrapper whose makeMask returns it. The wrapper delegates every KV operation
13
- // to a real cache (KVCache/RotatingKVCache are already shape-generic over B),
14
- // so only the mask differs from the stock path. Window/offset handling is
15
- // preserved: training always runs offset-0 full-sequence forwards, so the
16
- // causal-with-window mask is built at offset 0, exactly matching what the
17
- // underlying cache's makeMask would have produced for the unpadded case.
18
-
19
- import { MlxArray } from "../mlx/array";
20
- import { Dtype } from "../mlx/ffi";
21
- import * as ops from "../mlx/ops";
22
- import { createCausalMask, type Cache, type Mask } from "../model/gemma4-base";
23
- import { SSMCache } from "../model/qwen3-delta";
24
- import type { RuntimeModel } from "../model/factory";
25
-
26
- /** Stateless cache for the training forward. Training is always a single
27
- * offset-0 full-sequence pass, so `updateAndFetch` is a pure pass-through (no
28
- * buffer, no state mutation) and `makeMask` is the offset-0 causal/windowed
29
- * mask — identical to KVCache at offset 0. The statelessness is REQUIRED for
30
- * gradient checkpointing: mlx_checkpoint re-runs the layer closure during the
31
- * backward recompute, and a stateful (appending) cache would corrupt on the
32
- * second run. */
33
- export class TrainingCache implements Cache {
34
- /** Training-only adapter; never admitted, merged, or persisted. */
35
- signature(): string { return "train:training"; }
36
- offset = 0;
37
- updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
38
- // Return fresh views: the caller disposes its input k/v right after (its
39
- // contract with the real KVCache, which returns buffer slices). A
40
- // full-range slice is a cheap view node; mlx keeps the source alive
41
- // through it, so disposing the inputs is safe.
42
- return [k.slice([0, 0, 0, 0], k.shape), v.slice([0, 0, 0, 0], v.shape)];
43
- }
44
- makeMask(N: number, windowSize: number | null): Mask {
45
- if (N === 1) return { mode: "", arr: null };
46
- if (windowSize === null || N <= windowSize) return { mode: "causal", arr: null };
47
- return { mode: "array", arr: createCausalMask(N, 0, windowSize) };
48
- }
49
- state(): MlxArray[] {
50
- return [];
51
- }
52
- isTrimmable(): boolean {
53
- return true;
54
- }
55
- trim(_n: number): void {
56
- /* offset is pinned at 0 */
57
- }
58
- dispose(): void {
59
- /* owns no arrays */
60
- }
61
- }
62
-
63
- /** Stateless DeltaNet stand-in (same re-runnability contract as
64
- * TrainingCache): state WRITES are discarded immediately, so a gradient-
65
- * checkpoint recompute re-enters with the identical null-state conditions;
66
- * reads always yield null (zeros inside the kernel) — a full-sequence
67
- * forward from t=0, exactly the training/perplexity semantics. Before this
68
- * existed, qwen3_5 routed its DeltaNet layers into plain TrainingCache and
69
- * threw on the missing conv/recurrent/advance surface (the recorded
70
- * "mlx-bun perplexity cannot score qwen3_5" gap). */
71
- export class TrainingSSMCache implements Cache {
72
- /** Training-only adapter; never admitted, merged, or persisted. */
73
- signature(): string { return "train:training-ssm"; }
74
- offset = 0;
75
- specRound: null = null;
76
- get conv(): MlxArray | null { return null; }
77
- set conv(v: MlxArray | null) { v?.dispose(); }
78
- get recurrent(): MlxArray | null { return null; }
79
- set recurrent(v: MlxArray | null) { v?.dispose(); }
80
- advance(_n: number): void { /* offset pinned at 0 */ }
81
- rowOffset(_i: number): number { return 0; }
82
- updateAndFetch(): [MlxArray, MlxArray] {
83
- throw new Error("TrainingSSMCache: DeltaNet layers do not use the KV path");
84
- }
85
- makeMask(_N: number, _w: number | null): Mask {
86
- return { mode: "", arr: null }; // ssm_mask is None at B=1 (SSMCache parity)
87
- }
88
- state(): MlxArray[] { return []; }
89
- isTrimmable(): boolean { return true; }
90
- trim(_n: number): void { /* stateless */ }
91
- dispose(): void { /* owns no arrays */ }
92
- }
93
-
94
- /** Run a full-sequence forward for training.
95
- * @param ids int32 array [B, L].
96
- * @returns logits [B, L, V] (caller owns; dispose when done).
97
- *
98
- * When `validLengths` is omitted (or every row is full-length L) the batch
99
- * has no padding and the model's stock mask path is used unchanged — the
100
- * B=1 result is bit-identical to the original trainer. When some row is
101
- * shorter than L, a padding-aware batched mask is injected. */
102
- export function trainForwardHidden(
103
- model: RuntimeModel,
104
- ids: MlxArray,
105
- validLengths?: number[],
106
- ): MlxArray {
107
- const [B, L] = ids.shape as [number, number];
108
- const needsPadMask =
109
- validLengths !== undefined && validLengths.some((v) => v < L);
110
-
111
- if (needsPadMask) {
112
- // B>1 padded path: dense cache wrapped with the per-row padding mask
113
- // (unchanged; not yet checkpointed — B=1 is the checkpointing target).
114
- const realCache = model.makeCache();
115
- const cache = wrapWithBatchedMask(realCache, B, L, validLengths!);
116
- try {
117
- return model.forwardHidden(ids, cache);
118
- } finally {
119
- for (const c of realCache) c.dispose();
120
- for (const c of cache) (c as BatchedMaskCache).disposeOwnMask();
121
- }
122
- }
123
-
124
- // B=1 / no padding: stateless offset-0 caches (one per donor layer), so the
125
- // forward is pure and safe to recompute under gradient checkpointing.
126
- const probe = model.makeCache();
127
- const cache: Cache[] = probe.map((c) =>
128
- c instanceof SSMCache ? new TrainingSSMCache() : new TrainingCache());
129
- for (const c of probe) c.dispose();
130
- try {
131
- return model.forwardHidden(ids, cache); // [B, L, hidden]
132
- } finally {
133
- for (const c of cache) c.dispose();
134
- }
135
- }
136
-
137
- /** Full-sequence forward returning logits [B, L, V]. Convenience over
138
- * trainForwardHidden + logitsFromHidden. The SFT B=1 path deliberately skips
139
- * this and applies the LM head only at response positions (responseOnlyCe)
140
- * to avoid materializing the [B, L, vocab] logits — the dominant training
141
- * memory term at long context with a 262k vocab. */
142
- export function trainForward(
143
- model: RuntimeModel,
144
- ids: MlxArray,
145
- validLengths?: number[],
146
- ): MlxArray {
147
- const h = trainForwardHidden(model, ids, validLengths);
148
- const logits = model.logitsFromHidden(h);
149
- h.dispose();
150
- return logits;
151
- }
152
-
153
- /** Build the padding-aware additive-style boolean mask for one window type.
154
- * Shape `[B, 1, L, L]`, true where query i may attend to key j:
155
- * allow[b,0,i,j] = (j <= i) (causal, offset 0)
156
- * AND (j < validLen[b]) (key within row b's true length)
157
- * [AND (i < j + window) if a sliding window applies]
158
- * Built from createCausalMask(L, 0, window) (the model's own [L, L] bool
159
- * matrix) broadcast over B, masked by a per-row [B,1,1,L] key-validity mask.
160
- * Caller owns the returned array. */
161
- export function buildBatchedPadMask(
162
- B: number,
163
- L: number,
164
- validLengths: number[],
165
- windowSize: number | null,
166
- ): MlxArray {
167
- // [L, L] bool causal (+window) matrix at offset 0 — exactly the model's.
168
- const causal2d = createCausalMask(L, 0, windowSize); // [L, L] bool
169
- const causal4d = ops.reshape(causal2d, [1, 1, L, L]);
170
- causal2d.dispose();
171
-
172
- // Per-row key validity: keyValid[b, 0, 0, j] = j < validLengths[b].
173
- // Build as int32 {0,1} then compare to get bool, broadcasting over q (i).
174
- const keyValidData = new Int32Array(B * L);
175
- for (let b = 0; b < B; b++) {
176
- const vlen = validLengths[b]!;
177
- for (let j = 0; j < L; j++) keyValidData[b * L + j] = j < vlen ? 1 : 0;
178
- }
179
- const keyValidI = MlxArray.fromInt32(keyValidData, [B, 1, 1, L]);
180
- const zero = MlxArray.fromInt32(new Int32Array([0]), []);
181
- const keyValid = ops.less(zero, keyValidI); // 0 < x → bool true where x==1
182
- zero.dispose();
183
- keyValidI.dispose();
184
-
185
- const allow = ops.logicalAnd(causal4d, keyValid); // broadcasts [1,1,L,L] & [B,1,1,L]
186
- causal4d.dispose();
187
- keyValid.dispose();
188
- return allow; // [B, 1, L, L] bool
189
- }
190
-
191
- /** Cache wrapper that forwards every KV op to a real cache but overrides
192
- * makeMask to return a precomputed padding-aware batched mask. One mask is
193
- * built per distinct window size (null = full attention, else the sliding
194
- * window) and shared across all caches of that type; ownership of the mask
195
- * arrays stays with the wrapper (disposeOwnMask), NOT the model — the model's
196
- * forwardLayers only disposes the mask it pulled from makeMask, so makeMask
197
- * hands back a fresh view each call and keeps the canonical copy itself. */
198
- class BatchedMaskCache implements Cache {
199
- /** Training-only adapter; never admitted, merged, or persisted. */
200
- signature(): string { return "train:batched-mask"; }
201
- constructor(
202
- private readonly inner: Cache,
203
- private readonly window: number | null,
204
- private readonly maskStore: Map<number | null, MlxArray>,
205
- ) {}
206
-
207
- get offset(): number {
208
- return this.inner.offset;
209
- }
210
-
211
- updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
212
- return this.inner.updateAndFetch(k, v);
213
- }
214
-
215
- makeMask(N: number, _windowSize: number | null): Mask {
216
- const canonical = this.maskStore.get(this.window)!;
217
- // Hand back a fresh full-range view so the model can dispose it without
218
- // touching our canonical copy. slice over all 4 axes = cheap view node.
219
- const [B, , L] = canonical.shape as [number, number, number, number];
220
- const view = canonical.slice([0, 0, 0, 0], [B, 1, L, L]);
221
- return { mode: "array", arr: view };
222
- }
223
-
224
- state(): MlxArray[] {
225
- return this.inner.state();
226
- }
227
- isTrimmable(): boolean {
228
- return this.inner.isTrimmable();
229
- }
230
- trim(n: number): void {
231
- this.inner.trim(n);
232
- }
233
- dispose(): void {
234
- this.inner.dispose();
235
- }
236
- /** Dispose the canonical mask copies (shared store; idempotent via clear). */
237
- disposeOwnMask(): void {
238
- for (const m of this.maskStore.values()) m.dispose();
239
- this.maskStore.clear();
240
- }
241
- }
242
-
243
- /** Wrap each real cache in a BatchedMaskCache, precomputing one padding-aware
244
- * mask per distinct window. Caches without a sliding window get the
245
- * full-attention (null) mask. */
246
- function wrapWithBatchedMask(
247
- realCache: Cache[],
248
- B: number,
249
- L: number,
250
- validLengths: number[],
251
- ): Cache[] {
252
- const windows = realCache.map((c) => cacheWindow(c));
253
- const distinct = new Set(windows);
254
- const store = new Map<number | null, MlxArray>();
255
- for (const w of distinct) store.set(w, buildBatchedPadMask(B, L, validLengths, w));
256
- return realCache.map((c, i) => new BatchedMaskCache(c, windows[i]!, store));
257
- }
258
-
259
- /** The sliding-window size a cache enforces, or null for full attention.
260
- * RotatingKVCache stores its window as `maxSize`; KVCache has none. Mirrors
261
- * gemma4 forwardLayers, where sliding layers use windowSize == the rotating
262
- * cache's maxSize. (Quantized KV caches are not used on the training path.) */
263
- function cacheWindow(c: Cache): number | null {
264
- const maxSize = (c as { maxSize?: number }).maxSize;
265
- return typeof maxSize === "number" ? maxSize : null;
266
- }
package/src/train/fuse.ts DELETED
@@ -1,243 +0,0 @@
1
- // Fuse a LoRA adapter into base model weights → a standalone model snapshot
2
- // (mlx_lm.fuse counterpart, minus GGUF/de-quantize export).
3
- //
4
- // Math (matches mlx-lm's LoRALinear.fuse and our QuantizedLinear.forward):
5
- // forward: y = x @ W_dequantᵀ + scale · ((x @ A) @ B)
6
- // fused: W'[out,in] = W + (scale · B)ᵀ.astype(bf16) @ Aᵀ.astype(bf16)
7
- // mlx-lm casts the (scaled) LoRA factors to the base dtype BEFORE the matmul,
8
- // then re-quantizes with the module's OWN source spec (bits/group unchanged) —
9
- // we do the same, so the output config.json is the source config verbatim.
10
- //
11
- // Everything runs on the CPU stream (same discipline as src/quantize/quantizer:
12
- // the safetensors loader lives there, and one module is materialized, fused,
13
- // and disposed at a time). Modules the adapter does not touch pass through
14
- // bit-identical (no dequant→requant round trip).
15
-
16
- import { existsSync, mkdirSync } from "node:fs";
17
- import { join } from "node:path";
18
- import { clearCache, Dtype } from "../mlx/ffi";
19
- import { cpuStream, type MlxArray } from "../mlx/array";
20
- import * as ops from "../mlx/ops";
21
- import { Weights } from "../weights";
22
- import { loadModelConfig, quantFor } from "../config";
23
- import { loadAdapterTensors, readAdapterScale, adapterWeightsFile } from "../lora";
24
- import { writeShardedSafetensors, type NamedTensor } from "../quantize/safetensors-writer";
25
- import { copyAuxFiles } from "../quantize/config-writer";
26
-
27
- /** Quantize the allocator cache this often (in modules processed). */
28
- const CACHE_CLEAR_EVERY = 16;
29
-
30
- export interface FuseStats {
31
- outDir: string;
32
- /** Modules whose weights had an adapter delta folded in. */
33
- fusedModules: number;
34
- /** Adapter tensors that matched no base weight (skipped). */
35
- skippedAdapterTensors: number;
36
- totalTensors: number;
37
- }
38
-
39
- export interface FuseProgress {
40
- stage: string;
41
- message: string;
42
- /** 0..1 fraction of modules processed. */
43
- progress: number;
44
- }
45
-
46
- /** Resolve the adapter safetensors file: `<dir>/adapters.safetensors`,
47
- * `<dir>/best/adapters.safetensors`, `adapter_model.safetensors`, or an
48
- * explicit file path. */
49
- function resolveAdapterFile(dir: string): string {
50
- const best = join(dir, "best");
51
- if (existsSync(join(best, "adapters.safetensors"))) return join(best, "adapters.safetensors");
52
- if (dir.endsWith(".safetensors") && existsSync(dir)) return dir;
53
- return adapterWeightsFile(dir);
54
- }
55
-
56
- /** Group adapter tensors into (modulePath → {a, b}); disposes non-LoRA
57
- * tensors and returns how many were skipped. */
58
- function groupPairs(tensors: Map<string, MlxArray>): {
59
- pairs: Map<string, { a?: MlxArray; b?: MlxArray }>;
60
- skipped: number;
61
- } {
62
- const pairs = new Map<string, { a?: MlxArray; b?: MlxArray }>();
63
- let skipped = 0;
64
- for (const [name, arr] of tensors) {
65
- const m = name.match(/^(.*)\.(lora_a|lora_A|lora_b|lora_B)(\.weight)?$/);
66
- if (!m) { skipped++; arr.dispose(); continue; }
67
- const slot = pairs.get(m[1]!) ?? {};
68
- if (m[2]!.toLowerCase() === "lora_a") slot.a = arr; else slot.b = arr;
69
- pairs.set(m[1]!, slot);
70
- }
71
- return { pairs, skipped };
72
- }
73
-
74
- /**
75
- * Fuse the adapter at `adapterDir` into the base model at `modelDir` and
76
- * write a complete, loadable snapshot to `outDir` (weights + config.json +
77
- * tokenizer aux files). Quantization layout is preserved verbatim: fused
78
- * modules are dequantized, updated, and re-quantized with their own source
79
- * spec; untouched modules pass through bit-identical.
80
- */
81
- export async function fuseAdapter(
82
- modelDir: string,
83
- adapterDir: string,
84
- outDir: string,
85
- onProgress?: (e: FuseProgress) => void,
86
- ): Promise<FuseStats> {
87
- const progress = (stage: string, message: string, frac: number) =>
88
- onProgress?.({ stage, message, progress: frac });
89
-
90
- progress("loading", `Reading ${modelDir}`, 0);
91
- const config = await loadModelConfig(modelDir);
92
- const weights = await Weights.open(modelDir);
93
-
94
- const adapterFile = resolveAdapterFile(adapterDir);
95
- const { scale, rsLora } = await readAdapterScale(adapterDir);
96
- const { pairs, skipped } = groupPairs(loadAdapterTensors(adapterFile));
97
-
98
- // Map each adapter module path to the base weight-file module path,
99
- // probing the VLM-wrapped/unwrapped `language_model.` prefixes both ways
100
- // (same normalization as AdapterManager.mount).
101
- const weightNames = new Set(weights.tensorNames);
102
- const byModule = new Map<string, { a: MlxArray; b: MlxArray }>();
103
- let unmatched = 0;
104
- const disposeAll = () => {
105
- for (const { a, b } of pairs.values()) { a?.dispose(); b?.dispose(); }
106
- };
107
- try {
108
- for (const [path, { a, b }] of pairs) {
109
- if (!a || !b) throw new Error(`${path}: adapter has only one of lora_a/lora_b`);
110
- const candidates = [
111
- path,
112
- `language_model.${path}`,
113
- path.startsWith("language_model.") ? path.slice("language_model.".length) : null,
114
- ].filter((c): c is string => c !== null);
115
- const hit = candidates.find((c) => weightNames.has(`${c}.weight`));
116
- if (!hit) { unmatched++; continue; }
117
- byModule.set(hit, { a, b });
118
- }
119
- if (byModule.size === 0)
120
- throw new Error(
121
- `no adapter tensors match the base model's weights — was ` +
122
- `${adapterFile} trained for a different base?`,
123
- );
124
- } catch (e) {
125
- disposeAll();
126
- weights.dispose();
127
- throw e;
128
- }
129
-
130
- mkdirSync(outDir, { recursive: true });
131
-
132
- // Walk tensors in source order (diff-friendly like the quantizer): group
133
- // names into module bases so a fused module's weight/scales/biases are
134
- // regenerated together.
135
- const names = weights.tensorNames;
136
- const moduleSuffixes = new Map<string, Set<string>>();
137
- const passthroughNames: string[] = [];
138
- for (const name of names) {
139
- const m = /^(.*)\.(weight|scales|biases)$/.exec(name);
140
- if (!m) { passthroughNames.push(name); continue; }
141
- let set = moduleSuffixes.get(m[1]!);
142
- if (!set) { set = new Set(); moduleSuffixes.set(m[1]!, set); }
143
- set.add(m[2]!);
144
- }
145
-
146
- const out: NamedTensor[] = [];
147
- let fusedModules = 0;
148
- let processed = 0;
149
- const totalModules = moduleSuffixes.size;
150
-
151
- try {
152
- for (const [base, suffixes] of moduleSuffixes) {
153
- processed++;
154
- if (processed % CACHE_CLEAR_EVERY === 0) clearCache();
155
- const pair = suffixes.has("weight") ? byModule.get(base) : undefined;
156
- if (!pair) {
157
- for (const suf of suffixes) {
158
- const n = `${base}.${suf}`;
159
- out.push({ name: n, array: weights.tensor(n) });
160
- }
161
- continue;
162
- }
163
- progress("fusing", `Module ${processed}/${totalModules}: ${base}`, processed / totalModules);
164
-
165
- // Materialize the full-precision weight [out, in].
166
- const srcSpec = suffixes.has("scales") ? quantFor(config.quantization, base) : null;
167
- let w: MlxArray;
168
- if (srcSpec) {
169
- const packed = weights.tensor(`${base}.weight`);
170
- const scales = weights.tensor(`${base}.scales`);
171
- const biases = suffixes.has("biases") ? weights.tensor(`${base}.biases`) : null;
172
- w = ops.dequantize(packed, scales, biases, srcSpec, cpuStream);
173
- } else {
174
- w = weights.tensor(`${base}.weight`);
175
- }
176
- const [outF, inF] = w.shape as [number, number];
177
- const { a, b } = pair;
178
- const [aIn, rank] = a.shape as [number, number];
179
- const [bRank, bOut] = b.shape as [number, number];
180
- if (aIn !== inF || bRank !== rank || bOut !== outF)
181
- throw new Error(
182
- `${base}: shape mismatch — lora_a [${a.shape}] / lora_b [${b.shape}] ` +
183
- `vs base [out ${outF}, in ${inF}]`,
184
- );
185
-
186
- // delta[out,in] = (scale·B)ᵀ.astype(bf16) @ Aᵀ.astype(bf16)
187
- // (rsLoRA: effective per-layer scale is α/√rank, matching mount/training.)
188
- const effScale = rsLora ? scale / Math.sqrt(rank) : scale;
189
- const bScaled = ops.mulScalar(b, effScale, cpuStream);
190
- const bT = ops.transposeAxes(bScaled, [1, 0], cpuStream);
191
- bScaled.dispose();
192
- const bT16 = bT.astype(Dtype.bfloat16, cpuStream);
193
- bT.dispose();
194
- const aT = ops.transposeAxes(a, [1, 0], cpuStream);
195
- const aT16 = aT.astype(Dtype.bfloat16, cpuStream);
196
- aT.dispose();
197
- const delta = ops.matmul(bT16, aT16, cpuStream);
198
- bT16.dispose();
199
- aT16.dispose();
200
-
201
- const w16 = w.dtype === Dtype.bfloat16 ? w : w.astype(Dtype.bfloat16, cpuStream);
202
- if (w16 !== w) w.dispose();
203
- const fusedW = ops.add(w16, delta, cpuStream);
204
- w16.dispose();
205
- delta.dispose();
206
-
207
- if (srcSpec) {
208
- const q = ops.quantize(fusedW, srcSpec.groupSize, srcSpec.bits, srcSpec.mode, cpuStream);
209
- fusedW.dispose();
210
- out.push({ name: `${base}.weight`, array: q.packed });
211
- out.push({ name: `${base}.scales`, array: q.scales });
212
- out.push({ name: `${base}.biases`, array: q.biases });
213
- } else {
214
- out.push({ name: `${base}.weight`, array: fusedW });
215
- }
216
- fusedModules++;
217
- }
218
- for (const name of passthroughNames) out.push({ name, array: weights.tensor(name) });
219
-
220
- progress("writing", `Writing ${out.length} tensors to ${outDir}`, 1);
221
- writeShardedSafetensors(outDir, out);
222
-
223
- // Config is unchanged (same quantization layout) — copy it and the aux
224
- // files (tokenizer, chat template, …) verbatim for a loadable snapshot.
225
- await Bun.write(join(outDir, "config.json"), Bun.file(join(modelDir, "config.json")));
226
- await copyAuxFiles(modelDir, outDir);
227
-
228
- progress("done", `Fused ${fusedModules} module(s) → ${outDir}`, 1);
229
- return {
230
- outDir,
231
- fusedModules,
232
- skippedAdapterTensors: skipped + unmatched * 2,
233
- totalTensors: out.length,
234
- };
235
- } finally {
236
- for (const t of out) {
237
- try { t.array.dispose(); } catch {}
238
- }
239
- disposeAll();
240
- weights.dispose();
241
- clearCache();
242
- }
243
- }
@@ -1,46 +0,0 @@
1
- // Public surface for native LoRA + DPO fine-tuning.
2
-
3
- export { ValueAndGrad } from "../mlx/autograd";
4
-
5
- export {
6
- loadSftDataset, iterateSftBatches, encodeSftRow, sftRowFormat,
7
- loadDpoDataset, iterateDpoBatches, encodeDpoRow, tokenizePair,
8
- probeFormat, rowLength, dpoChosenLength, dpoRejectedLength, PAD_TO,
9
- type SftExample, type SftBatch, type SftRowFormat,
10
- type DpoExample, type DpoBatch, type DatasetFormat,
11
- } from "./dataset";
12
-
13
- export { trainForward, buildBatchedPadMask } from "./forward";
14
-
15
- export {
16
- sftLoss, dpoLoss, dpoRefLogps, dpoMetrics,
17
- orpoLoss, orpoMetrics, type SftScope,
18
- } from "./loss";
19
-
20
- export {
21
- AdamW, warmupCosineSchedule, type AdamWOptions,
22
- } from "./optimizer";
23
-
24
- export {
25
- buildTrainableLora, attachForTraining, detachTraining, setLoraScale,
26
- flatA, flatB, flatParams, saveAdapter, disposeLora,
27
- type TrainableLora, type TrainableTarget, type SaveAdapterConfig,
28
- } from "./lora-params";
29
-
30
- export {
31
- resolveRanks, DEFAULT_TARGET_MODULES,
32
- type RankScaling, type ResolveRanksOptions,
33
- } from "./rank";
34
-
35
- export {
36
- trainLora, DEFAULT_TRAIN_CONFIG,
37
- type TrainConfig, type TrainResult,
38
- } from "./trainer";
39
-
40
- export { mergeAdapters, type MergeStats } from "./merge";
41
-
42
- export { fuseAdapter, type FuseStats, type FuseProgress } from "./fuse";
43
-
44
- export { exportAdapter, type ExportManifest } from "./export";
45
-
46
- export { finetuneRunner, inspectDataset } from "./job";