mlx-bun 0.3.0 → 0.4.0

package/src/model/qwen3-delta.ts

@@ -5,8 +5,8 @@
 // default (use_kernel = not training), and its float accumulation / simd_sum
 // reduction order differ from the pure-ops fallback — so BIT-EXACT parity with
 // mlx-lm requires the SAME kernel, dispatched with the SAME grid/threadgroup.
-// We port the non-vectorized, non-masked variant (g.ndim == 3, mask is None),
-// which is the B=1 single-stream path (ssm_mask is None at batch 1).
+// We port the non-vectorized variant (g.ndim == 3). Padded prefill selects
+// the oracle masked specialization; unpadded execution retains its kernel.
 //
 // Numerics (must match the reference dtypes exactly — mlx infers the kernel's
 // pointer element types from the input arrays):
@@ -20,10 +20,12 @@
 // GQA is handled inside the kernel (hk_idx = hv_idx / (Hv/Hk)); q/k stay at Hk.
 
 import { MlxArray } from "../mlx/array";
+import { materializeCopy } from "../mlx/materialize";
 import { Dtype } from "../mlx/ffi";
 import { MetalKernel } from "../mlx/metal-kernel";
 import { CompiledFunction } from "../mlx/compile";
 import * as ops from "../mlx/ops";
+import type { SsmPrefillPadding } from "./ssm-prefill-padding";
 import type { Cache, Mask } from "./gemma4-base";
 
 // Verbatim body of mlx-lm's gated_delta_step (has_mask=False, vectorized=False),
@@ -104,7 +106,15 @@ const SOURCE = String.raw`
 `;
 
 let kernel: MetalKernel | null = null;
