mlx-bun 0.0.9 → 0.0.11

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 (179) hide show
  1. package/README.md +66 -21
  2. package/THIRD_PARTY_LICENSES.md +1 -0
  3. package/docs/archive/mlx-bun-lab-report.html +275 -0
  4. package/docs/archive/wwdc-mlx-bun.html +848 -0
  5. package/docs/design/audio-input-plan.md +356 -0
  6. package/docs/design/batching-perf-path.md +255 -0
  7. package/docs/design/batching-v2-plan.md +23 -9
  8. package/docs/design/bucketing-stage.md +2 -1
  9. package/docs/design/compat-cli-surface-design.md +4 -3
  10. package/docs/design/decode-speed-program.md +172 -0
  11. package/docs/design/docs-reorg-plan.md +1 -1
  12. package/docs/design/dspark-serving-program.md +371 -0
  13. package/docs/design/dspark-speculative-decoding.md +167 -5
  14. package/docs/design/faithful-l1-consolidation.md +310 -0
  15. package/docs/design/generic-model-support.md +8 -5
  16. package/docs/design/grammar-spec-batching-integration.md +301 -0
  17. package/docs/design/minicpm5-decode-megakernel.md +16 -11
  18. package/docs/design/mlx-lm-tool-parity-plan.md +11 -1
  19. package/docs/design/omlx-adoption-map.md +98 -0
  20. package/docs/design/orpo-training.md +34 -30
  21. package/docs/design/paged-kv-cache.md +214 -0
  22. package/docs/design/parallel-slots.md +82 -49
  23. package/docs/design/parity-tier-dag.md +43 -33
  24. package/docs/design/repo-cleanup-plan.md +189 -0
  25. package/docs/design/runtime-isolation.md +19 -3
  26. package/docs/design/segmented-backward-training.md +18 -2
  27. package/docs/design/ssd-kv-cold-tier.md +288 -0
  28. package/docs/design/structured-output.md +484 -0
  29. package/docs/design/the-dreaming-master-plan.md +2 -3
  30. package/docs/design/turboquant-kv.md +235 -0
  31. package/docs/design/unified-engine-frontier-plan.md +774 -0
  32. package/docs/design/web-chat-beat-matrix.md +315 -0
  33. package/docs/design/web-chat-redesign.md +1112 -0
  34. package/docs/investigations/decode-roofline-lookagain.md +92 -6
  35. package/docs/investigations/dspark-handoff.md +182 -16
  36. package/docs/investigations/kernel-perf-review-2026-07.md +105 -26
  37. package/docs/investigations/lab-build-journal.md +2 -2
  38. package/docs/investigations/openwiki-evaluation.md +249 -0
  39. package/docs/planning/IDEAS.md +14 -0
  40. package/docs/planning/PRODUCT_ROADMAP.md +2 -3
  41. package/docs/planning/chat-ui-vision-memory-first.md +199 -0
  42. package/docs/planning/chat-ui-vision.md +50 -0
  43. package/docs/planning/memory-docs-and-dag-plan.md +15 -6
  44. package/docs/planning/release-notes-v0.0.10.md +57 -0
  45. package/docs/planning/release-notes-v0.0.11.md +239 -0
  46. package/docs/reference/cli.md +35 -16
  47. package/docs/reference/distribution.md +13 -6
  48. package/docs/reference/embedding.md +35 -14
  49. package/docs/reference/features-matrix.md +97 -0
  50. package/docs/reference/library-api.md +30 -9
  51. package/docs/reference/memory.md +26 -6
  52. package/docs/reference/models.md +11 -1
  53. package/docs/reference/orpo-quickstart.md +3 -4
  54. package/docs/reference/server-api.md +634 -22
  55. package/docs/reference/server-config.md +251 -119
  56. package/docs/reference/training.md +20 -16
  57. package/package.json +5 -3
  58. package/src/anthropic.ts +13 -0
  59. package/src/audio/conformer.ts +647 -0
  60. package/src/audio/decode.ts +133 -0
  61. package/src/audio/features.ts +202 -0
  62. package/src/audio/transcode.ts +67 -0
  63. package/src/cli.ts +402 -107
  64. package/src/config.ts +49 -2
  65. package/src/download.ts +40 -0
  66. package/src/eval/README.md +7 -6
  67. package/src/eval/kl.ts +59 -3
  68. package/src/eval/quality-db.ts +2 -2
  69. package/src/eval/runner.ts +3 -3
  70. package/src/expert-offload.ts +15 -2
  71. package/src/flags.ts +13 -0
  72. package/src/generate.ts +330 -16
  73. package/src/grammar.ts +585 -0
  74. package/src/hub-rest.ts +242 -0
  75. package/src/kv-store.ts +618 -122
  76. package/src/lora.ts +8 -0
  77. package/src/md-modules.d.ts +42 -0
  78. package/src/media-fetch.ts +284 -0
  79. package/src/memory/rest.ts +245 -0
  80. package/src/memory/vault.ts +74 -0
  81. package/src/mlx/array.ts +127 -19
  82. package/src/mlx/autograd.ts +2 -2
  83. package/src/mlx/ffi.ts +27 -3
  84. package/src/mlx/metal-kernel.ts +7 -4
  85. package/src/mlx/ops.ts +179 -23
  86. package/src/mlx/turboquant-ops.ts +831 -0
  87. package/src/mlx/turboquant-tables.ts +237 -0
  88. package/src/model/batched-mask.ts +62 -1
  89. package/src/model/batched-quant.ts +206 -0
  90. package/src/model/batched-rotating-quant.ts +346 -0
  91. package/src/model/batched-rotating.ts +46 -1
  92. package/src/model/compiled-decode.ts +4 -15
  93. package/src/model/factory.ts +10 -4
  94. package/src/model/gemma4-base.ts +391 -34
  95. package/src/model/gemma4.ts +122 -33
  96. package/src/model/generated/gemma4-12b.ts +9 -13
  97. package/src/model/generated/gemma4-26b.ts +10 -5
  98. package/src/model/generated/gemma4-e4b.ts +14 -9
  99. package/src/model/minicpm5.ts +58 -102
  100. package/src/model/paged-kv.ts +253 -0
  101. package/src/model/qwen3-delta.ts +123 -15
  102. package/src/model/qwen3-moe.ts +404 -0
  103. package/src/model/qwen3.ts +9 -8
  104. package/src/model/qwen3_5.ts +99 -25
  105. package/src/model/support.ts +12 -1
  106. package/src/model/universal/dense.ts +40 -5
  107. package/src/model/universal/rope.ts +36 -21
  108. package/src/pi-web.ts +1299 -82
  109. package/src/preflight.ts +27 -8
  110. package/src/prompt-cache.ts +338 -32
  111. package/src/quantize/quantizer.ts +16 -1
  112. package/src/registry.ts +61 -3
  113. package/src/serve/batch-scheduler.ts +1001 -76
  114. package/src/serve/generation-gateway.ts +292 -30
  115. package/src/serve/isolate.ts +417 -0
  116. package/src/serve/lane-registry.ts +44 -0
  117. package/src/serve/session-search.ts +326 -0
  118. package/src/server.ts +1460 -105
  119. package/src/spec/assistant-source.ts +150 -0
  120. package/src/spec/deepspec-source.ts +143 -0
  121. package/src/spec/dflash-source.ts +124 -0
  122. package/src/spec/dspark/ab-stats.ts +136 -0
  123. package/src/spec/dspark/calibration.ts +112 -0
  124. package/src/spec/dspark/deepspec-module.ts +801 -0
  125. package/src/spec/dspark/generate-dflash.ts +72 -20
  126. package/src/spec/dspark/loader.ts +23 -0
  127. package/src/spec/dspark/module-dflash.ts +278 -36
  128. package/src/spec/dspark/quantize-drafter.ts +91 -0
  129. package/src/spec/generate.ts +59 -25
  130. package/src/spec/ngram-source.ts +132 -0
  131. package/src/spec/serve-loop.ts +536 -0
  132. package/src/spec/source.ts +93 -0
  133. package/src/spec/two-model.ts +171 -0
  134. package/src/ssd-cache.ts +285 -0
  135. package/src/tokenizer.ts +288 -2
  136. package/src/tool-approvals.ts +131 -0
  137. package/src/tool-call.ts +408 -20
  138. package/src/train/flash-cce.ts +29 -26
  139. package/src/train/job.ts +0 -2
  140. package/src/train/loss.ts +53 -11
  141. package/src/train/prefix-shared.ts +30 -7
  142. package/src/train/segmented.ts +150 -46
  143. package/src/train/trainer.ts +22 -15
  144. package/src/vision/prompt.ts +217 -48
  145. package/src/web/app.html +1184 -1789
  146. package/src/web/app.js +5667 -0
  147. package/src/web/icon.svg +13 -0
  148. package/src/web/manifest.webmanifest +24 -0
  149. package/src/web/src/adapters-panel.ts +288 -0
  150. package/src/web/src/api.ts +57 -0
  151. package/src/web/src/assistant.ts +485 -0
  152. package/src/web/src/chat.ts +1050 -0
  153. package/src/web/src/composer.ts +959 -0
  154. package/src/web/src/dataset.ts +132 -0
  155. package/src/web/src/finetune.ts +267 -0
  156. package/src/web/src/hub.ts +428 -0
  157. package/src/web/src/main.ts +98 -0
  158. package/src/web/src/markdown.ts +0 -0
  159. package/src/web/src/memory-panel.ts +572 -0
  160. package/src/web/src/model-picker.ts +159 -0
  161. package/src/web/src/palette.ts +346 -0
  162. package/src/web/src/protocol.ts +62 -0
  163. package/src/web/src/quantize.ts +167 -0
  164. package/src/web/src/rag.ts +274 -0
  165. package/src/web/src/sessions.ts +469 -0
  166. package/src/web/src/shell.ts +804 -0
  167. package/src/web/src/status.ts +252 -0
  168. package/src/web/src/ui-catalog.ts +92 -0
  169. package/src/web/sw.js +71 -0
  170. package/src/web/vendor/README.md +61 -0
  171. package/src/web/vendor/hljs-LICENSE +29 -0
  172. package/src/web/vendor/hljs-theme.css +59 -0
  173. package/src/web/vendor/hljs.js +26 -0
  174. package/src/model/fused-decode-kernel.ts +0 -312
  175. package/src/model/fused-geglu-kernel.ts +0 -198
  176. package/src/model/fused-mlp-kernel.ts +0 -169
  177. package/src/model/fused-swiglu-kernel.ts +0 -240
  178. package/src/model/steel-linear-kernel.ts +0 -96
  179. package/src/status-page.html +0 -497
