mlx-bun 0.5.0 → 1.0.0

package/src/train/flash-cce.ts

@@ -1,1234 +0,0 @@
-// Flash Cut-Cross-Entropy Metal kernel for the ORPO log-prob head — the vocab
-// loop lives INSIDE the kernel, so neither the [M,V] logits nor a dequantized
-// [V,hidden] head ever touch HBM (the residency bound pure MLX can't reach; see
-// docs/design/orpo-training.md → "Vocab-blocked online-softmax"). Computes the
-// per-token response log-prob `logp = target_logit − logsumexp_v(h·W_v)`.
-//
-// TRANSPILE SOURCE — this is, in effect, Apple's Cut Cross Entropy + Liger's
-// FusedLinearCrossEntropy Triton kernels transpiled Triton → Metal:
-//   - Apple CCE  (github.com/apple/ml-cross-entropy):
-//       cce_lse_forward.py — forward: tile (BLOCK_B tokens × BLOCK_V vocab),
-//         accum = tl.dot(E_tile, Cᵀ_tile) over a D-loop, per-block
-//         this_lse = mx + log(Σ exp(logit−mx)), merged across vocab blocks into a
-//         per-token LSE via locked atomic logaddexp; mask offs_v<V → −inf.
-//       cce_backward.py — backward: same (token×vocab) tiling, recompute logits,
-//         d_accum = exp(accum − lse); d_accum += where(is_target, −1, 0); then
-//         dE += d_accum @ C and dC += d_accumᵀ @ E via locked atomic adds;
-//         optional skip when |d_accum| < filter_eps.
-//   - Liger FusedLinearCrossEntropy — token-chunking + the in-place softmax−onehot
-//       gradient written back over the logit buffer.
-//
-// DIVERGENCES (why this isn't a line-for-line port):
-//   1. Our classifier C is the QUANTIZED head (4/8-bit affine). Apple loads C
-//      dense (bf16); we keep the in-Metal dequant for `tl.dot` (the qdot pattern
-//      from mlx's quantized.h qdot pattern). This is the one substantive change.
-//   2. The head is FROZEN (not a LoRA target) → we need only dE (=dh), NOT dC.
-//   3. The LSE cross-vocab-block merge is done in cheap MLX ops on per-block
-//      partials, not Metal atomic-logaddexp spinlocks (simpler, same result).
-//
-// STATUS: v1 (one threadgroup per token, whole vocab serially) is correct +
-// memory-bounded (validated MiniCPM5→M=8192, e4b→M=512) but its single-token
-// threadgroup runs the full 262k vocab serially → trips the GPU watchdog on e4b
-// at higher M. The vocab-PARALLEL form below (grid over token × vocab blocks, à la
-// Apple CCE) bounds per-threadgroup work and is the production path.
-
-import { MetalKernel } from "../mlx/metal-kernel";
-import { MlxArray } from "../mlx/array";
-import { Dtype } from "../mlx/ffi";
-import * as ops from "../mlx/ops";
-import { STEEL_QMM_HEADER } from "./steel-qmm-header";
-import { runtimeValue } from "../runtime-config";
-
-/** The quantized head pieces the kernel needs (packed weight + per-group affine
- *  scale/bias + bits/group_size). Mirrors HeadQuant in loss.ts. */
-export interface FlashCceHead {
-  w: MlxArray; // packed quantized weight, uint32 [V, hidden*bits/32]
-  scales: MlxArray; // [V, hidden/groupSize]
-  biases: MlxArray; // [V, hidden/groupSize] (affine — required)
-  bits: number; // 4 or 8
-  groupSize: number; // 64
-  softcap: number | null;
-}
-
-const H_MAX = 4096; // covers MiniCPM5 (1536) and e4b (2560) hidden sizes
-const TG = 128; // threads per token-threadgroup
-// Vocab rows per threadgroup (Apple CCE's BLOCK_V analog). Bounds per-threadgroup
-// runtime so a single dispatch never loops the whole vocab — the e4b watchdog fix.
-const VOCAB_BLOCK = 8192;
-// Tokens processed together per threadgroup (Apple CCE's BLOCK_B). Each thread
-// dequantizes a vocab row's weights ONCE and reuses them across BLOCK_B tokens,
-// amortizing the ALU-bound dequant BLOCK_B× (the dominant cost in the v1
-// one-token-per-threadgroup kernel). Forward and backward use different tile
-// widths: the backward holds a per-token output accumulator `dacc[BLOCK_B·DPER]`
-// in registers, so it is register-bound to a smaller tile (BLOCK_B=8 spills).
-const BLOCK_B = 8;       // forward (online stats are only ~7·BLOCK_B floats/thread)
-const BLOCK_B_BWD = 4;   // backward (dacc[BLOCK_B·DPER] — 8 overflows the register file)
-// simdgroup_matrix forward GEMM tiles (Apple CCE's BLOCK_B×BLOCK_V×BLOCK_D, mapped
-// to Metal 8×8 simdgroup_matrix). BB tokens × BV vocab logit tile, accumulated over
-// BD hidden chunks. BV/8 = 4 column-blocks → one per simdgroup (TG/32 = 4).
-const SG_BB = 8;   // token tile (one 8-row simdgroup block)
-const SG_BV = 32;  // vocab tile (4 column-blocks)
-const SG_BD = 32;  // hidden chunk (4 k-blocks). BD=64 (GS-aligned, ½ the barriers)
-                   // measured SLOWER (fwd 844→1010, bwd 3311→3799): the 8KB Csh tile
-                   // cuts occupancy more than fewer barriers help. 32 is the sweet spot.
-
-// VOCAB-PARALLEL forward (Apple CCE cce_lse_forward, transpiled): one threadgroup
-// per (token=grid.z, vocab-block=grid.y) handles a bounded V/NBLK vocab slice → no
-// single threadgroup runs the whole vocab (the watchdog fix). Each writes per-block
-// partials (blockMax, blockSumexp-relative-to-blockMax, target logit) to [M,NBLK];
-// the cross-vocab-block LSE merge is done in cheap MLX ops (vs Apple's locked
-// atomic logaddexp). 128 threads stride the slice, per-thread online softmax,
-// merged across the threadgroup.
-const FWD_SOURCE = String.raw`
-  const uint tokBlk = threadgroup_position_in_grid.z; // token-block index
-  const uint vb  = threadgroup_position_in_grid.y;    // vocab-block index
-  const uint tid = thread_position_in_threadgroup.x;  // 0..TG-1
-  const uint M    = shape[0];
-  const uint V    = shape[1];
-  const uint H    = shape[2];
-  const uint GS   = shape[3];
-  const uint BITS = shape[4];
-  const uint WPR  = shape[5];   // uint32 words per packed weight row
-  const uint GR   = shape[6];   // groups per row (H/GS)
-  const uint NBLK = shape[7];   // number of vocab blocks
-  const float cap = capv[0];
-  const uint perWord = 32u / BITS;
-  const uint qmask = (BITS == 32u) ? 0xffffffffu : ((1u << BITS) - 1u);
-  const uint wordsPerGroup = GS / perWord;
-  const uint blockV = (V + NBLK - 1u) / NBLK;
-  const uint nStart = vb * blockV;
-  const uint nEnd = min(nStart + blockV, V);
-  const uint tokBase = tokBlk * BLOCK_B;
-
-  // Per-token online state (registers). Each thread keeps BLOCK_B tokens' running
-  // softmax stats so one vocab-row dequant feeds all BLOCK_B tokens.
-  float mx[BLOCK_B], sumexp[BLOCK_B], tgtLogit[BLOCK_B];
-  uint  tgtId[BLOCK_B];
-  for (uint b = 0u; b < BLOCK_B; ++b) {
-    mx[b] = -INFINITY; sumexp[b] = 0.0f; tgtLogit[b] = 0.0f;
-    tgtId[b] = (tokBase + b < M) ? targets[tokBase + b] : 0xffffffffu;
-  }
-
-  for (uint n = nStart + tid; n < nEnd; n += TG) {
-    const device uint* row = wp + (ulong)n * WPR;
-    const uint sbBase = n * GR;
-    float partial[BLOCK_B];
-    for (uint b = 0u; b < BLOCK_B; ++b) partial[b] = 0.0f;
-    // qdot factoring per group; the nibble dequant (q) is computed ONCE per
-    // element and reused across all BLOCK_B tokens (the amortization).
-    for (uint grp = 0u; grp < GR; ++grp) {
-      const float sc = scales[sbBase + grp];
-      const float bi = biases[sbBase + grp];
-      const uint d0 = grp * GS;
-      float gacc[BLOCK_B], ghsum[BLOCK_B];
-      for (uint b = 0u; b < BLOCK_B; ++b) { gacc[b] = 0.0f; ghsum[b] = 0.0f; }
-      for (uint wi = 0u; wi < wordsPerGroup; ++wi) {
-        const uint word = row[d0 / perWord + wi];
-        const uint dbase = d0 + wi * perWord;
-        for (uint k = 0u; k < perWord; ++k) {
-          const float q = (float)((word >> (BITS * k)) & qmask);
-          const uint d = dbase + k;
-          for (uint b = 0u; b < BLOCK_B; ++b) {
-            const uint tb = min(tokBase + b, M - 1u); // clamp padding lanes (last block) — result discarded in merge
-            // h is passed TRANSPOSED [H, M], so the BLOCK_B inner reads hit one
-            // cache line (h[d, tb..tb+BLOCK_B] contiguous) instead of BLOCK_B
-            // separate lines — the memory-access fix for the strided pattern.
-            const float hd = (float)h[(ulong)d * M + tb];
-            gacc[b] += hd * q; ghsum[b] += hd;
-          }
-        }
-      }
-      for (uint b = 0u; b < BLOCK_B; ++b) partial[b] += sc * gacc[b] + bi * ghsum[b];
-    }
-    for (uint b = 0u; b < BLOCK_B; ++b) {
-      float logit = partial[b];
-      if (cap > 0.0f) logit = cap * tanh(logit / cap);
-      const float nm = max(mx[b], logit);
-      sumexp[b] = sumexp[b] * exp(mx[b] - nm) + exp(logit - nm);
-      mx[b] = nm;
-      if (n == tgtId[b]) tgtLogit[b] = logit;
-    }
-  }
-
-  // Merge the TG threads' partials per token → per-(token, vocab-block) output.
-  threadgroup float mM[TG], mS[TG], mT[TG];
-  for (uint b = 0u; b < BLOCK_B; ++b) {
-    if (tokBase + b >= M) break;
-    mM[tid] = mx[b]; mS[tid] = sumexp[b]; mT[tid] = tgtLogit[b];
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    if (tid == 0u) {
-      float gm = -INFINITY, gs = 0.0f, gt = 0.0f;
-      for (uint i = 0u; i < TG; ++i) {
-        const float im = mM[i];
-        const float nm = max(gm, im);
-        if (nm > -INFINITY) { gs = gs * exp(gm - nm) + mS[i] * exp(im - nm); gm = nm; }
-        gt += mT[i]; // target logit (one thread nonzero), else 0
-      }
-      blockMax_out[(tokBase + b) * NBLK + vb] = gm;
-      blockSum_out[(tokBase + b) * NBLK + vb] = gs;
-      blockTgt_out[(tokBase + b) * NBLK + vb] = gt;
-    }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-  }
-`;
-
-// SIMDGROUP_MATRIX forward GEMM (the speed tier): per (token-block, vocab-block)
-// threadgroup, loop the vocab slice in BV-chunks; for each chunk compute the
-// [BB, BV] logit tile with 8×8 simdgroup_matrix mma over BD hidden chunks, then a
-// scalar online-softmax update. The classifier tile C [BV, BD] is dequantized into
-// threadgroup once per (vocab-chunk, D-chunk) and reused across all BB tokens by
-// the matmul — so the dequant is amortized AND the matmul runs at simdgroup-matrix
-// throughput (the scalar madd was the bottleneck after BLOCK_B tiling). BV/8 = 4
-// column-blocks → one per simdgroup (TG/32 = 4). logit[b,v] = Σ_d E[b,d]·C[v,d] =
-// E @ Cᵀ (the C block is loaded transposed).
-const FWD_SG_SOURCE = String.raw`
-  const uint tokBlk = threadgroup_position_in_grid.z;
-  const uint vb  = threadgroup_position_in_grid.y;
-  const uint tid = thread_position_in_threadgroup.x;
-  const uint sg  = tid / 32u;                 // simdgroup 0..3
-  const uint M=shape[0], V=shape[1], H=shape[2], GS=shape[3], BITS=shape[4], WPR=shape[5], GR=shape[6], NBLK=shape[7];
-  const float cap = capv[0];
-  const uint perWord = 32u / BITS;
-  const uint qmask = (BITS == 32u) ? 0xffffffffu : ((1u << BITS) - 1u);
-  const uint blockV = (V + NBLK - 1u) / NBLK;
-  const uint nStart = vb * blockV;
-  const uint nEnd = min(nStart + blockV, V);
-  const uint tokBase = tokBlk * BB;
-
-  // +4-float padded leading dims (MLX steel's BK_padded trick): a stride of 32 ==
-  // 32 banks → 32-way bank conflicts on every simdgroup_load; 36 breaks alignment.
-  threadgroup float Esh[BB * BDX];
-  threadgroup float Csh[BV * BDX];
-  threadgroup float Lsh[BB * BVX];
-  threadgroup float tgMax[BB], tgSum[BB], tgTgt[BB];
-  threadgroup uint  tgTid[BB];
-  if (tid < BB) {
-    tgMax[tid] = -INFINITY; tgSum[tid] = 0.0f; tgTgt[tid] = 0.0f;
-    tgTid[tid] = (tokBase + tid < M) ? targets[tokBase + tid] : 0xffffffffu;
-  }
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-
-  for (uint v0 = nStart; v0 < nEnd; v0 += BV) {
-    simdgroup_float8x8 acc = simdgroup_float8x8(0);
-    for (uint d0 = 0u; d0 < H; d0 += BD) {
-      // E tile [BB, BD] (token hiddens)
-      for (uint i = tid; i < BB * BD; i += TG) {
-        const uint b = i / BD, dd = i % BD;
-        const uint tb = min(tokBase + b, M - 1u);
-        Esh[b * BDX + dd] = (float)h[(ulong)tb * H + d0 + dd];
-      }
-      // C tile [BV, BD] (dequantized head, reused across all BB tokens).
-#if CCE_SG_SKIPDQ
-      // PROBE: no global loads — isolates whether the staging loads are the limiter.
-      for (uint i = tid; i < BV * BD; i += TG) {
-        const uint vrow = i / BD, dd = i % BD;
-        Csh[vrow * BDX + dd] = (v0 + vrow < nEnd) ? ((float)((i & 7u)) * 0.01f + 0.001f) : 0.0f;
-      }
-#else
-      // ONE THREAD PER PACKED WORD: each word loaded once (was perWord-redundant),
-      // and scale/bias once per word's group (was per-element — the whole BD chunk
-      // is one GS group, so this collapses ~perWord× of scale/bias traffic too).
-      const uint wordsPerRow = BD / perWord;
-      for (uint wi = tid; wi < BV * wordsPerRow; wi += TG) {
-        const uint vrow = wi / wordsPerRow;
-        const uint dloc = (wi % wordsPerRow) * perWord;
-        const uint n = v0 + vrow;
-        const uint d = d0 + dloc;
-        uint word = 0u; float sc = 0.0f, bi = 0.0f;
-        if (n < nEnd) {
-          word = wp[(ulong)n * WPR + d / perWord];
-          const uint grp = d / GS;
-          sc = scales[n * GR + grp]; bi = biases[n * GR + grp];
-        }
-        for (uint k = 0u; k < perWord; ++k) {
-          const float q = (float)((word >> (BITS * k)) & qmask);
-          Csh[vrow * BDX + dloc + k] = sc * q + bi;
-        }
-      }
-#endif
-      threadgroup_barrier(mem_flags::mem_threadgroup);
-      // acc(this sg's 8 cols) += E @ Cᵀ over BD/8 k-blocks (padded strides BDX)
-      for (uint kk = 0u; kk < BD / 8u; ++kk) {
-        simdgroup_float8x8 me, mc;
-        simdgroup_load(me, Esh + kk * 8u, BDX);
-        simdgroup_load(mc, Csh + (ulong)sg * 8u * BDX + kk * 8u, BDX, ulong2(0, 0), true);
-        simdgroup_multiply_accumulate(acc, me, mc, acc);
-      }
-      threadgroup_barrier(mem_flags::mem_threadgroup);
-    }
-    simdgroup_store(acc, Lsh + sg * 8u, BVX);  // [0:BB, sg*8 : sg*8+8], padded stride
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    // Online-softmax update for BB tokens over this BV chunk (thread b<BB → token b).
-    if (tid < BB) {
-      const uint b = tid;
-      for (uint v = 0u; v < BV; ++v) {
-        const uint n = v0 + v;
-        if (n >= nEnd) break;
-        float logit = Lsh[b * BVX + v];
-        if (cap > 0.0f) logit = cap * tanh(logit / cap);
-        const float nm = max(tgMax[b], logit);
-        tgSum[b] = tgSum[b] * exp(tgMax[b] - nm) + exp(logit - nm);
-        tgMax[b] = nm;
-        if (n == tgTid[b]) tgTgt[b] = logit;
-      }
-    }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-  }
-
-  if (tid < BB && tokBase + tid < M) {
-    blockMax_out[(tokBase + tid) * NBLK + vb] = tgMax[tid];
-    blockSum_out[(tokBase + tid) * NBLK + vb] = tgSum[tid];
-    blockTgt_out[(tokBase + tid) * NBLK + vb] = tgTgt[tid];
-  }
-`;
-
-// LANE-SLICED forward (the real memory fix): instead of one thread per vocab row
-// (adjacent threads read rows WPR words apart → uncoalesced weight reads, the
-// dominant ~43 GB stream at ~7% bandwidth), a SIMDGROUP of 32 lanes cooperates on
-// ONE row — lane l reads words {l, l+32, ...} so the 32 lanes hit CONSECUTIVE
-// memory each step (coalesced). The lane-partial dots reduce with simd_sum. The
-// threadgroup's 4 simdgroups process 4 rows in parallel; BLOCK_B tokens amortize
-// the dequant (each row dequantized once, reused across tokens). Same coalesced
-// pattern as mlx's quantized.h qdot. h is passed TRANSPOSED [H, M].
-const SGN = TG / 32; // simdgroups per threadgroup (4)
-const FWD_LANE_SOURCE = String.raw`
-  const uint tokBlk = threadgroup_position_in_grid.z;
-  const uint vb  = threadgroup_position_in_grid.y;
-  const uint lane = thread_index_in_simdgroup;       // 0..31 within the simdgroup
-  const uint sg   = thread_position_in_threadgroup.x / 32u; // 0..SGN-1
-  const uint M=shape[0], V=shape[1], H=shape[2], GS=shape[3], BITS=shape[4], WPR=shape[5], GR=shape[6], NBLK=shape[7];
-  const float cap = capv[0];
-  const uint perWord = 32u / BITS;
-  const uint qmask = (BITS == 32u) ? 0xffffffffu : ((1u << BITS) - 1u);
-  const uint wordsPerLane = WPR / 32u;               // WPR is a multiple of 32
-  const uint blockV = (V + NBLK - 1u) / NBLK;
-  const uint nStart = vb * blockV;
-  const uint nEnd = min(nStart + blockV, V);
-  const uint tokBase = tokBlk * BLOCK_B;
-
-  float mx[BLOCK_B], sumexp[BLOCK_B], tgtLogit[BLOCK_B];
-  uint  tgtId[BLOCK_B];
-  for (uint b = 0u; b < BLOCK_B; ++b) {
-    mx[b] = -INFINITY; sumexp[b] = 0.0f; tgtLogit[b] = 0.0f;
-    tgtId[b] = (tokBase + b < M) ? targets[tokBase + b] : 0xffffffffu;
-  }
-
-  // each simdgroup owns rows {nStart+sg, nStart+sg+SGN, ...}; its 32 lanes share the row
-  for (uint n = nStart + sg; n < nEnd; n += SGN) {
-    const device uint* row = wp + (ulong)n * WPR;
-    const uint sbBase = n * GR;
-    float partial[BLOCK_B];
-    for (uint b = 0u; b < BLOCK_B; ++b) partial[b] = 0.0f;
-    for (uint j = 0u; j < wordsPerLane; ++j) {
-      const uint wi = lane + 32u * j;                // coalesced across the 32 lanes
-      const uint word = row[wi];
-      const uint dbase = wi * perWord;
-      for (uint k = 0u; k < perWord; ++k) {
-        const uint d = dbase + k;
-        const uint grp = d / GS;
-        const float w = scales[sbBase + grp] * (float)((word >> (BITS * k)) & qmask) + biases[sbBase + grp];
-        for (uint b = 0u; b < BLOCK_B; ++b) {
-          const uint tb = min(tokBase + b, M - 1u);
-          partial[b] += (float)h[(ulong)d * M + tb] * w; // h transposed [H, M]
-        }
-      }
-    }
-    // reduce each token's lane-partials across the simdgroup → full logit
-    for (uint b = 0u; b < BLOCK_B; ++b) {
-      float logit = simd_sum(partial[b]);
-      if (cap > 0.0f) logit = cap * tanh(logit / cap);
-      const float nm = max(mx[b], logit);
-      sumexp[b] = sumexp[b] * exp(mx[b] - nm) + exp(logit - nm);
-      mx[b] = nm;
-      if (n == tgtId[b]) tgtLogit[b] = logit;
-    }
-  }
-
-  // simd_sum broadcasts, so each simdgroup's online state is replicated across its
-  // lanes; lane 0 of each sg publishes, then SGN states merge per token.
-  threadgroup float mM[SGN * BLOCK_B], mS[SGN * BLOCK_B], mT[SGN * BLOCK_B];
-  if (lane == 0u) {
-    for (uint b = 0u; b < BLOCK_B; ++b) {
-      mM[sg * BLOCK_B + b] = mx[b]; mS[sg * BLOCK_B + b] = sumexp[b]; mT[sg * BLOCK_B + b] = tgtLogit[b];
-    }
-  }
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-  const uint tid = thread_position_in_threadgroup.x;
-  if (tid < BLOCK_B && tokBase + tid < M) {
-    const uint b = tid;
-    float gm = -INFINITY, gs = 0.0f, gt = 0.0f;
-    for (uint s = 0u; s < SGN; ++s) {
-      const float im = mM[s * BLOCK_B + b];
-      const float nm = max(gm, im);
-      if (nm > -INFINITY) { gs = gs * exp(gm - nm) + mS[s * BLOCK_B + b] * exp(im - nm); gm = nm; }
-      gt += mT[s * BLOCK_B + b];
-    }
-    blockMax_out[(tokBase + b) * NBLK + vb] = gm;
-    blockSum_out[(tokBase + b) * NBLK + vb] = gs;
-    blockTgt_out[(tokBase + b) * NBLK + vb] = gt;
-  }
-`;
-
-let _fwdLaneKernel: MetalKernel | null = null;
-function fwdLaneKernel(): MetalKernel {
-  if (!_fwdLaneKernel) {
-    _fwdLaneKernel = new MetalKernel({
-      name: "flash_cce_fwd_lane",
-      inputNames: ["h", "wp", "scales", "biases", "targets", "shape", "capv"],
-      outputNames: ["blockMax_out", "blockSum_out", "blockTgt_out"],
-      source: FWD_LANE_SOURCE,
-      header: `#define TG ${TG}\n#define BLOCK_B ${BLOCK_B}\n#define SGN ${SGN}\n`,
-      ensureRowContiguous: true,
-    });
-  }
-  return _fwdLaneKernel;
-}
-
-let _fwdSgKernel: MetalKernel | null = null;
-function fwdSgKernel(): MetalKernel {
-  if (!_fwdSgKernel) {
-    _fwdSgKernel = new MetalKernel({
-      name: "flash_cce_fwd_sg",
-      inputNames: ["h", "wp", "scales", "biases", "targets", "shape", "capv"],
-      outputNames: ["blockMax_out", "blockSum_out", "blockTgt_out"],
-      source: FWD_SG_SOURCE,
-      header: `#include <metal_simdgroup_matrix>\n#define TG ${TG}\n#define BB ${SG_BB}\n#define BV ${SG_BV}\n#define BD ${SG_BD}\n#define BDX ${SG_BD + 4}\n#define BVX ${SG_BV + 4}\n#define CCE_SG_SKIPDQ ${runtimeValue("MLX_BUN_CCE_SG_SKIPDQ") === "1" ? 1 : 0}\n`,
-      ensureRowContiguous: true,
-    });
-  }
-  return _fwdSgKernel;
-}
-// simdgroup_matrix forward is now the DEFAULT (fastest): the earlier "not faster"
-// verdict was confounded — the SG path was bound by the Csh dequant STAGING loads
-// (packed weight + per-group scale/bias), not the matmul, so it never reached the
-// compute-bound regime. Removing the staging redundancy (one thread per packed
-// word; scale/bias once per group) cut it 1629→848 ms, beating the scalar
-// kernel's 1097 ms. Fall back to scalar (MLX_BUN_CCE_SCALAR=1) when H % SG_BD != 0.
-const USE_SIMDGROUP_FWD = runtimeValue("MLX_BUN_CCE_SCALAR") !== "1" && runtimeValue("MLX_BUN_CCE_LANE") !== "1";
-
-// Backward: dh[t,d] = Σ_v coeff_v · dequant(W_v)[d], coeff_v = g_t·(onehot_v −
-// softmax_v)·sech²_v, softmax_v = exp(softcap(logit_v) − lse_t). Two-phase per
-// 128-row vocab block, NO atomics: phase 1 each thread computes coeff for one
-// block row (full dot → into threadgroup coeff[]); phase 2 each thread owns
-// H/TG OUTPUT dims and accumulates them over the block's rows into per-thread
-// registers (dacc), so each output dim is written by exactly one thread. lse is
-// reused from the forward (no recompute of the global softmax normalizer).
-const DPER_MAX = H_MAX / TG; // 32
-const BWD_SOURCE = String.raw`
-  const uint tokBlk = threadgroup_position_in_grid.z; // token-block
-  const uint vb  = threadgroup_position_in_grid.y;    // vocab-block
-  const uint tid = thread_position_in_threadgroup.x;
-  const uint M   = shape[0];
-  const uint V   = shape[1];
-  const uint H   = shape[2];
-  const uint GS  = shape[3];
-  const uint BITS = shape[4];
-  const uint WPR = shape[5];
-  const uint GR  = shape[6];
-  const uint NBLK = shape[7];
-  const float cap = capv[0];
-  const uint perWord = 32u / BITS;
-  const uint qmask = (BITS == 32u) ? 0xffffffffu : ((1u << BITS) - 1u);
-  const uint wordsPerGroup = GS / perWord;
-  const uint DPER = H / TG;        // output dims this thread owns
-  const uint dStart = tid * DPER;
-  const uint blockV = (V + NBLK - 1u) / NBLK;
-  const uint nStart = vb * blockV;
-  const uint nEnd = min(nStart + blockV, V);
-  const uint tokBase = tokBlk * BLOCK_B;
-
-  float lseB[BLOCK_B], gB[BLOCK_B]; uint tgtB[BLOCK_B];
-  for (uint b = 0u; b < BLOCK_B; ++b) {
-    const uint tk = min(tokBase + b, M - 1u);
-    lseB[b] = lse[tk]; gB[b] = gv[tk];
-    tgtB[b] = (tokBase + b < M) ? targets[tokBase + b] : 0xffffffffu;
-  }
-
-  threadgroup float coeffB[TG * BLOCK_B];  // [row][token]
-  // BLOCK_B == 4 → one float4 per output dim holds all 4 tokens, so the phase-2
-  // accumulation (the backward's ~98% hotspot — a scalar coeffᵀ@W GEMM) runs as
-  // one vector FMA per (row, dim) instead of 4 scalar ones.
-  float4 dacc[DACC_N];                     // dacc[j] = (tok0..tok3) for dim dStart+j
-  for (uint i = 0u; i < DPER; ++i) dacc[i] = float4(0.0f);
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-
-  for (uint blk = nStart; blk < nEnd; blk += TG) {
-    // --- phase 1: row n = blk + tid → coeff for each of BLOCK_B tokens (the
-    // dequant-dot is computed once per token via shared nibble q) ---
-    const uint n = blk + tid;
-    float cf[BLOCK_B];
-    for (uint b = 0u; b < BLOCK_B; ++b) cf[b] = 0.0f;
-    if (n < nEnd) {
-      const device uint* row = wp + (ulong)n * WPR;
-      const uint sbBase = n * GR;
-      float raw[BLOCK_B];
-      for (uint b = 0u; b < BLOCK_B; ++b) raw[b] = 0.0f;
-      for (uint grp = 0u; grp < GR; ++grp) {
-        const float sc = scales[sbBase + grp];
-        const float bi = biases[sbBase + grp];
-        const uint d0 = grp * GS;
-        float gacc[BLOCK_B], ghsum[BLOCK_B];
-        for (uint b = 0u; b < BLOCK_B; ++b) { gacc[b] = 0.0f; ghsum[b] = 0.0f; }
-        for (uint wi = 0u; wi < wordsPerGroup; ++wi) {
-          const uint word = row[d0 / perWord + wi];
-          const uint dbase = d0 + wi * perWord;
-          for (uint k = 0u; k < perWord; ++k) {
-            const float q = (float)((word >> (BITS * k)) & qmask);
-            const uint d = dbase + k;
-            for (uint b = 0u; b < BLOCK_B; ++b) {
-              const uint tk = min(tokBase + b, M - 1u);
-              const float hd = (float)h[(ulong)tk * H + d]; // phase-1 h reads aren't the bwd bottleneck (phase 2 is)
-              gacc[b] += hd * q; ghsum[b] += hd;
-            }
-          }
-        }
-        for (uint b = 0u; b < BLOCK_B; ++b) raw[b] += sc * gacc[b] + bi * ghsum[b];
-      }
-      for (uint b = 0u; b < BLOCK_B; ++b) {
-        float logit = raw[b], sech2 = 1.0f;
-        if (cap > 0.0f) { const float th = tanh(raw[b] / cap); logit = cap * th; sech2 = 1.0f - th * th; }
-        const float sm = exp(logit - lseB[b]);
-        float c = gB[b] * ((n == tgtB[b] ? 1.0f : 0.0f) - sm);
-        if (cap > 0.0f) c *= sech2;
-        cf[b] = c;
-      }
-    }
-    for (uint b = 0u; b < BLOCK_B; ++b) coeffB[tid * BLOCK_B + b] = cf[b];
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-    // --- phase 2: the coeffᵀ@W accumulation — the backward's ~98% hotspot. Each
-    // weight dequant (wd) feeds all BLOCK_B tokens via one float4 FMA. (Loads are
-    // negligible — proven ~137 ms — so the per-element dequant stays; the win is
-    // vectorizing the arithmetic.) ---
-#if !CCE_BWD_SKIP_P2
-    for (uint r = 0u; r < TG; ++r) {
-      const uint n2 = blk + r;
-      if (n2 >= nEnd) break;
-      const device uint* row = wp + (ulong)n2 * WPR;
-      const uint sbBase = n2 * GR;
-      const float4 cf = *(const threadgroup float4*)(coeffB + r * BLOCK_B); // tok0..3
-#if CCE_BWD_SKIP_P2DQ
-      for (uint j = 0u; j < DPER; ++j) dacc[j] += cf * ((float)((dStart + j) & 7u) * 0.01f);
-#else
-      for (uint j = 0u; j < DPER; ++j) {
-        const uint d = dStart + j;
-        const uint grp = d / GS;
-        const uint word = row[d / perWord];
-        const float q = (float)((word >> (BITS * (d % perWord))) & qmask);
-        const float wd = scales[sbBase + grp] * q + biases[sbBase + grp];
-        dacc[j] += cf * wd;
-      }
-#endif
-    }
-#else
-    // PROBE: skip phase-2 accumulation but consume coeffB so phase 1 isn't DCE'd —
-    // isolates phase-1 (the coeff dot) cost from phase-2 accumulation.
-    dacc[0] += float4(coeffB[tid * BLOCK_B]);
-#endif
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-  }
-
-  // Atomic-accumulate each token's partial dh across the NBLK vocab-block programs.
-  // dacc[j] = (tok0..tok3) for dim dStart+j.
-  for (uint b = 0u; b < BLOCK_B; ++b) {
-    if (tokBase + b >= M) break;
-    for (uint j = 0u; j < DPER; ++j)
-      atomic_fetch_add_explicit(&dh[(ulong)(tokBase + b) * H + dStart + j], dacc[j][b], memory_order_relaxed);
-  }
-`;
-
-// dacc is sized to the ACTUAL per-thread output dims (ceil(H/TG)) rather than the
-// worst-case DPER_MAX — right-sizing the register footprint raises occupancy (the
-// backward's limiter). Keyed by daccN so a model with a different H rebuilds.
-const _bwdKernels = new Map<number, MetalKernel>();
-function bwdKernel(daccN: number): MetalKernel {
-  let k = _bwdKernels.get(daccN);
-  if (!k) {
-    k = new MetalKernel({
-      name: "flash_cce_bwd",
-      inputNames: ["h", "wp", "scales", "biases", "targets", "lse", "gv", "shape", "capv"],
-      outputNames: ["dh"],
-      source: BWD_SOURCE,
-      header: `#define TG ${TG}\n#define BLOCK_B ${BLOCK_B_BWD}\n#define DACC_N ${daccN}\n#define CCE_BWD_SKIP_P2DQ ${runtimeValue("MLX_BUN_CCE_BWD_SKIP_P2DQ") === "1" ? 1 : 0}\n#define CCE_BWD_SKIP_P2 ${runtimeValue("MLX_BUN_CCE_BWD_SKIP_P2") === "1" ? 1 : 0}\n`,
-      ensureRowContiguous: true,
-      atomicOutputs: true, // dh is summed across the NBLK vocab-block programs
-    });
-    _bwdKernels.set(daccN, k);
-  }
-  return k;
-}
-
-// SIMDGROUP_MATRIX backward (phase-2 GEMM): the phase-2 accumulation dacc = coeffᵀ@W
-// is ~69% of the backward (1733 ms phase-1 + 3779 ms phase-2 at e4b M=512) and is
-// compute-bound, NOT register-bound (right-sizing dacc gave 0%) — so we can afford
-// simdgroup_matrix's larger accumulator. Phase 1 unchanged (one thread per row →
-// coeff into threadgroup, BB=8 tokens). Phase 2: each simdgroup owns H/SGN dims as
-// DIMTILES persistent 8×8 acc fragments; per 8-row chunk, load coeffᵀ (transpose)
-// and the dequantized W tile, mma. dh atomic-added at the end. Needs H % (SGN*8)==0.
-const BB_BWD_SG = 8;
-const BWD_SG_SOURCE = String.raw`
-  const uint tokBlk = threadgroup_position_in_grid.z;
-  const uint vb  = threadgroup_position_in_grid.y;
-  const uint tid = thread_position_in_threadgroup.x;
-  const uint lane = thread_index_in_simdgroup;
-  const uint sg  = tid / 32u;
-  const uint M=shape[0],V=shape[1],H=shape[2],GS=shape[3],BITS=shape[4],WPR=shape[5],GR=shape[6],NBLK=shape[7];
-  const float cap = capv[0];
-  const uint perWord = 32u / BITS;
-  const uint qmask = (BITS == 32u) ? 0xffffffffu : ((1u << BITS) - 1u);
-  const uint HSG = H / SGN;          // dims owned by this simdgroup
-  const uint d0sg = sg * HSG;
-  const uint blockV = (V + NBLK - 1u) / NBLK;
-  const uint nStart = vb * blockV;
-  const uint nEnd = min(nStart + blockV, V);
-  const uint tokBase = tokBlk * BB;
-
-  float lseB[BB], gB[BB]; uint tgtB[BB];
-  for (uint b = 0u; b < BB; ++b) {
-    const uint tk = min(tokBase + b, M - 1u);
-    lseB[b] = lse[tk]; gB[b] = gv[tk];
-    tgtB[b] = (tokBase + b < M) ? targets[tokBase + b] : 0xffffffffu;
-  }
-
-  threadgroup float coeffB[TG * 8];   // [row][token], 8-wide for the simdgroup tile
-  threadgroup float Wsh[SGN * 64];    // per-simdgroup 8×8 dequantized W tile
-  threadgroup float Csh[SGN * 64];    // per-simdgroup 8×8 store-out buffer
-  threadgroup float Esh[BB * BDP];    // phase-1: token hiddens tile (logit GEMM)
-  threadgroup float Cstg[BVP * BDP];  // phase-1: dequantized W tile (logit GEMM)
-  threadgroup float Lsh[BB * BVP];    // phase-1: logit[token, row] tile
-
-#if CCE_BWD_BLOCKSKIP
-  // blockMax vocab-block early exit: this (token-block, vocab-block) program does
-  // ~0 work if, for EVERY token, the block's max softmax prob exp(blockMax−lse) is
-  // below eps AND the target isn't in this block (so no onehot term lives here).
-  // On peaked logits ~all but one vocab block per token is cold → skips phase 1+2.
-  // The early return is AFTER the threadgroup decls (uniform across the TG; declaring
-  // threadgroup arrays after an early return fails to compile).
-  bool warm = false;
-  for (uint b = 0u; b < BB; ++b) {
-    if (tokBase + b < M) {
-      if (exp(blockMax[(tokBase + b) * NBLK + vb] - lseB[b]) >= CCE_BWD_BLOCK_EPS) warm = true;
-      if (tgtB[b] >= nStart && tgtB[b] < nEnd) warm = true;
-    }
-  }
-  if (!warm) return;
-#endif
-
-  simdgroup_float8x8 acc[DIMTILES];
-  for (uint t = 0u; t < DIMTILES; ++t) acc[t] = simdgroup_float8x8(0);
-
-  for (uint blk = nStart; blk < nEnd; blk += TG) {
-    // --- phase 1 (SG-matrix): logit[tok, row] = E @ Wᵀ over H (the SAME tile GEMM
-    // as the forward, with the sub-block's TG rows as the "vocab"), then coeff. ---
-    for (uint vc = 0u; vc < TG; vc += BVP) {           // TG rows in BVP-row chunks
-      simdgroup_float8x8 lacc = simdgroup_float8x8(0);
-      for (uint d0 = 0u; d0 < H; d0 += BDP) {
-        for (uint i = tid; i < BB * BDP; i += TG) {    // E tile [BB tokens, BDP]
-          const uint b = i / BDP, dd = i % BDP;
-          const uint tk = min(tokBase + b, M - 1u);
-          Esh[i] = (float)h[(ulong)tk * H + d0 + dd];
-        }
-        const uint wordsPerRow = BDP / perWord;        // C tile [BVP rows, BDP] (one thread/word)
-        for (uint wi = tid; wi < BVP * wordsPerRow; wi += TG) {
-          const uint vrow = wi / wordsPerRow;
-          const uint dloc = (wi % wordsPerRow) * perWord;
-          const uint n2 = blk + vc + vrow;
-          const uint d = d0 + dloc;
-          uint word = 0u; float scl = 0.0f, bia = 0.0f;
-          if (n2 < nEnd) {
-            word = wp[(ulong)n2 * WPR + d / perWord];
-            const uint grp = d / GS; scl = scales[n2 * GR + grp]; bia = biases[n2 * GR + grp];
-          }
-          for (uint k = 0u; k < perWord; ++k) Cstg[vrow * BDP + dloc + k] = scl * (float)((word >> (BITS * k)) & qmask) + bia;
-        }
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        for (uint kk = 0u; kk < BDP / 8u; ++kk) {       // lacc(sg's 8 rows) += E @ Cᵀ
-          simdgroup_float8x8 me, mc;
-          simdgroup_load(me, Esh + kk * 8u, BDP);
-          simdgroup_load(mc, Cstg + (ulong)sg * 8u * BDP + kk * 8u, BDP, ulong2(0, 0), true);
-          simdgroup_multiply_accumulate(lacc, me, mc, lacc);
-        }
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-      }
-      simdgroup_store(lacc, Lsh + sg * 8u, BVP);        // Lsh[tok*BVP + row] = logit[tok][row]
-      threadgroup_barrier(mem_flags::mem_threadgroup);
-      for (uint i = tid; i < BB * BVP; i += TG) {        // derive coeff[row][tok] → coeffB
-        const uint tok = i / BVP, row = i % BVP;
-        const uint n2 = blk + vc + row;
-        float c = 0.0f;
-        if (n2 < nEnd && tokBase + tok < M) {
-          float logit = Lsh[tok * BVP + row], sech2 = 1.0f;
-          if (cap > 0.0f) { const float th = tanh(logit / cap); logit = cap * th; sech2 = 1.0f - th * th; }
-          const float sm = exp(logit - lseB[tok]);
-          c = gB[tok] * ((n2 == tgtB[tok] ? 1.0f : 0.0f) - sm);
-          if (cap > 0.0f) c *= sech2;
-        }
-        coeffB[(vc + row) * 8u + tok] = c;
-      }
-      threadgroup_barrier(mem_flags::mem_threadgroup);
-    }
-
-    // --- phase 2: acc[t][tok][dim] += Σ_row coeff[row][tok] · W[row][dim] ---
-    for (uint rc = 0u; rc < TG; rc += 8u) {
-      // Apple-CCE near-zero coeff filter: softmax mass concentrates on a few
-      // tokens, so coeff ≈ 0 for almost every vocab row → most 8-row chunks
-      // contribute nothing. Skip the chunk's dim-tile dequant+mma when its whole
-      // [8 rows × 8 tokens] coeff block is below eps. Safe because the SKIPPED
-      // softmax mass (not each element) is what's tiny. Uniform across the
-      // simdgroup (all lanes/sgs see the same coeffB[rc]).
-#if CCE_BWD_FILTER
-      float cmax = 0.0f;
-      for (uint i = lane; i < 64u; i += 32u) cmax = max(cmax, fabs(coeffB[rc * 8u + i]));
-      if (simd_max(cmax) < CCE_BWD_FILTER_EPS) continue;
-#endif
-      simdgroup_float8x8 mA;
-      simdgroup_load(mA, coeffB + rc * 8u, 8, ulong2(0, 0), true); // mA[tok][row]
-      threadgroup float* ws = Wsh + sg * 64u;
-      for (uint t = 0u; t < DIMTILES; ++t) {
-        const uint d0 = d0sg + t * 8u;
-        for (uint i = lane; i < 64u; i += 32u) {     // dequant W[8 rows, 8 dims]
-          const uint rr = i / 8u, dd = i % 8u;
-          const uint n2 = blk + rc + rr;
-          float w = 0.0f;
-          if (n2 < nEnd) {
-            const uint d = d0 + dd;
-            const uint word = wp[(ulong)n2 * WPR + d / perWord];
-            const float q = (float)((word >> (BITS * (d % perWord))) & qmask);
-            w = scales[n2 * GR + d / GS] * q + biases[n2 * GR + d / GS];
-          }
-          ws[i] = w;
-        }
-        simdgroup_barrier(mem_flags::mem_threadgroup);
-        simdgroup_float8x8 mB;
-        simdgroup_load(mB, ws, 8);                    // mB[row][dim]
-        simdgroup_multiply_accumulate(acc[t], mA, mB, acc[t]);
-        simdgroup_barrier(mem_flags::mem_threadgroup);
-      }
-    }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-  }
-
-  // write each dim-tile's acc → dh (atomic across the NBLK vocab-block programs)
-  threadgroup float* cs = Csh + sg * 64u;
-  for (uint t = 0u; t < DIMTILES; ++t) {
-    simdgroup_store(acc[t], cs, 8);
-    simdgroup_barrier(mem_flags::mem_threadgroup);
-    const uint d0 = d0sg + t * 8u;
-    for (uint i = lane; i < 64u; i += 32u) {
-      const uint tok = i / 8u, dd = i % 8u;
-      if (tokBase + tok < M)
-        atomic_fetch_add_explicit(&dh[(ulong)(tokBase + tok) * H + d0 + dd], cs[i], memory_order_relaxed);
-    }
-    simdgroup_barrier(mem_flags::mem_threadgroup);
-  }
-`;
-
-// 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
-// (flash-cce-filter-realdata.ts (deleted 2026-08-23; git history)) 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 = runtimeValue("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 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 = runtimeValue("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;
-  const blockOn = Number.parseFloat(blockEps) > 0 ? 1 : 0;
-  const key = `${dimTiles}:${filterEps}:${blockEps}`;
-  let k = _bwdSgKernels.get(key);
-  if (!k) {
-    const inputNames = ["h", "wp", "scales", "biases", "targets", "lse", "gv", "shape", "capv"];
-    if (blockOn) inputNames.push("blockMax"); // [M, NBLK] forward partial, for the skip
-    k = new MetalKernel({
-      name: "flash_cce_bwd_sg",
-      inputNames,
-      outputNames: ["dh"],
-      source: BWD_SG_SOURCE,
-      header: `#include <metal_simdgroup_matrix>\n#define TG ${TG}\n#define BB ${BB_BWD_SG}\n#define SGN ${SGN}\n#define DIMTILES ${dimTiles}\n#define BVP ${SG_BV}\n#define BDP ${SG_BD}\n#define CCE_BWD_FILTER ${filterOn}\n#define CCE_BWD_FILTER_EPS ${filterEps}f\n#define CCE_BWD_BLOCKSKIP ${blockOn}\n#define CCE_BWD_BLOCK_EPS ${blockEps}f\n`,
-      ensureRowContiguous: true,
-      atomicOutputs: true,
-    });
-    _bwdSgKernels.set(key, k);
-  }
-  return k;
-}
-
-// STEEL backward (the production path): MLX's verbatim BlockMMA GEMM (steel-qmm-
-// header.ts) for BOTH phases of dh = Σ_v coeff_v·dequant(W_v), at MLX speed AND
-// [M,V]-free. Per (token-block z = BM=8 tokens, vocab-block y) threadgroup:
-//   phase 1 — logit[8,32] = X[8,H] @ Wᵀ[H,32] via the SAME steel logit GEMM as the
-//     forward (mma_t + QuantizedBlockLoader, K-pipelined over H) → coeff epilogue.
-//   phase 2 — dh[8,H] += coeff[8,32] @ W[32,H], accumulated over the vocab tiles in a
-//     persistent register accumulator D = MMATile<1, H/32>. To keep threadgroup
-//     bounded (the full Wd[32,H] is 320 KB at e4b), H is tiled into HTw=128 columns:
-//     per H-tile a TEMP BlockMMA over a bounded Wd[32,HTw] dequant, whose Ctile is
-//     accumulated into D's frags LANE-LOCALLY (frag += frag — no get_coord math, the
-//     per-lane element mapping is identical between temp and D). dh atomic-added
-//     across the NBLK vocab-block programs. Validated dh ~1e-4 rel vs autograd.
-//     (Design + parity: steel-bwd-htile-test.ts (deleted 2026-08-23; git history).)
-//   PHASE-2 DEQUANT: the Wd[32,HTw] staging uses MLX's vectorized+fused
-//     QuantizedBlockLoader (the SAME loader the forward uses), NOT a manual scalar
-//     dequant. The wrinkle: phase-2's GEMM contracts over vocab (K=32) and outputs H
-//     (the dim the quant groups run along), so the loader must run with reduction_dim=0
-//     and one group per H-tile → HTw is pinned to group_size (so each H-tile is exactly
-//     one quant group, satisfying BCOLS<=group_size && group_size%BCOLS==0). The loader
-//     is built fresh per H-tile (correct src/scales/biases group offsets) + a single
-//     load_unsafe (no next()). TN = HTw/32 = group_size/32 frags per H-tile. The earlier
-//     manual scalar dequant (recomputing grp/scale/bias per element) was the bottleneck;
-//     this matches the forward's qmm staging. (Validated: steel-bwd-htile-test.ts.)
-// HTw (the phase-2 H-tile width) is pinned to the quant group_size in-kernel (HTw=GS_T),
-// so it is not a standalone constant here — the JS side derives NHT_T/NHT4_T from gs.
-const BWD_STEEL_SOURCE = String.raw`
-  const uint M=shape[0], V=shape[1], H=shape[2], NBLK=shape[7], WPR=shape[5];
-  const uint tokBlk=threadgroup_position_in_grid.z, vb=threadgroup_position_in_grid.y;
-  // HTw is pinned to the quant group_size so each phase-2 H-tile is exactly one group
-  // (lets QuantizedBlockLoader stage Wd with reduction_dim=0). TNv = HTw/32 frags/tile;
-  // NHTc = H/HTw tiles; total accumulator frags = NHTc*TNv = H/32 (= NHT4_T from JS).
-  constexpr int BM=8,BNv=32,BK=32,WM=1,WN=4,gs=GS_T,bts=BITS_T,HTw=GS_T,NHTc=NHT_T,TNv=GS_T/32;
-  constexpr int pf=get_pack_factor<bts,8>(), bpp=get_bytes_per_pack<bts>(), BKp=BK+16/sizeof(float);
-  threadgroup float Xs[BM*BKp]; threadgroup float Ws[BNv*BKp]; threadgroup float Ls[BM*BNv];
-  threadgroup float coeffS[BM*BNv]; threadgroup float Wd[BNv*HTw];
-  threadgroup float redMax[4]; threadgroup float tileMaxSh;   // coeff-filter reduction scratch
-  using lmma_t=mlx::steel::BlockMMA<float,float,BM,BNv,BK,WM,WN,false,true,BKp,BKp>;
-  using dmma_t=mlx::steel::BlockMMA<float,float,BM,HTw,BNv,WM,WN,false,false,BNv,HTw>;
-  using lx_t=mlx::steel::BlockLoader<float,BM,BK,BKp,1,WM*WN*SIMD_SIZE>;
-  using lw_t=QuantizedBlockLoader<float,BNv,BK,BKp,1,WM*WN*SIMD_SIZE,gs,bts>;
-  // phase-2 W loader: BROWS=32 vocab (GEMM K), BCOLS=HTw=group_size (one group/tile),
-  // reduction_dim=0, dst_ld=HTw — built fresh per H-tile, single load_unsafe.
-  using lw2_t=QuantizedBlockLoader<float,BNv,HTw,HTw,0,WM*WN*SIMD_SIZE,gs,bts>;
-  const uint sg=simdgroup_index_in_threadgroup, sl=thread_index_in_simdgroup, lid=thread_position_in_threadgroup.x;
-  const float cap=capv[0];
-  const int Hh=(int)H, K_w=Hh*bpp/pf, K_g=Hh/gs, tokBase=(int)tokBlk*BM;
-  const int blockV=((int)V+(int)NBLK-1)/(int)NBLK, nStart=(int)vb*blockV, nEnd=min(nStart+blockV,(int)V);
-  const short num_els=(short)min(BM,(int)M-tokBase);
-  short2 sc2=mlx::steel::BaseMMAFrag<float,8,8>::get_coord(sl); const short fm=sc2.y, fn=sc2.x;
-  const uint perWord=32u/bts, qmask=(bts==32u)?0xffffffffu:((1u<<bts)-1u);
-
-#if CCE_BWD_BLOCKSKIP
-  // Apple-CCE vocab-block early exit: this (token-block, vocab-block) program does ~0
-  // work if, for EVERY token, the block's max softmax prob exp(blockMax−lse) is below
-  // eps AND the target isn't in this block (no onehot term lives here). On peaked
-  // logits ~all but a couple of vocab blocks per token are cold → skips BOTH phases.
-  // Uniform across the threadgroup (all threads read the same blockMax/lse/targets), so
-  // the early return can't desync barriers; it is AFTER the threadgroup decls (required).
-  bool warm = false;
-  for (int b=0; b<BM; ++b){ int tk=tokBase+b; if (tk<(int)M){
-    if (exp(blockMax[tk*(int)NBLK + (int)vb] - lse[tk]) >= CCE_BWD_BLOCK_EPS) warm=true;
-    int tg2=(int)targets[tk]; if (tg2>=nStart && tg2<nEnd) warm=true; } }
-  if (!warm) return;
-#endif
-
-  mlx::steel::MMATile<float,1,NHT4_T> D;   // persistent dh accumulator (H/32 frags/sg)
-
-  for (int v0=nStart; v0<nEnd; v0+=BNv){
-    // phase 1: logit tile → coeff
-    lmma_t lmma(sg,sl);
-    lx_t loader_x(h+(int64_t)tokBase*Hh, Hh, Xs, sg, sl);
-    lw_t loader_w((const device uint8_t*)wp+v0*K_w, scales+v0*K_g, biases+v0*K_g, Hh, Ws, sg, sl);
-    for(int k=0;k<Hh;k+=BK){ threadgroup_barrier(mem_flags::mem_threadgroup);
-      if(num_els<BM) loader_x.load_safe(short2(BK,num_els)); else loader_x.load_unsafe();
-      loader_w.load_unsafe();
-      threadgroup_barrier(mem_flags::mem_threadgroup); lmma.mma(Xs,Ws); loader_x.next(); loader_w.next(); }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    lmma.Ctile.store<float,WM,WN,BNv,1>(Ls + lmma.sm*BNv + lmma.sn);
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    if (lid<(uint)BM){ uint b=lid;
-      if(tokBase+(int)b<(int)M){ const float lseb=lse[tokBase+b], gb=gv[tokBase+b]; const uint tgt=targets[tokBase+b];
-        for(int v=0;v<BNv;++v){ int n=v0+v; float logit=Ls[b*BNv+v],sech2=1.0f; if(cap>0.0f){float th=tanh(logit/cap);logit=cap*th;sech2=1.0f-th*th;} float sm=exp(logit-lseb); float c=(n<nEnd)?gb*(((uint)n==tgt?1.0f:0.0f)-sm):0.0f; if(cap>0.0f)c*=sech2; coeffS[b*BNv+v]=c; } }
-      else { for(int v=0;v<BNv;++v) coeffS[b*BNv+v]=0.0f; } }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
-#if CCE_BWD_FILTER
-    // Apple-CCE coeff filter: softmax mass concentrates on a few tokens, so coeff ≈ 0 for
-    // almost every vocab tile even within a warm block → skip the whole phase-2 (Wd staging
-    // + temp BlockMMA × NHT) when the entire coeff[8,32] tile is below eps. Parallel max
-    // over the 256 coeff elems (each thread 2, simd_max per simdgroup, combine 4). Uniform
-    // continue (all threads read the same tileMaxSh) so there is no barrier desync.
-    { float lm=max(fabs(coeffS[lid]), fabs(coeffS[lid+128u])); lm=simd_max(lm);
-      if (sl==0u) redMax[sg]=lm; }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    if (lid==0u) tileMaxSh=max(max(redMax[0],redMax[1]),max(redMax[2],redMax[3]));
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    if (tileMaxSh < CCE_BWD_FILTER_EPS) continue;
-#endif
-
-    // phase 2: per H-tile stage Wd[32,HTw] via QuantizedBlockLoader, temp BlockMMA, accum into D
-    for (int ht=0; ht<NHTc; ++ht){
-      const int hbase=ht*HTw;
-      // src at byte addr of W[v0, hbase]; scales/biases at group ht of vocab row v0.
-      lw2_t loader_w2((const device uint8_t*)wp + (int64_t)v0*K_w + (int64_t)hbase*bpp/pf,
-                      scales + (int64_t)v0*K_g + ht, biases + (int64_t)v0*K_g + ht, Hh, Wd, sg, sl);
-      loader_w2.load_unsafe();
-      threadgroup_barrier(mem_flags::mem_threadgroup);
-      dmma_t tmp(sg,sl); tmp.mma(coeffS, Wd);
-      for(int j=0;j<TNv;++j) D.frag_at(0, ht*TNv+j) += tmp.Ctile.frag_at(0,j);
-      threadgroup_barrier(mem_flags::mem_threadgroup);
-    }
-  }
-
-  // store D → dh (atomic-add across the NBLK vocab-block programs)
-  for(int gj=0; gj<NHT4_T; ++gj){
-    int ht=gj/TNv, j=gj%TNv; int row=tokBase+(int)fm; int col=ht*HTw + 8*(int)sg + (int)fn + j*32;
-    if(row<(int)M && col+1<=Hh){ thread auto& fr=D.frag_at(0,gj);
-      atomic_fetch_add_explicit(&dh[(int64_t)row*Hh + col], fr[0], memory_order_relaxed);
-      atomic_fetch_add_explicit(&dh[(int64_t)row*Hh + col+1], fr[1], memory_order_relaxed); }
-  }
-`;
-const BWD_STEEL_BM = 8;
-const _bwdSteelKernels = new Map<string, MetalKernel>();
-function bwdSteelKernel(gs: number, bits: number, H: number, filterEps: string, blockEps: string): MetalKernel {
-  // HTw == group_size (one quant group per phase-2 H-tile → QuantizedBlockLoader).
-  const nht = H / gs;            // # H-tiles
-  const nFrags = nht * (gs / 32); // total accumulator frags/sg = H/32 (NHT4_T name kept)
-  // Apple-CCE skips: coeff filter (phase-2 per-tile) + blockMax vocab-block early exit
-  // (whole-program). Both compile out when eps=0 → zero overhead, exact dh (the parity
-  // gate). blockMax skip needs the forward's [M,NBLK] blockMax as an extra input.
-  const filterOn = Number.parseFloat(filterEps) > 0 ? 1 : 0;
-  const blockOn = Number.parseFloat(blockEps) > 0 ? 1 : 0;
-  const key = `${gs}:${bits}:${H}:${filterEps}:${blockEps}`;
-  let k = _bwdSteelKernels.get(key);
-  if (!k) {
-    const inputNames = ["h", "wp", "scales", "biases", "targets", "lse", "gv", "shape", "capv"];
-    if (blockOn) inputNames.push("blockMax"); // [M, NBLK] forward partial, for the vocab-block skip
-    k = new MetalKernel({
-      name: "flash_cce_bwd_steel",
-      inputNames,
-      outputNames: ["dh"],
-      source: BWD_STEEL_SOURCE,
-      header: `${STEEL_QMM_HEADER}\n#define GS_T ${gs}\n#define BITS_T ${bits}\n#define NHT_T ${nht}\n#define NHT4_T ${nFrags}\n#define CCE_BWD_FILTER ${filterOn}\n#define CCE_BWD_FILTER_EPS ${filterEps}f\n#define CCE_BWD_BLOCKSKIP ${blockOn}\n#define CCE_BWD_BLOCK_EPS ${blockEps}f\n`,
-      ensureRowContiguous: true,
-      atomicOutputs: true,
-    });
-    _bwdSteelKernels.set(key, k);
-  }
-  return k;
-}
-
-let _fwdKernel: MetalKernel | null = null;
-function fwdKernel(): MetalKernel {
-  if (!_fwdKernel) {
-    _fwdKernel = new MetalKernel({
-      name: "flash_cce_fwd",
-      inputNames: ["h", "wp", "scales", "biases", "targets", "shape", "capv"],
-      outputNames: ["blockMax_out", "blockSum_out", "blockTgt_out"],
-      source: FWD_SOURCE,
-      header: `#define TG ${TG}\n#define BLOCK_B ${BLOCK_B}\n`,
-      ensureRowContiguous: true,
-    });
-  }
-  return _fwdKernel;
-}
-
-function u32(values: number[]): MlxArray {
-  // COPY into an mlx-owned array (fromBytesCopy = mlx_array_new_data), NOT the
-  // pinning fromView. These tiny kernel-arg arrays (shape, targets) feed a LAZY
-  // kernel and are disposed before it evaluates; pinning leaks one host buffer
-  // per call, never released (unbounded `pinned` map + a deferred native
-  // use-after-free under memory pressure — the seg/accum crash). A few-int copy
-  // is free; correctness is not.
-  return MlxArray.fromBytesCopy(new Uint8Array(new Uint32Array(values).buffer), [values.length], Dtype.uint32);
-}
-
-/** Flash-CCE forward: per-token response log-prob [M] (f32) + lse [M] (f32),
- *  computing logits in-kernel from the quantized head — no [M,V] materialized.
- *  `hResp` is [M, hidden] (any float dtype; cast to f32 for the kernel). `targets`
- *  are the response target token ids [M]. Caller owns the returned arrays. */
-// STEEL forward: MLX's verbatim quantized BlockMMA GEMM (steel-qmm-header.ts) for the
-// logit tile (BM=BN=BK=32, WM=WN=2 → 4 simdgroups × 16×16 register tiles, BK_padded),
-// with ONLY the ORPO epilogue diverging — softcap + online-softmax over each BN vocab
-// tile (no [M,V]). Per (token-block z, vocab-block y) threadgroup; partials [M,NBLK]
-// merged in MLX. Same logit math as quantizedMatmul (validated 1e-5) at MLX speed
-// (e4b logits 101 ms vs MLX 106 ms; 8.4× our old hand-rolled forward).
-const FWD_STEEL_SOURCE = String.raw`
-  const uint M=shape[0], V=shape[1], H=shape[2], NBLK=shape[7];
-  const uint tokBlk=threadgroup_position_in_grid.z, vb=threadgroup_position_in_grid.y;
-  constexpr int BM=32,BK=32,BN=32,WM=2,WN=2,gs=GS_T,bts=BITS_T;
-  constexpr int pf=get_pack_factor<bts,8>(), bpp=get_bytes_per_pack<bts>(), BKp=BK+16/sizeof(float);
-  threadgroup float Xs[BM*BKp]; threadgroup float Ws[BN*BKp]; threadgroup float Ls[BM*BN];
-  threadgroup float tgMax[BM], tgSum[BM], tgTgt[BM]; threadgroup uint tgTid[BM];
-  using mma_t=mlx::steel::BlockMMA<float,float,BM,BN,BK,WM,WN,false,true,BKp,BKp>;
-  using lx_t=mlx::steel::BlockLoader<float,BM,BK,BKp,1,WM*WN*SIMD_SIZE>;
-  using lw_t=QuantizedBlockLoader<float,BN,BK,BKp,1,WM*WN*SIMD_SIZE,gs,bts>;
-  const uint sg=simdgroup_index_in_threadgroup, sl=thread_index_in_simdgroup, lid=thread_position_in_threadgroup.x;
-  const float cap=capv[0];
-  const int Hh=(int)H, K_w=Hh*bpp/pf, K_g=Hh/gs;
-  const int blockV=((int)V+(int)NBLK-1)/(int)NBLK;
-  const int nStart=(int)vb*blockV, nEnd=min(nStart+blockV,(int)V), tokBase=(int)tokBlk*BM;
-  const short num_els=(short)min(BM,(int)M-tokBase);
-  if (lid<(uint)BM){ tgMax[lid]=-INFINITY; tgSum[lid]=0.0f; tgTgt[lid]=0.0f; tgTid[lid]=(tokBase+(int)lid<(int)M)?targets[tokBase+lid]:0xffffffffu; }
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-  for (int v0=nStart; v0<nEnd; v0+=BN){
-    mma_t mma_op(sg,sl);
-    lx_t loader_x(h+(int64_t)tokBase*Hh, Hh, Xs, sg, sl);
-    lw_t loader_w((const device uint8_t*)wp+v0*K_w, scales+v0*K_g, biases+v0*K_g, Hh, Ws, sg, sl);
-    for(int k=0;k<Hh;k+=BK){ threadgroup_barrier(mem_flags::mem_threadgroup);
-      if(num_els<BM) loader_x.load_safe(short2(BK,num_els)); else loader_x.load_unsafe();
-      loader_w.load_unsafe();
-      threadgroup_barrier(mem_flags::mem_threadgroup); mma_op.mma(Xs,Ws); loader_x.next(); loader_w.next(); }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    mma_op.Ctile.store<float,WM,WN,BN,1>(Ls + mma_op.sm*BN + mma_op.sn);
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    if (lid<(uint)BM && tokBase+(int)lid<(int)M){ uint b=lid;
-      for (int v=0; v<BN; ++v){ int n=v0+v; if(n>=nEnd) break;
-        float logit=Ls[b*BN+v]; if(cap>0.0f) logit=cap*tanh(logit/cap);
-        float nm=max(tgMax[b],logit); tgSum[b]=tgSum[b]*exp(tgMax[b]-nm)+exp(logit-nm); tgMax[b]=nm;
-        if((uint)n==tgTid[b]) tgTgt[b]=logit; }
-    }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-  }
-  if (lid<(uint)BM && tokBase+(int)lid<(int)M){ uint t=tokBase+lid;
-    blockMax_out[t*NBLK+vb]=tgMax[lid]; blockSum_out[t*NBLK+vb]=tgSum[lid]; blockTgt_out[t*NBLK+vb]=tgTgt[lid]; }
-`;
-const STEEL_BM = 32;
-const _fwdSteelKernels = new Map<string, MetalKernel>();
-function fwdSteelKernel(gs: number, bits: number): MetalKernel {
-  const key = `${gs}:${bits}`;
-  let k = _fwdSteelKernels.get(key);
-  if (!k) {
-    k = new MetalKernel({
-      name: "flash_cce_fwd_steel",
-      inputNames: ["h", "wp", "scales", "biases", "targets", "shape", "capv"],
-      outputNames: ["blockMax_out", "blockSum_out", "blockTgt_out"],
-      source: FWD_STEEL_SOURCE,
-      header: `${STEEL_QMM_HEADER}\n#define GS_T ${gs}\n#define BITS_T ${bits}\n`,
-      ensureRowContiguous: true,
-    });
-    _fwdSteelKernels.set(key, k);
-  }
-  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 } {
-  const M = hResp.shape[0]!;
-  const H = hResp.shape[1]!;
-  const V = head.scales.shape[0]!;
-  const GR = head.scales.shape[1]!;
-  const GS = head.groupSize;
-  const WPR = head.w.shape[1]!;
-  if (H > H_MAX) throw new Error(`flashCceForward: hidden ${H} > H_MAX ${H_MAX}`);
-  if (GR * GS !== H) throw new Error(`flashCceForward: GR*GS ${GR * GS} != H ${H}`);
-
-  // Vocab blocks sized so each threadgroup handles a bounded slice (≈ VOCAB_BLOCK
-  // rows) — independent of V, so a single threadgroup never runs the whole vocab
-  // (the e4b watchdog fix). NBLK = ceil(V / VOCAB_BLOCK).
-  const NBLK = Math.max(1, Math.ceil(V / VOCAB_BLOCK));
-  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);
-  const biF = head.biases.dtype === Dtype.float32 ? head.biases : head.biases.astype(Dtype.float32);
-  const tgtArr = u32(targets);
-  const shape = u32([M, V, H, GS, head.bits, WPR, GR, NBLK]);
-  const capv = MlxArray.fromFloat32(new Float32Array([head.softcap ?? 0]), [1]);
-
-  // DEFAULT = simdgroup_matrix kernel with efficient dequant staging (848 ms, the
-  // fastest measured — matmul throughput once the staging loads stopped bounding
-  // it). MLX_BUN_CCE_SCALAR=1 → one-thread-per-row scalar+coalesced-h (1097 ms;
-  // also the H % SG_BD != 0 fallback). MLX_BUN_CCE_LANE=1 → lane-sliced (2.6×
-  // slower; scatters the h broadcast — recorded dead end).
-  // DEFAULT = STEEL (verbatim MLX BlockMMA GEMM, 8.4× the old SG forward) when the
-  // tiling is clean (H,V,blockV all % 32). MLX_BUN_CCE_NOSTEEL=1 falls back to the
-  // old kernels (simdgroup / scalar / lane).
-  const blockV0 = Math.ceil(V / NBLK);
-  const useSteel = runtimeValue("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 && runtimeValue("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 && runtimeValue("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;
-  // steel + simdgroup read h as [M, H] (M on axis 0); lane + scalar read h TRANSPOSED
-  // [H, M] (M on axis 1) — they stage E into threadgroup themselves.
-  const transposedH = !useSteel && !useSimd;
-  const hIn = transposedH ? ops.transposeAxes(hF, [1, 0]) : hF;
-  let bMax: MlxArray | undefined, bSum: MlxArray | undefined, bTgt: MlxArray | undefined;
-  try {
-    const outs = kernel.apply([hIn, head.w, scF, biF, tgtArr, shape, capv], {
-      // Derive outputs + grid from the input shapes (hResp = inputs[0]) so the
-      // kernel composes inside an mx.compile'd closure: [M,NBLK] f32 partials, one
-      // threadgroup per (token-block, vocab-block). NBLK is fixed (from V).
-      outputShapeFn: (ins) => {
-        const m = transposedH ? ins[0]!.shape[1]! : ins[0]!.shape[0]!;
-        return [
-          { shape: [m, NBLK], dtype: Dtype.float32 },
-          { shape: [m, NBLK], dtype: Dtype.float32 },
-          { shape: [m, NBLK], dtype: Dtype.float32 },
-        ];
-      },
-      grid: (ins) => [TG, NBLK, Math.ceil((transposedH ? ins[0]!.shape[1]! : ins[0]!.shape[0]!) / tokTile)],
-      threadGroup: [TG, 1, 1],
-    });
-    [bMax, bSum, bTgt] = outs as [MlxArray, MlxArray, MlxArray];
-    // Cross-vocab-block merge in MLX: gMax = max_b; gSum = Σ_b sum_b·exp(max_b−gMax);
-    // lse = gMax + log(gSum); logp = target − lse. (Apple CCE does this with a
-    // locked atomic logaddexp in-kernel; the partials are tiny [M,NBLK] so an MLX
-    // reduction is simpler and just as exact.)
-    const gMax = ops.maxAxis(bMax, 1, false); // [M]
-    const gMaxCol = ops.reshape(gMax, [M, 1]);
-    // Bind every nested intermediate so its MlxArray wrapper gets disposed — a
-    // bare `ops.exp(ops.sub(...))` leaks the inner `sub` (stays live in `active`
-    // until the GC finalizer backstop), and this runs TWICE per step (chosen +
-    // rejected). gMax was also missing from the dispose list below.
-    const sub = ops.sub(bMax, gMaxCol);
-    const shifted = ops.exp(sub); // [M, NBLK]
-    const prod = ops.mul(bSum, shifted);
-    const gSum = ops.sumAxis(prod, 1, false); // [M]
-    const logG = ops.log(gSum);
-    const lse = ops.add(gMax, logG); // [M]
-    const tgt = ops.sumAxis(bTgt, 1, false); // [M] (owning block only is nonzero)
-    const logp = ops.sub(tgt, lse);
-    for (const a of [gMax, gMaxCol, sub, shifted, prod, gSum, logG, tgt]) a.dispose();
-    const blockMax = bMax!; bMax = undefined; // hand off to caller (backward block-skip)
-    return { logp, lse, blockMax };
-  } finally {
-    bMax?.dispose(); bSum?.dispose(); bTgt?.dispose();
-    if (hIn !== hF) hIn.dispose();
-    if (hF !== hResp) hF.dispose();
-    if (scF !== head.scales) scF.dispose();
-    if (biF !== head.biases) biF.dispose();
-    tgtArr.dispose();
-    shape.dispose();
-    capv.dispose();
-  }
-}
-
-/** Flash-CCE backward: dh [M, hidden] (f32) = ∂(Σ_t cotangent_t · logp_t)/∂h.
- *  `lse` is the forward's saved log-sum-exp [M]; `cot` is the per-token cotangent
- *  ∂loss/∂logp_t [M]. Recomputes logits in-kernel (no [M,V]); writes dh directly
- *  (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 — 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 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]!;
-  const H = hResp.shape[1]!;
-  const V = head.scales.shape[0]!;
-  const GR = head.scales.shape[1]!;
-  const GS = head.groupSize;
-  const WPR = head.w.shape[1]!;
-  if (H > H_MAX) throw new Error(`flashCceBackward: hidden ${H} > H_MAX ${H_MAX}`);
-  // 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);
-  const biF = head.biases.dtype === Dtype.float32 ? head.biases : head.biases.astype(Dtype.float32);
-  const lseF = lse.dtype === Dtype.float32 ? lse : lse.astype(Dtype.float32);
-  const NBLK = Math.max(1, Math.ceil(V / VOCAB_BLOCK));
-  const tgtArr = u32(targets);
-  const gv = MlxArray.fromFloat32(new Float32Array(cot), [M]);
-  const shape = u32([M, V, H, GS, head.bits, WPR, GR, NBLK]);
-  const capv = MlxArray.fromFloat32(new Float32Array([head.softcap ?? 0]), [1]);
-
-  // 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 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).
-  const useSteel = runtimeValue("MLX_BUN_CCE_BWD_NOSTEEL") !== "1" && H % 32 === 0 && V % 32 === 0 && blockV0 % 32 === 0 && H % GS === 0 && GS % 32 === 0;
-  // 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 && runtimeValue("MLX_BUN_CCE_BWD_SCALAR") !== "1" && H % (SGN * 8) === 0 && H % SG_BD === 0;
-  if (!useSteel && runtimeValue("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;
-  const bmF = useBlockSkip ? (blockMax!.dtype === Dtype.float32 ? blockMax! : blockMax!.astype(Dtype.float32)) : null;
-  const kernel = useSteel
-    ? bwdSteelKernel(GS, head.bits, H, filterEps, useBlockSkip ? blockEps : "0")
-    : useSg ? bwdSgKernel(H / SGN / 8, filterEps, useBlockSkip ? blockEps : "0") : bwdKernel(Math.ceil(H / TG));
-  const tokTile = useSteel ? BWD_STEEL_BM : useSg ? BB_BWD_SG : BLOCK_B_BWD;
-  const baseInputs = [hF, head.w, scF, biF, tgtArr, lseF, gv, shape, capv];
-  try {
-    const outs = kernel.apply(useBlockSkip ? [...baseInputs, bmF!] : baseInputs, {
-      outputShapeFn: (ins) => [{ shape: [ins[0]!.shape[0]!, H], dtype: Dtype.float32 }], // dh [M, hidden]
-      grid: (ins) => [TG, NBLK, Math.ceil(ins[0]!.shape[0]! / tokTile)],
-      threadGroup: [TG, 1, 1],
-      initValue: 0, // dh starts at 0; vocab-block programs atomic-add into it
-    });
-    return outs[0]!;
-  } finally {
-    if (hF !== hResp) hF.dispose();
-    if (bmF && bmF !== blockMax) bmF.dispose();
-    if (scF !== head.scales) scF.dispose();
-    if (biF !== head.biases) biF.dispose();
-    if (lseF !== lse) lseF.dispose();
-    tgtArr.dispose();
-    gv.dispose();
-    shape.dispose();
-    capv.dispose();
-  }
-}