-function getKernel(): MetalKernel {
+let maskedKernel: MetalKernel | null = null;
+function getKernel(masked = false): MetalKernel {
+  if (masked) return maskedKernel ??= new MetalKernel({
+    name: "gated_delta_step_mask",
+    inputNames: ["q", "k", "v", "g", "beta", "state_in", "Tin", "mask"],
+    outputNames: ["y", "state_out"],
+    source: SOURCE.replace("if (true)", "if (mask[b_idx * T + t])"),
+    ensureRowContiguous: true,
+  });
   if (!kernel)
     kernel = new MetalKernel({
       name: "gated_delta_step",
@@ -154,7 +164,7 @@ export function computeG(aLog: MlxArray, a: MlxArray, dtBias: MlxArray): MlxArra
  *   state: [B, Hv, Dv, Dk] f32 (or null → zeros) */
 export function gatedDeltaUpdate(
   q: MlxArray, k: MlxArray, v: MlxArray, a: MlxArray, b: MlxArray,
-  aLog: MlxArray, dtBias: MlxArray, state: MlxArray | null,
+  aLog: MlxArray, dtBias: MlxArray, state: MlxArray | null, mask: MlxArray | null = null,
 ): [MlxArray, MlxArray] {
   const [B, , Hk, Dk] = q.shape as [number, number, number, number];
   const [, , Hv, Dv] = v.shape as [number, number, number, number];
@@ -171,7 +181,9 @@ export function gatedDeltaUpdate(
 
   const T = q.shape[1]!;
   const tArr = MlxArray.fromInt32(new Int32Array([T]), [1]);
-  const [y, stateOut] = getKernel().apply([q, k, v, g, beta, stateIn, tArr], {
+  const inputs = [q, k, v, g, beta, stateIn, tArr];
+  if (mask) inputs.push(mask);
+  const [y, stateOut] = getKernel(mask !== null).apply(inputs, {
     outputs: [
       { shape: [B, T, Hv, Dv], dtype: q.dtype },
       { shape: [B, Hv, Dv, Dk], dtype: Dtype.float32 },
@@ -207,6 +219,7 @@ export interface SsmSpecRound {
   prevConv: MlxArray | null;
   prevRecurrent: MlxArray | null;
   prevOffset: number;
+  prevOffsets: number[] | null;
   /** Pre-conv in_proj output [B,S,convDim] — position-local. */
   qkv: MlxArray | null;
   /** Gate/beta projections [B,S,Hv] — position-local. */
@@ -216,14 +229,15 @@ export interface SsmSpecRound {
   S: number;
   /** Layer-bound prefix replay: called AFTER the snapshot is restored onto
    *  the cache; advances conv/recurrent/offset by `keep` tokens. */
-  replay: ((cache: SSMCache, keep: number) => void) | null;
+  replay: ((cache: SSMCache, keep: number | readonly number[]) => void) | null;
 }
 
 export class SSMCache implements Cache {
   conv: MlxArray | null = null;
   recurrent: MlxArray | null = null;
   offset = 0;
-  /** Live speculative verify round (serial spec lane only; null otherwise). */
+  prefillPadding: SsmPrefillPadding | null = null;
+  /** Live speculative verify round; null outside a verify transaction. */
   specRound: SsmSpecRound | null = null;
   /** Per-row token coverage (batch lane only; null on serial B=1 caches,
    *  where `offset` IS the row's count). mlx-lm's ArraysCache tracks no
@@ -248,15 +262,13 @@ export class SSMCache implements Cache {
     throw new Error("SSMCache has no KV updateAndFetch (gated-DeltaNet layer)");
   }
 
-  makeMask(_N: number, _windowSize: number | null): Mask {
-    // ssm_mask is None both single-stream and in the batch lane: rows
-    // solo-prefill UNPADDED at B=1 (state never sees pad tokens), and
-    // batched decode feeds one real token per row — so unlike mlx-lm's
-    // left-padded batch-prefill there is never a pad position to mask.
-    return { mode: "", arr: null };
+  makeMask(N: number, _windowSize: number | null): Mask {
+    const arr = this.prefillPadding?.makeMask(N) ?? null;
+    return { mode: arr ? "array" : "", arr };
   }
 
   advance(n: number): void {
+    if (this.prefillPadding) { this.advanceRows(this.prefillPadding.advance(n)); return; }
     this.offset += n;
     // Batched rows step together (one forward advances every row), so the
     // per-row counts move in lockstep; they DIFFER only in their merge-time
@@ -264,6 +276,13 @@ export class SSMCache implements Cache {
     if (this.offsets) for (let i = 0; i < this.offsets.length; i++) this.offsets[i]! += n;
   }
 
+  /** Methods may retain different verified prefixes for each active request. */
+  advanceRows(counts: readonly number[]): void {
+    if (!this.offsets) { this.offset += counts[0]!; return; }
+    for (let row = 0; row < this.offsets.length; row++) this.offsets[row]! += counts[row]!;
+    this.offset = Math.max(...this.offsets);
+  }
+
   state(): MlxArray[] {
     const out: MlxArray[] = [];
     if (this.conv) out.push(this.conv);
@@ -279,7 +298,7 @@ export class SSMCache implements Cache {
     throw new Error("SSMCache is not trimmable");
   }
 
-  // --- speculative verify-round support (serial spec lane, B=1) ------------
+  // --- speculative verify-round support -----------------------------------
   // The serve loop arms a round before the verify forward and resolves it
   // after the accept walk. The layer's forward, seeing an armed round, hands
   // its replaced state slots to the round instead of disposing them, records
@@ -294,6 +313,7 @@ export class SSMCache implements Cache {
       prevConv: null,
       prevRecurrent: null,
       prevOffset: this.offset,
+      prevOffsets: this.offsets ? [...this.offsets] : null,
       qkv: null,
       a: null,
       b: null,
@@ -306,12 +326,12 @@ export class SSMCache implements Cache {
     this.#dropSpecRound();
   }
 
-  specRoundRollback(keep: number): void {
+  specRoundRollback(keep: number | readonly number[]): void {
     const r = this.specRound;
     if (!r) throw new Error("SSMCache.specRoundRollback without an armed round");
     if (r.armed)
       throw new Error("SSMCache.specRoundRollback before the verify forward recorded");
-    if (keep < 0 || keep > r.S)
+    if (typeof keep === "number" ? keep < 0 || keep > r.S : keep.some(count => count < 0 || count > r.S))
       throw new Error(`SSMCache.specRoundRollback keep=${keep} outside window S=${r.S}`);
     // Restore the pre-round snapshot (ownership moves back to the cache).
     this.conv?.dispose();
@@ -321,9 +341,8 @@ export class SSMCache implements Cache {
     r.prevConv = null;
     r.prevRecurrent = null;
     this.offset = r.prevOffset;
-    if (this.offsets)
-      throw new Error("SSMCache.specRoundRollback on a batched cache (serial lane only)");
-    if (keep > 0) {
+    this.offsets = r.prevOffsets;
+    if (typeof keep === "number" ? keep > 0 : keep.some(count => count > 0)) {
       if (!r.replay) throw new Error("SSMCache.specRoundRollback with no recorded replay");
       r.replay(this, keep); // advances conv/recurrent/offset by `keep`
     }
@@ -347,6 +366,7 @@ export class SSMCache implements Cache {
     this.recurrent?.dispose();
     this.conv = null;
     this.recurrent = null;
+    this.prefillPadding = null;
   }
 
   // --- batch-lane dynamic-B ops (BatchScheduler only) ----------------------
@@ -386,25 +406,30 @@ export class SSMCache implements Cache {
 
   /** Evict rows not in `keep` (ascending row indices), in place. */
   filter(keep: number[]): void {
-    if (!this.conv || !this.recurrent) return;
-    const idx = ops.fromInt32(keep, [keep.length]);
-    const conv = ops.takeAxis(this.conv, idx, 0);
-    const recurrent = ops.takeAxis(this.recurrent, idx, 0);
-    idx.dispose();
-    this.conv.dispose();
-    this.recurrent.dispose();
-    this.conv = conv;
-    this.recurrent = recurrent;
-    if (this.offsets) this.offsets = keep.map((i) => this.offsets![i]!);
+    if (this.conv && this.recurrent) {
+      const idx = ops.fromInt32(keep, [keep.length]);
+      const conv = ops.takeAxis(this.conv, idx, 0);
+      const recurrent = ops.takeAxis(this.recurrent, idx, 0);
+      idx.dispose();
+      this.conv.dispose();
+      this.recurrent.dispose();
+      this.conv = conv;
+      this.recurrent = recurrent;
+    }
+    this.prefillPadding?.filter(keep);
+    if (this.offsets) {
+      this.offsets = keep.map((i) => this.offsets![i]!);
+      this.offset = Math.max(0, ...this.offsets);
+    }
   }
 
   filterRows(keep: readonly number[]): void { this.filter([...keep]); }
 
   /** Row `i` as a fresh SERIAL cache — port of mlx-lm ArraysCache.extract
    *  (cache.py:673-676: `cache.cache = [c[idx : idx + 1] for c in self.cache]`,
-   *  a B-axis slice of every state slot). Ours are OWNED contiguous copies
-   *  (the entry outlives the batch; the batched buffers must free when the
-   *  batch moves on — same rule as extractKVRow, batched-mask.ts). Offset =
+   *  a B-axis slice of every state slot). A proper subset gets compact storage
+   *  so retaining one row does not pin the other rows. A singleton retains an
+   *  independent handle over its complete allocation. Offset =
    *  the row's OWN coverage, not the shared max. */
   extractRow(i: number): SSMCache {
     if (!this.conv || !this.recurrent)
@@ -414,10 +439,8 @@ export class SSMCache implements Cache {
       lo[0] = i;
       const hi = [...a.shape];
       hi[0] = i + 1;
-      const view = a.slice(lo, hi);
-      const own = ops.copyOf(view); // TRUE copy: contiguous(view) is a no-op VIEW when already contiguous — pins the source buffer (2026-08-20 DeltaNet conv leak class)
-      view.dispose();
-      return own;
+      using view = a.slice(lo, hi);
+      return a.shape[0] === 1 ? ops.contiguous(view) : materializeCopy(view);
     };
     const out = new SSMCache();
     out.conv = cut(this.conv);

package/src/model/qwen3_5.ts

@@ -10,6 +10,7 @@
 // shared maybeQuantizeKv path).
 
 import type { ModelConfig } from "../config";
+import { gatedDeltaState } from "./qwen3-delta-state";
 import type { Weights } from "../weights";
 import { MlxArray } from "../mlx/array";
 import { Dtype, deviceArchitecture } from "../mlx/ffi";
@@ -23,7 +24,6 @@ import {
   disposing,
   KVCache,
   LoraState,
-  QuantizedKVCache,
   QuantizedEmbedding,
   QuantizedLinear,
   quantizedSdpa,
@@ -155,16 +155,20 @@ export class GatedDeltaNet {
     this.normWeight = weights.tensor(`${prefix}.norm.weight`);
   }
 
-  #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;
+  #convolve(qkv: MlxArray, state: MlxArray, rowLengths?: readonly number[]): [MlxArray, MlxArray] {
+    if (this.convolution) {
+      const result = this.convolution(qkv, state, this.convWeight);
+      if (!rowLengths) return result;
+      using input = ops.concatAxis([state, qkv], 1);
+      const tail = this.#convTail(input, qkv.shape[1]!, rowLengths);
+      result[1].dispose();
+      return [result[0], tail];
+    }
+    const [, S, D] = qkv.shape as [number, number, number];
     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 tail = this.#convTail(input, S, rowLengths);
     const conv = ops.conv1d(input, this.convWeight, 1, 0, 1, D);
     input.dispose();
     const out = compiledSilu(conv);
@@ -172,7 +176,19 @@ export class GatedDeltaNet {
     return [out, tail];
   }
 
-  forward(x: MlxArray, cache: SSMCache, independentRows = false): MlxArray {
+  #convTail(input: MlxArray, processed: number, rowLengths?: readonly number[]): MlxArray {
+    const [B, , D] = input.shape as [number, number, number];
+    const nKeep = this.convKernel - 1;
+    if (rowLengths) {
+      using indices = ops.fromInt32(rowLengths.flatMap(length =>
+        Array.from({ length: nKeep }, (_, position) => length + position)), [B, nKeep, 1]);
+      return ops.takeAlongAxis(input, indices, 1);
+    }
+    using view = input.slice([0, processed, 0], [B, processed + nKeep, D]);
+    return materializeCopy(view);
+  }
+
+  forward(x: MlxArray, cache: SSMCache, independentRows = false, mask?: MlxArray | null): MlxArray {
     const [B, S] = x.shape as [number, number, number];
     const convDim = this.keyDim * 2 + this.valueDim;
     const nKeep = this.convKernel - 1;
@@ -194,7 +210,7 @@ export class GatedDeltaNet {
     }
 
     const t0 = prof ? performance.now() : 0;
-    const qkv = this.inProjQkv.forward(x, independentRows); // [B,S,convDim]
+    let 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, independentRows); // [B,S,numVHeads]
@@ -202,10 +218,19 @@ export class GatedDeltaNet {
     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).
+    // The model shares one padding mask across recurrent layers. Direct layer
+    // callers can request the same cache-owned mask without a model wrapper.
+    using localMask = mask === undefined ? cache.prefillPadding?.makeMask(S) ?? null : null;
+    const ssmMask = mask ?? localMask;
+    if (ssmMask) {
+      using expanded = ops.reshape(ssmMask, [B, S, 1]);
+      using zero = MlxArray.fromBytesCopy(new Uint8Array(qkv.dtype === Dtype.float32 ? 4 : 2), [], qkv.dtype);
+      qkv = disposing(qkv, ops.where(expanded, qkv, zero));
+    }
+    // Convolution keeps each row's last real-token tail, including empty rows.
     const convState =
       cache.conv ?? ops.zeros([B, nKeep, convDim], x.dtype);
-    const [convOut, newConv] = this.#convolve(qkv, convState);
+    const [convOut, newConv] = this.#convolve(qkv, convState, cache.prefillPadding?.convolutionLengths(S));
     if (!cache.conv) convState.dispose();
     if (spec) spec.qkv = qkv;
     else qkv.dispose();
@@ -236,7 +261,7 @@ export class GatedDeltaNet {
     const tk = prof ? performance.now() : 0;
 
     const [out, newState] = gatedDeltaUpdate(
-      q, k, v, a, b, this.aLog, this.dtBias, cache.recurrent,
+      q, k, v, a, b, this.aLog, this.dtBias, cache.recurrent, ssmMask,
     );
     if (prof) { ops.evalAll([out]); prof.kernel = (prof.kernel ?? 0) + performance.now() - tk; }
     const to = prof ? performance.now() : 0;
@@ -275,18 +300,21 @@ export class GatedDeltaNet {
    *  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 {
+  #replaySpecPrefix(cache: SSMCache, keep: number | readonly number[]): void {
     const r = cache.specRound;
-    if (!r || !r.qkv || !r.a || !r.b)      throw new Error("GatedDeltaNet replay without recorded round inputs");
+    if (!r || !r.qkv || !r.a || !r.b)
+      throw new Error("GatedDeltaNet replay without recorded round inputs");
     const B = r.qkv.shape[0]!;
+    const count = typeof keep === "number" ? keep : Math.max(...keep);
+    const rowLengths = typeof keep === "number" ? undefined : keep;
     const convDim = this.keyDim * 2 + this.valueDim;
     const nKeep = this.convKernel - 1;
 
-    const qkvPfxView = r.qkv.slice([0, 0, 0], [B, keep, convDim]);
+    const qkvPfxView = r.qkv.slice([0, 0, 0], [B, count, 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);
+    const [convOut, newConv] = this.#convolve(qkvPfx, convState, rowLengths);
     if (!cache.conv) convState.dispose();
     qkvPfx.dispose();
     cache.conv?.dispose();
@@ -296,36 +324,45 @@ export class GatedDeltaNet {
       convOut, [this.keyDim, 2 * this.keyDim], -1,
     ) as [MlxArray, MlxArray, MlxArray];
     convOut.dispose();
-    let q = ops.reshape(qFlat, [B, keep, this.numKHeads, this.headKDim]);
+    let q = rowLengths ? null : ops.reshape(qFlat, [B, count, this.numKHeads, this.headKDim]);
     qFlat.dispose();
-    let k = ops.reshape(kFlat, [B, keep, this.numKHeads, this.headKDim]);
+    let k = ops.reshape(kFlat, [B, count, this.numKHeads, this.headKDim]);
     kFlat.dispose();
-    const v = disposing(vFlat, ops.reshape(vFlat, [B, keep, this.numVHeads, this.headVDim]));
+    const v = disposing(vFlat, ops.reshape(vFlat, [B, count, 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));
+    if (q) {
+      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 aPfxView = r.a.slice([0, 0, 0], [B, count, this.numVHeads]);
     const aPfx = ops.contiguous(aPfxView);
     aPfxView.dispose();
-    const bPfxView = r.b.slice([0, 0, 0], [B, keep, this.numVHeads]);
+    const bPfxView = r.b.slice([0, 0, 0], [B, count, 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();
+    let newState: MlxArray;
+    if (rowLengths) {
+      // Unequal prefixes need a row-length-aware recurrence. The state-only
+      // kernel omits query/output work. Uniform replay retains its measured
+      // kernel until the separate optimization demonstrates a serving win.
+      newState = gatedDeltaState(k, v, aPfx, bPfx, this.aLog, this.dtBias, cache.recurrent, rowLengths);
+    } else {
+      const [output, state] = gatedDeltaUpdate(q!, k, v, aPfx, bPfx, this.aLog, this.dtBias, cache.recurrent);
+      output.dispose(); newState = state;
+    }
+    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);
+    if (typeof keep === "number") cache.advance(keep);
+    else cache.advanceRows(keep);
   }
 
   private rmsNormGated(hidden: MlxArray, gate: MlxArray): MlxArray {
@@ -415,11 +452,16 @@ export class Qwen3Attention {
     }
 
     let attn: MlxArray;
-    if (cache instanceof QuantizedKVCache) {
-      const [keys, values] = cache.updateAndFetchQuantized(k, v);
+    const quantized = cache.quantizedAttention;
+    if (cache.attentionState) {
+      const view = cache.attentionState.appendAndFetch(k, v);
+      k.dispose(); v.dispose();
+      try { attn = view.attend(q, this.scale, mask); } finally { view.dispose(); }
+    } else if (quantized) {
+      const [keys, values] = quantized.updateAndFetchQuantized(k, v);
       k.dispose();
       v.dispose();
-      attn = quantizedSdpa(q, keys, values, this.scale, mask, cache.groupSize, cache.bits);
+      attn = quantizedSdpa(q, keys, values, this.scale, mask, quantized.groupSize, quantized.bits);
       disposeTriple(keys);
       disposeTriple(values);
     } else {
@@ -517,10 +559,10 @@ export class Qwen3Layer {
     this.postAttnNorm = new RMSNorm(weights.tensor(`${prefix}.post_attention_layernorm.weight`), config.text.rmsNormEps);
   }
 
-  forward(x: MlxArray, faMask: Mask, cache: Cache, independentRows = false): MlxArray {
+  forward(x: MlxArray, faMask: Mask, cache: Cache, independentRows = false, ssmMask?: MlxArray | null): MlxArray {
     const xn = this.inputNorm.forward(x);
     const r = this.isLinear
-      ? this.linearAttn!.forward(xn, cache as SSMCache, independentRows)
+      ? this.linearAttn!.forward(xn, cache as SSMCache, independentRows, ssmMask)
       : this.selfAttn!.forward(xn, faMask, cache, independentRows);
     xn.dispose();
     const h = ops.add(x, r);
@@ -647,8 +689,6 @@ export class Qwen35Model {
    *  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
@@ -687,6 +727,7 @@ export class Qwen35Model {
     // 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);
+    using ssmMask = (cache[0] as SSMCache).prefillPadding?.makeMask(L) ?? null;
     // 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;
@@ -706,14 +747,13 @@ export class Qwen35Model {
       | 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);
+        const next = this.layers[i]!.forward(h, faMask, cache[i]!, independentRows, ssmMask);
         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)) {
+        if (L > TRELLIS_MATVEC_MAX_M) {
+          // Affine prefills also need bounded evaluation: a retained prefix
+          // history can leave too little room for the entire deferred graph.
           // 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.

package/src/model/rotating-kv-donor.ts

@@ -0,0 +1,35 @@
+import type { KvDonorAttention } from "./gemma4-base";
+import { BatchedRotatingState } from "./batched-rotating-state";
+import { BatchedRotatingQuantCache } from "./batched-rotating-quant";
+import { plainRowStorage, quantizedRowStorage, temporalStorageView } from "./batched-row-storage";
+import { quantizedDonorAttention } from "./kv-attention-view";
+import { decodedKvDonorAttention } from "./decoded-kv-donor";
+import type { RotatingRing } from "./rotating-row-transaction";
+
+/** Capture the same chronological donor rectangle for plain and encoded rings. */
+export function captureRotatingDonorAttention(ring: RotatingRing): KvDonorAttention {
+  const position = new BatchedRotatingState(ring.maxSize, []);
+  position.restore(ring.positionSnapshot);
+  const offsets = position.offsets.map((_, row) => position.validOffset(row));
+  const ends = offsets.map((offset, row) => position.activeLength - (position.offsets[row]! - offset));
+  // Pending prefill keeps a shared physical rectangle, including when mixed
+  // rows are captured separately. Cropping from each row's valid end would
+  // give siblings different widths and can discard an earlier valid window.
+  const from = position.hasPendingPadding ? 0 : Math.max(0, position.activeLength - ring.maxSize);
+  const range = { from, to: position.activeLength };
+  const validity = { offsets,
+    starts: position.leftPad.map((pad, row) => Math.max(0, pad - from, ends[row]! - from - ring.maxSize)),
+    ends: ends.map(end => Math.max(0, end - from)) };
+  if (ring instanceof BatchedRotatingQuantCache) {
+    const keys = temporalStorageView(quantizedRowStorage, ring.keys!, position, range);
+    let values: import("../mlx/ops").QuantizedTensor;
+    try { values = temporalStorageView(quantizedRowStorage, ring.values!, position, range); }
+    catch (error) { quantizedRowStorage.dispose(keys); throw error; }
+    return quantizedDonorAttention(keys, values, validity, ring.groupSize, ring.bits);
+  }
+  const keys = temporalStorageView(plainRowStorage, ring.keys!, position, range);
+  let values: import("../mlx/array").MlxArray;
+  try { values = temporalStorageView(plainRowStorage, ring.values!, position, range); }
+  catch (error) { keys.dispose(); throw error; }
+  return decodedKvDonorAttention({ keys, values, ...validity });
+}