mlx-bun 0.0.9 → 0.0.11

package/src/model/gemma4.ts

@@ -45,15 +45,77 @@ import {
   RMSNorm,
   RotatingKVCache,
   RotatingQuantizedKVCache,
+  TurboQuantKVCache,
   quantizedSdpa,
   type Cache,
   type Mask,
   type SharedKv,
   type LoraWeights,
 } from "./gemma4-base";
-import { fusedGeglu, fusedGegluDifferentiable, fusedGeluEnabled, fusedGeluSupported, fusedGeluTraining } from "./fused-geglu-kernel";
 import { Checkpoint } from "../mlx/checkpoint";
 import { flashAttention, getTrainingAttn, flashSupported } from "./flash-attention";
+import { unrotateValues as tqUnrotateValues } from "../mlx/turboquant-ops";
+import { CompiledFunction } from "../mlx/compile";
+import { flagOn } from "../flags";
+
+/** Verbatim port of mlx_lm/models/gemma4_text.py:
+ *    `@partial(mx.compile, shapeless=True) def geglu(gate, x): return nn.gelu_approx(gate) * x`
+ *  The oracle @mx.compile's this and uses it for EVERY MLP shape (dense MLP
+ *  __call__ line 114, and the SwitchGLU GeGLU activation line 149-150). Our
+ *  default MLP runs the SAME math UNFUSED (standalone ops.geluApprox + ops.mul),
+ *  a different dispatched kernel set + the per-op host tax the MiniCPM5 swiglu
+ *  fix measured at ~5%. This closure fuses geluApprox's ~9 element-wise ops + the
+ *  final mul into one traced-once/replayed graph, matching the oracle's kernel
+ *  set. `geluApprox` is composed op-for-op from the same primitives as the
+ *  oracle's nn.gelu_approx (0.5·x·(1+tanh(√(2/π)(x+0.044715·x³)))), so the
+ *  compiled graph is the oracle's geglu graph. Autograd-safe (verified: the vjp
+ *  through this closure is bit-identical to the plain composition). */
+/** True when gemma's geglu activation runs through the compiled closure (matching
+ *  the oracle's @mx.compile geglu) instead of the spelled-out ops.geluApprox+ops.mul.
+ *  DEFAULT ON: compiled geglu is bit-exact to mlx-lm's @mx.compile geglu (same
+ *  libmlx graph → same kernel) AND dispatches one fused kernel set instead of the
+ *  ~9-op path. `MLX_BUN_COMPILED_GEGLU=0` (--compiled-activations off) selects the
+ *  uncompiled composition — same L1 parity, slower (the A/B opt-out). */
+export function compiledGegluActive(): boolean {
+  return flagOn("MLX_BUN_COMPILED_GEGLU", true);
+}
+
+let _gegluClosure: CompiledFunction | null = null;
+export function compiledGeglu(gate: MlxArray, up: MlxArray): MlxArray {
+  if (!_gegluClosure) {
+    _gegluClosure = new CompiledFunction((inputs) => {
+      const g = inputs[0]!, u = inputs[1]!; // nn.gelu_approx(gate) * x
+      const act = ops.geluApprox(g);
+      const out = ops.mul(act, u); act.dispose();
+      return [out];
+    });
+  }
+  return _gegluClosure.apply([gate, up])[0]!;
+}
+
+/** gemma4_text.py `@partial(mx.compile) def logit_softcap(softcap, x):
+ *    return mx.tanh(x / softcap) * softcap`
+ *  The oracle fuses divide+tanh+multiply into ONE kernel; our default
+ *  `logitSoftcap` runs them UNFUSED (standalone div + tanh + mul → a separate
+ *  `vn_Tanh` + `vsn_Divide`). Keyed by cap so the softcap is baked as a graph
+ *  constant (matching the oracle's weak-scalar → the `…Divide…Tanh…Multiply…_VC_`
+ *  kernel). The default (matches mlx-lm's dispatched kernel set). */
+const _softcapClosures = new Map<number, CompiledFunction>();
+export function compiledLogitSoftcap(x: MlxArray, cap: number): MlxArray {
+  let fn = _softcapClosures.get(cap);
+  if (!fn) {
+    fn = new CompiledFunction((inputs) => {
+      const xx = inputs[0]!;                    // tanh(x / softcap) * softcap
+      const capArr = ops.scalarLike(cap, xx);
+      const scaled = ops.div(xx, capArr);
+      const t = ops.tanh(scaled); scaled.dispose();
+      const out = ops.mul(t, capArr); t.dispose(); capArr.dispose();
+      return [out];
+    });
+    _softcapClosures.set(cap, fn);
+  }
+  return fn.apply([x])[0]!;
+}
 
 /** Optional gradient-checkpointing context for the training forward. When set,
  *  forwardLayers wraps each layer in a Checkpoint so its interior activations
@@ -275,6 +337,14 @@ class Attention {
           groupSize: cache.groupSize, bits: cache.bits,
           offsetArr,
         };
+      } else if (cache instanceof TurboQuantKVCache) {
+        // Deferred-V read: values stay rotated; the attention output is
+        // un-rotated below (shared.vRotated) — including by KV-shared
+        // consumer layers, which see the flag through sharedIn.
+        const [keys, values] = cache.updateAndFetchDeferredV(kRoped, vT);
+        kRoped.dispose();
+        vT.dispose();
+        shared = { kind: "plain", keys, values, offset, offsetArr, vRotated: true };
       } else {
         const [keys, values] = cache.updateAndFetch(kRoped, vT);
         kRoped.dispose();
@@ -310,6 +380,11 @@ class Attention {
       attn = ops.sdpa(q, shared.keys, shared.values, 1.0, mask.mode, mask.arr);
     }
     q.dispose();
+    if (shared.kind === "plain" && shared.vRotated) {
+      // Attention is linear in V: one InvFWHT on the [B,H,L,D] output
+      // replaces one per cached token (TurboQuant deferred-V).
+      attn = disposing(attn, tqUnrotateValues(attn));
+    }
     const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
     attn.dispose();
     const merged = ops.reshape(attnT, [B, L, -1]);
@@ -334,14 +409,11 @@ class MLP {
   forward(x: MlxArray): MlxArray {
     const g = this.gate.forward(x);
     const u = this.up.forward(x);
-    // Fused GeGLU perf kernel (opt-in, not in a compiled trace): one pass
-    // instead of ~9 element-wise kernels. Off → the bit-exact spelled path.
-    // In training mode the differentiable wrapper (kernel forward + hand-derived
-    // vjp) is used so autograd can flow AND the backward recomputes the gelu
-    // from the primal instead of retaining the spelled-out intermediates.
     let h: MlxArray;
-    if (fusedGeluEnabled() && !isCompiledTrace() && fusedGeluSupported(g)) {
-      h = fusedGeluTraining() ? fusedGegluDifferentiable(g, u) : fusedGeglu(g, u);
+    if (compiledGegluActive() && !isCompiledTrace()) {
+      // Oracle geglu(gate_proj(x), up_proj(x)): @mx.compile'd gelu_approx·mul in
+      // one kernel set (the default; = mlx-lm). Autograd-safe → covers training.
+      h = compiledGeglu(g, u);
       g.dispose();
       u.dispose();
     } else {
@@ -464,12 +536,21 @@ class SwitchGLU {
     const xUp = this.up.forward(h, idx, doSort);
     const xGate = this.gate.forward(h, idx, doSort);
     h.dispose();
-    // GeGLU activation: gelu_approx(gate) · up (same composition as MLP)
-    const act = ops.geluApprox(xGate);
-    xGate.dispose();
-    const mid = ops.mul(act, xUp);
-    act.dispose();
-    xUp.dispose();
+    // GeGLU activation: gelu_approx(gate) · up (same composition as MLP). The
+    // oracle SwitchGLU uses the SAME @mx.compile'd geglu (activation=GeGLU(),
+    // switch_layers.py) → the compiled path routes through the shared closure.
+    let mid: MlxArray;
+    if (compiledGegluActive() && !isCompiledTrace()) {
+      mid = compiledGeglu(xGate, xUp);
+      xGate.dispose();
+      xUp.dispose();
+    } else {
+      const act = ops.geluApprox(xGate);
+      xGate.dispose();
+      mid = ops.mul(act, xUp);
+      act.dispose();
+      xUp.dispose();
+    }
     let y = this.down.forward(mid, idx, doSort);
     mid.dispose();
     if (idx !== indices) idx.dispose();
@@ -625,17 +706,12 @@ export class DecoderLayer {
     f.dispose();
 
     // per-layer input gating (reference gemma4_text DecoderLayer). gelu(gate)·pli
-    // is the second GeGLU site — fused (differentiable) on the training path, the
-    // bit-exact spelled-out gelu+mul on the inference path (unchanged).
+    // is the second GeGLU site — the bit-exact spelled-out gelu+mul.
     if (this.perLayerGate && this.perLayerProjection && this.postPerLayerNorm && perLayerInput) {
       const res2 = h;
       let gate = this.perLayerGate.forward(h);
-      if (fusedGeluTraining() && fusedGeluEnabled() && !isCompiledTrace() && fusedGeluSupported(gate)) {
-        gate = disposing(gate, fusedGegluDifferentiable(gate, perLayerInput));
-      } else {
-        gate = disposing(gate, ops.geluApprox(gate));
-        gate = disposing(gate, ops.mul(gate, perLayerInput));
-      }
+      gate = disposing(gate, ops.geluApprox(gate));
+      gate = disposing(gate, ops.mul(gate, perLayerInput));
       gate = disposing(gate, this.perLayerProjection.forward(gate));
       gate = disposing(gate, this.postPerLayerNorm.forward(gate));
       h = ops.add(res2, gate);
@@ -805,25 +881,35 @@ export class Gemma4Model {
   }
 
   /** Pre-merged (unscaled) input embeddings → hidden states. Used by the
-   *  vision path; `bidir` (bool [L]) marks image tokens, which attend
-   *  bidirectionally among themselves (use_bidirectional_attention:
-   *  "vision" — text stays causal). `ids` ([1, L], the spliced token ids)
-   *  is required for per-layer-input models (e2b/e4b): image-token
-   *  positions get token 0's per-layer embedding, and their vision content
+   *  multimodal (vision/audio) path; `bidir` (bool [L]) marks image tokens,
+   *  which attend bidirectionally among themselves (use_bidirectional_attention:
+   *  "vision" — text stays causal; audio prompts pass NO bidir mask, and a
+   *  mixed image+audio prompt drops the image overlay too, §3.3 Q1 of
+   *  docs/design/audio-input-plan.md). `ids` ([1, L], the spliced token ids)
+   *  is required for per-layer-input models (e2b/e4b): multimodal soft-token
+   *  positions get token 0's per-layer embedding, and their tower content
    *  enters only through the projection term (the merged hidden). Matches
-   *  optiq's frontend: `zeroed = where(text_mask, input_ids, 0)`. */
+   *  optiq's frontend: `zeroed = where(text_mask, input_ids, 0)`.
+   *
+   *  `multimodal` (bool [L]) is the per-layer-input ZEROING mask — the UNION
+   *  of multimodal soft tokens (image | audio, §3.3 Q2), decoupled from the
+   *  attention overlay: an audio-only prompt has no bidir mask but its audio
+   *  positions must still be zeroed. When null (the legacy vision call
+   *  shape), the zeroing mask falls back to `bidir` — vision passes the same
+   *  image mask for both roles, so that path is unchanged. */
   forwardEmbeddings(
     embeds: MlxArray, cache: Cache[], bidir: MlxArray | null,
-    ids: MlxArray | null = null,
+    ids: MlxArray | null = null, multimodal: MlxArray | null = null,
   ): MlxArray {
     const h = ops.mulScalar(embeds, this.embedScale);
     if (this.perLayerWidth === 0) return this.forwardLayers(h, cache, bidir, null);
 
     if (!ids)
-      throw new Error("per-layer-input vision models require token ids");
+      throw new Error("per-layer-input multimodal models require token ids");
+    const zeroing = multimodal ?? bidir;
     let plIds = ids;
-    if (bidir) {
-      const cond = ops.reshape(bidir, ids.shape);
+    if (zeroing) {
+      const cond = ops.reshape(zeroing, ids.shape);
       const zero = ops.zeros(ids.shape, ids.dtype);
       plIds = ops.where(cond, zero, ids);
       cond.dispose();
@@ -1180,7 +1266,10 @@ export class Gemma4Model {
   logitsFromHidden(h: MlxArray): MlxArray {
     let logits = this.embed.asLinear(h);
     const softcap = this.config.text.finalLogitSoftcapping;
-    if (softcap !== null) logits = disposing(logits, logitSoftcap(logits, softcap));
+    if (softcap !== null)
+      logits = disposing(logits, compiledGegluActive() && !isCompiledTrace()
+        ? compiledLogitSoftcap(logits, softcap)   // fused, matches @mx.compile logit_softcap
+        : logitSoftcap(logits, softcap));
     return logits;
   }
 

package/src/model/generated/gemma4-12b.ts

@@ -24,9 +24,6 @@ import {
   type Mask,
   type SharedKv,
 } from "../gemma4-base";
-import {
-  fusedDecodeKernelSupported, fusedDecodeSdpa, perfKernelEnabled,
-} from "../fused-decode-kernel";
 import { Gemma4Model, type DecoderLayer } from "../gemma4";
 
 export const FINGERPRINT = "9f812d2eb461fcbe";
@@ -73,9 +70,7 @@ function donorSlidQuantG64B4VScal(layer: DecoderLayer, x: MlxArray, mask: Mask,
   q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
   q = disposing(q, a.rope(q, cache.ropeOffsetArr ?? offset));
   // dispatch-site constants: bits=4 group_size=64 nRep=2 head_dim=256
-  const attn = perfKernelEnabled() && mask.mode === "" && fusedDecodeKernelSupported(q, 4, 64)
-    ? fusedDecodeSdpa(q, kq, vq, 64, 4)
-    : (L > 1 || process.env.MLX_BUN_FUSED_DECODE === "1") && fusedSdpaRuntimeOk(q, mask)
+  const attn = L > 1 && fusedSdpaRuntimeOk(q, mask)
       ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 4)
       : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 4);
   q.dispose();
@@ -129,9 +124,7 @@ function donorSlidQuantG64B8VScal(layer: DecoderLayer, x: MlxArray, mask: Mask,
   q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
   q = disposing(q, a.rope(q, cache.ropeOffsetArr ?? offset));
   // dispatch-site constants: bits=8 group_size=64 nRep=2 head_dim=256
-  const attn = perfKernelEnabled() && mask.mode === "" && fusedDecodeKernelSupported(q, 8, 64)
-    ? fusedDecodeSdpa(q, kq, vq, 64, 8)
-    : (L > 1 || process.env.MLX_BUN_FUSED_DECODE === "1") && fusedSdpaRuntimeOk(q, mask)
+  const attn = L > 1 && fusedSdpaRuntimeOk(q, mask)
       ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 8)
       : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 8);
   q.dispose();
@@ -182,9 +175,7 @@ function donorFullQuantG64B4KeqVScal(layer: DecoderLayer, x: MlxArray, mask: Mas
   q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
   q = disposing(q, a.rope(q, cache.ropeOffsetArr ?? offset));
   // dispatch-site constants: bits=4 group_size=64 nRep=16 head_dim=512
-  const attn = perfKernelEnabled() && mask.mode === "" && fusedDecodeKernelSupported(q, 4, 64)
-    ? fusedDecodeSdpa(q, kq, vq, 64, 4)
-    : (L > 1 || process.env.MLX_BUN_FUSED_DECODE === "1") && fusedSdpaRuntimeOk(q, mask)
+  const attn = L > 1 && fusedSdpaRuntimeOk(q, mask)
       ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 4)
       : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 4);
   q.dispose();
@@ -366,7 +357,12 @@ export class GeneratedGemma4 extends Gemma4Model {
   protected override forwardLayers(
     h0: MlxArray, cache: Cache[], bidir: MlxArray | null, ids: MlxArray | null,
   ): MlxArray {
-    if (bidir !== null || !this.#matches(cache))
+    // hiddenTap (DSpark H_ctx capture) lives in the monolith's forwardLayers —
+    // this unrolled forward never calls captureLayer, so a tapped call MUST
+    // fall back or the spec path silently captures nothing and throws. Tapped
+    // forwards are only the spec path's prefill/verify; the monolith is the
+    // bit-exact reference, so this is correct-but-unspecialized.
+    if (bidir !== null || this.hiddenTap !== null || !this.#matches(cache))
       return super.forwardLayers(h0, cache, bidir, ids);
     generatedForwardUses++;
     const L = h0.shape[1]!;

package/src/model/generated/gemma4-26b.ts

@@ -70,7 +70,7 @@ function donorSlidQuantG64B8VMoeScal(layer: DecoderLayer, x: MlxArray, mask: Mas
   q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
   q = disposing(q, a.rope(q, cache.ropeOffsetArr ?? offset));
   // dispatch-site constants: bits=8 group_size=64 nRep=2 head_dim=256
-  const attn = (L > 1 || process.env.MLX_BUN_FUSED_DECODE === "1") && fusedSdpaRuntimeOk(q, mask)
+  const attn = L > 1 && fusedSdpaRuntimeOk(q, mask)
     ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 8)
     : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 8);
   q.dispose();
@@ -134,7 +134,7 @@ function donorSlidQuantG64B4VMoeScal(layer: DecoderLayer, x: MlxArray, mask: Mas
   q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
   q = disposing(q, a.rope(q, cache.ropeOffsetArr ?? offset));
   // dispatch-site constants: bits=4 group_size=64 nRep=2 head_dim=256
-  const attn = (L > 1 || process.env.MLX_BUN_FUSED_DECODE === "1") && fusedSdpaRuntimeOk(q, mask)
+  const attn = L > 1 && fusedSdpaRuntimeOk(q, mask)
     ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 4)
     : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 4);
   q.dispose();
@@ -195,7 +195,7 @@ function donorFullQuantG64B8KeqVMoeScal(layer: DecoderLayer, x: MlxArray, mask:
   q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
   q = disposing(q, a.rope(q, cache.ropeOffsetArr ?? offset));
   // dispatch-site constants: bits=8 group_size=64 nRep=8 head_dim=512
-  const attn = (L > 1 || process.env.MLX_BUN_FUSED_DECODE === "1") && fusedSdpaRuntimeOk(q, mask)
+  const attn = L > 1 && fusedSdpaRuntimeOk(q, mask)
     ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 8)
     : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 8);
   q.dispose();
@@ -256,7 +256,7 @@ function donorFullQuantG64B4KeqVMoeScal(layer: DecoderLayer, x: MlxArray, mask:
   q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
   q = disposing(q, a.rope(q, cache.ropeOffsetArr ?? offset));
   // dispatch-site constants: bits=4 group_size=64 nRep=8 head_dim=512
-  const attn = (L > 1 || process.env.MLX_BUN_FUSED_DECODE === "1") && fusedSdpaRuntimeOk(q, mask)
+  const attn = L > 1 && fusedSdpaRuntimeOk(q, mask)
     ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 4)
     : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 4);
   q.dispose();
@@ -394,7 +394,12 @@ export class GeneratedGemma4 extends Gemma4Model {
   protected override forwardLayers(
     h0: MlxArray, cache: Cache[], bidir: MlxArray | null, ids: MlxArray | null,
   ): MlxArray {
-    if (bidir !== null || !this.#matches(cache))
+    // hiddenTap (DSpark H_ctx capture) lives in the monolith's forwardLayers —
+    // this unrolled forward never calls captureLayer, so a tapped call MUST
+    // fall back or the spec path silently captures nothing and throws. Tapped
+    // forwards are only the spec path's prefill/verify; the monolith is the
+    // bit-exact reference, so this is correct-but-unspecialized.
+    if (bidir !== null || this.hiddenTap !== null || !this.#matches(cache))
       return super.forwardLayers(h0, cache, bidir, ids);
     generatedForwardUses++;
     const L = h0.shape[1]!;

package/src/model/generated/gemma4-e4b.ts

@@ -88,7 +88,7 @@ function donorSlidQuantG64B4VPliScal(layer: DecoderLayer, x: MlxArray, mask: Mas
   q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
   q = disposing(q, a.rope(q, cache.ropeOffsetArr ?? offset));
   // dispatch-site constants: bits=4 group_size=64 nRep=4 head_dim=256
-  const attn = (L > 1 || process.env.MLX_BUN_FUSED_DECODE === "1") && fusedSdpaRuntimeOk(q, mask)
+  const attn = L > 1 && fusedSdpaRuntimeOk(q, mask)
     ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 4)
     : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 4);
   q.dispose();
@@ -151,7 +151,7 @@ function donorSlidQuantG64B8VPliScal(layer: DecoderLayer, x: MlxArray, mask: Mas
   q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
   q = disposing(q, a.rope(q, cache.ropeOffsetArr ?? offset));
   // dispatch-site constants: bits=8 group_size=64 nRep=4 head_dim=256
-  const attn = (L > 1 || process.env.MLX_BUN_FUSED_DECODE === "1") && fusedSdpaRuntimeOk(q, mask)
+  const attn = L > 1 && fusedSdpaRuntimeOk(q, mask)
     ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 8)
     : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 8);
   q.dispose();
@@ -214,7 +214,7 @@ function donorFullQuantG64B4VPliScal(layer: DecoderLayer, x: MlxArray, mask: Mas
   q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
   q = disposing(q, a.rope(q, cache.ropeOffsetArr ?? offset));
   // dispatch-site constants: bits=4 group_size=64 nRep=4 head_dim=512
-  const attn = (L > 1 || process.env.MLX_BUN_FUSED_DECODE === "1") && fusedSdpaRuntimeOk(q, mask)
+  const attn = L > 1 && fusedSdpaRuntimeOk(q, mask)
     ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 4)
     : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 4);
   q.dispose();
@@ -277,7 +277,7 @@ function donorFullQuantG64B8VPliScal(layer: DecoderLayer, x: MlxArray, mask: Mas
   q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
   q = disposing(q, a.rope(q, cache.ropeOffsetArr ?? offset));
   // dispatch-site constants: bits=8 group_size=64 nRep=4 head_dim=512
-  const attn = (L > 1 || process.env.MLX_BUN_FUSED_DECODE === "1") && fusedSdpaRuntimeOk(q, mask)
+  const attn = L > 1 && fusedSdpaRuntimeOk(q, mask)
     ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 8)
     : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 8);
   q.dispose();
@@ -340,7 +340,7 @@ function donorSlidQuantG64B4VPliScalExp(layer: DecoderLayer, x: MlxArray, mask:
   q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
   q = disposing(q, a.rope(q, cache.ropeOffsetArr ?? offset));
   // dispatch-site constants: bits=4 group_size=64 nRep=4 head_dim=256
-  const attn = (L > 1 || process.env.MLX_BUN_FUSED_DECODE === "1") && fusedSdpaRuntimeOk(q, mask)
+  const attn = L > 1 && fusedSdpaRuntimeOk(q, mask)
     ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 4)
     : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 4);
   q.dispose();
