mlx-bun 0.0.10 → 0.0.11

package/src/tool-call.ts

@@ -28,6 +28,12 @@ export interface ParsedToolCall {
   name: string;
   /** Parsed argument object (callers serialize for OpenAI's string field). */
   arguments: Record<string, unknown>;
+  /** Set only when strict parsing failed and a repair pass recovered the
+   *  call — never silent (Axis 7 self-healing beat row). `repairs` lists
+   *  every malformation class that had to be corrected, in application
+   *  order, so failures are observable (logs/telemetry) instead of masked. */
+  repaired?: true;
+  repairs?: RepairKind[];
 }
 
 interface ToolSpec {
@@ -60,7 +66,18 @@ export function gemmaArgsToJson(text: string): string {
   return out;
 }
 
-/** Parse every `call:name{...}` block in a decoded tool-call segment. */
+/** Parse every `call:name{...}` block in a decoded tool-call segment.
+ *
+ *  Strict scan first (unchanged behavior on well-formed input). On the two
+ *  ways a small model mangles THIS family's syntax specifically —
+ *  (a) generation truncated mid-call (an unterminated `<|"|>` string or a
+ *  brace that never closes, both hit_max_tokens artifacts) and (b) a
+ *  trailing comma surviving `gemmaArgsToJson`'s conversion to JSON — the
+ *  repair pass recovers by treating "ran off the end of the text" as an
+ *  implicit close (mirrors closeUnbalanced's brace-stack unwind, adapted
+ *  for the `<|"|>`-quoted syntax) and re-running stripTrailingCommas on the
+ *  converted JSON. Recovered calls are tagged repaired/repairs; a segment
+ *  that still can't be recovered throws exactly as before. */
 export function parseToolCalls(text: string): ParsedToolCall[] {
   const calls: ParsedToolCall[] = [];
   let i = 0;
@@ -71,10 +88,11 @@ export function parseToolCalls(text: string): ParsedToolCall[] {
     // balanced-brace scan, skipping <|"|>-delimited strings
     let depth = 0;
     let j = nameEnd;
+    let unterminatedString = false;
     while (j < text.length) {
       if (text.startsWith(QUOTE, j)) {
         const close = text.indexOf(QUOTE, j + QUOTE.length);
-        if (close === -1) throw new Error("unterminated <|\"|> string");
+        if (close === -1) { unterminatedString = true; break; }
         j = close + QUOTE.length;
         continue;
       }
@@ -86,11 +104,55 @@ export function parseToolCalls(text: string): ParsedToolCall[] {
       }
       j++;
     }
-    if (depth !== 0) throw new Error("unbalanced braces in tool call");
-    const argsBlock = text.slice(nameEnd, j + 1);
+    const truncated = unterminatedString || depth !== 0;
+    let argsBlock: string;
+    const repairs: RepairKind[] = [];
+    if (truncated) {
+      // Ran off the end of the text without closing — truncated generation.
+      // Close the dangling <|"|> string (if any) then append enough `}` to
+      // balance whatever braces are still open, using the SAME depth
+      // counter's remaining value (skip <|"|>-quoted spans again so a
+      // brace character inside a string doesn't miscount).
+      let closeDepth = 0;
+      let k = nameEnd;
+      let inQuote = false;
+      while (k < text.length) {
+        if (text.startsWith(QUOTE, k) && !inQuote) { inQuote = true; k += QUOTE.length; continue; }
+        if (inQuote) {
+          if (text.startsWith(QUOTE, k)) { inQuote = false; k += QUOTE.length; continue; }
+          k++;
+          continue;
+        }
+        const c = text[k];
+        if (c === "{") closeDepth++;
+        else if (c === "}") closeDepth--;
+        k++;
+      }
+      argsBlock = text.slice(nameEnd) + (inQuote ? QUOTE : "") + "}".repeat(Math.max(closeDepth, 1));
+      repairs.push("unbalanced_braces");
+      j = text.length - 1;
+    } else {
+      argsBlock = text.slice(nameEnd, j + 1);
+    }
+    let converted = gemmaArgsToJson(argsBlock);
+    let parsed: unknown;
+    try {
+      parsed = JSON.parse(converted);
+    } catch {
+      const commas = stripTrailingCommas(converted);
+      if (!commas.changed) throw new Error("unbalanced braces in tool call");
+      try {
+        parsed = JSON.parse(commas.out);
+      } catch {
+        throw new Error("unbalanced braces in tool call");
+      }
+      converted = commas.out;
+      repairs.push("trailing_commas");
+    }
     calls.push({
       name: m[1]!,
-      arguments: JSON.parse(gemmaArgsToJson(argsBlock)) as Record<string, unknown>,
+      arguments: parsed as Record<string, unknown>,
+      ...(repairs.length ? { repaired: true, repairs } : {}),
     });
     i = j + 1;
   }
@@ -159,28 +221,327 @@ function decodeToolValue(value: string, schema: Record<string, unknown> | null):
   }
 }
 
