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