mlx-bun 0.0.9 → 0.0.10

package/src/ssd-cache.ts

@@ -0,0 +1,213 @@
+// SSD cold tier for the prompt/KV cache (docs/design/ssd-kv-cold-tier.md).
+//
+// Files ARE the database: no sidecar index. Layout
+//   <dir>/<configFingerprint>/<nsHash>/<uuid>.mlxkv
+// so two models (or two kv-quant schemes of one model) share a directory
+// without ever colliding, and an incompatible fingerprint dir is simply
+// ignored (never deleted — it may belong to the user's other model).
+// Startup recovery = a header-only scan of OUR fingerprint dir: header hash
+// verified (cheap), corrupt or metadata-mismatched files unlinked, `.tmp`
+// write orphans reaped. LRU state = file mtime (bumped on hit); the byte
+// cap is enforced at write time by oldest-mtime eviction. A cache is always
+// droppable: every failure path degrades to "no hit" or "not stored" — the
+// tier must never take serving down (oMLX paged_ssd_cache lesson, Apache-2.0
+// idea port; ours spills whole prefix entries, not content-hashed blocks —
+// see the design doc's D1 for why).
+
+import {
+  existsSync, mkdirSync, readdirSync, rmSync, statSync, utimesSync,
+} from "node:fs";
+import { join } from "node:path";
+import { randomUUID } from "node:crypto";
+import { commonPrefixLength } from "./prompt-cache";
+import {
+  saveKvCache, loadKvCache, readKvHeader,
+  type KvSaveMeta, type KvLoadExpect, type LoadedKvCache,
+} from "./kv-store";
+import type { Cache } from "./model/gemma4-base";
+
+export interface SsdIndexEntry {
+  path: string;
+  ns: string;
+  tokens: number[];
+  bytes: number;
+  mtimeMs: number;
+}
+
+export interface SsdStoreOptions {
+  dir: string;
+  maxBytes: number;
+  /** Server identity — written into every file, enforced on scan/load. */
+  configFingerprint: string;
+  tokenizerHash: string;
+  modelId: string;
+  /** Verify every tensor hash on load (--ssd-cache-verify). */
+  verify?: boolean;
+}
+
+const nsHash = (ns: string): string =>
+  ns === "" ? "base" : Bun.hash(ns).toString(16);
+
+export class SsdCacheStore {
+  readonly #opts: SsdStoreOptions;
+  readonly #root: string; // <dir>/<configFingerprint>
+  #index: SsdIndexEntry[] = [];
+  #warnedWriteFailure = false;
+  stats = { restores: 0, spills: 0, restoreMsLast: 0 };
+
+  constructor(opts: SsdStoreOptions) {
+    this.#opts = opts;
+    this.#root = join(opts.dir, opts.configFingerprint);
+  }
+
+  get entries(): number {
+    return this.#index.length;
+  }
+
+  get totalBytes(): number {
+    return this.#index.reduce((a, e) => a + e.bytes, 0);
+  }
+
+  get maxBytes(): number {
+    return this.#opts.maxBytes;
+  }
+
+  /** Startup recovery: rebuild the in-memory index from disk. Corrupt
+   *  headers and metadata mismatches are unlinked; `.tmp` orphans reaped.
+   *  Only OUR fingerprint dir is touched. Returns entries indexed. */
+  scan(): number {
+    this.#index = [];
+    if (!existsSync(this.#root)) return 0;
+    for (const nsDir of readdirSync(this.#root)) {
+      const nsPath = join(this.#root, nsDir);
+      let files: string[];
+      try { files = readdirSync(nsPath); } catch { continue; }
+      for (const f of files) {
+        const path = join(nsPath, f);
+        if (f.endsWith(".tmp")) {
+          // A crash mid-write: the rename never happened. Ours by
+          // construction (uuid names), safe to reap.
+          try { rmSync(path, { force: true }); } catch {}
+          continue;
+        }
+        if (!f.endsWith(".mlxkv")) continue;
+        try {
+          const h = readKvHeader(path); // verifies the header hash
+          if (
+            h.configFingerprint !== this.#opts.configFingerprint ||
+            h.tokenizerHash !== this.#opts.tokenizerHash
+          )
+            throw new Error("metadata mismatch");
+          const st = statSync(path);
+          this.#index.push({
+            path, ns: h.ns ?? "", tokens: h.tokens, bytes: st.size, mtimeMs: st.mtimeMs,
+          });
+        } catch {
+          try { rmSync(path, { force: true }); } catch {}
+        }
+      }
+    }
+    return this.#index.length;
+  }
+
+  /** Best stored entry for `prompt` in `ns`: the longest common prefix
+   *  (capped at prompt.length−1 — at least one token must forward). An
+   *  entry LONGER than the matched prefix is still returned (the divergent-
+   *  tail pattern: same long context, different question) — the caller
+   *  trims after restore when every cache kind is trimmable, and falls
+   *  back to a fresh prefill when not (ring post-wrap, SSM). Pure index
+   *  lookup; no I/O. */
+  find(prompt: number[], ns = ""): { entry: SsdIndexEntry; prefixLen: number } | null {
+    let best: SsdIndexEntry | null = null;
+    let bestLen = 0;
+    for (const e of this.#index) {
+      if (e.ns !== ns) continue;
+      const p = Math.min(commonPrefixLength(e.tokens, prompt), prompt.length - 1);
+      if (p === 0 || p <= bestLen) continue;
+      bestLen = p;
+      best = e;
+    }
+    return best ? { entry: best, prefixLen: bestLen } : null;
+  }
+
+  /** Restore an entry zero-copy (COW mmap; pages fault in lazily). Bumps
+   *  LRU mtime. On ANY failure the file is dropped and null returned. */
+  restore(entry: SsdIndexEntry, model: { makeCache(): Cache[] }): LoadedKvCache | null {
+    const t0 = performance.now();
+    try {
+      const expect: KvLoadExpect = {
+        configFingerprint: this.#opts.configFingerprint,
+        tokenizerHash: this.#opts.tokenizerHash,
+        ns: entry.ns,
+        verify: this.#opts.verify,
+      };
+      const loaded = loadKvCache(entry.path, model, expect);
+      const now = new Date();
+      try { utimesSync(entry.path, now, now); entry.mtimeMs = now.getTime(); } catch {}
+      this.stats.restores++;
+      this.stats.restoreMsLast = performance.now() - t0;
+      return loaded;
+    } catch {
+      this.remove(entry.path);
+      return null;
+    }
+  }
+
+  /** Persist an entry. Synchronous (the tier calls it on the idle serial
+   *  lane between requests); atomic via kv-store's tmp+fsync+rename. An
+   *  entry bigger than the cap is refused; disk/write failure is a warn-once
+   *  soft-fail — serving never stalls on the cold tier. Existing entries in
+   *  the same ns whose tokens are a PREFIX of the new entry are superseded
+   *  (the agent-conversation pattern: one growing entry, not N generations
+   *  of it). Returns true if stored. */
+  store(tokens: number[], caches: Cache[], ns = ""): boolean {
+    const dir = join(this.#root, nsHash(ns));
+    const path = join(dir, `${randomUUID()}.mlxkv`);
+    try {
+      mkdirSync(dir, { recursive: true });
+      const meta: KvSaveMeta = {
+        modelId: this.#opts.modelId,
+        configFingerprint: this.#opts.configFingerprint,
+        tokenizerHash: this.#opts.tokenizerHash,
+        ns,
+      };
+      saveKvCache(path, tokens, caches, meta);
+      const st = statSync(path);
+      if (st.size > this.#opts.maxBytes) {
+        rmSync(path, { force: true });
+        return false;
+      }
+      // Supersede shadowed ancestors of this prefix (same ns).
+      for (const e of [...this.#index]) {
+        if (e.ns !== ns || e.tokens.length >= tokens.length) continue;
+        if (commonPrefixLength(e.tokens, tokens) === e.tokens.length) this.remove(e.path);
+      }
+      this.#index.push({ path, ns, tokens, bytes: st.size, mtimeMs: Date.now() });
+      this.stats.spills++;
+      this.evictToCap();
+      return true;
+    } catch (err) {
+      try { rmSync(path, { force: true }); rmSync(`${path}.tmp`, { force: true }); } catch {}
+      if (!this.#warnedWriteFailure) {
+        this.#warnedWriteFailure = true;
+        console.warn(`[ssd-cache] store failed (disk full or unwritable?) — cold tier disabled for this entry: ${(err as Error).message}`);
+      }
+      return false;
+    }
+  }
+
+  /** Enforce the byte cap: unlink oldest-mtime entries until under it. */
+  evictToCap(): void {
+    while (this.totalBytes > this.#opts.maxBytes && this.#index.length > 0) {
+      let oldest = 0;
+      for (let i = 1; i < this.#index.length; i++)
+        if (this.#index[i]!.mtimeMs < this.#index[oldest]!.mtimeMs) oldest = i;
+      this.remove(this.#index[oldest]!.path);
+    }
+  }
+
+  remove(path: string): void {
+    try { rmSync(path, { force: true }); } catch {}
+    this.#index = this.#index.filter((e) => e.path !== path);
+  }
+}

package/src/train/flash-cce.ts

@@ -723,24 +723,28 @@ const BWD_SG_SOURCE = String.raw`
   }
 `;
 
-// Apple-CCE coeff filter epsilon (phase-2 per-tile skip). DEFAULT OFF ("0") today —
-// no production path sets the env — but the VERDICT is measured and positive: the old
+// Apple-CCE coeff filter epsilon (phase-2 per-tile skip). DEFAULT ON at 1e-5
+// (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, commit 3d259ef, 2026-06-22)
-// measured eps 1e-5 = 0.16% dh (under the bf16 floor) for a 1.35× faster backward,
-// and PLAN.md's recorded decision is "enable the filter at eps ~1e-5". (The ~3.5×
-// figure floating around older notes was the retired SG kernel, not steel.) Flipping
-// the default is the open backlog item; until then opt in with
-// MLX_BUN_CCE_BWD_FILTER_EPS=1e-5 on real (peaked) training data.
-const BWD_FILTER_EPS = process.env.MLX_BUN_CCE_BWD_FILTER_EPS ?? "0";
+// (scripts/experiments/flash-cce-filter-realdata.ts) 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
+// exists (CPM5 0.850→0.913% relnorm, e4b 1.207→1.220%, cosine ≥0.99993 both).
+// (The ~3.5× figure floating around older notes was the retired SG kernel, not
+// steel.) MLX_BUN_CCE_BWD_FILTER_EPS=0 restores exact gradients — eps "0"
+// compiles the filter OUT (#if CCE_BWD_FILTER 0), the identical pre-flip kernel.
+const BWD_FILTER_EPS = process.env.MLX_BUN_CCE_BWD_FILTER_EPS ?? "1e-5";
 // blockMax vocab-block early-exit eps (skip whole cold (token-block,vocab-block) programs
-// — attacks BOTH phases). DEFAULT OFF ("0"). It is LOSSLESS when it fires (only skips
-// blocks whose max prob < eps AND which lack the target → ~0 contribution), so it is the
-// SAFE skip; but it only pays off on genuinely peaked real text where whole 8192-vocab
-// blocks go cold (the M=512 synthetic probe found nothing cold → pure overhead). Opt in
-// with MLX_BUN_CCE_BWD_BLOCK_EPS=1e-5 for long real-text training; needs blockMax supplied.
-const BWD_BLOCK_EPS = process.env.MLX_BUN_CCE_BWD_BLOCK_EPS ?? "0";
+// — attacks BOTH phases). DEFAULT ON at 1e-5 (flipped 2026-07-02). It is ~LOSSLESS when
+// it fires (only skips blocks whose max prob < eps AND which lack the target → ~0
+// contribution). The M=512 synthetic probe found nothing cold (pure overhead), but the
+// real-data gate proved real text DOES go cold: alone it's 1.23× (CPM5) / 2.02× (e4b)
+// at ≤0.004% dh; COMBINED with the coeff filter at 1e-5/1e-5 the backward is 1.71×
+// (CPM5) / 3.16× (e4b) vs exact. Fires only when the caller supplies blockMax (the
+// production head does); MLX_BUN_CCE_BWD_BLOCK_EPS=0 disables (compiled out).
+const BWD_BLOCK_EPS = process.env.MLX_BUN_CCE_BWD_BLOCK_EPS ?? "1e-5";
 const _bwdSgKernels = new Map<string, MetalKernel>();
 function bwdSgKernel(dimTiles: number, filterEps: string, blockEps: string): MetalKernel {
   const filterOn = Number.parseFloat(filterEps) > 0 ? 1 : 0;
@@ -1132,14 +1136,13 @@ export function flashCceForward(
  *  (no atomics). Caller owns the result. */
 export function flashCceBackward(
   hResp: MlxArray, head: FlashCceHead, targets: number[], lse: MlxArray, cot: number[],
-  // Apple-CCE coeff filter eps. Defaults to BWD_FILTER_EPS, which is OFF ("0")
-  // unless MLX_BUN_CCE_BWD_FILTER_EPS is set — exact gradients by default. The
-  // measured real-data case FOR enabling it (eps 1e-5 = 0.16% dh, 1.35×; PLAN.md
-  // "filter-on-real-data") is a pending default flip — see BWD_FILTER_EPS's note.
-  // Pass "0" explicitly for guaranteed-exact (the parity gate does).
+  // Apple-CCE coeff filter eps. Defaults to BWD_FILTER_EPS — ON at 1e-5 since
+  // 2026-07-02 (real-data gates in BWD_FILTER_EPS's note). Pass "0" explicitly
+  // for guaranteed-exact gradients (the parity gate does).
   filterEps: string = BWD_FILTER_EPS,
   // The forward's `blockMax [M, NBLK]` enables the vocab-block early exit (skips
-  // whole cold programs — attacks phase 1). Omit (or set blockEps "0") for exact.
+  // whole cold programs — attacks BOTH phases); ON at 1e-5 by default when
+  // blockMax is supplied. Omit blockMax (or set blockEps "0") for exact.
   blockMax?: MlxArray, blockEps: string = BWD_BLOCK_EPS,
 ): MlxArray {
   const M = hResp.shape[0]!;
@@ -1169,9 +1172,9 @@ export function flashCceBackward(
   // MLX_BUN_CCE_BWD_NOSTEEL=1 falls back to the old kernels. The steel backward carries
   // the SAME Apple-CCE skips as the SG kernel — coeff filter (phase-2 per-tile) +
   // blockMax vocab-block early exit (whole-program) — both compiled out at eps=0.
-  // BOTH skips are OFF by default (the envs default "0", nothing in production sets
-  // them): exact gradients unless explicitly opted in. Enabling the coeff filter at
-  // ~1e-5 is the measured-and-decided pending flip (PLAN.md "filter-on-real-data").
+  // BOTH skips default ON at 1e-5 (2026-07-02 flip; measured gates in the
+  // BWD_FILTER_EPS / BWD_BLOCK_EPS notes). MLX_BUN_CCE_BWD_{FILTER,BLOCK}_EPS=0
+  // restores exact gradients (identical pre-flip kernels).
   const blockV0 = Math.ceil(V / NBLK);
   // HTw == group_size for the phase-2 QuantizedBlockLoader (one group per H-tile); needs
   // H % GS == 0 (always, GR*GS=H) and GS % 32 == 0 (TNv = GS/32 frags per H-tile integral).

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); }