mlx-bun 0.4.0 → 1.0.0

package/src/model/diffusion-gemma.ts

@@ -1,1026 +0,0 @@
-import { cleanupFailure, disposeResources } from "../engine/resources";
-// DiffusionGemma-26B-A4B-it (model_type "diffusion_gemma") — the first
-// NON-autoregressive model in the codebase. Instead of a left-to-right AR loop
-// it runs an ENCODER prefill over the prompt (builds a KV cache) and then a
-// DECODER pass over a fixed 256-token "canvas" that it denoises over a few
-// steps. This file is the D1 static graph: one full forward
-//   encoder(prompt) -> cache ;  decoder(canvas, cache) -> hidden -> tied head
-// held bit-exact (KL/argmax where the quantized GEMV is 1-ULP off) against the
-// mlx-optiq reference (optiq IS the oracle here — stock mlx-lm can't load it).
-//
-// Ported verbatim from optiq/vlm/_mlxvlm/models/diffusion_gemma/language.py.
-// Reuses mlx-bun's quantized primitives (QuantizedLinear/Embedding/SwitchLinear,
-// RMSNorm, KVCache/RotatingKVCache) from gemma4-base. Architecture deltas vs
-// gemma4 (verified against the checkpoint, see docs/design/generic-model-support.md):
-//   - attention scale=1.0, NO attention softcap (only final-logit softcap 30 fp32)
-//   - QK/V-norm POST-proj / PRE-RoPE; v_norm is RMSNormNoScale; V gets no RoPE
-//   - full-attention layers (5,11,17,23,29) reuse k as v (no v_proj), hd=512 kv=2,
-//     partial-rotary 0.25; sliding layers hd=256 kv=8
-//   - per layer: parallel dense-MLP + 128-expert top-8 MoE, summed via 7 RMSNorms
-//     and a per-layer `layer_scalar`
-//   - experts use a FUSED gate_up_proj SwitchLinear (split at moe_intermediate)
-//   - plain nn.RMSNorm everywhere (NO Gemma 1+weight offset) — weights used as-is
-//   - tied LM head via the 4-bit QuantizedEmbedding.asLinear
-//   - SelfConditioning MLP wraps the canvas embeddings every step (zero signal on
-//     the first step still applies its RMSNormNoScale post-norm)
-
-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 {
-  disposing,
-  KVCache,
-  LoraState,
-  QuantizedEmbedding,
-  QuantizedLinear,
-  QuantizedSwitchLinear,
-  RMSNorm,
-  RotatingKVCache,
-  type Cache,
-} from "./gemma4-base";
-import { DiffusionVisionTower } from "../vision/diffusion-vision";
-
-/** A mask in the shape ops.sdpa wants: a mode string + optional bool array.
- *  `null` array with mode "" means "no mask" (full attention). */
-interface SdpaMask {
-  mode: "" | "causal" | "array";
-  arr: MlxArray | null;
-}
-const NO_MASK: SdpaMask = { mode: "", arr: null };
-
-/** KVCache / RotatingKVCache both expose temporalView() (chronological K/V
- *  sliced to offset); the base Cache interface does not declare it. The decoder
- *  only ever uses those two concrete caches. */
-type TemporalCache = Cache & { temporalView(): [MlxArray, MlxArray] };
-
-/** geglu(gate, x) = gelu_approx(gate) * x  (reference language.py:geglu). */
-function geglu(gate: MlxArray, x: MlxArray): MlxArray {
-  const act = ops.geluApprox(gate);
-  const out = ops.mul(act, x);
-  act.dispose();
-  return out;
-}
-
-/** Fused fp32-upcast tanh softcap: tanh(x_fp32 / cap) * cap. Stays f32 (the
- *  reference make_compiled_softcap returns f32; the golden is dumped f32). */
-function softcapFp32(logits: MlxArray, cap: number): MlxArray {
-  const xf = logits.astype(Dtype.float32);
-  const capArr = ops.scalarLike(cap, xf);
-  const scaled = ops.div(xf, capArr);
-  xf.dispose();
-  const t = ops.tanh(scaled);
-  scaled.dispose();
-  const out = ops.mul(t, capArr);
-  t.dispose();
-  capArr.dispose();
-  return out;
-}
-
-/** Dense MLP: down(geglu(gate(x), up(x))). */
-class DiffMLP {
-  readonly gate: QuantizedLinear;
-  readonly up: QuantizedLinear;
-  readonly down: QuantizedLinear;
-  constructor(weights: Weights, config: ModelConfig, prefix: string) {
-    this.gate = QuantizedLinear.load(weights, `${prefix}.gate_proj`, config);
-    this.up = QuantizedLinear.load(weights, `${prefix}.up_proj`, config);
-    this.down = QuantizedLinear.load(weights, `${prefix}.down_proj`, config);
-  }
-  forward(x: MlxArray): MlxArray {
-    const g = this.gate.forward(x);
-    const u = this.up.forward(x);
-    const m = geglu(g, u);
-    g.dispose();
-    u.dispose();
-    const out = this.down.forward(m);
-    m.dispose();
-    return out;
-  }
-}
-
-/** Router: pre-projection RMSNorm (no-scale) * scale * hidden**-0.5, top-8
- *  argpartition, softmax(precise) * per_expert_scale[idx]. */
-class DiffRouter {
-  readonly proj: QuantizedLinear;
-  readonly scale: MlxArray; // per-hidden-dim scale (applied AFTER the no-scale norm)
-  readonly rootSize: number; // hidden**-0.5
-  readonly perExpertScale: MlxArray;
-  readonly eps: number;
-  readonly numExperts: number;
-  readonly topK: number;
-  constructor(weights: Weights, config: ModelConfig, prefix: string) {
-    const t = config.text;
-    this.proj = QuantizedLinear.load(weights, `${prefix}.proj`, config);
-    this.scale = weights.tensor(`${prefix}.scale`);
-    this.rootSize = Math.pow(t.hiddenSize, -0.5);
-    this.perExpertScale = weights.tensor(`${prefix}.per_expert_scale`);
-    this.eps = t.rmsNormEps;
-    this.numExperts = t.numExperts;
-    this.topK = t.topKExperts;
-  }
-  /** x: [T, H] -> { indices: [T, k], weights: [T, k] }. */
-  forward(x: MlxArray): { indices: MlxArray; weights: MlxArray } {
-    // Match the reference EXACTLY (do NOT fold the post-scale into the norm
-    // weight — that changes the bf16 rounding and shifts the softmax weights):
-    //   x = rms_norm(x, None, eps); x = x * scale * hidden**-0.5
-    let normed = ops.rmsNorm(x, null, this.eps);
-    normed = disposing(normed, ops.mul(normed, this.scale));
-    normed = disposing(normed, ops.mulScalar(normed, this.rootSize));
-    const scores = this.proj.forward(normed);
-    normed.dispose();
-    const part = ops.argpartitionAxis(scores, this.numExperts - this.topK, -1);
-    const [T] = scores.shape as [number, number];
-    const indices = part.slice([0, this.numExperts - this.topK], [T, this.numExperts]);
-    part.dispose();
-    let w = ops.takeAlongAxis(scores, indices, -1);
-    scores.dispose();
-    w = disposing(w, ops.softmaxAxis(w, -1, true));
-    const gathered = ops.takeAxis(this.perExpertScale, indices, 0);
-    w = disposing(w, ops.mul(w, gathered));
-    gathered.dispose();
-    return { indices, weights: w };
-  }
-}
-
-/** 128-expert top-8 MoE with a FUSED gate_up SwitchLinear (split at
- *  moe_intermediate_size) + a down SwitchLinear. _gather_sort / _scatter_unsort
- *  exactly as mlx_lm.models.switch_layers (sort threshold idx.size >= 64). */
-class DiffExperts {
-  readonly gateUp: QuantizedSwitchLinear;
-  readonly down: QuantizedSwitchLinear;
-  readonly hiddenDims: number; // moe_intermediate_size (split point)
-  constructor(weights: Weights, config: ModelConfig, prefix: string) {
-    this.gateUp = QuantizedSwitchLinear.load(weights, `${prefix}.gate_up_proj`, config);
-    this.down = QuantizedSwitchLinear.load(weights, `${prefix}.down_proj`, config);
-    this.hiddenDims = config.text.moeIntermediateSize;
-  }
-  /** x: [T, H], indices/weights: [T, k] -> [T, H]. */
-  forward(x: MlxArray, indices: MlxArray, weights: MlxArray): MlxArray {
-    let h = ops.expandDims(x, -2); // [T, 1, H]
-    h = disposing(h, ops.expandDims(h, -3)); // [T, 1, 1, H]
-    const doSort = indices.size >= 64;
-    let idx = indices;
-    let order: MlxArray | null = null;
-    let invOrder: MlxArray | null = null;
-    if (doSort) {
-      const k = indices.shape[indices.ndim - 1]!;
-      const idxFlat = ops.reshape(indices, [indices.size]);
-      order = ops.argsortAxis(idxFlat, 0);
-      invOrder = ops.argsortAxis(order, 0);
-      const kS = ops.scalarLike(k, order);
-      const rowIdx = ops.floorDivide(order, kS);
-      kS.dispose();
-      const H = h.shape[h.ndim - 1]!;
-      const flat = ops.reshape(h, [-1, 1, H]);
-      h.dispose();
-      h = ops.takeAxis(flat, rowIdx, 0);
-      flat.dispose();
-      rowIdx.dispose();
-      idx = ops.takeAxis(idxFlat, order, 0);
-      idxFlat.dispose();
-    }
-
-    const gateUp = this.gateUp.forward(h, idx, doSort);
-    h.dispose();
-    const parts = ops.split(gateUp, [this.hiddenDims], -1);
-    gateUp.dispose();
-    const gate = parts[0]!;
-    const up = parts[1]!;
-    const mid = geglu(gate, up);
-    gate.dispose();
-    up.dispose();
-    let y = this.down.forward(mid, idx, doSort);
-    mid.dispose();
-    if (idx !== indices) idx.dispose();
-
-    if (doSort) {
-      y = disposing(y, ops.takeAxis(y, invOrder!, 0));
-      invOrder!.dispose();
-      order!.dispose();
-      const shape = [...indices.shape, 1, y.shape[y.ndim - 1]!]; // [T, k, 1, H]
-      y = disposing(y, ops.reshape(y, shape));
-    }
-    // squeeze(-2): [T, k, 1, H] -> [T, k, H]
-    {
-      const shape = y.shape;
-      shape.splice(shape.length - 2, 1);
-      y = disposing(y, ops.reshape(y, shape));
-    }
-    // (y * weights[..., None]).sum(axis=-2) -> [T, H]
-    const w = ops.expandDims(weights, -1);
-    const wy = ops.mul(y, w);
-    w.dispose();
-    y.dispose();
-    return disposing(wy, ops.sumAxis(wy, -2, false));
-  }
-}
-
-/** SelfConditioning: post_norm(inputs_embeds + down(geglu(gate(pre_norm(signal)),
- *  up(pre_norm(signal))))). pre_norm is RMSNorm, post_norm is RMSNormNoScale. */
-class SelfConditioning {
-  readonly preNorm: RMSNorm;
-  readonly postNorm: RMSNorm; // RMSNormNoScale (null weight)
-  readonly gate: QuantizedLinear;
-  readonly up: QuantizedLinear;
-  readonly down: QuantizedLinear;
-  constructor(weights: Weights, config: ModelConfig, prefix: string) {
-    const eps = config.text.rmsNormEps;
-    this.preNorm = new RMSNorm(weights.tensor(`${prefix}.pre_norm.weight`), eps);
-    this.postNorm = new RMSNorm(null, eps);
-    this.gate = QuantizedLinear.load(weights, `${prefix}.gate_proj`, config);
-    this.up = QuantizedLinear.load(weights, `${prefix}.up_proj`, config);
-    this.down = QuantizedLinear.load(weights, `${prefix}.down_proj`, config);
-  }
-  forward(inputsEmbeds: MlxArray, signal: MlxArray): MlxArray {
-    const normed = this.preNorm.forward(signal);
-    const g = this.gate.forward(normed);
-    const u = this.up.forward(normed);
-    normed.dispose();
-    const m = geglu(g, u);
-    g.dispose();
-    u.dispose();
-    const s = this.down.forward(m);
-    m.dispose();
-    const sum = ops.add(inputsEmbeds, s);
-    s.dispose();
-    const out = this.postNorm.forward(sum);
-    sum.dispose();
-    return out;
-  }
-}
-
-/** Attention with the encoder (decoder=false, cache.updateAndFetch) and decoder
- *  (decoder=true, reads encoder K/V from cache temporal state and concatenates
- *  the canvas K/V) paths. scale=1.0, no softcap. */
-class DiffAttention {
-  readonly isSliding: boolean;
-  readonly headDim: number;
-  readonly nHeads: number;
-  readonly nKvHeads: number;
-  readonly slidingWindow: number;
-  readonly qProj: QuantizedLinear;
-  readonly kProj: QuantizedLinear;
-  readonly vProj: QuantizedLinear | null;
-  readonly oProj: QuantizedLinear;
-  readonly qNorm: RMSNorm;
-  readonly kNorm: RMSNorm;
-  readonly vNorm: RMSNorm; // RMSNormNoScale (null weight)
-  readonly ropeBase: number | null;
-  readonly ropeFreqs: MlxArray | null;
-
-  constructor(weights: Weights, config: ModelConfig, prefix: string, layerType: string) {
-    const t = config.text;
-    this.isSliding = layerType === "sliding_attention";
-    this.slidingWindow = t.slidingWindow;
-    this.headDim = this.isSliding ? t.headDim : t.globalHeadDim;
-    this.nHeads = t.numAttentionHeads;
-    this.nKvHeads = this.isSliding ? t.numKeyValueHeads : t.numGlobalKeyValueHeads;
-
-    this.qProj = QuantizedLinear.load(weights, `${prefix}.q_proj`, config);
-    this.kProj = QuantizedLinear.load(weights, `${prefix}.k_proj`, config);
-    // Full-attention layers have no v_proj — they reuse the (pre-norm) k tensor.
-    this.vProj = this.isSliding
-      ? QuantizedLinear.load(weights, `${prefix}.v_proj`, config)
-      : null;
-    this.oProj = QuantizedLinear.load(weights, `${prefix}.o_proj`, config);
-    this.qNorm = new RMSNorm(weights.tensor(`${prefix}.q_norm.weight`), t.rmsNormEps);
-    this.kNorm = new RMSNorm(weights.tensor(`${prefix}.k_norm.weight`), t.rmsNormEps);
-    this.vNorm = new RMSNorm(null, t.rmsNormEps);
-
-    // RoPE: default (sliding, theta 1e4) or proportional (full, partial 0.25,
-    // theta 1e6). Proportional rotates only floor(head_dim*factor) dims; the
-    // unrotated tail is freq=Infinity (cos=1, sin=0 → identity). Mirrors
-    // gemma4.ts initialize_rope.
-    const rp = t.ropeParameters[this.isSliding ? "sliding_attention" : "full_attention"]!;
-    if (rp.ropeType === "default") {
-      this.ropeBase = rp.ropeTheta;
-      this.ropeFreqs = null;
-    } else {
-      const rotated = Math.floor(this.headDim * rp.partialRotaryFactor);
-      const n = this.headDim / 2;
-      const exponentsRaw = ops.arange(0, rotated, 2, Dtype.float32);
-      const dims = ops.scalarLike(this.headDim, exponentsRaw);
-      const exponents = ops.div(exponentsRaw, dims);
-      exponentsRaw.dispose();
-      dims.dispose();
-      const base = ops.scalarLike(rp.ropeTheta, exponents);
-      let rotFreqs = ops.pow(base, exponents);
-      base.dispose();
-      exponents.dispose();
-      if (rp.factor !== 1.0) {
-        const f = ops.scalarLike(rp.factor, rotFreqs);
-        rotFreqs = disposing(rotFreqs, ops.mul(f, rotFreqs));
-        f.dispose();
-      }
-      this.ropeBase = null;
-      const tailLen = n - rotated / 2;
-      if (tailLen > 0) {
-        const tail = MlxArray.fromFloat32(new Float32Array(tailLen).fill(Infinity), [tailLen]);
-        this.ropeFreqs = ops.concatAxis([rotFreqs, tail], 0);
-        rotFreqs.dispose();
-        tail.dispose();
-      } else {
-        this.ropeFreqs = rotFreqs;
-      }
-    }
-  }
-
-  #rope(x: MlxArray, offset: number): MlxArray {
-    return ops.rope(x, this.headDim, this.ropeBase, offset, this.ropeFreqs);
-  }
-
-  forward(
-    x: MlxArray,
-    mask: SdpaMask,
-    cache: Cache | null,
-    decoder: boolean,
-    offset: number,
-  ): MlxArray {
-    const [B, L] = x.shape as [number, number, number];
-
-    const qFlat = this.qProj.forward(x);
-    let q = ops.reshape(qFlat, [B, L, this.nHeads, this.headDim]);
-    qFlat.dispose();
-    q = disposing(q, this.qNorm.forward(q));
-    q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
-    q = disposing(q, this.#rope(q, offset));
-
-    const kFlat = this.kProj.forward(x);
-    const kReshaped = ops.reshape(kFlat, [B, L, this.nKvHeads, this.headDim]);
-    kFlat.dispose();
-    // v reuses the pre-norm k tensor on full layers (v_proj is None there).
-    let vReshaped: MlxArray;
-    if (this.vProj) {
-      const vFlat = this.vProj.forward(x);
-      vReshaped = ops.reshape(vFlat, [B, L, this.nKvHeads, this.headDim]);
-      vFlat.dispose();
-    } else {
-      vReshaped = kReshaped;
-    }
-
-    let k = this.kNorm.forward(kReshaped);
-    k = disposing(k, ops.transposeAxes(k, [0, 2, 1, 3]));
-    k = disposing(k, this.#rope(k, offset));
-    let v = this.vNorm.forward(vReshaped); // no RoPE on V
-    v = disposing(v, ops.transposeAxes(v, [0, 2, 1, 3]));
-    if (vReshaped !== kReshaped) vReshaped.dispose();
-    kReshaped.dispose();
-
-    let maskMode = mask.mode;
-    let maskArr = mask.arr;
-    let maskArrOwned = false;
-
-    if (decoder) {
-      // Read the encoder K/V (chronological, sliced to offset) from the cache
-      // and concatenate the canvas K/V. No cache update on the canvas pass.
-      if (cache) {
-        let [encK, encV] = (cache as TemporalCache).temporalView();
-        if (this.isSliding) {
-          const window = Math.max(this.slidingWindow - 1, 0);
-          const encoderLen = encK.shape[2]!;
-          if (window && encoderLen > window && offset >= encoderLen) {
-            const trimK = encK.slice([0, 0, encoderLen - window, 0], encK.shape);
-            const trimV = encV.slice([0, 0, encoderLen - window, 0], encV.shape);
-            encK.dispose();
-            encV.dispose();
-            encK = trimK;
-            encV = trimV;
-            if (maskArr) {
-              const ks = maskArr.shape;
-              const keep = window + L;
-              const trimmed = maskArr.slice(
-                [0, 0, 0, ks[3]! - keep],
-                ks,
-              );
-              maskArr = trimmed;
-              maskArrOwned = true;
-              maskMode = "array";
-            }
-          }
-        }
-        const kCat = ops.concatAxis([encK, k], 2);
-        const vCat = ops.concatAxis([encV, v], 2);
-        encK.dispose();
-        encV.dispose();
-        k.dispose();
-        v.dispose();
-        k = kCat;
-        v = vCat;
-      }
-    } else if (cache) {
-      const [fk, fv] = cache.updateAndFetch(k, v);
-      k.dispose();
-      v.dispose();
-      k = fk;
-      v = fv;
-    }
-
-    const attn = ops.sdpa(q, k, v, 1.0, maskMode, maskArr);
-    q.dispose();
-    k.dispose();
-    v.dispose();
-    if (maskArrOwned) maskArr!.dispose();
-
-    let out = ops.transposeAxes(attn, [0, 2, 1, 3]);
-    attn.dispose();
-    out = disposing(out, ops.reshape(out, [B, L, this.nHeads * this.headDim]));
-    const o = this.oProj.forward(out);
-    out.dispose();
-    return o;
-  }
-}
-
-/** One decoder layer: attention residual + a PARALLEL dense-MLP + MoE branch,
- *  combined through 7 RMSNorms and scaled by `layer_scalar` (the encoder pass
- *  overrides it with the per-layer encoder scalar). */
-class DiffDecoderLayer {
-  readonly layerType: string;
-  readonly attn: DiffAttention;
-  readonly mlp: DiffMLP;
-  readonly router: DiffRouter;
-  readonly experts: DiffExperts;
-  readonly inputNorm: RMSNorm;
-  readonly postAttnNorm: RMSNorm;
-  readonly preFfNorm: RMSNorm;
-  readonly postFfNorm: RMSNorm;
-  readonly postFfNorm1: RMSNorm;
-  readonly postFfNorm2: RMSNorm;
-  readonly preFfNorm2: RMSNorm;
-  readonly layerScalar: MlxArray;
-
-  constructor(weights: Weights, config: ModelConfig, idx: number) {
-    const t = config.text;
-    const prefix = `model.decoder.layers.${idx}`;
-    this.layerType = t.layerTypes[idx]!;
-    const norm = (n: string) => new RMSNorm(weights.tensor(`${prefix}.${n}.weight`), t.rmsNormEps);
-    this.attn = new DiffAttention(weights, config, `${prefix}.self_attn`, this.layerType);
-    this.mlp = new DiffMLP(weights, config, `${prefix}.mlp`);
-    this.router = new DiffRouter(weights, config, `${prefix}.router`);
-    this.experts = new DiffExperts(weights, config, `${prefix}.experts`);
-    this.inputNorm = norm("input_layernorm");
-    this.postAttnNorm = norm("post_attention_layernorm");
-    this.preFfNorm = norm("pre_feedforward_layernorm");
-    this.postFfNorm = norm("post_feedforward_layernorm");
-    this.postFfNorm1 = norm("post_feedforward_layernorm_1");
-    this.postFfNorm2 = norm("post_feedforward_layernorm_2");
-    this.preFfNorm2 = norm("pre_feedforward_layernorm_2");
-    this.layerScalar = weights.tensor(`${prefix}.layer_scalar`);
-  }
-
-  forward(
-    x: MlxArray,
-    mask: SdpaMask,
-    cache: Cache | null,
-    decoder: boolean,
-    offset: number,
-    layerScalar: MlxArray | null,
-  ): MlxArray {
-    // attention
-    let residual = x;
-    let h = this.inputNorm.forward(x);
-    h = disposing(h, this.attn.forward(h, mask, cache, decoder, offset));
-    h = disposing(h, this.postAttnNorm.forward(h));
-    h = disposing(h, ops.add(residual, h));
-
-    // parallel dense-MLP + MoE feedforward
-    residual = h;
-    const dense = this.preFfNorm.forward(h);
-    let h1 = this.mlp.forward(dense);
-    dense.dispose();
-    h1 = disposing(h1, this.postFfNorm1.forward(h1));
-
-    const shape = residual.shape;
-    const hidden = shape[shape.length - 1]!;
-    const flat = ops.reshape(residual, [-1, hidden]);
-    const { indices: rawIdx, weights: topw } = this.router.forward(flat);
-    // stop_gradient the routing indices so the MoE backward (gather_qmm vjp)
-    // doesn't try to differentiate the non-differentiable top-k indices (the
-    // reference does this in training mode). Identity in the forward, so D1/D2
-    // stay bit-exact; required for D5 LoRA grads to flow through the experts.
-    const indices = ops.stopGradient(rawIdx);
-    rawIdx.dispose();
-    let h2 = this.preFfNorm2.forward(flat);
-    flat.dispose();
-    h2 = disposing(h2, this.experts.forward(h2, indices, topw));
-    indices.dispose();
-    topw.dispose();
-    h2 = disposing(h2, ops.reshape(h2, shape));
-    h2 = disposing(h2, this.postFfNorm2.forward(h2));
-
-    const sum = ops.add(h1, h2);
-    h1.dispose();
-    h2.dispose();
-    h = this.postFfNorm.forward(sum);
-    sum.dispose();
-    h = disposing(h, ops.add(residual, h));
-    residual.dispose();
-
-    const scalar = layerScalar ?? this.layerScalar;
-    h = disposing(h, ops.mul(h, scalar));
-    return h;
-  }
-}
-
-export class DiffusionGemmaModel {
-  readonly config: ModelConfig;
-  readonly weightsBytes: number;
-  readonly embed: QuantizedEmbedding;
-  readonly embedScale: number;
-  readonly layers: DiffDecoderLayer[];
-  readonly finalNorm: RMSNorm;
-  readonly selfConditioning: SelfConditioning;
-  readonly encoderLayerScalars: MlxArray[];
-  readonly canvasLength: number;
-  readonly softcap: number;
-  readonly slidingWindow: number;
-  // Shared RuntimeModel surface. LoRA mounts onto the (quantized) decoder
-  // linears like the AR models; the diffusion engine activates it per request.
-  readonly loraState = new LoraState();
-  /** Dedicated DiffusionGemma SigLIP vision tower (image-text-to-text). */
-  readonly visionTower: DiffusionVisionTower | null;
-  readonly imageTokenId: number;
-
-  constructor(weights: Weights, config: ModelConfig) {
-    const t = config.text;
-    this.config = config;
-    this.weightsBytes = [...weights.shards.files.values()].reduce((a, f) => a + f.mmap.size, 0);
-    this.embed = QuantizedEmbedding.load(weights, "model.decoder.embed_tokens", config);
-    this.embedScale = Math.pow(t.hiddenSize, 0.5);
-    this.layers = Array.from(
-      { length: t.numHiddenLayers },
-      (_, i) => new DiffDecoderLayer(weights, config, i),
-    );
-    this.finalNorm = new RMSNorm(weights.tensor("model.decoder.norm.weight"), t.rmsNormEps);
-    this.selfConditioning = new SelfConditioning(weights, config, "model.decoder.self_conditioning");
-    this.encoderLayerScalars = this.layers.map((_, i) =>
-      weights.tensor(`model.encoder.language_model.layers.${i}.layer_scalar`),
-    );
-    this.canvasLength = t.canvasLength ?? 256;
-    this.softcap = t.finalLogitSoftcapping ?? 30.0;
-    this.slidingWindow = t.slidingWindow;
-    const raw = (config as unknown as { raw: Record<string, any> }).raw;
-    this.imageTokenId = (raw.image_token_id as number) ?? 258880;
-    this.visionTower = raw.vision_config ? new DiffusionVisionTower(weights, config) : null;
-  }
-
-  /** Build the merged prompt embeddings for an image+text turn: embed the text
-   *  ids (image positions → pad) * embed_scale, then scatter the SigLIP vision
-   *  features in at the contiguous <|image|> run. `pixels` is channel-first
-   *  [1,3,H,W]. Reference EncoderModel._embed_inputs + get_image_features. */
-  #embedInputsVision(splicedIds: number[], pixels: MlxArray): MlxArray {
-    if (!this.visionTower) throw new Error("this checkpoint has no vision tower");
-    const padId = 0; // pad_token_id (reference where(vision_mask, pad, ids))
-    const llm = splicedIds.map((id) => (id === this.imageTokenId ? padId : id));
-    const ids = MlxArray.fromInt32(Int32Array.from(llm), [1, llm.length]);
-    let embeds = this.embed.encode(ids);
-    ids.dispose();
-    embeds = disposing(embeds, ops.mulScalar(embeds, this.embedScale)); // [1, L, hidden]
-
-    const feats = this.visionTower.getImageFeatures(pixels); // [1, numSoft, hidden]
-    const featsCast = feats.astype(embeds.dtype);
-    feats.dispose();
-    // image tokens are one contiguous run (boi + image*N + eoi)
-    const start = splicedIds.indexOf(this.imageTokenId);
-    const count = splicedIds.filter((id) => id === this.imageTokenId).length;
-    const hidden = embeds.shape[embeds.shape.length - 1]!;
-    embeds = disposing(
-      embeds,
-      ops.sliceUpdate(embeds, featsCast, [0, start, 0], [1, start + count, hidden]),
-    );
-    featsCast.dispose();
-    return embeds;
-  }
-
-  /** Encoder attention mask for the vision prompt: causal, OR'd with the
-   *  bidirectional image-block overlay (tokens in the same <|image|> block
-   *  attend both ways). For N < sliding_window the sliding window is a no-op, so
-   *  one mask serves all layers. Reference _make_encoder_masks + _vision_block_overlay. */
-  #encoderVisionMask(splicedIds: number[]): MlxArray {
-    const N = splicedIds.length;
-    const isVision = splicedIds.map((id) => id === this.imageTokenId);
-    const bits = new Int32Array(N * N);
-    for (let i = 0; i < N; i++)
-      for (let j = 0; j < N; j++)
-        bits[i * N + j] = i >= j || (isVision[i] && isVision[j]) ? 1 : 0;
-    return MlxArray.fromInt32(bits, [1, 1, N, N]).astype(Dtype.bool);
-  }
-
-  /** Image-text-to-text encoder prefill: merged embeddings + the causal/overlay
-   *  mask through all layers (decoder=false), filling `cache`. The canvas decoder
-   *  then reads it like the text path. */
-  encodeVisionPrefill(splicedIds: number[], pixels: MlxArray, cache: Cache[]): void {
-    let h = this.#embedInputsVision(splicedIds, pixels); // [1, L, hidden] (*embed_scale)
-    const maskArr = this.#encoderVisionMask(splicedIds);
-    const mask: SdpaMask = { mode: "array", arr: maskArr };
-    for (let i = 0; i < this.layers.length; i++) {
-      const layer = this.layers[i]!;
-      const next = layer.forward(h, mask, cache[i]!, false, 0, this.encoderLayerScalars[i]!);
-      h.dispose();
-      h = next;
-    }
-    maskArr.dispose();
-    h.dispose();
-  }
-
-  /** Prefill from a pre-spliced image+text prompt → populated cache. */
-  prefillVision(splicedIds: number[], pixels: MlxArray): Cache[] {
-    const cache = this.makeCache();
-    try { this.encodeVisionPrefill(splicedIds, pixels, cache); return cache; }
-    catch (error) { cleanupFailure(error, () => disposeResources(cache)); }
-  }
-
-  /** One KVCache (full) / RotatingKVCache (sliding) per layer — the encoder
-   *  prefill fills these, the decoder canvas pass reads them. */
-  makeCache(): Cache[] {
-    return this.layers.map((l) =>
-      l.layerType === "sliding_attention"
-        ? new RotatingKVCache(this.slidingWindow)
-        : new KVCache(),
-    );
-  }
-
-  /** Encoder prefill: run all layers over the prompt (decoder=false), filling
-   *  the cache. Text-only (no vision merge yet — D3). Discards the hidden. */
-  encodePrompt(promptIds: number[], cache: Cache[]): void {
-    const ids = MlxArray.fromInt32(Int32Array.from(promptIds), [1, promptIds.length]);
-    let h = this.embed.encode(ids);
-    ids.dispose();
-    h = disposing(h, ops.mulScalar(h, this.embedScale));
-    const N = promptIds.length;
-    for (let i = 0; i < this.layers.length; i++) {
-      const layer = this.layers[i]!;
-      const window = layer.layerType === "sliding_attention" ? this.slidingWindow : null;
-      const m = cache[i]!.makeMask(N, window) as SdpaMask;
-      const next = layer.forward(h, m, cache[i]!, false, 0, this.encoderLayerScalars[i]!);
-      if (m.arr) m.arr.dispose();
-      h.dispose();
-      h = next;
-    }
-    h.dispose();
-  }
-
-  /** Diagnostic: run the encoder, invoking `onLayer(i, h)` after each layer
-   *  (h is the raw layer output, pre-final-norm) and `onEmbed(h)` on the
-   *  embedding. Returns the final pre-norm hidden. Fills `cache`. */
-  encoderTrace(
-    promptIds: number[],
-    cache: Cache[],
-    onEmbed: (h: MlxArray) => void,
-    onLayer: (i: number, h: MlxArray) => void,
-  ): MlxArray {
-    const ids = MlxArray.fromInt32(Int32Array.from(promptIds), [1, promptIds.length]);
-    let h = this.embed.encode(ids);
-    ids.dispose();
-    h = disposing(h, ops.mulScalar(h, this.embedScale));
-    onEmbed(h);
-    const N = promptIds.length;
-    for (let i = 0; i < this.layers.length; i++) {
-      const layer = this.layers[i]!;
-      const window = layer.layerType === "sliding_attention" ? this.slidingWindow : null;
-      const m = cache[i]!.makeMask(N, window) as SdpaMask;
-      const next = layer.forward(h, m, cache[i]!, false, 0, this.encoderLayerScalars[i]!);
-      if (m.arr) m.arr.dispose();
-      h.dispose();
-      h = next;
-      onLayer(i, h);
-    }
-    return h;
-  }
-
-  /** Diagnostic: run encoder layer 0 sub-components (attention vs parallel
-   *  dense/MoE) and dump each via `cb(name, h)`. */
-  diagLayer0(promptIds: number[], cache: Cache[], cb: (name: string, h: MlxArray) => void): void {
-    const ids = MlxArray.fromInt32(Int32Array.from(promptIds), [1, promptIds.length]);
-    let emb = this.embed.encode(ids);
-    ids.dispose();
-    emb = disposing(emb, ops.mulScalar(emb, this.embedScale));
-    const layer = this.layers[0]!;
-    const N = promptIds.length;
-    const window = layer.layerType === "sliding_attention" ? this.slidingWindow : null;
-    const m = cache[0]!.makeMask(N, window) as SdpaMask;
-    const ls0 = this.encoderLayerScalars[0]!;
-
-    const residual = emb;
-    const hn = layer.inputNorm.forward(emb);
-    const attnOut = layer.attn.forward(hn, m, cache[0]!, false, 0);
-    hn.dispose();
-    if (m.arr) m.arr.dispose();
-    cb("attn_out", attnOut);
-    const postAttn = layer.postAttnNorm.forward(attnOut);
-    attnOut.dispose();
-    const hMid = ops.add(residual, postAttn);
-    postAttn.dispose();
-    cb("h_mid", hMid);
-
-    const dense = layer.preFfNorm.forward(hMid);
-    let h1 = layer.mlp.forward(dense);
-    dense.dispose();
-    h1 = disposing(h1, layer.postFfNorm1.forward(h1));
-    cb("h1", h1);
-    h1.dispose();
-
-    const shape = hMid.shape;
-    const hidden = shape[shape.length - 1]!;
-    const flat = ops.reshape(hMid, [-1, hidden]);
-    const { indices, weights } = layer.router.forward(flat);
-    cb("router_idx", indices.astype(Dtype.float32));
-    cb("router_w", weights.astype(Dtype.float32));
-    let h2 = layer.preFfNorm2.forward(flat);
-    flat.dispose();
-    h2 = disposing(h2, layer.experts.forward(h2, indices, weights));
-    indices.dispose();
-    weights.dispose();
-    h2 = disposing(h2, ops.reshape(h2, shape));
-    cb("experts_raw", h2);
-    h2 = disposing(h2, layer.postFfNorm2.forward(h2));
-    cb("h2", h2);
-    h2.dispose();
-    void ls0;
-    emb.dispose();
-    hMid.dispose();
-  }
-
-  /** Diagnostic: encoder prefill then finalNorm(h) — matches the reference
-   *  EncoderModel.__call__ output (decoder.norm(h)) for per-stage localization. */
-  encodePromptHidden(promptIds: number[], cache: Cache[]): MlxArray {
-    const ids = MlxArray.fromInt32(Int32Array.from(promptIds), [1, promptIds.length]);
-    let h = this.embed.encode(ids);
-    ids.dispose();
-    h = disposing(h, ops.mulScalar(h, this.embedScale));
-    const N = promptIds.length;
-    for (let i = 0; i < this.layers.length; i++) {
-      const layer = this.layers[i]!;
-      const window = layer.layerType === "sliding_attention" ? this.slidingWindow : null;
-      const m = cache[i]!.makeMask(N, window) as SdpaMask;
-      const next = layer.forward(h, m, cache[i]!, false, 0, this.encoderLayerScalars[i]!);
-      if (m.arr) m.arr.dispose();
-      h.dispose();
-      h = next;
-    }
-    return disposing(h, this.finalNorm.forward(h));
-  }
-
-  /** Embed the canvas + self-conditioning. For the first step both signals are
-   *  None → zero signal (which still applies SelfConditioning's post RMSNorm). */
-  #embedCanvas(
-    canvasIds: MlxArray,
-    scLogits: MlxArray | null,
-    scEmbeddings: MlxArray | null,
-  ): MlxArray {
-    let inputsEmbeds = this.embed.encode(canvasIds);
-    inputsEmbeds = disposing(inputsEmbeds, ops.mulScalar(inputsEmbeds, this.embedScale));
-    let soft: MlxArray;
-    let softOwned = true;
-    if (scEmbeddings) {
-      soft = scEmbeddings.astype(inputsEmbeds.dtype);
-    } else if (scLogits === null) {
-      soft = ops.zeros(inputsEmbeds.shape, inputsEmbeds.dtype);
-    } else {
-      // probs @ embed_tokens.weight (quantized, transpose=false), * embed_scale.
-      const probs = ops.softmaxAxis(scLogits, -1, true);
-      const probsCast = probs.astype(inputsEmbeds.dtype);
-      probs.dispose();
-      const se = ops.quantizedMatmul(
-        probsCast,
-        this.embed.w,
-        this.embed.scales,
-        this.embed.biases,
-        this.embed.spec,
-        false,
-      );
-      probsCast.dispose();
-      soft = ops.mulScalar(se, this.embedScale);
-      se.dispose();
-    }
-    const out = this.selfConditioning.forward(inputsEmbeds, soft);
-    inputsEmbeds.dispose();
-    if (softOwned) soft.dispose();
-    return out;
-  }
-
-  /** Build the per-layer-type decoder masks (reference _make_decoder_masks).
-   *  Returns a map layerType -> SdpaMask. For the dynamic caches used here
-   *  encoder_len == valid_encoder_len always, so the only non-None mask is the
-   *  sliding-window case once the prompt exceeds sliding_window-1 (1023). */
-  #makeDecoderMasks(canvasLength: number, cache: Cache[]): Map<string, SdpaMask> {
-    const masks = new Map<string, SdpaMask>();
-    const types = new Set(this.layers.map((l) => l.layerType));
-    for (const layerType of types) {
-      const li = this.layers.findIndex((l) => l.layerType === layerType);
-      const c = cache[li]!;
-      const offset = c.offset;
-      // encoder_len = chronological cache length; for dynamic caches this equals
-      // min(offset, maxSize), and valid_encoder_len = min(offset, encoder_len)
-      // == encoder_len, so the "trailing invalid" branches never fire here.
-      const [encK, encV] = (c as TemporalCache).temporalView();
-      const encoderLen = encK.shape[2]!;
-      encK.dispose();
-      encV.dispose();
-      const validEncoderLen = Math.min(offset, encoderLen);
-      const keyLen = encoderLen + canvasLength;
-
-      if (layerType === "full_attention") {
-        masks.set(layerType, NO_MASK); // encoder_len == valid_encoder_len
-        continue;
-      }
-      const windowPrefix = Math.max(this.slidingWindow - 1, 0);
-      if (encoderLen <= windowPrefix) {
-        masks.set(layerType, NO_MASK);
-        continue;
-      }
-      // Sliding + wrapped prompt: keep encoder positions [start, valid) + canvas.
-      const start = Math.max(0, validEncoderLen - windowPrefix);
-      const positions = ops.arange(0, encoderLen, 1, Dtype.int32);
-      const startArr = ops.scalarLike(start, positions);
-      const validArr = ops.scalarLike(validEncoderLen, positions);
-      const geStart = ops.greaterEqual(positions, startArr);
-      const ltValid = ops.less(positions, validArr);
-      const encMask = ops.logicalAnd(geStart, ltValid);
-      positions.dispose();
-      startArr.dispose();
-      validArr.dispose();
-      geStart.dispose();
-      ltValid.dispose();
-      const canvasMask = ops.fromInt32(new Array(canvasLength).fill(1), [canvasLength]).astype(
-        encMask.dtype,
-      );
-      const row = ops.concatAxis([encMask, canvasMask], 0);
-      encMask.dispose();
-      canvasMask.dispose();
-      const reshaped = ops.reshape(row, [1, 1, 1, keyLen]);
-      row.dispose();
-      masks.set(layerType, { mode: "array", arr: reshaped });
-    }
-    return masks;
-  }
-
-  /** Decoder canvas pass over an on-device canvas: embed canvas -> 30 layers
-   *  (decoder=true) -> final norm. Returns hidden [1, canvasLength, hidden].
-   *  The D2 engine drives this per denoising step (canvas stays on-device). */
-  decodeCanvasArr(
-    canvasArr: MlxArray,
-    cache: Cache[],
-    scLogits: MlxArray | null = null,
-    scEmbeddings: MlxArray | null = null,
-  ): MlxArray {
-    const canvasLength = canvasArr.shape[canvasArr.shape.length - 1]!;
-    let h = this.#embedCanvas(canvasArr, scLogits, scEmbeddings);
-    const masks = this.#makeDecoderMasks(canvasLength, cache);
-    const offset = cache[0]!.offset;
-    for (let i = 0; i < this.layers.length; i++) {
-      const layer = this.layers[i]!;
-      const m = masks.get(layer.layerType)!;
-      const next = layer.forward(h, m, cache[i]!, true, offset, null);
-      h.dispose();
-      h = next;
-    }
-    for (const m of masks.values()) if (m.arr) m.arr.dispose();
-    h = disposing(h, this.finalNorm.forward(h));
-    return h;
-  }
-
-  /** Decoder canvas pass from host ids (D1 path). */
-  decodeCanvas(
-    canvasIds: number[],
-    cache: Cache[],
-    scLogits: MlxArray | null = null,
-    scEmbeddings: MlxArray | null = null,
-  ): MlxArray {
-    const ids = MlxArray.fromInt32(Int32Array.from(canvasIds), [1, canvasIds.length]);
-    const h = this.decodeCanvasArr(ids, cache, scLogits, scEmbeddings);
-    ids.dispose();
-    return h;
-  }
-
-  // ---- D2 denoising-engine surface ----
-
-  /** Encoder prefill over the prompt -> a fresh populated cache (the denoising
-   *  loop reuses it across all steps of a canvas block). */
-  prefill(promptIds: number[]): Cache[] {
-    const cache = this.makeCache();
-    try { this.encodePrompt(promptIds, cache); return cache; }
-    catch (error) { cleanupFailure(error, () => disposeResources(cache)); }
-  }
-
-  /** One denoising-step decoder pass -> softcapped logits [1, canvas, vocab]. */
-  decoderLogits(canvasArr: MlxArray, cache: Cache[], scEmbeddings: MlxArray | null): MlxArray {
-    const h = this.decodeCanvasArr(canvasArr, cache, null, scEmbeddings);
-    try { return this.logitsFromHidden(h); }
-    finally { h.dispose(); }
-  }
-
-  /** Dequantized embedding table (once per call), for self-conditioning soft
-   *  embeddings — matches optiq's `_diffusion_soft_embedding_weight`. */
-  dequantEmbedWeight(): MlxArray {
-    return ops.dequantize(this.embed.w, this.embed.scales, this.embed.biases, this.embed.spec);
-  }
-
-  /** Self-conditioning soft embeddings for the NEXT step (confidence sampler):
-   *  softmax(processed, precise) @ dequant(embed) * embed_scale. */
-  softEmbeddings(processedLogits: MlxArray, dequantWeight: MlxArray): MlxArray {
-    const probs = ops.softmaxAxis(processedLogits, -1, true);
-    const probsCast = probs.astype(dequantWeight.dtype);
-    probs.dispose();
-    const se = ops.matmul(probsCast, dequantWeight);
-    probsCast.dispose();
-    const cast = se.astype(dequantWeight.dtype);
-    se.dispose();
-    const out = ops.mulScalar(cast, this.embedScale);
-    cast.dispose();
-    return out;
-  }
-
-  /** Re-prefill the encoder over an accepted canvas block (continuation), so the
-   *  next block's decoder sees it as context. Appends to the existing cache. */
-  extendPrefill(canvasArr: MlxArray, cache: Cache[]): void {
-    let h = this.embed.encode(canvasArr);
-    h = disposing(h, ops.mulScalar(h, this.embedScale));
-    const N = canvasArr.shape[canvasArr.shape.length - 1]!;
-    for (let i = 0; i < this.layers.length; i++) {
-      const layer = this.layers[i]!;
-      const window = layer.layerType === "sliding_attention" ? this.slidingWindow : null;
-      const m = cache[i]!.makeMask(N, window) as SdpaMask;
-      const next = layer.forward(h, m, cache[i]!, false, 0, this.encoderLayerScalars[i]!);
-      if (m.arr) m.arr.dispose();
-      h.dispose();
-      h = next;
-    }
-    h.dispose();
-  }
-
-  /** Tied quantized head + fp32 softcap. hidden [1, L, hidden] -> [1, L, vocab]. */
-  logitsFromHidden(h: MlxArray): MlxArray {
-    const logits = this.embed.asLinear(h);
-    const capped = softcapFp32(logits, this.softcap);
-    logits.dispose();
-    return capped;
-  }
-
-  /** D1 entry point: prompt + canvas ids -> softcapped logits [1, canvas, vocab].
-   *  Builds a fresh cache, prefills the encoder, runs one decoder pass. */
-  forwardCanvasLogits(promptIds: number[], canvasIds: number[]): MlxArray {
-    const cache = this.makeCache();
-    try {
-      this.encodePrompt(promptIds, cache);
-      const h = this.decodeCanvas(canvasIds, cache);
-      const logits = this.logitsFromHidden(h);
-      h.dispose();
-      return logits;
-    } finally {
-      for (const c of cache) c.dispose();
-    }
-  }
-
-  /** Training forward over an on-device canvas (D5 LoRA): prompt + canvas array
-   *  -> softcapped logits [1, L, vocab], differentiable w.r.t. mounted LoRA. */
-  forwardCanvasLogitsArr(promptIds: number[], canvasArr: MlxArray): MlxArray {
-    const cache = this.makeCache();
-    try {
-      this.encodePrompt(promptIds, cache);
-      const h = this.decodeCanvasArr(canvasArr, cache);
-      const logits = this.logitsFromHidden(h);
-      h.dispose();
-      return logits;
-    } finally {
-      for (const c of cache) c.dispose();
-    }
-  }
-
-  // ---- AR-only RuntimeModel surface (never called on a diffusion model) ----
-  // DiffusionGemma is non-autoregressive: generate() detects it and routes to
-  // the denoising engine (src/diffusion/diffusion-generate.ts). These exist so
-  // the model fits the shared RuntimeModel union; calling them is a bug.
-  forwardHidden(_ids: MlxArray, _cache: Cache[]): MlxArray {
-    throw new Error("DiffusionGemma is non-autoregressive — use the diffusion engine, not forwardHidden");
-  }
-  forward(_tokens: number[] | MlxArray, _cache: Cache[]): MlxArray {
-    throw new Error("DiffusionGemma is non-autoregressive — use the diffusion engine, not forward");
-  }
-  generate(_promptTokens: number[], _maxTokens: number, _eosIds?: number[]): number[] {
-    throw new Error("DiffusionGemma is non-autoregressive — use diffusionGenerate, not generate");
-  }
-  forwardEmbeddings(_embeds: MlxArray, _cache: Cache[], _bidir: MlxArray | null): MlxArray {
-    throw new Error("DiffusionGemma vision runs through its own encoder merge, not forwardEmbeddings");
-  }
-
-  readonly prefixBase = "model.decoder";
-
-  /** LoRA-able decoder linears (the optiq diffusion DEFAULT_LORA_KEYS: attention
-   *  q/k/v/o + dense MLP gate/up/down per layer; full layers have no v_proj). */
-  loraTargets(): Map<string, QuantizedLinear> {
-    const m = new Map<string, QuantizedLinear>();
-    this.layers.forEach((layer, i) => {
-      const p = `model.decoder.layers.${i}`;
-      m.set(`${p}.self_attn.q_proj`, layer.attn.qProj);
-      m.set(`${p}.self_attn.k_proj`, layer.attn.kProj);
-      if (layer.attn.vProj) m.set(`${p}.self_attn.v_proj`, layer.attn.vProj);
-      m.set(`${p}.self_attn.o_proj`, layer.attn.oProj);
-      m.set(`${p}.mlp.gate_proj`, layer.mlp.gate);
-      m.set(`${p}.mlp.up_proj`, layer.mlp.up);
-      m.set(`${p}.mlp.down_proj`, layer.mlp.down);
-    });
-    return m;
-  }
-}