mlx-bun 0.0.9 → 0.0.11

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

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

@@ -1,240 +0,0 @@
-// Fused SwiGLU MLP expansion kernel (cpm5-generate-and-fuse-inference): compute
-//   hidden = silu(gate_proj(x)) · up_proj(x)
-// in ONE Metal dispatch, reusing MLX's verbatim steel quantized BlockMMA GEMM
-// (src/train/steel-qmm-header.ts — the same building block the ORPO flash-CCE
-// head uses). Replaces the spelled-out path's two quantized_matmul dispatches +
-// the gate/up/silu/product global-memory materialization (minicpm5.ts LlamaMLP):
-//   gate = qmm(x, Wg);  up = qmm(x, Wu);  silu = gate·σ(gate);  hidden = silu·up
-// down_proj is LEFT as a separate qmm so the `hidden` boundary stays bit-exact —
-// the existing minicpm5 L1/L2 parity goldens (maxDiff===0) gate this path directly
-// (their 6-token prompt prefill is an M>1 MLP call).
-//
-// WHERE THE WIN IS: prefill / M>1, where the BlockMMA GEMM is the right tool and
-// the [M, intermediate] materialization is real. At decode (M=1) the matmuls are
-// memory-bound GEMVs (reading Wg/Wu dominates), so the fused kernel only saves the
-// per-token dispatch + a tiny [1,I] round-trip — the caller falls back to unfused
-// below M_MIN. Benchmarks (prefill AND decode) decide the default.
-//
-// PRECISION (two variants, env-flagged — "flag and try both"):
-//  - MATCH (default): T = the activation dtype (bf16/f16). The gate/up tiles are
-//    stored bf16-rounded out of the float BlockMMA accumulator, exactly as MLX's
-//    quantized_matmul casts its float accum to the input dtype; σ/·/· then run
-//    op-for-op with ops.sigmoid+ops.mul+ops.mul on bf16. AIM: maxDiff===0 vs the
-//    unfused path (like the fused-gelu kernel, kl=0). Only residual: the metal
-//    exp() vs MLX's math lib — quantified by the parity test.
-//  - F32 (MLX_BUN_FUSED_SWIGLU_F32=1): keep gate/up/silu in f32 through the fusion,
-//    round to T only at the final store. More accurate, NOT bit-exact → L3, gated
-//    by teacher-forced KL like the fused-decode kernel.
-//
-// Never inside a compiled-decode trace: a CustomKernel has no output_shapes (same
-// rule as fusedDecodeSdpa / fusedGeglu) — the caller guards on !isCompiledTrace().
-
-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";
-
-/** OFF by default until the prefill/decode benchmark + parity delta are in (then
- *  flip to on like FUSED_GELU if it lands bit-exact). MLX_BUN_FUSED_SWIGLU=1 opts in. */
-export function fusedSwigluEnabled(): boolean {
-  return process.env.MLX_BUN_FUSED_SWIGLU === "1";
-}
-
-/** f32-internal variant (NOT bit-exact). Default off → the bit-exact MATCH path. */
-function f32Internal(): boolean {
-  return process.env.MLX_BUN_FUSED_SWIGLU_F32 === "1";
-}
-
-/** Dispatch counter — a gate test asserts the kernel actually ran. */
-export let fusedSwigluCalls = 0;
-
-/** Debug bisect: cap how many fused calls actually run (rest fall back). */
-function maxCalls(): number {
-  const v = Number(process.env.MLX_BUN_SWIGLU_MAXCALLS);
-  return Number.isFinite(v) && v >= 0 ? v : Infinity;
-}
-export function swigluUnderCap(): boolean {
-  return fusedSwigluCalls < maxCalls();
-}
-
-// Below this many rows (M) the matmuls are memory-bound GEMVs and the fusion win
-// is just dispatch overhead; the caller keeps the unfused path. Tunable via env
-// for the decode-vs-prefill sweep.
-export function swigluMMin(): number {
-  const v = Number(process.env.MLX_BUN_FUSED_SWIGLU_MMIN);
-  return Number.isFinite(v) && v > 0 ? v : 2;
-}
-
-const BM = 32, BN = 32, BK = 32, WM = 2, WN = 2;
-const TGN = WM * WN * 32; // 128 threads / threadgroup
-
-// One BM×BN output tile per threadgroup (grid.y = I-tiles, grid.z = M-tiles), the
-// K=H reduction looped inside. Xs is loaded once and fed to BOTH the gate and up
-// BlockMMA (the whole point of fusing — x is read from global once for two GEMMs).
-const SOURCE = String.raw`
-  // M from the x_shape helper (varies per call); I, H BAKED as template literals
-  // (model-constant) — no runtime `+"`shape`"+` input array to thread through the
-  // lazy graph. The per-loader quant constants (GS_G/BITS_G gate, GS_U/BITS_U up)
-  // are baked too — the de-branch from the known per-tensor mixed-precision config,
-  // so the gate=4/up=8 layers fuse instead of falling back.
-  const uint M=x_shape[0]; constexpr int I=I_T, H=H_T;
-  constexpr int gsG=GS_G, btsG=BITS_G, gsU=GS_U, btsU=BITS_U;
-  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 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];
-  using mma_t=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>;
-  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 K_wG=Hh*bppG/pfG, K_gG=Hh/gsG;   // gate: packed-row bytes, group count
-  const int K_wU=Hh*bppU/pfU, K_gU=Hh/gsU;   // up
-  const int m0=(int)threadgroup_position_in_grid.z*BM_T;
-  const int n0=(int)threadgroup_position_in_grid.y*BN_T;
-  const short num_els=(short)min(BM_T,(int)M-m0);
-  const short num_outs=(short)min(BN_T,(int)I-n0);
-
-#if SWIGLU_TGZERO
-  // Diagnostic: zero threadgroup memory so a back-to-back dispatch can't read the
-  // previous kernel's leftover staging (the cross-dispatch hazard hypothesis).
-  for (uint z=lid; z<(uint)(BM_T*BKp); z+=TGN_T) Xs[z]=T(0);
-  for (uint z=lid; z<(uint)(BN_T*BKp); z+=TGN_T){ Wgs[z]=T(0); Wus[z]=T(0); }
-  for (uint z=lid; z<(uint)(BM_T*BN_T); z+=TGN_T){ Lg[z]=T(0); Lu[z]=T(0); }
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-#endif
-  mma_t 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)n0*K_wG, sg_w + (int64_t)n0*K_gG, bg_w + (int64_t)n0*K_gG, Hh, Wgs, sg, sl);
-  lwu_t  lwu((const device uint8_t*)wu + (int64_t)n0*K_wU, su_w + (int64_t)n0*K_gU, bu_w + (int64_t)n0*K_gU, 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();
-    if (num_outs<BN_T){ lwg.load_safe(short2(BK_T,num_outs)); lwu.load_safe(short2(BK_T,num_outs)); }
-    else { 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);
-  // Stage both accumulators to threadgroup. store<T,...> casts the float accum to
-  // T (bf16) — the SAME rounding MLX's quantized_matmul applies to its output, so
-  // Lg/Lu hold the bf16 gate/up the unfused path would have written to global.
-  // Xs/Wgs/Wus alias Lg/Lu only if a kernel reused the same tgmem — they don't here,
-  // but the stores write threadgroup that the epilogue reads, so bracket with barriers.
-  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);
-
-  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 || n0+c>=(int)I) continue;
-#if SWIGLU_DEBUG_GATE
-    hidden[(int64_t)(m0+r)*(int64_t)I + n0+c] = Lg[idx]; // raw gate proj (debug)
-#elif SWIGLU_F32
-    // f32-internal: gate/up stayed bf16 in Lg/Lu (the matmul output is the same),
-    // but σ·/· run in f32 with no intermediate rounding → NOT bit-exact.
-    float g=(float)Lg[idx], u=(float)Lu[idx];
-    float sig=1.0f/(1.0f+metal::precise::exp(-g));
-    hidden[(int64_t)(m0+r)*(int64_t)I + n0+c] = (T)(g*sig*u);
-#else
-    // MATCH: round σ, silu, product to T op-for-op with ops.sigmoid+mul+mul on bf16.
-    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);
-    hidden[(int64_t)(m0+r)*(int64_t)I + n0+c] = (T)((float)siluT*(float)uT);
-#endif
-  }
-`;
-
-const _kernels = new Map<string, MetalKernel>();
-function getKernel(gsG: number, btsG: number, gsU: number, btsU: number, f32: boolean): MetalKernel {
-  const key = `${gsG}:${btsG}:${gsU}:${btsU}:${f32 ? 1 : 0}`;
-  let k = _kernels.get(key);
-  if (!k) {
-    k = new MetalKernel({
-      name: "mlx_bun_fused_swiglu",
-      inputNames: ["x", "wg", "sg_w", "bg_w", "wu", "su_w", "bu_w"],
-      outputNames: ["hidden"],
-      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 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` +
-        `#define SWIGLU_F32 ${f32 ? 1 : 0}\n` +
-        `#define SWIGLU_TGZERO ${process.env.MLX_BUN_SWIGLU_TGZERO === "1" ? 1 : 0}\n` +
-        `#define SWIGLU_DEBUG_GATE ${process.env.MLX_BUN_SWIGLU_DEBUG_GATE === "1" ? 1 : 0}\n`,
-      ensureRowContiguous: true,
-    });
-    _kernels.set(key, k);
-  }
-  return k;
-}
-
-/** Each loader's (bits, groupSize) is baked independently, so gate and up may
- *  differ (the per-tensor mixed-precision layers). Requires bf16/f16 activations,
- *  affine biases, group size a multiple of BK, and H % BK == 0. */
-export function fusedSwigluSupported(
-  x: MlxArray, gateSpec: QuantSpec, upSpec: QuantSpec,
-  gateBias: MlxArray | null, upBias: MlxArray | null,
-): boolean {
-  const H = x.shape[x.shape.length - 1]!;
-  const specOk = (s: QuantSpec) =>
-    (s.bits === 2 || s.bits === 4 || s.bits === 8) && s.groupSize % BK === 0 && BK <= s.groupSize;
-  return (
-    (x.dtype === Dtype.bfloat16 || x.dtype === Dtype.float16) &&
-    gateBias != null && upBias != null &&
-    specOk(gateSpec) && specOk(upSpec) &&
-    H % BK === 0
-  );
-}
-
-/** silu(gate_proj(x)) · up_proj(x), fused. `x` is [..., H] (collapsed to [M, H]);
- *  wg/wu are the packed quantized weights [I, H*bits/32], with per-group
- *  scales/biases [I, H/groupSize]. gate and up may use different quant specs.
- *  Returns hidden [M, I] in x.dtype. Caller owns it. */
-export function fusedSwiglu(
-  x: MlxArray,
-  wg: MlxArray, sg: MlxArray, bg: MlxArray,
-  wu: MlxArray, su: MlxArray, bu: MlxArray,
-  gateSpec: QuantSpec, upSpec: QuantSpec = gateSpec,
-): MlxArray {
-  fusedSwigluCalls++;
-  const H = x.shape[x.shape.length - 1]!;
-  const M = x.shape.reduce((a, b) => a * b, 1) / H;
-  const I = wg.shape[0]!;
-  const x2 = x.shape.length === 2 ? x : ops.reshape(x, [M, H]);
-  try {
-    const [hidden] = getKernel(gateSpec.groupSize, gateSpec.bits, upSpec.groupSize, upSpec.bits, f32Internal()).apply(
-      [x2, wg, sg, bg, wu, su, bu],
-      {
-        outputShapeFn: (ins) => [{ shape: [ins[0]!.shape[0]!, I], dtype: x.dtype }],
-        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 },
-        templateDtypes: { T: x.dtype },
-      },
-    );
-    // Restore the leading dims the caller passed (e.g. [B, T, I]).
-    if (x.shape.length !== 2) {
-      const out = ops.reshape(hidden!, [...x.shape.slice(0, -1), I]);
-      hidden!.dispose();
-      return out;
-    }
-    return hidden!;
-  } finally {
-    if (x2 !== x && process.env.MLX_BUN_SWIGLU_NODISPOSE !== "1") x2.dispose();
-  }
-}