package/src/tool-call.ts CHANGED
@@ -28,6 +28,12 @@ export interface ParsedToolCall {
28
28
  name: string;
29
29
  /** Parsed argument object (callers serialize for OpenAI's string field). */
30
30
  arguments: Record<string, unknown>;
31
+ /** Set only when strict parsing failed and a repair pass recovered the
32
+ * call — never silent (Axis 7 self-healing beat row). `repairs` lists
33
+ * every malformation class that had to be corrected, in application
34
+ * order, so failures are observable (logs/telemetry) instead of masked. */
35
+ repaired?: true;
36
+ repairs?: RepairKind[];
31
37
  }
32
38
 
33
39
  interface ToolSpec {
@@ -60,7 +66,18 @@ export function gemmaArgsToJson(text: string): string {
60
66
  return out;
61
67
  }
62
68
 
63
- /** Parse every `call:name{...}` block in a decoded tool-call segment. */
69
+ /** Parse every `call:name{...}` block in a decoded tool-call segment.
70
+ *
71
+ * Strict scan first (unchanged behavior on well-formed input). On the two
72
+ * ways a small model mangles THIS family's syntax specifically —
73
+ * (a) generation truncated mid-call (an unterminated `<|"|>` string or a
74
+ * brace that never closes, both hit_max_tokens artifacts) and (b) a
75
+ * trailing comma surviving `gemmaArgsToJson`'s conversion to JSON — the
76
+ * repair pass recovers by treating "ran off the end of the text" as an
77
+ * implicit close (mirrors closeUnbalanced's brace-stack unwind, adapted
78
+ * for the `<|"|>`-quoted syntax) and re-running stripTrailingCommas on the
79
+ * converted JSON. Recovered calls are tagged repaired/repairs; a segment
80
+ * that still can't be recovered throws exactly as before. */
64
81
  export function parseToolCalls(text: string): ParsedToolCall[] {
65
82
  const calls: ParsedToolCall[] = [];
66
83
  let i = 0;
@@ -71,10 +88,11 @@ export function parseToolCalls(text: string): ParsedToolCall[] {
71
88
  // balanced-brace scan, skipping <|"|>-delimited strings
72
89
  let depth = 0;
73
90
  let j = nameEnd;
91
+ let unterminatedString = false;
74
92
  while (j < text.length) {
75
93
  if (text.startsWith(QUOTE, j)) {
76
94
  const close = text.indexOf(QUOTE, j + QUOTE.length);
77
- if (close === -1) throw new Error("unterminated <|\"|> string");
95
+ if (close === -1) { unterminatedString = true; break; }
78
96
  j = close + QUOTE.length;
79
97
  continue;
80
98
  }
@@ -86,11 +104,55 @@ export function parseToolCalls(text: string): ParsedToolCall[] {
86
104
  }
87
105
  j++;
88
106
  }
89
- if (depth !== 0) throw new Error("unbalanced braces in tool call");
90
- const argsBlock = text.slice(nameEnd, j + 1);
107
+ const truncated = unterminatedString || depth !== 0;
108
+ let argsBlock: string;
109
+ const repairs: RepairKind[] = [];
110
+ if (truncated) {
111
+ // Ran off the end of the text without closing — truncated generation.
112
+ // Close the dangling <|"|> string (if any) then append enough `}` to
113
+ // balance whatever braces are still open, using the SAME depth
114
+ // counter's remaining value (skip <|"|>-quoted spans again so a
115
+ // brace character inside a string doesn't miscount).
116
+ let closeDepth = 0;
117
+ let k = nameEnd;
118
+ let inQuote = false;
119
+ while (k < text.length) {
120
+ if (text.startsWith(QUOTE, k) && !inQuote) { inQuote = true; k += QUOTE.length; continue; }
121
+ if (inQuote) {
122
+ if (text.startsWith(QUOTE, k)) { inQuote = false; k += QUOTE.length; continue; }
123
+ k++;
124
+ continue;
125
+ }
126
+ const c = text[k];
127
+ if (c === "{") closeDepth++;
128
+ else if (c === "}") closeDepth--;
129
+ k++;
130
+ }
131
+ argsBlock = text.slice(nameEnd) + (inQuote ? QUOTE : "") + "}".repeat(Math.max(closeDepth, 1));
132
+ repairs.push("unbalanced_braces");
133
+ j = text.length - 1;
134
+ } else {
135
+ argsBlock = text.slice(nameEnd, j + 1);
136
+ }
137
+ let converted = gemmaArgsToJson(argsBlock);
138
+ let parsed: unknown;
139
+ try {
140
+ parsed = JSON.parse(converted);
141
+ } catch {
142
+ const commas = stripTrailingCommas(converted);
143
+ if (!commas.changed) throw new Error("unbalanced braces in tool call");
144
+ try {
145
+ parsed = JSON.parse(commas.out);
146
+ } catch {
147
+ throw new Error("unbalanced braces in tool call");
148
+ }
149
+ converted = commas.out;
150
+ repairs.push("trailing_commas");
151
+ }
91
152
  calls.push({
92
153
  name: m[1]!,
93
- arguments: JSON.parse(gemmaArgsToJson(argsBlock)) as Record<string, unknown>,
154
+ arguments: parsed as Record<string, unknown>,
155
+ ...(repairs.length ? { repaired: true, repairs } : {}),
94
156
  });
95
157
  i = j + 1;
96
158
  }
@@ -159,28 +221,327 @@ function decodeToolValue(value: string, schema: Record<string, unknown> | null):
159
221
  }
160
222
  }
