mlx-bun 0.2.0 → 0.4.0

package/src/ssd-durability.ts

@@ -1,10 +1,12 @@
+import { cloneAttachments } from "./backends/mlx/checkpoint-state";
+import { cleanupFailure, disposeResources } from "./engine/resources";
 import type { Cache } from "./model/gemma4";
 import type { PromptCacheEntry } from "./prompt-cache";
 import type { SpillItem, SpillQueue } from "./kv-store";
 
 export interface DurabilityGateway {
   readonly busy: boolean;
-  runExclusive<T>(fn: () => Promise<T>): Promise<T>;
+  runWhenIdle<T>(fn: () => Promise<T>): Promise<T>;
 }
 
 export interface DurabilityPromptCache {
@@ -115,9 +117,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 {
@@ -128,14 +130,17 @@ export class SsdDurabilityCoordinator {
   async #store(rec: DirtySnapshot): Promise<StoreOutcome> {
     let snap: SpillItem | null = null;
     try {
-      snap = await this.gateway.runExclusive(async () => {
+      snap = await this.gateway.runWhenIdle(async () => {
+        if (this.isAlreadyDurable(rec.tokens, rec.ns)) return null;
         const entry = this.promptCache.findExact(rec.tokens, rec.ns);
         if (!entry) return null;
-        return {
-          tokens: [...entry.tokens],
-          caches: this.cloneCaches(entry.caches),
-          ns: rec.ns,
-        };
+        const caches = this.cloneCaches(entry.caches);
+        try {
+          return { tokens: [...entry.tokens], caches, ns: rec.ns,
+            attachments: cloneAttachments(entry.attachments) };
+        } catch (error) {
+          return cleanupFailure(error, () => disposeResources(caches));
+        }
       });
     } catch {
       return "failed";
@@ -150,7 +155,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 +165,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 +202,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);

package/src/vision/qwen3vl-tower.ts

@@ -115,9 +115,42 @@ export class Qwen3VLVisionTower {
     return names;
   }
 
+  /** Checkpoint-generation naming for the tower. Pre-5.8 artifacts store
+   *  `vision_tower.*`; the transformers-5.8 export stores `model.visual.*`
+   *  (mlx-lm's text-only qwen3_5 drops BOTH prefixes in `Model.sanitize`, so
+   *  the mapping between them is ours). Verified 2026-09-01 across
+   *  Qwen3.8-27B-tqalloc-norot vs the staged pre-5.8 artifact: 331/333
+   *  tensors byte-identical under the rename, patch_embed.proj.weight
+   *  identical under the layout permutation below, and merger.linear_fc2
+   *  differs only because the staged artifact is rotation-folded. */
+  static #sourceName(canonical: string, prefix: string): string {
+    return prefix + canonical.slice("vision_tower.".length);
+  }
+
+  /** HF stores the Conv3d patch embed as (out, in, kT, kH, kW); MLX (and this
+   *  tower) want (out, kT, kH, kW, in). Proven a pure permutation on the
+   *  artifacts above — transpose(0,2,3,4,1) of the 5.8 tensor is bit-equal to
+   *  the pre-5.8 one. */
+  static #fixPatchEmbed(w: Map<string, MlxArray>): void {
+    const key = "vision_tower.patch_embed.proj.weight";
+    const a = w.get(key);
+    if (!a || a.shape.length !== 5) return;
+    if (a.shape[4] === 3) return; // already MLX layout
+    const t = ops.transposeAxes(a, [0, 2, 3, 4, 1]);
+    const c = ops.contiguous(t);
+    t.dispose();
+    a.dispose();
+    w.set(key, c);
+  }
+
   /** Load one safetensors file's map and move `wanted` tensors into `out`
-   *  (missing names are skipped; the map is always freed). */
-  static #takeFromFile(path: string, wanted: string[], out: Map<string, MlxArray>): void {
+   *  under their CANONICAL (`vision_tower.*`) names; `source` maps each
+   *  canonical name to the stored one. Missing names are skipped; the map is
+   *  always freed. */
+  static #takeFromFile(
+    path: string, wanted: string[], out: Map<string, MlxArray>,
+    source: (canonical: string) => string = (n) => n,
+  ): void {
     const arrMap = new BigUint64Array([C.mlx_map_string_to_array_new()]);
     const metaMap = new BigUint64Array([C.mlx_map_string_to_string_new()]);
     const arrMapPtr = ptr(arrMap);
@@ -136,7 +169,7 @@ export class Qwen3VLVisionTower {
         if (out.has(name)) continue;
         const slot = new BigUint64Array([C.mlx_array_new()]);
         const slotPtr = ptr(slot);
-        if (C.mlx_map_string_to_array_get(slotPtr, handle, ptr(cstr(name))) !== 0) {
+        if (C.mlx_map_string_to_array_get(slotPtr, handle, ptr(cstr(source(name)))) !== 0) {
           C.mlx_array_free(read.u64(slotPtr, 0));
           continue;
         }
@@ -156,28 +189,37 @@ export class Qwen3VLVisionTower {
         // OptiQ-convention sidecar (their artifacts + ours pre-2026-08-18).
         Qwen3VLVisionTower.#takeFromFile(sidecar, names, weights);
       } else {
-        // In-main fallback: the bf16 vision_tower.* tensors live in the main
+        // In-main fallback: the bf16 tower tensors live in the main
         // safetensors shards (the mlx-vlm convention — one copy serves both
-        // ecosystems; sidecar duplication retired 2026-08-18).
+        // ecosystems; sidecar duplication retired 2026-08-18). The 5.8-family
+        // export names them `model.visual.*` instead of `vision_tower.*`.
         const idxPath = `${modelDir}/model.safetensors.index.json`;
         let files: string[];
+        let prefixes = ["vision_tower.", "model.visual."];
         if (existsSync(idxPath)) {
           const idx = JSON.parse(readFileSync(idxPath, "utf8")) as
             { weight_map: Record<string, string> };
-          files = [...new Set(
-            names.map((n) => idx.weight_map[n]).filter((f): f is string => !!f),
-          )];
+          prefixes = prefixes.filter((p) =>
+            Qwen3VLVisionTower.#sourceName(names[0]!, p) in idx.weight_map);
+          files = [...new Set(prefixes.flatMap((p) => names
+            .map((n) => idx.weight_map[Qwen3VLVisionTower.#sourceName(n, p)])
+            .filter((f): f is string => !!f)))];
           if (files.length === 0)
-            throw new Error(`no vision_tower tensors in ${modelDir} (index has none)`);
+            throw new Error(`no vision tower tensors in ${modelDir} (index has none)`);
         } else {
           files = ["model.safetensors"];
         }
         for (const f of files)
-          Qwen3VLVisionTower.#takeFromFile(`${modelDir}/${f}`, names, weights);
+          for (const p of prefixes)
+            Qwen3VLVisionTower.#takeFromFile(
+              `${modelDir}/${f}`, names, weights,
+              (n) => Qwen3VLVisionTower.#sourceName(n, p),
+            );
       }
       for (const name of names)
         if (!weights.has(name))
           throw new Error(`qwen vision weights missing tensor ${name} (sidecar or in-main)`);
+      Qwen3VLVisionTower.#fixPatchEmbed(weights);
     } catch (e) {
       for (const [, a] of weights) a.dispose();
       throw e;

package/src/web/app.html

@@ -1216,7 +1216,7 @@
 <!-- PWA installability (plan §9 Phase 3, beat-matrix Axis 10): manifest +
      inline SVG icon (src/web/manifest.webmanifest, src/web/icon.svg — no
      binary PNGs, per the hygiene gate; some browsers won't render an SVG
-     app icon and that's an accepted tradeoff, see features-matrix.md).
+     app icon and that's an accepted tradeoff, see docs/reference/server-config.md).
      theme-color matches the app's --bg token so the OS chrome doesn't
      flash a mismatched color on launch. -->
 <link rel="manifest" href="/manifest.webmanifest">