mlx-bun 0.0.8 → 0.0.10

package/src/spec/generate.ts

@@ -1,11 +1,17 @@
 // Speculative decoding loop — port of optiq runtime/spec/runtime.py
 // (γ-draft / verify-in-one-pass / accept-longest-prefix / rollback).
 //
-// Correctness contract: greedy spec output is token-for-token identical
-// to greedy non-spec decode — the verify step re-derives every emitted
-// token with the SAME pick function as the stock path (logprobs argmax
-// over [1, V]); drafts only ever skip target forwards, never change
-// what the target would emit.
+// Correctness oracle: OPTIQ's spec_generate, NOT stock decode. Assistant-
+// drafter speculation is an optiq feature (mlx-lm has only generic
+// two-model spec, and can't drive a KV-borrowing drafter). Both optiq AND
+// mlx-lm BATCH the verify lm-head (one matmul over the whole γ+1 window,
+// then argmax per position) — so neither is bit-exact to stock token-at-
+// a-time decode; they diverge from it at bf16 knife-edges. We do the same
+// (picksBatched): greedy spec is bit-exact to optiq's spec, and agrees
+// with stock only on tie-free prompts (where batched == per-position).
+// An earlier version verified per-position to stay bit-exact to STOCK —
+// that matched no real oracle and cost an extra γ× read of the lm-head
+// weight per step; superseded. See docs/design/spec-decode-larger-targets.md.
 //
 // Limitation (same as the reference): partial-accept rollback requires
 // trimmable caches; RotatingKVCache loses trimability once its ring
@@ -14,7 +20,6 @@
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
 import { Gemma4Model, type Cache } from "../model/gemma4";
