@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/models/diffusion-gemma/vision.ts

@@ -0,0 +1,549 @@
+// DiffusionGemma vision tower — DEDICATED, parity-exact port of optiq's
+// gemma4 `VisionModel` (optiq/vlm/_mlxvlm/models/gemma4/vision.py) at the
+// DiffusionGemma geometry (hidden 1152, head_dim 72, 16 heads, 27 layers,
+// position_embedding_size 10240, standardize ON). This is its OWN module, not a
+// reuse of the e4b SigLIP path — the two are separate models with separate
+// towers (e4b ships a bf16 optiq_vision.safetensors sidecar; DiffusionGemma's
+// vision is inline-quantized in the main shards). Ported op-for-op so it is
+// bit-exact to the optiq reference.
+//
+// Forward (reference VisionModel.__call__):
+//   pixels [1,3,H,W] → patchify (reshape/transpose, 2*(x-0.5)) → input_proj
+//     + 2D position embedding (gather table[0,x]+table[1,y]) → PAD to
+//     max_patches=2520 (zeros) → 27× transformer blocks (manual-f32 q/k/v RMS,
+//     2D multidimensional RoPE, bidirectional -1e4 mask, scale=1.0, GeGLU) →
+//     position-grouped 3×3 avg-pool to 280 (×√hidden) → trim to the real soft
+//     tokens → standardize (h - std_bias)*std_scale → MultimodalEmbedder
+//     (RMSNormNoScale → projection 1152→2816).
+// Output = `getImageFeatures`: [1, softTokens, textHidden], scattered AS-IS into
+// the prompt embeddings (NO /embed_scale — the diffusion merge adds it once via
+// embed_tokens*embed_scale).
+//
+// The vision linears are 4-bit quantized inline; we run them as QuantizedLinear
+// (== the reference's nn.quantize'd modules), not dequant+matmul.
+
+import type { ModelConfig } from "../../artifacts/config";
+import type { Weights } from "../../artifacts/weights";
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import { disposing } from "../../layers/helpers";
+import { QuantizedLinear } from "../../layers/quantized-linear";
+import { RMSNorm } from "../../layers/normalization";
+import { applyRope2d, rope2dTables, type Rope2dTables } from "../../layers/rope-2d";
+import { parseSiglipConfig, type SiglipVisionConfig } from "../vision/siglip";
+import { decodeImage, resizeBicubic } from "../../input/vision/preprocess";
+
+const VP = "model.encoder.vision_tower.";
+
+/** Expand each `<|image|>` placeholder id into `boi + image*softCount + eoi`,
+ *  matching optiq's `_build_image_inputs`. Returns the spliced ids (the
+ *  diffusion engine's encoder prefills the merged vision features at the
+ *  contiguous image run). */
+export function spliceImageTokens(
+  ids: number[],
+  softCounts: number[],
+  tokenIds: { image: number; boi: number; eoi: number },
+): number[] {
+  const out: number[] = [];
+  let k = 0;
+  for (const id of ids) {
+    if (id === tokenIds.image) {
+      out.push(tokenIds.boi);
+      const n = softCounts[k] ?? softCounts[softCounts.length - 1] ?? 0;
+      for (let i = 0; i < n; i++) out.push(tokenIds.image);
+      out.push(tokenIds.eoi);
+      k++;
+    } else {
+      out.push(id);
+    }
+  }
+  return out;
+}
+
+/** A vision linear that is EITHER 4-bit quantized (q/k/v/o, gate/up, input_proj,
+ *  embedding_projection) OR plain bf16 (down_proj is full-precision). Matches the
+ *  reference, where nn.quantize skips the unquantized `down_proj`. */
+class VisionLinear {
+  private constructor(
+    readonly q: QuantizedLinear | null,
+    readonly w: MlxArray | null,
+  ) {}
+  static load(weights: Weights, config: ModelConfig, path: string): VisionLinear {
+    if (weights.has(`${path}.scales`))
+      return new VisionLinear(QuantizedLinear.load(weights, path, config), null);
+    return new VisionLinear(null, weights.tensor(`${path}.weight`));
+  }
+  /** The dtype of the underlying weight tensor — uint32 for a quantized linear
+   *  (the packed weight), the float dtype for a plain one. The reference
+   *  patch-embedder casts its input to THIS dtype (a uint32 truncation for the
+   *  quantized diffusion input_proj — see #patchify). */
+  get weightDtype(): Dtype {
+    return (this.q ? this.q.w : this.w!).dtype;
+  }
+  forward(x: MlxArray): MlxArray {
+    if (this.q) return this.q.forward(x);
+    const xc = x.astype(this.w!.dtype);
+    const wT = ops.transposeAxes(this.w!, [1, 0]);
+    const out = ops.matmul(xc, wT);
+    xc.dispose();
+    wT.dispose();
+    return out;
+  }
+}
+
+/** Manual fp32 RMS norm matching the reference VisionRMSNorm / NoScale:
+ *  x_f32 * rsqrt(mean(x²)+eps) [* weight_f32], cast back. (NOT the fused kernel
+ *  — the reference q/k/v norms are decomposed f32.) */
+function manualRms(x: MlxArray, weight: MlxArray | null, eps: number): MlxArray {
+  const xf = x.astype(Dtype.float32);
+  const sq = ops.mul(xf, xf);
+  const varr = ops.meanAxis(sq, -1, true);
+  sq.dispose();
+  const epsArr = ops.scalarLike(eps, varr);
+  const denom = ops.add(varr, epsArr);
+  varr.dispose();
+  epsArr.dispose();
+  const rs = ops.rsqrt(denom);
+  denom.dispose();
+  let normed = ops.mul(xf, rs);
+  xf.dispose();
+  rs.dispose();
+  if (weight) {
+    const wf = weight.astype(Dtype.float32);
+    normed = disposing(normed, ops.mul(normed, wf));
+    wf.dispose();
+  }
+  const out = normed.astype(x.dtype);
+  normed.dispose();
+  return out;
+}
+
+/** Pad the last dim of x up to `target` with zeros (for the fused-SDPA head_dim
+ *  padding, reference ensure_fused_sdpa). */
+function padLastDim(x: MlxArray, target: number): MlxArray {
+  const sh = x.shape;
+  const D = sh[sh.length - 1]!;
+  if (D >= target) return x;
+  const padShape = [...sh];
+  padShape[padShape.length - 1] = target - D;
+  const zeros = ops.zeros(padShape, x.dtype);
+  const out = ops.concatAxis([x, zeros], -1);
+  zeros.dispose();
+  return out;
+}
+
+class DiffVisionAttention {
+  readonly qProj: VisionLinear;
+  readonly kProj: VisionLinear;
+  readonly vProj: VisionLinear;
+  readonly oProj: VisionLinear;
+  readonly qNorm: MlxArray;
+  readonly kNorm: MlxArray;
+  readonly nHeads: number;
+  readonly nKvHeads: number;
+  readonly headDim: number;
+  readonly eps: number;
+  constructor(w: Weights, config: ModelConfig, cfg: SiglipVisionConfig, prefix: string) {
+    this.qProj = VisionLinear.load(w, config, `${prefix}.q_proj.linear`);
+    this.kProj = VisionLinear.load(w, config, `${prefix}.k_proj.linear`);
+    this.vProj = VisionLinear.load(w, config, `${prefix}.v_proj.linear`);
+    this.oProj = VisionLinear.load(w, config, `${prefix}.o_proj.linear`);
+    this.qNorm = w.tensor(`${prefix}.q_norm.weight`);
+    this.kNorm = w.tensor(`${prefix}.k_norm.weight`);
+    this.nHeads = cfg.numHeads;
+    this.nKvHeads = cfg.numKvHeads;
+    this.headDim = cfg.headDim;
+    this.eps = cfg.rmsNormEps;
+  }
+  forward(x: MlxArray, rope: Rope2dTables, mask: MlxArray): MlxArray {
+    const [B, L] = x.shape as [number, number, number];
+    const qf = this.qProj.forward(x);
+    let q = ops.reshape(qf, [B, L, this.nHeads, this.headDim]);
+    qf.dispose();
+    const kf = this.kProj.forward(x);
+    let k = ops.reshape(kf, [B, L, this.nKvHeads, this.headDim]);
+    kf.dispose();
+    const vf = this.vProj.forward(x);
+    let v = ops.reshape(vf, [B, L, this.nKvHeads, this.headDim]);
+    vf.dispose();
+
+    q = disposing(q, manualRms(q, this.qNorm, this.eps));
+    k = disposing(k, manualRms(k, this.kNorm, this.eps));
+    v = disposing(v, manualRms(v, null, this.eps));
+
+    q = disposing(q, applyRope2d(q, rope));
+    k = disposing(k, applyRope2d(k, rope));
+
+    q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
+    k = disposing(k, ops.transposeAxes(k, [0, 2, 1, 3]));
+    v = disposing(v, ops.transposeAxes(v, [0, 2, 1, 3]));
+
+    // Reference ensure_fused_sdpa: head_dim 72 isn't a fused-SDPA size, so it
+    // PADS q/k/v to 80 (zeros, no effect on scores — the extra dims are 0),
+    // runs the fused kernel, and slices back to 72. Replicate exactly: mlx's
+    // SDPA path/numerics differ for an unpadded 72 vs the padded-80 fused path.
+    const D = this.headDim;
+    const target = D <= 64 ? 64 : D <= 80 ? 80 : D <= 128 ? 128 : D;
+    if (target !== D) {
+      q = disposing(q, padLastDim(q, target));
+      k = disposing(k, padLastDim(k, target));
+      v = disposing(v, padLastDim(v, target));
+    }
+    let attn = ops.sdpa(q, k, v, 1.0, "array", mask);
+    q.dispose();
+    k.dispose();
+    v.dispose();
+    if (target !== D) {
+      const sh = attn.shape;
+      attn = disposing(attn, attn.slice(sh.map(() => 0), sh.map((s, i) => (i === sh.length - 1 ? D : s))));
+    }
+    let o = ops.transposeAxes(attn, [0, 2, 1, 3]);
+    attn.dispose();
+    o = disposing(o, ops.reshape(o, [B, L, this.nHeads * this.headDim]));
+    const out = this.oProj.forward(o);
+    o.dispose();
+    return out;
+  }
+}
+
+/** GeGLU stays uncompiled here: the reference vision MLP is plain
+ *  `nn.gelu_approx(gate) * up`, not the language model's @mx.compile'd geglu. */
+class DiffVisionMLP {
+  readonly gate: VisionLinear;
+  readonly up: VisionLinear;
+  readonly down: VisionLinear;
+  constructor(w: Weights, config: ModelConfig, prefix: string) {
+    this.gate = VisionLinear.load(w, config, `${prefix}.gate_proj.linear`);
+    this.up = VisionLinear.load(w, config, `${prefix}.up_proj.linear`);
+    this.down = VisionLinear.load(w, config, `${prefix}.down_proj.linear`);
+  }
+  forward(x: MlxArray): MlxArray {
+    const g = this.gate.forward(x);
+    const act = ops.geluApprox(g);
+    g.dispose();
+    const u = this.up.forward(x);
+    const m = ops.mul(act, u);
+    act.dispose();
+    u.dispose();
+    const out = this.down.forward(m);
+    m.dispose();
+    return out;
+  }
+}
+
+class DiffVisionBlock {
+  readonly attn: DiffVisionAttention;
+  readonly mlp: DiffVisionMLP;
+  readonly inputNorm: RMSNorm;
+  readonly postAttnNorm: RMSNorm;
+  readonly preFfNorm: RMSNorm;
+  readonly postFfNorm: RMSNorm;
+  constructor(w: Weights, config: ModelConfig, cfg: SiglipVisionConfig, prefix: string) {
+    this.attn = new DiffVisionAttention(w, config, cfg, `${prefix}.self_attn`);
+    this.mlp = new DiffVisionMLP(w, config, `${prefix}.mlp`);
+    const norm = (n: string) => new RMSNorm(w.tensor(`${prefix}.${n}.weight`), cfg.rmsNormEps);
+    this.inputNorm = norm("input_layernorm");
+    this.postAttnNorm = norm("post_attention_layernorm");
+    this.preFfNorm = norm("pre_feedforward_layernorm");
+    this.postFfNorm = norm("post_feedforward_layernorm");
+  }
+  forward(x: MlxArray, rope: Rope2dTables, mask: MlxArray): MlxArray {
+    const normed = this.inputNorm.forward(x);
+    let attnOut = this.attn.forward(normed, rope, mask);
+    normed.dispose();
+    attnOut = disposing(attnOut, this.postAttnNorm.forward(attnOut));
+    const h = ops.add(x, attnOut);
+    attnOut.dispose();
+    const normedH = this.preFfNorm.forward(h);
+    let ffw = this.mlp.forward(normedH);
+    normedH.dispose();
+    ffw = disposing(ffw, this.postFfNorm.forward(ffw));
+    const out = ops.add(h, ffw);
+    h.dispose();
+    ffw.dispose();
+    return out;
+  }
+}
+
+export class DiffusionVisionTower {
+  readonly cfg: SiglipVisionConfig;
+  readonly inputProj: VisionLinear;
+  readonly posTable: MlxArray; // [2, posSize, hidden]
+  readonly layers: DiffVisionBlock[];
+  readonly stdBias: MlxArray | null;
+  readonly stdScale: MlxArray | null;
+  readonly embedProj: VisionLinear; // MultimodalEmbedder projection
+  readonly embedEps: number;
+  readonly maxPatches: number;
+  readonly patchSize: number;
+  readonly defaultOutputLength: number;
+  readonly rootHidden: number;
+
+  constructor(weights: Weights, config: ModelConfig) {
+    const raw = (config as unknown as { raw: Record<string, any> }).raw.vision_config;
+    const cfg = parseSiglipConfig(raw);
+    this.cfg = cfg;
+    this.patchSize = cfg.patchSize;
+    this.defaultOutputLength = cfg.defaultOutputLength;
+    this.maxPatches = cfg.defaultOutputLength * cfg.poolingKernelSize * cfg.poolingKernelSize;
+    this.rootHidden = Math.sqrt(cfg.hiddenSize);
+    this.embedEps = (raw.rms_norm_eps as number) ?? 1e-6;
+
+    this.inputProj = VisionLinear.load(weights, config, `${VP}patch_embedder.input_proj`);
+    this.posTable = weights.tensor(`${VP}patch_embedder.position_embedding_table`);
+    this.layers = Array.from(
+      { length: cfg.numLayers },
+      (_, i) => new DiffVisionBlock(weights, config, cfg, `${VP}encoder.layers.${i}`),
+    );
+    this.stdBias = cfg.standardize ? weights.tensor(`${VP}std_bias`) : null;
+    this.stdScale = cfg.standardize ? weights.tensor(`${VP}std_scale`) : null;
+    this.embedProj = VisionLinear.load(weights, config, "model.encoder.embed_vision.embedding_projection");
+  }
+
+  /** Decode + aspect-preserving resize (multiples of pooling·patch=48, within the
+   *  max_patches budget) + rescale 1/255 + channel-first. Mirrors optiq's
+   *  gemma4 `preprocess_images` (BICUBIC, do_rescale, do_normalize=False).
+   *  Returns pixel_values [1,3,H,W] f32 and the per-image soft-token count. */
+  async preprocess(bytes: Uint8Array): Promise<{ pixels: MlxArray; softTokens: number }> {
+    let img = await decodeImage(bytes);
+    const p = this.patchSize;
+    const k = this.cfg.poolingKernelSize;
+    const sideMult = p * k;
+    const targetPx = this.maxPatches * p * p;
+    const factor = Math.sqrt(targetPx / (img.height * img.width));
+    let th = Math.floor((factor * img.height) / sideMult) * sideMult;
+    let tw = Math.floor((factor * img.width) / sideMult) * sideMult;
+    const maxSide = Math.floor(this.maxPatches / (k * k)) * sideMult;
+    if (th === 0 && tw === 0) {
+      th = sideMult;
+      tw = sideMult;
+    } else if (th === 0) {
+      th = sideMult;
+      tw = Math.min(Math.floor(img.width / img.height) * sideMult, maxSide);
+    } else if (tw === 0) {
+      tw = sideMult;
+      th = Math.min(Math.floor(img.height / img.width) * sideMult, maxSide);
+    }
+    if (th !== img.height || tw !== img.width) img = resizeBicubic(img, tw, th);
+
+    const { width: W, height: H, data } = img;
+    // HWC uint8 -> CHW f32 / 255
+    const chw = new Float32Array(3 * H * W);
+    for (let y = 0; y < H; y++)
+      for (let x = 0; x < W; x++)
+        for (let c = 0; c < 3; c++)
+          chw[c * H * W + y * W + x] = data[(y * W + x) * 3 + c]! / 255;
+    const pixels = MlxArray.fromFloat32(chw, [1, 3, H, W]);
+    const softTokens = Math.floor((Math.floor(H / p) * Math.floor(W / p)) / (k * k));
+    return { pixels, softTokens };
+  }
+
+  /** Host-side patch positions (x,y) for the real patches + the padding mask,
+   *  matching VisionModel._patch_positions_single (meshgrid xy, row-major). */
+  #positions(H: number, W: number): { real: Int32Array; numReal: number } {
+    const p = this.patchSize;
+    const pH = Math.floor(H / p);
+    const pW = Math.floor(W / p);
+    const numReal = pH * pW;
+    const real = new Int32Array(numReal * 2);
+    for (let y = 0; y < pH; y++)
+      for (let x = 0; x < pW; x++) {
+        const L = y * pW + x;
+        real[L * 2] = x;
+        real[L * 2 + 1] = y;
+      }
+    return { real, numReal };
+  }
+
+  /** patchify: [1,3,H,W] -> [1, numReal, 3*p*p] (order p_row,p_col,C), 2*(x-0.5). */
+  #patchify(pixels: MlxArray): MlxArray {
+    const [B, C, H, W] = pixels.shape as [number, number, number, number];
+    const p = this.patchSize;
+    const pH = Math.floor(H / p);
+    const pW = Math.floor(W / p);
+    let t = ops.reshape(pixels, [B, C, pH, p, pW, p]);
+    t = disposing(t, ops.transposeAxes(t, [0, 2, 4, 3, 5, 1])); // [B,pH,pW,p,p,C]
+    t = disposing(t, ops.reshape(t, [B, pH * pW, C * p * p]));
+    // 2*(x - 0.5)
+    const half = ops.scalarLike(0.5, t);
+    t = disposing(t, ops.sub(t, half));
+    half.dispose();
+    t = disposing(t, ops.mulScalar(t, 2));
+    // Reference quirk: `patches.astype(self.input_proj.weight.dtype)`. The
+    // DiffusionGemma input_proj is QUANTIZED, so weight.dtype is uint32 and this
+    // TRUNCATES the patches to integers (mostly 0) before the matmul. The model
+    // was trained with this, so replicate it exactly (e4b's bf16 input_proj made
+    // it a no-op there). Back to a float the quantized matmul consumes.
+    t = disposing(t, t.astype(this.inputProj.weightDtype));
+    t = disposing(t, t.astype(Dtype.bfloat16));
+    return t;
+  }
+
+  /** Full tower forward over channel-first pixel_values [1,3,H,W] -> standardized
+   *  pooled hidden [1, numRealSoft, hidden] (BEFORE the MultimodalEmbedder). */
+  visionTower(pixels: MlxArray, dbg?: (n: string, a: MlxArray) => void): MlxArray {
+    const [, , H, W] = pixels.shape as [number, number, number, number];
+    const { real, numReal } = this.#positions(H, W);
+    const k = this.cfg.poolingKernelSize;
+    const p = this.patchSize;
+    const pW = Math.floor(W / p);
+    const pH = Math.floor(H / p);
+
+    // --- patch embed (real patches only) ---
+    const patches = this.#patchify(pixels);
+    let embeds = this.inputProj.forward(patches); // [1, numReal, hidden]
+    patches.dispose();
+
+    // + 2D position embedding: table[0, x] + table[1, y]
+    const hidden = this.cfg.hiddenSize;
+    const posSize = this.cfg.positionEmbeddingSize;
+    const pxIdx = new Int32Array(numReal);
+    const pyIdx = new Int32Array(numReal);
+    for (let i = 0; i < numReal; i++) {
+      pxIdx[i] = real[i * 2]!;
+      pyIdx[i] = real[i * 2 + 1]!;
+    }
+    const plane0 = this.posTable.slice([0, 0, 0], [1, posSize, hidden]);
+    const t0r = ops.reshape(plane0, [posSize, hidden]);
+    plane0.dispose();
+    const plane1 = this.posTable.slice([1, 0, 0], [2, posSize, hidden]);
+    const t1r = ops.reshape(plane1, [posSize, hidden]);
+    plane1.dispose();
+    const pxArr = MlxArray.fromInt32(pxIdx, [numReal]);
+    const pyArr = MlxArray.fromInt32(pyIdx, [numReal]);
+    const e0 = ops.takeAxis(t0r, pxArr, 0);
+    const e1 = ops.takeAxis(t1r, pyArr, 0);
+    t0r.dispose();
+    t1r.dispose();
+    pxArr.dispose();
+    pyArr.dispose();
+    let pos = ops.add(e0, e1);
+    e0.dispose();
+    e1.dispose();
+    pos = disposing(pos, ops.reshape(pos, [1, numReal, hidden]));
+    pos = disposing(pos, pos.astype(embeds.dtype));
+    embeds = disposing(embeds, ops.add(embeds, pos)); // [1, numReal, hidden]
+    pos.dispose();
+    dbg?.("patchembed", embeds);
+
+    // --- pad to maxPatches (zeros) ---
+    const numPad = this.maxPatches - numReal;
+    if (numPad > 0) {
+      const padE = ops.zeros([1, numPad, hidden], embeds.dtype);
+      embeds = disposing(embeds, ops.concatAxis([embeds, padE], 1));
+      padE.dispose();
+    }
+
+    // --- positions array [1, maxPatches, 2] (real then -1 pad) for RoPE ---
+    const posInts = new Int32Array(this.maxPatches * 2).fill(-1);
+    posInts.set(real, 0);
+    const positions = MlxArray.fromInt32(posInts, [1, this.maxPatches, 2]);
+    const rope = rope2dTables(positions, this.cfg.headDim, this.cfg.ropeTheta, embeds.dtype);
+    positions.dispose();
+
+    // --- bidirectional -1e4 mask [1,1,maxPatches,maxPatches] ---
+    // valid_i AND valid_j → 0 else -1e4. valid = patch index < numReal.
+    const validInts = new Float32Array(this.maxPatches);
+    for (let i = 0; i < this.maxPatches; i++) validInts[i] = i < numReal ? 1 : 0;
+    const validRow = MlxArray.fromFloat32(validInts, [1, 1, 1, this.maxPatches]); // keys
+    const validCol = MlxArray.fromFloat32(validInts, [1, 1, this.maxPatches, 1]); // queries
+    const both = ops.mul(validRow, validCol); // [1,1,mp,mp] 1 if both valid
+    validRow.dispose();
+    validCol.dispose();
+    // mask_fill: where(both, 0, -1e4) = (both-1)*1e4
+    const oneArr = ops.scalarLike(1, both);
+    const bm1 = ops.sub(both, oneArr); // both-1: 0 or -1
+    both.dispose();
+    oneArr.dispose();
+    let mask = ops.mulScalar(bm1, 1e4); // 0 or -1e4
+    bm1.dispose();
+    mask = disposing(mask, mask.astype(embeds.dtype));
+
+    // --- transformer ---
+    let h = embeds;
+    for (let li = 0; li < this.layers.length; li++) {
+      const next = this.layers[li]!.forward(h, rope, mask);
+      h.dispose();
+      h = next;
+      if (li < 3) dbg?.(`vlayer${li}`, h);
+    }
+    rope.dispose();
+    mask.dispose();
+    dbg?.("transformer", h);
+
+    // --- pooler: zero padding, position-grouped 3x3 avg-pool to 280, ×√hidden ---
+    // zero padding rows
+    if (numPad > 0) {
+      const keep = new Float32Array(this.maxPatches);
+      for (let i = 0; i < this.maxPatches; i++) keep[i] = i < numReal ? 1 : 0;
+      const keepArr = MlxArray.fromFloat32(keep, [1, this.maxPatches, 1]).astype(h.dtype);
+      h = disposing(h, ops.mul(h, keepArr));
+      keepArr.dispose();
+    }
+    // pooling matrix [maxPatches, length] : kernel_idx = (x//k) + (pW_blocks)*(y//k)
+    // pW_blocks = (max_x+1)//k where max_x = pW-1 → pW//k... reference: (max_x//k)
+    // where max_x = max(clamped x)+1 = pW. So blocksPerRow = pW // k.
+    const length = this.defaultOutputLength;
+    const blocksPerRow = Math.floor(pW / k);
+    const poolW = new Float32Array(this.maxPatches * length);
+    const invK2 = 1 / (k * k);
+    for (let i = 0; i < this.maxPatches; i++) {
+      // clamped positions: pad (-1) → 0
+      const x = i < numReal ? real[i * 2]! : 0;
+      const y = i < numReal ? real[i * 2 + 1]! : 0;
+      const idx = Math.floor(x / k) + blocksPerRow * Math.floor(y / k);
+      if (idx < length) poolW[i * length + idx] = invK2;
+    }
+    const hf = h.astype(Dtype.float32);
+    h.dispose();
+    const poolWArr = MlxArray.fromFloat32(poolW, [this.maxPatches, length]);
+    // output[l,d] = sum_L poolW[L,l] * hf[L,d] = poolW.T @ hf
+    const hf2 = ops.reshape(hf, [this.maxPatches, hidden]);
+    hf.dispose();
+    const poolWT = ops.transposeAxes(poolWArr, [1, 0]); // [length, maxPatches]
+    poolWArr.dispose();
+    let pooled = ops.matmul(poolWT, hf2); // [length, hidden] f32
+    poolWT.dispose();
+    hf2.dispose();
+    // ×√hidden, back to bf16, reshape [1, length, hidden]
+    const dtype = this.posTable.dtype;
+    pooled = disposing(pooled, pooled.astype(dtype));
+    pooled = disposing(pooled, ops.mulScalar(pooled, this.rootHidden));
+    pooled = disposing(pooled, ops.reshape(pooled, [1, length, hidden]));
+
+    // --- trim to the real soft tokens (kernel_idx < numRealSoft) ---
+    const softW = blocksPerRow;
+    const softH = Math.floor(pH / k);
+    const numSoft = softW * softH;
+    const trimmed = pooled.slice([0, 0, 0], [1, numSoft, hidden]);
+    pooled.dispose();
+    let out = trimmed;
+    dbg?.("pooled", out);
+
+    // --- standardize: (h - std_bias) * std_scale ---
+    if (this.stdBias && this.stdScale) {
+      const sb = this.stdBias.astype(out.dtype);
+      const ss = this.stdScale.astype(out.dtype);
+      const sub = ops.sub(out, sb);
+      out.dispose();
+      sb.dispose();
+      out = ops.mul(sub, ss);
+      sub.dispose();
+      ss.dispose();
+    }
+    dbg?.("toweroutput", out);
+    return out;
+  }
+
+  /** get_image_features = embed_vision(vision_tower(pixels)):
+   *  MultimodalEmbedder (RMSNormNoScale → projection 1152→2816). NO /embed_scale. */
+  getImageFeatures(pixels: MlxArray, dbg?: (n: string, a: MlxArray) => void): MlxArray {
+    const tower = this.visionTower(pixels, dbg);
+    const normed = manualRms(tower, null, this.embedEps); // RMSNormNoScale
+    tower.dispose();
+    const out = this.embedProj.forward(normed);
+    normed.dispose();
+    return out;
+  }
+}

