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,318 +0,0 @@
1
- # DSpark speculative decoding for mlx-bun
2
-
3
- Implementation of DeepSeek's **DSpark** ("Confidence-Scheduled Speculative
4
- Decoding with Semi-Autoregressive Generation", `DSpark_paper.pdf`) against this
5
- repo's models. DSpark = a **DFlash** parallel backbone (KV-injected, all γ draft
6
- tokens in one pass) **+ a lightweight sequential head** (Markov, Eq 5, or RNN,
7
- Eq 6) that adds intra-block dependency (the *semi-autoregressive* part) **+ a
8
- confidence head → STS calibration → hardware-aware prefix scheduler**. Decode
9
- becomes draft → verify → accept; output is lossless (bit-identical to the target
10
- at any temperature); only decode speed (τ) changes.
11
-
12
- > **DFlash vs DSpark:** DFlash is *only* the parallel backbone (prior work) and
13
- > suffers suffix decay. DSpark adds the sequential head + confidence-scheduled
14
- > verification. The sequential Markov head is what makes this DSpark.
15
-
16
- ## Two implementations (flag-selected, both in `src/spec/dspark/`)
17
-
18
- | | `module.ts` (v1) | `module-dflash.ts` (v2, **the real one**) |
19
- |---|---|---|
20
- | context | **single vector** (final hidden at the anchor) fused as one token | faithful **multi-layer KV injection** (Eq 2–3): H_ctx = RMSNorm(W_c·[H^{l1};…;H^{lm}]) over the FULL context, re-projected into every draft layer's K/V |
21
- | layers | 2 | 5 |
22
- | status | **superseded** — dropped the paper's core idea (kept as baseline) | the actual DSpark model (Markov variant) |
23
-
24
- v1 was a shortcut that removed the load-bearing mechanism; it never tested the
25
- paper. v2 is faithful. Checkpoints stamp `variant` in `dspark.json`; the runtime
26
- loads the matching module. (`variant:"dflash"` should be renamed to `"dspark"` —
27
- see Open items.)
28
-
29
- ## The mechanism (v2, paper Eq 2–3)
30
-
31
- - **Eq 2 (context):** tap the target's hidden states at layers `{l1…lm}` (e4b
32
- default {20,31,41,42}, the last = post-finalNorm), concat on the feature axis,
33
- project to the draft width: `H_ctx ∈ [L_ctx, d]`. Extraction via a parity-safe
34
- `model.hiddenTap` (no-op when off) added to `gemma4.ts` `forwardLayers`.
35
- - **Eq 3 (injection):** every draft layer forms `K_i=[W_i^K H_ctx; W_i^K H_d]`,
36
- `V_i` likewise, `Q_i=W_i^Q H_d`; the γ-block queries attend bidirectionally
37
- over `[context ++ block]`. Context is read-only memory re-projected fresh each
38
- layer. The draft projects the target's HIDDEN STATES with its OWN Wk/Wv (not
39
- the target's K/V — that's the separate `GemmaAssistantDrafter` baseline).
40
- - **Markov head (Eq 5):** `B_k = W1[x_{k-1}]·W2`, r=256, W2=0 init → starts as
41
- pure DFlash, τ climbs as intra-block dependency is learned.
42
- - **Confidence head (Eq 7):** `c_k = σ(wᵀ[h_k; W1[x_{k-1}]])`.
43
- - **Loss (Eq 9–12):** `0.1·CE + 0.9·TV + 1.0·BCE`, position-weighted.
44
- - **Lossless verify** (`generate-dflash.ts`): greedy = longest-prefix match;
45
- temp>0 = `min(1,p/q)` accept + residual `norm(relu(p−q))` resample. Growing
46
- multi-layer H_ctx buffer; rollback via `trim(n, bypass=true)` (physically
47
- slices the rejected tips past the sliding window — see `gemma4-base.ts`).
48
-
49
- ## Results & findings (measured, M1 Max 32GB, e4b OptiQ-4bit)
50
-
51
- **The architecture is CORRECT — proven by an overfit test.** Trained on 3
52
- articles, eval on the same 3: per-position acceptance **~0.75**, τ **3.24** —
53
- squarely paper-range (0.6–0.9 / τ 3–4). So the faithful DFlash, correctly
54
- implemented, reaches the paper's numbers.
55
-
56
- **Getting there required finding a real bug — an off-by-one in the p^t target
57
- alignment** (the TV loss, weight 0.9). Block position k predicts `x_{t+1+k}`,
58
- whose target distribution is `softmax(LM_head(h_{t+k}))`; the data gathered the
59
- hidden at `t+k+1` (shifted forward one). Present in **both v1 and v2** — so every
60
- τ number before the fix was trained against a shifted target. Fixed in
61
- `data-dflash.ts` (`blockIdx = g+k`, k=0..γ-1). Before fix: per-pos ~0.08. After:
62
- ~0.17 generalizing, ~0.75 overfit.
63
-
64
- **Two things stand between "correct architecture" and "net speedup," now
65
- separated:**
66
- 1. **Data.** 160 articles → generalizes only to per-pos ~0.17 (plateaus, doesn't
67
- climb); overfit on 3 proves capacity. The paper uses 1.3M samples × 10 epochs.
68
- → needs real data scale.
69
- 2. **Target speed (the decisive one).** Spec decode amortizes the fixed draft
70
- overhead only when the *target forward* is expensive. Measured decode:
71
- **e4b 45.9 tok/s, 12b 27.5 tok/s, 27B ~15 tok/s.** On fast e4b the draft
72
- overhead dominates → even τ≈3 nets ~0.4–0.5× (SLOWER). On a slow target (27B,
73
- ~67 ms/forward) that overhead is a small fraction → the same τ≈3 could land
74
- ~2–3× (15 → ~35–45 tok/s). **e4b is close to the worst case for spec decode;
75
- the 27B agentic workload is the ideal target and the real motivation.**
76
-
77
- Also confirmed: width is NOT the ceiling (`dDraft=1024` and `2560` both plateau
78
- ~0.17 pre-scale); the parity-safe `hiddenTap` doesn't change target numerics;
79
- long-context spec rollback works via the `trim(n, bypass)` physical-slice.
80
-
81
- ## Serve integration (Phase 1, 2026-07-06)
82
-
83
- DSpark is now serve-loadable behind the same `--draft-model` seam as every
84
- other drafter — one verify/accept executor, no lane routing. The seam
85
- (`src/spec/source.ts`) was extended from its two-model-only shape to carry the
86
- target state KV-borrowing drafters need (it previously could NOT host these,
87
- despite the header claiming so):
88
-
89
- - `DraftProvider.open({sampler, target:{model,caches}})` — the source gets the
90
- target model + its live caches.
91
- - `DraftSource.tapLayers?` — when set (DSpark), the serve loop taps the
92
- target's prefill + every verify forward and hands the captured `[1,L,m*H]`
93
- context to `prefill(ids, ctxML)` / `commit(n, kAccept, vCtxML)` (ownership
94
- transfers). `draft(feed, n, stepBase, anchorHidden?)` — the assistant borrows
95
- the target anchor hidden; DSpark drafts from its growing `H_ctx`.
96
-
97
- Provider selection is by artifact kind (`src/server.ts` `detectDraftKind`):
98
- `dspark.json` → **DflashProvider** (`src/spec/dflash-source.ts`), a `*_assistant`
99
- config → **AssistantProvider** (`src/spec/assistant-source.ts`, the optiq
100
- KV-borrowing Gemma drafter — ships a real ~1.09× γ=1 decode win with no
101
- training), else **TwoModelProvider**. `--draft-kind` overrides. The server pins
102
- `--num-draft-tokens` to the DSpark checkpoint's trained block width
103
- (`cfg.gamma`) so the serve loop never asks for more positions than were trained.
104
-
105
- **Unified-engine alignment (per-slot readiness).** This is the
106
- `DraftSource` seam the unified-engine plan names as the substrate for its
107
- frontier "per-slot drafting behind the `DraftSource` seam" row
108
- ([docs/design/unified-engine-frontier-plan.md](unified-engine-frontier-plan.md)).
109
- Spec still forces the serial lane today (upstream `is_batchable = draft is
110
- None`); Phase 1 did NOT change that — it made the seam *carry per-source target
111
- state* (each source binds its own `TargetView{model,caches}`, anchor hidden,
112
- and tapped H_ctx), which is exactly what a batched executor would need to open
113
- one `DraftSource` per row. So the seam is forward-compatible with per-slot
114
- drafting; the hard part that remains is the executor, not the seam — variable
115
- accept-length per slot breaks a uniform-B step (see
116
- [spec-decode-larger-targets.md](spec-decode-larger-targets.md) caveats), so
117
- per-slot spec is a real batched-scheduler project, not free composition.
118
-
119
- **v1 acceptance caveat:** the serve loop verifies with TOKEN-MATCH acceptance
120
- (mlx-lm/optiq style, lossless at any temperature), NOT the paper's
121
- distribution-level rejection-sampling verify. Greedy is identical; temp>0 is
122
- still lossless but lower-acceptance. The richer rejection-sampling verify stays
123
- in the standalone `generate-dflash.ts` (the measure script). See
124
- [[dspark-seam-kv-borrowing]].
125
-
126
- ## Files
127
- - `src/spec/source.ts` — the extended DraftSource/DraftProvider/TargetView seam.
128
- - `src/spec/assistant-source.ts` / `dflash-source.ts` — the two KV-borrowing providers.
129
- - `src/spec/serve-loop.ts` — the shared verify/accept executor (now tap-aware).
130
- - `src/model/gemma4.ts` — `hiddenTap`/`captureLayer` (parity-safe multi-layer tap).
131
- - `src/model/gemma4-base.ts` — `trim(n, bypass)` on rotating caches (spec rollback past the window).
132
- - `src/spec/dspark/module-dflash.ts` — faithful DFlash+Markov+confidence module.
133
- - `src/spec/dspark/data-dflash.ts` — multi-layer shards, variable-length prefix context, left-pad masking.
134
- - `src/spec/dspark/generate-dflash.ts` — growing-H_ctx lossless decode loop.
135
- - `src/spec/dspark/{loss,sample}.ts` — loss (Eq 9–12) + temp>0 sampling (shared with v1).
136
- - `src/spec/dspark/{module,data,generate}.ts` — v1 single-vector (superseded baseline).
137
- - `scripts/dspark-regen-dflash.ts` / `dspark-train-dflash.ts` (`--resume`) / `dspark-measure-dflash.ts`.
138
- - `scripts/dspark-dflash-smoke.ts` — CPU smoke (16/16). `scripts/dspark-smoke.ts` — v1 smoke (33/33).
139
-
140
- ## Run sequence (faithful DFlash; GPU = Josh runs)
141
- ```
142
- bun scripts/dspark-regen-dflash.ts --topics <topics.txt> --out <data> --max-resp 320
143
- bun scripts/dspark-train-dflash.ts --data <data> --out <ckpt> --iters 8000 --batch 8 [--resume] [--ddraft 2560]
144
- bun scripts/dspark-measure-dflash.ts --drafter <ckpt> --data <prompts.jsonl> # τ + tok/s vs vanilla
145
- ```
146
-
147
- ## Paper components — Phase 2 (2026-07-06): the paper is CODE-COMPLETE
148
-
149
- All remaining components landed the same day (multi-agent build, adversarially
150
- reviewed; smoke 21/21, dspark test files 17/17, real-weights serve gate 3/3):
151
-
152
- - **Alg 1, single-user form — confidence-scheduled draft-length pruning.**
153
- `forwardInfer` drops position k (and truncates the block) when
154
- `c_k < thresholds[k] ?? minConf`; position 0 always survives. Activation is
155
- checkpoint-driven: STS thresholds in `dspark.json` (`cfg.sts`) or the
156
- `MLX_BUN_DSPARK_MINCONF` env override — **uncalibrated checkpoints draft
157
- fixed-γ exactly as before** (zero behavior change). Losslessness invariant:
158
- pruning changes how many positions the target verifies, never what's emitted
159
- (gated: truncation-never-redraw, prefix-identical to the unpruned block).
160
- - **Variable-length draft contract.** `DraftSource.draft()` returns 1..n —
161
- the RETURN LENGTH is authoritative; the serve loop and `dflashGenerate`
162
- verify/trim/re-feed over `d = drafts.length`. Two-model/assistant return
163
- exactly n (unchanged).
164
- - **STS calibration (§3.2.1).** `src/spec/dspark/calibration.ts`
165
- (`fitStsThresholds`: per position, smallest τ whose Laplace-smoothed
166
- `P(accepted | conf ≥ τ)` ≥ target; pos 0 → 0, under-sampled → 0 (don't
167
- prune on thin evidence), unreachable target → 1.0) +
168
- `scripts/dspark-calibrate.ts` (GPU: greedy unpruned rounds via
169
- `dflashGenerate`'s `onRound` hook → fit → patch `config.sts` into
170
- dspark.json in place). ⚠ estimator shape is our reading — paper PDF absent.
171
- - **RNN sequential head (Eq 6 — ⚠ design-doc-faithful shape, paper PDF
172
- absent, flagged in code).** `cfg.seqHead: "markov"|"rnn"`; Elman recurrence
173
- `s_k = tanh(s_{k-1}·wH + E[x_{k-1}] + bH)`, `B_k = s_k·wO`, `wO` zero-init
174
- (starts as pure DFlash, like markov.w2); shares `markov.w1` as the token
175
- embedding (also feeds the confidence head). Same loss, autograd-gated;
176
- init-equivalence gate: rnn ≡ markov token-for-token at init.
177
- Train with `scripts/dspark-train-dflash.ts --seq-head rnn`.
178
- - **Draft loop TIGHTENED** (handoff item 3): greedy tokens chain ON-DEVICE
179
- (one concat + one host read after the loop, was γ syncs); confidence reads
180
- deferred+batched when pruning is inactive (per-position only when Alg 1
181
- needs the answer mid-loop — inherent); `collectLogits: false` on the serve
182
- path skips the per-position draft-logits materialization entirely
183
- (`draftLogits` now optional). Gate: bit-identical to the pre-tightening
184
- loop (pinned reference tokens, tests/dspark-infer-loop.test.ts).
185
- - **Rename + central loader.** Canonical variant is now `"dspark"`
186
- (`save()` stamps it; `load()` accepts legacy `"dflash"`);
187
- `src/spec/dspark/loader.ts` `loadDsparkDrafter()` dispatches by variant and
188
- refuses v1 single-vector checkpoints with a pointer at the v2 trainer. The
189
- serve path and the bench (`bench-feature-matrix.ts` via the server's
190
- exported `detectDraftKind`) both go through it.
191
-
192
- ## DeepSpec ground truth (2026-07-06 audit — paper arXiv:2607.05147 + source)
193
-
194
- DeepSeek open-sourced the reference implementation (**DeepSpec**, MIT,
195
- github.com/deepseek-ai/DeepSpec) + trained drafters incl.
196
- **`deepseek-ai/dspark_gemma4_12b_block7`** (for google/gemma-4-12B-it — our
197
- primary target; 6.86 GB bf16 single safetensors, γ=7). Two audit agents
198
- verified our build against the paper AND the source. Verdicts:
199
-
200
- - **MATCH the paper exactly:** Eq 2-3 (context + KV injection), the
201
- anchor++masks block, Eq 5 Markov (r=256), Eq 7 confidence (we add a
202
- zero-init bias — superset), **Eq 8-12 loss to the letter** (0.1/0.9/1.0,
203
- w_k=exp(−(k−1)/γ), analytic c*). Our defaults = their production config
204
- (γ=5, 5 layers, Markov head).
205
- - **Eq 6 RNN differs:** paper is a gated cell over
206
- `z_k=[s_{k−1}; W₁[x_{k−1}]; h_k]`, fused (2r+d)→3r projection,
207
- `B_k=W₂ᵀtanh(W_o z_k)` (conditioned on the backbone hidden). Ours is
208
- ungated Elman without h_k — kept as variant `"rnn"` (paper calls the gains
209
- "marginal"; production uses Markov); DeepSpec also ships an RNNHead of the
210
- gated form if we ever want the faithful one.
211
- - **§3.2.1 STS + Alg 1: the paper's machinery is NOT in the released code.**
212
- STS = Sequential Temperature Scaling (per-position temperatures by ECE grid
213
- search on cumulative survival products); Alg 1 = expected-throughput
214
- maximization over survival products with a profiled SPS(B) cost table —
215
- both live in their UNRELEASED production serving layer. The released
216
- reference ships **per-position threshold truncation on sigmoid confidence
217
- (default 0 = off), ℓ=0 allowed** — i.e. the released reference is shaped
218
- like OUR scheduler. Our threshold fit (`calibration.ts`) is therefore a
219
- reasonable calibrator for the reference-shaped scheduler, NOT the paper's
220
- STS; renamed framing in code comments. Paper-faithful STS/Alg-1 = a future
221
- Lab item if the survival-product scheduler ever earns a default.
222
- - **The reference drafter is a DIFFERENT module than ours:** a full
223
- Gemma4-shaped 5-layer transformer — hidden 3840, 16 Q heads × head_dim 512,
224
- **1 KV head with K≡V (`attention_k_eq_v`, no v_proj tensor; scale-less
225
- v_norm)**, attention scale 1.0 (QK-norm), RoPE partial_rotary 0.25 θ=1e6,
226
- per-layer learned `layer_scalar`, scaled embed (×√3840, not stored), UNTIED
227
- lm_head, final softcap 30, and an **incremental context-KV cache** (stores
228
- projected target-context K/V; crops the block's noise K/V each round).
229
- Sequential Markov sampling + confidence head (input [h; w1[prev]] = 4096)
230
- as in the paper. Checkpoint: 74 tensors, bf16, **no key prefix**;
231
- config.json `architectures: ["Gemma4DSparkModel"]` (loader detection key).
232
- **12B tap layers: `target_layer_ids [5,17,29,41,46]`** (their layer_id+1
233
- output convention == our tapLayers index; no final-norm sentinel).
234
- - **THE ORACLE:** their verify is leaky rejection sampling that at
235
- temperature 0 **degenerates to exact argmax token-match — RNG-free,
236
- token-for-token deterministic**. Their eval emits per-round
237
- proposal/accepted/next-token traces. Protocol: run their eval at temp 0
238
- (+ a threshold arm) on fixed prompts → JSONL round fixtures → our port
239
- must match round-for-round. Same discipline as the D2 diffusion oracle,
240
- minus RNG parity entirely.
241
-
242
- **Consequence:** DSpark is no longer oracle-less L3. The port of their module
243
- (the `deepspec` variant) + their checkpoint gives a bit-exact-gateable path to
244
- a REAL 12B speedup with zero training. Our own trainable module remains the
245
- research path for custom targets. Training recipe reference: open-perfectblend
246
- regenerated by the target, 10 epochs, batch 512, lr 6e-4, position decay
247
- exp(−pos/4), ~38 TB cached hiddens — i.e. our "thousands of generations" data
248
- estimate was directionally right and their scale is far beyond a laptop run.
249
-
250
- ## 27B program (2026-08-19): the TQ quant unlocks the real target
251
-
252
- The turboquant-weights campaign shipped rot-sens48 Qwen3.8-27B at 17 GB
253
- (docs/design/turboquant-weights.md) — and it changes both DSpark blockers
254
- this doc recorded:
255
-
256
- - **Target speed:** measured ~11.8 tok/s decode. This is the slow-target
257
- regime where the doc's own analysis says τ≈3 nets 2–3× (≈25–35 tok/s) —
258
- "the ideal target and the real motivation," now actually runnable.
259
- - **Training feasibility:** "27B memory-infeasible to train on 24 GB" (open
260
- items below) assumed the bf16 trunk (54 GB). With the quantized trunk
261
- FROZEN (inference-only: hidden taps + teacher logits) and only the small
262
- drafter (~5 layers) holding grads/optimizer state in bf16, 27B-target
263
- drafter training fits a 32 GB machine. The premise is dead; the line in
264
- the gated item is superseded by this section.
265
- - **Rollback risk retired:** the stock Qwen-MTP spec path already serves
266
- this model losslessly (68% per-token acceptance), so verify/rollback
267
- across the 48 DeltaNet recurrent-state layers is proven engine behavior,
268
- not new work.
269
-
270
- Two tracks, cheap-first:
271
-
272
- - **Track A — confidence-scheduled verification on the stock MTP head (no
273
- drafter training).** Dynamic per-step block length instead of fixed
274
- `--num-draft-tokens`: v0 confidence = the draft's own token probability
275
- (no trained head), mapped through STS calibration
276
- (`scripts/dspark-calibrate.ts` pattern) to expected acceptance; the
277
- scheduler stops drafting when marginal expected-accept drops below the
278
- verify amortization point. Calibration needs thousands of
279
- (confidence, accepted?) pairs from served traffic — hours, not weeks.
280
- Gate: paired A/B decode tok/s vs fixed-γ MTP on the same prompts;
281
- lossless by construction (verify unchanged).
282
- - **Track B — true DSpark drafter for the 27B (multi-token single-pass).**
283
- The faithful v2 (KV injection, Markov head) retargeted: `hiddenTap` added
284
- to `qwen3_5.ts` forwardLayers (gemma4 has it; parity-safe no-op when
285
- off), tap-layer selection for the 64-layer trunk (bias toward the 16
286
- full-attention layers + post-finalNorm), data generated by TEACHER-FORCED
287
- passes over corpus text with the frozen quantized trunk (every position
288
- of a forward yields a sample — data scale is forward-passes over tokens,
289
- not generate calls). Open idea: distill from / initialize against the
290
- shipped MTP head to shortcut the acceptance floor. Gate: generalizing
291
- per-position acceptance must beat the MTP head's 0.68 or Track B loses to
292
- Track A + MTP and gets dropped.
293
-
294
- Sequencing: after the benchmark campaign closes (GPU is owned by the
295
- certification suite). Track A first — it's measurable in a day and its
296
- calibration data collection is also Track B's baseline instrumentation.
297
- Adapters/DSpark drafts trained on the TQ artifact are married to its
298
- rotation basis (fine for our bundle; note on the card if published).
299
-
300
- ## Open items / next
301
-
302
- - [x] Serve integration (Phase 1) · [x] Alg-1-shaped scheduler (matches the
303
- RELEASED reference; paper Alg-1 unreleased) · [x] threshold calibration ·
304
- [x] RNN head (Elman variant; paper's gated cell differs — documented above)
305
- · [x] loop tightening · [x] `dspark` rename + loader —
306
- **everything buildable without a GPU is DONE.**
307
- - [ ] **DeepSpec-variant port** (`Gemma4DSparkModel`) + loader detection +
308
- oracle scripts — the no-training path to the 12B speedup (in progress).
309
- - [ ] **Josh-gated GPU (the payoff):** real data scale (thousands of
310
- on-distribution generations) + **retarget to 12B** (`--model` +
311
- `--tap-layers 5,17,29,41,46` (DeepSeek's trained taps); ~~27B memory-infeasible to train on 24 GB~~ (superseded — the 17 GB TQ trunk makes it feasible; see "27B program" above), kept
312
- dim-generic) + train (`--seq-head` A/B optional) + **calibrate**
313
- (`scripts/dspark-calibrate.ts`) + live-τ measure. Recipe:
314
- [docs/archive/investigations/dspark-handoff.md](../archive/investigations/dspark-handoff.md).
315
- - [ ] Verify the Eq 6 / §3.2.1 shapes against the actual paper when the PDF
316
- is available (both flagged in code).
317
-
318
- Full session handoff: `docs/archive/investigations/dspark-handoff.md`.
@@ -1,301 +0,0 @@
1
- # Grammar × spec decode × batching — integration plan + feature-matrix benchmark
2
-
3
- Status: **ALL PHASES EXECUTED** (A/B/C/E 2026-07-03; D completed
4
- 2026-07-04: KV-budget admission + vectorized homogeneous sampling +
5
- extend-join, each oracle/A-B gated — see the execution log). Remaining
6
- follow-ups live in "Later / parked" with their triggers.
7
- Execution log at the bottom of this doc. Successor plan tying together
8
- three workstreams so they compose deliberately instead of colliding:
9
-
10
- - **Structured output** — SHIPPED on main (serial + batched lanes;
11
- docs/design/structured-output.md). Open debt: B2 model-gated batch
12
- tests + bench, F4 compiler cache, F5 regex, F6 escaping, F7 structural
13
- tags, U1/U2 engine upgrades.
14
- - **Batching perf path** — P5 done; P0–P4 open
15
- (docs/design/batching-perf-path.md).
16
- - **Speculative decoding** — research proven (assistant drafter bit-exact
17
- vs optiq, 12B γ=1 ≈ 1.09× win, docs/design/spec-decode-larger-targets.md;
18
- DSpark architecture verified, docs/archive/investigations/dspark-handoff.md) but
19
- **not served**: no `--draft-model`. Build plan already written:
20
- docs/design/mlx-lm-tool-parity-plan.md §7 (the `DraftSource` interface,
21
- two-model L1 source, serve loop, admission, prompt-cache composition).
22
-
23
- This doc does NOT restate those plans — it sequences them, defines the
24
- composition rules where they meet, and specifies the benchmark that
25
- measures every combination. The per-feature designs stay canonical in
26
- their own docs.
27
-
28
- ## The composition matrix (v1 target state)
29
-
30
- `GenerationGateway.willBatch` is the single routing point; every rule
31
- below is one predicate there. "Serial" routing is a v1 rule, not a
32
- limitation to hide — upstream mlx-lm routes spec serial too
33
- (`is_batchable = draft_model is None`, server.py:371).
34
-
35
- | Request shape | Lane | Mechanism | Status |
36
- | --- | --- | --- | --- |
37
- | plain | batch (or serial) | shipped engine | DONE |
38
- | grammar only | batch (or serial) | B1 per-row matchers / serial mask loop | DONE (B2 gates open) |
39
- | spec only | **serial** | `hasDraft` routes serial (upstream parity) | Phase B |
40
- | grammar + spec | **serial** | constrained verify walk (novel — §Phase C) | Phase C |
41
- | batch + spec | serial per request | `hasDraft` wins over batch | Phase B (routing) |
42
- | batch + grammar + spec | serial per request | `hasDraft` wins | falls out of B+C |
43
-
44
- Non-goals (explicitly out of scope, revisit only on measured demand):
45
- per-row spec inside the batch scheduler (variable accept-lengths break
46
- uniform-B; a research project, not a feature gap), DSpark as a served
47
- source before its 27B retrain/data-scale lands, grammar on the
48
- `/v1/responses` shim beyond what chat inheritance already gives.
49
-
50
- ## Phases
51
-
52
- Ordering rationale: A is small and closes the trust gap on what already
53
- shipped. B is the flagship parity feature and lays the `DraftSource`
54
- seam C needs. C is the novel composition. D runs parallel to B/C
55
- (scheduler files vs serial-loop files — near-zero overlap). E needs all
56
- of them and produces the numbers.
57
-
58
- ### Phase A — structured-output debt (B2 + F4) [S, 2–3 days]
59
-
60
- 1. **B2 model-gated scheduler tests** (spec'd in structured-output.md):
61
- mixed batch mask isolation; **all-grammar B=4 with four different
62
- schemas** (THE bug-class gate — row↔matcher misalignment after
63
- eviction/join); early termination + churn + joiner; max_tokens
64
- mid-JSON truncation under batch; prefill-terminated row never merges.
65
- 2. **F4 compiler cache**: one `GrammarCompiler` per `TokenizerInfo`
66
- (both vocab-structural; `ownsCompiler=false` path already exists).
67
- Measure compile time cold vs cached on a real JSON schema — this is
68
- the agentic-replay win and a benchmark row in Phase E.
69
- 3. Opportunistic: F6 choice escaping (tiny), F3 explicit padded-vocab
70
- comment. F5/F7 stay parked for the structural-tag pass (see Later).
71
-
72
- Exit: gated suites green on cpm5 + e4b; compile-cache A/B number
73
- recorded.
74
-
75
- ### Phase B — `serve --draft-model` (two-model, serial lane) [M, ~1 wk]
76
-
77
- Execute mlx-lm-tool-parity-plan.md §7 as written (it is already the
78
- hand-to-an-implementer plan). Integration deltas this plan adds:
79
-
80
- - **`RequestShape.hasDraft`** beside `hasGrammar`; `willBatch` returns
81
- false for it. Lane-debug line includes it. (The B0-grammar routing
82
- pattern, reused verbatim.)
83
- - **The serve verify loop lands grammar-aware from day one**: the loop
84
- emits bursts of ≤ n+1 tokens; the stop-sequence/detokenizer audit in
85
- §7 ("risky parts") must treat the grammar mask hook as a first-class
86
- participant even while it's unused in this phase — i.e. the per-round
87
- walk is structured as `for pos in verifyWindow: sample(logits[pos],
88
- maskHook?)` so Phase C is a hook implementation, not a loop rewrite.
89
- - **Kill switch** `MLX_BUN_SPEC=0` (degrade: ignore the flag, log once),
90
- house style.
91
-
92
- Oracle: L1 spec-vs-spec token equality with `mlx_lm.server
93
- --draft-model` from the oracle venv (§7.8). Acceptance telemetry
94
- (`drafted/accepted/targetCalls`) in stats — Phase E consumes it.
95
-
96
- Exit: greedy token-for-token vs mlx-lm spec on a fixed pair; rotating
97
- wrap falls back non-spec (no 500); `--batch 4` + `--draft-model` routes
98
- spec requests serial and batches the rest; admission accounts for draft
99
- weights + KV.
100
-
101
- ### Phase C — grammar × spec: the constrained verify walk [M, ~1 wk]
102
-
103
- Novel territory — **no oracle serves grammar+spec** (mlx-lm has no
104
- grammar; oMLX has no spec). No-oracle (Lab-class) gating per the
105
- parity-tier contract: equivalence + validity gates, not bit-parity to an
106
- ancestor.
107
-
108
- Design (v1 = verify-side masking, drafter free-running):
109
-
110
- 1. **Drafter is unconstrained.** Drafts are proposals; correctness never
111
- depends on them (a grammar-invalid draft is simply rejected at
112
- verify). No matcher state touches the draft path.
113
- 2. **The verify walk applies the mask sequentially.** One batched target
114
- forward over `[pending, ...drafts]` (n+1 positions) as in Phase B;
115
- then the host-side accept walk goes position by position:
116
- `mask_i = matcher.bitmask` → apply to `logits[i]` → sample → if
117
- sampled == draft_i: `matcher.accept(draft_i)`, continue; else emit
118
- the masked correction token, `matcher.accept(correction)`, stop.
119
- The matcher only ever advances on **emitted** tokens — no rollback
120
- needed in v1. Sequential bitmask fills cost 0.004–0.19 ms each
121
- (measured, structured-output.md); at γ≤3 that's ≤0.8 ms/round against
122
- a 12B target step of ~40 ms — noise.
123
- 3. **Grammar termination mid-burst** truncates the accepted prefix at
124
- the terminating token (same all-`-inf` guarantee as serial/batched:
125
- nothing is sampled past termination), finish_reason "stop".
126
- 4. **The measured fork** ([[flag-and-try-both-on-forks]]): on
127
- grammar-tight spans (e.g. `guided_choice`, closing braces) a
128
- free-running drafter's acceptance may crater — every draft off-grammar
129
- is a wasted round. If Phase E shows that, the alternative is built
130
- behind `MLX_BUN_SPEC_GRAMMAR_DRAFT=1`: mask the drafter's logits too
131
- (matcher advances per draft, `rollBack(n − kAccept)` on rejection —
132
- the WASM API exposes rollBack; this is its use case). Build the flag
133
- only if the default's acceptance number demands it; don't pre-build.
134
-
135
- Gates (no external oracle):
136
- - **Equivalence**: greedy grammar+spec long-prefix-equal to greedy
137
- grammar-only serial on fixed prompts (exact-match acceptance means
138
- spec must reproduce the non-spec masked-greedy stream, modulo the
139
- documented batched-verify-head knife-edges — same caveat class as L1
140
- spec-vs-spec).
141
- - **Validity**: every grammar+spec output parses / conforms to its
142
- schema (100%, no exceptions), matcher never rejects an emitted token.
143
- - **Composition regression fence**: grammar-only and spec-only suites
144
- byte-unchanged.
145
-
146
- ### Phase D — batching P0 (+ admission slice of P3) [M, ~1 wk, parallel with B/C]
147
-
148
- From batching-perf-path.md, the subset that the benchmark needs and that
149
- doesn't collide with B/C's serial-loop work:
150
-
151
- 1. **P0**: `scripts/bench-serving-load.ts` (N concurrent vs live server,
152
- aggregate tok/s + TTFT p50/p95 at the SSE first-token event) —
153
- Phase E builds on this harness; `extend` join (kill the O(B·S)
154
- re-merge); vectorized homogeneous-sampler fast path.
155
- - `extend` + grammar rows: the B1 read-before-build step must hold
156
- across an extend-join — add one B2-style churn test with a joiner
157
- landing while a grammar row is live.
158
- 2. **P3 admission slice only**: `projectKvBytes` + `--kv-budget`
159
- queueing (the OOM backstop the combined benchmark needs before it
160
- pushes B=4 × long prompts on 24 GB). Prompt-cache-under-batching,
161
- adapters, and the default flip stay in batching-perf-path P3.
162
-
163
- P1 (quantized KV at B>1) is NOT in this plan's critical path — it's the
164
- next wave after the matrix benchmark exists, and the benchmark is exactly
165
- what will size its payoff. (P2 — perf kernel at B>1 — is OBSOLETE: the
166
- kernel was deleted 2026-07-05, docs/design/unified-engine-frontier-plan.md §6.)
167
-
168
- ### Phase E — the feature-matrix benchmark [S–M, 3–4 days]
169
-
170
- One harness answering "what does each feature — and each legal
171
- combination — cost or buy, per model?" Extends P0's
172
- `bench-serving-load.ts`; preflight-gated like benchmark.sh (dirty-machine
173
- numbers are garbage; quotable rows need the clean-machine protocol).
174
-
175
- **Cells** (skip-if-unavailable, like the model-gated suites):
176
-
177
- - Lanes/features: `serial`, `batch4`, `serial+grammar`, `batch4+grammar`,
178
- `serial+spec`, `serial+grammar+spec` — the six legal shapes from the
179
- matrix. (batch+spec cells route serial by design; the harness records
180
- them as routing checks, not perf rows.)
181
- - Models: cpm5 + e4b (fast movers, both lanes), 12B (+ its assistant
182
- drafter — the measured spec win), Qwen3.5 (batch + thinking-model
183
- grammar canary; grammar cells expected to expose F7 until structural
184
- tags land — record, don't hide).
185
- - Workloads: short-chat (128 tok), JSON-extraction (a real schema,
186
- grammar cells), long-form (512 tok). 4 concurrent for batch cells,
187
- 1 for serial cells; N=3 median.
188
-
189
- **Metrics per cell**: aggregate + per-stream decode tok/s · TTFT
190
- p50/p95 (SSE first token, wall-clock) · spec acceptance
191
- (drafted/accepted/targetCalls) · grammar overhead vs the same cell
192
- without grammar (%) · schema-conformance rate (**must be 100% in every
193
- grammar cell — a conformance failure fails the run, it's a correctness
194
- gate wearing a benchmark costume**) · peak RSS.
195
-
196
- **Denominator discipline** (spec-decode-larger-targets.md): every
197
- feature cell is compared to the SAME model+lane baseline cell from the
198
- SAME run. No cross-model or cross-run speedup claims.
199
-
200
- **Output**: eval-DB rows tagged `feature-matrix[<host>]` + a dated
201
- markdown artifact (gitignored, like benchmark.sh's); promoted rows land
202
- in benchmarks/RESULTS.md under a new "composition" table. Josh-gated
203
- clean-machine run for anything quotable.
204
-
205
- Exit: the matrix runs end-to-end on this machine; a first
206
- (loaded-machine, unquotable) sweep sanity-checks every cell; the
207
- clean-machine run is queued in STATUS Josh-gated.
208
-
209
- ## Later / parked (pointers, so they aren't lost)
210
-
211
- - **F7 + F5 structural-tag pass** (thinking-model grammar + real regex)
212
- — try the installed WASM package's `fromStructuralTag`/`RegexFormat`
213
- first; U1 rebuild only if its Dec-2025 core predates the semantics
214
- (structured-output.md addendum). Qwen3.5's grammar cells in Phase E
215
- are the forcing function.
216
- - **U2 native XGrammar-2 via TVM-FFI** — trigger conditions unchanged
217
- (B≫8, spec+grammar draft-tree masking via `traverse_draft_tree`, or
218
- compile time visible despite F4). Phase C's v1 walk deliberately
219
- avoids needing it.
220
- - **Batching P1** (quant-KV at B>1) — next wave, sized by the matrix.
221
- (P2, perf kernel at B>1, is obsolete — kernel deleted 2026-07-05.)
222
- - **P4 device-side step chaining** — orthogonal serial-decode work; the
223
- matrix's serial baselines will show its headroom.
224
- - **Assistant/DSpark sources behind `--draft-model`** —
225
- mlx-lm-tool-parity-plan §7.9; DSpark stays research until 27B
226
- retrain + data scale (dspark-handoff.md).
227
- - **Spec inside the batch lane** — research question, not scheduled.
228
-
229
- ## Execution log (2026-07-03)
230
-
231
- - **Phase A** ✅ — B2 gates in `tests/batch-grammar.test.ts` (all five
232
- scenarios + a sixth found later, below); F4 compiler cache went through a
233
- bring-up fix (single-flight rebuild — concurrent compiles after a
234
- tokenizer switch double-disposed the previous compiler); F6, F3 landed.
235
- - **Phase B** ✅ — `serve --draft-model`/`--num-draft-tokens`;
236
- `src/spec/{source,two-model,serve-loop}.ts`. **L1 gate: 48/48
237
- token-for-token vs mlx-lm spec** (Llama 3B+1B greedy, 65% acceptance,
238
- `scripts/oracle-spec-two-model.py`). Deviations as designed (ring-wrap
239
- degrades pre-pollution via an offset+n+1 gate — upstream's post-hoc check
240
- would already have polluted the ring; prompt cache bypassed v1).
241
- - **Phase C** ✅ — the constrained verify walk in `samplePos`; gates green
242
- (100% validity; 12/12 token-identical to grammar-only serial; choice
243
- terminates mid-spec). Drafter-masking flag NOT built (acceptance 6/15 on
244
- a tight schema — usable; revisit only if Phase E's clean run says so).
245
- - **Phase E** ✅ (harness) — `scripts/bench-feature-matrix.ts` (the planned
246
- bench-serving-load extension became its own script: that file already
247
- exists as a client-only stack-vs-stack tool and stays one). Six cells
248
- green end-to-end on Llama 3B+1B; `usage.speculation` telemetry on all
249
- four API usage blocks. Clean-machine run for RESULTS.md is Josh-gated.
250
- - **The conformance gate paid for the whole phase in one smoke run** —
251
- three real bugs: (1) snake_case wire fields never reached the grammar
252
- resolver (structured output was DEAD over HTTP since it landed);
253
- (2) `#flushPipeline` didn't advance matchers → stale masks on mid-decode
254
- joins (regression test added); (3) **UniversalDenseModel batched RoPE
255
- uses the scalar cache.offset** → uneven-row batches decode joiners at
256
- wrong positions — latent for ALL Tier-0 archs since v0.0.9. **Fixed
257
- same day**: `UniversalRope.applyDynamic` (+ `ops.ropeScaledDynamic`)
258
- routes the wrapper's per-row `ropeOffsetArr`; gated **token-exact vs
259
- mlx-lm B=2 on Llama-3.2-3B** — static uneven rows AND the dynamic
260
- join/leave protocol (tests/batched-decode-parity.test.ts "Llama 3B
261
- Tier-0"). Plain full-attention universal archs now batch; maskArray
262
- (gemma2-family, pad-blind causal mask in forwardLayers) and
263
- sliding-window universal archs remain serial (unvalidated cells).
264
- Lesson, per [[per-model-quant-specialization]]: the gateway's
265
- capability gate admitted a model family no batched-parity suite ever
266
- validated.
267
-
268
- ### Phase D completion (2026-07-04)
269
-
270
- - **KV-budget admission** — `--kv-budget <GB>`: aggregate projected-KV gate
271
- in the scheduler (queue-don't-OOM, oversized-alone rejects,
272
- `/stats.batch.{pending_rows,kv_bytes,kv_budget_bytes}`;
273
- tests/batch-kv-budget.test.ts).
274
- - **Vectorized homogeneous sampling** — all-plain-greedy batches sample via
275
- ONE log-softmax+argmax over [B,V]; gateway flags `plainGreedy`; BIT-equal
276
- A/B vs the per-row path (`MLX_BUN_BATCH_VEC_SAMPLE=0` kill switch;
277
- tests/batch-vec-sample.test.ts).
278
- - **Extend-join** — `extendKVRows` (mlx-lm `BatchKVCache.extend` semantics:
279
- one pad + one B-axis concat, running rows untouched, pads grow) replaces
280
- the O(B·S) re-merge for full-attention layers in `#mergeJoiner`
281
- (`MLX_BUN_BATCH_EXTEND=0` = re-merge). **Own oracle** (extend's pad
282
- layout ≠ re-merge's): scripts/gen-batched-extend-golden.py →
283
- token-for-token on CPM AND Llama-3B (tests/batched-decode-parity.test.ts
284
- "EXTEND-JOIN ORACLE"); full scheduler suite incl. the grammar-churn join
285
- test green through the extend path AND the kill-switch path.
286
- Rotating-layer extend (Gemma joins re-merge their sliding layers, extend
287
- their full layers) is a follow-up.
288
-
289
- ## Risks
290
-
291
- - **Grammar+spec has no oracle** — mitigated by the three-gate stack
292
- (equivalence to grammar-only greedy, 100% validity, regression fence)
293
- and by keeping the drafter out of the grammar state in v1.
294
- - **Multi-token bursts through stop-sequence/detokenizer state** — the
295
- known §7 risky part, now load-bearing for two phases; land the burst
296
- audit with Phase B tests, not after.
297
- - **Memory at 12B+drafter / B=4 on 24 GB** — Phase D's admission slice
298
- is the backstop; benchmark cells must go through it, not around it.
299
- - **Scope creep via the matrix** — the benchmark measures what exists;
300
- it must not become a reason to build P1/P2/drafter-masking
301
- speculatively. Each parked item has its trigger written above.