mlx-bun 0.0.10 → 0.0.12

package/src/model/fused-decode-kernel.ts

@@ -1,312 +0,0 @@
-// Fused decode-SDPA Metal kernel, v2: flash-decoding split-N
-// (docs/design/optimization_plan.md Phase E). L=1 attention over a QUANTIZED KV
-// cache with the dequant inlined — neither the score row nor the
-// dequantized cache is ever materialized. Dispatched from Bun via
-// mx.fast.metal_kernel (no MLX fork), always OUTSIDE compiled-decode
-// closures (CustomKernel has no output_shapes; compiled segments
-// reclassify perf-kernel layers as JS layers instead).
-//
-// v1 post-mortem (git 48c9b10): numerics landed at the accepted tier-b
-// envelope but ONE THREADGROUP PER QUERY HEAD (16 total on the 12B)
-// could not occupy the GPU — 0.72× @8k. v2 splits the KV rows into
-// G = ceil(N/BLOCK) blocks (grid [128, H, G]):
-//   kernel A: per-block (max, sumexp) over T-rounded scores
-//   kernel C: per-block Σ T(exp(s−M)/denom)·dequant(V) partials, with
-//             the full-row (M, denom) derived inline from A's stats —
-//             deterministic (no atomics: atomic add order would make
-//             rounding nondeterministic)
-//   merge:    one mlx sum over G + cast to T (deterministic order)
-// Rounding points still MATCH the compat path where it matters (scores
-// and normalized probs round to T; f32 accumulation) — per-layer
-// divergence stays at the stray-ulp scale the teacher-forced gate is
-// calibrated for (tests/perf-kernel-oracle.test.ts).
-
-import { MlxArray } from "../mlx/array";
-import { Dtype } from "../mlx/ffi";
-import { MetalKernel } from "../mlx/metal-kernel";
-import * as ops from "../mlx/ops";
-
-/** Perf-mode lever: DEFAULT ON (2026-06-11 cleared-machine pass, ~1.02x @8k)
- *  — but an L3/explicit-only lever, OFF in the bare --l2 tier. This is an
- *  mlx-bun ORIGINAL kernel, NOT a port of optiq's fused decode, and the
- *  quantized-KV parity goldens do NOT track it: the optiq-golden decode
- *  composition is stock UNFUSED L=1 (scripts/regen-kvq-goldens.ts — the fused
- *  N-tiled SDPA covers L>1 prefill only, where this kernel never dispatches).
- *  It is NOT bit-exact vs the -O0 reference (online softmax, tier-b); its gate
- *  is ENVELOPE quality — ≥56/64 teacher-forced argmax agreement vs mlx-bun's
- *  own frozen compat trajectory (tests/perf-kernel-oracle.test.ts,
- *  scripts/freeze-perf-oracle.ts) plus bounded divergence vs the unfused port
- *  (tests/fused-decode-kernel.test.ts). MLX_BUN_PERF_KERNEL=0 selects the
- *  compat path we bench against the python libs. (B/C-style bit-exact
- *  optimizations stay on either way — this lever only toggles the
- *  parity-breaking decode kernel.) */
-export function perfKernelEnabled(): boolean {
-  return process.env.MLX_BUN_PERF_KERNEL !== "0";
-}
-
-/** Dispatch counter (the oracle gate asserts the kernel actually ran). */
-export let fusedKernelCalls = 0;
-
-const SIMD_GROUPS = 4;
-const LANES = 32;
-const TG_THREADS = SIMD_GROUPS * LANES;
-
-// Shared per-row score computation: lane-sliced dot product with the
-// dequant FACTORED out of the inner loop (the mlx qdot pattern,
-// quantized.h, Apache-2.0 — translated into this kernel's dialect):
-//   Σ (s·w_e + b)·q_e = s·Σ(w_e·q_e) + b·Σ q_e
-// and for 4-bit the nibbles multiply MASKED-BUT-UNSHIFTED against a
-// query pre-divided by 16^k — the inner loop is mask+madd only, via
-// 16-bit reads. Lane slices are group-aligned for D ∈ {256, 512},
-// GS ∈ {32, 64, 128}: one scale/bias per lane per row.
-const SCORE_PREAMBLE = String.raw`
-  constexpr int SG = ${SIMD_GROUPS};
-  constexpr int ELEMS = D / ${LANES};
-  constexpr int PER_WORD = 32 / BITS;
-  constexpr int WORDS = ELEMS / PER_WORD;
-
-  const uint lane = thread_index_in_simdgroup;
-  const uint sg = simdgroup_index_in_threadgroup;
-  const uint h = thread_position_in_grid.y;
-  const uint g = thread_position_in_grid.z;
-  const uint kv = h / NREP;
-  const int N = kp_shape[2];
-  const int G = (N + BLOCK - 1) / BLOCK;
-  const int nStart = (int)g * BLOCK;
-  const int nEnd = metal::min(nStart + BLOCK, N);
-  const int kRowWords = D / PER_WORD;
-  const int gRow = D / GS;
-
-  // lane's q slice: raw sum for the bias term, and (4-bit) nibble-place
-  // pre-scaled copies for the no-shift masked dot
-  float qv[ELEMS];
-  float qSum = 0.0f;
-  const uint qBase = h * D + lane * ELEMS;
-  for (int e = 0; e < ELEMS; e++) {
-    const float v = float(q[qBase + e]);
-    qSum += v;
-    qv[e] = BITS == 4
-      ? v * ((e & 3) == 0 ? 1.0f : (e & 3) == 1 ? 0.0625f : (e & 3) == 2 ? 0.00390625f : 0.000244140625f)
-      : v;
-  }
-
-  const int grpIdx = (lane * ELEMS) / GS;
-`;
-
-const SCORE_ROW = String.raw`
-    const ulong rowOff = ((ulong)kv * (ulong)N + (ulong)n);
-    const float kScale = float(ks[rowOff * gRow + grpIdx]);
-    const float kBias = float(kb[rowOff * gRow + grpIdx]);
-    float accum = 0.0f;
-    if (BITS == 4) {
-      const device uint16_t* kw =
-        ((const device uint16_t*)(kp + rowOff * kRowWords)) + lane * (ELEMS / 4);
-      for (int i = 0; i < ELEMS / 4; i++) {
-        const uint16_t wv = kw[i];
-        accum += qv[4 * i]     * (float)(wv & 0x000f)
-               + qv[4 * i + 1] * (float)(wv & 0x00f0)
-               + qv[4 * i + 2] * (float)(wv & 0x0f00)
-               + qv[4 * i + 3] * (float)(wv & 0xf000);
-      }
-    } else {
-      const device uint8_t* kw8 =
-        ((const device uint8_t*)(kp + rowOff * kRowWords)) + lane * ELEMS;
-      for (int e = 0; e < ELEMS; e++) accum += qv[e] * (float)kw8[e];
-    }
-    const float partial = kScale * accum + qSum * kBias;
-    const float s = float(T(simd_sum(partial)));
-`;
-
-// kernel: single-pass per-block ONLINE softmax — K and V each read
-// exactly once (the two-pass v2.0 read K twice and measured 0.945);
-// per-block (max, sumexp) and unnormalized o land in stats/opart for
-// the merge kernel.
-const BLOCK_SOURCE = SCORE_PREAMBLE + String.raw`
-  float mx = -INFINITY;
-  float sumexp = 0.0f;
-  // V accumulates RAW masked nibbles per slot (factor 16^k folds once at
-  // the end) with the bias term collapsed to one scalar: Σ p·b — the
-  // qdot factoring applied to the value side.
-  float oacc[ELEMS];
-  for (int e = 0; e < ELEMS; e++) oacc[e] = 0.0f;
-  float pbSum = 0.0f;
-
-  for (int n = nStart + (int)sg; n < nEnd; n += SG) {
-${SCORE_ROW}
-    const float newmax = metal::max(mx, s);
-    const float factor = metal::exp(mx - newmax);
-    const float p = metal::exp(s - newmax);
-    sumexp = sumexp * factor + p;
-    mx = newmax;
-
-    const float vScale = float(vs[rowOff * gRow + grpIdx]);
-    const float vBias = float(vb[rowOff * gRow + grpIdx]);
-    const float ps = p * vScale;
-    pbSum = pbSum * factor + p * vBias;
-    if (BITS == 4) {
-      const device uint16_t* vw =
-        ((const device uint16_t*)(vp + rowOff * kRowWords)) + lane * (ELEMS / 4);
-      for (int i = 0; i < ELEMS / 4; i++) {
-        const uint16_t wv = vw[i];
-        oacc[4 * i]     = oacc[4 * i]     * factor + ps * (float)(wv & 0x000f);
-        oacc[4 * i + 1] = oacc[4 * i + 1] * factor + ps * (float)(wv & 0x00f0);
-        oacc[4 * i + 2] = oacc[4 * i + 2] * factor + ps * (float)(wv & 0x0f00);
-        oacc[4 * i + 3] = oacc[4 * i + 3] * factor + ps * (float)(wv & 0xf000);
-      }
-    } else {
-      const device uint8_t* vw8 =
-        ((const device uint8_t*)(vp + rowOff * kRowWords)) + lane * ELEMS;
-      for (int e = 0; e < ELEMS; e++)
-        oacc[e] = oacc[e] * factor + ps * (float)vw8[e];
-    }
-  }
-  // fold the nibble-place factors and the bias sum (once per block)
-  for (int e = 0; e < ELEMS; e++) {
-    const float inv = BITS == 4
-      ? ((e & 3) == 0 ? 1.0f : (e & 3) == 1 ? 0.0625f : (e & 3) == 2 ? 0.00390625f : 0.000244140625f)
-      : 1.0f;
-    oacc[e] = oacc[e] * inv + pbSum;
-  }
-
-  // cross-simdgroup online merge within the block
-  threadgroup float tgMax[SG];
-  threadgroup float tgSum[SG];
-  threadgroup float tgO[SG * D];
-  if (lane == 0) {
-    tgMax[sg] = mx;
-    tgSum[sg] = sumexp;
-  }
-  for (int e = 0; e < ELEMS; e++) tgO[sg * D + lane * ELEMS + e] = oacc[e];
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-  if (sg == 0) {
-    float M = -INFINITY;
-    for (int gg = 0; gg < SG; gg++) M = metal::max(M, tgMax[gg]);
-    float S = 0.0f;
-    float wgt[SG];
-    for (int gg = 0; gg < SG; gg++) {
-      wgt[gg] = metal::exp(tgMax[gg] - M);
-      S += tgSum[gg] * wgt[gg];
-    }
-    for (int e = 0; e < ELEMS; e++) {
-      const int d = lane * ELEMS + e;
-      float o = 0.0f;
-      for (int gg = 0; gg < SG; gg++) o += tgO[gg * D + d] * wgt[gg];
-      opart[((ulong)h * (ulong)G + (ulong)g) * (ulong)D + (ulong)d] = o;
-    }
-    if (lane == 0) {
-      stats[(h * (uint)G + g) * 2 + 0] = M;
-      stats[(h * (uint)G + g) * 2 + 1] = S;
-    }
-  }
-`;
-
-// merge kernel: out[h,d] = Σ_g opart[h,g,d]·exp(m_g−M) / denom, cast T.
-// Deterministic (fixed g order), one tiny dispatch.
-const MERGE_SOURCE = String.raw`
-  const uint d = thread_position_in_grid.x;
-  const uint h = thread_position_in_grid.y;
-  const int G = stats_shape[2];
-  if ((int)d >= D) return;
-  float M = -INFINITY;
-  for (int g = 0; g < G; g++)
-    M = metal::max(M, stats[(h * (uint)G + (uint)g) * 2 + 0]);
-  float denom = 0.0f;
-  for (int g = 0; g < G; g++)
-    denom += stats[(h * (uint)G + (uint)g) * 2 + 1] *
-             metal::exp(stats[(h * (uint)G + (uint)g) * 2 + 0] - M);
-  float o = 0.0f;
-  for (int g = 0; g < G; g++)
-    o += opart[((ulong)h * (ulong)G + (ulong)g) * (ulong)D + (ulong)d] *
-         metal::exp(stats[(h * (uint)G + (uint)g) * 2 + 0] - M);
-  out[h * (uint)D + d] = T(o / denom);
-`;
-
-let blockKernel: MetalKernel | null = null;
-let mergeKernel: MetalKernel | null = null;
-
-function kernels(): { block: MetalKernel; merge: MetalKernel } {
-  if (!blockKernel) {
-    blockKernel = new MetalKernel({
-      name: "mlx_bun_fqsdpa_block",
-      inputNames: ["q", "kp", "ks", "kb", "vp", "vs", "vb"],
-      outputNames: ["opart", "stats"],
-      source: BLOCK_SOURCE,
-      ensureRowContiguous: true,
-    });
-    mergeKernel = new MetalKernel({
-      name: "mlx_bun_fqsdpa_merge",
-      inputNames: ["opart", "stats"],
-      outputNames: ["out"],
-      source: MERGE_SOURCE,
-      ensureRowContiguous: true,
-    });
-  }
-  return { block: blockKernel, merge: mergeKernel! };
-}
-
-/** Split width: enough blocks to occupy the GPU at small N without
- *  shrinking blocks into per-dispatch overhead at large N. */
-function blockSize(n: number): number {
-  return n <= 2048 ? 128 : 512;
-}
-
-/** Is this dispatch site servable by the fused kernel? (Phase D bakes
- *  bits/gs at generation time; shape checks guard runtime drift.) */
-export function fusedDecodeKernelSupported(
-  q: MlxArray, bits: number, groupSize: number,
-): boolean {
-  if (bits !== 4 && bits !== 8) return false;
-  if (groupSize !== 32 && groupSize !== 64 && groupSize !== 128) return false;
-  if (q.dtype !== Dtype.bfloat16) return false;
-  const [B, , L, D] = q.shape as [number, number, number, number];
-  if (B !== 1 || L !== 1) return false;
-  if (D % LANES !== 0) return false;
-  const elems = D / LANES;
-  if ((elems * bits) % 32 !== 0) return false; // whole u32 words per lane
-  // lane slice must sit inside one quantization group
-  if (elems > groupSize || groupSize % elems !== 0) return false;
-  return true;
-}
-
-/** Fused decode attention over quantized KV (flash-decoding split-N).
- *  Caller guarantees fusedDecodeKernelSupported and mask mode "" (decode
- *  attends the whole fetched window). q [1,H,1,D] bf16. */
-export function fusedDecodeSdpa(
-  q: MlxArray, kq: ops.QuantizedTensor, vq: ops.QuantizedTensor,
-  groupSize: number, bits: number,
-): MlxArray {
-  const [, H, , D] = q.shape as [number, number, number, number];
-  const KV = kq.packed.shape[1]!;
-  const N = kq.packed.shape[2]!;
-  const nRep = H / KV;
-  const BLOCK = blockSize(N);
-  const G = Math.ceil(N / BLOCK);
-  fusedKernelCalls++;
-
-  const { block, merge } = kernels();
-  const tmpl = { D, GS: groupSize, BITS: bits, NREP: nRep, BLOCK };
-
-  const [opart, statsArr] = block.apply(
-    [q, kq.packed, kq.scales, kq.biases, vq.packed, vq.scales, vq.biases],
-    {
-      outputs: [
-        { shape: [1, H, G, D], dtype: Dtype.float32 },
-        { shape: [1, H, G, 2], dtype: Dtype.float32 },
-      ],
-      grid: [TG_THREADS, H, G],
-      threadGroup: [TG_THREADS, 1, 1],
-      templateInts: tmpl,
-      templateDtypes: { T: q.dtype },
-    },
-  );
-  const [out] = merge.apply([opart!, statsArr!], {
-    outputs: [{ shape: [1, H, 1, D], dtype: q.dtype }],
-    grid: [D, H, 1],
-    threadGroup: [Math.min(D, 256), 1, 1],
-    templateInts: { D },
-    templateDtypes: { T: q.dtype },
-  });
-  opart!.dispose();
-  statsArr!.dispose();
-  return out!;
-}

