mlx-bun 0.3.0 → 0.4.0

package/src/model/batched-rotating-state.ts

@@ -1,3 +1,27 @@
+import type { PrefillPadding } from "./gemma4-base";
+
+export interface RotatingPositionSnapshot {
+  readonly maxSize: number;
+  readonly offsets: readonly number[];
+  readonly leftPad: readonly number[];
+  readonly ringIndex: number;
+  readonly totalOffset: number;
+  readonly rotated: boolean;
+  readonly prefillEnds?: readonly number[];
+}
+
+/** Acceptance changes row positions, independently of plain or encoded planes. */
+export function rotatingRollbackPlan(snapshot: RotatingPositionSnapshot, before: readonly number[],
+  keep: number | readonly number[], preserveColumns = false): { trim: number } | { shifts: number[]; position: RotatingPositionSnapshot } {
+  const counts = before.map((_, row) => typeof keep === "number" ? keep : keep[row]!);
+  if (!preserveColumns && before.length === 1 && snapshot.totalOffset < snapshot.maxSize)
+    return { trim: snapshot.offsets[0]! - before[0]! - counts[0]! };
+  const offsets = before.map((offset, row) => offset + counts[row]!);
+  const shifts = snapshot.offsets.map((offset, row) => offset - offsets[row]!);
+  return { shifts, position: { ...snapshot, offsets,
+    leftPad: snapshot.leftPad.map((pad, row) => pad + shifts[row]!) } };
+}
+
 /** Host-side row and ring bookkeeping shared by plain and quantized storage. */
 export class BatchedRotatingState {
   readonly maxSize: number;
@@ -6,6 +30,7 @@ export class BatchedRotatingState {
   ringIndex = 0;
   totalOffset = 0;
   rotated = false;
+  #prefillEnds?: number[];
 
   constructor(maxSize: number, leftPad: readonly number[], offsets?: readonly number[]) {
     this.maxSize = maxSize;
@@ -13,8 +38,45 @@ export class BatchedRotatingState {
     this.offsets = offsets ? [...offsets] : leftPad.map((pad) => -pad);
   }
 
+  snapshot(): RotatingPositionSnapshot {
+    return { maxSize: this.maxSize, offsets: [...this.offsets], leftPad: [...this.leftPad],
+      ringIndex: this.ringIndex, totalOffset: this.totalOffset, rotated: this.rotated,
+      ...(this.#prefillEnds ? { prefillEnds: [...this.#prefillEnds] } : {}) };
+  }
+  restore(snapshot: RotatingPositionSnapshot): void {
+    this.offsets = [...snapshot.offsets]; this.leftPad = [...snapshot.leftPad];
+    this.ringIndex = snapshot.ringIndex; this.totalOffset = snapshot.totalOffset; this.rotated = snapshot.rotated;
+    this.#prefillEnds = snapshot.prefillEnds ? [...snapshot.prefillEnds] : undefined;
+  }
+
+  preparePrefill(padding: PrefillPadding, rightPaddedBatch = padding.rightPadding?.some(pad => pad > 0)): void {
+    if (!this.batchSize) { this.leftPad = padding.lengths.map(() => 0); this.offsets = [...this.leftPad]; }
+    if (padding.leftPadding) {
+      this.leftPad = this.leftPad.map((pad, row) => pad + padding.leftPadding![row]!);
+      this.offsets = this.offsets.map((offset, row) => offset - padding.leftPadding![row]!);
+    }
+    this.#prefillEnds = rightPaddedBatch
+      ? padding.lengths.map((length, row) => length + this.offsets[row]!) : undefined;
+  }
+  get hasPendingPadding(): boolean { return this.#prefillEnds !== undefined; }
+  validOffset(row: number): number {
+    return Math.min(this.offsets[row]!, this.#prefillEnds?.[row] ?? this.offsets[row]!);
+  }
+
+  prefillRoll(): number[] | undefined {
+    return this.#prefillEnds?.map((end, row) => Math.max(0, this.offsets[row]! - end));
+  }
+  #rollPadding(): void {
+    const shifts = this.prefillRoll();
+    if (!shifts) return;
+    this.leftPad = this.leftPad.map((pad, row) => pad + shifts[row]!);
+    this.offsets = this.offsets.map((offset, row) => offset - shifts[row]!);
+  }
+  finalizePrefill(): void { this.#rollPadding(); this.#prefillEnds = undefined; }
+
   get batchSize(): number { return this.leftPad.length; }
   get validLength(): number { return Math.min(this.totalOffset, this.maxSize); }
+  get activeLength(): number { return this.rotated ? this.maxSize : this.ringIndex; }
   get trimmable(): boolean { return this.totalOffset < this.maxSize; }
 
   markGrown(previousOffset: number): void {
@@ -43,9 +105,22 @@ export class BatchedRotatingState {
     this.ringIndex += tokens;
   }
 
+  /** mlx-lm BatchRotatingKVCache._update_concat: the block keeps up to
+   * maxSize - 1 historical columns so each query has a complete window. */
+  commitConcat(tokens: number, historicalLength: number): void {
+    if (historicalLength) this.#rollPadding();
+    const trim = Math.max(0, historicalLength - this.maxSize + 1);
+    this.leftPad = this.leftPad.map(pad => pad - trim);
+    this.totalOffset += tokens;
+    this.offsets = this.offsets.map(offset => offset + tokens);
+    this.ringIndex = historicalLength - trim + tokens;
+    this.rotated = false;
+  }
+
   filter(keep: readonly number[]): void {
     this.offsets = keep.map((row) => this.offsets[row]!);
     this.leftPad = keep.map((row) => this.leftPad[row]!);
+    if (this.#prefillEnds) this.#prefillEnds = keep.map(row => this.#prefillEnds![row]!);
   }
 
   trim(tokens: number): number {
@@ -61,6 +136,7 @@ export class BatchedRotatingState {
     this.ringIndex = width;
     this.offsets = [...offsets];
     this.rotated = false;
+    this.#prefillEnds = undefined;
   }
 
   /** Physical ranges that expose a ring buffer in temporal order. */

package/src/model/batched-rotating.ts

@@ -6,12 +6,10 @@
 // column is the same for every row), with PER-ROW `offset` (absolute position,
 // drives RoPE) and `leftPad` (padding columns in the buffer, masked out).
 //
-// SCOPE: only the pieces the scheduler needs — `merge` (assemble from solo
-// prefills), the N=1 decode update (`_update_in_place`), `make_mask`, `filter`,
-// and temporalView (extract). The N>1 batched-PREFILL path
-// (`_update_concat`/`finalize`/`_lengths`) is NOT ported: the scheduler
-// solo-prefills each request (single-stream RotatingKVCache) then merges, so
-// this cache only ever sees N=1 updates. Positions (offset/leftPad/_idx/_offset/
+// Includes merge, one-token writes, block prefill, padding finalization, masks
+// and row extraction. Prefill preparation is a storage capability; the serving
+// scheduler's use of whole prompt cohorts is a separate concern.
+// Positions (offset/leftPad/_idx/_offset/
 // rotated) are tracked on the HOST (small deterministic ints) so make_mask is
 // built in JS like buildBatchedDecodeMask; only K/V live on device.
 //
@@ -24,12 +22,14 @@
 
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
-import { RotatingKVCache, type Cache, type Mask } from "./gemma4-base";
-import { BatchedRotatingState } from "./batched-rotating-state";
+import { RotatingKVCache, type Cache, type Mask, type PaddedPrefillCache, type PrefillPadding } from "./gemma4-base";
+import { BatchedRotatingState, type RotatingPositionSnapshot } from "./batched-rotating-state";
 import {
+  appendRotatingStorage,
   mergeStorageRows,
   plainRowStorage,
   temporalStorageView,
+  rowRollIndices,
 } from "./batched-row-storage";
 
 const STEP = 256;
@@ -45,14 +45,14 @@ const STEP = 256;
  *  update's decrement is the persistent one, applied separately by the cache). */
 export function buildBatchedRotatingMask(
   B: number, N: number, leftPad: number[],
-  maxSize: number, window: number, idx: number, offsetScalar: number, rotated: boolean,
+  maxSize: number, window: number, idx: number, offsetScalar: number, rotated: boolean, block = N > 1,
 ): MlxArray {
   const off = Math.min(maxSize - 1, offsetScalar);
   const S = off + N;
 
   // Local (non-persisted) leftPad for the mask: trim + rotation shrink it.
-  const trimSize = idx - maxSize + (N > 1 ? 1 : 0);
-  const isRot = N === 1 && (rotated || idx >= maxSize);
+  const trimSize = idx - maxSize + (block ? 1 : 0);
+  const isRot = !block && (rotated || idx >= maxSize);
   const lp = leftPad.map((x) => x - (trimSize > 0 ? trimSize : 0) - (isRot ? 1 : 0));
 
   // roll(shift): physical column = (temporal column + shift) mod S.
@@ -80,8 +80,9 @@ export function buildBatchedRotatingMask(
   return mask;
 }
 
-/** Faithful port of mlx-lm BatchRotatingKVCache (decode-only — see file header). */
-export class BatchedRotatingCache implements Cache {
+/** Port of mlx-lm BatchRotatingKVCache; padded singleton chunks use its
+ * concat operation with the matching block mask until finalization. */
+export class BatchedRotatingCache implements Cache, PaddedPrefillCache {
   keys: MlxArray | null = null;
   values: MlxArray | null = null;
   readonly #rows: BatchedRotatingState;
@@ -103,6 +104,29 @@ export class BatchedRotatingCache implements Cache {
    *  misroutes a batched cache into a serial-only path. */
   signature(): string { return "kv:rotating-plain"; }
 
+  get positionSnapshot(): RotatingPositionSnapshot { return this.#rows.snapshot(); }
+  validOffset(row: number): number { return this.#rows.validOffset(row); }
+  preparePrefill(padding: PrefillPadding, rightPaddedBatch?: boolean): void { this.#rows.preparePrefill(padding, rightPaddedBatch); this.releaseRopeArr(); }
+  finalizePrefill(): void {
+    const shifts = this.#rows.prefillRoll();
+    if (shifts?.some(shift => shift !== 0) && this.keys && this.values) {
+      using indices = rowRollIndices(this.keys.shape[2]!, shifts);
+      const keys = plainRowStorage.rollRows(this.keys, indices);
+      let values: MlxArray;
+      try { values = plainRowStorage.rollRows(this.values, indices); }
+      catch (error) { keys.dispose(); throw error; }
+      this.keys.dispose(); this.values.dispose(); this.keys = keys; this.values = values;
+    }
+    this.#rows.finalizePrefill(); this.releaseRopeArr();
+  }
+  /** Adopt owned tensors without changing their physical ring columns. */
+  static adoptPhysical(keys: MlxArray | null, values: MlxArray | null, position: RotatingPositionSnapshot): BatchedRotatingCache {
+    const cache = new BatchedRotatingCache(position.maxSize, [...position.leftPad]);
+    cache.#rows.restore(position); cache.keys = keys; cache.values = values;
+
+    return cache;
+  }
+
   get offsetArr(): number[] { return this.#rows.offsets; }
   get leftPad(): number[] { return this.#rows.leftPad; }
   get batchSize(): number { return this.#rows.batchSize; }
@@ -133,7 +157,7 @@ export class BatchedRotatingCache implements Cache {
       mode: "array",
       arr: buildBatchedRotatingMask(
         this.#B, N, this.leftPad, this.maxSize, window,
-        this.#rows.ringIndex, this.offset, this.#rows.rotated,
+        this.#rows.ringIndex, this.offset, this.#rows.rotated, N > 1 || this.#rows.hasPendingPadding,
       ),
     };
   }
@@ -142,8 +166,7 @@ export class BatchedRotatingCache implements Cache {
   updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
     const [B, H, S, D] = k.shape as [number, number, number, number];
     const vD = v.shape[3]!;
-    if (S !== 1)
-      throw new Error("BatchedRotatingCache supports N=1 decode updates only (solo-prefill then merge)");
+    if (S !== 1 || this.#rows.hasPendingPadding) return this.#updateConcat(k, v);
     const prev = this.offset;
 
     // Grow the buffer (in STEP chunks) until it reaches maxSize.
@@ -202,6 +225,21 @@ export class BatchedRotatingCache implements Cache {
     ];
   }
 
+  #updateConcat(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    const keys = appendRotatingStorage(plainRowStorage, this.keys, k, this.#rows);
+    let values: MlxArray;
+    try { values = appendRotatingStorage(plainRowStorage, this.values, v, this.#rows); }
+    catch (error) { keys.dispose(); throw error; }
+    const history = this.keys ? this.#rows.activeLength : 0;
+    this.keys?.dispose(); this.values?.dispose();
+    this.keys = keys; this.values = values;
+    this.#rows.commitConcat(k.shape[2]!, history);
+    return [
+      plainRowStorage.slice(keys, 0, this.#B, 0, this.#rows.ringIndex),
+      plainRowStorage.slice(values, 0, this.#B, 0, this.#rows.ringIndex),
+    ];
+  }
+
   /** Free the per-step RoPE array without disposing KV (wrapper-rebuild path). */
   releaseRopeArr(): void {
     this.#ropeArr?.dispose();
@@ -213,8 +251,12 @@ export class BatchedRotatingCache implements Cache {
   temporalView(): [MlxArray, MlxArray] {
     if (!this.keys || !this.values) throw new Error("cache is empty");
     return [
-      temporalStorageView(plainRowStorage, this.keys, this.#rows),
-      temporalStorageView(plainRowStorage, this.values, this.#rows),
+      temporalStorageView(plainRowStorage, this.keys, this.#rows, {
+        from: Math.max(0, this.#rows.activeLength - this.maxSize), to: this.#rows.activeLength,
+      }),
+      temporalStorageView(plainRowStorage, this.values, this.#rows, {
+        from: Math.max(0, this.#rows.activeLength - this.maxSize), to: this.#rows.activeLength,
+      }),
     ];
   }
 
@@ -228,15 +270,15 @@ export class BatchedRotatingCache implements Cache {
    *  oracle's `cache._idx = cache.keys.shape[2]`. Bit-exact vs a solo run:
    *  merge/decode/filter keep each row's ring bytes identical to the serial
    *  cache's (tests/batched-rotating) and this is a pure slice+copy. */
-  extractRow(i: number): RotatingKVCache | null {
+  extractRow(i: number, limit?: number): RotatingKVCache | null {
     if (!this.keys || !this.values) return null;
-    const pad = Math.max(0, this.leftPad[i]!);
+    const pad = Math.max(0, this.leftPad[i]!, limit === undefined ? 0 : this.#rows.activeLength - limit);
     const c = new RotatingKVCache(this.maxSize);
     const k = temporalStorageView(plainRowStorage, this.keys, this.#rows, {
-      row: i, from: pad, copy: true,
+      row: i, from: pad, to: this.#rows.activeLength, copy: true,
     });
     const v = temporalStorageView(plainRowStorage, this.values, this.#rows, {
-      row: i, from: pad, copy: true,
+      row: i, from: pad, to: this.#rows.activeLength, copy: true,
     });
     c.restoreState(k, v, this.offsetArr[i]!, k.shape[2]!);
     return c;

package/src/model/batched-row-storage.ts

@@ -1,6 +1,7 @@
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
 import { BatchedRotatingState } from "./batched-rotating-state";
+import { Dtype } from "../mlx/ffi";
 
 export interface RowStorage<T> {
   shape(value: T): readonly [number, number, number, number];
@@ -10,6 +11,7 @@ export interface RowStorage<T> {
   padLeft(value: T, tokens: number): T;
   takeRows(value: T, keep: readonly number[]): T;
   copy(value: T): T;
+  rollRows(value: T, indices: MlxArray): T;
   dispose(value: T): void;
 }
 
@@ -36,6 +38,7 @@ export const plainRowStorage: RowStorage<MlxArray> = {
     return out;
   },
   copy: (value) => ops.copyOf(value),
+  rollRows: (value, indices) => ops.takeAlongAxis(value, indices, 2),
   dispose: (value) => value.dispose(),
 };
 
@@ -77,6 +80,14 @@ export const quantizedRowStorage: RowStorage<ops.QuantizedTensor> = {
       biases: ops.copyOf(value.biases),
     };
   },
+  rollRows(value, indices) {
+    const packed = ops.takeAlongAxis(value.packed, indices, 2);
+    let scales: MlxArray | undefined;
+    try {
+      scales = ops.takeAlongAxis(value.scales, indices, 2);
+      return { packed, scales, biases: ops.takeAlongAxis(value.biases, indices, 2) };
+    } catch (error) { packed.dispose(); scales?.dispose(); throw error; }
+  },
   dispose(value) {
     value.packed.dispose();
     value.scales.dispose();
@@ -84,6 +95,19 @@ export const quantizedRowStorage: RowStorage<ops.QuantizedTensor> = {
   },
 };
 
+/** Per-row cyclic shift, shared by plain and encoded storage planes.
+ * This is mlx-lm's dynamic_roll indexing; only positions cross the FFI. */
+export function rowRollIndices(length: number, shifts: readonly number[]): MlxArray {
+  using columns = ops.arange(0, length, 1, Dtype.int32);
+  using positions = ops.reshape(columns, [1, 1, length, 1]);
+  using starts = ops.fromInt32(shifts.map(shift => ((-shift % length) + length) % length), [shifts.length, 1, 1, 1]);
+  using raw = ops.add(positions, starts);
+  using width = ops.fromInt32([length], []);
+  using wrapped = ops.sub(raw, width);
+  using over = ops.greaterEqual(raw, width);
+  return ops.where(over, wrapped, raw);
+}
+
 export function mergeStorageRows<T>(
   storage: RowStorage<T>,
   rows: readonly T[],
@@ -125,3 +149,31 @@ export function temporalStorageView<T>(
   storage.dispose(cut);
   return owned;
 }
+
+/** Borrow storage, return an owned chronological block without decoding the
+ * representation. Port of mlx-lm's BatchRotatingKVCache._update_concat
+ * (MIT): discard spare capacity, retain the full window for the first query. */
+export function appendRotatingStorage<T>(
+  storage: RowStorage<T>, previous: T | null, incoming: T, state: BatchedRotatingState,
+): T {
+  if (previous === null) return storage.slice(incoming, 0, storage.shape(incoming)[0], 0, storage.shape(incoming)[2]);
+  const active = state.activeLength;
+  const shifts = state.prefillRoll();
+  if (shifts?.some(shift => shift !== 0)) {
+    const chronological = temporalStorageView(storage, previous, state, { to: active });
+    let rolled: T | undefined, history: T | undefined;
+    try {
+      using indices = rowRollIndices(active, shifts);
+      rolled = storage.rollRows(chronological, indices);
+      history = storage.slice(rolled, 0, state.batchSize, Math.max(0, active - state.maxSize + 1), active);
+      return storage.concatTokens([history, incoming]);
+    } finally {
+      storage.dispose(chronological); if (rolled) storage.dispose(rolled); if (history) storage.dispose(history);
+    }
+  }
+  const history = temporalStorageView(storage, previous, state, {
+    from: Math.max(0, active - state.maxSize + 1), to: active,
+  });
+  try { return storage.concatTokens([history, incoming]); }
+  finally { storage.dispose(history); }
+}

package/src/model/batched-ssm.ts

@@ -0,0 +1,52 @@
+import { SsmPrefillPadding } from "./ssm-prefill-padding";
+import type { PaddedPrefillCache, PrefillPadding } from "./gemma4-base";
+import * as ops from "../mlx/ops";
+import type { MlxArray } from "../mlx/array";
+import type { BatchableCache, Cache } from "./gemma4-base";
+import { SSMCache } from "./qwen3-delta";
+
+/** Recurrent storage behind the same row-layout port as attention KV.
+ * Inputs to merge are borrowed; request membership is owned by the caller. */
+export class BatchedSSMCache extends SSMCache implements BatchableCache, PaddedPrefillCache {
+  constructor() { super(); this.offsets = []; }
+  preparePrefill(padding: PrefillPadding): void {
+    if (!this.offsets!.length) this.offsets = padding.lengths.map(() => 0);
+    const padded = padding.leftPadding?.some(n => n > 0) || padding.rightPadding?.some(n => n > 0);
+    this.prefillPadding = padded ? new SsmPrefillPadding(padding) : null;
+  }
+  finalizePrefill(): void { this.prefillPadding = null; }
+  get rowOffsets(): readonly number[] { return this.offsets!; }
+  get leftPad(): readonly number[] { return this.rowOffsets.map(() => 0); }
+  makeEmptyBatch(): BatchedSSMCache { return new BatchedSSMCache(); }
+  projectedBytes(_tokens: number): number {
+    return this.batchSize ? this.state().reduce((bytes, array) => bytes + array.nbytes, 0) / this.batchSize : 0;
+  }
+
+  mergeRows(rows: readonly Cache[]): void {
+    if (!rows.length) { this.dispose(); return; }
+    const sources = rows as readonly SSMCache[];
+    const offsets = sources.flatMap(source => source.offsets ?? [source.offset]);
+    const merge = (field: "conv" | "recurrent"): MlxArray | null => {
+      const prototype = sources.find(source => source[field])?.[field];
+      if (!prototype) return null;
+      if (sources.length === 1) return ops.copyOf(prototype);
+      const empty: MlxArray[] = [];
+      try {
+        return ops.concatAxis(sources.map(source => {
+          if (source[field]) return source[field]!;
+          const zeros = ops.zeros([source.offsets?.length ?? 1, ...prototype.shape.slice(1)], prototype.dtype);
+          empty.push(zeros); return zeros;
+        }), 0);
+      } finally { for (const array of empty) array.dispose(); }
+    };
+    const conv = merge("conv");
+    let recurrent;
+    try { recurrent = merge("recurrent"); }
+    catch (error) { conv?.dispose(); throw error; }
+    this.dispose();
+    this.conv = conv; this.recurrent = recurrent;
+    this.offsets = offsets; this.offset = Math.max(0, ...offsets);
+  }
+
+  override dispose(): void { super.dispose(); this.offset = 0; this.offsets = []; }
+}

package/src/model/batched-turboquant-kv.ts

@@ -0,0 +1,93 @@
+import { decodedKvDonorAttention } from "./decoded-kv-donor";
+import type { MlxArray } from "../mlx/array";
+import { KvTensorRows, type KvTensorRowView } from "../backends/mlx/kv-tensor-rows";
+import { disposeResources } from "../engine/resources";
+import { TurboQuantKVCache, type BatchableCache, type Cache, type Mask, type RotatedValueAttentionState, type PaddedPrefillCache, type PrefillPadding } from "./gemma4-base";
+import { TurboQuantCodec, disposeTurboQuant, type TurboQuantTensor } from "./turboquant-codec";
+
+const fields = ["kIdx", "kScales", "kZeros", "vPacked", "vScales"] as const;
+function encoded(planes: readonly MlxArray[]): TurboQuantTensor {
+  return { kIdx: planes[0]!, kScales: planes[1]!, kZeros: planes[2]!, vPacked: planes[3]!, vScales: planes[4]! };
+}
+
+/** The existing TurboQuant codec over shared tensor-row storage. Row changes
+ * copy encoded bytes without decoding, rotating or requantizing their values. */
+export class BatchedTurboQuantKVCache implements BatchableCache, PaddedPrefillCache {
+  readonly #storage = new KvTensorRows();
+  readonly #codec: TurboQuantCodec;
+  #reuseOffsets: number[] = [];
+  #headDim: number | null = null;
+  constructor(readonly kBits: number, readonly vBits: number) {
+    this.#codec = new TurboQuantCodec(kBits, vBits, process.env.MLX_BUN_TURBOQUANT_FUSED_DECODE === "1");
+  }
+  restorePrefillEnds(ends: readonly number[] | undefined): void { this.#storage.restorePrefillEnds(ends); }
+  preparePrefill(padding: PrefillPadding): void {
+    this.#storage.preparePrefill(padding);
+    if (!this.#reuseOffsets.length) this.#reuseOffsets = padding.lengths.map(() => 0);
+  }
+  finalizePrefill(): void { this.#storage.finalizePrefill(); }
+  captureDonorRows(): import("./gemma4-base").KvDonorRows {
+    if (!this.#storage.planes.length || this.#headDim === null) throw new Error("cache is empty");
+    const [keys, values] = this.#codec.decode(encoded(this.#storage.planes), this.offset, this.#headDim, false);
+    return { keys, values, ...this.#storage.captureDonorValidity() };
+  }
+  captureDonorAttention(): import("./gemma4-base").KvDonorAttention {
+    return decodedKvDonorAttention(this.captureDonorRows());
+  }
+  get rotatedValueAttention(): RotatedValueAttentionState { return this; }
+  signature(): string { return `kv:batched-turboquant:${this.kBits}:${this.vBits}`; }
+  get minimumReusableOffset(): number { return Math.max(0, ...this.#reuseOffsets); }
+  get headDim(): number | null { return this.#headDim; }
+  get rowOffsets(): readonly number[] { return this.#storage.rowOffsets; }
+  get leftPad(): readonly number[] { return this.#storage.leftPad; }
+  get batchSize(): number | null { return this.#storage.batchSize; }
+  get offset(): number { return this.#storage.offset; }
+  get ropeOffsetArr(): MlxArray | undefined { return this.#storage.ropeOffsetArr; }
+  makeEmptyBatch(): BatchedTurboQuantKVCache { return new BatchedTurboQuantKVCache(this.kBits, this.vBits); }
+  bytesPerToken(): number { return this.#storage.bytesPerToken(); }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * tokens; }
+  makeMask(tokens: number, window: number | null): Mask { return this.#storage.makeMask(tokens, window); }
+  isTrimmable(): boolean { return true; }
+  trim(count: number): void { this.#storage.trim(count); }
+  specRoundBegin(): void { this.#storage.specRoundBegin(); }
+  specRoundCommit(): void { this.#storage.specRoundCommit(); }
+  specRoundRollback(keep: number | readonly number[]): void { this.#storage.specRoundRollback(keep); }
+  #append(k: MlxArray, v: MlxArray, defer: boolean): [MlxArray, MlxArray] {
+    const packed = this.#codec.encode(k, v);
+    try {
+      this.#storage.append(fields.map(field => packed[field]));
+      this.#headDim = k.shape[3]!;
+      return this.#codec.decode(encoded(this.#storage.planes), this.offset, this.#headDim, defer);
+    } finally { disposeTurboQuant(packed); }
+  }
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.#append(k, v, false); }
+  updateAndFetchDeferredV(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.#append(k, v, true); }
+  mergeRows(rows: readonly Cache[]): void {
+    const views: KvTensorRowView[] = [], held: MlxArray[] = [];
+    const reuseOffsets: number[] = [];
+    let dim = this.#headDim;
+    try {
+      for (const row of rows) {
+        if (row instanceof BatchedTurboQuantKVCache) {
+          views.push(row.#storage); dim ??= row.headDim; reuseOffsets.push(...row.#reuseOffsets);
+        } else if (row instanceof TurboQuantKVCache) {
+          const planes = row.state(); held.push(...planes); reuseOffsets.push(row.minimumReusableOffset);
+          views.push({ planes, rowOffsets: [row.offset], leftPad: [0] }); dim ??= row.headDim;
+        } else throw new Error(`TurboQuant row layout cannot merge ${row.signature()}`);
+      }
+      this.#storage.mergeRows(views); this.#headDim = dim; this.#reuseOffsets = reuseOffsets;
+    } finally { disposeResources(held); }
+  }
+  alignRows(leftPad: readonly number[]): void { this.#storage.alignRows(leftPad); }
+  filterRows(keep: readonly number[]): void {
+    this.#storage.filterRows(keep); this.#reuseOffsets = keep.map(row => this.#reuseOffsets[row]!);
+  }
+  state(): MlxArray[] { return this.#storage.planes; }
+  extractRow(row: number): TurboQuantKVCache {
+    const result = new TurboQuantKVCache(this.kBits, this.vBits), count = this.rowOffsets[row]!;
+    if (count) result.restoreState(encoded(this.#storage.extractRow(row)), count, this.#headDim!);
+    result.minimumReusableOffset = this.#reuseOffsets[row] ?? 0;
+    return result;
+  }
+  dispose(): void { this.#storage.dispose(); this.#headDim = null; this.#reuseOffsets = []; }
+}

package/src/model/compiled-decode.ts

@@ -277,7 +277,10 @@ function decodePhase(c: AnyCache): "concat" | "ring" {
 }
 
 function disposeShared(s: SharedKv): void {
-  if (s.kind === "plain") {
+  if (s.kind === "view") {
+    s.attention.dispose();
+    s.offsetArr?.dispose();
+  } else if (s.kind === "plain") {
     s.keys.dispose();
     s.values.dispose();
   } else {
@@ -381,6 +384,9 @@ export class CompiledDecode {
   readonly #segmented: boolean;
 
   private constructor(readonly model: Gemma4Model) {
+    // A fresh resume has no prefill to evaluate lazy FP32 constants. Capture
+    // their values, not their construction graph, in every compiled closure.
+    model.materializeGraphConstants();
     const t = model.config.text;
     this.#segmented =
       t.numKvSharedLayers === 0 && !t.enableMoeBlock && model.perLayerWidth === 0;

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

@@ -0,0 +1,12 @@
+import * as ops from "../mlx/ops";
+import type { KvDonorAttention, KvDonorRows } from "./gemma4-base";
+
+/** Takes ownership of decoded donor planes. Assistant attention uses decoded
+ * values, preserving its SDPA arithmetic rather than moving inverse rotation. */
+export function decodedKvDonorAttention(view: KvDonorRows): KvDonorAttention {
+  const { keys, values, offsets, starts, ends } = view;
+  return { width: keys.shape[2]!, dtype: keys.dtype, offsets, starts, ends,
+    attend: (q, scale, mask) => ops.sdpa(q, keys, values, scale, mask.mode, mask.arr),
+    dispose() { keys.dispose(); values.dispose(); },
+  };
+}

package/src/model/delayed-quantized-kv.ts

@@ -0,0 +1,72 @@
+import { captureFullKvDonorAttention } from "../backends/mlx/full-kv-row-donor";
+import { appendFullKvRows } from "../backends/mlx/full-kv-row-append";
+import { fullRowPadding, fullRowPhysicalLength } from "./full-prefill-row";
+import { captureKvAttention } from "./kv-attention-view";
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { FullTransitioningKvRows } from "../backends/mlx/full-transitioning-kv-rows";
+import { BatchedQuantizedKVCache } from "./batched-quantized-kv";
+import { QuantizedKVCache, type Cache, type Mask, type KvAttentionState, type KvAttentionView } from "./gemma4-base";
+
+/** Affine storage keeps its native quantized-attention arithmetic while rows
+ * cross the conversion boundary independently. Once all rows convert, their
+ * packed planes use the existing batched attention implementation. */
+export class DelayedQuantizedKVCache extends FullTransitioningKvRows<BatchedQuantizedKVCache> implements KvAttentionState {
+  constructor(readonly groupSize: number, readonly bits: number, readonly start: number,
+    readonly maintain: (rows: Cache[]) => void, row?: Cache) {
+    super({ signature: `kv:delayed-quant:${bits}:${groupSize}:${start}`, maintain,
+      converted: row => row instanceof QuantizedKVCache,
+      makeLayout: () => new BatchedQuantizedKVCache(groupSize, bits) }, row);
+  }
+  get attentionState(): KvAttentionState { return this; }
+  captureDonorAttention() {
+    return this.packed?.captureDonorAttention() ?? captureFullKvDonorAttention(this.rows, this.leftPad, this.offset);
+  }
+  makeEmptyBatch(): DelayedQuantizedKVCache { return new DelayedQuantizedKVCache(this.groupSize, this.bits, this.start, this.maintain); }
+  updateAndFetch(): [MlxArray, MlxArray] { throw new Error("mixed precision rows use their attention state"); }
+  appendAndFetch(k: MlxArray, v: MlxArray): KvAttentionView {
+    this.advance();
+    if (this.packed) return captureKvAttention(this.packed, k, v);
+    if (this.rows.every(row => !row.quantizedAttention)) {
+      const [keys, values] = appendFullKvRows(this.rows, this.leftPad, this.offset + k.shape[2]!, k, v);
+      try { this.syncPositions(true); } catch (error) { keys.dispose(); values.dispose(); throw error; }
+      return { attend: (q, scale, mask) => ops.sdpa(q, keys, values, scale, mask.mode, mask.arr),
+        dispose() { keys.dispose(); values.dispose(); } };
+    }
+    const views: KvAttentionView[] = [], pads = this.rows.map((row, index) => this.leftPad[index]! - fullRowPadding(row));
+    const lengths = this.rows.map(cache => fullRowPhysicalLength(cache) + k.shape[2]!);
+    try {
+      for (const [row, cache] of this.rows.entries()) {
+        const slice = (a: MlxArray) => a.slice([row, 0, 0, 0], [row + 1, a.shape[1]!, a.shape[2]!, a.shape[3]!]);
+        using kr = slice(k), vr = slice(v);
+        views.push(captureKvAttention(cache, kr, vr));
+      }
+      this.syncPositions(true);
+    } catch (error) { for (const view of views) view.dispose(); throw error; }
+    return {
+      attend(q, scale, mask) {
+        const outputs: MlxArray[] = [];
+        try {
+          for (const [row, view] of views.entries()) {
+            using qr = q.slice([row, 0, 0, 0], [row + 1, q.shape[1]!, q.shape[2]!, q.shape[3]!]);
+            // Capture alignment with the KV tensors, never consult live rows.
+            const start = mask.arr?.shape.map(() => 0), end = mask.arr ? [...mask.arr.shape] : undefined;
+            if (start && end) {
+              if (end.length === 4 && end[0]! > 1) { start[0] = row; end[0] = row + 1; }
+              start[start.length - 1] = pads[row]!;
+              end[end.length - 1] = pads[row]! + lengths[row]!;
+            }
+            using arr = mask.arr ? mask.arr.slice(start!, end!) : null;
+            outputs.push(view.attend(qr, scale, { mode: mask.mode, arr }));
+          }
+          return outputs.length === 1 ? ops.contiguous(outputs[0]!) : ops.concatAxis(outputs, 0);
+        } finally { for (const output of outputs) output.dispose(); }
+      },
+      dispose() { for (const view of views) view.dispose(); },
+    };
+  }
+  updateAndAttend(q: MlxArray, k: MlxArray, v: MlxArray, scale: number, mask: Mask): MlxArray {
+    const view = this.appendAndFetch(k, v);
+    try { return view.attend(q, scale, mask); } finally { view.dispose(); }
+  }
+}