package/src/models/drafter-loaders.ts

@@ -0,0 +1,49 @@
+// Drafter models as the layers above build them. Each loader imports its
+// implementation on first use and hands the drafter over as the port the draft
+// sources consume (`contracts/mlx/drafter`); no source names a class.
+
+import type { ModelConfig } from "../artifacts/config";
+import type { Weights } from "../artifacts/weights";
+import type { AssistantDrafterModel, DeepspecDrafterModel, Dflash2DrafterModel, DsparkDrafterModel, NativeMtpHead, RecurrentMtpModule }
+  from "../contracts/mlx/drafter";
+import type { MlxDeclaredGraph } from "../contracts/mlx/graph";
+
+/** The KV-borrowing assistant drafter in `dir`. */
+export async function loadAssistantDrafter(dir: string): Promise<AssistantDrafterModel> {
+  const { GemmaAssistantDrafter } = await import("./gemma4/assistant");
+  return GemmaAssistantDrafter.load(dir);
+}
+
+/** A DeepSpec drafter in `dir`. */
+export async function loadDeepspecDrafter(dir: string): Promise<DeepspecDrafterModel> {
+  const { DeepspecDrafter } = await import("./speculative/deepspec");
+  return DeepspecDrafter.load(dir);
+}
+
+/** A DFlash 2 drafter in `dir`, in its trained basis (projections quantized to `bits`, 0 = BF16). */
+export async function loadDflash2Drafter(dir: string, opts: { bits: number }): Promise<Dflash2DrafterModel> {
+  const { Dflash2Drafter } = await import("./speculative/dflash2");
+  return Dflash2Drafter.load(dir, opts);
+}
+
+/** Our trained DSpark drafter in `dir` (`dspark.json` beside the weights), dispatching on its variant. */
+export async function loadDsparkDrafter(dir: string): Promise<DsparkDrafterModel> {
+  const { loadDsparkDrafter: load } = await import("./speculative/loader");
+  return load(dir);
+}
+
+/** The recurrent MTP block of a Qwen companion checkpoint. Its tensors belong to `weights`; the
+ *  module's own views belong to `resources`. */
+export async function loadQwenMtpModule(weights: Weights, config: ModelConfig,
+  resources: DisposableStack): Promise<RecurrentMtpModule> {
+  const { MtpModule } = await import("./qwen/mtp");
+  return new MtpModule(weights, config, resources);
+}
+
+/** The native MTP head a graph declares (`nativeDraft`). Only a graph that provides its head binds. */
+export function nativeMtpHeadOf(graph: object): NativeMtpHead {
+  const head = (graph as Partial<MlxDeclaredGraph>).nativeDraftHead;
+  if (typeof head !== "function")
+    throw new Error("native GLM-5.2 MTP binds only to the GLM-5.2 graph that declares it");
+  return head.call(graph);
+}

