mlx-bun 0.2.0 → 0.3.0

package/src/ssd-cache.ts

@@ -1,4 +1,4 @@
-// SSD cold tier for the prompt/KV cache (docs/design/ssd-kv-cold-tier.md).
+// SSD cold tier for the prompt/KV cache (docs/design/kv-cache.md).
 //
 // Files ARE the database: no sidecar index. Layout
 //   <dir>/<configFingerprint>/<nsHash>/<uuid>.mlxkv
@@ -22,7 +22,7 @@ import { randomUUID } from "node:crypto";
 import { commonPrefixLength } from "./prompt-cache";
 import {
   saveKvCache, saveKvCacheAsync, loadKvCache, readKvHeader,
-  cacheHeadersTrimmable,
+  cacheHeadersTrimmable, legacyCacheCodecs, type CacheCodecProvider,
   type KvSaveMeta, type KvLoadExpect, type LoadedKvCache,
 } from "./kv-store";
 import type { Cache } from "./model/gemma4-base";
@@ -39,9 +39,11 @@ export interface SsdIndexEntry {
    *  over a shorter usable one (the 2026-07-06 restart-0 defect: the
    *  [prompt+gen] file always won, always got rejected after restore). */
   trimmable: boolean;
+  generationCheckpoint?: NonNullable<KvSaveMeta["generationCheckpoint"]>;
 }
 
 export interface SsdStoreOptions {
+  codecs?: CacheCodecProvider;
   dir: string;
   maxBytes: number;
   /** Server identity — written into every file, enforced on scan/load. */
@@ -57,6 +59,7 @@ const nsHash = (ns: string): string =>
 
 export class SsdCacheStore {
   readonly #opts: SsdStoreOptions;
+  readonly #codecs: CacheCodecProvider;
   readonly #root: string; // <dir>/<configFingerprint>
   #index: SsdIndexEntry[] = [];
   #warnedWriteFailure = false;
@@ -64,6 +67,7 @@ export class SsdCacheStore {
 
   constructor(opts: SsdStoreOptions) {
     this.#opts = opts;
+    this.#codecs = opts.codecs ?? legacyCacheCodecs;
     this.#root = join(opts.dir, opts.configFingerprint);
   }
 
@@ -121,10 +125,14 @@ export class SsdCacheStore {
             h.tokenizerHash !== this.#opts.tokenizerHash
           )
             throw new Error("metadata mismatch");
+          if ((h.codecProvider ?? legacyCacheCodecs.id) !== this.#codecs.id) continue;
           const st = statSync(path);
           this.#index.push({
             path, ns: h.ns ?? "", tokens: h.tokens, bytes: st.size, mtimeMs: st.mtimeMs,
-            trimmable: cacheHeadersTrimmable(h.caches),
+            trimmable: cacheHeadersTrimmable(h.caches, this.#codecs),
+            ...(h.generationCheckpoint
+              ? { generationCheckpoint: h.generationCheckpoint }
+              : {}),
           });
         } catch {
           try { rmSync(path, { force: true }); } catch {}
@@ -171,7 +179,7 @@ export class SsdCacheStore {
         ns: entry.ns,
         verify: this.#opts.verify,
       };
-      const loaded = loadKvCache(entry.path, model, expect);
+      const loaded = loadKvCache(entry.path, model, expect, this.#codecs);
       const now = new Date();
       try { utimesSync(entry.path, now, now); entry.mtimeMs = now.getTime(); } catch {}
       this.stats.restores++;
@@ -198,7 +206,7 @@ export class SsdCacheStore {
     const path = join(dir, `${randomUUID()}.mlxkv`);
     try {
       mkdirSync(dir, { recursive: true });
-      saveKvCache(path, tokens, caches, this.#meta(ns));
+      saveKvCache(path, tokens, caches, this.#meta(ns), this.#codecs);
       return this.#indexStored(path, tokens, caches, ns);
     } catch (err) {
       return this.#storeFailed(path, err);
@@ -210,21 +218,75 @@ export class SsdCacheStore {
    *  yields the event loop between tensors so serving interleaves.
    *  Caller passes zero-copy CLONES it owns (a consistent snapshot no
    *  matter what the live entry does meanwhile) and disposes them after.
-   *  `waitTurn` gates every per-tensor step (see saveKvCacheAsync) — the
-   *  server passes the gateway's onIdle so the flush only progresses while
-   *  the engine is idle, never mid-decode. */
-  async storeAsync(tokens: number[], caches: Cache[], ns = "", waitTurn?: () => Promise<void>): Promise<boolean> {
+   *  `runStep` owns every per-tensor step (see saveKvCacheAsync) — the
+   *  server passes the gateway's exclusive runner so a request cannot start
+   *  between an idle check and the blocking MLX readback. */
+  async storeAsync(
+    tokens: number[], caches: Cache[], ns = "",
+    runStep?: <T>(step: () => T) => Promise<T>,
+  ): Promise<boolean> {
     const dir = join(this.#root, nsHash(ns));
     const path = join(dir, `${randomUUID()}.mlxkv`);
     try {
       mkdirSync(dir, { recursive: true });
-      await saveKvCacheAsync(path, tokens, caches, this.#meta(ns), waitTurn);
+      await saveKvCacheAsync(path, tokens, caches, this.#meta(ns), runStep, this.#codecs);
       return this.#indexStored(path, tokens, caches, ns);
     } catch (err) {
       return this.#storeFailed(path, err);
     }
   }
 
+  /** Longest durable in-flight generation for an identical request. The
+   *  checkpoint namespace is separate from the ordinary prompt-cache tier,
+   *  so incomplete assistant output can never seed an unrelated prefill. */
+  findGenerationCheckpoint(
+    prompt: number[], key: string, cacheNs = "",
+  ): SsdIndexEntry | null {
+    let best: SsdIndexEntry | null = null;
+    for (const entry of this.#index) {
+      const checkpoint = entry.generationCheckpoint;
+      if (!checkpoint || checkpoint.key !== key || checkpoint.cacheNs !== cacheNs)
+        continue;
+      if (checkpoint.originalPromptTokens !== prompt.length) continue;
+      if (entry.tokens.length <= prompt.length) continue;
+      if (commonPrefixLength(entry.tokens, prompt) !== prompt.length) continue;
+      if (!best || entry.tokens.length > best.tokens.length) best = entry;
+    }
+    return best;
+  }
+
+  /** Persist one in-flight generation atomically. A completed newer
+   *  checkpoint supersedes older checkpoints for the same request only after
+   *  its rename succeeds, so a crash during a write leaves the prior point. */
+  async storeGenerationCheckpoint(
+    tokens: number[], caches: Cache[], checkpoint: NonNullable<KvSaveMeta["generationCheckpoint"]>,
+  ): Promise<boolean> {
+    const ns = `__generation_checkpoint__:${checkpoint.key}`;
+    const dir = join(this.#root, nsHash(ns));
+    const path = join(dir, `${randomUUID()}.mlxkv`);
+    try {
+      mkdirSync(dir, { recursive: true });
+      await saveKvCacheAsync(path, tokens, caches, {
+        ...this.#meta(ns), generationCheckpoint: checkpoint,
+      }, undefined, this.#codecs);
+      const stored = this.#indexStored(path, tokens, caches, ns, checkpoint);
+      if (stored) {
+        for (const entry of [...this.#index]) {
+          if (entry.path === path) continue;
+          if (entry.generationCheckpoint?.key === checkpoint.key) this.remove(entry.path);
+        }
+      }
+      return stored;
+    } catch (err) {
+      return this.#storeFailed(path, err);
+    }
+  }
+
+  removeGenerationCheckpoints(key: string): void {
+    for (const entry of [...this.#index])
+      if (entry.generationCheckpoint?.key === key) this.remove(entry.path);
+  }
+
   #meta(ns: string): KvSaveMeta {
     return {
       modelId: this.#opts.modelId,
@@ -234,7 +296,10 @@ export class SsdCacheStore {
     };
   }
 
-  #indexStored(path: string, tokens: number[], caches: Cache[], ns: string): boolean {
+  #indexStored(
+    path: string, tokens: number[], caches: Cache[], ns: string,
+    generationCheckpoint?: NonNullable<KvSaveMeta["generationCheckpoint"]>,
+  ): boolean {
     const st = statSync(path);
     if (st.size > this.#opts.maxBytes) {
       rmSync(path, { force: true });
@@ -260,7 +325,10 @@ export class SsdCacheStore {
         if (commonPrefixLength(e.tokens, tokens) === e.tokens.length) this.remove(e.path);
       }
     }
-    this.#index.push({ path, ns, tokens, bytes: st.size, mtimeMs: Date.now(), trimmable });
+    this.#index.push({
+      path, ns, tokens, bytes: st.size, mtimeMs: Date.now(), trimmable,
+      ...(generationCheckpoint ? { generationCheckpoint } : {}),
+    });
     this.stats.spills++;
     this.evictToCap();
     return true;

package/src/ssd-durability.ts

@@ -115,9 +115,9 @@ export class SsdDurabilityCoordinator {
     this.#attempts.set(key, task);
     try {
       const outcome = await task;
-      if (outcome !== "failed" && this.#dirty.get(key) === rec)
+      if (outcome === "stored" && this.#dirty.get(key) === rec)
         this.#dirty.delete(key);
-      if (outcome === "failed" && !force && this.#dirty.get(key) === rec)
+      if (outcome !== "stored" && !force && this.#dirty.get(key) === rec)
         this.#arm(key, this.busyRetryMs);
       return outcome;
     } finally {
@@ -129,6 +129,7 @@ export class SsdDurabilityCoordinator {
     let snap: SpillItem | null = null;
     try {
       snap = await this.gateway.runExclusive(async () => {
+        if (this.isAlreadyDurable(rec.tokens, rec.ns)) return null;
         const entry = this.promptCache.findExact(rec.tokens, rec.ns);
         if (!entry) return null;
         return {
@@ -150,7 +151,7 @@ export class SsdDurabilityCoordinator {
     this.#timers.clear();
 
     let flushedSnapshots = 0;
-    let missingSnapshots = 0;
+    const missing = new Map<string, DirtySnapshot>();
     const droppedBefore = this.spillQueue.droppedCount;
     const failedBefore = this.spillQueue.failedCount;
 
@@ -160,22 +161,33 @@ export class SsdDurabilityCoordinator {
 
     // Flush one entry at a time. This prevents the queue cap from dropping a
     // boundary snapshot while a large final snapshot is already in flight.
+    // Attempt each record version once. Missing state stays dirty, so a
+    // second flush cannot report success without a real durable snapshot.
+    const attempted = new Set<DirtySnapshot>();
     while (this.#dirty.size > 0) {
-      const keys = [...this.#dirty.keys()];
-      let progressed = false;
-      for (const key of keys) {
-        if (!this.#dirty.has(key)) continue;
+      const records = [...this.#dirty.entries()].filter(([, rec]) => !attempted.has(rec));
+      if (records.length === 0) break;
+      for (const [key, rec] of records) {
+        if (this.#dirty.get(key) !== rec) continue;
+        attempted.add(rec);
         const outcome = await this.#attempt(key, true);
         await this.spillQueue.drain();
         if (outcome === "stored") {
           flushedSnapshots++;
-          progressed = true;
         } else if (outcome === "missing") {
-          missingSnapshots++;
-          progressed = true;
+          missing.set(key, rec);
         }
       }
-      if (!progressed) break;
+    }
+
+    // A later trimmable snapshot may now cover a superseded RAM ancestor.
+    // Reconcile against committed SSD coverage after all writes settle;
+    // never clear a missing prefix merely because another write succeeded.
+    for (const [key, rec] of missing) {
+      if (!this.isAlreadyDurable(rec.tokens, rec.ns)) continue;
+      if (this.#dirty.get(key) === rec) this.#dirty.delete(key);
+      missing.delete(key);
+      flushedSnapshots++;
     }
 
     const stats = this.stats;
@@ -186,9 +198,9 @@ export class SsdDurabilityCoordinator {
         stats.pendingSpills === 0 &&
         stats.droppedSpills === droppedBefore &&
         stats.failedSpills === failedBefore &&
-        missingSnapshots === 0,
+        missing.size === 0,
       flushedSnapshots,
-      missingSnapshots,
+      missingSnapshots: missing.size,
       elapsedMs: performance.now() - started,
     };
   }

package/src/tool-call.ts

@@ -161,7 +161,7 @@ export function parseToolCalls(text: string): ParsedToolCall[] {
 
 const TOOL_CALL_BLOCK_RE = /<tool_call>\s*([\s\S]*?)\s*<\/tool_call>/gi;
 const XML_FUNCTION_EQUALS_RE = /^\s*<function=([^>\s]+)>\s*([\s\S]*?)\s*<\/function>\s*$/i;
-const XML_PARAMETER_EQUALS_RE = /<parameter=([^>\s]+)>\s*([\s\S]*?)\s*<\/parameter>/gi;
+const XML_PARAMETER_EQUALS_RE = /<parameter=([^>\s]+)>([\s\S]*?)<\/parameter>/gi;
 const GLM52_ARG_KEY_RE = /<arg_key>/i;
 const GLM52_ARG_PAIR_RE = /<arg_key>\s*([\s\S]*?)\s*<\/arg_key>\s*<arg_value>\s*([\s\S]*?)\s*<\/arg_value>/gi;
 const XML_FUNCTION_ATTR_RE = /<function\s+name=["']([^"']+)["']\s*>\s*([\s\S]*?)\s*<\/function>/gi;
@@ -208,11 +208,9 @@ function toolParameterSchema(
   return null;
 }
 
-/** Oracle `_decode_tool_parameter_value`: string-typed params stay raw
- *  text (a path like "2025" must not become a number); everything else
- *  is JSON-decoded when possible. */
-function decodeToolValue(value: string, schema: Record<string, unknown> | null): unknown {
-  const text = decodeXml(value);
+/** Decode a value after its caller removes protocol-specific delimiters.
+ *  String-typed params stay raw; everything else is JSON-decoded when possible. */
+function decodeToolValue(text: string, schema: Record<string, unknown> | null): unknown {
   if (!text) return "";
   const types = schemaTypeNames(schema);
   if (types.size && [...types].every((t) => t === "string" || t === "null")) return text;
@@ -546,23 +544,59 @@ function parseJsonToolCall(block: string): ParsedToolCall | null {
   return repairs.length ? { ...call, repaired: true, repairs } : call;
 }
 
-function parseXmlEqualsToolCall(block: string, tools: ToolSpec[]): ParsedToolCall | null {
+function parseXmlEqualsToolCall(block: string, tools: ToolSpec[], strict = false): ParsedToolCall | null {
   const match = XML_FUNCTION_EQUALS_RE.exec(block);
   XML_FUNCTION_EQUALS_RE.lastIndex = 0;
   if (!match) return null;
+  if (strict && (match.index !== 0 || match[0].length !== block.length)) return null;
   const name = match[1]!.trim();
   const body = match[2]!;
   const args: Record<string, unknown> = {};
+  let cursor = 0;
   for (const param of body.matchAll(XML_PARAMETER_EQUALS_RE)) {
+    if (strict && body.slice(cursor, param.index).trim()) return null;
     const key = param[1]!.trim();
     if (!key) throw new Error(`tool '${name}' contains an empty parameter name`);
-    args[key] = decodeToolValue(param[2]!, toolParameterSchema(tools, name, key));
+    if (strict && Object.hasOwn(args, key)) return null;
+    // Qwen's format carries raw text, not XML entity-encoded text. Match
+    // mlx_lm.tool_parsers.qwen3_coder: remove one framing newline at each
+    // end, preserving indentation, literal entities and JSON string escapes.
+    const value = param[2]!.replace(/^\n|\n$/g, "");
+    args[key] = decodeToolValue(value, toolParameterSchema(tools, name, key));
+    cursor = param.index + param[0].length;
   }
   XML_PARAMETER_EQUALS_RE.lastIndex = 0;
+  if (strict && body.slice(cursor).trim()) return null;
   return { name, arguments: args };
 }
 
-function parseGlm52ToolCall(block: string, tools: ToolSpec[]): ParsedToolCall | null {
+/** One complete envelope, with full consumption and no syntax repair. Used
+ * to establish structural fill boundaries, never to verify model predictions. */
+export function parseStrictToolCall(text: string, tools: ToolSpec[]): ParsedToolCall | null {
+  const open = "<tool_call>", close = "</tool_call>";
+  if (!text.startsWith(open) || !text.endsWith(close)) return null;
+  const body = text.slice(open.length, -close.length).trim();
+  if (body.includes(open) || body.includes(close)) return null;
+  try {
+    let call: ParsedToolCall | null;
+    if (body.startsWith("{")) {
+      // JSON.parse must succeed before the existing payload interpretation.
+      JSON.parse(body);
+      call = parseJsonToolCall(body);
+    } else {
+      call = parseXmlEqualsToolCall(body, tools, true) ?? parseGlm52ToolCall(body, tools, true);
+    }
+    if (!call || call.repaired || !tools.some(tool => toolSpecName(tool) === call.name)) return null;
+    return call;
+  } catch { return null; }
+}
+
+/** A textual argument delimiter is complete only after its closing bracket. */
+export function hasCompleteToolValueDelimiter(text: string): boolean {
+  return /^(?:<\/parameter>|<\/arg_value>)/.test(text.trimStart());
+}
+
+function parseGlm52ToolCall(block: string, tools: ToolSpec[], strict = false): ParsedToolCall | null {
   const firstArg = block.search(GLM52_ARG_KEY_RE);
   const name = (firstArg < 0 ? block : block.slice(0, firstArg)).trim();
   if (!name || !/^[^\s<>]+$/.test(name)) return null;
@@ -576,7 +610,8 @@ function parseGlm52ToolCall(block: string, tools: ToolSpec[]): ParsedToolCall |
     if (body.slice(cursor, pair.index).trim()) return null;
     const key = pair[1]!.trim();
     if (!key) throw new Error(`tool '${name}' contains an empty argument name`);
-    args[key] = decodeToolValue(pair[2]!, toolParameterSchema(tools, name, key));
+    if (strict && Object.hasOwn(args, key)) return null;
+    args[key] = decodeToolValue(decodeXml(pair[2]!), toolParameterSchema(tools, name, key));
     cursor = pair.index + pair[0].length;
     matched = true;
   }
@@ -594,7 +629,7 @@ function parseXmlAttrToolCalls(text: string, tools: ToolSpec[]): ParsedToolCall[
     for (const param of body.matchAll(XML_PARAM_ATTR_RE)) {
       const key = param[1]!.trim();
       if (!key) throw new Error(`tool '${name}' contains an empty parameter name`);
-      args[key] = decodeToolValue(param[2]!, toolParameterSchema(tools, name, key));
+      args[key] = decodeToolValue(decodeXml(param[2]!), toolParameterSchema(tools, name, key));
     }
     XML_PARAM_ATTR_RE.lastIndex = 0;
     calls.push({ name, arguments: args });
@@ -616,9 +651,11 @@ export function parseGeneratedToolCalls(text: string, tools: ToolSpec[]): Parsed
   const calls: ParsedToolCall[] = [];
   for (const block of text.matchAll(TOOL_CALL_BLOCK_RE)) {
     const body = block[1]!.trim();
-    const parsed = parseJsonToolCall(body) ??
-      parseXmlEqualsToolCall(body, tools) ??
-      parseGlm52ToolCall(body, tools);
+    // A valid Qwen parameter may itself contain JSON. Dispatch by the outer
+    // format before JSON repair can mistake that value for the tool envelope.
+    const parsed = /^<function=/i.test(body)
+      ? parseXmlEqualsToolCall(body, tools)
+      : parseJsonToolCall(body) ?? parseGlm52ToolCall(body, tools);
     if (!parsed) throw new Error("unsupported tool_call payload format");
     calls.push(parsed);
   }

package/src/train/flash-cce.ts

@@ -728,7 +728,7 @@ const BWD_SG_SOURCE = String.raw`
 // (flipped 2026-07-02 — the kernel-perf-review backlog #1 flip). The old
 // 0.66%→2.7% dh-accuracy scare was on RANDOM targets (flat softmax, the filter's
 // worst case — see PLAN.md "filter-on-real-data"); the real-data gate
-// (scripts/experiments/flash-cce-filter-realdata.ts) measured, on REAL chunk-ORPO
+// (flash-cce-filter-realdata.ts (deleted 2026-08-23; git history)) measured, on REAL chunk-ORPO
 // data (M1 Max, 2026-07-02): CPM5 eps 1e-5 = 0.343% dh for 1.41×; e4b = 0.158% dh
 // for 1.70× — and the teacher-forced full-logits fidelity shows the filter's added
 // error is far below the flash-vs-full fp-reassociation scatter that already
@@ -781,7 +781,7 @@ function bwdSgKernel(dimTiles: number, filterEps: string, blockEps: string): Met
 //     accumulated into D's frags LANE-LOCALLY (frag += frag — no get_coord math, the
 //     per-lane element mapping is identical between temp and D). dh atomic-added
 //     across the NBLK vocab-block programs. Validated dh ~1e-4 rel vs autograd.
-//     (Design + parity: scripts/experiments/steel-bwd-htile-test.ts.)
+//     (Design + parity: steel-bwd-htile-test.ts (deleted 2026-08-23; git history).)
 //   PHASE-2 DEQUANT: the Wd[32,HTw] staging uses MLX's vectorized+fused
 //     QuantizedBlockLoader (the SAME loader the forward uses), NOT a manual scalar
 //     dequant. The wrinkle: phase-2's GEMM contracts over vocab (K=32) and outputs H

package/src/train/forward.ts

@@ -31,6 +31,8 @@ import type { RuntimeModel } from "../model/factory";
  *  backward recompute, and a stateful (appending) cache would corrupt on the
  *  second run. */
 export class TrainingCache implements Cache {
+  /** Training-only adapter; never admitted, merged, or persisted. */
+  signature(): string { return "train:training"; }
   offset = 0;
   updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
     // Return fresh views: the caller disposes its input k/v right after (its
@@ -67,6 +69,8 @@ export class TrainingCache implements Cache {
  *  threw on the missing conv/recurrent/advance surface (the recorded
  *  "mlx-bun perplexity cannot score qwen3_5" gap). */
 export class TrainingSSMCache implements Cache {
+  /** Training-only adapter; never admitted, merged, or persisted. */
+  signature(): string { return "train:training-ssm"; }
   offset = 0;
   specRound: null = null;
   get conv(): MlxArray | null { return null; }
@@ -192,6 +196,8 @@ export function buildBatchedPadMask(
  *  forwardLayers only disposes the mask it pulled from makeMask, so makeMask
  *  hands back a fresh view each call and keeps the canonical copy itself. */
 class BatchedMaskCache implements Cache {
+  /** Training-only adapter; never admitted, merged, or persisted. */
+  signature(): string { return "train:batched-mask"; }
   constructor(
     private readonly inner: Cache,
     private readonly window: number | null,

package/src/train/loss.ts

@@ -1190,7 +1190,7 @@ export function fusedRespLogpMean(
   // [chunk,vocabBlock]. Otherwise the whole-vocab analytic head below.
   const vb = vocabBlock > 0 && vocabBlock < V ? vocabBlock : 0;
   // Auto-dispatch by M (kernel backlog #3; measured 2026-07-02 on the M1 Max,
-  // scripts/experiments/head-dispatch-sweep.ts): the fused QM head is FASTER
+  // head-dispatch-sweep.ts (deleted 2026-08-23; git history)): the fused QM head is FASTER
   // at short M (e4b M=64 1.6×, M=512 1.13×; CPM5 faster at EVERY M) — the
   // flash head's win is residency, plus time at long M on e4b (M≥2048 after
   // the 2026-07-02 filter/blockMax flip; 8K: 10.7 vs 13.2 s). So a

package/src/train/prefix-shared.ts

@@ -11,7 +11,7 @@
 // prompt dominates (our chunk/document fine-tunes), ~0 when the response does.
 //
 // Bit-exact with the two-forward path (validated in
-// scripts/experiments/prefix-shared-parity.ts): each predicted token uses a hidden
+// prefix-shared-parity.ts (deleted 2026-08-23; git history)): each predicted token uses a hidden
 // conditioned on exactly the same context + RoPE positions —
 //   chosen[k]   from H[P-1+k]                 (k=0..Rc-1; H[P-1] = prompt's last)
 //   rejected[k] from H[P-1] (k=0) / H[P+Rc+k-1] (k≥1)
@@ -62,6 +62,8 @@ export function blockSparseMask(P: number, Rc: number, Rr: number): MlxArray {
  *  TrainingCache) but `makeMask` returns the block-sparse mask. One per layer;
  *  MiniCPM5's runLayerRange calls makeMask once and disposes the returned arr. */
 export class PrefixSharedCache implements Cache {
+  /** Training-only adapter; never admitted, merged, or persisted. */
+  signature(): string { return "train:prefix-shared"; }
   offset = 0;
   constructor(
     private readonly P: number, private readonly Rc: number, private readonly Rr: number,
@@ -371,6 +373,8 @@ export function blockSparsePrefixMaskGemma(P: number, Rc: number, Rr: number, wi
  *  type serves both (the window arg selects the sliding term). One per DONOR
  *  layer (sharers reuse donors' fetched KV). */
 class Gemma4PrefixSharedCache implements Cache {
+  /** Training-only adapter; never admitted, merged, or persisted. */
+  signature(): string { return "train:gemma4prefix-shared"; }
   offset = 0;
   constructor(private readonly P: number, private readonly Rc: number, private readonly Rr: number) {}
   updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {

package/src/train/segmented.ts

@@ -2,7 +2,7 @@
 // segment-by-segment so only ONE segment's activations are ever live, instead
 // of holding every layer's recompute activations at once (the spike that makes
 // naive gradient checkpointing crash at long context). Design + proof:
-// docs/design/segmented-backward-training.md (mechanism validated on a toy in
+// docs/design/orpo-training.md (mechanism validated on a toy in
 // scripts/ckpt-mem-test.ts, SEG mode: 0.000% grad error, peak below per-layer
 // checkpointing; bit-exact on MiniCPM5 under flash attention).
 //
@@ -24,7 +24,7 @@
 // `value_and_grad(sum(stop_grad(dh) (.) segment_forward(...)))` is numerically
 // identical but LEAKS ~one activation buffer per segment per step at the mlx
 // level (measured; not GC/cache/synchronize/reuse reclaimable — see
-// docs/design/segmented-backward-training.md §9.4). mlx_vjp takes the cotangent
+// docs/design/orpo-training.md §9.4). mlx_vjp takes the cotangent
 // directly, needs no surrogate, and does not leak.
 //
 // IMPORTANT lifetime fact: mlx `eval` does NOT detach — an eval'd array retains
@@ -130,7 +130,7 @@ export function planSegmentsBySize(nLayers: number, segSize: number): [number, n
 // segment_size 1 measures 14.59 GB @8K (+3% step time) vs 17.5 @seg2. Below
 // that needs the [M,V]-bounded SFT head (backlog #8, ~3 GB head-vjp spike)
 // and/or an O(L)-memory attention backward. Evidence:
-// scripts/experiments/seg-isolation-smoke.ts + MLX_BUN_SEG_MEM_LOG=1 probes.
+// seg-isolation-smoke.ts (deleted 2026-08-23; git history) + MLX_BUN_SEG_MEM_LOG=1 probes.
 
 /** Detach `a` into a graph-free LEAF holding its exact (evaluated) bytes, then
  *  dispose `a`. mlx retains the upstream graph after `eval`, so a segment
@@ -1517,6 +1517,8 @@ export class SegmentedBackwardOrpoGemma4 {
  *  and null for full layers; the window arg selects the sliding term. One per DONOR
  *  layer (sharers reuse donors' fetched KV inside runLayerRange). */
 class Gemma4PrefixSharedSegCache implements Cache {
+  /** Training-only adapter; never admitted, merged, or persisted. */
+  signature(): string { return "train:gemma4-prefix-shared-seg"; }
   offset = 0;
   constructor(private readonly P: number, private readonly Rc: number, private readonly Rr: number) {}
   updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {

package/src/train/trainer.ts

@@ -82,7 +82,7 @@ export interface TrainConfig {
    *  beats per-layer checkpointing at long context, where naive checkpointing
    *  holds every layer's recompute activations at once and crashes. B=1 SFT
    *  only (the responseOnlyCe path); mutually exclusive with gradCheckpoint.
-   *  Numerically identical to off. See docs/design/segmented-backward-training.md. */
+   *  Numerically identical to off. See docs/design/orpo-training.md. */
   segmentSize: number;
   /** Keep ALL eval-step checkpoints (full mountable adapters under
    *  checkpoints/step-<NNNNN>-val<loss>/) + write metrics.json. Off by default
@@ -440,12 +440,12 @@ export async function trainLora(
 
 // ---------------------------------------------------------------------------
 // Segment planning: uniform planSegmentsBySize only. A full-attention
-// ISOLATION planner (kernel-review backlog #2 / segmented-backward-training.md
+// ISOLATION planner (kernel-review backlog #2 / docs/design/orpo-training.md
 // §5) was built and A/B-measured 2026-07-02: ZERO peak win — mlx's sdpa
 // backward is O(L²) for EVERY layer (~3.5 GB/layer @8K e4b), sliding included,
 // so the worst segment is set by layer count alone and segment_size is the
 // whole knob (seg1 = 14.59 GB @8K, +3% step time — fits the 24 GB M4 Pro).
-// Evidence: scripts/experiments/seg-isolation-smoke.ts; note in segmented.ts.
+// Evidence: seg-isolation-smoke.ts (deleted 2026-08-23; git history); note in segmented.ts.
 // ---------------------------------------------------------------------------
 
 // ---------------------------------------------------------------------------
@@ -473,7 +473,7 @@ async function sftLoop(
   // Segmented backward: stream the SFT backward segment-by-segment (only one
   // segment's activations live at a time). Phase A is MiniCPM5 SFT B=1 only;
   // it replaces the single value_and_grad below and is mutually exclusive with
-  // gradient checkpointing (see docs/design/segmented-backward-training.md).
+  // gradient checkpointing (see docs/design/orpo-training.md).
   const useSegmented = cfg.segmentSize > 0;
   let segmented: SegmentedBackward | SegmentedBackwardGemma4 | null = null;
   if (useSegmented) {
@@ -532,13 +532,13 @@ async function sftLoop(
   // hand-rolled flash kernel is an opt-in (MLX_BUN_TRAIN_ATTN=flash) O(L)-memory
   // path for memory-bound long context — now correct (its dK-transpose and
   // dQ-causal-barrier bugs were fixed; FD-validated in
-  // scripts/experiments/flash-fd-check.ts), though ~30× slower than ops.sdpa.
+  // flash-fd-check.ts (deleted 2026-08-23; git history)), though ~30× slower than ops.sdpa.
   // Cleared in finally.
   //
   // 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)
+  // segmented-grad-test-e4b.ts (deleted 2026-08-23; git history); docs/reference/training.md)
   // and has NOT been re-validated at that scale since the two kernel fixes —
   // the regression tests stop at T<=256. cli.ts even defaults e4b seq to 8192,
   // so a stale `export MLX_BUN_TRAIN_ATTN=flash` from a MiniCPM5 experiment

package/src/vision/prompt.ts

@@ -8,7 +8,7 @@
 // n = min(ceil(duration_ms/40), 750) (processing_gemma4.py) — embed the full
 // sequence, then overwrite the soft-token rows with the towers' features.
 //
-// Attention semantics (audio-input-plan.md §3.3, resolved vs the oracle):
+// Attention semantics (docs/design/generic-model-support.md §3.3, resolved vs the oracle):
 //   - image runs attend bidirectionally among themselves (bidirMask), BUT
 //     any audio token in the prompt disables the vision overlay entirely —
 //     mixed image+audio prompts run fully causal (bidirMask = null);