mlx-bun 0.0.8 → 0.0.10

package/src/train/segmented.ts

@@ -36,14 +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, orpoLossFromLogps, type ChunkCtx } from "./loss";
-import { PrefixSharedCache, prefixGatherIdx, branchLogpMeanGathered, blockSparsePrefixMaskGemma } from "./prefix-shared";
+import {
+  responseOnlyLogpMean, fusedLogpMeanFromHidden, chunkedLogpMeanFromHidden,
+  orpoLossFromLogps, spanLogpMeanFromHidden, respSpanFromMask, combineFullNll,
+  type ChunkCtx, type SftScope, type HeadSpan,
+} from "./loss";
+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);
@@ -54,14 +71,43 @@ import { PrefixSharedCache, prefixGatherIdx, branchLogpMeanGathered, blockSparse
 const SEG_HEAD = process.env.MLX_BUN_SEG_HEAD ?? "checkpoint";
 function boundedHeadFromHidden(
   model: MiniCPM5Model | Gemma4Model, h: MlxArray, ids: number[], mask: number[],
-  chunkSize: number, sink: Array<{ dispose(): void }>,
+  chunkSize: number, sink: Array<{ dispose(): void }>, span?: HeadSpan,
 ): MlxArray {
   return SEG_HEAD === "fused"
-    ? fusedLogpMeanFromHidden(model, h, ids, mask, chunkSize, sink)
-    : chunkedLogpMeanFromHidden(model, h, ids, mask, chunkSize, sink);
+    ? fusedLogpMeanFromHidden(model, h, ids, mask, chunkSize, sink, 0, false, span)
+    : 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:
@@ -73,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+)\./);
@@ -116,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(
@@ -136,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);
@@ -200,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;
 
@@ -216,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;
       }
@@ -234,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) {
@@ -291,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(
@@ -333,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);
@@ -433,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--) {
@@ -462,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();
@@ -484,6 +588,7 @@ export class SegmentedBackwardGemma4 {
         }
         boundaries[k]!.dispose();
         boundaries[k] = null;
+        segMem(`bwd seg ${k} [${ranges[k]![0]},${ranges[k]![1]})`);
       }
 
       transferred = true;
@@ -494,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(); }
@@ -582,6 +689,7 @@ export class SegmentedBackwardOrpo {
     private readonly ranges: [number, number][],
     private readonly lambda: number,
     private readonly fusedChunkSize = 0,
+    private readonly sftScope: SftScope = "response",
   ) {
     this.nLayers = model.layers.length;
     const lastHi = ranges[ranges.length - 1]?.[1] ?? 0;
@@ -609,9 +717,26 @@ export class SegmentedBackwardOrpo {
       const lr = cs > 0 // [1]
         ? boundedHeadFromHidden(model, hnR, this.curRejectedIds!, this.curRejectedMask!, cs, this.headSink)
         : responseOnlyLogpMean(model, hnR, this.curRejectedIds!, this.curRejectedMask!);
+      // sft_scope:"full": the chosen head additionally runs over the PROMPT span
+      // [0, startT) of the SAME hnC (same tier: bounded or whole-vocab), and the
+      // full-scope chosen NLL replaces mean(-ℓw). The prompt-position gradient
+      // rides into this vjp's dh_chosen automatically (the loss graph includes
+      // the prompt-span head). ℓw/ℓr stay response-only for the odds ratio.
+      let nllFull: MlxArray | null = null;
+      if (this.sftScope === "full") {
+        const { startT, M } = respSpanFromMask(this.curChosenMask!, this.curChosenIds!.length - 1);
+        let pm: MlxArray | null = null;
+        if (M > 0 && startT > 0)
+          pm = cs > 0
+            ? boundedHeadFromHidden(model, hnC, this.curChosenIds!, this.curChosenMask!, cs, this.headSink, { startT: 0, M: startT })
+            : spanLogpMeanFromHidden(model, hnC, this.curChosenIds!, 0, startT);
+        nllFull = combineFullNll(pm, lw, M > 0 ? startT : 0, M);
+        pm?.dispose();
+      }
       hnC.dispose(); hnR.dispose();
-      const loss = orpoLossFromLogps(lw, lr, this.lambda);
+      const loss = orpoLossFromLogps(lw, lr, this.lambda, nllFull ?? undefined);
       lw.dispose(); lr.dispose();
+      nllFull?.dispose();
       return [loss];
     }, 1);
 
