mlx-bun 0.0.7 → 0.0.9

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/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/diffusion-lora.ts

@@ -0,0 +1,216 @@
+// DiffusionGemma LoRA fine-tuning (port D5). DiffusionGemma is a discrete
+// block-diffusion LM, so the AR cross-entropy trainer doesn't apply — train LoRA
+// with the model's NATIVE denoising objective (port of optiq
+// vlm/diffusion_gemma/lora.py `diffusion_loss` + `train_diffusion_lora`):
+//   corrupt the target canvas to a random noise level t∈[t_min,t_max] (replace
+//   each token with prob t by a uniform-random token — the same corruption the
+//   inference canvas starts from), run one encoder+decoder forward, and minimise
+//   cross-entropy on the CORRUPTED positions (predict the clean token).
+//
+// LoRA is injected on the decoder blocks (DEFAULT_LORA_KEYS = attn q/k/v/o +
+// dense MLP gate/up/down). The encoder reuses those same blocks (weight-tied,
+// only per-layer scalars differ), so one injection trains both paths. Experts/
+// router stay frozen (their routing indices are stop_gradient'd so the MoE
+// backward differentiates only the activations). Adapter saved in the
+// AdapterManager.mount layout (lora_a/lora_b), mountable for inference.
+
+import * as ops from "../mlx/ops";
+import { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import { ValueAndGrad } from "../mlx/autograd";
+import { DiffusionGemmaModel } from "../model/diffusion-gemma";
+import { AdamW } from "./optimizer";
+import {
+  attachForTraining,
+  buildTrainableLora,
+  detachTraining,
+  flatParams,
+  type TrainableLora,
+} from "./lora-params";
+
+export const DEFAULT_LORA_KEYS = [
+  "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj",
+  "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj",
+];
+
+export interface DiffusionPair {
+  promptIds: number[];
+  targetIds: number[];
+}
+
+export interface DiffusionLoraConfig {
+  rank?: number; // default 8
+  scale?: number; // default 8.0 (diffusion is scale-sensitive; AR's 20 collapses)
+  iters?: number; // default 200
+  learningRate?: number; // default 1e-4
+  tMin?: number; // default 0.4
+  tMax?: number; // default 0.8
+  seed?: number; // default 0
+  reportEvery?: number; // default 10
+  onReport?: (iter: number, avgLoss: number) => void;
+}
+
+/** swap/restore the differentiated primals into the LoRA leaves (port of the
+ *  trainer.ts private helpers — primals order is [...A, ...B]). */
+function swapPrimals(lora: TrainableLora, primals: MlxArray[]): MlxArray[] {
+  const n = lora.targets.length;
+  const saved: MlxArray[] = [];
+  for (let i = 0; i < n; i++) {
+    saved.push(lora.targets[i]!.lw.a, lora.targets[i]!.lw.b);
+    lora.targets[i]!.lw.a = primals[i]!;
+    lora.targets[i]!.lw.b = primals[n + i]!;
+  }
+  return saved;
+}
+function restorePrimals(lora: TrainableLora, saved: MlxArray[]): void {
+  for (let i = 0; i < lora.targets.length; i++) {
+    lora.targets[i]!.lw.a = saved[2 * i]!;
+    lora.targets[i]!.lw.b = saved[2 * i + 1]!;
+  }
+}
+function writeParam(lora: TrainableLora, i: number, p: MlxArray): void {
+  const n = lora.targets.length;
+  if (i < n) lora.targets[i]!.lw.a = p;
+  else lora.targets[i - n]!.lw.b = p;
+}
+
+/** sum all elements -> scalar []. */
+function sumAll(a: MlxArray): MlxArray {
+  const flat = ops.reshape(a, [a.size]);
+  const s = ops.sumAxis(flat, 0, false);
+  flat.dispose();
+  return s;
+}
+
+/** Denoising cross-entropy on corrupted canvas positions (the reference
+ *  diffusion_loss). Returns a scalar loss MlxArray, differentiable w.r.t. the
+ *  mounted LoRA. */
+export function diffusionLoss(
+  model: DiffusionGemmaModel,
+  promptIds: number[],
+  targetIds: number[],
+  tMin: number,
+  tMax: number,
+  vocab: number,
+): MlxArray {
+  const L = targetIds.length;
+  // t ~ U(t_min, t_max) scalar; corrupt[i] = U(0,1) < t ; noise ~ randint
+  const t = ops.randomUniform([1], Dtype.float32, tMin, tMax, null);
+  const u = ops.randomUniform([1, L], Dtype.float32, 0, 1, null);
+  const corrupt = ops.less(u, t); // bool [1, L]
+  u.dispose();
+  t.dispose();
+  const noise = ops.randint(0, vocab, [1, L], Dtype.int32);
+  const target = MlxArray.fromInt32(Int32Array.from(targetIds), [1, L]);
+  const canvas = ops.where(corrupt, noise, target);
+  noise.dispose();
+
+  const logits = model.forwardCanvasLogitsArr(promptIds, canvas); // [1, L, vocab] f32
+  canvas.dispose();
+
+  // cross_entropy(reduction=none) = logsumexp(logits) - logit_at_target
+  const lse = ops.logsumexpAxis(logits, -1, false); // [1, L]
+  const tExp = ops.expandDims(target, -1); // [1, L, 1]
+  const gathered = ops.takeAlongAxis(logits, tExp, -1); // [1, L, 1]
+  tExp.dispose();
+  logits.dispose();
+  const tgt = ops.reshape(gathered, [1, L]);
+  gathered.dispose();
+  const ce = ops.sub(lse, tgt); // [1, L]
+  lse.dispose();
+  tgt.dispose();
+  target.dispose();
+
+  const corruptF = corrupt.astype(Dtype.float32);
+  corrupt.dispose();
+  const masked = ops.mul(ce, corruptF);
+  ce.dispose();
+  const num = sumAll(masked);
+  masked.dispose();
+  const cnt = sumAll(corruptF);
+  corruptF.dispose();
+  const one = ops.scalarLike(1, cnt);
+  const denom = ops.maximum(cnt, one);
+  cnt.dispose();
+  one.dispose();
+  const loss = ops.div(num, denom);
+  num.dispose();
+  denom.dispose();
+  return loss;
+}
+
+/** Train a LoRA adapter on DiffusionGemma with the denoising objective. Returns
+ *  the TrainableLora (caller saves it via saveAdapter) + the loss history. */
+export function trainDiffusionLora(
+  model: DiffusionGemmaModel,
+  pairs: DiffusionPair[],
+  cfg: DiffusionLoraConfig = {},
+): { lora: TrainableLora; losses: number[] } {
+  const rank = cfg.rank ?? 8;
+  const scale = cfg.scale ?? 8.0;
+  const iters = cfg.iters ?? 200;
+  const lr = cfg.learningRate ?? 1e-4;
+  const tMin = cfg.tMin ?? 0.4;
+  const tMax = cfg.tMax ?? 0.8;
+  const seed = cfg.seed ?? 0;
+  const reportEvery = cfg.reportEvery ?? 10;
+  const vocab = model.config.text.vocabSize;
+  if (pairs.length === 0) throw new Error("no training pairs");
+
+  ops.randomSeed(BigInt(seed >>> 0));
+
+  // LoRA on every decoder block's attn + dense MLP (skips full-layer v_proj
+  // automatically — buildTrainableLora filters to existing loraTargets).
+  const ranks = new Map<string, number>();
+  for (let i = 0; i < model.layers.length; i++)
+    for (const key of DEFAULT_LORA_KEYS) ranks.set(`model.decoder.layers.${i}.${key}`, rank);
+  const lora = buildTrainableLora(model, ranks, scale, seed);
+  attachForTraining(model, lora, "train");
+  model.loraState.active = ["train"];
+
+  const params = flatParams(lora);
+  const opt = new AdamW(
+    params,
+    { lr, betas: [0.9, 0.999], eps: 1e-8, weightDecay: 0.01 },
+    (i, p) => writeParam(lora, i, p),
+  );
+
+  let current: DiffusionPair = pairs[0]!;
+  const vag = new ValueAndGrad((primals) => {
+    const saved = swapPrimals(lora, primals);
+    try {
+      return diffusionLoss(model, current.promptIds, current.targetIds, tMin, tMax, vocab);
+    } finally {
+      restorePrimals(lora, saved);
+    }
+  }, params.map((_, i) => i));
+
+  const losses: number[] = [];
+  try {
+    for (let it = 0; it < iters; it++) {
+      current = pairs[it % pairs.length]!;
+      if (current.targetIds.length === 0) continue;
+      const { value, grads } = vag.apply(flatParams(lora));
+      opt.step(grads); // takes ownership of grads + writes via writeParam
+      ops.evalAll([...flatParams(lora), value]);
+      const lossNum = value.astype(Dtype.float32).toFloat32()[0]!;
+      value.dispose();
+      losses.push(lossNum);
+      if ((it + 1) % reportEvery === 0) {
+        const window = losses.slice(-reportEvery);
+        cfg.onReport?.(it + 1, window.reduce((a, b) => a + b, 0) / window.length);
+      }
+    }
+  } catch (e) {
+    // On failure, undo the mount so the model is left clean.
+    detachTraining(model, lora);
+    model.loraState.active = [];
+    vag.dispose();
+    throw e;
+  }
+  // On success leave the trained LoRA ATTACHED + ACTIVE so the model is
+  // immediately usable; the caller saves it (saveAdapter) and may detach.
+  vag.dispose();
+  model.loraState.active = ["train"];
+  return { lora, losses };
+}

package/src/train/flash-cce.ts

@@ -723,12 +723,16 @@ 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.
+// 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
+// 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";
 // 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
@@ -1005,6 +1009,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 +1059,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,10 +1132,11 @@ 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, 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).
   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.
@@ -1119,7 +1149,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 +1166,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 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").
   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 +1179,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;