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.
- package/README.md +66 -21
- package/THIRD_PARTY_LICENSES.md +1 -0
- package/docs/archive/mlx-bun-lab-report.html +275 -0
- package/docs/archive/wwdc-mlx-bun.html +848 -0
- package/docs/design/audio-input-plan.md +356 -0
- package/docs/design/batching-perf-path.md +255 -0
- package/docs/design/batching-v2-plan.md +23 -9
- package/docs/design/bucketing-stage.md +2 -1
- package/docs/design/compat-cli-surface-design.md +4 -3
- package/docs/design/decode-speed-program.md +172 -0
- package/docs/design/docs-reorg-plan.md +1 -1
- package/docs/design/dspark-serving-program.md +371 -0
- package/docs/design/dspark-speculative-decoding.md +167 -5
- package/docs/design/faithful-l1-consolidation.md +310 -0
- package/docs/design/generic-model-support.md +8 -5
- package/docs/design/grammar-spec-batching-integration.md +301 -0
- package/docs/design/minicpm5-decode-megakernel.md +16 -11
- package/docs/design/mlx-lm-tool-parity-plan.md +11 -1
- package/docs/design/omlx-adoption-map.md +98 -0
- package/docs/design/orpo-training.md +34 -30
- package/docs/design/paged-kv-cache.md +214 -0
- package/docs/design/parallel-slots.md +82 -49
- package/docs/design/parity-tier-dag.md +43 -33
- package/docs/design/repo-cleanup-plan.md +189 -0
- package/docs/design/runtime-isolation.md +19 -3
- package/docs/design/segmented-backward-training.md +18 -2
- package/docs/design/ssd-kv-cold-tier.md +288 -0
- package/docs/design/structured-output.md +484 -0
- package/docs/design/the-dreaming-master-plan.md +2 -3
- package/docs/design/turboquant-kv.md +235 -0
- package/docs/design/unified-engine-frontier-plan.md +774 -0
- package/docs/design/web-chat-beat-matrix.md +315 -0
- package/docs/design/web-chat-redesign.md +1112 -0
- package/docs/investigations/decode-roofline-lookagain.md +92 -6
- package/docs/investigations/dspark-handoff.md +182 -16
- package/docs/investigations/kernel-perf-review-2026-07.md +105 -26
- package/docs/investigations/lab-build-journal.md +2 -2
- package/docs/investigations/openwiki-evaluation.md +249 -0
- package/docs/planning/IDEAS.md +14 -0
- package/docs/planning/PRODUCT_ROADMAP.md +2 -3
- package/docs/planning/chat-ui-vision-memory-first.md +199 -0
- package/docs/planning/chat-ui-vision.md +50 -0
- package/docs/planning/memory-docs-and-dag-plan.md +15 -6
- package/docs/planning/release-notes-v0.0.10.md +57 -0
- package/docs/planning/release-notes-v0.0.11.md +239 -0
- package/docs/reference/cli.md +35 -16
- package/docs/reference/distribution.md +13 -6
- package/docs/reference/embedding.md +35 -14
- package/docs/reference/features-matrix.md +97 -0
- package/docs/reference/library-api.md +30 -9
- package/docs/reference/memory.md +26 -6
- package/docs/reference/models.md +11 -1
- package/docs/reference/orpo-quickstart.md +3 -4
- package/docs/reference/server-api.md +634 -22
- package/docs/reference/server-config.md +251 -119
- package/docs/reference/training.md +20 -16
- package/package.json +5 -3
- package/src/anthropic.ts +13 -0
- package/src/audio/conformer.ts +647 -0
- package/src/audio/decode.ts +133 -0
- package/src/audio/features.ts +202 -0
- package/src/audio/transcode.ts +67 -0
- package/src/cli.ts +402 -107
- package/src/config.ts +49 -2
- package/src/download.ts +40 -0
- package/src/eval/README.md +7 -6
- package/src/eval/kl.ts +59 -3
- package/src/eval/quality-db.ts +2 -2
- package/src/eval/runner.ts +3 -3
- package/src/expert-offload.ts +15 -2
- package/src/flags.ts +13 -0
- package/src/generate.ts +330 -16
- package/src/grammar.ts +585 -0
- package/src/hub-rest.ts +242 -0
- package/src/kv-store.ts +618 -122
- package/src/lora.ts +8 -0
- package/src/md-modules.d.ts +42 -0
- package/src/media-fetch.ts +284 -0
- package/src/memory/rest.ts +245 -0
- package/src/memory/vault.ts +74 -0
- package/src/mlx/array.ts +127 -19
- package/src/mlx/autograd.ts +2 -2
- package/src/mlx/ffi.ts +27 -3
- package/src/mlx/metal-kernel.ts +7 -4
- package/src/mlx/ops.ts +179 -23
- package/src/mlx/turboquant-ops.ts +831 -0
- package/src/mlx/turboquant-tables.ts +237 -0
- package/src/model/batched-mask.ts +62 -1
- package/src/model/batched-quant.ts +206 -0
- package/src/model/batched-rotating-quant.ts +346 -0
- package/src/model/batched-rotating.ts +46 -1
- package/src/model/compiled-decode.ts +4 -15
- package/src/model/factory.ts +10 -4
- package/src/model/gemma4-base.ts +391 -34
- package/src/model/gemma4.ts +122 -33
- package/src/model/generated/gemma4-12b.ts +9 -13
- package/src/model/generated/gemma4-26b.ts +10 -5
- package/src/model/generated/gemma4-e4b.ts +14 -9
- package/src/model/minicpm5.ts +58 -102
- package/src/model/paged-kv.ts +253 -0
- package/src/model/qwen3-delta.ts +123 -15
- package/src/model/qwen3-moe.ts +404 -0
- package/src/model/qwen3.ts +9 -8
- package/src/model/qwen3_5.ts +99 -25
- package/src/model/support.ts +12 -1
- package/src/model/universal/dense.ts +40 -5
- package/src/model/universal/rope.ts +36 -21
- package/src/pi-web.ts +1299 -82
- package/src/preflight.ts +27 -8
- package/src/prompt-cache.ts +338 -32
- package/src/quantize/quantizer.ts +16 -1
- package/src/registry.ts +61 -3
- package/src/serve/batch-scheduler.ts +1001 -76
- package/src/serve/generation-gateway.ts +292 -30
- package/src/serve/isolate.ts +417 -0
- package/src/serve/lane-registry.ts +44 -0
- package/src/serve/session-search.ts +326 -0
- package/src/server.ts +1460 -105
- package/src/spec/assistant-source.ts +150 -0
- package/src/spec/deepspec-source.ts +143 -0
- package/src/spec/dflash-source.ts +124 -0
- package/src/spec/dspark/ab-stats.ts +136 -0
- package/src/spec/dspark/calibration.ts +112 -0
- package/src/spec/dspark/deepspec-module.ts +801 -0
- package/src/spec/dspark/generate-dflash.ts +72 -20
- package/src/spec/dspark/loader.ts +23 -0
- package/src/spec/dspark/module-dflash.ts +278 -36
- package/src/spec/dspark/quantize-drafter.ts +91 -0
- package/src/spec/generate.ts +59 -25
- package/src/spec/ngram-source.ts +132 -0
- package/src/spec/serve-loop.ts +536 -0
- package/src/spec/source.ts +93 -0
- package/src/spec/two-model.ts +171 -0
- package/src/ssd-cache.ts +285 -0
- package/src/tokenizer.ts +288 -2
- package/src/tool-approvals.ts +131 -0
- package/src/tool-call.ts +408 -20
- package/src/train/flash-cce.ts +29 -26
- package/src/train/job.ts +0 -2
- package/src/train/loss.ts +53 -11
- package/src/train/prefix-shared.ts +30 -7
- package/src/train/segmented.ts +150 -46
- package/src/train/trainer.ts +22 -15
- package/src/vision/prompt.ts +217 -48
- package/src/web/app.html +1184 -1789
- package/src/web/app.js +5667 -0
- package/src/web/icon.svg +13 -0
- package/src/web/manifest.webmanifest +24 -0
- package/src/web/src/adapters-panel.ts +288 -0
- package/src/web/src/api.ts +57 -0
- package/src/web/src/assistant.ts +485 -0
- package/src/web/src/chat.ts +1050 -0
- package/src/web/src/composer.ts +959 -0
- package/src/web/src/dataset.ts +132 -0
- package/src/web/src/finetune.ts +267 -0
- package/src/web/src/hub.ts +428 -0
- package/src/web/src/main.ts +98 -0
- package/src/web/src/markdown.ts +0 -0
- package/src/web/src/memory-panel.ts +572 -0
- package/src/web/src/model-picker.ts +159 -0
- package/src/web/src/palette.ts +346 -0
- package/src/web/src/protocol.ts +62 -0
- package/src/web/src/quantize.ts +167 -0
- package/src/web/src/rag.ts +274 -0
- package/src/web/src/sessions.ts +469 -0
- package/src/web/src/shell.ts +804 -0
- package/src/web/src/status.ts +252 -0
- package/src/web/src/ui-catalog.ts +92 -0
- package/src/web/sw.js +71 -0
- package/src/web/vendor/README.md +61 -0
- package/src/web/vendor/hljs-LICENSE +29 -0
- package/src/web/vendor/hljs-theme.css +59 -0
- package/src/web/vendor/hljs.js +26 -0
- package/src/model/fused-decode-kernel.ts +0 -312
- package/src/model/fused-geglu-kernel.ts +0 -198
- package/src/model/fused-mlp-kernel.ts +0 -169
- package/src/model/fused-swiglu-kernel.ts +0 -240
- package/src/model/steel-linear-kernel.ts +0 -96
- 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)
|
|
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
|
-
|
|
90
|
-
|
|
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:
|
|
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
|
-
|
|
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
|
|
171
|
-
|
|
172
|
-
|
|
173
|
-
|
|
174
|
-
|
|
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
|
}
|
package/src/train/flash-cce.ts
CHANGED
|
@@ -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
|
|
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
|
|
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
|
|
727
|
-
//
|
|
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,
|
|
731
|
-
//
|
|
732
|
-
//
|
|
733
|
-
//
|
|
734
|
-
//
|
|
735
|
-
//
|
|
736
|
-
|
|
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
|
|
739
|
-
// blocks whose max prob < eps AND which lack the target → ~0
|
|
740
|
-
//
|
|
741
|
-
//
|
|
742
|
-
// with
|
|
743
|
-
|
|
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
|
|
1136
|
-
//
|
|
1137
|
-
//
|
|
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
|
|
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
|
|
1173
|
-
//
|
|
1174
|
-
//
|
|
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
|
|