package/src/model/fused-geglu-kernel.ts

@@ -1,198 +0,0 @@
-// Fused GeGLU Metal kernel (optimization_plan Phase E, opportunity B/C):
-// gelu_approx(a)·b in ONE pass, replacing the ~9 element-wise kernels the
-// spelled-out path emits (power, ×3, +, ×, tanh, +, ×, × + a final mul).
-// Both GeGLU sites are this shape — MLP (gelu(gate)·up) and the e4b
-// per-layer gate (gelu(gate)·pli).
-//
-// MATCH-COMPAT precision: bf16 in/out, each intermediate rounded to bf16
-// op-for-op with the spelled-out path (the mlx-lm parity constraint). A
-// first full-f32 version drifted 0.22 KL on e4b (small per-element error
-// compounding over 42 MLPs) — match-compat keeps the distribution on the
-// trained path; the only residual is the pow/tanh math-lib difference, which
-// the KL gate quantifies. Still a perf kernel (not bit-exact) and never
-// inside a compiled-decode trace: CustomKernel has no output_shapes (same
-// rule as fusedDecodeSdpa) — the caller guards on !isCompiledTrace().
-
-import { MlxArray } from "../mlx/array";
-import { Dtype } from "../mlx/ffi";
-import { MetalKernel } from "../mlx/metal-kernel";
-import { CustomVjp } from "../mlx/custom-vjp";
-import * as ops from "../mlx/ops";
-
-/** ON by default: the match-compat kernel is BIT-EXACT with the spelled-out
- *  path (kl --decode = 0 on e4b), so per the bit-parity-floor principle it
- *  ships on — every optimization we can keep while staying bit-exact.
- *  MLX_BUN_FUSED_GELU=0 opts out (e.g. the compile arm of the A/B). */
-export function fusedGeluEnabled(): boolean {
-  return process.env.MLX_BUN_FUSED_GELU !== "0";
-}
-
-/** Dispatch counter (a gate test asserts the kernel actually ran). */
-export let fusedGeluCalls = 0;
-
-// Training-mode flag: when set, model GeGLU sites use the DIFFERENTIABLE wrapper
-// (Metal-kernel forward + hand-derived vjp) instead of the plain `fusedGeglu`
-// (a bare CustomKernel with no vjp — feeding it to autograd throws
-// "[Primitive::vjp] Not implemented for CustomKernel"). The differentiable
-// wrapper also recomputes the gelu in the backward from the primal `a` instead
-// of retaining the ~9 forward intermediates, so it lowers the MLP activation
-// memory the backward holds — the dominant resident at long context under
-// ops.sdpa (which never materializes the O(T²) scores). Set by the trainer
-// around the train loop; cleared in its finally. Off → inference behaviour
-// (plain `fusedGeglu`) is byte-for-byte unchanged.
-let _trainingGeglu = false;
-export function setFusedGeluTraining(on: boolean): void {
-  _trainingGeglu = on;
-}
-export function fusedGeluTraining(): boolean {
-  return _trainingGeglu;
-}
-
-// Pointwise: one thread per element. gelu_approx =
-//   0.5 x (1 + tanh(√(2/π)·(x + 0.044715 x³)))  — composed in f32.
-const SOURCE = String.raw`
-  const uint i = thread_position_in_grid.x;
-  if (i >= SIZE) return;
-  // Match-compat: round each intermediate (and the weak scalars) to T, op
-  // for op with the spelled-out bf16 path, so the fused distribution tracks
-  // the TRAINED path rather than drifting to "more accurate but different".
-  // x^3 via pow (the reference uses mx.power, NOT x*x*x — they round
-  // differently in bf16). The single remaining residual vs compat is the
-  // pow/tanh math-lib difference, which the KL gate quantifies.
-  const float C_CUBE = float(T(0.044715f));
-  const float C_SQRT = float(T(0.7978845608028654f)); // √(2/π)
-  const float C_HALF = float(T(0.5f));
-  const float x   = float(a[i]);
-  const float x3  = float(T(metal::pow(x, 3.0f)));
-  const float cx3 = float(T(x3 * C_CUBE));
-  const float inn = float(T(x + cx3));
-  const float scl = float(T(inn * C_SQRT));
-  const float tnh = float(T(metal::precise::tanh(scl)));
-  const float t1  = float(T(1.0f + tnh));
-  const float hx  = float(T(x * C_HALF));
-  const float g   = float(T(hx * t1));
-  out[i] = T(g * float(b[i]));
-`;
-
-let kernel: MetalKernel | null = null;
-function getKernel(): MetalKernel {
-  if (!kernel)
-    kernel = new MetalKernel({
-      name: "mlx_bun_fused_geglu",
-      inputNames: ["a", "b"],
-      outputNames: ["out"],
-      source: SOURCE,
-      ensureRowContiguous: true,
-    });
-  return kernel;
-}
-
-/** bf16/f16 only (the activation dtype the kernel reads/writes). */
-export function fusedGeluSupported(a: MlxArray): boolean {
-  return a.dtype === Dtype.bfloat16 || a.dtype === Dtype.float16;
-}
-
-/** gelu_approx(a) · b, fused. a and b must share shape/dtype; out is a.dtype. */
-export function fusedGeglu(a: MlxArray, b: MlxArray): MlxArray {
-  fusedGeluCalls++;
-  const size = a.shape.reduce((x, y) => x * y, 1);
-  const [out] = getKernel().apply([a, b], {
-    outputs: [{ shape: a.shape, dtype: a.dtype }],
-    grid: [size, 1, 1],
-    threadGroup: [Math.min(size, 256), 1, 1],
-    templateInts: { SIZE: size },
-    templateDtypes: { T: a.dtype },
-  });
-  return out!;
-}
-
-// ---------------------------------------------------------------------------
-// Differentiable wrapper: Metal kernel forward + pure-MLX backward
-// ---------------------------------------------------------------------------
-
-// GELU tanh-approx constants
-const SQRT_2_PI = 0.7978845608028654;
-const C_CUBE = 0.044715;
-
-/** Compute gelu_approx(a) = 0.5 * a * (1 + tanh(sqrt(2/pi) * (a + 0.044715*a^3)))
- *  and gelu'(a) using the MLX ops layer. All intermediates are disposed
- *  before returning; caller owns the two returned arrays. */
-function geluAndGrad(a: MlxArray): { gelu: MlxArray; geluGrad: MlxArray } {
-  // a^3 = a * a * a
-  const a2 = ops.mul(a, a);
-  const a3 = ops.mul(a2, a);
-  // inner = a + 0.044715 * a^3
-  const cx3 = ops.mulScalar(a3, C_CUBE);
-  const inner = ops.add(a, cx3);
-  // z = sqrt(2/pi) * inner
-  const z = ops.mulScalar(inner, SQRT_2_PI);
-  // t = tanh(z)
-  const t = ops.tanh(z);
-  // gelu = 0.5 * a * (1 + t)
-  const one = ops.scalarLike(1, a);
-  const t1 = ops.add(one, t);
-  const halfA = ops.mulScalar(a, 0.5);
-  const gelu = ops.mul(halfA, t1);
-
-  // gelu'(a) = 0.5*(1+t) + 0.5*a*(1-t^2)*dz/da
-  // dz/da = sqrt(2/pi) * (1 + 3*0.044715*a^2)
-  // 1 - t^2 = sech^2(z)
-  const halfT1 = ops.mulScalar(t1, 0.5);            // 0.5*(1+t)
-  const t2 = ops.mul(t, t);                         // t^2
-  const sech2 = ops.sub(one, t2);                   // 1 - t^2
-  const halfASech2 = ops.mul(halfA, sech2);         // 0.5*a*(1-t^2)
-  // dz/da
-  const a2scaled = ops.mulScalar(a2, 3 * C_CUBE);   // 3*0.044715*a^2
-  const dzda = ops.add(one, a2scaled);               // 1 + 3*0.044715*a^2
-  const dzdaScaled = ops.mulScalar(dzda, SQRT_2_PI);// sqrt(2/pi)*dz_unnorm
-  // chain: 0.5*a*(1-t^2)*dz/da
-  const chainTerm = ops.mul(halfASech2, dzdaScaled);
-  const geluGrad = ops.add(halfT1, chainTerm);       // 0.5*(1+t) + chain
-
-  // dispose all intermediates
-  for (const arr of [a2, a3, cx3, inner, z, t, one, t1, halfA, halfT1, t2, sech2,
-                      halfASech2, a2scaled, dzda, dzdaScaled, chainTerm]) {
-    arr.dispose();
-  }
-
-  return { gelu, geluGrad };
-}
-
-let _differentiableKernel: CustomVjp | null = null;
-
-function getDifferentiableKernel(): CustomVjp {
-  if (!_differentiableKernel) {
-    _differentiableKernel = new CustomVjp(
-      // forward: use the Metal kernel
-      ([a, b]: MlxArray[]) => [fusedGeglu(a!, b!)],
-      // backward
-      (primals: MlxArray[], cotangents: MlxArray[], _outputs: MlxArray[]) => {
-        const [a, b] = primals as [MlxArray, MlxArray];
-        const [dc] = cotangents as [MlxArray];
-
-        const { gelu, geluGrad } = geluAndGrad(a);
-
-        // grad_b = dc * gelu(a)
-        const grad_b = ops.mul(dc, gelu);
-
-        // grad_a = dc * b * gelu'(a)
-        const dcb = ops.mul(dc, b);
-        const grad_a = ops.mul(dcb, geluGrad);
-
-        gelu.dispose();
-        geluGrad.dispose();
-        dcb.dispose();
-
-        return [grad_a, grad_b];
-      },
-    );
-  }
-  return _differentiableKernel;
-}
-
-/** gelu_approx(a) · b with a hand-derived vjp. Uses the Metal kernel forward,
- *  MLX-ops backward. For training only (inference uses the non-differentiable
- *  fusedGeglu). */
-export function fusedGegluDifferentiable(a: MlxArray, b: MlxArray): MlxArray {
-  return getDifferentiableKernel().apply([a, b])[0]!;
-}