-import { toLogprobs } from "../sampler";
 import type { GemmaAssistantDrafter } from "./drafter";
 
 export interface SpecOptions {
@@ -39,12 +44,8 @@ export interface SpecResult {
   stats: SpecStats;
 }
 
-/** The stock path's greedy pick, bit-identical shapes: slice ONE hidden
- *  position [1,1,H], lm-head on it (exactly what generate() does for a
- *  decode step), then logprobs argmax over [1, V]. Computing the lm-head
- *  per position (instead of batched over the verify window) keeps every
- *  pick in the same kernel shapes as stock decode — batched lm-head
- *  rounds differently and flips bf16 knife-edges. */
+/** Single greedy pick from one hidden position [1,1,H] (lm-head + argmax).
+ *  Used for the NON-speculative prefill token (matches stock there). */
 function pickFromHidden(model: Gemma4Model, hidden: MlxArray, pos: number): number {
   const H = hidden.shape[2]!;
   const hSl = hidden.slice([0, pos, 0], [1, pos + 1, H]);
@@ -53,15 +54,36 @@ function pickFromHidden(model: Gemma4Model, hidden: MlxArray, pos: number): numb
   const V = logits.shape[2]!;
   const flat = ops.reshape(logits, [1, V]);
   logits.dispose();
-  const lp = toLogprobs(flat);
+  const am = ops.argmaxAxis(flat, -1);
   flat.dispose();
-  const am = ops.argmaxAxis(lp, -1);
-  lp.dispose();
   const t = ops.itemUint32(am);
   am.dispose();
   return t;
 }
 
+/** optiq-faithful batched verify: ONE lm-head over the whole γ+1 window
+ *  (reads the lm-head weight ONCE, as optiq and mlx-lm both do), then
+ *  argmax per position. This is what makes greedy spec bit-exact to
+ *  optiq's spec_generate. argmax is invariant to the logsumexp shift and
+ *  to the monotone final-logit softcap, so raw-logit argmax here picks the
+ *  same token optiq does. */
+function picksBatched(model: Gemma4Model, hidden: MlxArray, count: number): number[] {
+  const logits = model.logitsFromHidden(hidden); // [1, L, V] — one matmul
+  const V = logits.shape[2]!;
+  const out: number[] = [];
+  for (let k = 0; k < count; k++) {
+    const sl = logits.slice([0, k, 0], [1, k + 1, V]);
+    const flat = ops.reshape(sl, [1, V]);
+    sl.dispose();
+    const am = ops.argmaxAxis(flat, -1);
+    flat.dispose();
+    out.push(ops.itemUint32(am));
+    am.dispose();
+  }
+  logits.dispose();
+  return out;
+}
+
 /** Donor caches: the target's LAST sliding and LAST full cache owners
  *  (port of optiq kv_view.find_donor_layers). */
 function donorIndices(model: Gemma4Model): { sliding: number; full: number } {
@@ -163,10 +185,9 @@ export function specGenerate(
       vIds.dispose();
       stats.targetCalls++;
 
-      // 2c. accept the longest matching prefix (per-position lm-head —
-      // see pickFromHidden)
-      const gt: number[] = [];
-      for (let k = 0; k <= gamma; k++) gt.push(pickFromHidden(model, vHidden, k));
+      // 2c. accept the longest matching prefix. Batched verify lm-head
+      // (optiq-faithful — see picksBatched).
+      const gt = picksBatched(model, vHidden, gamma + 1);
       let kAccept = 0;
       while (kAccept < gamma && drafts[kAccept] === gt[kAccept]) kAccept++;
       stats.accepted += kAccept;

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/status-page.html

@@ -285,7 +285,7 @@
           <thead><tr><th>Command</th><th style="text-align:left">What it does</th></tr></thead>
           <tbody>
             <tr><td><code>mlx-bun pi</code></td><td style="text-align:left">Agent session on a local model — picks/downloads the right model, serves, drops into pi. Flags: <code>--query</code>, <code>--port</code>; everything else passes to pi (<code>-p</code>, <code>--continue</code>, <code>@file</code>…)</td></tr>
-            <tr><td><code>mlx-bun serve &lt;query&gt;</code></td><td style="text-align:left">Persistent server + this page. <code>--port 8090</code>, <code>--memory-budget GB</code></td></tr>
+            <tr><td><code>mlx-bun serve &lt;query&gt;</code></td><td style="text-align:left">Persistent server + this page. <code>--port 8080</code>, <code>--memory-budget GB</code></td></tr>
             <tr><td><code>mlx-bun get &lt;org/repo&gt;</code></td><td style="text-align:left">Resumable, sha256-verified model download</td></tr>
             <tr><td><code>mlx-bun ls</code> / <code>fit</code></td><td style="text-align:left">List downloaded models / will-it-fit assessment (<code>--ctx</code>, <code>--skus</code>)</td></tr>
             <tr><td><code>mlx-bun harness pi</code></td><td style="text-align:left">Register this server as a provider in your own pi install (<code>--remove</code> undoes)</td></tr>

package/src/tokenizer.ts

@@ -11,6 +11,10 @@ export interface LoadedTokenizer {
    *  you get a corrupting double-BOS `<s><s>…`). */
   encode(text: string, addSpecialTokens?: boolean): number[];
   decode(ids: number[], skipSpecialTokens?: boolean): string;
+  /** Token id → raw vocab token string (mlx-lm's convert_ids_to_tokens —
+   *  the undecoded piece, e.g. "▁Hello", used in logprobs responses).
+   *  Falls back to decode() for ids outside the base vocab (added tokens). */
+  idToToken(id: number): string;
   readonly bosTokenId: number | null;
   readonly eosTokenId: number | null;
 }
@@ -36,6 +40,9 @@ export async function loadTokenizer(modelDir: string): Promise<LoadedTokenizer>
     decode: (ids, skipSpecialTokens = false) =>
       // python's decode([]) === ""; the JS lib throws on empty input
       ids.length === 0 ? "" : tok.decode(ids, { skip_special_tokens: skipSpecialTokens }),
+    idToToken: (id) =>
+      (tok as unknown as { id_to_token(id: number): string | undefined })
+        .id_to_token(id) ?? tok.decode([id], { skip_special_tokens: false }),
     bosTokenId: idOf("bos_token"),
     eosTokenId: idOf("eos_token"),
   };

package/src/train/flash-cce.ts

@@ -723,20 +723,28 @@ const BWD_SG_SOURCE = String.raw`
   }
 `;
 
-// Apple-CCE coeff filter epsilon (phase-2 per-tile skip). DEFAULT OFF ("0") — the
-// Apple-production shortcut skips vocab tiles whose whole coeff block < eps, which gave
-// the OLD (slow, phase-2-dominated) SG kernel ~3.5×. On the steel kernel — whose phase-2
-// is now QuantizedBlockLoader-fast — it is a BAD trade: measured e4b M=512 REAL it cut dh
-// accuracy 0.66%→2.7% for only ~7% speedup (the win was already captured by the fast
-// dequant). Opt in with MLX_BUN_CCE_BWD_FILTER_EPS=1e-5 only on genuinely peaked data.
-const BWD_FILTER_EPS = process.env.MLX_BUN_CCE_BWD_FILTER_EPS ?? "0";
+// 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) 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;
@@ -1005,6 +1013,22 @@ function fwdSteelKernel(gs: number, bits: number): MetalKernel {
   return k;
 }
 
+// One-time dispatch-route warnings. A failed tiling guard silently reroutes the
+// whole flash head onto the legacy kernels (~5× slower class: e4b M=512 backward
+// 754 → 3687+ ms) with no observable signal — the exact "forward passes are
+// flag-gated, verify in code" hazard. Name the failed guard once per route so a
+// new model port that misses steel is visible in the training log.
+const _steelFallbackWarned = new Set<string>();
+function warnSteelFallback(which: "forward" | "backward", detail: string, fallback: string): void {
+  if (_steelFallbackWarned.has(which)) return;
+  _steelFallbackWarned.add(which);
+  console.warn(
+    `[flash-cce] ${which}: steel kernel guard failed (${detail}) — dispatching the ` +
+      `${fallback} instead (~5x-slower class). If this model should be steel-capable, ` +
+      `check the tiling guards in flash-cce.ts.`,
+  );
+}
+
 export function flashCceForward(
   hResp: MlxArray, head: FlashCceHead, targets: number[],
 ): { logp: MlxArray; lse: MlxArray; blockMax: MlxArray } {
@@ -1039,6 +1063,15 @@ export function flashCceForward(
   const blockV0 = Math.ceil(V / NBLK);
   const useSteel = process.env.MLX_BUN_CCE_NOSTEEL !== "1" && H % 32 === 0 && V % 32 === 0 && blockV0 % 32 === 0;
   const useSimd = !useSteel && USE_SIMDGROUP_FWD && H % SG_BD === 0;
+  if (!useSteel && process.env.MLX_BUN_CCE_NOSTEEL !== "1") {
+    // NOTE blockV0 = ceil(V/NBLK) can violate %32 even when V%32==0 (e.g. V=50304)
+    const fails = [
+      H % 32 !== 0 ? `H%32!=0 (H=${H})` : null,
+      V % 32 !== 0 ? `V%32!=0 (V=${V})` : null,
+      blockV0 % 32 !== 0 ? `blockV0%32!=0 (blockV0=${blockV0}, V=${V}, NBLK=${NBLK})` : null,
+    ].filter(Boolean).join(", ");
+    warnSteelFallback("forward", fails, useSimd ? "simdgroup kernel" : "scalar/lane kernel");
+  }
   const useLane = !useSteel && !useSimd && process.env.MLX_BUN_CCE_LANE === "1" && WPR % 32 === 0;
   const kernel = useSteel ? fwdSteelKernel(GS, head.bits) : useSimd ? fwdSgKernel() : useLane ? fwdLaneKernel() : fwdKernel();
   const tokTile = useSteel ? STEEL_BM : useSimd ? SG_BB : BLOCK_B;
@@ -1103,13 +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. Default ON (BWD_FILTER_EPS) for the ORPO training
-  // path — a pretrained model's softmax is peaked, so skipping near-zero-coeff rows
-  // is near-lossless and ~3.5× faster. Pass "0" for exact gradients (the parity gate
-  // does this, since synthetic-flat softmax has no negligible rows to skip).
+  // 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]!;
@@ -1119,7 +1152,10 @@ export function flashCceBackward(
   const GS = head.groupSize;
   const WPR = head.w.shape[1]!;
   if (H > H_MAX) throw new Error(`flashCceBackward: hidden ${H} > H_MAX ${H_MAX}`);
-  if (H % TG !== 0) throw new Error(`flashCceBackward: hidden ${H} not divisible by TG ${TG}`);
+  // NOTE: H % TG is a LEGACY-SCALAR-kernel requirement only (it owns DPER = H/TG
+  // dims per thread — floor division, so a remainder would leave dims silently
+  // zero). The steel/simdgroup kernels have their own (weaker) guards below; the
+  // throw is applied only when dispatch actually lands on the scalar kernel.
 
   const hF = hResp.dtype === Dtype.float32 ? hResp : hResp.astype(Dtype.float32);
   const scF = head.scales.dtype === Dtype.float32 ? head.scales : head.scales.astype(Dtype.float32);
@@ -1133,10 +1169,12 @@ export function flashCceBackward(
 
   // STEEL is the DEFAULT (verbatim MLX BlockMMA GEMM for both phases, [M,V]-free) when
   // the tiling is clean (H,V,blockV all %32 AND H % HTw for the H-tiled accumulator).
-  // MLX_BUN_CCE_BWD_NOSTEEL=1 falls back to the old kernels. The steel backward now
-  // 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
-  // (exact, the parity gate) and default-on for training (peaked softmax → ~all vocab cold).
+  // 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 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).
@@ -1144,6 +1182,26 @@ export function flashCceBackward(
   // simdgroup_matrix phase-2 GEMM is the next tier (5512→3913 ms, 1.41×) when the H
   // tiling is clean; MLX_BUN_CCE_BWD_SCALAR=1 forces the float4 scalar kernel.
   const useSg = !useSteel && process.env.MLX_BUN_CCE_BWD_SCALAR !== "1" && H % (SGN * 8) === 0 && H % SG_BD === 0;
+  if (!useSteel && process.env.MLX_BUN_CCE_BWD_NOSTEEL !== "1") {
+    // NOTE blockV0 = ceil(V/NBLK) can violate %32 even when V%32==0 (e.g. V=50304)
+    const fails = [
+      H % 32 !== 0 ? `H%32!=0 (H=${H})` : null,
+      V % 32 !== 0 ? `V%32!=0 (V=${V})` : null,
+      blockV0 % 32 !== 0 ? `blockV0%32!=0 (blockV0=${blockV0}, V=${V}, NBLK=${NBLK})` : null,
+      H % GS !== 0 ? `H%GS!=0 (H=${H}, GS=${GS})` : null,
+      GS % 32 !== 0 ? `GS%32!=0 (GS=${GS})` : null,
+    ].filter(Boolean).join(", ");
+    warnSteelFallback("backward", fails, useSg ? "simdgroup kernel" : "scalar kernel");
+  }
+  // Legacy scalar kernel: each thread owns DPER = H/TG dims (floor) — a remainder
+  // would leave the last H%TG dims silently zero, so it hard-requires H % TG == 0.
+  // Only enforced when dispatch actually lands there (steel/SG have weaker guards).
+  if (!useSteel && !useSg && H % TG !== 0)
+    throw new Error(
+      `flashCceBackward: hidden ${H} not divisible by TG ${TG} — required by the legacy ` +
+        `scalar kernel, which is the selected route because the steel/simdgroup tiling ` +
+        `guards also failed for this shape (see the [flash-cce] warning above).`,
+    );
   // Block-skip is on (either kernel) only when blockMax is supplied AND blockEps>0 (the
   // kernel reads it as an extra input). The coeff filter needs no extra input (just eps).
   const useBlockSkip = (useSteel || useSg) && blockMax != null && Number.parseFloat(blockEps) > 0;