mlx-bun 0.0.8 → 0.0.10

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();
+  }
+}

package/src/model/gemma4-base.ts

@@ -87,6 +87,11 @@ export class QuantizedLinear {
     readonly scales: MlxArray,
     readonly biases: MlxArray | null,
     readonly spec: ops.QuantSpec,
+    /** ADDITIVE bias term (`.bias`, mlx nn.QuantizedLinear's optional bias —
+     *  qwen2 qkv, starcoder2, …). Distinct from `biases`, the quantization
+     *  zero-points. Applied after the matmul, before any LoRA residual,
+     *  exactly like mlx's QuantizedLinear.__call__. */
+    readonly bias: MlxArray | null = null,
   ) {}
 
   static load(weights: Weights, path: string, config: ModelConfig): QuantizedLinear {
@@ -99,6 +104,7 @@ export class QuantizedLinear {
       weights.tensor(`${path}.scales`),
       weights.has(`${path}.biases`) ? weights.tensor(`${path}.biases`) : null,
       spec,
+      weights.has(`${path}.bias`) ? weights.tensor(`${path}.bias`) : null,
     );
   }
 
@@ -113,6 +119,9 @@ export class QuantizedLinear {
 
   forward(x: MlxArray): MlxArray {
     let out = ops.quantizedMatmul(x, this.w, this.scales, this.biases, this.spec, true);
+    // Additive bias (mlx QuantizedLinear: `x = x + self["bias"]`), before
+    // the LoRA residual (LoRALinear calls the base linear bias-inclusive).
+    if (this.bias) out = disposing(out, ops.add(out, this.bias));
     // LoRA residual — composition is mlx-lm LoRALinear / optiq apply.py:
     //   y + (scale · ((x @ A) @ B)).astype(x.dtype)
     // (optiq mount.py omits the astype, leaking the f32 residual into the
@@ -213,8 +222,11 @@ export interface Cache {
   /** Can `trim(n)` drop the last n tokens? (Ring caches lose trimability
    *  once wrapped.) */
   isTrimmable(): boolean;
-  /** Drop the last n tokens (future writes overwrite them). */
-  trim(n: number): void;
+  /** Drop the last n tokens (future writes overwrite them). `bypass` (default
+   *  false → unchanged behavior) skips the isTrimmable guard for the
+   *  speculative-decode rollback case (trim the last n ≤ γ tips right after a
+   *  >1 concat write); rotating caches then physically shrink the buffer tail. */
+  trim(n: number, bypass?: boolean): void;
   dispose(): void;
 }
 
@@ -578,6 +590,15 @@ export class QuantizedKVCache implements Cache {
     this.offset = Math.max(0, this.offset - n);
   }
 
+  /** Adopt persisted state (takes ownership of the triples' arrays) —
+   *  the quantized twin of KVCache.restoreState (kv-store persistence). */
+  restoreState(keys: ops.QuantizedTensor, values: ops.QuantizedTensor, offset: number): void {
+    this.dispose();
+    this.keys = keys;
+    this.values = values;
+    this.offset = offset;
+  }
+
   dispose(): void {
     for (const a of this.state()) a.dispose();
     this.keys = this.values = null;
@@ -749,10 +770,32 @@ export class RotatingKVCache implements Cache {
     return this.offset < this.maxSize;
   }
 
-  trim(n: number): void {
+  trim(n: number, bypass = false): void {
     const k = Math.min(this.offset, n);
+    if (!bypass) {
+      // mlx-lm semantics: only valid pre-wrap (guarded by isTrimmable upstream).
+      this.offset -= k;
+      this.#idx -= k;
+      return;
+    }
+    // Spec-decode rollback after a >1 #updateConcat write: the buffer is in
+    // temporal order with #idx === bufferLen and the newest tokens are the last
+    // rows, so physically slice the last k off and keep #idx === new length.
+    // (Only sound right after a concat write with k ≤ γ — DSpark's only caller.)
     this.offset -= k;
-    this.#idx -= k;
+    if (k > 0 && this.keys && this.values) {
+      const [B, H, S, D] = this.keys.shape as [number, number, number, number];
+      const vD = this.values.shape[3]!;
+      const nk = this.keys.slice([0, 0, 0, 0], [B, H, S - k, D]);
+      const nv = this.values.slice([0, 0, 0, 0], [B, H, S - k, vD]);
+      this.keys.dispose();
+      this.values.dispose();
+      this.keys = nk;
+      this.values = nv;
+      this.#idx = nk.shape[2]!;
+    } else {
+      this.#idx -= k;
+    }
   }
 
   get ringIdx(): number {
@@ -1100,10 +1143,42 @@ export class RotatingQuantizedKVCache implements Cache {
     return this.offset < this.maxSize;
   }
 
-  trim(n: number): void {
+  trim(n: number, bypass = false): void {
     const k = Math.min(this.offset, n);
+    if (!bypass) {
+      this.offset -= k;
+      this.ringIdx -= k;
+      return;
+    }
+    // Spec-decode rollback after a >1 concat write — physically slice the last k
+    // off each (packed, scales, biases) component and keep ringIdx === seqLen.
     this.offset -= k;
-    this.ringIdx -= k;
+    if (k > 0 && this.keys && this.values) {
+      const S = this.#seqLen();
+      const cut = (a: MlxArray): MlxArray => {
+        const [B, H, , D] = a.shape as [number, number, number, number];
+        return a.slice([0, 0, 0, 0], [B, H, S - k, D]);
+      };
+      const nk = mapTriple(this.keys, cut);
+      const nv = mapTriple(this.values, cut);
+      disposeTriple(this.keys);
+      disposeTriple(this.values);
+      this.keys = nk;
+      this.values = nv;
+      this.ringIdx = this.#seqLen();
+    } else {
+      this.ringIdx -= k;
+    }
+  }
+
+  /** Adopt persisted state (takes ownership of the triples' arrays) —
+   *  ring order as-laid-out, ringIdx carried with it (kv-store persistence). */
+  restoreState(keys: ops.QuantizedTensor, values: ops.QuantizedTensor, offset: number, idx: number): void {
+    this.dispose();
+    this.keys = keys;
+    this.values = values;
+    this.offset = offset;
+    this.ringIdx = idx;
   }
 
   /** Oracle: to_quantized on an already-quantized rotating cache is
@@ -1516,8 +1591,22 @@ export function quantizedSdpa(
   )
     return fusedDecodeSdpa(q, kq, vq, groupSize, bits);
   const tile = q.shape[2]! > 1 || process.env.MLX_BUN_FUSED_DECODE === "1";
-  if (tile && fusedSdpaSupported(q, mask, groupSize, bits))
+  if (tile && fusedSdpaSupported(q, mask, groupSize, bits)) {
+    // The N-tiled loop CANNOT be traced: its tile count and slice extents
+    // bake at trace-time N, so under shapeless replay a growing (concat-
+    // phase) quantized cache's newest rows are silently never attended —
+    // reproduced on e4b (whole-graph form, FUSED_DECODE=1: diverges within
+    // 3 tokens, then repetition-loops). generate() refuses to compile when
+    // MLX_BUN_FUSED_DECODE=1; this throw is the backstop for any other
+    // compiled caller — it lands in the trace's transactional fallback
+    // instead of producing silently wrong attention.
+    if (isCompiledTrace())
+      throw new Error(
+        "quantizedSdpaTiled inside a compiled trace: the tile loop bakes trace-time N " +
+          "(growing KV rows would be silently dropped) — run this generation uncompiled",
+      );
     return quantizedSdpaTiled(q, kq, vq, scale, mask, groupSize, bits);
+  }
   return quantizedSdpaUnfused(q, kq, vq, scale, mask, groupSize, bits);
 }
 

package/src/model/gemma4.ts

@@ -295,7 +295,14 @@ class Attention {
     } else if (
       ta === "flash" && flashSupported(q) && (mask.mode === "causal" || mask.mode === "array")
     ) {
-      // Training: ops.sdpa's dK vjp is wrong, so use the validated flash kernel.
+      // Training, OPT-IN flash path (MLX_BUN_TRAIN_ATTN=flash) for O(L)-memory
+      // long context. NOTE the default (ops.sdpa, below) is the CORRECT and
+      // verified route: its dQ/dK/dV vjps are exact (finite-difference checked;
+      // it's what mlx-lm's tuner differentiates). The two historical bugs were
+      // in THIS flash kernel (spurious dK transpose, divergent dQ barrier),
+      // both fixed and gated by tests/flash-attention.test.ts vs the ops.sdpa
+      // oracle — see flash-attention.ts's routing comment. Flash also remains
+      // Gemma-guarded in the trainer (e4b >=2K SIGTRAP not yet re-validated).
       // Full layers → window 0 (pure causal); sliding layers → their window.
       const window = this.isSliding ? this.windowSize : 0;
       attn = flashAttention(q, shared.keys, shared.values, 1.0, true, window);
@@ -663,6 +670,13 @@ export class Gemma4Model {
   /** Set by the trainer to enable gradient checkpointing in forwardLayers
    *  (null = off; the non-checkpointed path is untouched). */
   gradCkpt: GradCheckpointCtx | null = null;
+  /** Optional per-forward hidden-state tap for DSpark's DFlash H_ctx (paper
+   *  Eq 2). null = off → byte-identical forward. `layers`: residual-output
+   *  indices to capture (0..nLayers-1 = post-layer residual; nLayers =
+   *  post-finalNorm). `pos`: if set, capture only [1,1,H] at that seq position
+   *  (live decode); else the full [1,L,H] (offline regen). The caller owns and
+   *  disposes the captured arrays. Additive/read-only → parity untouched. */
+  hiddenTap: { layers: Set<number>; pos?: number; captured: Map<number, MlxArray> } | null = null;
   // e2b/e4b per-layer input machinery
   readonly perLayerEmbed: QuantizedEmbedding | null;
   readonly perLayerModelProjection: QuantizedLinear | null;
@@ -820,6 +834,20 @@ export class Gemma4Model {
     return out;
   }
 
+  /** Store layer `i`'s residual stream `h` ([1,L,H]) if the hidden tap requests
+   *  it. Copies (astype bf16, optional slice) so the loop's `h.dispose()` can't
+   *  free the capture. No-op and no graph ops when the tap is null → the forward
+   *  is byte-identical to an untapped run. (DSpark DFlash H_ctx.) */
+  protected captureLayer(i: number, h: MlxArray): void {
+    const tap = this.hiddenTap;
+    if (!tap || !tap.layers.has(i)) return;
+    const H = h.shape[2]!;
+    const src = tap.pos !== undefined ? h.slice([0, tap.pos, 0], [1, tap.pos + 1, H]) : h;
+    const copy = src.astype(this.embed.scales.dtype); // bf16, matches regen dump dtype
+    if (src !== h) src.dispose();
+    tap.captured.set(i, copy);
+  }
+
   /** Consumes h. */
   protected forwardLayers(
     h0: MlxArray, cache: Cache[], bidir: MlxArray | null, ids: MlxArray | null,
@@ -878,6 +906,7 @@ export class Gemma4Model {
       pls?.dispose();
       h.dispose();
       h = next;
+      this.captureLayer(i, h); // DFlash H_ctx tap (no-op unless hiddenTap set)
       intermediates[i] = shared;
     }
     perLayer?.dispose();
@@ -895,7 +924,9 @@ export class Gemma4Model {
     }
     for (const m of masks.values()) m.arr?.dispose();
 
-    return disposing(h, this.finalNorm.forward(h));
+    h = disposing(h, this.finalNorm.forward(h));
+    this.captureLayer(this.layers.length, h); // post-finalNorm sentinel (index = nLayers)
+    return h;
   }
 
   // --- Segmented-backward support (docs/design/segmented-backward-training.md

package/src/model/minicpm5.ts

@@ -4,6 +4,7 @@
 import type { ModelConfig } from "../config";
 import type { Weights } from "../weights";
 import { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
 import {
   argmaxLastPosition,
@@ -16,9 +17,56 @@ import {
   QuantizedLinear,
   quantizedSdpa,
   RMSNorm,
+  isCompiledTrace,
   type Cache,
   type Mask,
 } from "./gemma4-base";
+import {
+  fusedSwiglu,
+  fusedSwigluEnabled,
+  fusedSwigluSupported,
+  swigluMMin,
+  swigluUnderCap,
+} from "./fused-swiglu-kernel";
+import { steelLinear } from "./steel-linear-kernel";
+import { fusedMlp } from "./fused-mlp-kernel";
+
+/** Break the graph at a boundary: contiguous → eval → fresh graph-free leaf — the
+ *  exact primitive the segmented TRAINER uses so MLX never plans across boundaries
+ *  and prior-segment activations are freed (copy of train/segmented.ts detachLeaf). */
+function detachLeaf(a: MlxArray): MlxArray {
+  const c = ops.contiguous(a);
+  c.eval();
+  const leaf = MlxArray.fromBytesCopy(c.rawBytes(), c.shape, c.dtype);
+  c.dispose();
+  a.dispose();
+  return leaf;
+}
+
+/** A quantized projection as OUR steel dispatch when the controlled path is live
+ *  (no adapter, affine, clean tiling, M>1 so it's not a decode GEMV); otherwise
+ *  the normal linear. x is [B,L,H] or [M,H]. */
+function ourLinear(lin: QuantizedLinear, x: MlxArray): MlxArray {
+  const H = x.shape[x.shape.length - 1]!;
+  const M = x.shape.reduce((a, b) => a * b, 1) / H;
+  if (
+    !fusedSwigluEnabled() || isCompiledTrace() || M < swigluMMin() ||
+    lin.adapters != null || lin.biases == null ||
+    (x.dtype !== Dtype.bfloat16 && x.dtype !== Dtype.float16) ||
+    H % 32 !== 0 || lin.spec.groupSize % 32 !== 0
+  ) {
+    return lin.forward(x);
+  }
+  const x2 = x.shape.length === 2 ? x : ops.reshape(x, [M, H]);
+  const y = steelLinear(x2, lin.w, lin.scales, lin.biases, lin.spec);
+  if (x2 !== x) x2.dispose();
+  return x.shape.length === 2
+    ? y
+    : disposing(y, ops.reshape(y, [...x.shape.slice(0, -1), lin.w.shape[0]!]));
+}
+
+/** How often the inference fused path materializes the residual stream (layers). */
+const EVAL_EVERY = Math.max(1, Number(process.env.MLX_BUN_SWIGLU_EVAL_EVERY) || 1);
 import { flashAttention, getTrainingAttn, flashSupported } from "./flash-attention";
 
 // Shared prompt-prefix plan (lever 7). When set, the attention ropes the
@@ -156,7 +204,40 @@ class LlamaMLP {
     this.down = QuantizedLinear.load(weights, `${prefix}.down_proj`, config);
   }
 
-  forward(x: MlxArray): MlxArray {
+  /** MLP. When `residual` (the layer input h) is given AND the fused path applies,
+   *  the kernel computes h + mlp(x) directly — the residual is ABSORBED, so nothing
+   *  stays live across the kernel for MLX's planner to alias (it just runs our one
+   *  dispatch). Otherwise returns mlp(x) and the caller adds the residual. */
+  forward(x: MlxArray, residual?: MlxArray): MlxArray {
+    const H = x.shape[x.shape.length - 1]!;
+    const M = x.shape.reduce((a, b) => a * b, 1) / H;
+    if (
+      residual != null &&
+      fusedSwigluEnabled() && !isCompiledTrace() && M >= swigluMMin() && swigluUnderCap() &&
+      this.gate.adapters == null && this.up.adapters == null && this.down.adapters == null &&
+      this.down.biases != null && this.down.spec.groupSize % 32 === 0 &&
+      fusedSwigluSupported(x, this.gate.spec, this.up.spec, this.gate.biases, this.up.biases)
+    ) {
+      // The WHOLE MLP + residual in ONE kernel — gate+up+silu+down then +h, `hidden`
+      // never materialized, residual absorbed. MLX runs one dispatch; nothing for
+      // its planner to alias → chains across layers with no boundary.
+      const Hin = this.gate.scales.shape[1]! * this.gate.spec.groupSize;
+      const Ho = this.down.w.shape[0]!;
+      const x2 = x.shape.length === 2 ? x : ops.reshape(x, [M, Hin]);
+      const h2 = residual.shape.length === 2 ? residual : ops.reshape(residual, [M, Ho]);
+      const out2 = fusedMlp(
+        x2,
+        this.gate.w, this.gate.scales, this.gate.biases!, this.gate.spec,
+        this.up.w, this.up.scales, this.up.biases!, this.up.spec,
+        this.down.w, this.down.scales, this.down.biases!, this.down.spec,
+        h2,
+      );
+      if (x2 !== x) x2.dispose();
+      if (h2 !== residual) h2.dispose();
+      return x.shape.length === 2
+        ? out2
+        : disposing(out2, ops.reshape(out2, [...x.shape.slice(0, -1), Ho]));
+    }
     const gate = this.gate.forward(x);
     const up = this.up.forward(x);
     const sig = ops.sigmoid(gate);
@@ -166,9 +247,10 @@ class LlamaMLP {
     const hidden = ops.mul(silu, up);
     silu.dispose();
     up.dispose();
-    const out = this.down.forward(hidden);
+    const m = this.down.forward(hidden);
     hidden.dispose();
-    return out;
+    if (residual == null) return m;
+    return disposing(m, ops.add(residual, m));
   }
 }
 
@@ -192,11 +274,10 @@ class LlamaLayer {
     const h = ops.add(x, a);
     a.dispose();
     const hn = this.postAttnNorm.forward(h);
-    const m = this.mlp.forward(hn);
+    // Pass the residual h INTO the MLP so the fused kernel absorbs it (out = h + mlp).
+    const out = this.mlp.forward(hn, h);
     hn.dispose();
-    const out = ops.add(h, m);
     h.dispose();
-    m.dispose();
     return out;
   }
 }
@@ -256,7 +337,15 @@ export class MiniCPM5Model {
   }
 
   protected forwardLayers(h0: MlxArray, cache: Cache[]): MlxArray {
-    const h = this.runLayerRange(h0, 0, this.layers.length, cache);
+    // When the fused SwiGLU path is live, WE own the eval schedule: materialize
+    // each layer's output so the custom kernel sits in a shallow per-layer graph
+    // rather than a 24-deep lazy graph whose buffer planner aliases the kernel's
+    // output (the deep-graph corruption). Training/segmented-backward call
+    // runLayerRange directly and never set this, so autograd is untouched.
+    const h = this.runLayerRange(
+      h0, 0, this.layers.length, cache,
+      fusedSwigluEnabled() && process.env.MLX_BUN_SWIGLU_NOEVAL !== "1",
+    );
     h0.dispose(); // runLayerRange keeps its input alive; we own h0 here
     return disposing(h, this.finalNorm.forward(h));
   }
@@ -269,7 +358,7 @@ export class MiniCPM5Model {
    *  segmented-backward training path (docs/design/segmented-backward-training.md).
    *  MiniCPM5 is plain full-attention with no KV-sharing, so each layer reads
    *  its own stateless `cache[i]` and a single causal mask. */
-  runLayerRange(h: MlxArray, aIdx: number, bIdx: number, cache: Cache[]): MlxArray {
+  runLayerRange(h: MlxArray, aIdx: number, bIdx: number, cache: Cache[], evalPerLayer = false): MlxArray {
     if (aIdx >= bIdx) return h;
     const L = h.shape[1]!;
     const mask = cache[aIdx]!.makeMask(L, null);
@@ -278,6 +367,15 @@ export class MiniCPM5Model {
       const next = this.layers[i]!.forward(cur, mask, cache[i]!);
       if (i > aIdx) cur.dispose(); // keep the input boundary leaf; drop interiors
       cur = next;
+      // Per-layer boundary (segmented-trainer pattern): materialize each layer's
+      // output so MLX never plans across layers — our single-dispatch fused MLP is
+      // the only custom kernel it runs per layer. evalAll (GPU sync, no host copy)
+      // suffices for one kernel/boundary; MLX_BUN_SWIGLU_DETACH=1 forces the
+      // stronger graph-breaking detachLeaf copy.
+      if (evalPerLayer) {
+        if (process.env.MLX_BUN_SWIGLU_DETACH === "1") cur = detachLeaf(cur);
+        else ops.evalAll([cur]);
+      }
     }
     mask.arr?.dispose();
     return cur;