package/src/model/fused-mlp-kernel.ts

@@ -1,169 +0,0 @@
-// The WHOLE MLP in ONE dispatch — down = (silu(x·Wgᵀ)·(x·Wuᵀ)) · Wdownᵀ — with no
-// intermediate `hidden` ever exposed to MLX. MLX runs this one kernel; there is no
-// boundary between our kernels for its planner to alias (the corruption only lives
-// at those boundaries). Built from the steel building blocks in steel-qmm-header.ts,
-// structured like the flash-CCE backward: each threadgroup owns one (M-block,I-block),
-// computes silu·up for that I-block ONCE (no recompute), then atomic-accumulates its
-// contribution to every down output — exactly how flash-CCE BWD atomic-adds dh.
-//
-//   grid = (TGN, I/BI, M/BM); each tg: i-block i0 = tid.y*BI, m-block m0 = tid.z*BM
-//   phase A: gate/up [BM,BI] = x[BM,H] @ {Wg,Wu}[i0,H]ᵀ  (reduce H) → silu·up = Hs[BM,BI]
-//   phase B: for each Ho-block no: D[BM,BN] = Hs @ Wdown[no,i0]ᵀ; atomic-add D → y[M,Ho]
-//
-// Per-weight quant constants (gate/up/down) are baked literals; H/I/Ho are templateInts;
-// M from x_shape. Output is f32 (atomic accumulation); the caller casts to the model
-// dtype at the MLP boundary. L3 (atomic-add ordering + bf16-rounded silu); gated by
-// teacher-forced KL like the rest of the controlled path.
-
-import { MlxArray } from "../mlx/array";
-import { Dtype } from "../mlx/ffi";
-import { MetalKernel } from "../mlx/metal-kernel";
-import { STEEL_QMM_HEADER } from "../train/steel-qmm-header";
-import * as ops from "../mlx/ops";
-import type { QuantSpec } from "../mlx/ops";
-
-const BM = 32, BN = 32, BK = 32, WM = 2, WN = 2;
-const TGN = WM * WN * 32; // 128 threads / threadgroup
-
-const SOURCE = String.raw`
-  const uint M=x_shape[0];
-  constexpr int I=I_T, H=H_T, Ho=HO_T;
-  constexpr int gsG=GS_G,btsG=BITS_G, gsU=GS_U,btsU=BITS_U, gsD=GS_D,btsD=BITS_D;
-  constexpr int pfG=get_pack_factor<btsG,8>(), bppG=get_bytes_per_pack<btsG>();
-  constexpr int pfU=get_pack_factor<btsU,8>(), bppU=get_bytes_per_pack<btsU>();
-  constexpr int pfD=get_pack_factor<btsD,8>(), bppD=get_bytes_per_pack<btsD>();
-  constexpr int BKp=BK_T+16/sizeof(T);
-  threadgroup T Xs[BM_T*BKp];
-  threadgroup T Wgs[BN_T*BKp];
-  threadgroup T Wus[BN_T*BKp];
-  threadgroup T Lg[BM_T*BN_T];
-  threadgroup T Lu[BM_T*BN_T];
-  threadgroup T Hs[BM_T*BKp];   // silu·up for this I-block, padded leading dim BKp
-  threadgroup T Wds[BN_T*BKp];  // a down-weight [BN, BI] tile
-  using gu_mma=mlx::steel::BlockMMA<T,T,BM_T,BN_T,BK_T,WM_T,WN_T,false,true,BKp,BKp>;
-  using lx_t  =mlx::steel::BlockLoader<T,BM_T,BK_T,BKp,1,WM_T*WN_T*SIMD_SIZE>;
-  using lwg_t =QuantizedBlockLoader<T,BN_T,BK_T,BKp,1,WM_T*WN_T*SIMD_SIZE,gsG,btsG>;
-  using lwu_t =QuantizedBlockLoader<T,BN_T,BK_T,BKp,1,WM_T*WN_T*SIMD_SIZE,gsU,btsU>;
-  using d_mma =mlx::steel::BlockMMA<T,T,BM_T,BN_T,BK_T,WM_T,WN_T,false,true,BKp,BKp>;
-  using lwd_t =QuantizedBlockLoader<T,BN_T,BK_T,BKp,1,WM_T*WN_T*SIMD_SIZE,gsD,btsD>;
-  const uint sg=simdgroup_index_in_threadgroup, sl=thread_index_in_simdgroup;
-  const uint lid=thread_position_in_threadgroup.x;
-  const int Hh=(int)H;
-  const int KwG=Hh*bppG/pfG, KgG=Hh/gsG;
-  const int KwU=Hh*bppU/pfU, KgU=Hh/gsU;
-  const int KwD=I*bppD/pfD,  KgD=I/gsD;          // down weight rows are length I
-  const int m0=(int)threadgroup_position_in_grid.z*BM_T;
-  const int i0=(int)threadgroup_position_in_grid.y*BN_T;
-  const short num_els=(short)min(BM_T,(int)M-m0);
-
-  // ---- Phase A: hidden[BM,BI] = silu(x@Wg[i0]ᵀ)·(x@Wu[i0]ᵀ), reduce over H ----
-  gu_mma mg(sg,sl), mu(sg,sl);
-  lx_t  lx (x  + (int64_t)m0*Hh, Hh, Xs, sg, sl);
-  lwg_t lwg((const device uint8_t*)wg + (int64_t)i0*KwG, sg_w + (int64_t)i0*KgG, bg_w + (int64_t)i0*KgG, Hh, Wgs, sg, sl);
-  lwu_t lwu((const device uint8_t*)wu + (int64_t)i0*KwU, su_w + (int64_t)i0*KgU, bu_w + (int64_t)i0*KgU, Hh, Wus, sg, sl);
-  for (int k=0;k<Hh;k+=BK_T){
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    if (num_els<BM_T) lx.load_safe(short2(BK_T,num_els)); else lx.load_unsafe();
-    lwg.load_unsafe(); lwu.load_unsafe();
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    mg.mma(Xs,Wgs); mu.mma(Xs,Wus);
-    lx.next(); lwg.next(); lwu.next();
-  }
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-  mg.Ctile.template store<T,WM_T,WN_T,BN_T,1>(Lg + mg.sm*BN_T + mg.sn);
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-  mu.Ctile.template store<T,WM_T,WN_T,BN_T,1>(Lu + mu.sm*BN_T + mu.sn);
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-  // silu·up into Hs (padded BKp leading dim, so the phase-B BlockMMA can read it as A)
-  for (uint idx=lid; idx<(uint)(BM_T*BN_T); idx+=TGN_T){
-    int r=(int)idx/BN_T, c=(int)idx%BN_T;
-    T gT=Lg[idx], uT=Lu[idx];
-    T sigT=(T)(1.0f/(1.0f+metal::precise::exp(-(float)gT)));
-    T siluT=(T)((float)gT*(float)sigT);
-    Hs[r*BKp + c] = (T)((float)siluT*(float)uT);
-  }
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-
-  // ---- Phase B: for each Ho-block, D[BM,BN] = Hs @ Wdown[no,i0]ᵀ; atomic-add → y ----
-  for (int no=0; no<Ho; no+=BN_T){
-    lwd_t lwd((const device uint8_t*)wd + (int64_t)no*KwD + (int64_t)i0*bppD/pfD,
-              sd_w + (int64_t)no*KgD + i0/gsD, bd_w + (int64_t)no*KgD + i0/gsD, I, Wds, sg, sl);
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    lwd.load_unsafe();               // one BI-wide tile (BI==BK, single step)
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    d_mma D(sg,sl);
-    D.mma(Hs, Wds);                  // D[BM,BN] = Hs[BM,BI] @ Wds[BN,BI]ᵀ
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    D.Ctile.template store<T,WM_T,WN_T,BN_T,1>(Lg + D.sm*BN_T + D.sn); // reuse Lg to stage D
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    for (uint idx=lid; idx<(uint)(BM_T*BN_T); idx+=TGN_T){
-      int r=(int)idx/BN_T, c=(int)idx%BN_T;
-      if (m0+r<(int)M){
-        const int64_t off=(int64_t)(m0+r)*(int64_t)Ho + (no+c);
-        // Absorb the residual: add h exactly once (only the i0==0 threadgroups,
-        // which cover every (m,ho) once) so the kernel outputs h + mlp(x) and
-        // nothing is left live across it for MLX to plan a buffer around.
-        float add = (float)Lg[idx];
-        if (i0==0) add += (float)h[off];
-        atomic_fetch_add_explicit(&y[off], add, memory_order_relaxed);
-      }
-    }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-  }
-`;
-
-const _kernels = new Map<string, MetalKernel>();
-function getKernel(
-  gsG: number, btsG: number, gsU: number, btsU: number, gsD: number, btsD: number,
-): MetalKernel {
-  const key = `${gsG}:${btsG}:${gsU}:${btsU}:${gsD}:${btsD}`;
-  let k = _kernels.get(key);
-  if (!k) {
-    k = new MetalKernel({
-      name: "mlx_bun_fused_mlp",
-      inputNames: ["x", "wg", "sg_w", "bg_w", "wu", "su_w", "bu_w", "wd", "sd_w", "bd_w", "h"],
-      outputNames: ["y"],
-      source: SOURCE,
-      header:
-        `${STEEL_QMM_HEADER}\n` +
-        `#define GS_G ${gsG}\n#define BITS_G ${btsG}\n#define GS_U ${gsU}\n#define BITS_U ${btsU}\n` +
-        `#define GS_D ${gsD}\n#define BITS_D ${btsD}\n` +
-        `#define BM_T ${BM}\n#define BN_T ${BN}\n#define BK_T ${BK}\n` +
-        `#define WM_T ${WM}\n#define WN_T ${WN}\n#define TGN_T ${TGN}\n`,
-      ensureRowContiguous: true,
-      atomicOutputs: true,
-    });
-    _kernels.set(key, k);
-  }
-  return k;
-}
-
-/** The whole MLP, fused, as one dispatch. x is [M, H] (bf16/f16). Returns the MLP
- *  output [M, Ho] in x.dtype (cast from the f32 atomic accumulator). gate/up have
- *  shape [I, H*bits/32]; down has [Ho, I*bits/32]. Caller owns the result. */
-export function fusedMlp(
-  x: MlxArray,
-  wg: MlxArray, sg: MlxArray, bg: MlxArray, gateSpec: QuantSpec,
-  wu: MlxArray, su: MlxArray, bu: MlxArray, upSpec: QuantSpec,
-  wd: MlxArray, sd: MlxArray, bd: MlxArray, downSpec: QuantSpec,
-  h: MlxArray,  // the residual (layer input) — absorbed: output = h + mlp(x)
-): MlxArray {
-  const M = x.shape[0]!, I = wg.shape[0]!, H = x.shape[1]!, Ho = wd.shape[0]!;
-  const [yF] = getKernel(
-    gateSpec.groupSize, gateSpec.bits, upSpec.groupSize, upSpec.bits, downSpec.groupSize, downSpec.bits,
-  ).apply([x, wg, sg, bg, wu, su, bu, wd, sd, bd, h], {
-    // f32 atomic accumulator, zero-initialised; grid tiles (M, I) — every I-block
-    // contributes to the output via atomic-add.
-    outputShapeFn: (ins) => [{ shape: [ins[0]!.shape[0]!, Ho], dtype: Dtype.float32 }],
-    grid: (ins) => [TGN, Math.ceil(I / BN), Math.ceil(ins[0]!.shape[0]! / BM)],
-    threadGroup: [TGN, 1, 1],
-    templateInts: { I_T: I, H_T: H, HO_T: Ho },
-    templateDtypes: { T: x.dtype },
-    initValue: 0,
-  });
-  void M;
-  if (yF!.dtype === x.dtype) return yF!;
-  const y = yF!.astype(x.dtype);
-  yF!.dispose();
-  return y;
-}