mlx-bun 0.2.0 → 0.3.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (364) hide show
  1. package/README.md +71 -588
  2. package/docs/README.md +91 -0
  3. package/docs/{design → archive/investigations}/audio-input-plan.md +4 -4
  4. package/docs/{design → archive/investigations}/diffusion-gemma-port.md +3 -3
  5. package/docs/{design → archive/investigations}/faithful-l1-consolidation.md +4 -4
  6. package/docs/{design → archive/investigations}/minicpm5-decode-megakernel.md +11 -11
  7. package/docs/{design → archive/investigations}/optimization_plan.md +1 -1
  8. package/docs/{design → archive/investigations}/pre-colibri-stabilization.md +5 -5
  9. package/docs/{design → archive/investigations}/the-dreaming-handoff.md +2 -2
  10. package/docs/design/batching.md +569 -0
  11. package/docs/design/decode-speed-program.md +3178 -169
  12. package/docs/design/dreaming-nightly-pipeline.md +675 -282
  13. package/docs/design/generic-model-support.md +579 -334
  14. package/docs/design/kv-cache.md +680 -0
  15. package/docs/design/orpo-training.md +673 -1239
  16. package/docs/design/speculative-decoding.md +1627 -0
  17. package/docs/design/turboquant.md +2102 -0
  18. package/docs/design/unified-engine-frontier-plan.md +1831 -727
  19. package/docs/design/web-chat-redesign.md +666 -1083
  20. package/docs/planning/PRODUCT_ROADMAP.md +34 -1
  21. package/docs/planning/release-notes-v0.3.0.md +35 -0
  22. package/docs/planning/video-series.md +2 -2
  23. package/docs/reference/benchmarks.md +1162 -0
  24. package/docs/reference/cli.md +394 -159
  25. package/docs/reference/distribution.md +321 -182
  26. package/docs/reference/environment.md +348 -0
  27. package/docs/reference/glossary.md +39 -0
  28. package/docs/reference/library-api.md +566 -168
  29. package/docs/reference/memory.md +227 -64
  30. package/docs/reference/models.md +104 -6
  31. package/docs/reference/server-api.md +886 -425
  32. package/docs/reference/server-config.md +638 -330
  33. package/docs/reference/training.md +393 -254
  34. package/docs/reference/troubleshooting.md +166 -0
  35. package/package.json +9 -4
  36. package/src/audio/decode.ts +2 -2
  37. package/src/audio/features.ts +2 -2
  38. package/src/backends/mlx/autoregressive.ts +108 -0
  39. package/src/backends/mlx/batch-group.ts +1646 -0
  40. package/src/backends/mlx/diffusion.ts +39 -0
  41. package/src/backends/mlx/draft-target.ts +47 -0
  42. package/src/backends/mlx/evaluation.ts +105 -0
  43. package/src/backends/mlx/fill-append.ts +45 -0
  44. package/src/backends/mlx/gateway-binding.ts +71 -0
  45. package/src/backends/mlx/graph.ts +38 -0
  46. package/src/backends/mlx/memory-client.ts +189 -0
  47. package/src/backends/mlx/memory-guard.ts +40 -0
  48. package/src/backends/mlx/methods.ts +90 -0
  49. package/src/backends/mlx/model-prompt.ts +164 -0
  50. package/src/backends/mlx/model-serving.ts +79 -0
  51. package/src/backends/mlx/prefill.ts +33 -0
  52. package/src/backends/mlx/request-policy.ts +20 -0
  53. package/src/backends/mlx/rollback.ts +51 -0
  54. package/src/backends/mlx/serial-executor.ts +273 -0
  55. package/src/backends/mlx/speculative.ts +114 -0
  56. package/src/backends/mlx/state-views.ts +26 -0
  57. package/src/backends/mlx/text-engine.ts +73 -0
  58. package/src/cli.ts +58 -25
  59. package/src/client.ts +66 -0
  60. package/src/config.ts +21 -4
  61. package/src/contracts/completion.ts +16 -0
  62. package/src/contracts/execution.ts +49 -0
  63. package/src/contracts/generation.ts +103 -0
  64. package/src/contracts/host.ts +7 -0
  65. package/src/contracts/jobs.ts +66 -0
  66. package/src/contracts/pi-provider.ts +30 -0
  67. package/src/contracts/pi.ts +259 -0
  68. package/src/contracts/resources.ts +18 -0
  69. package/src/contracts/scheduling.ts +27 -0
  70. package/src/contracts/tsconfig.json +14 -0
  71. package/src/diffusion/diffusion-generate.ts +218 -153
  72. package/src/embed.ts +1 -1
  73. package/src/engine/admission.ts +59 -0
  74. package/src/engine/cancellation.ts +38 -0
  75. package/src/engine/engine.ts +56 -0
  76. package/src/engine/execution-coordinator.ts +70 -0
  77. package/src/engine/execution-plan.ts +51 -0
  78. package/src/engine/resources.ts +42 -0
  79. package/src/engine/scheduler.ts +67 -0
  80. package/src/engine/session.ts +247 -0
  81. package/src/eval/README.md +44 -1
  82. package/src/eval/kl.ts +236 -3
  83. package/src/eval/reference-logits.ts +207 -0
  84. package/src/eval/runner.ts +8 -96
  85. package/src/fill/echo-index.ts +188 -0
  86. package/src/fill/fill-session.ts +452 -0
  87. package/src/fill/proposal.ts +68 -0
  88. package/src/fill/schema-rows.ts +506 -0
  89. package/src/fill/tool-boundary.ts +103 -0
  90. package/src/fit.ts +1 -1
  91. package/src/generate.ts +656 -268
  92. package/src/grammar.ts +12 -7
  93. package/src/harness-pi.ts +14 -33
  94. package/src/hub-rest.ts +4 -4
  95. package/src/index.ts +7 -0
  96. package/src/inference/denoising.ts +17 -0
  97. package/src/inference/graph.ts +33 -0
  98. package/src/inference/prefill.ts +39 -0
  99. package/src/inference/rollback.ts +7 -0
  100. package/src/jobs/db.ts +4 -0
  101. package/src/jobs/index.ts +2 -0
  102. package/src/jobs/job-entry.ts +12 -12
  103. package/src/jobs/runner.ts +98 -13
  104. package/src/jobs/task-client.ts +19 -0
  105. package/src/jobs/types.ts +2 -58
  106. package/src/kv-store.ts +109 -49
  107. package/src/lab/README.md +18 -0
  108. package/src/{assets → lab/curve}/curve-designer.html +2 -2
  109. package/src/{curve-sampler.ts → lab/curve/curve-sampler.ts} +2 -2
  110. package/src/{expert-trace.ts → lab/expert-trace/expert-trace.ts} +5 -5
  111. package/src/{model → lab/paged-kv}/paged-kv.ts +9 -6
  112. package/src/library.ts +40 -0
  113. package/src/lora.ts +18 -0
  114. package/src/media-fetch.ts +1 -1
  115. package/src/memory/article.ts +1 -1
  116. package/src/memory/chunk.ts +1 -1
  117. package/src/memory/model.ts +30 -169
  118. package/src/memory/prompts.ts +1 -1
  119. package/src/memory/rest.ts +1 -1
  120. package/src/memory/schedule.ts +1 -1
  121. package/src/memory/stages.ts +4 -1
  122. package/src/memory/vault.ts +1 -3
  123. package/src/mlx/array.ts +9 -10
  124. package/src/mlx/compile.ts +12 -2
  125. package/src/mlx/ffi.ts +58 -11
  126. package/src/mlx/materialize.ts +23 -0
  127. package/src/mlx/metal-kernel.ts +10 -11
  128. package/src/mlx/ops.ts +33 -13
  129. package/src/mlx/quantized-rows.ts +36 -0
  130. package/src/mlx/token-bitmask.ts +52 -0
  131. package/src/mlx/turboquant-kv-decode.ts +82 -0
  132. package/src/mlx/turboquant-kv-inverse.ts +122 -0
  133. package/src/mlx/turboquant-ops.ts +33 -3
  134. package/src/mlx/turboquant-tables.ts +2 -2
  135. package/src/model/batched-mask.ts +8 -2
  136. package/src/model/batched-quant.ts +1 -1
  137. package/src/model/batched-rotating-quant.ts +2 -2
  138. package/src/model/batched-rotating.ts +1 -1
  139. package/src/model/compiled-decode.ts +1 -1
  140. package/src/model/diffusion-gemma.ts +8 -8
  141. package/src/model/factory.ts +76 -33
  142. package/src/model/fingerprint.ts +1 -1
  143. package/src/model/flash-attention.ts +3 -3
  144. package/src/model/gemma4-base.ts +26 -9
  145. package/src/model/gemma4.ts +6 -6
  146. package/src/model/generated/gemma4-12b.ts +2 -2
  147. package/src/model/generated/gemma4-26b.ts +2 -2
  148. package/src/model/generated/gemma4-e4b.ts +2 -2
  149. package/src/model/glm52-dsa.ts +13 -9
  150. package/src/model/glm52-memory.ts +1 -1
  151. package/src/model/implementation.ts +69 -0
  152. package/src/model/minicpm5.ts +6 -8
  153. package/src/model/profile.ts +12 -2
  154. package/src/model/qwen-append.ts +12 -0
  155. package/src/model/qwen-conv.ts +52 -0
  156. package/src/model/qwen3_5-checkpoint.ts +164 -0
  157. package/src/model/qwen3_5.ts +140 -56
  158. package/src/model/support.ts +1 -1
  159. package/src/model/trellis-balanced-scatter.ts +51 -0
  160. package/src/model/trellis-linear.ts +672 -0
  161. package/src/model/trellis-shared-m.ts +151 -0
  162. package/src/model/trellis-shared-scatter.ts +144 -0
  163. package/src/model/trellis-splitk-prefill.ts +158 -0
  164. package/src/model/trellis-tiled-prefill.ts +136 -0
  165. package/src/model/trellis-vector-expand.ts +62 -0
  166. package/src/model/trellis-wide-prefill.ts +94 -0
  167. package/src/model/universal/dense.ts +1 -1
  168. package/src/model/universal/modules.ts +8 -1
  169. package/src/native-pack.ts +6 -5
  170. package/src/pi-provider.ts +6 -64
  171. package/src/pi-web.ts +8 -315
  172. package/src/prompt-cache.ts +96 -53
  173. package/src/quantize/allocator.ts +1 -1
  174. package/src/quantize/config-writer.ts +10 -3
  175. package/src/quantize/rotate.ts +4 -4
  176. package/src/quantize/trellis.ts +586 -0
  177. package/src/responses.ts +28 -1
  178. package/src/runtime-config.ts +18 -4
  179. package/src/sampler.ts +4 -4
  180. package/src/serve/batch-scheduler.ts +8 -1671
  181. package/src/serve/chat-request.ts +391 -0
  182. package/src/serve/chat-stage.ts +195 -0
  183. package/src/serve/checkpoint-identity.ts +64 -0
  184. package/src/serve/completion-executor.ts +15 -7
  185. package/src/serve/discovery-routes.ts +12 -13
  186. package/src/serve/generation-gateway.ts +154 -193
  187. package/src/serve/http.ts +159 -0
  188. package/src/serve/inference-request.ts +79 -0
  189. package/src/serve/isolate.ts +297 -105
  190. package/src/serve/lab-routes.ts +3 -0
  191. package/src/serve/lane-registry.ts +2 -1
  192. package/src/serve/model-admin-routes.ts +7 -5
  193. package/src/serve/model-binding.ts +43 -0
  194. package/src/serve/model-host.ts +512 -0
  195. package/src/serve/openai-wire.ts +103 -0
  196. package/src/serve/parent-application.ts +56 -0
  197. package/src/serve/pipeline.ts +33 -0
  198. package/src/serve/preparation.ts +33 -0
  199. package/src/serve/request-plan.ts +19 -15
  200. package/src/serve/request-prep.ts +346 -0
  201. package/src/serve/response-lease.ts +24 -0
  202. package/src/serve/responses-client.ts +69 -0
  203. package/src/serve/session-completion-engine.ts +110 -0
  204. package/src/serve/session-search.ts +1 -1
  205. package/src/serve/static-routes.ts +5 -5
  206. package/src/serve/text-completion-stage.ts +104 -0
  207. package/src/serve/token-streams.ts +519 -0
  208. package/src/server.ts +317 -2649
  209. package/src/spec/assistant-source.ts +10 -44
  210. package/src/spec/deepspec-source.ts +6 -6
  211. package/src/spec/dflash-source.ts +8 -8
  212. package/src/spec/dspark/ab-stats.ts +3 -3
  213. package/src/spec/dspark/calibration.ts +2 -2
  214. package/src/spec/dspark/deepspec-module.ts +2 -2
  215. package/src/spec/dspark/generate-dflash.ts +2 -2
  216. package/src/spec/dspark/generate.ts +2 -2
  217. package/src/spec/dspark/loader.ts +1 -1
  218. package/src/spec/dspark/module-dflash.ts +4 -3
  219. package/src/spec/dspark/module.ts +3 -3
  220. package/src/spec/dspark/quantize-drafter.ts +1 -1
  221. package/src/spec/generate.ts +1 -1
  222. package/src/spec/glm52-mtp-source.ts +1 -1
  223. package/src/spec/ngram-source.ts +1 -1
  224. package/src/spec/prefix-state.ts +45 -0
  225. package/src/spec/qwen-mtp-source.ts +166 -69
  226. package/src/spec/serve-loop.ts +154 -156
  227. package/src/spec/source.ts +42 -10
  228. package/src/spec/two-model.ts +1 -1
  229. package/src/ssd-cache.ts +80 -12
  230. package/src/ssd-durability.ts +25 -13
  231. package/src/tool-call.ts +51 -14
  232. package/src/train/flash-cce.ts +2 -2
  233. package/src/train/forward.ts +6 -0
  234. package/src/train/loss.ts +1 -1
  235. package/src/train/prefix-shared.ts +5 -1
  236. package/src/train/segmented.ts +5 -3
  237. package/src/train/trainer.ts +6 -6
  238. package/src/vision/prompt.ts +1 -1
  239. package/src/vision/qwen3vl-tower.ts +52 -10
  240. package/src/web/app.html +1 -1
  241. package/src/web/app.js +5 -3
  242. package/src/web/src/adapters-panel.ts +3 -3
  243. package/src/web/src/assistant.ts +2 -36
  244. package/src/web/src/chat.ts +4 -4
  245. package/src/web/src/composer.ts +4 -4
  246. package/src/web/src/hub.ts +1 -1
  247. package/src/web/src/markdown.ts +0 -0
  248. package/src/web/src/memory-panel.ts +1 -1
  249. package/src/web/src/model-picker.ts +1 -1
  250. package/src/web/src/palette.ts +1 -1
  251. package/src/web/src/protocol.ts +3 -39
  252. package/src/web/src/sessions.ts +1 -1
  253. package/src/web/src/ui-catalog.ts +4 -18
  254. package/src/web-assets.ts +47 -0
  255. package/src/weights.ts +55 -9
  256. package/docs/archive/investigations/curve-runs/bisector-route.json +0 -98
  257. package/docs/archive/investigations/curve-runs/bisector-trace.json +0 -1393
  258. package/docs/archive/investigations/curve-runs/curve-terrain.json +0 -1179
  259. package/docs/archive/investigations/curve-runs/curve-terrain.svg +0 -42
  260. package/docs/archive/investigations/hlg-ab.html +0 -34
  261. package/docs/archive/investigations/hlg-explorer.html +0 -238
  262. package/docs/archive/investigations/hlg-figs/cliff.svg +0 -1
  263. package/docs/archive/investigations/hlg-figs/cross.svg +0 -1
  264. package/docs/archive/investigations/hlg-figs/edges.svg +0 -1
  265. package/docs/archive/investigations/hlg-figs/k20-grids.svg +0 -1
  266. package/docs/archive/investigations/hlg-figs/k20-peaks.svg +0 -1
  267. package/docs/archive/investigations/hlg-figs/lw-curve.svg +0 -1
  268. package/docs/archive/investigations/hlg-figs/terrain.svg +0 -1
  269. package/docs/archive/investigations/hlg-report.html +0 -262
  270. package/docs/archive/investigations/hlg-runs/hlg-ab.json +0 -149
  271. package/docs/archive/investigations/hlg-runs/hlg-blind-prompts.md +0 -87
  272. package/docs/archive/investigations/hlg-runs/hlg-blind.json +0 -107
  273. package/docs/archive/investigations/hlg-runs/hlg-cliff.json +0 -197
  274. package/docs/archive/investigations/hlg-runs/hlg-cross.json +0 -505
  275. package/docs/archive/investigations/hlg-runs/hlg-edge.json +0 -388
  276. package/docs/archive/investigations/hlg-runs/hlg-eval1.json +0 -25
  277. package/docs/archive/investigations/hlg-runs/hlg-frontier.json +0 -86
  278. package/docs/archive/investigations/hlg-runs/hlg-grid.json +0 -242
  279. package/docs/archive/investigations/hlg-runs/hlg-grid2.json +0 -343
  280. package/docs/archive/investigations/hlg-runs/hlg-lw-coarse.json +0 -110
  281. package/docs/archive/investigations/hlg-runs/hlg-map.json +0 -383
  282. package/docs/archive/investigations/hlg-runs/run0-original.log +0 -159
  283. package/docs/archive/investigations/hlg-runs/runA-recommended.log +0 -131
  284. package/docs/archive/investigations/hlg-runs/runAA-cliff-cpm5.log +0 -26
  285. package/docs/archive/investigations/hlg-runs/runAB-grid2-e4b.log +0 -18
  286. package/docs/archive/investigations/hlg-runs/runAC-peak-e4b-partial.log +0 -25
  287. package/docs/archive/investigations/hlg-runs/runAD-ab-samples.log +0 -8
  288. package/docs/archive/investigations/hlg-runs/runAE-edge-e4b.log +0 -51
  289. package/docs/archive/investigations/hlg-runs/runAF-logit-calibration.log +0 -10
  290. package/docs/archive/investigations/hlg-runs/runAG-eval-config2.log +0 -16
  291. package/docs/archive/investigations/hlg-runs/runB-gain-variants.log +0 -192
  292. package/docs/archive/investigations/hlg-runs/runC-pivots.log +0 -162
  293. package/docs/archive/investigations/hlg-runs/runD-oetf.log +0 -132
  294. package/docs/archive/investigations/hlg-runs/runE-eotf.log +0 -162
  295. package/docs/archive/investigations/hlg-runs/runF-shaper-working.log +0 -86
  296. package/docs/archive/investigations/hlg-runs/runG-shaper-sweep.log +0 -101
  297. package/docs/archive/investigations/hlg-runs/runH-wsweep.log +0 -176
  298. package/docs/archive/investigations/hlg-runs/runI-asweep.log +0 -176
  299. package/docs/archive/investigations/hlg-runs/runJ-knobmap.log +0 -49
  300. package/docs/archive/investigations/hlg-runs/runK-frontier.log +0 -19
  301. package/docs/archive/investigations/hlg-runs/runL-widemap.log +0 -65
  302. package/docs/archive/investigations/hlg-runs/runM-frontier-emb.log +0 -15
  303. package/docs/archive/investigations/hlg-runs/runN-verify.log +0 -116
  304. package/docs/archive/investigations/hlg-runs/runO-frontier-seed5000.log +0 -15
  305. package/docs/archive/investigations/hlg-runs/runP-grid-AxSm.log +0 -15
  306. package/docs/archive/investigations/hlg-runs/runQ-grid-tgW.log +0 -15
  307. package/docs/archive/investigations/hlg-runs/runR-blind-transcript.log +0 -79
  308. package/docs/archive/investigations/hlg-runs/runS-cross-coherence.log +0 -33
  309. package/docs/archive/investigations/hlg-runs/runT-cross-cpm5.log +0 -33
  310. package/docs/archive/investigations/hlg-runs/runU-cross-12b.log +0 -33
  311. package/docs/archive/investigations/hlg-runs/runV-lw-e4b.log +0 -22
  312. package/docs/archive/investigations/hlg-runs/runW-lw-12b.log +0 -22
  313. package/docs/archive/investigations/hlg-runs/runX-lw-cpm5.log +0 -22
  314. package/docs/archive/investigations/hlg-runs/runY-cliff-e4b.log +0 -26
  315. package/docs/archive/investigations/hlg-runs/runZ-cliff-12b.log +0 -26
  316. package/docs/archive/mlx-bun-lab-report.html +0 -275
  317. package/docs/archive/wwdc-mlx-bun.html +0 -848
  318. package/docs/design/adapters-end-to-end.md +0 -205
  319. package/docs/design/batching-perf-path.md +0 -255
  320. package/docs/design/batching-v2-plan.md +0 -529
  321. package/docs/design/bucketing-stage.md +0 -197
  322. package/docs/design/compat-cli-surface-design.md +0 -230
  323. package/docs/design/docs-reorg-plan.md +0 -165
  324. package/docs/design/dspark-serving-program.md +0 -371
  325. package/docs/design/dspark-speculative-decoding.md +0 -318
  326. package/docs/design/grammar-spec-batching-integration.md +0 -301
  327. package/docs/design/memory-inference-path.md +0 -422
  328. package/docs/design/memory-synthesis.md +0 -287
  329. package/docs/design/memory-system.md +0 -237
  330. package/docs/design/mlx-lm-serving-execution-seam.md +0 -81
  331. package/docs/design/mlx-lm-tool-parity-plan.md +0 -545
  332. package/docs/design/omlx-adoption-map.md +0 -98
  333. package/docs/design/orpo-dynamic-lambda.md +0 -220
  334. package/docs/design/orpo-future-enhancements.md +0 -340
  335. package/docs/design/paged-kv-cache.md +0 -214
  336. package/docs/design/parallel-slots.md +0 -567
  337. package/docs/design/parity-tier-dag.md +0 -161
  338. package/docs/design/repo-cleanup-plan.md +0 -231
  339. package/docs/design/runtime-isolation.md +0 -263
  340. package/docs/design/segmented-backward-training.md +0 -543
  341. package/docs/design/spec-decode-larger-targets.md +0 -229
  342. package/docs/design/ssd-kv-cold-tier.md +0 -415
  343. package/docs/design/structured-output.md +0 -484
  344. package/docs/design/tauri-desktop-app.md +0 -127
  345. package/docs/design/the-dreaming-master-plan.md +0 -445
  346. package/docs/design/turboquant-kv.md +0 -235
  347. package/docs/design/turboquant-weights.md +0 -467
  348. package/docs/design/w4a16-compute-precision-spike.md +0 -232
  349. package/docs/design/web-chat-beat-matrix.md +0 -315
  350. package/docs/design/write-pipeline-entity.md +0 -468
  351. package/docs/planning/IDEAS.md +0 -69
  352. package/docs/reference/embedding.md +0 -197
  353. package/docs/reference/features-matrix.md +0 -101
  354. package/docs/reference/orpo-quickstart.md +0 -135
  355. package/src/flags.ts +0 -11
  356. /package/docs/{design → archive/investigations}/colibri-glm52-port.md +0 -0
  357. /package/docs/{planning → archive/planning}/chunk-task-roadmap.md +0 -0
  358. /package/docs/{planning → archive/planning}/curve-sampler-research-plan.md +0 -0
  359. /package/docs/{planning → archive/planning}/memory-docs-and-dag-plan.md +0 -0
  360. /package/docs/{planning → archive/planning}/pi-integration-review.md +0 -0
  361. /package/docs/{planning → archive/planning}/release-notes-v0.2.0.md +0 -0
  362. /package/src/{memory-plan.ts → admission-plan.ts} +0 -0
  363. /package/{docs/archive/investigations → src/serve/assets}/curve-terrain.html +0 -0
  364. /package/{docs/dag → src/serve/assets}/training-inference-map.html +0 -0