@@ -403,7 +403,7 @@ function donorFullQuantG64B4VPliScalExp(layer: DecoderLayer, x: MlxArray, mask:
   q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
   q = disposing(q, a.rope(q, cache.ropeOffsetArr ?? offset));
   // dispatch-site constants: bits=4 group_size=64 nRep=4 head_dim=512
-  const attn = (L > 1 || process.env.MLX_BUN_FUSED_DECODE === "1") && fusedSdpaRuntimeOk(q, mask)
+  const attn = L > 1 && fusedSdpaRuntimeOk(q, mask)
     ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 4)
     : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 4);
   q.dispose();
@@ -449,7 +449,7 @@ function sharerSlidQuantG64B4VPliScal(layer: DecoderLayer, x: MlxArray, mask: Ma
   q = disposing(q, a.rope(q, shared.offsetArr ?? shared.offset));
   if (shared.kind !== "quant") throw new Error("generated sharer expected quant shared KV");
   // dispatch-site constants: bits=4 group_size=64 nRep=4 head_dim=256
-  const attn = (L > 1 || process.env.MLX_BUN_FUSED_DECODE === "1") && fusedSdpaRuntimeOk(q, mask)
+  const attn = L > 1 && fusedSdpaRuntimeOk(q, mask)
     ? quantizedSdpaTiled(q, shared.keys, shared.values, 1.0, mask, 64, 4)
     : quantizedSdpaUnfused(q, shared.keys, shared.values, 1.0, mask, 64, 4);
   q.dispose();
@@ -494,7 +494,7 @@ function sharerFullQuantG64B4VPliScal(layer: DecoderLayer, x: MlxArray, mask: Ma
   q = disposing(q, a.rope(q, shared.offsetArr ?? shared.offset));
   if (shared.kind !== "quant") throw new Error("generated sharer expected quant shared KV");
   // dispatch-site constants: bits=4 group_size=64 nRep=4 head_dim=512
-  const attn = (L > 1 || process.env.MLX_BUN_FUSED_DECODE === "1") && fusedSdpaRuntimeOk(q, mask)
+  const attn = L > 1 && fusedSdpaRuntimeOk(q, mask)
     ? quantizedSdpaTiled(q, shared.keys, shared.values, 1.0, mask, 64, 4)
     : quantizedSdpaUnfused(q, shared.keys, shared.values, 1.0, mask, 64, 4);
   q.dispose();
@@ -612,7 +612,12 @@ export class GeneratedGemma4 extends Gemma4Model {
   protected override forwardLayers(
     h0: MlxArray, cache: Cache[], bidir: MlxArray | null, ids: MlxArray | null,
   ): MlxArray {
-    if (bidir !== null || ids === null || !this.#matches(cache))
+    // hiddenTap (DSpark H_ctx capture) lives in the monolith's forwardLayers —
+    // this unrolled forward never calls captureLayer, so a tapped call MUST
+    // fall back or the spec path silently captures nothing and throws. Tapped
+    // forwards are only the spec path's prefill/verify; the monolith is the
+    // bit-exact reference, so this is correct-but-unspecialized.
+    if (bidir !== null || ids === null || this.hiddenTap !== null || !this.#matches(cache))
       return super.forwardLayers(h0, cache, bidir, ids);
     generatedForwardUses++;
     const L = h0.shape[1]!;