161
223
 
224
+ // ---- Repair layer (Axis 7 self-healing beat row) ---------------------
225
+ //
226
+ // pi-ai (node_modules/@earendil-works/pi-ai/dist/utils/json-parse.js)
227
+ // already ships a generic JSON repair pass for cloud-model streaming
228
+ // tool calls: `repairJson`/`parseJsonWithRepair` fix invalid backslash
229
+ // escapes and raw control characters inside strings, and
230
+ // `parseStreamingJson` falls back to the `partial-json` package for
231
+ // truncated mid-stream JSON. Its `validateToolArguments`
232
+ // (utils/validation.js) additionally coerces JSON-Schema primitive
233
+ // TYPES (string "3" -> number 3, etc.) once a call is already
234
+ // structurally valid JSON. None of that touches the STRUCTURAL
235
+ // malformations small local models actually produce here: single
236
+ // quotes instead of double, trailing commas, unbalanced/truncated
237
+ // braces, a whole call wrapped in a markdown ```json fence, stray
238
+ // prose ("Sure, here's the call:") glued in front of the JSON, or the
239
+ // function name landing inside `arguments` instead of the envelope
240
+ // (toolCallFromPayload above already handles that last one for
241
+ // already-valid JSON — the repair pass below reuses it once the text
242
+ // is coerced to valid JSON). This layer is therefore purely additive
243
+ // to pi-ai's, applied ONLY when strict parsing has already failed, and
244
+ // every repair is tagged on the returned ParsedToolCall (`repaired`,
245
+ // `repairs`) rather than silently swallowed — see ParsedToolCall above.
246
+ export type RepairKind =
247
+ | "fenced_wrapper" // ```json ... ``` or ``` ... ``` wrapper stripped
248
+ | "prose_prefix" // leading non-JSON prose trimmed to the first { or [
249
+ | "single_quotes" // 'like this' -> "like this" (keys and string values)
250
+ | "trailing_commas" // {"a":1,} / [1,2,] -> valid JSON
251
+ | "unbalanced_braces" // truncated generation: missing closing }/]/" appended
252
+ | "python_literals" // True/False/None -> true/false/null
253
+ | "name_in_args"; // function name found inside `arguments` instead of the envelope
254
+
255
+ /** Strip a ```json fenced block (or bare ```) wrapping the whole payload. */
256
+ function stripFence(text: string): { out: string; changed: boolean } {
257
+ const m = /^\s*```(?:json)?\s*\n?([\s\S]*?)\n?```\s*$/i.exec(text.trim());
258
+ if (!m) return { out: text, changed: false };
259
+ return { out: m[1] ?? "", changed: true };
260
+ }
261
+
262
+ /** Trim stray prose before the first `{`/`[` and after the matching
263
+ * document's last `}`/`]` (a small model narrating "Sure, here's the
264
+ * call: {...}" or trailing chatter after the JSON). Leaves the text
265
+ * alone if it doesn't start with prose (no-op safe to always run). */
266
+ function stripProsePrefix(text: string): { out: string; changed: boolean } {
267
+ const trimmed = text.trim();
268
+ const firstBrace = trimmed.search(/[{[]/);
269
+ if (firstBrace <= 0) return { out: trimmed, changed: false };
270
+ // Only strip when what precedes genuinely looks like prose (no braces
271
+ // of its own) — never chew into legitimate leading whitespace-only JSON.
272
+ const prefix = trimmed.slice(0, firstBrace);
273
+ if (/[{}[\]]/.test(prefix)) return { out: trimmed, changed: false };
274
+ const lastClose = Math.max(trimmed.lastIndexOf("}"), trimmed.lastIndexOf("]"));
275
+ const out = lastClose >= firstBrace ? trimmed.slice(firstBrace, lastClose + 1) : trimmed.slice(firstBrace);
276
+ return { out, changed: true };
277
+ }
278
+
279
+ /** Convert Python-ish literals (True/False/None) used outside string
280
+ * literals into JSON ones. Scans char-by-char tracking string state so
281
+ * a value like "Nonetheless" or a string containing "True" is untouched. */
282
+ function pythonLiteralsToJson(text: string): { out: string; changed: boolean } {
283
+ let out = "";
284
+ let changed = false;
285
+ let inString: '"' | "'" | null = null;
286
+ for (let i = 0; i < text.length; i++) {
287
+ const ch = text[i]!;
288
+ if (inString) {
289
+ out += ch;
290
+ if (ch === "\\") { out += text[++i] ?? ""; continue; }
291
+ if (ch === inString) inString = null;
292
+ continue;
293
+ }
294
+ if (ch === '"' || ch === "'") { inString = ch; out += ch; continue; }
295
+ const rest = text.slice(i);
296
+ const wordMatch = /^(True|False|None)\b/.exec(rest);
297
+ if (wordMatch) {
298
+ const replacement = wordMatch[1] === "True" ? "true" : wordMatch[1] === "False" ? "false" : "null";
299
+ out += replacement;
300
+ i += wordMatch[1]!.length - 1;
301
+ changed = true;
302
+ continue;
303
+ }
304
+ out += ch;
305
+ }
306
+ return { out, changed };
307
+ }
308
+
309
+ /** Convert single-quoted JSON (keys and string values) to double-quoted,
310
+ * tracking string state so an apostrophe INSIDE a double-quoted string
311
+ * (e.g. "don't") is left alone. Only fires when the text has no double
312
+ * quotes at all in "quote position" — i.e. looks like it was written
313
+ * with single quotes throughout — to avoid mangling a string that
314
+ * legitimately contains a single quote inside double-quoted JSON. */
315
+ function singleQuotesToDouble(text: string): { out: string; changed: boolean } {
316
+ if (!/'/.test(text)) return { out: text, changed: false };
317
+ // Heuristic gate: real JSON never opens a key/value with a bare `'`.
318
+ // If the text already parses with plain JSON.parse it wouldn't have
319
+ // reached here, so any '"'-delimited string is assumed intentional and
320
+ // left byte-for-byte; only '...'-delimited runs are retargeted.
321
+ let out = "";
322
+ let changed = false;
323
+ let inDouble = false;
324
+ let inSingle = false;
325
+ for (let i = 0; i < text.length; i++) {
326
+ const ch = text[i]!;
327
+ if (inDouble) {
328
+ out += ch;
329
+ if (ch === "\\") { out += text[++i] ?? ""; continue; }
330
+ if (ch === '"') inDouble = false;
331
+ continue;
332
+ }
333
+ if (inSingle) {
334
+ if (ch === "\\") { out += ch + (text[++i] ?? ""); continue; }
335
+ if (ch === "'") { out += '"'; inSingle = false; continue; }
336
+ if (ch === '"') { out += '\\"'; continue; } // embedded " must be escaped now
337
+ out += ch;
338
+ continue;
339
+ }
340
+ if (ch === '"') { inDouble = true; out += ch; continue; }
341
+ if (ch === "'") { inSingle = true; changed = true; out += '"'; continue; }
342
+ out += ch;
343
+ }
344
+ return { out, changed };
345
+ }
346
+
347
+ /** Remove trailing commas before a closing `}`/`]` (whitespace/newlines
348
+ * between the comma and the close are preserved otherwise). */
349
+ function stripTrailingCommas(text: string): { out: string; changed: boolean } {
350
+ let changed = false;
351
+ const out = text.replace(/,(\s*[}\]])/g, (_m, close) => { changed = true; return close; });
352
+ return { out, changed };
353
+ }
354
+
355
+ /** Append the minimal closers to balance braces/brackets/an unterminated
356
+ * string — a small model that got truncated (or just stopped early)
357
+ * mid-object. Walks once tracking a stack of open delimiters and string
358
+ * state; on EOF, closes the open string (if any) then unwinds the stack
359
+ * in reverse. No-op (unchanged) when already balanced. */
360
+ function closeUnbalanced(text: string): { out: string; changed: boolean } {
361
+ const stack: ("}" | "]")[] = [];
362
+ let inString = false;
363
+ for (let i = 0; i < text.length; i++) {
364
+ const ch = text[i]!;
365
+ if (inString) {
366
+ if (ch === "\\") { i++; continue; }
367
+ if (ch === '"') inString = false;
368
+ continue;
369
+ }
370
+ if (ch === '"') { inString = true; continue; }
371
+ if (ch === "{") stack.push("}");
372
+ else if (ch === "[") stack.push("]");
373
+ else if (ch === "}" || ch === "]") stack.pop();
374
+ }
375
+ if (!inString && stack.length === 0) return { out: text, changed: false };
376
+ let out = text;
377
+ if (inString) out += '"';
378
+ out += stack.reverse().join("");
379
+ return { out, changed: true };
380
+ }
381
+
382
+ /** Try strict JSON.parse first (pi-ai's own escape-level repair already
383
+ * ran upstream where applicable — see the comment above); on failure,
384
+ * apply the structural repair passes in order, re-testing after each so
385
+ * only the ones that were actually needed get tagged. Returns null (no
386
+ * throw) when nothing recovers a parseable object — callers decide how
387
+ * to surface that. */
388
+ function repairJsonToolText(text: string): { value: unknown; repairs: RepairKind[] } | null {
389
+ const repairs: RepairKind[] = [];
390
+ let candidate = text;
391
+
392
+ const fence = stripFence(candidate);
393
+ if (fence.changed) { candidate = fence.out; repairs.push("fenced_wrapper"); }
394
+
395
+ const prose = stripProsePrefix(candidate);
396
+ if (prose.changed) { candidate = prose.out; repairs.push("prose_prefix"); }
397
+
398
+ const tryParse = (s: string): unknown | undefined => {
399
+ try { return JSON.parse(s); } catch { return undefined; }
400
+ };
401
+
402
+ let parsed = tryParse(candidate);
403
+ if (parsed !== undefined) return { value: parsed, repairs };
404
+
405
+ const literals = pythonLiteralsToJson(candidate);
406
+ if (literals.changed) {
407
+ parsed = tryParse(literals.out);
408
+ if (parsed !== undefined) return { value: parsed, repairs: [...repairs, "python_literals"] };
409
+ candidate = literals.out;
410
+ }
411
+
412
+ const commas = stripTrailingCommas(candidate);
413
+ if (commas.changed) {
414
+ parsed = tryParse(commas.out);
415
+ if (parsed !== undefined) return { value: parsed, repairs: [...repairs, "trailing_commas"] };
416
+ candidate = commas.out;
417
+ }
418
+
419
+ const closed = closeUnbalanced(candidate);
420
+ if (closed.changed) {
421
+ parsed = tryParse(closed.out);
422
+ if (parsed !== undefined) return { value: parsed, repairs: [...repairs, "unbalanced_braces"] };
423
+ candidate = closed.out;
424
+ }
425
+
426
+ const quotes = singleQuotesToDouble(candidate);
427
+ if (quotes.changed) {
428
+ parsed = tryParse(quotes.out);
429
+ if (parsed !== undefined) return { value: parsed, repairs: [...repairs, "single_quotes"] };
430
+ candidate = quotes.out;
431
+ // Single-quote conversion can itself leave a trailing comma or an
432
+ // unbalanced document (e.g. `{'a': 1,}` truncated) — retry those two
433
+ // passes once more on the requoted text before giving up.
434
+ const commas2 = stripTrailingCommas(candidate);
435
+ if (commas2.changed) {
436
+ parsed = tryParse(commas2.out);
437
+ if (parsed !== undefined) return { value: parsed, repairs: [...repairs, "single_quotes", "trailing_commas"] };
438
+ candidate = commas2.out;
439
+ }
440
+ const closed2 = closeUnbalanced(candidate);
441
+ if (closed2.changed) {
442
+ parsed = tryParse(closed2.out);
443
+ if (parsed !== undefined) return { value: parsed, repairs: [...repairs, "single_quotes", "unbalanced_braces"] };
444
+ }
445
+ }
446
+
447
+ return null;
448
+ }
449
+
450
+ /** Lenient JSON.parse used for argument STRINGS nested inside an already
451
+ * JSON-valid envelope (e.g. `"arguments": "{'a': 1}"`, a small model
452
+ * double-encoding with the wrong quote style). Falls back through the
453
+ * same repair passes; throws (not silently {}) when nothing recovers,
454
+ * same contract as JSON.parse for its callers. */
455
+ function parseJsonLoose(text: string): unknown {
456
+ try {
457
+ return JSON.parse(text);
458
+ } catch (err) {
459
+ const repaired = repairJsonToolText(text);
460
+ if (repaired) return repaired.value;
461
+ throw err;
462
+ }
463
+ }
464
+
465
+ /** Build a ParsedToolCall from an already-parsed JSON payload (object or
466
+ * {function:{name,arguments}} envelope). Shared by the strict and
467
+ * repaired paths so name/arguments extraction logic lives once. Handles
468
+ * the "name-in-args confusion" malformation: a small model sometimes
469
+ * emits the function name INSIDE the arguments object instead of (or in
470
+ * addition to) the envelope's own `name` field — e.g.
471
+ * {"arguments":{"name":"read","path":"x"}} with no top-level name at
472
+ * all. When the envelope has no usable name, an args-embedded `name` (or
473
+ * `function`) string is promoted and stripped out of the arguments. */
474
+ function toolCallFromPayload(
475
+ payload: unknown,
476
+ ): { name: string; arguments: Record<string, unknown>; nameInArgs: boolean } | null {
477
+ if (!payload || typeof payload !== "object") return null;
478
+ const obj = payload as Record<string, any>;
479
+ const fn = obj.function && typeof obj.function === "object" ? obj.function : obj;
480
+ let name = typeof fn.name === "string" ? fn.name.trim() : "";
481
+ let args = fn.arguments ?? obj.arguments ?? {};
482
+ if (typeof args === "string") {
483
+ const parsed = args.trim() ? parseJsonLoose(args) : {};
484
+ args = parsed && typeof parsed === "object" && !Array.isArray(parsed) ? parsed : {};
485
+ }
486
+ if (!args || typeof args !== "object" || Array.isArray(args)) return null;
487
+ const argsObj = { ...(args as Record<string, unknown>) };
488
+ let nameInArgs = false;
489
+ if (!name) {
490
+ const embedded = argsObj.name ?? argsObj.function;
491
+ if (typeof embedded === "string" && embedded.trim()) {
492
+ name = embedded.trim();
493
+ delete argsObj.name;
494
+ delete argsObj.function;
495
+ nameInArgs = true;
496
+ }
497
+ }
498
+ if (!name) return null;
499
+ return { name, arguments: argsObj, nameInArgs };
500
+ }
501
+
502
+ /** Shape guard for the wrapper-less last-resort repair path only (see its
503
+ * call site in parseGeneratedToolCalls). A `<tool_call>`/`<function...>`
504
+ * tag already proves the model intended a tool call, so toolCallFromPayload
505
+ * can be lenient there; a bare JSON object with NO such tag is ambiguous —
506
+ * it could just as easily be an ordinary conversational reply that happens
507
+ * to be JSON-shaped (e.g. `{"name": "read", "age": 42}` answering "give me
508
+ * an example record"). Require an actual envelope shape: a `function`
509
+ * wrapper object, or a top-level/nested `arguments` object (including the
510
+ * name-in-args case, where the name lives inside that same `arguments`
511
+ * object) — never just a bare `name`/`function` string sitting next to
512
+ * unrelated fields. */
513
+ function looksLikeToolEnvelope(payload: unknown): boolean {
514
+ if (!payload || typeof payload !== "object" || Array.isArray(payload)) return false;
515
+ const obj = payload as Record<string, unknown>;
516
+ if (obj.function && typeof obj.function === "object" && !Array.isArray(obj.function)) return true;
517
+ const args = obj.arguments;
518
+ if (args && typeof args === "object" && !Array.isArray(args)) return true;
519
+ if (typeof args === "string") return true; // double-encoded arguments string
520
+ return false;
521
+ }
522
+
162
523
  function parseJsonToolCall(block: string): ParsedToolCall | null {
163
524
  let payload: unknown;
525
+ let repairs: RepairKind[] = [];
164
526
  try {
165
527
  payload = JSON.parse(block);
166
528
  } catch {
167
- return null;
529
+ // Strict JSON.parse failed — this is the ONLY point the syntax-repair
530
+ // pass engages (never attempted when the model emitted valid JSON). A
531
+ // repair failure here returns null (not throw): parseGeneratedToolCalls
532
+ // falls through to the XML parsers next, which is the right behavior
533
+ // for a `<tool_call>` block that isn't actually JSON-shaped at all.
534
+ const repaired = repairJsonToolText(block);
535
+ if (!repaired) return null;
536
+ payload = repaired.value;
537
+ repairs = repaired.repairs;
168
538
  }
169
539
  if (!payload || typeof payload !== "object") throw new Error("JSON tool_call payload must be an object");
170
- const obj = payload as Record<string, any>;
171
- const fn = obj.function && typeof obj.function === "object" ? obj.function : obj;
172
- const name = typeof fn.name === "string" ? fn.name.trim() : "";
173
- if (!name) throw new Error("JSON tool_call is missing a function name");
174
- const args = fn.arguments ?? obj.arguments ?? {};
175
- if (typeof args === "string") {
176
- const parsed = args.trim() ? JSON.parse(args) : {};
177
- if (!parsed || typeof parsed !== "object" || Array.isArray(parsed))
178
- throw new Error("tool_call arguments must be a JSON object");
179
- return { name, arguments: parsed as Record<string, unknown> };
180
- }
181
- if (!args || typeof args !== "object" || Array.isArray(args))
182
- throw new Error("tool_call arguments must be a JSON object");
183
- return { name, arguments: args as Record<string, unknown> };
540
+ const result = toolCallFromPayload(payload);
541
+ if (!result) throw new Error("tool_call arguments must be a JSON object");
542
+ if (result.nameInArgs) repairs = [...repairs, "name_in_args"];
543
+ const { nameInArgs: _nameInArgs, ...call } = result;
544
+ return repairs.length ? { ...call, repaired: true, repairs } : call;
184
545
  }
185
546
 
186
547
  function parseXmlEqualsToolCall(block: string, tools: ToolSpec[]): ParsedToolCall | null {
@@ -235,6 +596,33 @@ export function parseGeneratedToolCalls(text: string, tools: ToolSpec[]): Parsed
235
596
  }
236
597
  TOOL_CALL_BLOCK_RE.lastIndex = 0;
237
598
  if (calls.length === 0) calls.push(...parseXmlAttrToolCalls(text, tools));
599
+ // Last-resort repair: no `<tool_call>` tags AND no `<function name=...>`
600
+ // XML matched at all — a small model sometimes drops the wrapper tags
601
+ // entirely and emits just (possibly fenced/prose-prefixed/malformed)
602
+ // JSON. Only attempted when everything else found nothing, so a
603
+ // genuinely tool-free assistant reply never risks being misread as a
604
+ // call. Tagged repaired like every other path here.
605
+ //
606
+ // Unlike the `<tool_call>`-wrapped path, there is no surrounding markup
607
+ // here to prove the model actually intended a tool call — an ordinary
608
+ // conversational reply can itself BE a bare JSON object (e.g. "give me
609
+ // an example record" -> `{"name": "read", "age": 42}`), which happens to
610
+ // collide with a real tool name. Require the recovered payload to
611
+ // actually look like a tool-call envelope (an `arguments`/
612
+ // `function.arguments` object, or a `function` wrapper, or the
613
+ // name-in-args shape) before accepting it — a flat object whose only
614
+ // tool-shaped feature is a `name` key that matches a tool is NOT enough
615
+ // signal on its own.
616
+ if (calls.length === 0) {
617
+ const repaired = repairJsonToolText(text.trim());
618
+ if (repaired && looksLikeToolEnvelope(repaired.value)) {
619
+ const result = toolCallFromPayload(repaired.value);
620
+ if (result) {
621
+ const repairs = result.nameInArgs ? [...repaired.repairs, "name_in_args" as const] : repaired.repairs;
622
+ calls.push({ name: result.name, arguments: result.arguments, repaired: true, repairs });
623
+ }
624
+ }
625
+ }
238
626
  for (const c of calls) {
239
627
  if (!known.has(c.name)) throw new Error(`unknown tool '${c.name}'`);
240
628
  }
@@ -21,7 +21,7 @@
21
21
  // DIVERGENCES (why this isn't a line-for-line port):
22
22
  // 1. Our classifier C is the QUANTIZED head (4/8-bit affine). Apple loads C
23
23
  // dense (bf16); we keep the in-Metal dequant for `tl.dot` (the qdot pattern
24
- // from src/model/fused-decode-kernel.ts). This is the one substantive change.
24
+ // from mlx's quantized.h qdot pattern). This is the one substantive change.
25
25
  // 2. The head is FROZEN (not a LoRA target) → we need only dE (=dh), NOT dC.
26
26
  // 3. The LSE cross-vocab-block merge is done in cheap MLX ops on per-block
27
27
  // partials, not Metal atomic-logaddexp spinlocks (simpler, same result).
@@ -289,7 +289,7 @@ const FWD_SG_SOURCE = String.raw`
289
289
  // memory each step (coalesced). The lane-partial dots reduce with simd_sum. The
290
290
  // threadgroup's 4 simdgroups process 4 rows in parallel; BLOCK_B tokens amortize
291
291
  // the dequant (each row dequantized once, reused across tokens). Same coalesced
292
- // pattern as fused-decode-kernel.ts. h is passed TRANSPOSED [H, M].
292
+ // pattern as mlx's quantized.h qdot. h is passed TRANSPOSED [H, M].
293
293
  const SGN = TG / 32; // simdgroups per threadgroup (4)
294
294
  const FWD_LANE_SOURCE = String.raw`
295
295
  const uint tokBlk = threadgroup_position_in_grid.z;
@@ -723,24 +723,28 @@ const BWD_SG_SOURCE = String.raw`
723
723
  }
724
724
  `;
725
725
 
726
- // Apple-CCE coeff filter epsilon (phase-2 per-tile skip). DEFAULT OFF ("0") today —
727
- // no production path sets the env — but the VERDICT is measured and positive: the old
726
+ // Apple-CCE coeff filter epsilon (phase-2 per-tile skip). DEFAULT ON at 1e-5
727
+ // (flipped 2026-07-02 — the kernel-perf-review backlog #1 flip). The old
728
728
  // 0.66%→2.7% dh-accuracy scare was on RANDOM targets (flat softmax, the filter's
729
729
  // worst case — see PLAN.md "filter-on-real-data"); the real-data gate
730
- // (scripts/experiments/flash-cce-filter-realdata.ts, commit 3d259ef, 2026-06-22)
731
- // measured eps 1e-5 = 0.16% dh (under the bf16 floor) for a 1.35× faster backward,
732
- // and PLAN.md's recorded decision is "enable the filter at eps ~1e-5". (The ~3.5×
733
- // figure floating around older notes was the retired SG kernel, not steel.) Flipping
734
- // the default is the open backlog item; until then opt in with
735
- // MLX_BUN_CCE_BWD_FILTER_EPS=1e-5 on real (peaked) training data.
736
- const BWD_FILTER_EPS = process.env.MLX_BUN_CCE_BWD_FILTER_EPS ?? "0";
730
+ // (scripts/experiments/flash-cce-filter-realdata.ts) measured, on REAL chunk-ORPO
731
+ // data (M1 Max, 2026-07-02): CPM5 eps 1e-5 = 0.343% dh for 1.41×; e4b = 0.158% dh
732
+ // for 1.70× — and the teacher-forced full-logits fidelity shows the filter's added
733
+ // error is far below the flash-vs-full fp-reassociation scatter that already
734
+ // exists (CPM5 0.850→0.913% relnorm, e4b 1.207→1.220%, cosine ≥0.99993 both).
735
+ // (The ~3.5× figure floating around older notes was the retired SG kernel, not
736
+ // steel.) MLX_BUN_CCE_BWD_FILTER_EPS=0 restores exact gradients — eps "0"
737
+ // compiles the filter OUT (#if CCE_BWD_FILTER 0), the identical pre-flip kernel.
738
+ const BWD_FILTER_EPS = process.env.MLX_BUN_CCE_BWD_FILTER_EPS ?? "1e-5";
737
739
  // blockMax vocab-block early-exit eps (skip whole cold (token-block,vocab-block) programs
738
- // — attacks BOTH phases). DEFAULT OFF ("0"). It is LOSSLESS when it fires (only skips
739
- // blocks whose max prob < eps AND which lack the target → ~0 contribution), so it is the
740
- // SAFE skip; but it only pays off on genuinely peaked real text where whole 8192-vocab
741
- // blocks go cold (the M=512 synthetic probe found nothing cold → pure overhead). Opt in
742
- // with MLX_BUN_CCE_BWD_BLOCK_EPS=1e-5 for long real-text training; needs blockMax supplied.
743
- const BWD_BLOCK_EPS = process.env.MLX_BUN_CCE_BWD_BLOCK_EPS ?? "0";
740
+ // — attacks BOTH phases). DEFAULT ON at 1e-5 (flipped 2026-07-02). It is ~LOSSLESS when
741
+ // it fires (only skips blocks whose max prob < eps AND which lack the target → ~0
742
+ // contribution). The M=512 synthetic probe found nothing cold (pure overhead), but the
743
+ // real-data gate proved real text DOES go cold: alone it's 1.23× (CPM5) / 2.02× (e4b)
744
+ // at ≤0.004% dh; COMBINED with the coeff filter at 1e-5/1e-5 the backward is 1.71×
745
+ // (CPM5) / 3.16× (e4b) vs exact. Fires only when the caller supplies blockMax (the
746
+ // production head does); MLX_BUN_CCE_BWD_BLOCK_EPS=0 disables (compiled out).
747
+ const BWD_BLOCK_EPS = process.env.MLX_BUN_CCE_BWD_BLOCK_EPS ?? "1e-5";
744
748
  const _bwdSgKernels = new Map<string, MetalKernel>();
745
749
  function bwdSgKernel(dimTiles: number, filterEps: string, blockEps: string): MetalKernel {
746
750
  const filterOn = Number.parseFloat(filterEps) > 0 ? 1 : 0;
@@ -1132,14 +1136,13 @@ export function flashCceForward(
1132
1136
  * (no atomics). Caller owns the result. */
1133
1137
  export function flashCceBackward(
1134
1138
  hResp: MlxArray, head: FlashCceHead, targets: number[], lse: MlxArray, cot: number[],
1135
- // Apple-CCE coeff filter eps. Defaults to BWD_FILTER_EPS, which is OFF ("0")
1136
- // unless MLX_BUN_CCE_BWD_FILTER_EPS is set — exact gradients by default. The
1137
- // measured real-data case FOR enabling it (eps 1e-5 = 0.16% dh, 1.35×; PLAN.md
1138
- // "filter-on-real-data") is a pending default flip — see BWD_FILTER_EPS's note.
1139
- // Pass "0" explicitly for guaranteed-exact (the parity gate does).
1139
+ // Apple-CCE coeff filter eps. Defaults to BWD_FILTER_EPS — ON at 1e-5 since
1140
+ // 2026-07-02 (real-data gates in BWD_FILTER_EPS's note). Pass "0" explicitly
1141
+ // for guaranteed-exact gradients (the parity gate does).
1140
1142
  filterEps: string = BWD_FILTER_EPS,
1141
1143
  // The forward's `blockMax [M, NBLK]` enables the vocab-block early exit (skips
1142
- // whole cold programs — attacks phase 1). Omit (or set blockEps "0") for exact.
1144
+ // whole cold programs — attacks BOTH phases); ON at 1e-5 by default when
1145
+ // blockMax is supplied. Omit blockMax (or set blockEps "0") for exact.
1143
1146
  blockMax?: MlxArray, blockEps: string = BWD_BLOCK_EPS,
1144
1147
  ): MlxArray {
1145
1148
  const M = hResp.shape[0]!;
@@ -1169,9 +1172,9 @@ export function flashCceBackward(
1169
1172
  // MLX_BUN_CCE_BWD_NOSTEEL=1 falls back to the old kernels. The steel backward carries
1170
1173
  // the SAME Apple-CCE skips as the SG kernel — coeff filter (phase-2 per-tile) +
1171
1174
  // blockMax vocab-block early exit (whole-program) — both compiled out at eps=0.
1172
- // BOTH skips are OFF by default (the envs default "0", nothing in production sets
1173
- // them): exact gradients unless explicitly opted in. Enabling the coeff filter at
1174
- // ~1e-5 is the measured-and-decided pending flip (PLAN.md "filter-on-real-data").
1175
+ // BOTH skips default ON at 1e-5 (2026-07-02 flip; measured gates in the
1176
+ // BWD_FILTER_EPS / BWD_BLOCK_EPS notes). MLX_BUN_CCE_BWD_{FILTER,BLOCK}_EPS=0
1177
+ // restores exact gradients (identical pre-flip kernels).
1175
1178
  const blockV0 = Math.ceil(V / NBLK);
1176
1179
  // HTw == group_size for the phase-2 QuantizedBlockLoader (one group per H-tile); needs
1177
1180
  // H % GS == 0 (always, GR*GS=H) and GS % 32 == 0 (TNv = GS/32 frags per H-tile integral).
package/src/train/job.ts CHANGED
@@ -130,8 +130,6 @@ export const finetuneRunner: JobRunner = async (emit, config) => {
130
130
  try {
131
131
  const isGemma = (await Bun.file(`${modelDir}/config.json`).text()).toLowerCase().includes("gemma");
132
132
  if (isGemma) {
133
- process.env.MLX_BUN_PERF_KERNEL ??= "0";
134
- process.env.MLX_BUN_FUSED_GELU ??= "0";
135
133
  }
136
134
  } catch {}
137
135