@@ -703,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)
@@ -720,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
@@ -801,6 +928,7 @@ export class SegmentedBackwardOrpoPrefix {
   private curP = 0;
   private curRc = 0;
   private curRr = 0;
+  private curPromptIds: number[] | null = null;
   private curChosenResp: number[] | null = null;
   private curRejectedResp: number[] | null = null;
   private curChunk: ChunkCtx | undefined = undefined;
@@ -816,6 +944,7 @@ export class SegmentedBackwardOrpoPrefix {
     private readonly ranges: [number, number][],
     private readonly lambda: number,
     private readonly chunkCtx?: ChunkCtx,
+    private readonly sftScope: SftScope = "response",
   ) {
     this.nLayers = model.layers.length;
     const lastHi = ranges[ranges.length - 1]?.[1] ?? 0;
@@ -838,9 +967,16 @@ export class SegmentedBackwardOrpoPrefix {
       const { chosenIdx, rejectedIdx } = prefixGatherIdx(this.curP, this.curRc, this.curRr);
       const lw = branchLogpMeanGathered(model, hn, chosenIdx, this.curChosenResp!, this.curChunk);
       const lr = branchLogpMeanGathered(model, hn, rejectedIdx, this.curRejectedResp!, this.curChunk);
+      // sft_scope:"full": prompt predictions (H[0..P-2] → prompt[1..P-1]) from the
+      // same concat hiddens; the gather's takeAxis vjp scatter-adds the prompt
+      // cotangent into dh alongside the responses'. ℓw/ℓr stay response-only.
+      const nllFull = this.sftScope === "full"
+        ? prefixFullNll(model, hn, this.curPromptIds!, lw, this.curRc, this.curChunk)
+        : null;
       hn.dispose();
-      const loss = orpoLossFromLogps(lw, lr, this.lambda);
+      const loss = orpoLossFromLogps(lw, lr, this.lambda, nllFull ?? undefined);
       lw.dispose(); lr.dispose();
+      nllFull?.dispose();
       return [loss];
     }, 1);
 
@@ -873,6 +1009,7 @@ export class SegmentedBackwardOrpoPrefix {
     if (P < 1 || Rc < 1 || Rr < 1) throw new Error("SegmentedBackwardOrpoPrefix: need P,Rc,Rr >= 1");
     const T = P + Rc + Rr;
     this.curP = P; this.curRc = Rc; this.curRr = Rr;
+    this.curPromptIds = promptIds;
     this.curChosenResp = chosenResp; this.curRejectedResp = rejectedResp;
     // Per-step chunk ctx pointing the head's CustomVjp/slice disposables at this
     // step's headSink (the flash/fused head recompute runs during headVjp.apply;
@@ -880,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);
@@ -925,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.
@@ -942,7 +1080,7 @@ export class SegmentedBackwardOrpoPrefix {
       inputIds.dispose();
       for (const c of this.curCaches) c.dispose();
       this.curCaches = [];
-      this.curChosenResp = this.curRejectedResp = null;
+      this.curPromptIds = this.curChosenResp = this.curRejectedResp = null;
       for (const d of this.headSink) d.dispose();
       this.headSink = [];
       dhOut?.dispose();
@@ -1024,6 +1162,7 @@ export class SegmentedBackwardOrpoGemma4 {
     private readonly ranges: [number, number][],
     private readonly lambda: number,
     private readonly fusedChunkSize = 0,
+    private readonly sftScope: SftScope = "response",
   ) {
     this.nLayers = model.layers.length;
     const lastHi = ranges[ranges.length - 1]?.[1] ?? 0;
@@ -1072,9 +1211,23 @@ export class SegmentedBackwardOrpoGemma4 {
       const lr = cs > 0
         ? boundedHeadFromHidden(this.model, hnR, this.curRejectedIds!, this.curRejectedMask!, cs, this.headSink)
         : responseOnlyLogpMean(this.model, hnR, this.curRejectedIds!, this.curRejectedMask!);
+      // sft_scope:"full": prompt-span head over the same hnC (see
+      // SegmentedBackwardOrpo); ℓw/ℓr stay response-only for the odds ratio.
+      let nllFull: MlxArray | null = null;
+      if (this.sftScope === "full") {
+        const { startT, M } = respSpanFromMask(this.curChosenMask!, this.curChosenIds!.length - 1);
+        let pm: MlxArray | null = null;
+        if (M > 0 && startT > 0)
+          pm = cs > 0
+            ? boundedHeadFromHidden(this.model, hnC, this.curChosenIds!, this.curChosenMask!, cs, this.headSink, { startT: 0, M: startT })
+            : spanLogpMeanFromHidden(this.model, hnC, this.curChosenIds!, 0, startT);
+        nllFull = combineFullNll(pm, lw, M > 0 ? startT : 0, M);
+        pm?.dispose();
+      }
       hnC.dispose(); hnR.dispose();
-      const loss = orpoLossFromLogps(lw, lr, this.lambda);
+      const loss = orpoLossFromLogps(lw, lr, this.lambda, nllFull ?? undefined);
       lw.dispose(); lr.dispose();
+      nllFull?.dispose();
       return [loss];
     }, 1);
 
@@ -1185,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();
@@ -1394,6 +1546,7 @@ export class SegmentedBackwardOrpoPrefixGemma4 {
   private curP = 0;
   private curRc = 0;
   private curRr = 0;
+  private curPromptIds: number[] | null = null;
   private curChosenResp: number[] | null = null;
   private curRejectedResp: number[] | null = null;
   private curChunk: ChunkCtx | undefined = undefined;
@@ -1408,6 +1561,7 @@ export class SegmentedBackwardOrpoPrefixGemma4 {
     private readonly ranges: [number, number][],
     private readonly lambda: number,
     private readonly chunkCtx?: ChunkCtx,
+    private readonly sftScope: SftScope = "response",
   ) {
     this.nLayers = model.layers.length;
     const lastHi = ranges[ranges.length - 1]?.[1] ?? 0;
@@ -1455,9 +1609,15 @@ export class SegmentedBackwardOrpoPrefixGemma4 {
       const { chosenIdx, rejectedIdx } = prefixGatherIdx(this.curP, this.curRc, this.curRr);
       const lw = branchLogpMeanGathered(model, hn, chosenIdx, this.curChosenResp!, this.curChunk);
       const lr = branchLogpMeanGathered(model, hn, rejectedIdx, this.curRejectedResp!, this.curChunk);
+      // sft_scope:"full": prompt predictions from the same concat hiddens (see
+      // SegmentedBackwardOrpoPrefix); ℓw/ℓr stay response-only.
+      const nllFull = this.sftScope === "full"
+        ? prefixFullNll(model, hn, this.curPromptIds!, lw, this.curRc, this.curChunk)
+        : null;
       hn.dispose();
-      const loss = orpoLossFromLogps(lw, lr, this.lambda);
+      const loss = orpoLossFromLogps(lw, lr, this.lambda, nllFull ?? undefined);
       lw.dispose(); lr.dispose();
+      nllFull?.dispose();
       return [loss];
     }, 1);
 
@@ -1525,6 +1685,7 @@ export class SegmentedBackwardOrpoPrefixGemma4 {
     if (P < 1 || Rc < 1 || Rr < 1) throw new Error("SegmentedBackwardOrpoPrefixGemma4: need P,Rc,Rr >= 1");
     const T = P + Rc + Rr;
     this.curP = P; this.curRc = Rc; this.curRr = Rr;
+    this.curPromptIds = promptIds;
     this.curChosenResp = chosenResp; this.curRejectedResp = rejectedResp;
     this.headSink = [];
     this.curChunk = this.chunkCtx ? { ...this.chunkCtx, sink: this.headSink } : undefined;
@@ -1607,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();
@@ -1640,7 +1800,7 @@ export class SegmentedBackwardOrpoPrefixGemma4 {
       this.curPerLayer?.dispose();
       this.curPerLayer = null;
       if (this.curMasks) { for (const m of this.curMasks.values()) m.arr?.dispose(); this.curMasks = null; }
-      this.curChosenResp = this.curRejectedResp = null;
+      this.curPromptIds = this.curChosenResp = this.curRejectedResp = null;
       for (const d of this.headSink) d.dispose();
       this.headSink = [];
       dhOut?.dispose();