mlx-bun 0.0.4

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

@@ -0,0 +1,757 @@
+// GENERATED by scripts/gen-model.ts — DO NOT EDIT (regenerate instead):
+//   bun scripts/gen-model.ts <model-dir> gemma4-e4b
+// source: models--mlx-community--gemma-4-e4b-it-OptiQ-4bit/snapshots/fcdb12d740cd813634064567fc7cb51159b34253
+// fingerprint: 418e9adc386ea67c
+//
+// Branch-resolved, unrolled forward pass for this architecture
+// (docs/design/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
+// generator's gate: tests/generated-parity.test.ts. The cache-signature
+// guard falls back to the monolith for anything this file wasn't
+// generated for (bf16 compat runs, vision bidir prefill).
+
+import { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import {
+  disposing,
+  fusedSdpaRuntimeOk,
+  KVCache,
+  QuantizedKVCache,
+  quantizedSdpaTiled,
+  quantizedSdpaUnfused,
+  RotatingKVCache,
+  RotatingQuantizedKVCache,
+  type Cache,
+  type Mask,
+  type SharedKv,
+} from "../gemma4-base";
+import { Gemma4Model, type DecoderLayer } from "../gemma4";
+import { isCompiledTrace } from "../gemma4-base";
+
+export const FINGERPRINT = "418e9adc386ea67c";
+
+/** Forward passes served by the generated fast path (vs monolith
+ *  fallback) — the parity gate asserts the fast path actually ran. */
+export let generatedForwardUses = 0;
+
+function disposeSharedKv(s: SharedKv): void {
+  if (s.kind === "plain") {
+    s.keys.dispose();
+    s.values.dispose();
+  } else {
+    for (const t of [s.keys, s.values])
+      for (const a of [t.packed, t.scales, t.biases]) a.dispose();
+  }
+}
+
+/** Per-layer-input split with the layer index baked; DynamicSlice under
+ *  a compiled-decode trace (mlx Slice lacks output_shapes), the exact
+ *  monolith ops otherwise. */
+function perLayerSlice(perLayer: MlxArray, i: number, L: number, width: number): MlxArray {
+  let pls: MlxArray;
+  if (isCompiledTrace()) {
+    const start = ops.fromInt32([i], [1]);
+    pls = ops.sliceDynamic(perLayer, start, [2], [1, L, 1, width]);
+    start.dispose();
+  } else {
+    pls = perLayer.slice([0, 0, i, 0], [1, L, i + 1, width]);
+  }
+  const r = ops.reshape(pls, [1, L, width]);
+  pls.dispose();
+  return r;
+}
+
+function donorSlidQuantG64B4VPliScal(layer: DecoderLayer, x: MlxArray, mask: Mask, cache: RotatingQuantizedKVCache, pli: MlxArray): MlxArray {
+  const a = layer.attn;
+  const [B, L] = x.shape as [number, number, number];
+  let h = layer.inputNorm.forward(x);
+  let q = a.qProj.forward(h);
+  q = disposing(q, ops.reshape(q, [B, L, a.nHeads, a.headDim]));
+  q = disposing(q, a.qNorm.forward(q));
+  const offset = cache.offset;
+  let k = a.kProj!.forward(h);
+  k = disposing(k, ops.reshape(k, [B, L, a.nKvHeads, a.headDim]));
+  let v = a.vProj!.forward(h);
+  v = disposing(v, ops.reshape(v, [B, L, a.nKvHeads, a.headDim]));
+  const kNormed = a.kNorm!.forward(k);
+  const kT = ops.transposeAxes(kNormed, [0, 2, 1, 3]);
+  kNormed.dispose();
+  const kRoped = a.rope(kT, cache.ropeOffsetArr ?? offset);
+  kT.dispose();
+  const vNormed = a.vNorm!.forward(v);
+  const vT = ops.transposeAxes(vNormed, [0, 2, 1, 3]);
+  vNormed.dispose();
+  v.dispose();
+  k.dispose();
+  const [kq, vq] = cache.updateAndFetchQuantized(kRoped, vT);
+  kRoped.dispose();
+  vT.dispose();
+  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)
+    ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 4)
+    : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 4);
+  q.dispose();
+  const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
+  attn.dispose();
+  const merged = ops.reshape(attnT, [B, L, -1]);
+  attnT.dispose();
+  const out = a.oProj.forward(merged);
+  merged.dispose();
+  h.dispose();
+  h = out;
+  h = disposing(h, layer.postAttnNorm.forward(h));
+  h = disposing(h, ops.add(x, h));
+  const residual = h;
+  let f = layer.preFfNorm.forward(h);
+  f = disposing(f, layer.mlp.forward(f));
+  f = disposing(f, layer.postFfNorm.forward(f));
+  h = ops.add(residual, f);
+  residual.dispose();
+  f.dispose();
+  const res2 = h;
+  let gate = layer.perLayerGate!.forward(h);
+  gate = disposing(gate, ops.geluApprox(gate));
+  gate = disposing(gate, ops.mul(gate, pli));
+  gate = disposing(gate, layer.perLayerProjection!.forward(gate));
+  gate = disposing(gate, layer.postPerLayerNorm!.forward(gate));
+  h = ops.add(res2, gate);
+  res2.dispose();
+  gate.dispose();
+  h = disposing(h, ops.mul(h, layer.layerScalar!));
+  for (const t of [kq, vq]) for (const c of [t.packed, t.scales, t.biases]) c.dispose();
+  return h;
+}
+
+function donorSlidQuantG64B8VPliScal(layer: DecoderLayer, x: MlxArray, mask: Mask, cache: RotatingQuantizedKVCache, pli: MlxArray): MlxArray {
+  const a = layer.attn;
+  const [B, L] = x.shape as [number, number, number];
+  let h = layer.inputNorm.forward(x);
+  let q = a.qProj.forward(h);
+  q = disposing(q, ops.reshape(q, [B, L, a.nHeads, a.headDim]));
+  q = disposing(q, a.qNorm.forward(q));
+  const offset = cache.offset;
+  let k = a.kProj!.forward(h);
+  k = disposing(k, ops.reshape(k, [B, L, a.nKvHeads, a.headDim]));
+  let v = a.vProj!.forward(h);
+  v = disposing(v, ops.reshape(v, [B, L, a.nKvHeads, a.headDim]));
+  const kNormed = a.kNorm!.forward(k);
+  const kT = ops.transposeAxes(kNormed, [0, 2, 1, 3]);
+  kNormed.dispose();
+  const kRoped = a.rope(kT, cache.ropeOffsetArr ?? offset);
+  kT.dispose();
+  const vNormed = a.vNorm!.forward(v);
+  const vT = ops.transposeAxes(vNormed, [0, 2, 1, 3]);
+  vNormed.dispose();
+  v.dispose();
+  k.dispose();
+  const [kq, vq] = cache.updateAndFetchQuantized(kRoped, vT);
+  kRoped.dispose();
+  vT.dispose();
+  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)
+    ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 8)
+    : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 8);
+  q.dispose();
+  const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
+  attn.dispose();
+  const merged = ops.reshape(attnT, [B, L, -1]);
+  attnT.dispose();
+  const out = a.oProj.forward(merged);
+  merged.dispose();
+  h.dispose();
+  h = out;
+  h = disposing(h, layer.postAttnNorm.forward(h));
+  h = disposing(h, ops.add(x, h));
+  const residual = h;
+  let f = layer.preFfNorm.forward(h);
+  f = disposing(f, layer.mlp.forward(f));
+  f = disposing(f, layer.postFfNorm.forward(f));
+  h = ops.add(residual, f);
+  residual.dispose();
+  f.dispose();
+  const res2 = h;
+  let gate = layer.perLayerGate!.forward(h);
+  gate = disposing(gate, ops.geluApprox(gate));
+  gate = disposing(gate, ops.mul(gate, pli));
+  gate = disposing(gate, layer.perLayerProjection!.forward(gate));
+  gate = disposing(gate, layer.postPerLayerNorm!.forward(gate));
+  h = ops.add(res2, gate);
+  res2.dispose();
+  gate.dispose();
+  h = disposing(h, ops.mul(h, layer.layerScalar!));
+  for (const t of [kq, vq]) for (const c of [t.packed, t.scales, t.biases]) c.dispose();
+  return h;
+}
+
+function donorFullQuantG64B4VPliScal(layer: DecoderLayer, x: MlxArray, mask: Mask, cache: QuantizedKVCache, pli: MlxArray): MlxArray {
+  const a = layer.attn;
+  const [B, L] = x.shape as [number, number, number];
+  let h = layer.inputNorm.forward(x);
+  let q = a.qProj.forward(h);
+  q = disposing(q, ops.reshape(q, [B, L, a.nHeads, a.headDim]));
+  q = disposing(q, a.qNorm.forward(q));
+  const offset = cache.offset;
+  let k = a.kProj!.forward(h);
+  k = disposing(k, ops.reshape(k, [B, L, a.nKvHeads, a.headDim]));
+  let v = a.vProj!.forward(h);
+  v = disposing(v, ops.reshape(v, [B, L, a.nKvHeads, a.headDim]));
+  const kNormed = a.kNorm!.forward(k);
+  const kT = ops.transposeAxes(kNormed, [0, 2, 1, 3]);
+  kNormed.dispose();
+  const kRoped = a.rope(kT, cache.ropeOffsetArr ?? offset);
+  kT.dispose();
+  const vNormed = a.vNorm!.forward(v);
+  const vT = ops.transposeAxes(vNormed, [0, 2, 1, 3]);
+  vNormed.dispose();
+  v.dispose();
+  k.dispose();
+  const [kq, vq] = cache.updateAndFetchQuantized(kRoped, vT);
+  kRoped.dispose();
+  vT.dispose();
+  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)
+    ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 4)
+    : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 4);
+  q.dispose();
+  const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
+  attn.dispose();
+  const merged = ops.reshape(attnT, [B, L, -1]);
+  attnT.dispose();
+  const out = a.oProj.forward(merged);
+  merged.dispose();
+  h.dispose();
+  h = out;
+  h = disposing(h, layer.postAttnNorm.forward(h));
+  h = disposing(h, ops.add(x, h));
+  const residual = h;
+  let f = layer.preFfNorm.forward(h);
+  f = disposing(f, layer.mlp.forward(f));
+  f = disposing(f, layer.postFfNorm.forward(f));
+  h = ops.add(residual, f);
+  residual.dispose();
+  f.dispose();
+  const res2 = h;
+  let gate = layer.perLayerGate!.forward(h);
+  gate = disposing(gate, ops.geluApprox(gate));
+  gate = disposing(gate, ops.mul(gate, pli));
+  gate = disposing(gate, layer.perLayerProjection!.forward(gate));
+  gate = disposing(gate, layer.postPerLayerNorm!.forward(gate));
+  h = ops.add(res2, gate);
+  res2.dispose();
+  gate.dispose();
+  h = disposing(h, ops.mul(h, layer.layerScalar!));
+  for (const t of [kq, vq]) for (const c of [t.packed, t.scales, t.biases]) c.dispose();
+  return h;
+}
+
+function donorFullQuantG64B8VPliScal(layer: DecoderLayer, x: MlxArray, mask: Mask, cache: QuantizedKVCache, pli: MlxArray): MlxArray {
+  const a = layer.attn;
+  const [B, L] = x.shape as [number, number, number];
+  let h = layer.inputNorm.forward(x);
+  let q = a.qProj.forward(h);
+  q = disposing(q, ops.reshape(q, [B, L, a.nHeads, a.headDim]));
+  q = disposing(q, a.qNorm.forward(q));
+  const offset = cache.offset;
+  let k = a.kProj!.forward(h);
+  k = disposing(k, ops.reshape(k, [B, L, a.nKvHeads, a.headDim]));
+  let v = a.vProj!.forward(h);
+  v = disposing(v, ops.reshape(v, [B, L, a.nKvHeads, a.headDim]));
+  const kNormed = a.kNorm!.forward(k);
+  const kT = ops.transposeAxes(kNormed, [0, 2, 1, 3]);
+  kNormed.dispose();
+  const kRoped = a.rope(kT, cache.ropeOffsetArr ?? offset);
+  kT.dispose();
+  const vNormed = a.vNorm!.forward(v);
+  const vT = ops.transposeAxes(vNormed, [0, 2, 1, 3]);
+  vNormed.dispose();
+  v.dispose();
+  k.dispose();
+  const [kq, vq] = cache.updateAndFetchQuantized(kRoped, vT);
+  kRoped.dispose();
+  vT.dispose();
+  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)
+    ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 8)
+    : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 8);
+  q.dispose();
+  const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
+  attn.dispose();
+  const merged = ops.reshape(attnT, [B, L, -1]);
+  attnT.dispose();
+  const out = a.oProj.forward(merged);
+  merged.dispose();
+  h.dispose();
+  h = out;
+  h = disposing(h, layer.postAttnNorm.forward(h));
+  h = disposing(h, ops.add(x, h));
+  const residual = h;
+  let f = layer.preFfNorm.forward(h);
+  f = disposing(f, layer.mlp.forward(f));
+  f = disposing(f, layer.postFfNorm.forward(f));
+  h = ops.add(residual, f);
+  residual.dispose();
+  f.dispose();
+  const res2 = h;
+  let gate = layer.perLayerGate!.forward(h);
+  gate = disposing(gate, ops.geluApprox(gate));
+  gate = disposing(gate, ops.mul(gate, pli));
+  gate = disposing(gate, layer.perLayerProjection!.forward(gate));
+  gate = disposing(gate, layer.postPerLayerNorm!.forward(gate));
+  h = ops.add(res2, gate);
+  res2.dispose();
+  gate.dispose();
+  h = disposing(h, ops.mul(h, layer.layerScalar!));
+  for (const t of [kq, vq]) for (const c of [t.packed, t.scales, t.biases]) c.dispose();
+  return h;
+}
+
+function donorSlidQuantG64B4VPliScalExp(layer: DecoderLayer, x: MlxArray, mask: Mask, cache: RotatingQuantizedKVCache, pli: MlxArray): { h: MlxArray; shared: SharedKv } {
+  const a = layer.attn;
+  const [B, L] = x.shape as [number, number, number];
+  let h = layer.inputNorm.forward(x);
+  let q = a.qProj.forward(h);
+  q = disposing(q, ops.reshape(q, [B, L, a.nHeads, a.headDim]));
+  q = disposing(q, a.qNorm.forward(q));
+  const offset = cache.offset;
+  let k = a.kProj!.forward(h);
+  k = disposing(k, ops.reshape(k, [B, L, a.nKvHeads, a.headDim]));
+  let v = a.vProj!.forward(h);
+  v = disposing(v, ops.reshape(v, [B, L, a.nKvHeads, a.headDim]));
+  const kNormed = a.kNorm!.forward(k);
+  const kT = ops.transposeAxes(kNormed, [0, 2, 1, 3]);
+  kNormed.dispose();
+  const kRoped = a.rope(kT, cache.ropeOffsetArr ?? offset);
+  kT.dispose();
+  const vNormed = a.vNorm!.forward(v);
+  const vT = ops.transposeAxes(vNormed, [0, 2, 1, 3]);
+  vNormed.dispose();
+  v.dispose();
+  k.dispose();
+  const [kq, vq] = cache.updateAndFetchQuantized(kRoped, vT);
+  kRoped.dispose();
+  vT.dispose();
+  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)
+    ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 4)
+    : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 4);
+  q.dispose();
+  const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
+  attn.dispose();
+  const merged = ops.reshape(attnT, [B, L, -1]);
+  attnT.dispose();
+  const out = a.oProj.forward(merged);
+  merged.dispose();
+  h.dispose();
+  h = out;
+  h = disposing(h, layer.postAttnNorm.forward(h));
+  h = disposing(h, ops.add(x, h));
+  const residual = h;
+  let f = layer.preFfNorm.forward(h);
+  f = disposing(f, layer.mlp.forward(f));
+  f = disposing(f, layer.postFfNorm.forward(f));
+  h = ops.add(residual, f);
+  residual.dispose();
+  f.dispose();
+  const res2 = h;
+  let gate = layer.perLayerGate!.forward(h);
+  gate = disposing(gate, ops.geluApprox(gate));
+  gate = disposing(gate, ops.mul(gate, pli));
+  gate = disposing(gate, layer.perLayerProjection!.forward(gate));
+  gate = disposing(gate, layer.postPerLayerNorm!.forward(gate));
+  h = ops.add(res2, gate);
+  res2.dispose();
+  gate.dispose();
+  h = disposing(h, ops.mul(h, layer.layerScalar!));
+  const shared: SharedKv = { kind: "quant", keys: kq, values: vq, offset, groupSize: 64, bits: 4, offsetArr: cache.ropeOffsetArr };
+  return { h, shared };
+}
+
+function donorFullQuantG64B4VPliScalExp(layer: DecoderLayer, x: MlxArray, mask: Mask, cache: QuantizedKVCache, pli: MlxArray): { h: MlxArray; shared: SharedKv } {
+  const a = layer.attn;
+  const [B, L] = x.shape as [number, number, number];
+  let h = layer.inputNorm.forward(x);
+  let q = a.qProj.forward(h);
+  q = disposing(q, ops.reshape(q, [B, L, a.nHeads, a.headDim]));
+  q = disposing(q, a.qNorm.forward(q));
+  const offset = cache.offset;
+  let k = a.kProj!.forward(h);
+  k = disposing(k, ops.reshape(k, [B, L, a.nKvHeads, a.headDim]));
+  let v = a.vProj!.forward(h);
+  v = disposing(v, ops.reshape(v, [B, L, a.nKvHeads, a.headDim]));
+  const kNormed = a.kNorm!.forward(k);
+  const kT = ops.transposeAxes(kNormed, [0, 2, 1, 3]);
+  kNormed.dispose();
+  const kRoped = a.rope(kT, cache.ropeOffsetArr ?? offset);
+  kT.dispose();
+  const vNormed = a.vNorm!.forward(v);
+  const vT = ops.transposeAxes(vNormed, [0, 2, 1, 3]);
+  vNormed.dispose();
+  v.dispose();
+  k.dispose();
+  const [kq, vq] = cache.updateAndFetchQuantized(kRoped, vT);
+  kRoped.dispose();
+  vT.dispose();
+  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)
+    ? quantizedSdpaTiled(q, kq, vq, 1.0, mask, 64, 4)
+    : quantizedSdpaUnfused(q, kq, vq, 1.0, mask, 64, 4);
+  q.dispose();
+  const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
+  attn.dispose();
+  const merged = ops.reshape(attnT, [B, L, -1]);
+  attnT.dispose();
+  const out = a.oProj.forward(merged);
+  merged.dispose();
+  h.dispose();
+  h = out;
+  h = disposing(h, layer.postAttnNorm.forward(h));
+  h = disposing(h, ops.add(x, h));
+  const residual = h;
+  let f = layer.preFfNorm.forward(h);
+  f = disposing(f, layer.mlp.forward(f));
+  f = disposing(f, layer.postFfNorm.forward(f));
+  h = ops.add(residual, f);
+  residual.dispose();
+  f.dispose();
+  const res2 = h;
+  let gate = layer.perLayerGate!.forward(h);
+  gate = disposing(gate, ops.geluApprox(gate));
+  gate = disposing(gate, ops.mul(gate, pli));
+  gate = disposing(gate, layer.perLayerProjection!.forward(gate));
+  gate = disposing(gate, layer.postPerLayerNorm!.forward(gate));
+  h = ops.add(res2, gate);
+  res2.dispose();
+  gate.dispose();
+  h = disposing(h, ops.mul(h, layer.layerScalar!));
+  const shared: SharedKv = { kind: "quant", keys: kq, values: vq, offset, groupSize: 64, bits: 4, offsetArr: cache.ropeOffsetArr };
+  return { h, shared };
+}
+
+function sharerSlidQuantG64B4VPliScal(layer: DecoderLayer, x: MlxArray, mask: Mask, shared: SharedKv, pli: MlxArray): MlxArray {
+  const a = layer.attn;
+  const [B, L] = x.shape as [number, number, number];
+  let h = layer.inputNorm.forward(x);
+  let q = a.qProj.forward(h);
+  q = disposing(q, ops.reshape(q, [B, L, a.nHeads, a.headDim]));
+  q = disposing(q, a.qNorm.forward(q));
+  q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
+  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)
+    ? quantizedSdpaTiled(q, shared.keys, shared.values, 1.0, mask, 64, 4)
+    : quantizedSdpaUnfused(q, shared.keys, shared.values, 1.0, mask, 64, 4);
+  q.dispose();
+  const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
+  attn.dispose();
+  const merged = ops.reshape(attnT, [B, L, -1]);
+  attnT.dispose();
+  const out = a.oProj.forward(merged);
+  merged.dispose();
+  h.dispose();
+  h = out;
+  h = disposing(h, layer.postAttnNorm.forward(h));
+  h = disposing(h, ops.add(x, h));
+  const residual = h;
+  let f = layer.preFfNorm.forward(h);
+  f = disposing(f, layer.mlp.forward(f));
+  f = disposing(f, layer.postFfNorm.forward(f));
+  h = ops.add(residual, f);
+  residual.dispose();
+  f.dispose();
+  const res2 = h;
+  let gate = layer.perLayerGate!.forward(h);
+  gate = disposing(gate, ops.geluApprox(gate));
+  gate = disposing(gate, ops.mul(gate, pli));
+  gate = disposing(gate, layer.perLayerProjection!.forward(gate));
+  gate = disposing(gate, layer.postPerLayerNorm!.forward(gate));
+  h = ops.add(res2, gate);
+  res2.dispose();
+  gate.dispose();
+  h = disposing(h, ops.mul(h, layer.layerScalar!));
+  return h;
+}
+
+function sharerFullQuantG64B4VPliScal(layer: DecoderLayer, x: MlxArray, mask: Mask, shared: SharedKv, pli: MlxArray): MlxArray {
+  const a = layer.attn;
+  const [B, L] = x.shape as [number, number, number];
+  let h = layer.inputNorm.forward(x);
+  let q = a.qProj.forward(h);
+  q = disposing(q, ops.reshape(q, [B, L, a.nHeads, a.headDim]));
+  q = disposing(q, a.qNorm.forward(q));
+  q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
+  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)
+    ? quantizedSdpaTiled(q, shared.keys, shared.values, 1.0, mask, 64, 4)
+    : quantizedSdpaUnfused(q, shared.keys, shared.values, 1.0, mask, 64, 4);
+  q.dispose();
+  const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
+  attn.dispose();
+  const merged = ops.reshape(attnT, [B, L, -1]);
+  attnT.dispose();
+  const out = a.oProj.forward(merged);
+  merged.dispose();
+  h.dispose();
+  h = out;
+  h = disposing(h, layer.postAttnNorm.forward(h));
+  h = disposing(h, ops.add(x, h));
+  const residual = h;
+  let f = layer.preFfNorm.forward(h);
+  f = disposing(f, layer.mlp.forward(f));
+  f = disposing(f, layer.postFfNorm.forward(f));
+  h = ops.add(residual, f);
+  residual.dispose();
+  f.dispose();
+  const res2 = h;
+  let gate = layer.perLayerGate!.forward(h);
+  gate = disposing(gate, ops.geluApprox(gate));
+  gate = disposing(gate, ops.mul(gate, pli));
+  gate = disposing(gate, layer.perLayerProjection!.forward(gate));
+  gate = disposing(gate, layer.postPerLayerNorm!.forward(gate));
+  h = ops.add(res2, gate);
+  res2.dispose();
+  gate.dispose();
+  h = disposing(h, ops.mul(h, layer.layerScalar!));
+  return h;
+}
+
+export class GeneratedGemma4 extends Gemma4Model {
+  /** Per-layer cache classes must match what this file was generated
+   *  for; anything else runs the monolith path. */
+  #matches(cache: Cache[]): boolean {
+    return (
+      cache.length === 24 &&
+      cache[0] instanceof RotatingQuantizedKVCache &&
+      cache[0].bits === 4 &&
+      cache[0].groupSize === 64 &&
+      cache[1] instanceof RotatingQuantizedKVCache &&
+      cache[1].bits === 4 &&
+      cache[1].groupSize === 64 &&
+      cache[2] instanceof RotatingQuantizedKVCache &&
+      cache[2].bits === 4 &&
+      cache[2].groupSize === 64 &&
+      cache[3] instanceof RotatingQuantizedKVCache &&
+      cache[3].bits === 8 &&
+      cache[3].groupSize === 64 &&
+      cache[4] instanceof RotatingQuantizedKVCache &&
+      cache[4].bits === 8 &&
+      cache[4].groupSize === 64 &&
+      cache[5] instanceof QuantizedKVCache &&
+      cache[5].bits === 4 &&
+      cache[5].groupSize === 64 &&
+      cache[6] instanceof RotatingQuantizedKVCache &&
+      cache[6].bits === 4 &&
+      cache[6].groupSize === 64 &&
+      cache[7] instanceof RotatingQuantizedKVCache &&
+      cache[7].bits === 4 &&
+      cache[7].groupSize === 64 &&
+      cache[8] instanceof RotatingQuantizedKVCache &&
+      cache[8].bits === 4 &&
+      cache[8].groupSize === 64 &&
+      cache[9] instanceof RotatingQuantizedKVCache &&
+      cache[9].bits === 8 &&
+      cache[9].groupSize === 64 &&
+      cache[10] instanceof RotatingQuantizedKVCache &&
+      cache[10].bits === 8 &&
+      cache[10].groupSize === 64 &&
+      cache[11] instanceof QuantizedKVCache &&
+      cache[11].bits === 8 &&
+      cache[11].groupSize === 64 &&
+      cache[12] instanceof RotatingQuantizedKVCache &&
+      cache[12].bits === 4 &&
+      cache[12].groupSize === 64 &&
+      cache[13] instanceof RotatingQuantizedKVCache &&
+      cache[13].bits === 4 &&
+      cache[13].groupSize === 64 &&
+      cache[14] instanceof RotatingQuantizedKVCache &&
+      cache[14].bits === 8 &&
+      cache[14].groupSize === 64 &&
+      cache[15] instanceof RotatingQuantizedKVCache &&
+      cache[15].bits === 4 &&
+      cache[15].groupSize === 64 &&
+      cache[16] instanceof RotatingQuantizedKVCache &&
+      cache[16].bits === 4 &&
+      cache[16].groupSize === 64 &&
+      cache[17] instanceof QuantizedKVCache &&
+      cache[17].bits === 4 &&
+      cache[17].groupSize === 64 &&
+      cache[18] instanceof RotatingQuantizedKVCache &&
+      cache[18].bits === 4 &&
+      cache[18].groupSize === 64 &&
+      cache[19] instanceof RotatingQuantizedKVCache &&
+      cache[19].bits === 4 &&
+      cache[19].groupSize === 64 &&
+      cache[20] instanceof RotatingQuantizedKVCache &&
+      cache[20].bits === 4 &&
+      cache[20].groupSize === 64 &&
+      cache[21] instanceof RotatingQuantizedKVCache &&
+      cache[21].bits === 4 &&
+      cache[21].groupSize === 64 &&
+      cache[22] instanceof RotatingQuantizedKVCache &&
+      cache[22].bits === 4 &&
+      cache[22].groupSize === 64 &&
+      cache[23] instanceof QuantizedKVCache &&
+      cache[23].bits === 4 &&
+      cache[23].groupSize === 64
+    );
+  }
+
+  protected override forwardLayers(
+    h0: MlxArray, cache: Cache[], bidir: MlxArray | null, ids: MlxArray | null,
+  ): MlxArray {
+    if (bidir !== null || ids === null || !this.#matches(cache))
+      return super.forwardLayers(h0, cache, bidir, ids);
+    generatedForwardUses++;
+    const L = h0.shape[1]!;
+    const maskS = cache[0]!.makeMask(L, this.windowSize);
+    const maskF = cache[5]!.makeMask(L, null);
+    const perLayer = this.computePerLayerInputs(ids!, h0);
+    let h = h0;
+    let next: MlxArray;
+    next = donorSlidQuantG64B4VPliScal(this.layers[0]!, h, maskS, cache[0] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 0, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B4VPliScal(this.layers[1]!, h, maskS, cache[1] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 1, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B4VPliScal(this.layers[2]!, h, maskS, cache[2] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 2, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B8VPliScal(this.layers[3]!, h, maskS, cache[3] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 3, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B8VPliScal(this.layers[4]!, h, maskS, cache[4] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 4, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorFullQuantG64B4VPliScal(this.layers[5]!, h, maskF, cache[5] as QuantizedKVCache, perLayerSlice(perLayer, 5, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B4VPliScal(this.layers[6]!, h, maskS, cache[6] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 6, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B4VPliScal(this.layers[7]!, h, maskS, cache[7] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 7, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B4VPliScal(this.layers[8]!, h, maskS, cache[8] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 8, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B8VPliScal(this.layers[9]!, h, maskS, cache[9] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 9, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B8VPliScal(this.layers[10]!, h, maskS, cache[10] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 10, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorFullQuantG64B8VPliScal(this.layers[11]!, h, maskF, cache[11] as QuantizedKVCache, perLayerSlice(perLayer, 11, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B4VPliScal(this.layers[12]!, h, maskS, cache[12] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 12, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B4VPliScal(this.layers[13]!, h, maskS, cache[13] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 13, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B8VPliScal(this.layers[14]!, h, maskS, cache[14] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 14, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B4VPliScal(this.layers[15]!, h, maskS, cache[15] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 15, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B4VPliScal(this.layers[16]!, h, maskS, cache[16] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 16, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorFullQuantG64B4VPliScal(this.layers[17]!, h, maskF, cache[17] as QuantizedKVCache, perLayerSlice(perLayer, 17, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B4VPliScal(this.layers[18]!, h, maskS, cache[18] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 18, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B4VPliScal(this.layers[19]!, h, maskS, cache[19] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 19, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B4VPliScal(this.layers[20]!, h, maskS, cache[20] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 20, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = donorSlidQuantG64B4VPliScal(this.layers[21]!, h, maskS, cache[21] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 21, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    const r22 = donorSlidQuantG64B4VPliScalExp(this.layers[22]!, h, maskS, cache[22] as RotatingQuantizedKVCache, perLayerSlice(perLayer, 22, L, this.perLayerWidth));
+    h.dispose();
+    h = r22.h;
+    const r23 = donorFullQuantG64B4VPliScalExp(this.layers[23]!, h, maskF, cache[23] as QuantizedKVCache, perLayerSlice(perLayer, 23, L, this.perLayerWidth));
+    h.dispose();
+    h = r23.h;
+    next = sharerSlidQuantG64B4VPliScal(this.layers[24]!, h, maskS, r22.shared, perLayerSlice(perLayer, 24, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerSlidQuantG64B4VPliScal(this.layers[25]!, h, maskS, r22.shared, perLayerSlice(perLayer, 25, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerSlidQuantG64B4VPliScal(this.layers[26]!, h, maskS, r22.shared, perLayerSlice(perLayer, 26, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerSlidQuantG64B4VPliScal(this.layers[27]!, h, maskS, r22.shared, perLayerSlice(perLayer, 27, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerSlidQuantG64B4VPliScal(this.layers[28]!, h, maskS, r22.shared, perLayerSlice(perLayer, 28, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerFullQuantG64B4VPliScal(this.layers[29]!, h, maskF, r23.shared, perLayerSlice(perLayer, 29, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerSlidQuantG64B4VPliScal(this.layers[30]!, h, maskS, r22.shared, perLayerSlice(perLayer, 30, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerSlidQuantG64B4VPliScal(this.layers[31]!, h, maskS, r22.shared, perLayerSlice(perLayer, 31, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerSlidQuantG64B4VPliScal(this.layers[32]!, h, maskS, r22.shared, perLayerSlice(perLayer, 32, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerSlidQuantG64B4VPliScal(this.layers[33]!, h, maskS, r22.shared, perLayerSlice(perLayer, 33, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerSlidQuantG64B4VPliScal(this.layers[34]!, h, maskS, r22.shared, perLayerSlice(perLayer, 34, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerFullQuantG64B4VPliScal(this.layers[35]!, h, maskF, r23.shared, perLayerSlice(perLayer, 35, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerSlidQuantG64B4VPliScal(this.layers[36]!, h, maskS, r22.shared, perLayerSlice(perLayer, 36, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerSlidQuantG64B4VPliScal(this.layers[37]!, h, maskS, r22.shared, perLayerSlice(perLayer, 37, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerSlidQuantG64B4VPliScal(this.layers[38]!, h, maskS, r22.shared, perLayerSlice(perLayer, 38, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerSlidQuantG64B4VPliScal(this.layers[39]!, h, maskS, r22.shared, perLayerSlice(perLayer, 39, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerSlidQuantG64B4VPliScal(this.layers[40]!, h, maskS, r22.shared, perLayerSlice(perLayer, 40, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    next = sharerFullQuantG64B4VPliScal(this.layers[41]!, h, maskF, r23.shared, perLayerSlice(perLayer, 41, L, this.perLayerWidth));
+    h.dispose();
+    h = next;
+    disposeSharedKv(r22.shared);
+    disposeSharedKv(r23.shared);
+    perLayer.dispose();
+    maskS.arr?.dispose();
+    maskF.arr?.dispose();
+    return disposing(h, this.finalNorm.forward(h));
+  }
+}