mlx-bun 0.0.13 → 0.3.0

package/src/model/qwen3_5.ts

@@ -12,9 +12,11 @@
 import type { ModelConfig } from "../config";
 import type { Weights } from "../weights";
 import { MlxArray } from "../mlx/array";
-import { Dtype } from "../mlx/ffi";
+import { Dtype, deviceArchitecture } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
 import { CompiledFunction } from "../mlx/compile";
+import { qwenAppendChunkSize } from "./qwen-append";
+import { TrellisLinear, fusedGateUpEligible, fusedGateUpSwiglu, TRELLIS_MATVEC_MAX_M } from "./trellis-linear";
 import {
   argmaxLastPosition,
   disposeTriple,
@@ -29,7 +31,14 @@ import {
   type Cache,
   type Mask,
 } from "./gemma4-base";
+import { qwen35WeightsView } from "./qwen3_5-checkpoint";
+import { materializeCopy } from "../mlx/materialize";
 import { gatedDeltaUpdate, SSMCache } from "./qwen3-delta";
+import type { QwenConvolution } from "./qwen-conv";
+import {
+  activeMrope, applyInterleavedRope, buildMropePositions, mropeInvFreq,
+  setActiveMrope, type MropeRequestState,
+} from "./qwen3-mrope";
 
 const PREFIX = "language_model";
 
@@ -103,6 +112,10 @@ export function compiledSilu(x: MlxArray): MlxArray {
 
 /** Gated-DeltaNet linear-attention layer (mlx-lm GatedDeltaNet). */
 export class GatedDeltaNet {
+  /** Per-model experiment seam; null retains the oracle graph. The same
+   * implementation advances speculative rollback prefixes. Borrowed inputs,
+   * owned activation and independent state tail; no global runtime mutation. */
+  convolution: QwenConvolution | null = null;
   readonly inProjQkv: QuantizedLinear;
   readonly inProjZ: QuantizedLinear;
   readonly inProjB: QuantizedLinear;
@@ -142,34 +155,66 @@ export class GatedDeltaNet {
     this.normWeight = weights.tensor(`${prefix}.norm.weight`);
   }
 
-  forward(x: MlxArray, cache: SSMCache): MlxArray {
+  #convolve(qkv: MlxArray, state: MlxArray): [MlxArray, MlxArray] {
+    if (this.convolution) return this.convolution(qkv, state, this.convWeight);
+    const [B, S, D] = qkv.shape as [number, number, number];
+    const nKeep = this.convKernel - 1;
+    const input = ops.concatAxis([state, qkv], 1);
+    // MLX copy and contiguous can both alias the whole prefill buffer.
+    // Materialize only the tail so cache residency follows its logical size.
+    const view = input.slice([0, S, 0], [B, S + nKeep, D]);
+    const tail = materializeCopy(view);
+    view.dispose();
+    const conv = ops.conv1d(input, this.convWeight, 1, 0, 1, D);
+    input.dispose();
+    const out = compiledSilu(conv);
+    conv.dispose();
+    return [out, tail];
+  }
+
+  forward(x: MlxArray, cache: SSMCache, independentRows = false): MlxArray {
     const [B, S] = x.shape as [number, number, number];
     const convDim = this.keyDim * 2 + this.valueDim;
     const nKeep = this.convKernel - 1;
+    const prof = (globalThis as Record<string, unknown>).__deltaProf as
+      Record<string, number>
+      | undefined;
+
+    // Armed speculative verify round (serve loop, serial lane): this forward
+    // must be rewindable. Snapshot = hand the REPLACED state slots to the
+    // round instead of disposing (free — arrays are immutable); record the
+    // position-local kernel inputs (qkv/a/b) and install the prefix replay.
+    const spec = cache.specRound;
+    if (spec) {
+      if (!spec.armed)
+        throw new Error("GatedDeltaNet: spec round already recorded this round");
+      spec.armed = false;
+      spec.S = S;
+      spec.replay = (c, keep) => this.#replaySpecPrefix(c, keep);
+    }
 
-    const qkv = this.inProjQkv.forward(x); // [B,S,convDim]
-    let z = this.inProjZ.forward(x);
+    const t0 = prof ? performance.now() : 0;
+    const qkv = this.inProjQkv.forward(x, independentRows); // [B,S,convDim]
+    let z = this.inProjZ.forward(x, independentRows);
     z = disposing(z, ops.reshape(z, [B, S, this.numVHeads, this.headVDim]));
-    const b = this.inProjB.forward(x); // [B,S,numVHeads]
-    const a = this.inProjA.forward(x); // [B,S,numVHeads]
+    const b = this.inProjB.forward(x, independentRows); // [B,S,numVHeads]
+    const a = this.inProjA.forward(x, independentRows);
+    if (prof) { ops.evalAll([qkv, z, b, a]); prof.proj = (prof.proj ?? 0) + performance.now() - t0; }
+    const tc = prof ? performance.now() : 0; // [B,S,numVHeads]
 
     // Causal depthwise conv with the conv-state prefix (B=1: no ssm mask).
     const convState =
       cache.conv ?? ops.zeros([B, nKeep, convDim], x.dtype);
-    const convInput = ops.concatAxis([convState, qkv], 1); // [B,S+nKeep,convDim]
+    const [convOut, newConv] = this.#convolve(qkv, convState);
     if (!cache.conv) convState.dispose();
-    qkv.dispose();
-    // New conv state = last nKeep rows (contiguous, the array is sliced).
-    const newConv = ops.contiguous(
-      convInput.slice([0, S, 0], [B, S + nKeep, convDim]),
-    );
-    cache.conv?.dispose();
+    if (spec) spec.qkv = qkv;
+    else qkv.dispose();
+    if (spec) spec.prevConv = cache.conv;
+    else cache.conv?.dispose();
     cache.conv = newConv;
 
-    const conv = ops.conv1d(convInput, this.convWeight, 1, 0, 1, convDim);
-    convInput.dispose();
-    const convOut = compiledSilu(conv); // nn.silu (mx.compile), matches mlx-lm
-    conv.dispose();
+    if (prof) { ops.evalAll([convOut]); prof.conv = (prof.conv ?? 0) + performance.now() - tc; }
+    const tn = prof ? performance.now() : 0;
 
     const [qFlat, kFlat, vFlat] = ops.split(
       convOut, [this.keyDim, 2 * this.keyDim], -1,
@@ -187,16 +232,26 @@ export class GatedDeltaNet {
     q = disposing(q, ops.mulScalar(q, invScale * invScale));
     k = disposing(k, ops.rmsNorm(k, null, 1e-6));
     k = disposing(k, ops.mulScalar(k, invScale));
+    if (prof) { ops.evalAll([q, k]); prof.norms = (prof.norms ?? 0) + performance.now() - tn; }
+    const tk = prof ? performance.now() : 0;
 
     const [out, newState] = gatedDeltaUpdate(
       q, k, v, a, b, this.aLog, this.dtBias, cache.recurrent,
     );
+    if (prof) { ops.evalAll([out]); prof.kernel = (prof.kernel ?? 0) + performance.now() - tk; }
+    const to = prof ? performance.now() : 0;
     q.dispose();
     k.dispose();
     v.dispose();
-    a.dispose();
-    b.dispose();
-    cache.recurrent?.dispose();
+    if (spec) {
+      spec.a = a;
+      spec.b = b;
+      spec.prevRecurrent = cache.recurrent;
+    } else {
+      a.dispose();
+      b.dispose();
+      cache.recurrent?.dispose();
+    }
     cache.recurrent = newState;
     cache.advance(S);
 
@@ -206,11 +261,73 @@ export class GatedDeltaNet {
     z.dispose();
     const merged = ops.reshape(gated, [B, S, this.valueDim]);
     gated.dispose();
-    const result = this.outProj.forward(merged);
+    const result = this.outProj.forward(merged, independentRows);
     merged.dispose();
+    if (prof) { ops.evalAll([result]); prof.out = (prof.out ?? 0) + performance.now() - to; }
     return result;
   }
 
+  /** Speculative rollback replay: with the pre-round snapshot RESTORED onto
+   *  `cache`, re-advance the first `keep` window tokens from the recorded
+   *  position-local inputs. Every op mirrors forward() on the same values —
+   *  conv windows, silu, per-position norms, and the kernel's serial prefix
+   *  are all independent of the rejected tail — so the resulting conv +
+   *  recurrent state is BIT-EXACTLY what a forward over only the accepted
+   *  prefix would have produced. Output projections (z / out_proj) are state-
+   *  free and skipped: only the states matter here. */
+  #replaySpecPrefix(cache: SSMCache, keep: number): void {
+    const r = cache.specRound;
+    if (!r || !r.qkv || !r.a || !r.b)      throw new Error("GatedDeltaNet replay without recorded round inputs");
+    const B = r.qkv.shape[0]!;
+    const convDim = this.keyDim * 2 + this.valueDim;
+    const nKeep = this.convKernel - 1;
+
+    const qkvPfxView = r.qkv.slice([0, 0, 0], [B, keep, convDim]);
+    const qkvPfx = ops.contiguous(qkvPfxView);
+    qkvPfxView.dispose();
+    const convState = cache.conv ?? ops.zeros([B, nKeep, convDim], qkvPfx.dtype);
+    const [convOut, newConv] = this.#convolve(qkvPfx, convState);
+    if (!cache.conv) convState.dispose();
+    qkvPfx.dispose();
+    cache.conv?.dispose();
+    cache.conv = newConv;
+
+    const [qFlat, kFlat, vFlat] = ops.split(
+      convOut, [this.keyDim, 2 * this.keyDim], -1,
+    ) as [MlxArray, MlxArray, MlxArray];
+    convOut.dispose();
+    let q = ops.reshape(qFlat, [B, keep, this.numKHeads, this.headKDim]);
+    qFlat.dispose();
+    let k = ops.reshape(kFlat, [B, keep, this.numKHeads, this.headKDim]);
+    kFlat.dispose();
+    const v = disposing(vFlat, ops.reshape(vFlat, [B, keep, this.numVHeads, this.headVDim]));
+    const invScale = Math.pow(this.headKDim, -0.5);
+    q = disposing(q, ops.rmsNorm(q, null, 1e-6));
+    q = disposing(q, ops.mulScalar(q, invScale * invScale));
+    k = disposing(k, ops.rmsNorm(k, null, 1e-6));
+    k = disposing(k, ops.mulScalar(k, invScale));
+
+    const aPfxView = r.a.slice([0, 0, 0], [B, keep, this.numVHeads]);
+    const aPfx = ops.contiguous(aPfxView);
+    aPfxView.dispose();
+    const bPfxView = r.b.slice([0, 0, 0], [B, keep, this.numVHeads]);
+    const bPfx = ops.contiguous(bPfxView);
+    bPfxView.dispose();
+
+    const [y, newState] = gatedDeltaUpdate(
+      q, k, v, aPfx, bPfx, this.aLog, this.dtBias, cache.recurrent,
+    );
+    q.dispose();
+    k.dispose();
+    v.dispose();
+    aPfx.dispose();
+    bPfx.dispose();
+    y.dispose(); // only the state advance matters on the rollback path
+    cache.recurrent?.dispose();
+    cache.recurrent = newState;
+    cache.advance(keep);
+  }
+
   private rmsNormGated(hidden: MlxArray, gate: MlxArray): MlxArray {
     const xn = ops.rmsNorm(hidden, this.normWeight, this.eps); // bf16
     // Oracle: _precise_swiglu(hidden, gate, xn) — ALWAYS one @mx.compile kernel
@@ -252,18 +369,18 @@ export class Qwen3Attention {
     this.kNorm = new RMSNorm(weights.tensor(`${prefix}.k_norm.weight`), t.rmsNormEps);
   }
 
-  forward(x: MlxArray, mask: Mask, cache: Cache): MlxArray {
+  forward(x: MlxArray, mask: Mask, cache: Cache, independentRows = false): MlxArray {
     const [B, L] = x.shape as [number, number, number];
 
     // q_proj emits 2× head_dim per head → split into queries + gate.
-    const qp = this.qProj.forward(x);
+    const qp = this.qProj.forward(x, independentRows);
     const qpr = disposing(qp, ops.reshape(qp, [B, L, this.nHeads, this.headDim * 2]));
     const [qHeads, gateHeads] = ops.split(qpr, [this.headDim], -1) as [MlxArray, MlxArray];
     qpr.dispose();
     const gate = disposing(gateHeads, ops.reshape(gateHeads, [B, L, this.nHeads * this.headDim]));
 
-    let k = this.kProj.forward(x);
-    let v = this.vProj.forward(x);
+    let k = this.kProj.forward(x, independentRows);
+    let v = this.vProj.forward(x, independentRows);
 
     // q/k norm over head_dim BEFORE transpose (reference order).
     let q = this.qNorm.forward(qHeads);
@@ -278,14 +395,24 @@ export class Qwen3Attention {
     // Batched decode: the scheduler's mask wrapper exposes each row's REAL
     // position as ropeOffsetArr (rows have different prompt lengths); the
     // dynamic-offset kernel is the same fast::rope, bit-exact vs the static
-    // form (tests/compile.test.ts). Serial lane: scalar offset, unchanged.
+    // form (tests/unit/compile.test.ts). Serial lane: scalar offset, unchanged.
+    // Vision requests (serial lane only) install activeMrope for the current
+    // forward: 3D interleaved positions via the manual apply — the reference
+    // never uses the fused fast-rope kernel when position_ids are supplied,
+    // so this IS the oracle's own arithmetic. Text-only stays on ops.rope.
+    const mr = activeMrope;
     const offArr = (cache as { ropeOffsetArr?: MlxArray }).ropeOffsetArr;
-    q = disposing(q, offArr
-      ? ops.ropeDynamic(q, this.ropeDims, this.ropeBase, offArr, null)
-      : ops.rope(q, this.ropeDims, this.ropeBase, cache.offset, null));
-    k = disposing(k, offArr
-      ? ops.ropeDynamic(k, this.ropeDims, this.ropeBase, offArr, null)
-      : ops.rope(k, this.ropeDims, this.ropeBase, cache.offset, null));
+    if (mr) {
+      q = disposing(q, applyInterleavedRope(q, mr));
+      k = disposing(k, applyInterleavedRope(k, mr));
+    } else {
+      q = disposing(q, offArr
+        ? ops.ropeDynamic(q, this.ropeDims, this.ropeBase, offArr, null)
+        : ops.rope(q, this.ropeDims, this.ropeBase, cache.offset, null));
+      k = disposing(k, offArr
+        ? ops.ropeDynamic(k, this.ropeDims, this.ropeBase, offArr, null)
+        : ops.rope(k, this.ropeDims, this.ropeBase, cache.offset, null));
+    }
 
     let attn: MlxArray;
     if (cache instanceof QuantizedKVCache) {
@@ -316,32 +443,56 @@ export class Qwen3Attention {
     const gated = ops.mul(merged, sig);
     merged.dispose();
     sig.dispose();
-    const out = this.oProj.forward(gated);
+    const out = this.oProj.forward(gated, independentRows);
     gated.dispose();
     return out;
   }
 }
 
+/** MLP projections are affine-quantized (QuantizedLinear) or, in a Q2b
+ *  packed-trellis artifact, TrellisLinear — chosen per module by the config's
+ *  quantization entry (`mode: "trellis"`). */
+export type MlpLinear = QuantizedLinear | TrellisLinear;
+function loadMlpLinear(weights: Weights, path: string, config: ModelConfig): MlpLinear {
+  if (!TrellisLinear.isTrellis(config, path)) return QuantizedLinear.load(weights, path, config);
+  // M4 Pro measurements support the 27B down projection at M3/4. The
+  // operation also checks variant, dtype and packed-code layout eligibility.
+  const useSharedScatterCodebook = config.text.hiddenSize === 5120 &&
+    config.text.intermediateSize === 17408 && path.endsWith(".down_proj");
+  return TrellisLinear.load(weights, path, config, useSharedScatterCodebook);
+}
+
 export class Qwen3MLP {
-  readonly gate: QuantizedLinear;
-  readonly up: QuantizedLinear;
-  readonly down: QuantizedLinear;
+  readonly gate: MlpLinear;
+  readonly up: MlpLinear;
+  readonly down: MlpLinear;
 
   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);
+    this.gate = loadMlpLinear(weights, `${prefix}.gate_proj`, config);
+    this.up = loadMlpLinear(weights, `${prefix}.up_proj`, config);
+    this.down = loadMlpLinear(weights, `${prefix}.down_proj`, config);
   }
 
-  forward(x: MlxArray): MlxArray {
-    const g = this.gate.forward(x);
-    const u = this.up.forward(x);
+  forward(x: MlxArray, inputRowContiguous = false, independentRows = false): MlxArray {
+    // Packed-trellis decode (M ≤ 4): gate, up and the swiglu in ONE kernel —
+    // x read once, no gate/up vectors materialized. This Lab path retains
+    // the packed activation arithmetic documented in turboquant.md.
+    if (this.gate instanceof TrellisLinear && this.up instanceof TrellisLinear &&
+        fusedGateUpEligible(this.gate, this.up) &&
+        x.shape.slice(0, -1).reduce((a, b) => a * b, 1) <= TRELLIS_MATVEC_MAX_M) {
+      const hidden = fusedGateUpSwiglu(x, this.gate, this.up);
+      const out = this.down instanceof QuantizedLinear ? this.down.forward(hidden, independentRows) : this.down.forward(hidden);
+      hidden.dispose();
+      return out;
+    }
+    const g = this.gate instanceof TrellisLinear ? this.gate.forward(x, inputRowContiguous) : this.gate.forward(x, independentRows);
+    const u = this.up instanceof TrellisLinear ? this.up.forward(x, inputRowContiguous) : this.up.forward(x, independentRows);
     // Oracle: down_proj(swiglu(gate_proj(x), up_proj(x))) — swiglu ALWAYS compiled
     // (mlx-lm's @mx.compile swiglu; no unfused silu+mul, matches its kernel set).
     const hidden = compiledSwiglu(g, u);
     g.dispose();
     u.dispose();
-    const out = this.down.forward(hidden);
+    const out = this.down instanceof QuantizedLinear ? this.down.forward(hidden, independentRows) : this.down.forward(hidden);
     hidden.dispose();
     return out;
   }
@@ -366,16 +517,18 @@ export class Qwen3Layer {
     this.postAttnNorm = new RMSNorm(weights.tensor(`${prefix}.post_attention_layernorm.weight`), config.text.rmsNormEps);
   }
 
-  forward(x: MlxArray, faMask: Mask, cache: Cache): MlxArray {
+  forward(x: MlxArray, faMask: Mask, cache: Cache, independentRows = false): MlxArray {
     const xn = this.inputNorm.forward(x);
     const r = this.isLinear
-      ? this.linearAttn!.forward(xn, cache as SSMCache)
-      : this.selfAttn!.forward(xn, faMask, cache);
+      ? this.linearAttn!.forward(xn, cache as SSMCache, independentRows)
+      : this.selfAttn!.forward(xn, faMask, cache, independentRows);
     xn.dispose();
     const h = ops.add(x, r);
     r.dispose();
     const hn = this.postAttnNorm.forward(h);
-    const m = this.mlp.forward(hn);
+    // RMSNorm allocates an aligned row-contiguous output. Pass that layout
+    // proof so Trellis can match native small-prefill matmul without an eval.
+    const m = this.mlp.forward(hn, true, independentRows);
     hn.dispose();
     const out = ops.add(h, m);
     h.dispose();
@@ -399,6 +552,14 @@ export class Qwen35Model {
   readonly faIdx: number;
 
   constructor(weights: Weights, config: ModelConfig) {
+    // Checkpoint-generation normalization BEFORE the first tensor() call: the
+    // 5.8-family export names the trunk `model.language_model.*` with a
+    // top-level `lm_head`, ships `mtp.*` in-repo, and stores RMSNorm gains as
+    // γ−1 with HF-layout conv1d. The view maps all of that onto the names and
+    // values this graph reads (mlx-lm's post-sanitize space). No-op — nothing
+    // installed at all — for artifacts already in that space.
+    const view = qwen35WeightsView(weights);
+    if (view) weights.setView(view);
     this.config = config;
     this.tied = config.text.tieWordEmbeddings;
     this.weightsBytes = [...weights.shards.files.values()]
@@ -432,9 +593,10 @@ export class Qwen35Model {
         out.set(`${p}.linear_attn.in_proj_a`, l.linearAttn.inProjA);
         out.set(`${p}.linear_attn.out_proj`, l.linearAttn.outProj);
       }
-      out.set(`${p}.mlp.gate_proj`, l.mlp.gate);
-      out.set(`${p}.mlp.up_proj`, l.mlp.up);
-      out.set(`${p}.mlp.down_proj`, l.mlp.down);
+      // Trellis-coded MLP tensors are not LoRA targets (no adapter seam yet).
+      if (l.mlp.gate instanceof QuantizedLinear) out.set(`${p}.mlp.gate_proj`, l.mlp.gate);
+      if (l.mlp.up instanceof QuantizedLinear) out.set(`${p}.mlp.up_proj`, l.mlp.up);
+      if (l.mlp.down instanceof QuantizedLinear) out.set(`${p}.mlp.down_proj`, l.mlp.down);
     }
     return out;
   }
@@ -443,28 +605,142 @@ export class Qwen35Model {
     return this.layers.map((l) => (l.isLinear ? new SSMCache() : new KVCache()));
   }
 
+  /** The 27B append path keeps the native M1 affine arithmetic while Trellis
+   * shares up to four rows. Other shapes retain single-token execution. */
+  createAppend(policy: { hasAdapters: boolean; pagedKv: boolean }) {
+    const t = this.config.text;
+    if (policy.hasAdapters || policy.pagedKv || this.loraState.active.length || this.mrope ||
+        t.hiddenSize !== 5120 || t.intermediateSize !== 17408 || this.layers.length !== 64 ||
+        t.headDim !== 256 || t.numAttentionHeads !== 24 || t.numKeyValueHeads !== 4)
+      return null;
+    if (deviceArchitecture() !== "applegpu_g16s") return null;
+    const qualified = (linear: QuantizedLinear | TrellisLinear) => {
+      if (linear instanceof TrellisLinear) return true;
+      const { mode, bits, groupSize } = linear.spec;
+      return mode === "affine" && [2, 3, 4, 6, 8].includes(bits) && [32, 64, 128].includes(groupSize);
+    };
+    for (const layer of this.layers) {
+      const a = layer.linearAttn, s = layer.selfAttn;
+      const attention = a ? [a.inProjQkv, a.inProjZ, a.inProjB, a.inProjA, a.outProj]
+        : [s!.qProj, s!.kProj, s!.vProj, s!.oProj];
+      if (!attention.every(qualified) || ![layer.mlp.gate, layer.mlp.up, layer.mlp.down].every(qualified))
+        return null;
+    }
+    return {
+      maxChunkSize: (state: readonly Cache[]) => qwenAppendChunkSize(state[0]!.offset),
+      forwardHidden: (ids: MlxArray, cache: Cache[]): MlxArray => {
+        if (ids.shape.length !== 2 || ids.shape[0] !== 1 || ids.shape[1]! > 4)
+          throw new Error("Qwen committed-token append supports one row and at most four positions");
+        const hidden = this.embed.encode(ids);
+        return this.forwardLayers(hidden, cache, true);
+      },
+    };
+  }
+
   forwardHidden(ids: MlxArray, cache: Cache[]): MlxArray {
     const h = this.embed.encode(ids);
     return this.forwardLayers(h, cache);
   }
 
-  forwardEmbeddings(_embeds: MlxArray, _cache: Cache[], _bidir: MlxArray | null): MlxArray {
-    throw new Error("qwen3_5 vision/input-embedding path is not supported");
+  /** Active vision mRoPE request state (serial lane; set by the generation
+   *  gateway around a vision request's run, null for text-only — which keeps
+   *  the bit-exact fast-rope path). While set, forwardLayers installs the
+   *  per-forward interleaved cos/sin consumed by every full-attn layer. */
+  mrope: MropeRequestState | null = null;
+  /** Serving's memory-pressure check, scoped to the serial request. */
+  prefillMemoryGuard: (() => void) | null = null;
+  #mropeInvFreq: MlxArray | null = null;
+
+  /** Vision prefill: spliced input embeddings [1, L, H] (image features
+   *  overwriting image-token rows — the caller builds them). Caller keeps
+   *  ownership of `embeds`. bidir/ids/multimodal are the Gemma-shaped extras
+   *  and are unused: Qwen3.5 vision attends fully causally. */
+  forwardEmbeddings(
+    embeds: MlxArray, cache: Cache[], _bidir: MlxArray | null,
+    _ids: MlxArray | null = null, _multimodal: MlxArray | null = null,
+  ): MlxArray {
+    const h = ops.contiguous(embeds); // fresh handle — the layer loop consumes it
+    return this.forwardLayers(h, cache);
+  }
+
+  /** Layer-output tap (same contract as Gemma4Model.hiddenTap): the serve
+   *  loop's forwardMaybeTap sets it around prefill/verify forwards so
+   *  KV-borrowing draft sources (native Qwen MTP) can read PRE-final-norm
+   *  hiddens — layer 63's output stream, what mlx-vlm captures with
+   *  skip_final_norm=True. No-op (and graph-identical) when null. */
+  hiddenTap: { layers: Set<number>; pos?: number; captured: Map<number, MlxArray> } | null = null;
+
+  /** Store layer `i`'s residual stream ([1,L,H]) if the tap requests it.
+   *  Copies so the loop's h.dispose() can't free the capture. */
+  protected captureLayer(i: number, h: MlxArray): void {
+    const tap = this.hiddenTap;
+    if (!tap || !tap.layers.has(i)) return;
+    const H = h.shape[2]!;
+    const src = tap.pos !== undefined ? h.slice([0, tap.pos, 0], [1, tap.pos + 1, H]) : h;
+    const copy = ops.contiguous(src);
+    if (src !== h) src.dispose();
+    tap.captured.set(i, copy);
   }
 
-  protected forwardLayers(h0: MlxArray, cache: Cache[]): MlxArray {
+  protected forwardLayers(h0: MlxArray, cache: Cache[], independentRows = false): MlxArray {
     const L = h0.shape[1]!;
     // One full-attention mask shared by all full layers (same offset); linear
     // layers see no ssm mask at B=1.
     const faMask = cache[this.faIdx]!.makeMask(L, null);
-    let h = h0;
-    for (let i = 0; i < this.layers.length; i++) {
-      const next = this.layers[i]!.forward(h, faMask, cache[i]!);
-      h.dispose();
-      h = next;
+    // Vision requests: one interleaved-mRoPE cos/sin table per forward,
+    // shared by all 12 full-attention layers (positions are layer-invariant).
+    let mropeFwd: ReturnType<typeof buildMropePositions> | null = null;
+    if (this.mrope) {
+      const t = this.config.text;
+      const ropeDims = Math.trunc(t.headDim * t.partialRotaryFactor);
+      const base = t.ropeParameters.full_attention?.ropeTheta ?? 10000;
+      this.#mropeInvFreq ??= mropeInvFreq(ropeDims, base);
+      mropeFwd = buildMropePositions(
+        this.mrope, cache[this.faIdx]!.offset, L, this.#mropeInvFreq, ropeDims,
+      );
+      setActiveMrope(mropeFwd);
+    }
+    let h: MlxArray | null = h0;
+    const prof = (globalThis as Record<string, unknown>).__deltaProf as
+      Record<string, number>
+      | undefined;
+    try {
+      for (let i = 0; i < this.layers.length; i++) {
+        if (L > 1) this.prefillMemoryGuard?.();
+        const tl = prof ? performance.now() : 0;
+        const next = this.layers[i]!.forward(h, faMask, cache[i]!, independentRows);
+        h.dispose();
+        h = next;
+        const mlp = this.layers[i]!.mlp;
+        if (L > TRELLIS_MATVEC_MAX_M &&
+            [mlp.gate, mlp.up, mlp.down].some((p) => p instanceof TrellisLinear && !p.fallback)) {
+          // Materialize the cache outputs too. In particular, the tiny
+          // copied conv tail is not a dependency of `h`; leaving it lazy
+          // pins the whole prefill conv buffer until the chunk ends.
+          const state = cache[i]!.state();
+          try { ops.evalAll([h, ...state]); }
+          finally { if (cache[i]!.stateNeedsDispose) for (const a of state) a.dispose(); }
+        }
+        if (prof) {
+          ops.evalAll([h]);
+          const key = this.layers[i]!.isLinear ? "layerLin" : "layerFull";
+          prof[key] = (prof[key] ?? 0) + performance.now() - tl;
+        }
+        this.captureLayer(i, h); // native-MTP pre-final-norm tap (no-op unless set)
+      }
+      const out = disposing(h, this.finalNorm.forward(h));
+      h = null; // consumed — the finally must not double-free
+      return out;
+    } finally {
+      if (mropeFwd) {
+        setActiveMrope(null);
+        mropeFwd.posIds.dispose();
+      }
+      // A mid-loop layer throw must not strand the mask array or the
+      // in-flight [1,L,H] residual (2026-08-18 review).
+      faMask.arr?.dispose();
+      h?.dispose();
     }
-    faMask.arr?.dispose();
-    return disposing(h, this.finalNorm.forward(h));
   }
 
   logitsFromHidden(h: MlxArray): MlxArray {

package/src/model/support.ts

@@ -56,7 +56,7 @@ export function isDrafterModelType(modelType: string): boolean {
 
 /** DiffusionGemma (model_type `diffusion_gemma`): block/masked-diffusion canvas
  *  model. Non-autoregressive — routed through the diffusion engine, not the AR
- *  loop. See docs/design/diffusion-gemma-port.md. */
+ *  loop. See docs/design/generic-model-support.md. */
 export function isDiffusionGemmaConfig(config: ModelConfig): boolean {
   return config.modelType === "diffusion_gemma";
 }

package/src/model/trellis-balanced-scatter.ts

@@ -0,0 +1,51 @@
+// Experimental variant 8, k=3 / T=256 / L<=12 only.
+// The word-per-lane kernel uses 24 of 32 SIMD lanes for a 3-bit block.
+// Assign eight outputs to every lane, load adjacent packed words, and align
+// them once so each unrolled decoder has constant shifts. Row accumulation,
+// split boundaries and the separate partial reduction stay identical to v6.
+export const BALANCED_SCATTER_SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint block = thread_position_in_grid.y * (uint)SG_TG + sg;
+  const uint nBlocks = (uint)C / (uint)BT;
+  const uint zi = thread_position_in_grid.z;
+  const uint sample = zi / (uint)SPLITS;
+  const uint split = zi - sample * (uint)SPLITS;
+  if (block >= nBlocks || sample >= (uint)M) return;
+  const uint wpb = (uint)(BT * K / 32);
+  const ulong rowWords = (ulong)(INTERLEAVE ? 2 : nBlocks) * wpb;
+  const uint pBase = lane * 8u * (uint)K;
+  const uint wi = pBase >> 5u;
+  const uint off = pBase & 31u;
+  const uint next = (wi + 1u == wpb) ? 0u : wi + 1u;
+  const device uint32_t* words = codes + (INTERLEAVE
+    ? (ulong)(block / 2) * R * 2 * wpb + (block % 2) * wpb : (ulong)block * wpb);
+  const uint rowsPer = ((uint)R + (uint)SPLITS - 1u) / (uint)SPLITS;
+  const uint r0 = split * rowsPer;
+  const uint r1 = metal::min((uint)R, r0 + rowsPer);
+  const device T* xs = x + (ulong)sample * (ulong)R;
+  float acc[8];
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i < 8u; ++i) acc[i] = 0.0f;
+  for (uint r = r0; r < r1; ++r) {
+    const uint word = words[(ulong)r * rowWords + wi];
+    const uint nxt = words[(ulong)r * rowWords + next];
+    const uint lo = off == 0u ? word : (word >> off) | (nxt << (32u - off));
+    const uint hi = off == 0u ? nxt : nxt >> off;
+    const float xr = float(xs[r]);
+    const float sr = float(scales[r]);
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < 8u; ++i) {
+      const uint shift = i * (uint)K;
+      uint win = lo >> shift;
+      if (shift + (uint)L > 32u) win |= hi << (32u - shift);
+      win &= (1u << (uint)L) - 1u;
+      const float value = (float)trellis_y(win) * (1.0f / 147.800537109375f);
+      acc[i] = metal::fma(value * sr, xr, acc[i]);
+    }
+  }
+  const ulong outBase = ((ulong)sample * (ulong)SPLITS + split) * (ulong)C;
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i < 8u; ++i)
+    partial[outBase + block * (uint)BT + (uint)BT - 1u - lane * 8u - i] = acc[i];
+`;