mlx-bun 0.0.13 → 0.3.0

package/src/train/forward.ts

@@ -20,6 +20,7 @@ import { MlxArray } from "../mlx/array";
 import { Dtype } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
 import { createCausalMask, type Cache, type Mask } from "../model/gemma4-base";
+import { SSMCache } from "../model/qwen3-delta";
 import type { RuntimeModel } from "../model/factory";
 
 /** Stateless cache for the training forward. Training is always a single
@@ -30,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
@@ -57,6 +60,37 @@ export class TrainingCache implements Cache {
   }
 }
 
+/** Stateless DeltaNet stand-in (same re-runnability contract as
+ *  TrainingCache): state WRITES are discarded immediately, so a gradient-
+ *  checkpoint recompute re-enters with the identical null-state conditions;
+ *  reads always yield null (zeros inside the kernel) — a full-sequence
+ *  forward from t=0, exactly the training/perplexity semantics. Before this
+ *  existed, qwen3_5 routed its DeltaNet layers into plain TrainingCache and
+ *  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; }
+  set conv(v: MlxArray | null) { v?.dispose(); }
+  get recurrent(): MlxArray | null { return null; }
+  set recurrent(v: MlxArray | null) { v?.dispose(); }
+  advance(_n: number): void { /* offset pinned at 0 */ }
+  rowOffset(_i: number): number { return 0; }
+  updateAndFetch(): [MlxArray, MlxArray] {
+    throw new Error("TrainingSSMCache: DeltaNet layers do not use the KV path");
+  }
+  makeMask(_N: number, _w: number | null): Mask {
+    return { mode: "", arr: null }; // ssm_mask is None at B=1 (SSMCache parity)
+  }
+  state(): MlxArray[] { return []; }
+  isTrimmable(): boolean { return true; }
+  trim(_n: number): void { /* stateless */ }
+  dispose(): void { /* owns no arrays */ }
+}
+
 /** Run a full-sequence forward for training.
  *  @param ids int32 array [B, L].
  *  @returns logits [B, L, V] (caller owns; dispose when done).
@@ -90,7 +124,8 @@ export function trainForwardHidden(
   // B=1 / no padding: stateless offset-0 caches (one per donor layer), so the
   // forward is pure and safe to recompute under gradient checkpointing.
   const probe = model.makeCache();
-  const cache: Cache[] = probe.map(() => new TrainingCache());
+  const cache: Cache[] = probe.map((c) =>
+    c instanceof SSMCache ? new TrainingSSMCache() : new TrainingCache());
   for (const c of probe) c.dispose();
   try {
     return model.forwardHidden(ids, cache); // [B, L, hidden]
@@ -161,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

@@ -27,6 +27,7 @@ import { logitSoftcap } from "../model/gemma4-base";
 import { trainForward, trainForwardHidden } from "./forward";
 import { setLoraScale, type TrainableLora } from "./lora-params";
 import { flashCceForward, flashCceBackward, type FlashCceHead } from "./flash-cce";
+import { runtimeValue } from "../runtime-config";
 
 /** The quantized LM-head weights + final-logit softcap, accessed uniformly across
  *  models for the fused linear-CE head: Gemma is tied (`embed.asLinear`, softcap
@@ -1189,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
@@ -1197,7 +1198,7 @@ export function fusedRespLogpMean(
   // e4b time-tie point where flash's memory win takes over) takes the EXACT
   // fused head instead — short rows get the faster exact head for free.
   // MLX_BUN_FLASH_MIN_M=0 always honors flash (the pre-dispatch behavior).
-  const flashMinM = Number(process.env.MLX_BUN_FLASH_MIN_M ?? "1024");
+  const flashMinM = Number(runtimeValue("MLX_BUN_FLASH_MIN_M") ?? "1024");
   const useFlash = flash && (flashMinM <= 0 || M >= flashMinM);
   let op: CustomVjp;
   if (useFlash) {

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
@@ -38,6 +38,7 @@ import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
 import { peakMemory } from "../mlx/ffi";
 import { Vjp } from "../mlx/autograd";
+import { runtimeValue } from "../runtime-config";
 
 // MLX_BUN_SEG_MEM_LOG=1: log the within-step peak after each phase of the
 // segmented step (forward+boundaries / head vjp / each segment backward).
@@ -46,7 +47,7 @@ import { Vjp } from "../mlx/autograd";
 // (This is how the backlog-#2 refutation below was measured. An inter-segment
 // clearCache was also tried and had NO effect — the watermark is live memory,
 // not allocator cache retention.)
-const SEG_MEM = process.env.MLX_BUN_SEG_MEM_LOG === "1";
+const SEG_MEM = runtimeValue("MLX_BUN_SEG_MEM_LOG") === "1";
 const segMem = (label: string): void => {
   if (SEG_MEM) console.log(`  [seg-mem] ${label}: peak=${(peakMemory() / 1e9).toFixed(2)} GB`);
 };
@@ -68,7 +69,7 @@ import { prefixSharedCaches, prefixGatherIdx, branchLogpMeanGathered, blockSpars
 // memory at the TOP level but not when nested in the segmented mlx_vjp (its
 // per-chunk graph isn't freed incrementally there); Checkpoint does. See
 // docs/design/orpo-training.md.
-const SEG_HEAD = process.env.MLX_BUN_SEG_HEAD ?? "checkpoint";
+const SEG_HEAD = runtimeValue("MLX_BUN_SEG_HEAD") ?? "checkpoint";
 function boundedHeadFromHidden(
   model: MiniCPM5Model | Gemma4Model, h: MlxArray, ids: number[], mask: number[],
   chunkSize: number, sink: Array<{ dispose(): void }>, span?: HeadSpan,
@@ -82,7 +83,7 @@ import { rowLength, type SftBatch, type DpoBatch } from "./dataset";
 
 // Token-chunk size for the bounded SFT head (the ORPO paths take theirs from
 // cfg.orpoChunkSize; SFT has no per-run knob, 512 matches the ORPO default).
-const SFT_HEAD_CHUNK = Number(process.env.MLX_BUN_SEG_HEAD_CHUNK ?? "512");
+const SFT_HEAD_CHUNK = Number(runtimeValue("MLX_BUN_SEG_HEAD_CHUNK") ?? "512");
 
 /** [M,V]-bounded response-only CE from post-finalNorm hidden — the SFT
  *  segmented head (kernel-review backlog #8). Same span arithmetic and
@@ -129,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
@@ -1516,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

@@ -10,8 +10,9 @@
 
 import { MlxArray, pinnedBufferCount } from "../mlx/array";
 import { clearCache, peakMemory, activeMemory, cacheMemory, resetPeakMemory } from "../mlx/ffi";
+import { runtimeValue } from "../runtime-config";
 
-const MEM_LOG = process.env.MLX_BUN_MEM_LOG === "1";
+const MEM_LOG = runtimeValue("MLX_BUN_MEM_LOG") === "1";
 const gbStr = (b: number) => `${(b / 1e9).toFixed(2)}GB`;
 import { ValueAndGrad } from "../mlx/autograd";
 import * as ops from "../mlx/ops";
@@ -81,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
@@ -439,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.
 // ---------------------------------------------------------------------------
 
 // ---------------------------------------------------------------------------
@@ -472,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) {
@@ -531,19 +532,19 @@ 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
   // would otherwise ride silently into a crash mid-run. Refuse it for Gemma
   // until the >=2K re-validation lands; MiniCPM5 stays allowed.
-  if (process.env.MLX_BUN_TRAIN_ATTN === "flash") {
+  if (runtimeValue("MLX_BUN_TRAIN_ATTN") === "flash") {
     if (model instanceof Gemma4Model)
       throw new Error(
         "MLX_BUN_TRAIN_ATTN=flash is disabled for Gemma models: e4b SIGTRAPs on this " +

package/src/vision/preprocess.ts

@@ -6,7 +6,7 @@
 // 1/255 → patchify into 48×48 patch vectors with 2D grid positions →
 // pad to the 280 soft-token budget.
 //
-// Decode strategy (Bun ≥1.3.14): PNG goes straight through fast-png
+// Decode strategy (Bun ≥1.4.0): PNG goes straight through fast-png
 // (exact, lossless — the parity-golden path). Everything else (JPEG,
 // HEIC, AVIF, WebP, TIFF, GIF, BMP) is transcoded to lossless PNG by
 // Bun.Image — native OS codecs (ImageIO on macOS), EXIF auto-orient,

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);
@@ -26,7 +26,7 @@ import type { LoadedTokenizer } from "../tokenizer";
 import type { AudioTower } from "../audio/conformer";
 import { audioSoftTokenCount, decodeAudio } from "../audio/decode";
 import { extractMelFeatures } from "../audio/features";
-import { fetchMediaBytes } from "../media-fetch";
+import { fetchMediaBytes, videoMediaFetchPolicy } from "../media-fetch";
 
 /** Common contract for both vision towers (encoder-free gemma4_unified in
  *  ./embedder.ts and the SigLIP encoder in ./siglip.ts): preprocess image
@@ -117,6 +117,52 @@ export async function extractImages(
   return { messages: out, images };
 }
 
+/** Extract video bytes from content parts, rewriting to the template's
+ *  {type:"video"} form. Accepts {type:"video_url", video_url:{url}} (data:
+ *  or http(s), same SSRF guard as images with the larger video body cap)
+ *  and {type:"video"} with base64 `data`. Decode to frames happens later
+ *  (vision/video-frames.ts — the AVFoundation sidecar). */
+export async function extractVideos(
+  messages: ChatMessage[],
+): Promise<{ messages: ChatMessage[]; videos: Uint8Array[] }> {
+  const videos: Uint8Array[] = [];
+  const out: ChatMessage[] = [];
+  for (const m of messages) {
+    if (!Array.isArray(m.content)) {
+      out.push(m);
+      continue;
+    }
+    const parts: Array<Record<string, unknown>> = [];
+    for (const part of m.content) {
+      if (part.type === "video_url") {
+        const url = (part.video_url as { url: string } | undefined)?.url
+          ?? (part.video_url as unknown as string);
+        if (typeof url !== "string") throw new Error("video_url part missing url");
+        videos.push(await fetchMediaBytes(url, "video"));
+        parts.push({ type: "video" });
+      } else if (part.type === "video") {
+        if (typeof part.data === "string") {
+          // Same cap as fetched/data:-URL video bodies — the inline base64
+          // form must not be the uncapped route.
+          const cap = videoMediaFetchPolicy().maxBytes;
+          if (part.data.length * 0.75 > cap)
+            throw new Error(
+              `video part exceeds the ${Math.round(cap / 1024 / 1024)} MB cap`,
+            );
+          videos.push(Uint8Array.from(Buffer.from(part.data, "base64")));
+          parts.push({ type: "video" });
+        } else {
+          throw new Error("video part requires base64 `data`");
+        }
+      } else {
+        parts.push(part);
+      }
+    }
+    out.push({ ...m, content: parts });
+  }
+  return { messages: out, videos };
+}
+
 /** Extract audio bytes from OpenAI-style content parts, rewriting the parts
  *  to the template's {type:"audio"} form. Accepts the OpenAI-canonical
  *  {type:"input_audio", input_audio:{data:<b64>, format:...}} plus optiq's