mlx-bun 0.0.9 → 0.0.11

package/src/train/loss.ts

@@ -138,6 +138,36 @@ export function sftLoss(model: RuntimeModel, batch: SftBatch): MlxArray {
   return scalar;
 }
 
+/** model.logitsFromHidden for TRAINING heads, with the upstream mlx
+ *  quantized_matmul bug worked around: transpose=false returns WRONG values at
+ *  row count 2 or 3 (reproduced on stock mlx 0.31.2; M=1 and M>=4 correct,
+ *  transpose=true correct at every M). mlx's autograd computes dh for a
+ *  transpose=true head qmm via an INTERNAL transpose=false call at the
+ *  cotangent's row count — so any differentiated [M,V] logits with M in {2,3}
+ *  (2-3-token responses, 2-3-token tail chunks, tiny prompt spans) got a
+ *  silently wrong dh (measured relnorm ~0.35-0.99 vs the f32 ground truth).
+ *  Fix: pad the span to 4 rows with zeros and slice the logits back — the pad
+ *  rows' cotangents are exactly zero (slice adjoint), and the internal vjp
+ *  then runs at M=4 where the kernel is correct. */
+export function logitsFromHiddenPadM(
+  model: { logitsFromHidden(h: MlxArray): MlxArray },
+  hSpan: MlxArray, // [1, M, hidden]
+): MlxArray {
+  const M = hSpan.shape[1]!;
+  if (M !== 2 && M !== 3) return model.logitsFromHidden(hSpan);
+  const [B, , H] = hSpan.shape as [number, number, number];
+  const pad = ops.zeros([B, 4 - M, H], hSpan.dtype);
+  const hPad = ops.concatAxis([hSpan, pad], 1);
+  pad.dispose();
+  const lgPad = model.logitsFromHidden(hPad); // [1, 4, V]
+  hPad.dispose();
+  const start = MlxArray.fromInt32(new Int32Array([0]), [1]);
+  const lg = ops.sliceDynamic(lgPad, start, [1], [B, M, lgPad.shape[2]!]);
+  start.dispose();
+  lgPad.dispose();
+  return lg;
+}
+
 /** B=1 masked CE that applies the LM head only at the supervised (response)
  *  span. The supervised input positions are the contiguous range
  *  [promptLen-1, len-1) (each predicts the next, response, token). The LM head
@@ -158,7 +188,7 @@ export function responseOnlyCe(model: RuntimeModel, h: MlxArray, batch: SftBatch
   const hidden = h.shape[2]!;
   const start = MlxArray.fromInt32(new Int32Array([startT]), [1]);
   const hResp = ops.sliceDynamic(h, start, [1], [1, M, hidden]);
-  const logits = model.logitsFromHidden(hResp); // [1, M, V]
+  const logits = logitsFromHiddenPadM(model, hResp); // [1, M, V] (M in {2,3} padded — mlx qmm bug)
   const V = logits.shape[2]!;
   const logits2d = ops.reshape(logits, [M, V]);
 
@@ -198,7 +228,7 @@ export function responseOnlyLogpMean(
   const hidden = h.shape[2]!;
   const start = MlxArray.fromInt32(new Int32Array([startT]), [1]);
   const hResp = ops.sliceDynamic(h, start, [1], [1, M, hidden]); // [1, M, hidden]
-  const logits = model.logitsFromHidden(hResp);                   // [1, M, V]
+  const logits = logitsFromHiddenPadM(model, hResp);              // [1, M, V] (M in {2,3} padded — mlx qmm bug)
   const V = logits.shape[2]!;
   const logits2d = ops.reshape(logits, [M, V]);
 
@@ -504,7 +534,7 @@ export function spanLogpMeanFromHidden(
   const hidden = h.shape[2]!;
   const start = MlxArray.fromInt32(new Int32Array([s0]), [1]);
   const hSpan = ops.sliceDynamic(h, start, [1], [1, n, hidden]); // [1, n, hidden]
-  const logits = model.logitsFromHidden(hSpan); // [1, n, V]
+  const logits = logitsFromHiddenPadM(model, hSpan); // [1, n, V] (n in {2,3} padded — mlx qmm bug)
   const V = logits.shape[2]!;
   const logits2d = ops.reshape(logits, [n, V]);
 
@@ -624,7 +654,7 @@ export function branchLogpMeanB1(
   const hidden = h.shape[2]!;
   const start = MlxArray.fromInt32(new Int32Array([startT]), [1]);
   const hResp = ops.sliceDynamic(h, start, [1], [1, M, hidden]); // [1, M, hidden]
-  const logits = model.logitsFromHidden(hResp); // [1, M, V] — head only on the response span
+  const logits = logitsFromHiddenPadM(model, hResp); // [1, M, V] — head only on the response span (M in {2,3} padded — mlx qmm bug)
   const V = logits.shape[2]!;
   const logits2d = ops.reshape(logits, [M, V]);
 
@@ -698,7 +728,7 @@ function chunkedLogpMeanB1(
 
     const ckpt = new Checkpoint((inp) => {
       const hc = inp[0]!;
-      const logits = model.logitsFromHidden(hc); // [1, Cc, V] — recomputed in backward
+      const logits = logitsFromHiddenPadM(model, hc); // [1, Cc, V] — recomputed in backward (Cc in {2,3} padded — mlx qmm bug)
       const V = logits.shape[2]!;
       const l2 = ops.reshape(logits, [Cc, V]);
       const targets = MlxArray.fromInt32(tHost, [Cc, 1]);
@@ -765,7 +795,7 @@ export function chunkedLogpMeanFromHidden(
     for (let i = 0; i < Cc; i++) tHost[i] = ids[c0 + 1 + i]!;
     const ckpt = new Checkpoint((inp) => {
       const hc = inp[0]!;
-      const logits = model.logitsFromHidden(hc); // [1, Cc, V] — recomputed in backward
+      const logits = logitsFromHiddenPadM(model, hc); // [1, Cc, V] — recomputed in backward (Cc in {2,3} padded — mlx qmm bug)
       const V = logits.shape[2]!;
       const l2 = ops.reshape(logits, [Cc, V]);
       const targets = MlxArray.fromInt32(tHost, [Cc, 1]);
@@ -1093,8 +1123,9 @@ function makeVocabBlockedHeadVjp(
 /** flash-CCE head as a CustomVjp: forward runs the Metal kernel → Σ logp (stashing
  *  `lse` for the backward, avoiding a forward recompute); backward runs the kernel's
  *  `dh` for unit cotangent then scales by the scalar cotangent `c` (loss = Σ logp →
- *  ∂loss/∂logp_t = c for all t). The coeff filter + vocab-block skip are OFF by
- *  default (env-gated MLX_BUN_CCE_BWD_*_EPS) — the backward computes the EXACT dh. */
+ *  ∂loss/∂logp_t = c for all t). The coeff filter + vocab-block skip default ON at
+ *  1e-5 (2026-07-02 flip; real-data gates in flash-cce.ts) — set
+ *  MLX_BUN_CCE_BWD_{FILTER,BLOCK}_EPS=0 for the EXACT dh. */
 function makeFlashCceHeadVjp(
   head: HeadQuant, targetsHost: Int32Array, M: number,
   sink: Array<{ dispose(): void }>,
@@ -1125,8 +1156,8 @@ function makeFlashCceHeadVjp(
       const c = cots[0]!;      // [1] = ∂loss/∂(Σ logp); uniform over tokens
       // dh for unit cotangent, then × c (avoids a host read of the cotangent).
       const ones = new Array(M).fill(1) as number[];
-      // filter + block-skip are OFF by default (env-gated) → EXACT dh. savedBlockMax is
-      // passed so the lossless vocab-block skip can be enabled via env without re-plumbing.
+      // filter + block-skip default ON at 1e-5 (env-overridable; =0 → exact dh).
+      // savedBlockMax feeds the ~lossless vocab-block skip.
       const dhUnit = flashCceBackward(hR, fch, targets, savedLse!, ones, undefined, savedBlockMax!); // [M, hidden] f32
       const cM = c.dtype === dhUnit.dtype ? c : c.astype(dhUnit.dtype);
       const dh = ops.mul(dhUnit, cM); // broadcast [M,hidden]·[1]
@@ -1157,8 +1188,19 @@ export function fusedRespLogpMean(
   // Vocab-blocked (full CCE) path: even [chunk,V] never materializes — only
   // [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
+  // 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
+  // requested-flash batch with M below MLX_BUN_FLASH_MIN_M (default 1024, the
+  // 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 useFlash = flash && (flashMinM <= 0 || M >= flashMinM);
   let op: CustomVjp;
-  if (flash) {
+  if (useFlash) {
     op = makeFlashCceHeadVjp(head, targetsHost, M, sink);
   } else if (vb > 0) {
     op = makeVocabBlockedHeadVjp(head, targetsHost, M, hidden, chunk, vb);

package/src/train/prefix-shared.ts

@@ -28,7 +28,7 @@ import * as ops from "../mlx/ops";
 import { createCausalMask, type Cache, type Mask } from "../model/gemma4-base";
 import { MiniCPM5Model, setMiniCpmPrefixPlan } from "../model/minicpm5";
 import { Gemma4Model, setGemmaPrefixPlan } from "../model/gemma4";
-import { orpoLossFromLogps, fusedRespLogpMean, combineFullNll, type ChunkCtx, type SftScope } from "./loss";
+import { orpoLossFromLogps, fusedRespLogpMean, combineFullNll, logitsFromHiddenPadM, type ChunkCtx, type SftScope } from "./loss";
 import type { DpoBatch } from "./dataset";
 
 /** A model that can project final-norm hidden states to vocab logits — the only
@@ -63,17 +63,40 @@ export function blockSparseMask(P: number, Rc: number, Rr: number): MlxArray {
  *  MiniCPM5's runLayerRange calls makeMask once and disposes the returned arr. */
 export class PrefixSharedCache implements Cache {
   offset = 0;
-  constructor(private readonly P: number, private readonly Rc: number, private readonly Rr: number) {}
+  constructor(
+    private readonly P: number, private readonly Rc: number, private readonly Rr: number,
+    // Step-shared mask memo (kernel-review backlog #9): the [T,T] block-sparse
+    // mask is IDENTICAL for every layer, segment forward, and backward
+    // recompute of a step, but each runLayerRange call rebuilt it from its
+    // elementwise ops (2·nSeg+ rebuilds per step in the segmented prefix
+    // class). With a memo it is built once; makeMask hands out cheap slice
+    // VIEWS (consumers dispose their view — runLayerRange does), and the
+    // first cache disposed frees the memo. Use prefixSharedCaches() to build
+    // a memo-sharing batch.
+    private readonly memo?: { mask: MlxArray | null },
+  ) {}
   updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
     return [k.slice([0, 0, 0, 0], k.shape), v.slice([0, 0, 0, 0], v.shape)];
   }
   makeMask(_N: number, _windowSize: number | null): Mask {
-    return { mode: "array", arr: blockSparseMask(this.P, this.Rc, this.Rr) };
+    if (!this.memo) return { mode: "array", arr: blockSparseMask(this.P, this.Rc, this.Rr) };
+    if (!this.memo.mask) this.memo.mask = blockSparseMask(this.P, this.Rc, this.Rr);
+    const m = this.memo.mask;
+    return { mode: "array", arr: m.slice([0, 0], m.shape) };
   }
   state(): MlxArray[] { return []; }
   isTrimmable(): boolean { return true; }
   trim(_n: number): void { /* offset pinned at 0 */ }
-  dispose(): void { /* owns no arrays */ }
+  dispose(): void {
+    if (this.memo?.mask) { this.memo.mask.dispose(); this.memo.mask = null; }
+  }
+}
+
+/** One PrefixSharedCache per layer, all sharing a single block-sparse-mask
+ *  memo for the step (see the constructor note). */
+export function prefixSharedCaches(nLayers: number, P: number, Rc: number, Rr: number): Cache[] {
+  const memo: { mask: MlxArray | null } = { mask: null };
+  return Array.from({ length: nLayers }, () => new PrefixSharedCache(P, Rc, Rr, memo));
 }
 
 /** Length-normalized mean log-prob over the predictions at hidden positions
@@ -87,7 +110,7 @@ function gatheredLogpMean(model: LogitProjector, h: MlxArray, gatherIdx: number[
   const hSel = ops.takeAxis(h, idxArr, 1); // [1, M, hidden]
   idxArr.dispose();
   if (hSel.shape[2] !== hidden) throw new Error("gatheredLogpMean: hidden mismatch");
-  const logits = model.logitsFromHidden(hSel); // [1, M, V]
+  const logits = logitsFromHiddenPadM(model, hSel); // [1, M, V] (M in {2,3} padded — mlx qmm bug)
   hSel.dispose();
   const V = logits.shape[2]!;
   const logits2d = ops.reshape(logits, [M, V]);
@@ -185,7 +208,7 @@ export function orpoLossPrefixShared(
   concat.set(chosenResp, P);
   concat.set(rejectedResp, P + Rc);
   const ids = MlxArray.fromInt32(concat, [1, T]);
-  const caches: Cache[] = model.layers.map(() => new PrefixSharedCache(P, Rc, Rr));
+  const caches: Cache[] = prefixSharedCaches(model.layers.length, P, Rc, Rr);
 
   let h: MlxArray;
   setMiniCpmPrefixPlan({ P, Rc, Rr });
@@ -230,7 +253,7 @@ export function prefixSharedLogps(
   const concat = new Int32Array(T);
   concat.set(promptIds, 0); concat.set(chosenResp, P); concat.set(rejectedResp, P + Rc);
   const ids = MlxArray.fromInt32(concat, [1, T]);
-  const caches: Cache[] = model.layers.map(() => new PrefixSharedCache(P, Rc, Rr));
+  const caches: Cache[] = prefixSharedCaches(model.layers.length, P, Rc, Rr);
   setMiniCpmPrefixPlan({ P, Rc, Rr });
   let h: MlxArray;
   try { h = model.forwardHidden(ids, caches); }

package/src/train/segmented.ts

@@ -36,18 +36,31 @@
 
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
+import { peakMemory } from "../mlx/ffi";
 import { Vjp } from "../mlx/autograd";
+
+// 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).
+// The trainer resets the peak at step start (MLX_BUN_MEM_LOG=1), so the phase
+// where the monotone peak JUMPS is where the step's memory actually lives.
+// (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 segMem = (label: string): void => {
+  if (SEG_MEM) console.log(`  [seg-mem] ${label}: peak=${(peakMemory() / 1e9).toFixed(2)} GB`);
+};
 import { MiniCPM5Model, setMiniCpmPrefixPlan } from "../model/minicpm5";
 import type { Gemma4Model } from "../model/gemma4";
 import { setGemmaPrefixPlan } from "../model/gemma4";
 import type { Cache, Mask, SharedKv } from "../model/gemma4-base";
 import { TrainingCache } from "./forward";
 import {
-  responseOnlyCe, responseOnlyLogpMean, fusedLogpMeanFromHidden, chunkedLogpMeanFromHidden,
+  responseOnlyLogpMean, fusedLogpMeanFromHidden, chunkedLogpMeanFromHidden,
   orpoLossFromLogps, spanLogpMeanFromHidden, respSpanFromMask, combineFullNll,
   type ChunkCtx, type SftScope, type HeadSpan,
 } from "./loss";
-import { PrefixSharedCache, prefixGatherIdx, branchLogpMeanGathered, blockSparsePrefixMaskGemma, prefixFullNll } from "./prefix-shared";
+import { prefixSharedCaches, prefixGatherIdx, branchLogpMeanGathered, blockSparsePrefixMaskGemma, prefixFullNll } from "./prefix-shared";
 
 // Experiment toggle: which bounded head to use inside the segmented mlx_vjp.
 // "checkpoint" = per-chunk Checkpoint recompute (bounds memory under nesting);
@@ -65,7 +78,36 @@ function boundedHeadFromHidden(
     : chunkedLogpMeanFromHidden(model, h, ids, mask, chunkSize, sink, span);
 }
 import type { TrainableLora } from "./lora-params";
-import { type SftBatch, type DpoBatch } from "./dataset";
+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");
+
+/** [M,V]-bounded response-only CE from post-finalNorm hidden — the SFT
+ *  segmented head (kernel-review backlog #8). Same span arithmetic and
+ *  token-mean reduction as loss.ts responseOnlyCe, but the logits only ever
+ *  exist per token-chunk ([chunk,V] Checkpoint recompute, or the fused
+ *  CustomVjp head via MLX_BUN_SEG_HEAD) instead of the full [1,M,V] + its
+ *  cotangent inside the head vjp. CE = −(mean response logp); per-position
+ *  grads are bit-exact vs the whole-vocab head (chunking is over TOKEN
+ *  positions; each position's full-V lse is computed unchanged — only the
+ *  scalar loss VALUE can differ in fp association across chunks). */
+export function boundedSftCe(
+  model: MiniCPM5Model | Gemma4Model, hn: MlxArray, batch: SftBatch,
+  sink: Array<{ dispose(): void }>, chunkSize = SFT_HEAD_CHUNK,
+): MlxArray {
+  const L = batch.ids[0]!.length;
+  const startT = Math.max(0, batch.promptLens[0]! - 1); // first supervised input position
+  const M = Math.min(rowLength(batch, 0), L) - 1 - startT;
+  if (M <= 0) throw new Error("sftLoss: no response tokens to supervise");
+  const mean = boundedHeadFromHidden(model, hn, batch.ids[0]!, [], chunkSize, sink, { startT, M });
+  const neg = ops.mulScalar(mean, -1); // CE = −(mean logp); exact sign flip
+  mean.dispose();
+  const loss = ops.sumAxis(neg, 0, false); // [1] -> scalar, matches the vjp cotangent `one`
+  neg.dispose();
+  return loss;
+}
 
 /** Split `nLayers` into contiguous `[lo, hi)` ranges of at most `segSize`
  *  layers each (the last range may be shorter). `segSize` is the memory knob:
@@ -77,24 +119,50 @@ export function planSegmentsBySize(nLayers: number, segSize: number): [number, n
   return ranges;
 }
 
+// NOTE (2026-07-02, kernel-review backlog #2 REFUTED BY MEASUREMENT): a
+// "full-attention isolation" planner (singleton segments for e4b's 7 full
+// layers) was built and A/B'd @8K — ZERO peak win (18.09 vs 18.02 GB). The
+// design §5 premise is false: mlx's sdpa BACKWARD materializes the O(L²)
+// scores for EVERY layer (~3.5 GB/layer @8K on e4b — sliding's window is just
+// an additive mask; measured: sliding pair +7.1 GB ≈ sliding+full +7.15). The
+// worst segment is set by LAYER COUNT alone, so segSize IS the whole knob:
+// 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.
+
 /** 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
  *  boundary (or a value_and_grad output) would otherwise keep its whole layer
  *  stack's activations alive. Copying the bytes into a fresh leaf frees that
  *  graph. The leaf is a valid differentiable input for the per-segment vag. */
 function detachLeaf(a: MlxArray): MlxArray {
-  // Force a row-major copy first: `rawBytes` reads the underlying buffer
-  // linearly, so a NON-contiguous input (e.g. a transposed donor K/V view) would
-  // be byte-scrambled. `ops.contiguous` is a no-op for already-contiguous arrays
+  // Force a row-major copy first: the buffer is read linearly, so a
+  // NON-contiguous input (e.g. a transposed donor K/V view) would be
+  // byte-scrambled. `ops.contiguous` is a no-op for already-contiguous arrays
   // (hidden boundaries) and materializes row-major for strided views.
+  // detachCopy = ONE mlx→mlx copy (backlog #9: the old
+  // fromBytesCopy(rawBytes()) route was two copies via the JS heap).
   const c = ops.contiguous(a);
-  c.eval();
-  const leaf = MlxArray.fromBytesCopy(c.rawBytes(), c.shape, c.dtype);
+  const leaf = c.detachCopy();
   c.dispose();
   a.dispose();
   return leaf;
 }
 
+/** Batch detach: ONE evalAll barrier for the whole set (backlog #9 — each
+ *  per-array eval was a separate blocking GPU drain; a segment's vjp returns
+ *  dh + 2·consumed-donors + 2·targets arrays, so the per-segment sync count
+ *  collapses from ~2+2n to 1), then one mlx→mlx copy each. Disposes inputs. */
+function detachLeaves(arrs: MlxArray[]): MlxArray[] {
+  const cs = arrs.map((a) => ops.contiguous(a));
+  ops.evalAll(cs);
+  const leaves = cs.map((c) => c.detachCopy());
+  for (const c of cs) c.dispose();
+  for (const a of arrs) a.dispose();
+  return leaves;
+}
+
 /** Layer index a target's modulePath belongs to, or -1 (e.g. embedding/head). */
 function targetLayer(modulePath: string): number {
   const m = modulePath.match(/\.layers\.(\d+)\./);
@@ -120,6 +188,10 @@ export class SegmentedBackward {
   private readonly one: MlxArray; // scalar cotangent 1.0 for the (scalar) head loss
   // Per-step context the (reused) closures read.
   private curBatch: SftBatch | null = null;
+  // Bounded-head per-step transients (chunk slices + Checkpoint/CustomVjp
+  // objects): the head backward recomputes through them, so they are disposed
+  // only AFTER the head vjp's outputs are detached (which forces the eval).
+  private headSink: Array<{ dispose(): void }> = [];
   private disposed = false;
 
   constructor(
@@ -140,11 +212,17 @@ export class SegmentedBackward {
     this.caches = Array.from({ length: this.nLayers }, () => new TrainingCache());
     this.one = MlxArray.fromFloat32(new Float32Array([1]), []); // d(loss)/d(loss) = 1
 
-    // Loss head: vjp of finalNorm + responseOnlyCe over the last boundary leaf.
+    // Loss head: vjp of finalNorm + the [M,V]-BOUNDED response CE over the last
+    // boundary leaf (kernel-review backlog #8: the old responseOnlyCe here
+    // materialized [1,M,V] logits + cotangents inside the head vjp — measured
+    // ~3 GB @8K on e4b, the exact transient the segmentation exists to kill).
+    // CE = −(mean response logp); span and reduction identical to
+    // responseOnlyCe, per-position grads bit-exact (chunking is over token
+    // positions — each position's full-V lse/softmax is computed unchanged).
     // outputs[0] = loss scalar; vjps[0] = dLoss/d(last residual stream) = dh.
     this.headVjp = new Vjp((p) => {
       const hn = model.finalNorm.forward(p[0]!);
-      const loss = responseOnlyCe(model, hn, this.curBatch!);
+      const loss = boundedSftCe(model, hn, this.curBatch!, this.headSink);
       hn.dispose();
       return [loss];
     }, 1);
@@ -204,9 +282,17 @@ export class SegmentedBackward {
       }
 
       // --- 2. Loss head over the last boundary (vjp with cotangent 1.0). ---
+      // The bounded head pushes per-step chunk slices + Checkpoints into
+      // headSink; the head BACKWARD recomputes through them, so materialize
+      // the full head-vjp graph BEFORE freeing the sink (lazy-eval UAF
+      // discipline, same as the ORPO classes).
+      this.headSink = [];
       const head = this.headVjp.apply([boundaries[nSeg]!], [this.one]);
+      ops.evalAll([head.outputs[0]!, head.vjps[0]!]);
       value = detachLeaf(head.outputs[0]!); // scalar loss, detached from the head graph (frees logits)
       dhOut = detachLeaf(head.vjps[0]!); // dLoss / d(last residual stream), graph-free
+      for (const d of this.headSink) d.dispose();
+      this.headSink = [];
       boundaries[nSeg]!.dispose();
       boundaries[nSeg] = null;
 
@@ -220,13 +306,15 @@ export class SegmentedBackward {
         res.outputs[0]!.dispose(); // segment output — unused (only its vjp matters)
         // vjps[0] = dh_in (grad w.r.t. this segment's input boundary) -> dh_out
         // for segment k-1; the rest scatter into the global flat vector. Detach
-        // each to a graph-free leaf (eval doesn't detach; keeps peak bounded).
+        // to graph-free leaves (eval doesn't detach; keeps peak bounded) — ONE
+        // barrier for the segment's whole vjp set.
+        const leaves = detachLeaves(res.vjps);
         idxs.forEach((j, s) => {
-          grads[j] = detachLeaf(res.vjps[1 + 2 * s]!); // dA
-          grads[n + j] = detachLeaf(res.vjps[1 + 2 * s + 1]!); // dB
+          grads[j] = leaves[1 + 2 * s]!; // dA
+          grads[n + j] = leaves[1 + 2 * s + 1]!; // dB
         });
         dhOut!.dispose(); // consumed as this segment's cotangent
-        dhOut = detachLeaf(res.vjps[0]!);
+        dhOut = leaves[0]!;
         boundaries[k]!.dispose();
         boundaries[k] = null;
       }
@@ -238,6 +326,8 @@ export class SegmentedBackward {
     } finally {
       inputIds.dispose();
       this.curBatch = null;
+      for (const d of this.headSink) d.dispose(); // no-op unless the head threw mid-step
+      this.headSink = [];
       dhOut?.dispose();
       for (const b of boundaries) b?.dispose();
       if (!transferred) {
@@ -295,6 +385,8 @@ export class SegmentedBackwardGemma4 {
   private curBatch: SftBatch | null = null;
   private curMasks: Map<string, Mask> | null = null;
   private curPerLayer: MlxArray | null = null;
+  // Bounded-head per-step transients — see SegmentedBackward.headSink.
+  private headSink: Array<{ dispose(): void }> = [];
   private disposed = false;
 
   constructor(
@@ -337,11 +429,11 @@ export class SegmentedBackwardGemma4 {
     this.caches = Array.from({ length: model.numDonors }, () => new TrainingCache());
     this.one = MlxArray.fromFloat32(new Float32Array([1]), []);
 
-    // Loss head (identical to MiniCPM5): finalNorm + responseOnlyCe over the last
+    // Loss head (identical to MiniCPM5): finalNorm + bounded response CE over the last
     // boundary, vjp with cotangent 1.0 -> [loss], [dh].
     this.headVjp = new Vjp((p) => {
       const hn = model.finalNorm.forward(p[0]!);
-      const loss = responseOnlyCe(model, hn, this.curBatch!);
+      const loss = boundedSftCe(model, hn, this.curBatch!, this.headSink);
       hn.dispose();
       return [loss];
     }, 1);
@@ -437,14 +529,22 @@ export class SegmentedBackwardGemma4 {
         for (const [d, kv] of donorKvOut)
           if (!produces[k]!.includes(d) && kv.kind === "plain") { kv.keys.dispose(); kv.values.dispose(); }
         boundaries.push(detachLeaf(h));
+        segMem(`fwd seg ${k} [${lo},${hi})`);
       }
 
-      // --- 2. Loss head (vjp, cotangent 1.0). ---
+      // --- 2. Loss head (vjp, cotangent 1.0). Bounded-head sink discipline:
+      // materialize the head-vjp graph BEFORE freeing the sink (see the
+      // MiniCPM5 class / ORPO classes — lazy-eval UAF). ---
+      this.headSink = [];
       const head = this.headVjp.apply([boundaries[nSeg]!], [this.one]);
+      ops.evalAll([head.outputs[0]!, head.vjps[0]!]);
       value = detachLeaf(head.outputs[0]!);
       dhOut = detachLeaf(head.vjps[0]!);
+      for (const d of this.headSink) d.dispose();
+      this.headSink = [];
       boundaries[nSeg]!.dispose();
       boundaries[nSeg] = null;
+      segMem("head vjp");
 
       // --- 3. Backward, reverse over segments (vjp, cotangents dh + dKV). ---
       for (let k = nSeg - 1; k >= 0; k--) {
@@ -466,17 +566,17 @@ export class SegmentedBackwardGemma4 {
         const res = this.segVjps[k]!.apply(primals, cots);
         for (const o of res.outputs) o.dispose();
 
-        // vjps[0] = dh_in; then (dk,dv) per consumed donor; then (da,db) per LoRA.
-        const newDh = detachLeaf(res.vjps[0]!);
+        // vjps[0] = dh_in; then (dk,dv) per consumed donor; then (da,db) per
+        // LoRA — detached in ONE barrier (backlog #9).
+        const leaves = detachLeaves(res.vjps);
+        const newDh = leaves[0]!;
         cons.forEach((d, ci) => {
-          const dk = detachLeaf(res.vjps[1 + 2 * ci]!);
-          const dv = detachLeaf(res.vjps[1 + 2 * ci + 1]!);
-          this.accumulateDKV(dKV, d, dk, dv);
+          this.accumulateDKV(dKV, d, leaves[1 + 2 * ci]!, leaves[1 + 2 * ci + 1]!);
         });
         const base = 1 + 2 * cons.length;
         idxs.forEach((j, s) => {
-          grads[j] = detachLeaf(res.vjps[base + 2 * s]!);
-          grads[n + j] = detachLeaf(res.vjps[base + 2 * s + 1]!);
+          grads[j] = leaves[base + 2 * s]!;
+          grads[n + j] = leaves[base + 2 * s + 1]!;
         });
 
         dh.dispose();
@@ -488,6 +588,7 @@ export class SegmentedBackwardGemma4 {
         }
         boundaries[k]!.dispose();
         boundaries[k] = null;
+        segMem(`bwd seg ${k} [${ranges[k]![0]},${ranges[k]![1]})`);
       }
 
       transferred = true;
@@ -498,6 +599,8 @@ export class SegmentedBackwardGemma4 {
       this.curPerLayer?.dispose();
       this.curPerLayer = null;
       if (this.curMasks) { for (const m of this.curMasks.values()) m.arr?.dispose(); this.curMasks = null; }
+      for (const d of this.headSink) d.dispose(); // no-op unless the head threw mid-step
+      this.headSink = [];
       dhOut?.dispose();
       for (const b of boundaries) b?.dispose();
       for (const kv of donorBnd.values()) { kv.keys.dispose(); kv.values.dispose(); }
@@ -725,12 +828,13 @@ export class SegmentedBackwardOrpo {
         for (const j of idxs) params.push(lora.targets[j]!.lw.a, lora.targets[j]!.lw.b);
         const res = this.segVjps[k]!.apply(params, [dhC!]);
         res.outputs[0]!.dispose();
+        const leaves = detachLeaves(res.vjps); // one barrier (backlog #9)
         idxs.forEach((j, s) => {
-          gradsC[j] = detachLeaf(res.vjps[1 + 2 * s]!); // dA
-          gradsC[n + j] = detachLeaf(res.vjps[1 + 2 * s + 1]!); // dB
+          gradsC[j] = leaves[1 + 2 * s]!; // dA
+          gradsC[n + j] = leaves[1 + 2 * s + 1]!; // dB
         });
         dhC!.dispose();
-        dhC = detachLeaf(res.vjps[0]!);
+        dhC = leaves[0]!;
         bndC[k]!.dispose(); bndC[k] = null;
       }
       dhC!.dispose(); dhC = null; // embedding grad, discarded (no LoRA on embed)
@@ -742,12 +846,13 @@ export class SegmentedBackwardOrpo {
         for (const j of idxs) params.push(lora.targets[j]!.lw.a, lora.targets[j]!.lw.b);
         const res = this.segVjps[k]!.apply(params, [dhR!]);
         res.outputs[0]!.dispose();
+        const leaves = detachLeaves(res.vjps); // one barrier (backlog #9)
         idxs.forEach((j, s) => {
-          gradsR[j] = detachLeaf(res.vjps[1 + 2 * s]!); // dA
-          gradsR[n + j] = detachLeaf(res.vjps[1 + 2 * s + 1]!); // dB
+          gradsR[j] = leaves[1 + 2 * s]!; // dA
+          gradsR[n + j] = leaves[1 + 2 * s + 1]!; // dB
         });
         dhR!.dispose();
-        dhR = detachLeaf(res.vjps[0]!);
+        dhR = leaves[0]!;
         bndR[k]!.dispose(); bndR[k] = null;
       }
       dhR!.dispose(); dhR = null; // embedding grad, discarded
@@ -912,7 +1017,7 @@ export class SegmentedBackwardOrpoPrefix {
     this.headSink = [];
     this.curChunk = this.chunkCtx ? { ...this.chunkCtx, sink: this.headSink } : undefined;
     // One prefix cache per layer (captures P/Rc/Rr; makeMask -> block-sparse mask).
-    this.curCaches = model.layers.map(() => new PrefixSharedCache(P, Rc, Rr));
+    this.curCaches = prefixSharedCaches(model.layers.length, P, Rc, Rr);
 
     const nSeg = ranges.length;
     const concat = new Int32Array(T);
@@ -957,12 +1062,13 @@ export class SegmentedBackwardOrpoPrefix {
         for (const j of idxs) params.push(lora.targets[j]!.lw.a, lora.targets[j]!.lw.b);
         const res = this.segVjps[k]!.apply(params, [dhOut!]);
         res.outputs[0]!.dispose();
+        const leaves = detachLeaves(res.vjps); // one barrier (backlog #9)
         idxs.forEach((j, s) => {
-          grads[j] = detachLeaf(res.vjps[1 + 2 * s]!); // dA
-          grads[n + j] = detachLeaf(res.vjps[1 + 2 * s + 1]!); // dB
+          grads[j] = leaves[1 + 2 * s]!; // dA
+          grads[n + j] = leaves[1 + 2 * s + 1]!; // dB
         });
         dhOut!.dispose();
-        dhOut = detachLeaf(res.vjps[0]!);
+        dhOut = leaves[0]!;
         boundaries[k]!.dispose(); boundaries[k] = null;
       }
       // Segment-0 dh_in = grad w.r.t. the embedding output (no LoRA on embed) — drop.
@@ -1232,16 +1338,15 @@ export class SegmentedBackwardOrpoGemma4 {
         const res = this.segVjps[k]!.apply(primals, cots);
         for (const o of res.outputs) o.dispose();
 
-        const newDh = detachLeaf(res.vjps[0]!);
+        const leaves = detachLeaves(res.vjps); // one barrier (backlog #9)
+        const newDh = leaves[0]!;
         cons.forEach((d, ci) => {
-          const dk = detachLeaf(res.vjps[1 + 2 * ci]!);
-          const dv = detachLeaf(res.vjps[1 + 2 * ci + 1]!);
-          this.accumulateDKV(dKV, d, dk, dv);
+          this.accumulateDKV(dKV, d, leaves[1 + 2 * ci]!, leaves[1 + 2 * ci + 1]!);
         });
         const base = 1 + 2 * cons.length;
         idxs.forEach((j, s) => {
-          grads[j] = detachLeaf(res.vjps[base + 2 * s]!);
-          grads[n + j] = detachLeaf(res.vjps[base + 2 * s + 1]!);
+          grads[j] = leaves[base + 2 * s]!;
+          grads[n + j] = leaves[base + 2 * s + 1]!;
         });
 
         dh.dispose();
@@ -1663,16 +1768,15 @@ export class SegmentedBackwardOrpoPrefixGemma4 {
         const res = this.segVjps[k]!.apply(primals, cots);
         for (const o of res.outputs) o.dispose();
 
-        const newDh = detachLeaf(res.vjps[0]!);
+        const leaves = detachLeaves(res.vjps); // one barrier (backlog #9)
+        const newDh = leaves[0]!;
         cons.forEach((d, ci) => {
-          const dk = detachLeaf(res.vjps[1 + 2 * ci]!);
-          const dv = detachLeaf(res.vjps[1 + 2 * ci + 1]!);
-          this.accumulateDKV(dKV, d, dk, dv);
+          this.accumulateDKV(dKV, d, leaves[1 + 2 * ci]!, leaves[1 + 2 * ci + 1]!);
         });
         const base = 1 + 2 * cons.length;
         idxs.forEach((j, s) => {
-          grads[j] = detachLeaf(res.vjps[base + 2 * s]!);
-          grads[n + j] = detachLeaf(res.vjps[base + 2 * s + 1]!);
+          grads[j] = leaves[base + 2 * s]!;
+          grads[n + j] = leaves[base + 2 * s + 1]!;
         });
 
         dh.dispose();