@@ -1,205 +0,0 @@
1
- # Adapters end-to-end — notes / TODO (2026-06-16)
2
-
3
- **Status legend:** ✅ DONE — shipped and verified in code | ⏳ PENDING — not yet implemented
4
-
5
- | Section | Status |
6
- | --- | --- |
7
- | §A Pi CLI extension (`extensions/mlx-bun-adapter.ts`) | ✅ DONE |
8
- | §B Web chat adapter selector (`src/web/app.html` + `before_provider_request` hook in `pi-web.ts`) | ✅ DONE |
9
- | §C `/v1/adapters/available` discovery (`src/server.ts:1182`, `src/lora.ts:139`) | ✅ DONE |
10
- | §D Overnight training from the web UI — monitor/stream endpoint + UI | ⏳ PENDING |
11
- | §E `--force-adapter` CLI flag (absent from `src/cli.ts`) | ⏳ PENDING |
12
- | Repo-id keying written into adapter configs at save time | ⏳ PENDING |
13
-
14
- Working notes for wiring trained LoRA adapters through the whole stack: Pi (CLI +
15
- web), the web Chat page UI, and the server's discovery surface. Captured from the
16
- 2026-06-16 session (serving verified working; pi-extension mechanism confirmed from
17
- a pi session example).
18
-
19
- ## The mental model — three distinct states
20
-
21
- Josh's framing, which the whole design should keep separate:
22
-
23
- 1. **Available** — adapters that exist *on disk* and are *compatible* with the
24
- currently-served model. Not in memory, cost nothing. (Discovery problem.)
25
- 2. **Loaded / mounted** — adapters mounted into the *running server* via
26
- `AdapterManager`; they cost memory. (`POST`/`DELETE /v1/adapters`.)
27
- 3. **Selected for this request** — the `adapter` field on a given request:
28
- `"none"` (default) | `"<id>"` | `"a+b"` (stacked). (`resolveSpec`.)
29
-
30
- You can be available-but-not-loaded, loaded-but-not-selected, etc. The UI and the
31
- API should let you move an adapter between these states explicitly.
32
-
33
- ## What already exists (server side)
34
-
35
- - `GET /v1/adapters` → lists **loaded** adapters (`id, path, rank, scale, size_bytes, mounted_layers`). `src/server.ts:1086`
36
- - `POST /v1/adapters {id, path}` → **load/mount**. `src/server.ts:1094`
37
- - `DELETE /v1/adapters/:id` → **unload/unmount**. `src/server.ts:1113`
38
- - Request `adapter` field → `resolveSpec` (`src/lora.ts:231`): `"none"`/empty → `[]`
39
- (genuine base); unknown id → **throws** (loud, not a silent base fallback);
40
- `"a+b"` stacks. Default behavior is already "none".
41
-
42
- So **loaded** and **selected** are done. The gap is **available** (discovery) plus
43
- the two front-ends (Pi extension, web Chat selector).
44
-
45
- ## Shipped (§A–§C) ✅
46
-
47
- ### A. Pi CLI/terminal — adapter via extension *(default none)* ✅ DONE
48
-
49
- Wired a Pi extension using the `before_provider_request` hook to inject the `adapter`
50
- field into the outgoing OpenAI-compatible payload. `models.json` can't add arbitrary
51
- body fields — the hook was the right tool.
52
-
53
- Shipped as `~/.pi/agent/extensions/mlx-bun-adapter.ts` (copy in-repo + one-liner
54
- installer):
55
-
56
- ```ts
57
- import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
58
-
59
- export default function (pi: ExtensionAPI) {
60
- pi.on("before_provider_request", (event) => {
61
- const payload = event.payload as Record<string, unknown>;
62
- if (payload.model !== "local") return; // only our local alias
63
- const adapter = process.env.PI_MLX_ADAPTER; // unset → default none
64
- if (!adapter) return; // no field → server serves base
65
- return { ...payload, adapter };
66
- });
67
- }
68
- ```
69
-
70
- - Defaults to **none** (env unset → hook returns `undefined` → payload unchanged).
71
- - Selection via `PI_MLX_ADAPTER=chunk pi --provider mlx-bun --model local -p "..."`
72
- (no native `-p --adapter` flag; unknown CLI flags get rejected). Optional shell
73
- helper: `pi-chunk() { PI_MLX_ADAPTER=chunk pi --provider mlx-bun --model local -p "$@"; }`.
74
- - Lives in the global `~/.pi/agent/extensions/` dir so it loads for non-interactive
75
- `-p` too. `/reload` picks it up.
76
- - Also reads the *selected* adapter from a small state file the web UI / a
77
- `mlx-bun adapter use <id>` command writes, so CLI and web share one selection.
78
-
79
- ### B. Web Chat page — selector + load/unload *(default none)* ✅ DONE
80
-
81
- Path: `src/web/app.html` (UI) ↔ `/ws/chat` (`src/pi-web.ts`) ↔ in-process Pi ↔ model.
82
-
83
- - **Selector** in `app.html`: a dropdown of *available* adapters with **"none"
84
- selected by default**. Changing it sets the active selection for subsequent
85
- messages. Multi-select → `"a+b"` (resolveSpec already supports stacking).
86
- - **Load / unload controls**: buttons that call `POST /v1/adapters` (load) and
87
- `DELETE /v1/adapters/:id` (unload). Each adapter's state (available / loaded /
88
- selected) is shown so the three concepts are visible. Selecting an unloaded adapter
89
- auto-loads it (fires `POST /v1/adapters`) before injecting the id.
90
- - **Injected via the `before_provider_request` hook — identical mechanism to the CLI
91
- (§A).** `pi-web.ts` registers the inline extension in `createAgentSession`
92
- (pi-web.ts:469/605), so `pi.on("before_provider_request", (e) => ({ ...e.payload,
93
- adapter }))` reads the current selection from WS session state. **No custom
94
- provider plumbing, no new model-call wiring.** The hook is the spine for BOTH UIs.
95
- - **Selection → hook**: the dropdown updates the session's selected-adapter value via
96
- a small client→server message (reusing the `set_thinking`-style toggle path). Default
97
- none → hook returns the payload unchanged → base model.
98
-
99
- ### C. Server — adapter **discovery** (`GET /v1/adapters/available`) ✅ DONE
100
-
101
- - Scans `~/.cache/mlx-bun-finetunes/` and `~/.cache/mlx-bun/adapters/`
102
- (`src/server.ts:1192`) and any sidecar next to the model dir.
103
- - `GET /v1/adapters/available` (`src/server.ts:1182`, `src/lora.ts:139`) →
104
- `[{ id, path, source_repo, rank, scale, compatible }]`, where each entry is read
105
- from the adapter's `optiq_lora_config.json` (`source_model`, `rank`, `scale`).
106
- - **Filters by the served model** on **repo id**
107
- (`mlx-community/MiniCPM5-1B-OptiQ-4bit`) rather than the fragile
108
- machine/snapshot-specific path. The server already tracks `repoId` per model;
109
- `compatible = adapter.source_repo === served.repoId`.
110
-
111
- ## TODO (§D–§E)
112
-
113
- ### D. Overnight training from the web UI
114
-
115
- Goal (Josh): *"start a new training session before bed, wake up to a new adapter."*
116
- Training must be fully drivable from the web UI — not just a script.
117
-
118
- - A **training page / panel** in the web UI: pick base model + dataset, set params
119
- (rank / scale / num_layers / LR / dropout / iters / seq / segmentSize — the
120
- configurable-params task, mlx-lm defaults), **start**, and leave it running.
121
- - **Monitor**: live loss / iter / ETA, and survives the browser closing (the run is
122
- server-side; the page reconnects). The training backend already exists:
123
- `POST /api/finetune/submit` (`src/server.ts:1781`). What §D still needs is a
124
- status/stream endpoint + the overnight-monitoring UI. (Note: `GET /fit` is the
125
- unrelated memory-fit solver in `src/fit.ts` — not a training endpoint.)
126
- - **Completion**: the finished adapter lands in the store and shows up as *available*
127
- (state A in the three-state model), ready to load/select or auto-serve.
128
- - This closes the loop with the chunk-task roadmap: iterate data → train overnight →
129
- wake to a scored adapter.
130
-
131
- ### E. `--force-adapter` for the nightly serving pipeline
132
-
133
- Goal (Josh): spin up a server and say *"use this model AND this adapter"* without
134
- per-request `adapter` fields — for an unattended nightly pipeline.
135
-
136
- - New CLI flag, e.g. `mlx-bun serve <model> --force-adapter <id|path>`:
137
- 1. **mount** the adapter at startup (load it), and
138
- 2. **force it active as the default** for every request — overriding the normal
139
- "selected defaults to none." A request could still override, or `--force-adapter`
140
- could be a hard pin (decide; pin is simpler for a pipeline).
141
- - Surface it on the ready card's `Perf`/serving summary so it's visible the server is
142
- adapter-pinned.
143
- - Pairs with D: train overnight via the UI → serve the result with `--force-adapter`
144
- in the nightly pipeline.
145
-
146
- ## Cross-cutting / already-open threads to fold in
147
-
148
- - **Repo-id keying for adapters** (from this session's discovery check): adapters
149
- currently record the model only as a fragile absolute path (`source_model` /
150
- `base_model_name_or_path`). Have the **trainer write the repo id** into the adapter
151
- config at save time, and match discovery on repo id. Robust across machines and
152
- snapshots. This is the prerequisite for B's selector and C's `compatible` filter.
153
- - **mlx-lm-aligned training defaults** (PARTIALLY DONE): `rank / scale /
154
- learning_rate` (plus `seq / segmentSize / iters / save_checkpoints`) are now
155
- run-configurable via env on the chunk-finetune script, threaded
156
- `chunk-finetune.ts` → `FinetuneSubmit` → `cfg`. In particular **`scale` is now
157
- wired end-to-end** (`SCALE` env → submit `scale` → `cfg.scale` →
158
- `buildTrainableLora`), and the script defaults to the **mlx-lm-aligned
159
- `scale 20 / LR 1e-5`** (== effective `scale 1 / LR 2e-4`). STILL PENDING:
160
- surface `num_layers` on the script (already in the submit contract +
161
- `TrainConfig`, just no env) and add `dropout` (not in `TrainConfig` yet).
162
- Note `DEFAULT_TRAIN_CONFIG` (`src/train/trainer.ts`) stays the conservative
163
- library default (`rank 8 / scale 1 / LR 2e-4`); the *script* overrides to the
164
- mlx-lm pair. (Older `minicpm5-chunk-segmented` predates this: rank 16 / scale 1
165
- / LR 2e-4.) Flow through
166
- `DEFAULT_TRAIN_CONFIG` (`src/train/trainer.ts`) → `FinetuneSubmit`/`parseConfig`
167
- (`src/train/job.ts`) → the `"finetune"` subprocess job (`src/server.ts:1678`,
168
- `submitSubprocess(store, "finetune", …)`). NOTE: `GET /fit` is NOT the training
169
- submit — it is the memory-fit solver (`src/fit.ts`): it solves weights + KV
170
- growth + prefill against this machine's usable memory and returns
171
- `{ fits, max_safe_context, predicted_decode_tps }` + a SKU matrix (used by
172
- `GET /fit`, serve-time admission, and the model listing).
173
- - **Adapter storage location**: prefer a dedicated `~/.cache/mlx-bun/adapters/` store
174
- with repo-id metadata over writing sidecars into the fragile HF snapshot cache.
175
-
176
- ## Inventory (scanned 2026-06-16)
177
-
178
- 5 adapters on this machine, **all local — none in the HF cache** (no LoRA repos
179
- downloaded, no `HF_HOME`/`HF_HUB_CACHE` override):
180
-
181
- - **3 production** in `~/.cache/mlx-bun-finetunes/`: `minicpm5-chunk-final`,
182
- `minicpm5-chunk-segmented`, `minicpm5-chunk-seq8192` — all
183
- `mlx-community/MiniCPM5-1B-OptiQ-4bit`, rank 16 / scale 1 / all layers / 44.9 MB.
184
- - **2 test fixtures** in `fixtures/adapters/`: `french`, `upper` (num_layers 4,
185
- 7 MB, **no `source_model` recorded** — model-agnostic).
186
-
187
- Design implications confirmed by the scan:
188
-
189
- 1. **HF cache is empty of adapters today** but the discovery scanner should still
190
- cover `~/.cache/huggingface/hub` (future LoRA pulls land there).
191
- 2. **`repo_id` / `source_repo` is absent in every real adapter** — only the
192
- snapshot *path* is stored. Repo-id keying must come first, plus a one-time
193
- backfill to stamp the existing three.
194
- 3. **The three MiniCPM5 adapters are indistinguishable by config** (same model /
195
- rank / scale / layers) — only the dir name differs. Discovery must surface the
196
- dir name as the id and ideally richer metadata (task, data, date) so the selector
197
- isn't three identical rows.
198
-
199
- ## Defaults & invariants to lock everywhere
200
-
201
- - Selected adapter defaults to **none** (Pi extension, web UI, server) — base model.
202
- - Unknown/typo'd id → **loud error** (already true server-side; surface it in the UI,
203
- don't silently fall back to base).
204
- - Stacking (`"a+b"`) is supported end-to-end; UI can expose it as multi-select.
205
- - An adapter must be **loaded** before it can be **selected**.
@@ -1,255 +0,0 @@
1
- # Batching perf path — promote continuous batching out of compat mode
2
-
3
- Status: P5 + two gate fixes LANDED (2026-07-02, same day). Since then
4
- (2026-07-03, executed via grammar-spec-batching-integration.md Phases D/E):
5
- the **P3 KV-budget admission slice LANDED** (`--kv-budget` aggregate
6
- projection gate, queue-don't-OOM,
7
- `/stats.batch.{pending_rows,kv_bytes,kv_budget_bytes}`,
8
- tests/batch-kv-budget.test.ts) and the **P0 bench harness LANDED** as
9
- `scripts/bench-feature-matrix.ts` (the six-cell composition matrix;
10
- `scripts/bench-serving-load.ts` stays a client-only stack-vs-stack tool).
11
- Also 2026-07-03: **Tier-0 universal models (plain full-attention, e.g.
12
- Llama) now batch** — per-row RoPE ported (`UniversalRope.applyDynamic`)
13
- and gated token-exact vs mlx-lm B=2
14
- (tests/batched-decode-parity.test.ts "Llama 3B Tier-0"); maskArray +
15
- sliding universal archs stay serial (unvalidated cells). Rest PLANNED —
16
- P0 COMPLETED 2026-07-04: extend-join (full-attention; own mlx-lm extend oracle) + vectorized all-greedy sampling (BIT-equal A/B).
17
-
18
- **Headline result (M1 Max 32 GB, loaded machine, 4 concurrent × 128 tok,
19
- median of 3, wall-clock aggregate; oMLX 0.4.5.dev1, same OptiQ snapshots):**
20
-
21
- | Model | oMLX | mlx-bun `--batch 4` | mean TTFT (theirs/ours) |
22
- | --- | --- | --- | --- |
23
- | MiniCPM5-1B | 339.0 | **345.4** | 397 / **119 ms** |
24
- | gemma-4-e4b | **89.8** | 87.1 (−3%) | 815 / **397 ms** |
25
- | Qwen3.5-4B | **100.9** | 99.5 (−1%) | 848 / **368 ms** |
26
-
27
- What landed (all bit-gated, `tsc` 0, suites green):
28
- 1. **P5 — SSM batched path** (`SSMCache.mergeRows`/`filter` in
29
- `src/model/qwen3-delta.ts`; "ssm" kind in `batch-scheduler.ts` merge/wrap/
30
- filter; per-row RoPE via `ropeOffsetArr` in `qwen3_5.ts`; gate relaxed in
31
- `generation-gateway.ts`, kill switch `MLX_BUN_BATCH_SSM=0`). Oracle:
32
- `tests/qwen35-batched-parity.test.ts` — solo-prefill+merge batched greedy
33
- is **token-for-token identical** to mlx-lm B=2 left-padded batch
34
- (`goldens/apple-m1-max/batched-golden-qwen35.json`). No ssm_mask needed:
35
- rows solo-prefill unpadded, decode feeds one real token per row.
36
- 2. **Logits processors batch** (`generation-gateway.ts`): the per-row sampler
37
- folds `makeLogitsProcessors` over a per-row device-side token history
38
- (generate()'s pushHistory), so repetition/presence/frequency penalties +
39
- logit_bias + min_p/XTC no longer route serial. Load-bearing discovery:
40
- **Qwen3.5 ships a default repetition penalty in generation_config.json**,
41
- which routed EVERY request serial under `--batch N` — the SSM port alone
42
- changed nothing over HTTP until this landed.
43
- 3. Diagnostic breadcrumb: `MLX_BUN_LANE_DEBUG=1` logs per-request lane
44
- decisions + shape (server.ts) — how (2) was found.
45
-
46
- Bench-method traps found while measuring — **both FIXED 2026-07-02**:
47
- - `serve`/`bench` accepted `--model` but IGNORED it (positional-only fuzzy
48
- query; auto-pick silently loaded the wrong model). Now: `--model
49
- <path-or-query>` is an explicit override — a directory with config.json
50
- loads via `scanSnapshot` (no registry dependency), anything else resolves
51
- like the positional. Precedence: `--model` > positional > `--query` >
52
- auto-pick.
53
- - Event-loop starvation during generation was the SERIAL lane, not batch
54
- (batch's drive loop yields per step — `/stats` answered in ~8 ms mid-
55
- generation once measured properly): serial decode is an unbroken microtask
56
- chain, so `/stats`//health/new accepts stalled for the whole generation
57
- (measured 2.5 s on a 512-token run). Fixed with a ≥25 ms rate-limited
58
- macrotask hop wrapped around onToken in `GenerationGateway.run`'s serial
59
- branch — covers all four API surfaces, streaming and not. After: `/stats`
60
- 10–44 ms mid-generation, decode unchanged (cpm5 serial 193 tok/s). Motivated by the oMLX head-to-head (M1 Max,
61
- 2026-07-02): at 4 concurrent requests oMLX's continuous batching gets
62
- 1.46–1.89× our aggregate throughput (Qwen3.5-4B 97.3 vs 66.5 tok/s;
63
- MiniCPM5 328 vs 173; e4b 85 vs 52) and our mean TTFT queue-inflates to
64
- 1.1–3.8 s. Single-stream we match or beat them (e4b 55.8 vs 53.7,
65
- Qwen3.5 66.3 vs 63.9, TTFT 3–7× better everywhere); the gap is entirely
66
- the serial lane going flat under parallel load. Full comparison + numbers:
67
- the 2026-07-02 oMLX report (Claude artifact) and `~/.omlx` bench scripts.
68
-
69
- Design stance: the batch lane grows the **same parity-tier DAG as the
70
- serial lane** (docs/design/parity-tier-dag.md). `--l1 --batch N` stays
71
- bit-exact mlx-lm-B=N compat mode forever (the oracle checkpoint); the
72
- default batch path becomes the optimized one. (`--l3` was deleted
73
- 2026-07-05 — docs/design/unified-engine-frontier-plan.md §6; output-changing
74
- experiments are now env-flag Lab items.) `--batch` remains a
75
- mode switch, not a load fallback.
76
-
77
- ## 1. Current state
78
-
79
- `src/serve/batch-scheduler.ts` (`BatchScheduler`) is a real
80
- continuous-batching engine: iteration-level scheduling, pipelined decode
81
- (`#step` via `ops.asyncEvalAll`, mlx-lm `GenerationBatch._step`
82
- semantics), chunked joiner prefill (2048-token chunks) interleaved with
83
- decode, per-row eviction (`filterKVRows` / `BatchedRotatingCache.filter`),
84
- per-row failure containment, drain-on-serial-waiter, `clearCache` every
85
- 256 steps. Bit-parity with mlx-lm B=N verified at B=2 across all 4 model
86
- paths (`tests/batched-decode-parity.test.ts`, `tests/batched-rotating.test.ts`,
87
- `tests/batch-scheduler.test.ts`) — plus, since 2026-07-03, Tier-0
88
- universal plain-full-attention archs (Llama 3B, token-exact vs mlx-lm B=2;
89
- see header).
90
-
91
- `src/serve/generation-gateway.ts` (`GenerationGateway.willBatch`) routes;
92
- one `AsyncMutex` is the GPU/`loraState` exclusion domain. The batch lane
93
- is deliberately compat mode: it calls `model.forwardHidden` /
94
- `logitsFromHidden` directly, bypassing `generate()` and every perf lever.
95
-
96
- ### Levers the batch lane lacks, and why
97
-
98
- | Lever | Serial mechanism | Why batch can't use it today |
99
- | --- | --- | --- |
100
- | Quantized KV (`kv_config`) | `maybeQuantizeKv` (`src/generate.ts:101`) | `QuantizedKVCache` (`src/model/gemma4-base.ts:419`) has no dynamic-B ops (no merge/filter/extend/temporalView over the packed/scales/biases triples); `willBatch` hard-gates `kvQuant` → serial |
101
- | ~~Perf-kernel fused decode SDPA~~ | ~~dispatched from `quantizedSdpa`~~ | **DELETED 2026-07-05** (kernel + `--perf-kernel` removed; unified-engine-frontier-plan.md §6) — no longer a lever, batched or serial |
102
- | ~~Fused N-tiled SDPA (`--fused-decode`)~~ | ~~`quantizedSdpaTiled`~~ | **DELETED 2026-07-05** (measured 1.00×; unified-engine-frontier-plan.md §6) |
103
- | Compiled decode | `CompiledDecode.for()` (`generate.ts:518-525`) | trace is `[1,1]`; B changes every join/evict |
104
- | Prompt-cache prefix reuse | `PromptCache.take/put` in `runGeneration` | bypassed by design in v1; every batched row solo-prefills (`cachedTokens: 0`) |
105
- | Vectorized sampling | n/a | `#step` samples per row (B slice+sample+concat graphs per token) |
106
-
107
- Two more gap sources specific to the benchmark:
108
-
109
- - **Qwen3.5 never batches at all** — **FIXED by P5 (2026-07-02, header
110
- item 1)**: `SSMCache.mergeRows`/`filter` landed; Qwen3.5 batches
111
- token-exact vs mlx-lm B=2.
112
- - Joins re-merge the whole batch (O(B·S) slice storm, `#mergeJoiner`)
113
- instead of mlx-lm's `extend`.
114
-
115
- ## 2. Gap analysis: quantized KV at B>1
116
-
117
- - `QuantizedKVCache` already stores `[B,H,S,D']` triples; missing is
118
- dynamic-B surgery. New `src/model/batched-quant.ts`:
119
- `BatchedQuantizedKVCache` = quantized twin of `BatchedRotatingCache`
120
- (merge/extend/filter/extract/temporalView = the `mergeKVRows`/`filterKVRows`
121
- B-axis ops from `src/model/batched-mask.ts:171,196`, applied to all three
122
- triples). Left-pad must align to quantization group boundaries — pad in
123
- whole 256-token STEP granules (both caches already grow in STEP=256).
124
- - Conversion point mirrors `maybeQuantizeKv`: row solo-prefills bf16,
125
- `toQuantized` per layer per `kv_config` before merge.
126
- - **Scope cut**: shipped `kv_config.json` quantize full-attention layers
127
- only ⇒ P1 = quantized full layers + bf16 `BatchedRotatingCache` sliding
128
- layers (the serial `config` composition). Batched rotating-quant deferred.
129
- - Correctness fallback: `quantizedSdpaUnfused` (`gemma4-base.ts:1273`) is
130
- already B-generic with `mask.mode === "array"`; needs the batched mask
131
- broadcast as `[B,1,1,1,N]` for the 5-D GQA score reshape.
132
- - ~~Perf kernel at B>1~~ — OBSOLETE: the perf kernel was deleted 2026-07-05
133
- (unified-engine-frontier-plan.md §6); `quantizedSdpaUnfused` is the
134
- quant-batched path, full stop.
135
- - RoPE per-row positions already solved (`ropeOffsetArr`, the S1b.1 fix).
136
-
137
- ## 3. Phases
138
-
139
- ### P0 — Baseline + engine quick wins (no new numerics) [S–M, 2–4 days]
140
- - Bench harness — **DONE 2026-07-03**: landed as
141
- `scripts/bench-feature-matrix.ts` (six-cell composition matrix over live
142
- SSE, TTFT p50/p95 + aggregate tok/s; integration-plan Phase E).
143
- `scripts/bench-serving-load.ts` exists as the client-only stack-vs-stack
144
- tool and stays one. Clean-machine run remains Josh-gated.
145
- - Land `extend` join (full-attention twin + `BatchedRotatingCache.extend`);
146
- oracle: extend `scripts/gen-batched-dynamic-golden.py` (mlx-lm `extend`).
147
- **LANDED 2026-07-04** for full-attention layers (`extendKVRows`,
148
- mlx-lm `BatchKVCache.extend` semantics; own oracle
149
- `scripts/gen-batched-extend-golden.py`, token-for-token CPM + Llama;
150
- `MLX_BUN_BATCH_EXTEND=0` = re-merge). `BatchedRotatingCache.extend`
151
- (sliding layers on Gemma joins) remains open.
152
- - Vectorize the homogeneous-sampler fast path in `#step` (all-greedy →
153
- one `ops.argmaxAxis` on `[B,V]`); mixed rows keep the loop. **LANDED
154
- 2026-07-04** (BIT-equal A/B, `MLX_BUN_BATCH_VEC_SAMPLE=0`;
155
- tests/batch-vec-sample.test.ts).
156
- - Exit: gated suites unchanged; B=4 baseline recorded; joins not O(B·S).
157
-
158
- ### P1 — Quantized KV at B>1 [M–L, 1–1.5 wk]
159
- - `src/model/batched-quant.ts` + scheduler wiring (layer kind
160
- `"full" | "rot" | "qfull"`); solo-prefill bf16 → `toQuantized` → extend.
161
- - Batched mask 5-D broadcast; validate `quantizedSdpaUnfused`.
162
- - `willBatch`: `kvQuant` stops being a serial gate when quant layers are
163
- full-attention-only.
164
- - Exit (parity ladder; no external oracle — mlx-lm quant batching NYI):
165
- (1) B=1-row batched-quant **bit-exact vs serial `--kv-quant config`**;
166
- (2) row-isolation within the calibrated bf16-reduction envelope;
167
- (3) KL vs batched-bf16 ≤ serial quant-vs-bf16 KL; 6-task eval unchanged;
168
- (4) `--l1 --batch N` suites untouched.
169
-
170
- ### P2 — Perf kernel at B>1 + fallback matrix [OBSOLETE — kernel deleted 2026-07-05]
171
- The perf kernel (and `--perf-kernel` / its frozen-oracle tests) was deleted
172
- 2026-07-05 (unified-engine-frontier-plan.md §6): it never beat the L1
173
- baseline in a paired A/B, so there is nothing to port to B>1. The planned
174
- "batched frozen-oracle envelope" gate is obsolete with it. Quant-batched
175
- decode stays on `quantizedSdpaUnfused` (P1).
176
-
177
- ### P3 — Admission, prompt cache, adapters, defaults [M–L, 1–1.5 wk]
178
- - KV-budget admission — **LANDED 2026-07-03** (integration-plan Phase D):
179
- `--kv-budget <GB>` caps the batch's aggregate projected KV
180
- (prompt+max_tokens per row, window-capped); over-budget joiners QUEUE
181
- (FIFO), never OOM; a single request over the budget alone rejects with
182
- an actionable error; `/stats.batch.{pending_rows,kv_bytes,kv_budget_bytes}`;
183
- gated by tests/batch-kv-budget.test.ts. `--prompt-concurrency` still open.
184
- - Prompt-cache reuse under batching: `take` at admit, forward uncached
185
- tail; on finish/evict `extract(row)` → solo caches → `put`. Quantized
186
- entries are non-trimmable mid-group — match on full-prefix or STEP
187
- boundaries. Real `prefillMs/decodeMs`; `BatchStats.cachedTokens` real.
188
- - Single-adapter batches: adapter-spec = batch compatibility key
189
- (join if match, drain if not; `loraState.active` set once per epoch).
190
- - Default review: if batch lane at 1 live row is within ~2% of serial,
191
- propose flipping the shipped default to `--batch 4` (Josh decision,
192
- numerics change vs serial default); else README recommends `--batch 4`
193
- for agent workloads.
194
- - Exit: TTFT p50 <500 ms / p95 <1.5 s at 4 concurrent short prompts;
195
- `cached_tokens>0` on repeats under `--batch 2`; oversized-request
196
- queueing test (no OOM); B=4 aggregate within 15% of oMLX on cpm5 + e4b.
197
-
198
- ### P4 — Burst decode [REVISED after reading + porting oMLX's source]
199
- - **oMLX-style event-loop bursting: BUILT AND REFUTED (2026-07-02).** Read
200
- their actual implementation (`omlx/engine_core.py _step_burst` — full
201
- source ships in `/Applications/oMLX.app/Contents/Resources/omlx/`): it is
202
- NOT a GPU trick — it runs K `scheduler.step()` calls per
203
- `run_in_executor` hand-off with adaptive time budgets (100 ms single /
204
- 30 ms concurrent, 64-step cap) because each per-token hand-off ping-pongs
205
- the GIL with asyncio/uvicorn (~1 ms/token; their comment: 80 vs 74 tok/s).
206
- Ported faithfully to `BatchScheduler.#drive` and measured: **regression**
207
- — cpm5 B=4 345→289, batch-lane B=1 149→121, TTFT +~100 ms (SSE flushes
208
- wait out the budget). Bun's `setImmediate` hop costs microseconds — no
209
- GIL to dodge — so the port kept the cost and lost the win. Reverted;
210
- numbers restored (cpm5 B=4 349.1). Breadcrumb comment left at the drive
211
- loop. This also explains their ~21-token SSE bursts (measurement trap in
212
- the comparison report).
213
- - **Still open (different mechanism, not refuted): device-side step
214
- chaining** — generalize `generateInner`'s pipeline to depth k (k chained
215
- step graphs fed device-side, one host sync per k tokens). That attacks
216
- OUR real per-token cost (FFI graph-build + sync readback), which is not
217
- what oMLX's burst attacks. Exit unchanged: bit-exact vs depth-1, cpm5
218
- single-stream gap vs oMLX (178 vs ~214) is the target; no regression at
219
- k=1.
220
-
221
- ### P5 — SSM batched path (Qwen3.5) [DONE 2026-07-02 — see header]
222
- - Landed in one day, much cheaper than the ~1 wk estimate: `ssm_mask` under
223
- left-padding turned out unnecessary (our join model solo-prefills unpadded,
224
- so SSM state never sees a pad token; mlx-lm's masked-pad prefill is
225
- equivalent — proven by the token-exact oracle match). Batch ops are plain
226
- B-axis concat/take on the two state slots.
227
- - Bonus fix that the exit benchmark forced out: logits processors now batch
228
- (per-row fold; Qwen3.5's default repetition penalty had been routing every
229
- request serial). Exit numbers in the header table.
230
-
231
- Ordering: P4 and P5 independent (P5 first if the Qwen3.5 headline matters
232
- most); P2 dropped (kernel deleted). Total ~4–5 wk serial, ~3 wk with P4/P5
233
- parallel.
234
-
235
- ## 4. Risks
236
-
237
- - **No external oracle for quant-batched** — mitigated by the P1 parity
238
- ladder; `--l1 --batch N` stays the always-reachable bit-exact checkpoint.
239
- - ~~Metal kernel at B>1~~ — risk retired with P2 (perf kernel deleted
240
- 2026-07-05).
241
- - **Memory**: B× quant caches + prompt cache + weights on 24 GB — P3
242
- admission is the backstop; mirror mlx-lm's
243
- `trim_to(total − activeBatchBytes)` so LRU and live batch share one
244
- ceiling. GPU OOM is uncatchable — queueing test before default changes.
245
- - **Ring-wrap at B>1 short-context-verified only**: land a long-context
246
- bf16 batched golden before any P3 default flip. Quantized rotating out
247
- of scope.
248
- - **Default flip changes numerics** vs today's serial default — explicit
249
- Josh decision with doc trail.
250
- - **Benchmark hygiene**: quotable numbers via `bench-serving-load.ts` on a
251
- clean machine; TTFT at the SSE first-token event, not server-side
252
- (oMLX's own server log over-reports by excluding ~180 ms/req cache
253
- bookkeeping — measure them wall-clock only).
254
- - **Regression fence**: all existing gated batch suites + the serial suite
255
- green unchanged at every phase.