+// ---- Repair layer (Axis 7 self-healing beat row) ---------------------
+//
+// pi-ai (node_modules/@earendil-works/pi-ai/dist/utils/json-parse.js)
+// already ships a generic JSON repair pass for cloud-model streaming
+// tool calls: `repairJson`/`parseJsonWithRepair` fix invalid backslash
+// escapes and raw control characters inside strings, and
+// `parseStreamingJson` falls back to the `partial-json` package for
+// truncated mid-stream JSON. Its `validateToolArguments`
+// (utils/validation.js) additionally coerces JSON-Schema primitive
+// TYPES (string "3" -> number 3, etc.) once a call is already
+// structurally valid JSON. None of that touches the STRUCTURAL
+// malformations small local models actually produce here: single
+// quotes instead of double, trailing commas, unbalanced/truncated
+// braces, a whole call wrapped in a markdown ```json fence, stray
+// prose ("Sure, here's the call:") glued in front of the JSON, or the
+// function name landing inside `arguments` instead of the envelope
+// (toolCallFromPayload above already handles that last one for
+// already-valid JSON — the repair pass below reuses it once the text
+// is coerced to valid JSON). This layer is therefore purely additive
+// to pi-ai's, applied ONLY when strict parsing has already failed, and
+// every repair is tagged on the returned ParsedToolCall (`repaired`,
+// `repairs`) rather than silently swallowed — see ParsedToolCall above.
+export type RepairKind =
+  | "fenced_wrapper" // ```json ... ``` or ``` ... ``` wrapper stripped
+  | "prose_prefix" // leading non-JSON prose trimmed to the first { or [
+  | "single_quotes" // 'like this' -> "like this" (keys and string values)
+  | "trailing_commas" // {"a":1,} / [1,2,] -> valid JSON
+  | "unbalanced_braces" // truncated generation: missing closing }/]/" appended
+  | "python_literals" // True/False/None -> true/false/null
+  | "name_in_args"; // function name found inside `arguments` instead of the envelope
+
+/** Strip a ```json fenced block (or bare ```) wrapping the whole payload. */
+function stripFence(text: string): { out: string; changed: boolean } {
+  const m = /^\s*```(?:json)?\s*\n?([\s\S]*?)\n?```\s*$/i.exec(text.trim());
+  if (!m) return { out: text, changed: false };
+  return { out: m[1] ?? "", changed: true };
+}
+
+/** Trim stray prose before the first `{`/`[` and after the matching
+ *  document's last `}`/`]` (a small model narrating "Sure, here's the
+ *  call: {...}" or trailing chatter after the JSON). Leaves the text
+ *  alone if it doesn't start with prose (no-op safe to always run). */
+function stripProsePrefix(text: string): { out: string; changed: boolean } {
+  const trimmed = text.trim();
+  const firstBrace = trimmed.search(/[{[]/);
+  if (firstBrace <= 0) return { out: trimmed, changed: false };
+  // Only strip when what precedes genuinely looks like prose (no braces
+  // of its own) — never chew into legitimate leading whitespace-only JSON.
+  const prefix = trimmed.slice(0, firstBrace);
+  if (/[{}[\]]/.test(prefix)) return { out: trimmed, changed: false };
+  const lastClose = Math.max(trimmed.lastIndexOf("}"), trimmed.lastIndexOf("]"));
+  const out = lastClose >= firstBrace ? trimmed.slice(firstBrace, lastClose + 1) : trimmed.slice(firstBrace);
+  return { out, changed: true };
+}
+
+/** Convert Python-ish literals (True/False/None) used outside string
+ *  literals into JSON ones. Scans char-by-char tracking string state so
+ *  a value like "Nonetheless" or a string containing "True" is untouched. */
+function pythonLiteralsToJson(text: string): { out: string; changed: boolean } {
+  let out = "";
+  let changed = false;
+  let inString: '"' | "'" | null = null;
+  for (let i = 0; i < text.length; i++) {
+    const ch = text[i]!;
+    if (inString) {
+      out += ch;
+      if (ch === "\\") { out += text[++i] ?? ""; continue; }
+      if (ch === inString) inString = null;
+      continue;
+    }
+    if (ch === '"' || ch === "'") { inString = ch; out += ch; continue; }
+    const rest = text.slice(i);
+    const wordMatch = /^(True|False|None)\b/.exec(rest);
+    if (wordMatch) {
+      const replacement = wordMatch[1] === "True" ? "true" : wordMatch[1] === "False" ? "false" : "null";
+      out += replacement;
+      i += wordMatch[1]!.length - 1;
+      changed = true;
+      continue;
+    }
+    out += ch;
+  }
+  return { out, changed };
+}
+
+/** Convert single-quoted JSON (keys and string values) to double-quoted,
+ *  tracking string state so an apostrophe INSIDE a double-quoted string
+ *  (e.g. "don't") is left alone. Only fires when the text has no double
+ *  quotes at all in "quote position" — i.e. looks like it was written
+ *  with single quotes throughout — to avoid mangling a string that
+ *  legitimately contains a single quote inside double-quoted JSON. */
+function singleQuotesToDouble(text: string): { out: string; changed: boolean } {
+  if (!/'/.test(text)) return { out: text, changed: false };
+  // Heuristic gate: real JSON never opens a key/value with a bare `'`.
+  // If the text already parses with plain JSON.parse it wouldn't have
+  // reached here, so any '"'-delimited string is assumed intentional and
+  // left byte-for-byte; only '...'-delimited runs are retargeted.
+  let out = "";
+  let changed = false;
+  let inDouble = false;
+  let inSingle = false;
+  for (let i = 0; i < text.length; i++) {
+    const ch = text[i]!;
+    if (inDouble) {
+      out += ch;
+      if (ch === "\\") { out += text[++i] ?? ""; continue; }
+      if (ch === '"') inDouble = false;
+      continue;
+    }
+    if (inSingle) {
+      if (ch === "\\") { out += ch + (text[++i] ?? ""); continue; }
+      if (ch === "'") { out += '"'; inSingle = false; continue; }
+      if (ch === '"') { out += '\\"'; continue; } // embedded " must be escaped now
+      out += ch;
+      continue;
+    }
+    if (ch === '"') { inDouble = true; out += ch; continue; }
+    if (ch === "'") { inSingle = true; changed = true; out += '"'; continue; }
+    out += ch;
+  }
+  return { out, changed };
+}
+
+/** Remove trailing commas before a closing `}`/`]` (whitespace/newlines
+ *  between the comma and the close are preserved otherwise). */
+function stripTrailingCommas(text: string): { out: string; changed: boolean } {
+  let changed = false;
+  const out = text.replace(/,(\s*[}\]])/g, (_m, close) => { changed = true; return close; });
+  return { out, changed };
+}
+
+/** Append the minimal closers to balance braces/brackets/an unterminated
+ *  string — a small model that got truncated (or just stopped early)
+ *  mid-object. Walks once tracking a stack of open delimiters and string
+ *  state; on EOF, closes the open string (if any) then unwinds the stack
+ *  in reverse. No-op (unchanged) when already balanced. */
+function closeUnbalanced(text: string): { out: string; changed: boolean } {
+  const stack: ("}" | "]")[] = [];
+  let inString = false;
+  for (let i = 0; i < text.length; i++) {
+    const ch = text[i]!;
+    if (inString) {
+      if (ch === "\\") { i++; continue; }
+      if (ch === '"') inString = false;
+      continue;
+    }
+    if (ch === '"') { inString = true; continue; }
+    if (ch === "{") stack.push("}");
+    else if (ch === "[") stack.push("]");
+    else if (ch === "}" || ch === "]") stack.pop();
+  }
+  if (!inString && stack.length === 0) return { out: text, changed: false };
+  let out = text;
+  if (inString) out += '"';
+  out += stack.reverse().join("");
+  return { out, changed: true };
+}
+
+/** Try strict JSON.parse first (pi-ai's own escape-level repair already
+ *  ran upstream where applicable — see the comment above); on failure,
+ *  apply the structural repair passes in order, re-testing after each so
+ *  only the ones that were actually needed get tagged. Returns null (no
+ *  throw) when nothing recovers a parseable object — callers decide how
+ *  to surface that. */
+function repairJsonToolText(text: string): { value: unknown; repairs: RepairKind[] } | null {
+  const repairs: RepairKind[] = [];
+  let candidate = text;
+
+  const fence = stripFence(candidate);
+  if (fence.changed) { candidate = fence.out; repairs.push("fenced_wrapper"); }
+
+  const prose = stripProsePrefix(candidate);
+  if (prose.changed) { candidate = prose.out; repairs.push("prose_prefix"); }
+
+  const tryParse = (s: string): unknown | undefined => {
+    try { return JSON.parse(s); } catch { return undefined; }
+  };
+
+  let parsed = tryParse(candidate);
+  if (parsed !== undefined) return { value: parsed, repairs };
+
+  const literals = pythonLiteralsToJson(candidate);
+  if (literals.changed) {
+    parsed = tryParse(literals.out);
+    if (parsed !== undefined) return { value: parsed, repairs: [...repairs, "python_literals"] };
+    candidate = literals.out;
+  }
+
+  const commas = stripTrailingCommas(candidate);
+  if (commas.changed) {
+    parsed = tryParse(commas.out);
+    if (parsed !== undefined) return { value: parsed, repairs: [...repairs, "trailing_commas"] };
+    candidate = commas.out;
+  }
+
+  const closed = closeUnbalanced(candidate);
+  if (closed.changed) {
+    parsed = tryParse(closed.out);
+    if (parsed !== undefined) return { value: parsed, repairs: [...repairs, "unbalanced_braces"] };
+    candidate = closed.out;
+  }
+
+  const quotes = singleQuotesToDouble(candidate);
+  if (quotes.changed) {
+    parsed = tryParse(quotes.out);
+    if (parsed !== undefined) return { value: parsed, repairs: [...repairs, "single_quotes"] };
+    candidate = quotes.out;
+    // Single-quote conversion can itself leave a trailing comma or an
+    // unbalanced document (e.g. `{'a': 1,}` truncated) — retry those two
+    // passes once more on the requoted text before giving up.
+    const commas2 = stripTrailingCommas(candidate);
+    if (commas2.changed) {
+      parsed = tryParse(commas2.out);
+      if (parsed !== undefined) return { value: parsed, repairs: [...repairs, "single_quotes", "trailing_commas"] };
+      candidate = commas2.out;
+    }
+    const closed2 = closeUnbalanced(candidate);
+    if (closed2.changed) {
+      parsed = tryParse(closed2.out);
+      if (parsed !== undefined) return { value: parsed, repairs: [...repairs, "single_quotes", "unbalanced_braces"] };
+    }
+  }
+
+  return null;
+}
+
+/** Lenient JSON.parse used for argument STRINGS nested inside an already
+ *  JSON-valid envelope (e.g. `"arguments": "{'a': 1}"`, a small model
+ *  double-encoding with the wrong quote style). Falls back through the
+ *  same repair passes; throws (not silently {}) when nothing recovers,
+ *  same contract as JSON.parse for its callers. */
+function parseJsonLoose(text: string): unknown {
+  try {
+    return JSON.parse(text);
+  } catch (err) {
+    const repaired = repairJsonToolText(text);
+    if (repaired) return repaired.value;
+    throw err;
+  }
+}
+
+/** Build a ParsedToolCall from an already-parsed JSON payload (object or
+ *  {function:{name,arguments}} envelope). Shared by the strict and
+ *  repaired paths so name/arguments extraction logic lives once. Handles
+ *  the "name-in-args confusion" malformation: a small model sometimes
+ *  emits the function name INSIDE the arguments object instead of (or in
+ *  addition to) the envelope's own `name` field — e.g.
+ *  {"arguments":{"name":"read","path":"x"}} with no top-level name at
+ *  all. When the envelope has no usable name, an args-embedded `name` (or
+ *  `function`) string is promoted and stripped out of the arguments. */
+function toolCallFromPayload(
+  payload: unknown,
+): { name: string; arguments: Record<string, unknown>; nameInArgs: boolean } | null {
+  if (!payload || typeof payload !== "object") return null;
+  const obj = payload as Record<string, any>;
+  const fn = obj.function && typeof obj.function === "object" ? obj.function : obj;
+  let name = typeof fn.name === "string" ? fn.name.trim() : "";
+  let args = fn.arguments ?? obj.arguments ?? {};
+  if (typeof args === "string") {
+    const parsed = args.trim() ? parseJsonLoose(args) : {};
+    args = parsed && typeof parsed === "object" && !Array.isArray(parsed) ? parsed : {};
+  }
+  if (!args || typeof args !== "object" || Array.isArray(args)) return null;
+  const argsObj = { ...(args as Record<string, unknown>) };
+  let nameInArgs = false;
+  if (!name) {
+    const embedded = argsObj.name ?? argsObj.function;
+    if (typeof embedded === "string" && embedded.trim()) {
+      name = embedded.trim();
+      delete argsObj.name;
+      delete argsObj.function;
+      nameInArgs = true;
+    }
+  }
+  if (!name) return null;
+  return { name, arguments: argsObj, nameInArgs };
+}
+
+/** Shape guard for the wrapper-less last-resort repair path only (see its
+ *  call site in parseGeneratedToolCalls). A `<tool_call>`/`<function...>`
+ *  tag already proves the model intended a tool call, so toolCallFromPayload
+ *  can be lenient there; a bare JSON object with NO such tag is ambiguous —
+ *  it could just as easily be an ordinary conversational reply that happens
+ *  to be JSON-shaped (e.g. `{"name": "read", "age": 42}` answering "give me
+ *  an example record"). Require an actual envelope shape: a `function`
+ *  wrapper object, or a top-level/nested `arguments` object (including the
+ *  name-in-args case, where the name lives inside that same `arguments`
+ *  object) — never just a bare `name`/`function` string sitting next to
+ *  unrelated fields. */
+function looksLikeToolEnvelope(payload: unknown): boolean {
+  if (!payload || typeof payload !== "object" || Array.isArray(payload)) return false;
+  const obj = payload as Record<string, unknown>;
+  if (obj.function && typeof obj.function === "object" && !Array.isArray(obj.function)) return true;
+  const args = obj.arguments;
+  if (args && typeof args === "object" && !Array.isArray(args)) return true;
+  if (typeof args === "string") return true; // double-encoded arguments string
+  return false;
+}
+
 function parseJsonToolCall(block: string): ParsedToolCall | null {
   let payload: unknown;
+  let repairs: RepairKind[] = [];
   try {
     payload = JSON.parse(block);
   } catch {
-    return null;
+    // Strict JSON.parse failed — this is the ONLY point the syntax-repair
+    // pass engages (never attempted when the model emitted valid JSON). A
+    // repair failure here returns null (not throw): parseGeneratedToolCalls
+    // falls through to the XML parsers next, which is the right behavior
+    // for a `<tool_call>` block that isn't actually JSON-shaped at all.
+    const repaired = repairJsonToolText(block);
+    if (!repaired) return null;
+    payload = repaired.value;
+    repairs = repaired.repairs;
   }
   if (!payload || typeof payload !== "object") throw new Error("JSON tool_call payload must be an object");
-  const obj = payload as Record<string, any>;
-  const fn = obj.function && typeof obj.function === "object" ? obj.function : obj;
-  const name = typeof fn.name === "string" ? fn.name.trim() : "";
-  if (!name) throw new Error("JSON tool_call is missing a function name");
-  const args = fn.arguments ?? obj.arguments ?? {};
-  if (typeof args === "string") {
-    const parsed = args.trim() ? JSON.parse(args) : {};
-    if (!parsed || typeof parsed !== "object" || Array.isArray(parsed))
-      throw new Error("tool_call arguments must be a JSON object");
-    return { name, arguments: parsed as Record<string, unknown> };
-  }
-  if (!args || typeof args !== "object" || Array.isArray(args))
-    throw new Error("tool_call arguments must be a JSON object");
-  return { name, arguments: args as Record<string, unknown> };
+  const result = toolCallFromPayload(payload);
+  if (!result) throw new Error("tool_call arguments must be a JSON object");
+  if (result.nameInArgs) repairs = [...repairs, "name_in_args"];
+  const { nameInArgs: _nameInArgs, ...call } = result;
+  return repairs.length ? { ...call, repaired: true, repairs } : call;
 }
 
 function parseXmlEqualsToolCall(block: string, tools: ToolSpec[]): ParsedToolCall | null {
@@ -235,6 +596,33 @@ export function parseGeneratedToolCalls(text: string, tools: ToolSpec[]): Parsed
   }
   TOOL_CALL_BLOCK_RE.lastIndex = 0;
   if (calls.length === 0) calls.push(...parseXmlAttrToolCalls(text, tools));
+  // Last-resort repair: no `<tool_call>` tags AND no `<function name=...>`
+  // XML matched at all — a small model sometimes drops the wrapper tags
+  // entirely and emits just (possibly fenced/prose-prefixed/malformed)
+  // JSON. Only attempted when everything else found nothing, so a
+  // genuinely tool-free assistant reply never risks being misread as a
+  // call. Tagged repaired like every other path here.
+  //
+  // Unlike the `<tool_call>`-wrapped path, there is no surrounding markup
+  // here to prove the model actually intended a tool call — an ordinary
+  // conversational reply can itself BE a bare JSON object (e.g. "give me
+  // an example record" -> `{"name": "read", "age": 42}`), which happens to
+  // collide with a real tool name. Require the recovered payload to
+  // actually look like a tool-call envelope (an `arguments`/
+  // `function.arguments` object, or a `function` wrapper, or the
+  // name-in-args shape) before accepting it — a flat object whose only
+  // tool-shaped feature is a `name` key that matches a tool is NOT enough
+  // signal on its own.
+  if (calls.length === 0) {
+    const repaired = repairJsonToolText(text.trim());
+    if (repaired && looksLikeToolEnvelope(repaired.value)) {
+      const result = toolCallFromPayload(repaired.value);
+      if (result) {
+        const repairs = result.nameInArgs ? [...repaired.repairs, "name_in_args" as const] : repaired.repairs;
+        calls.push({ name: result.name, arguments: result.arguments, repaired: true, repairs });
+      }
+    }
+  }
   for (const c of calls) {
     if (!known.has(c.name)) throw new Error(`unknown tool '${c.name}'`);
   }

package/src/train/flash-cce.ts

@@ -21,7 +21,7 @@
 // DIVERGENCES (why this isn't a line-for-line port):
 //   1. Our classifier C is the QUANTIZED head (4/8-bit affine). Apple loads C
 //      dense (bf16); we keep the in-Metal dequant for `tl.dot` (the qdot pattern
-//      from src/model/fused-decode-kernel.ts). This is the one substantive change.
+//      from mlx's quantized.h qdot pattern). This is the one substantive change.
 //   2. The head is FROZEN (not a LoRA target) → we need only dE (=dh), NOT dC.
 //   3. The LSE cross-vocab-block merge is done in cheap MLX ops on per-block
 //      partials, not Metal atomic-logaddexp spinlocks (simpler, same result).
@@ -289,7 +289,7 @@ const FWD_SG_SOURCE = String.raw`
 // memory each step (coalesced). The lane-partial dots reduce with simd_sum. The
 // threadgroup's 4 simdgroups process 4 rows in parallel; BLOCK_B tokens amortize
 // the dequant (each row dequantized once, reused across tokens). Same coalesced
-// pattern as fused-decode-kernel.ts. h is passed TRANSPOSED [H, M].
+// pattern as mlx's quantized.h qdot. h is passed TRANSPOSED [H, M].
 const SGN = TG / 32; // simdgroups per threadgroup (4)
 const FWD_LANE_SOURCE = String.raw`
   const uint tokBlk = threadgroup_position_in_grid.z;

package/src/train/job.ts

@@ -130,8 +130,6 @@ export const finetuneRunner: JobRunner = async (emit, config) => {
   try {
     const isGemma = (await Bun.file(`${modelDir}/config.json`).text()).toLowerCase().includes("gemma");
     if (isGemma) {
-      process.env.MLX_BUN_PERF_KERNEL ??= "0";
-      process.env.MLX_BUN_FUSED_GELU ??= "0";
     }
   } catch {}
 

package/src/train/trainer.ts

@@ -20,7 +20,6 @@ import type { RuntimeModel } from "../model/factory";
 import { Gemma4Model, type GradCheckpointCtx } from "../model/gemma4";
 import { MiniCPM5Model } from "../model/minicpm5";
 import { setTrainingAttn } from "../model/flash-attention";
-import { setFusedGeluTraining } from "../model/fused-geglu-kernel";
 import { SegmentedBackward, SegmentedBackwardGemma4, SegmentedBackwardOrpo, SegmentedBackwardOrpoGemma4, SegmentedBackwardOrpoPrefix, SegmentedBackwardOrpoPrefixGemma4, planSegmentsBySize } from "./segmented";
 import type { LoadedTokenizer } from "../tokenizer";
 import type { ChatTemplate } from "../chat-template";
@@ -260,7 +259,6 @@ export async function trainLora(
   // autograd flows and the backward recomputes the gelu from the primal instead
   // of retaining the spelled-out intermediates. No-op for non-Gemma models.
   // Cleared in the finally below.
-  setFusedGeluTraining(true);
 
   const saveCfg: SaveAdapterConfig = {
     rank: cfg.rank,
@@ -383,7 +381,6 @@ export async function trainLora(
     });
     return { adapterPath: cfg.adapterPath, appliedRanks, numIters: result.numIters };
   } finally {
-    setFusedGeluTraining(false);
     disposeLora(lora);
   }
 }
@@ -489,7 +486,7 @@ async function sftLoop(
   // scripts/experiments/flash-fd-check.ts), though ~30× slower than ops.sdpa.
   // Cleared in finally.
   //
-  // GEMMA GUARD (same convention as the PERF_KERNEL/FUSED_GELU sanitization in
+  // GEMMA GUARD (same convention as the launcher's env sanitization in
   // cli.ts / job.ts, but enforced at the trainer): e4b on this path SIGTRAPed
   // (uncatchable native crash) at multi-K sequence lengths (>=2K, reproduced in
   // scripts/experiments/segmented-grad-test-e4b.ts; docs/reference/training.md)