mlx-bun 0.0.8 → 0.0.9

package/src/mlx/ffi.ts

@@ -119,6 +119,9 @@ export const C = dlopen(LIBMLXC_PATH, {
   mlx_divide: { args: [P, u64, u64, u64], returns: i32 },
   mlx_tanh: { args: [P, u64, u64], returns: i32 },
   mlx_sigmoid: { args: [P, u64, u64], returns: i32 },
+  mlx_erf: { args: [P, u64, u64], returns: i32 },
+  // (res, c, a, b, alpha, beta, stream) — alpha·(a@b) + beta·c (nn.Linear bias path)
+  mlx_addmm: { args: [P, u64, u64, u64, f32, f32, u64], returns: i32 },
   mlx_power: { args: [P, u64, u64, u64], returns: i32 },
   mlx_take_axis: { args: [P, u64, u64, i32, u64], returns: i32 },
   mlx_reshape: { args: [P, u64, P, u64, u64], returns: i32 },

package/src/mlx/ops.ts

@@ -111,6 +111,23 @@ export function rope(
   );
 }
 
+/** fast::rope with every knob exposed (traditional + scale), for the
+ *  universal rope factory (rope_utils.initialize_rope port). The plain
+ *  `rope` above stays byte-identical for existing call sites. */
+export function ropeScaled(
+  x: MlxArray, dims: number, traditional: boolean, base: number | null,
+  scale: number, offset: number, freqs: MlxArray | null, s: S = gpuStream,
+): MlxArray {
+  return new MlxArray(
+    outArray("fast_rope", (o) =>
+      C.mlx_fast_rope(
+        o, x.handle, dims, traditional, optFloat(base), scale, offset,
+        freqs?.handle ?? 0n, s,
+      ),
+    ),
+  );
+}
+
 /** fast::rope with the position offset as an ARRAY (int32 scalar) instead
  *  of a baked int — required inside compiled decode graphs, where the
  *  offset changes every step but the graph must not. Same kernel as
@@ -267,6 +284,35 @@ export function geluApprox(x: MlxArray, s: S = gpuStream): MlxArray {
   return out;
 }
 
+export function erf(a: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("erf", (o) => C.mlx_erf(o, a.handle, s)));
+}
+
+/** nn.gelu (PRECISE, erf-based): x · (1 + erf(x/√2)) / 2 — composed exactly
+ *  as mlx's python source (weak scalars promote to x's dtype). Used by the
+ *  gemma-1 gated MLP and starcoder2's plain MLP; distinct from geluApprox
+ *  (gelu_pytorch_tanh). */
+export function geluPrecise(x: MlxArray, s: S = gpuStream): MlxArray {
+  const sqrt2 = scalarLike(Math.sqrt(2), x);
+  const xs = div(x, sqrt2, s); // x / math.sqrt(2) — division, not mul-by-reciprocal
+  const e = erf(xs, s);
+  const one = scalarLike(1, x);
+  const e1 = add(one, e, s);
+  const xe = mul(x, e1, s);
+  const two = scalarLike(2, x);
+  const out = div(xe, two, s); // (...) / 2
+  for (const a of [sqrt2, xs, e, one, e1, xe, two]) a.dispose();
+  return out;
+}
+
+/** mx.addmm(c, a, b): a@b + c fused (α=β=1) — nn.Linear's bias path
+ *  (`mx.addmm(bias, x, weight.T)`). */
+export function addmm(c: MlxArray, a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(
+    outArray("addmm", (o) => C.mlx_addmm(o, c.handle, a.handle, b.handle, 1.0, 1.0, s)),
+  );
+}
+
 export function exp(a: MlxArray, s: S = gpuStream): MlxArray {
   return new MlxArray(outArray("exp", (o) => C.mlx_exp(o, a.handle, s)));
 }

package/src/model/compiled-decode.ts

@@ -590,6 +590,29 @@ export class CompiledDecode {
     const evalWith: MlxArray[] = [];
     let carried = cur; // seg 0 consumes the token array; later, hidden state
     let carriedOwned = false; // cur belongs to the caller
+    // Mid-step failure safety (a later segment's closure.apply can throw at
+    // first trace AFTER earlier layers already advanced): the step must be
+    // TRANSACTIONAL, or the caller's uncompiled re-forward of the same token
+    // double-writes the committed layers (+1 offset skew forever, silently
+    // corrupted attention). Two mechanisms, chosen by what each phase allows:
+    //  - ring segments: adoptDecodeStep is a pure host-side handle swap with
+    //    no reader between here and step end (each cache is read only by its
+    //    own layer, inside the closure that produced the update), so the
+    //    adopts are STAGED and committed only once every segment and js layer
+    //    has succeeded. On failure the staged buffers are disposed and the
+    //    ring caches are exactly at their pre-step state (prepareDecodeStep's
+    //    growth/trim/rotate bookkeeping is benign for the uncompiled path —
+    //    the same invariant #stepWhole and generate.ts's fallback rely on).
+    //  - js (concat) layers: layer.forward commits via the cache's real
+    //    updateAndFetch mid-loop, so on failure the commit is ROLLED BACK by
+    //    reverting offset (and ring index) — trim(1) semantics. The written
+    //    KV row beyond the reverted offset is dead data the re-forward
+    //    overwrites at the same position (updateAndFetch writes at offset/
+    //    ringIdx), so post-recovery state equals a single clean write. This
+    //    holds for the post-wrap quantized-under-perf-kernel case too: the
+    //    in-place S=1 write slot is re-derived from the reverted ringIdx.
+    const stagedAdopts: { c: AnyCache; bufs: MlxArray[] }[] = [];
+    const committedJs: { c: AnyCache; offsetBefore: number }[] = [];
     try {
       for (let s = 0; s < segs.length; s++) {
         const { from, to } = segs[s]!;
@@ -636,17 +659,14 @@ export class CompiledDecode {
           carried = outs[0]!;
           carriedOwned = true;
           // ring updates are ancestors of the carried hidden state (the
-          // in-graph fetch reads the updated buffer): no eval roots
+          // in-graph fetch reads the updated buffer): no eval roots. The
+          // host-side adoption is deferred to the end of the step (see above).
           let oi = 1;
           for (let i = from; i < to; i++) {
             const c = anyCaches[i]!;
-            if (c instanceof RotatingQuantizedKVCache) {
-              c.adoptDecodeStep(outs.slice(oi, oi + 6));
-              oi += 6;
-            } else {
-              (c as RotatingKVCache).adoptDecodeStep(outs[oi]!, outs[oi + 1]!);
-              oi += 2;
-            }
+            const width = c instanceof RotatingQuantizedKVCache ? 6 : 2;
+            stagedAdopts.push({ c, bufs: outs.slice(oi, oi + width) });
+            oi += width;
           }
         }
         if (s < jsLayers.length) {
@@ -655,6 +675,7 @@ export class CompiledDecode {
           const c = anyCaches[li]!;
           const window = layer.layerType === "sliding_attention" ? this.model.windowSize : null;
           const mask = c.makeMask(1, window); // N=1 → mode ""
+          committedJs.push({ c, offsetBefore: c.offset });
           const { h: next, shared } = layer.forward(carried, mask, c, null, null);
           if (carriedOwned) carried.dispose();
           carried = next;
@@ -663,6 +684,20 @@ export class CompiledDecode {
           mask.arr?.dispose();
         }
       }
+      // every segment and js layer succeeded — commit the staged ring adopts
+      for (const { c, bufs } of stagedAdopts) {
+        if (c instanceof RotatingQuantizedKVCache) c.adoptDecodeStep(bufs);
+        else (c as RotatingKVCache).adoptDecodeStep(bufs[0]!, bufs[1]!);
+      }
+    } catch (e) {
+      // unwind: free staged (never-adopted) ring buffers and the owned hidden
+      // state, and revert the js layers' committed writes, so the caller's
+      // uncompiled re-forward of this token starts from the pre-step state.
+      for (const { bufs } of stagedAdopts) for (const b of bufs) b.dispose();
+      for (const { c, offsetBefore } of committedJs)
+        if (c.offset > offsetBefore) c.trim(c.offset - offsetBefore);
+      if (carriedOwned) carried.dispose();
+      throw e;
     } finally {
       for (const t of stepTemps) t.dispose();
     }

package/src/model/factory.ts

@@ -11,12 +11,21 @@ import { configFingerprint } from "./fingerprint";
 import { GENERATED } from "./generated";
 import { MiniCPM5Model } from "./minicpm5";
 import { Qwen35Model } from "./qwen3_5";
+import { Qwen3Model } from "./qwen3";
 import { DiffusionGemmaModel } from "./diffusion-gemma";
-import { isDiffusionGemmaConfig, isMiniCPM5Config, isQwen35Config } from "./support";
+import { UniversalDenseModel } from "./universal/dense";
+import { genericArgsFor, GENERIC_MODEL_TYPES, remapModelType } from "./universal/archs";
+import { isDiffusionGemmaConfig, isMiniCPM5Config, isQwen35Config, isQwen3Config } from "./support";
 
-export type RuntimeModel = Gemma4Model | MiniCPM5Model | Qwen35Model | DiffusionGemmaModel;
+export type RuntimeModel =
+  | Gemma4Model | MiniCPM5Model | Qwen35Model | Qwen3Model | DiffusionGemmaModel
+  | UniversalDenseModel;
 
+/** Dispatch ladder (docs/design/generic-model-support.md §3.3):
+ *  dedicated → generated/monolith (gemma4*) → generic (Tier-0 universal)
+ *  → reject with a helpful error. Generic never shadows a dedicated port. */
 export function createModel(weights: Weights, config: ModelConfig): RuntimeModel {
+  // 1. Dedicated classes.
   // DiffusionGemma is non-autoregressive — generate() detects it and routes to
   // the denoising engine instead of the AR decode loop. It exposes the shared
   // RuntimeModel surface (config/weightsBytes/loraState/makeCache) with the
@@ -24,10 +33,22 @@ export function createModel(weights: Weights, config: ModelConfig): RuntimeModel
   if (isDiffusionGemmaConfig(config)) return new DiffusionGemmaModel(weights, config);
   if (isMiniCPM5Config(config)) return new MiniCPM5Model(weights, config);
   if (isQwen35Config(config)) return new Qwen35Model(weights, config);
-  if (config.modelType === "llama")
-    throw new Error("unsupported llama config: only MiniCPM5-1B-OptiQ-4bit is wired");
-  if (config.modelType.startsWith("qwen3_5"))
-    throw new Error(`unsupported qwen3_5 config (MoE variants are deferred)`);
-  const cls = GENERATED.get(configFingerprint(config)) ?? Gemma4Model;
-  return new cls(weights, config);
+  if (isQwen3Config(config)) return new Qwen3Model(weights, config);
+  // 2. gemma4*: fingerprint-matched generated specialization, else monolith.
+  if (config.modelType.startsWith("gemma4")) {
+    const cls = GENERATED.get(configFingerprint(config)) ?? Gemma4Model;
+    return new cls(weights, config);
+  }
+  // 3. Generic Tier-0 fallback: the universal-dense descriptor table
+  // (replaces the old hard "llama = MiniCPM5 only" / unknown-type throws).
+  const generic = genericArgsFor(config);
+  if (generic) return new UniversalDenseModel(weights, config, generic);
+  // 4. Reject, naming the arch (post-remap) and the declared surface.
+  const arch = remapModelType(config.modelType);
+  throw new Error(
+    `unsupported model_type "${config.modelType}"` +
+    (arch !== config.modelType ? ` (mlx-lm remaps it to "${arch}")` : "") +
+    ` — targeted: gemma4*, diffusion_gemma, qwen3_5, qwen3, MiniCPM5;` +
+    ` generic (Tier-0): ${[...GENERIC_MODEL_TYPES].sort().join(", ")}`,
+  );
 }

package/src/model/flash-attention.ts

@@ -649,8 +649,10 @@ export function flashSupported(q: MlxArray): boolean {
 }
 
 /** Materialized differentiable attention (matmul → mask → softmax → matmul),
- *  GQA-aware. The correct-vjp fallback for masks flash doesn't take (sliding
- *  window) and for any case where ops.sdpa's broken dK can't be used. O(L²) —
+ *  GQA-aware. A correct-vjp fallback for shapes/masks the flash kernel doesn't
+ *  take. (ops.sdpa's dK vjp is CORRECT — FD-verified, the training default;
+ *  the old "broken dK" note here predated the flash-kernel root-cause: the
+ *  bugs were flash's own dK transpose + dQ barrier, both fixed.) O(L²) —
  *  pair with gradient checkpointing. */
 export function manualSdpa(q: MlxArray, k: MlxArray, v: MlxArray, scale: number, mask: Mask): MlxArray {
   const [B, Hq, T, Dh] = q.shape as [number, number, number, number];

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

@@ -26,13 +26,20 @@ 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).
- *  It matches optiq's fused-KV decode — which the quantized-KV parity goldens
- *  track — so it's the production default. It is NOT bit-exact vs the -O0
- *  reference (online softmax, tier-b), so it's the PERF side of the
- *  compat-vs-perf A/B: 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.) */
+/** 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";
 }

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

@@ -0,0 +1,169 @@
+// 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

@@ -0,0 +1,240 @@
+// 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();
+  }
+}