package/src/models/drafters.ts

@@ -0,0 +1,39 @@
+// Conventions by which speculative-decoding artifacts and graphs declare
+// themselves. Pure string facts, so the model index, the hub and the draft
+// provider registry all read one definition without loading MLX (the drafters
+// themselves load through `drafter-loaders`).
+
+import { isDeepspecArchitecture } from "./speculative/deepspec-artifact";
+import { isDflash2Architecture } from "./speculative/dflash2-artifact";
+
+/** A DeepSpec drafter (DeepSeek's released DSpark checkpoints): a plain HF config stamped with its
+ *  architecture. */
+export function isDeepspecDrafterConfig(config: Record<string, unknown> | null): boolean {
+  return isDeepspecArchitecture(config?.architectures);
+}
+
+/** A DFlash 2 block drafter (incoai/Qwen3.8-27B-DFlash2): a plain HF config stamped with its architecture. */
+export function isDflash2DrafterConfig(config: Record<string, unknown> | null): boolean {
+  return isDflash2Architecture(config?.architectures);
+}
+
+/** A drafter whose weights are Q-only and have no standalone LM head
+ *  (`gemma4_assistant`, `gemma4_unified_assistant`). */
+export function isAssistantModelType(modelType: string): boolean {
+  return modelType.includes("assistant");
+}
+
+/** A multi-token-prediction head split from a qwen3_5-family release
+ *  (`qwen3_5_mtp`). */
+export function isMtpModelType(modelType: string): boolean {
+  return modelType.endsWith("_mtp");
+}
+
+/** Drafters with no standalone graph: never servable or selectable as a model. */
+export function isUnservableDrafterModelType(modelType: string): boolean {
+  return isAssistantModelType(modelType);
+}
+
+/** The `GraphCapabilities.nativeDraft` kind of a graph whose checkpoint carries
+ *  its own multi-token-prediction row. */
+export const NATIVE_MTP_DRAFT = "checkpoint-mtp";