mlx-bun 0.3.0 → 0.4.0

package/src/model/batched-kv.ts

@@ -0,0 +1,235 @@
+import { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import * as ops from "../mlx/ops";
+import { materializeCopy } from "../mlx/materialize";
+import { KVCache, isPlainKvCache, type BatchableCache, type Cache, type Mask, type PaddedPrefillCache, type PrefillPadding } from "./gemma4-base";
+import { FullPrefillPadding } from "./full-prefill-padding";
+
+/** Full-attention storage with independently advancing row positions.
+ * Speculative rollback changes coverage, without moving the retained KV.
+ * Subsequent appends overwrite each row's rejected suffix in place when MLX
+ * can donate the buffer. Queue membership is outside this layout. */
+export class BatchedKVCache implements BatchableCache, PaddedPrefillCache {
+  keys: MlxArray | null = null;
+  values: MlxArray | null = null;
+  rowOffsets: number[] = [];
+  leftPad: number[] = [];
+  #rope: MlxArray | undefined;
+  #beforeRound: number[] | undefined;
+  readonly #padding = new FullPrefillPadding();
+
+  restorePrefillEnds(ends: readonly number[] | undefined): void { this.#padding.restoreEnds(ends); }
+  preparePrefill(padding: PrefillPadding): void {
+    this.#padding.prepare(this, padding); this.#positionsChanged();
+  }
+  finalizePrefill(): void {
+    const next = this.#padding.finalize(this.state(), this);
+    if (!next) return;
+    this.keys?.dispose(); this.values?.dispose();
+    this.keys = next[0] ?? null; this.values = next[1] ?? null;
+    this.#positionsChanged();
+  }
+
+  captureDonorRows(): import("./gemma4-base").KvDonorRows {
+    const B = this.rowOffsets.length, width = this.offset;
+    const keys = this.keys!.slice([0,0,0,0], [B,this.keys!.shape[1]!,width,this.keys!.shape[3]!]);
+    let values: MlxArray;
+    try { values = this.values!.slice([0,0,0,0], [B,this.values!.shape[1]!,width,this.values!.shape[3]!]); }
+    catch (error) { keys.dispose(); throw error; }
+    return { keys, values, offsets: [...this.rowOffsets], starts: [...this.leftPad], ends: this.#ends() };
+  }
+
+  signature(): string { return "kv:batched-positions"; }
+  get batchSize(): number | null { return this.rowOffsets.length || null; }
+  get offset(): number { return Math.max(0, ...this.#ends()); }
+  #ends(): number[] { return this.rowOffsets.map((offset, row) => offset + this.leftPad[row]!); }
+  get ropeOffsetArr(): MlxArray | undefined {
+    // A common unpadded position uses the existing scalar RoPE operation.
+    if (this.leftPad.every(pad => pad === 0) && this.rowOffsets.every(offset => offset === this.rowOffsets[0])) return undefined;
+    return this.#rope ??= ops.fromInt32(this.rowOffsets, [this.rowOffsets.length]);
+  }
+  #positionsChanged(): void { this.#rope?.dispose(); this.#rope = undefined; }
+  makeEmptyBatch(): BatchedKVCache { return new BatchedKVCache(); }
+  bytesPerToken(): number {
+    if (!this.keys || !this.values) return 0;
+    return (this.keys.nbytes + this.values.nbytes) / (this.rowOffsets.length * this.keys.shape[2]!);
+  }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * tokens; }
+  state(): MlxArray[] { return this.keys && this.values ? [this.keys, this.values] : []; }
+  isTrimmable(): boolean { return true; }
+  trim(count: number): void {
+    this.rowOffsets = this.rowOffsets.map(offset => offset - count);
+    this.#positionsChanged();
+  }
+  specRoundBegin(): void { this.#beforeRound = [...this.rowOffsets]; }
+  specRoundCommit(): void { this.#beforeRound = undefined; }
+  specRoundRollback(keep: number | readonly number[]): void {
+    this.rowOffsets = this.#beforeRound!.map((offset, row) => offset + (typeof keep === "number" ? keep : keep[row]!));
+    this.#beforeRound = undefined;
+    this.#positionsChanged();
+  }
+
+  makeMask(N: number, window: number | null): Mask {
+    const B = this.rowOffsets.length, ends = this.#ends(), S = this.offset + N;
+    if (window === null && this.leftPad.every(pad => pad === 0) && ends.every(end => end === ends[0]))
+      return { mode: N === 1 ? "" : "causal", arr: null };
+    using starts = ops.fromInt32(ends, [B, 1, 1, 1]);
+    using relative = ops.arange(0, N, 1, Dtype.int32);
+    using queries = ops.reshape(relative, [1, 1, N, 1]);
+    using position = ops.add(starts, queries);
+    using indices = ops.arange(0, S, 1, Dtype.int32);
+    using keys = ops.reshape(indices, [1, 1, 1, S]);
+    using pads = ops.fromInt32(this.leftPad, [B, 1, 1, 1]);
+    using causal = ops.lessEqual(keys, position);
+    using valid = ops.greaterEqual(keys, pads);
+    const mask = ops.logicalAnd(causal, valid);
+    if (window === null) return { mode: "array", arr: mask };
+    using windowSize = ops.fromInt32([window], []);
+    using limit = ops.add(keys, windowSize);
+    using within = ops.less(position, limit);
+    try { return { mode: "array", arr: ops.logicalAnd(mask, within) }; }
+    finally { mask.dispose(); }
+  }
+
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    const B = this.rowOffsets.length, N = k.shape[2]!, ends = this.#ends();
+    const needed = Math.max(...ends) + N;
+    const grow = (array: MlxArray) => {
+      const [rows, heads, capacity, width] = array.shape as [number, number, number, number];
+      if (capacity >= needed) return array;
+      using padding = ops.zeros([rows, heads, Math.ceil(needed / 256) * 256 - capacity, width], array.dtype);
+      const grown = ops.concatAxis([array, padding], 2);
+      array.dispose();
+      return grown;
+    };
+    const allocate = (input: MlxArray) => ops.zeros(
+      [B, input.shape[1]!, Math.ceil(needed / 256) * 256, input.shape[3]!], input.dtype);
+    this.keys = this.keys ? grow(this.keys) : allocate(k);
+    this.values = this.values ? grow(this.values) : allocate(v);
+    const common = ends.every(end => end === ends[0]);
+    using positions = common ? null : ops.fromInt32(ends.flatMap(end =>
+      Array.from({ length: N }, (_, position) => end + position)), [B, 1, N, 1]);
+    const append = (array: MlxArray, values: MlxArray) => common
+      ? ops.sliceUpdate(array, values, [0, 0, ends[0]!, 0], [B, array.shape[1]!, ends[0]! + N, array.shape[3]!])
+      : ops.putAlongAxis(array, positions!, values, 2);
+    const keys = append(this.keys, k);
+    let values: MlxArray;
+    try { values = append(this.values, v); }
+    catch (error) { keys.dispose(); throw error; }
+    this.keys.dispose(); this.values.dispose();
+    this.keys = keys; this.values = values;
+    this.rowOffsets = this.rowOffsets.map(offset => offset + N);
+    this.#positionsChanged();
+    return [keys.slice([0, 0, 0, 0], [B, keys.shape[1]!, needed, keys.shape[3]!]),
+      values.slice([0, 0, 0, 0], [B, values.shape[1]!, needed, values.shape[3]!])];
+  }
+
+  mergeRows(rows: readonly Cache[]): void {
+    const only = rows.length === 1 ? rows[0] : undefined;
+    if (only && (only instanceof BatchedKVCache || isPlainKvCache(only))) {
+      // Preserve allocated capacity when a request starts alone. Immutable
+      // aliases transfer ownership without copying its prefill storage.
+      const keys = only.keys ? ops.copyOf(only.keys) : null;
+      let values: MlxArray | null;
+      try { values = only.values ? ops.copyOf(only.values) : null; }
+      catch (error) { keys?.dispose(); throw error; }
+      const offsets = only instanceof BatchedKVCache ? [...only.rowOffsets] : [only.offset];
+      const pads = only instanceof BatchedKVCache ? [...only.leftPad] : [0];
+      this.dispose(); this.keys = keys; this.values = values;
+      this.rowOffsets = offsets; this.leftPad = pads;
+      return;
+    }
+    const chunks: Array<{ keys: MlxArray | null; values: MlxArray | null; offsets: number[]; pads: number[]; width: number }> = [];
+    const held: MlxArray[] = [];
+    try {
+      for (const cache of rows) {
+        let keys: MlxArray | null, values: MlxArray | null, offsets: number[], pads: number[];
+        if (cache instanceof BatchedKVCache && cache.offset === 0) {
+          keys = values = null; offsets = [...cache.rowOffsets]; pads = [...cache.leftPad];
+        } else if (cache instanceof BatchedKVCache) {
+          const B = cache.rowOffsets.length, width = cache.offset;
+          keys = cache.keys!.slice([0, 0, 0, 0], [B, cache.keys!.shape[1]!, width, cache.keys!.shape[3]!]);
+          held.push(keys);
+          values = cache.values!.slice([0, 0, 0, 0], [B, cache.values!.shape[1]!, width, cache.values!.shape[3]!]);
+          held.push(values);
+          offsets = [...cache.rowOffsets]; pads = [...cache.leftPad];
+        } else if (isPlainKvCache(cache)) {
+          if (cache.offset === 0) keys = values = null;
+          else { [keys, values] = cache.temporalView(); held.push(keys, values); }
+          offsets = [cache.offset]; pads = [0];
+        } else throw new Error(`full-attention batch cannot merge ${cache.signature()}`);
+        chunks.push({ keys, values, offsets, pads, width: keys?.shape[2] ?? 0 });
+      }
+      const width = Math.max(0, ...chunks.map(chunk => chunk.width));
+      const padded = (array: MlxArray, count: number) => {
+        if (!count) return array;
+        using padding = ops.zeros([array.shape[0]!, array.shape[1]!, count, array.shape[3]!], array.dtype);
+        const result = ops.concatAxis([padding, array], 2); held.push(result); return result;
+      };
+      const merge = (field: "keys" | "values") => {
+        if (!width) return null; // Geometry arrives with the first append.
+        const prototype = chunks.find(chunk => chunk[field])![field]!;
+        return ops.concatAxis(chunks.map(chunk => {
+          const array = chunk[field];
+          if (array) return padded(array, width - chunk.width);
+          const empty = ops.zeros([chunk.offsets.length, prototype.shape[1]!, width, prototype.shape[3]!], prototype.dtype);
+          held.push(empty); return empty;
+        }), 0);
+      };
+      const keys = merge("keys");
+      let values: MlxArray | null;
+      try { values = merge("values"); }
+      catch (error) { keys?.dispose(); throw error; }
+      this.dispose();
+      this.keys = keys; this.values = values;
+      this.rowOffsets = chunks.flatMap(chunk => chunk.offsets);
+      this.leftPad = chunks.flatMap(chunk => chunk.pads.map(pad => pad + width - chunk.width));
+    } finally { for (const array of held) array.dispose(); }
+  }
+
+  filterRows(rows: readonly number[]): void {
+    if (!rows.length) { this.dispose(); return; }
+    const pads = rows.map(row => this.leftPad[row]!);
+    const sharedPadding = Math.min(...pads);
+    using indices = this.keys ? ops.fromInt32([...rows], [rows.length]) : null;
+    const select = (array: MlxArray | null): MlxArray | null => {
+      if (!array) return null;
+      const selected = ops.takeAxis(array, indices!, 0);
+      if (!sharedPadding) return selected;
+      // mlx-lm BatchKVCache.filter removes padding common to all survivors.
+      // Keeping it changes attention reduction geometry after retirement.
+      try { return selected.slice([0, 0, sharedPadding, 0], [...selected.shape]); }
+      finally { selected.dispose(); }
+    };
+    const keys = select(this.keys);
+    let values: MlxArray | null;
+    try { values = select(this.values); }
+    catch (error) { keys?.dispose(); throw error; }
+    this.keys?.dispose(); this.values?.dispose();
+    this.keys = keys; this.values = values;
+    this.rowOffsets = rows.map(row => this.rowOffsets[row]!);
+    this.leftPad = pads.map(pad => pad - sharedPadding);
+    this.#padding.filter(rows);
+    this.#positionsChanged();
+  }
+
+  extractRow(row: number): KVCache {
+    const start = this.leftPad[row]!, count = this.rowOffsets[row]!;
+    if (count === 0) return new KVCache();
+    const copy = (array: MlxArray) => {
+      using view = array.slice([row, 0, start, 0], [row + 1, array.shape[1]!, start + count, array.shape[3]!]);
+      return materializeCopy(view);
+    };
+    const keys = copy(this.keys!);
+    let values: MlxArray;
+    try { values = copy(this.values!); }
+    catch (error) { keys.dispose(); throw error; }
+    const cache = new KVCache(); cache.restoreState(keys, values, count); return cache;
+  }
+
+  dispose(): void {
+    this.keys?.dispose(); this.values?.dispose(); this.#positionsChanged();
+    this.keys = this.values = null;
+    this.rowOffsets = []; this.leftPad = []; this.#beforeRound = undefined; this.#padding.clear();
+  }
+}

package/src/model/batched-mask.ts

@@ -259,13 +259,21 @@ export function extractKVRow(cache: KVCache, leftPad: number, i: number): KVCach
 
 /** mlx-lm `filter`: keep only `keep` (sorted row indices) along the batch
  *  axis — eviction of finished sequences from a batched [B,H,S,D] buffer.
- *  Caller owns the result; inputs are not disposed. */
+ *  `trimLeft` removes padding shared by all survivors (the caller updates
+ *  offsets and per-row padding). Caller owns the result; inputs stay owned. */
 export function filterKVRows(
-  keys: MlxArray, values: MlxArray, keep: number[],
+  keys: MlxArray, values: MlxArray, keep: number[], trimLeft = 0,
 ): { keys: MlxArray; values: MlxArray } {
   const idx = MlxArray.fromInt32(Int32Array.from(keep), [keep.length]);
-  const k = ops.takeAxis(keys, idx, 0);
-  const v = ops.takeAxis(values, idx, 0);
+  const take = (a: MlxArray) => {
+    const rows = ops.takeAxis(a, idx, 0);
+    if (trimLeft === 0) return rows;
+    const trimmed = rows.slice([0, 0, trimLeft, 0], rows.shape);
+    rows.dispose();
+    return trimmed;
+  };
+  const k = take(keys);
+  const v = take(values);
   idx.dispose();
   return { keys: k, values: v };
 }

package/src/model/batched-quant.ts

@@ -134,11 +134,17 @@ export function extractQuantRow(
 /** Quantized twin of filterKVRows: keep `keep` rows along the batch axis.
  *  Caller owns the result; inputs are not disposed. */
 export function filterQuantRows(
-  keys: ops.QuantizedTensor, values: ops.QuantizedTensor, keep: number[],
+  keys: ops.QuantizedTensor, values: ops.QuantizedTensor, keep: number[], trimLeft = 0,
 ): { keys: ops.QuantizedTensor; values: ops.QuantizedTensor } {
   const idx = MlxArray.fromInt32(Int32Array.from(keep), [keep.length]);
   const take = (t: ops.QuantizedTensor): ops.QuantizedTensor =>
-    tripleMap(t, (a) => ops.takeAxis(a, idx, 0));
+    tripleMap(t, (a) => {
+      const rows = ops.takeAxis(a, idx, 0);
+      if (trimLeft === 0) return rows;
+      const trimmed = rows.slice([0, 0, trimLeft, 0], rows.shape);
+      rows.dispose();
+      return trimmed;
+    });
   const k = take(keys);
   const v = take(values);
   idx.dispose();

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

@@ -0,0 +1,144 @@
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { cleanupFailure, disposeResources } from "../engine/resources";
+import { BatchedKVCache } from "./batched-kv";
+import { KVCache, QuantizedKVCache, isQuantizedKvCache,
+  type BatchableCache, type Cache, type Mask, type QuantizedAttentionState, type PaddedPrefillCache, type PrefillPadding } from "./gemma4-base";
+
+import { KvTensorRows } from "../backends/mlx/kv-tensor-rows";
+import { quantizedDonorAttention } from "./kv-attention-view";
+
+const FIELDS = ["packed", "scales", "biases"] as const;
+
+/** Affine KV uses the same row-position storage operations for its packed
+ * values, scales and biases. No dequantization or requantization is needed
+ * for merge, rollback, retirement or extraction. Inputs share this codec. */
+export class BatchedQuantizedKVCache implements BatchableCache, QuantizedAttentionState, PaddedPrefillCache {
+  #planes = FIELDS.map(() => new BatchedKVCache());
+  #minimum: number[] = [];
+  get minimumReusableOffset(): number { return Math.max(0, ...this.#minimum); }
+  constructor(readonly groupSize: number, readonly bits: number) {}
+  restorePrefillEnds(ends: readonly number[] | undefined): void { for (const plane of this.#planes) plane.restorePrefillEnds(ends); }
+  preparePrefill(padding: PrefillPadding): void {
+    for (const plane of this.#planes) plane.preparePrefill(padding);
+    if (!this.#minimum.length) this.#minimum = padding.lengths.map(() => 0);
+  }
+  finalizePrefill(): void { for (const plane of this.#planes) plane.finalizePrefill(); }
+  get quantizedAttention(): QuantizedAttentionState { return this; }
+  captureDonorAttention(): import("./gemma4-base").KvDonorAttention {
+    const views: import("./gemma4-base").KvDonorRows[] = [];
+    try {
+      for (const plane of this.#planes) views.push(plane.captureDonorRows());
+      const field = (name: "keys" | "values") => ({
+        packed: views[0]![name], scales: views[1]![name], biases: views[2]![name],
+      });
+      const { offsets, starts, ends } = views[0]!;
+      return quantizedDonorAttention(field("keys"), field("values"), { offsets, starts, ends }, this.groupSize, this.bits);
+    } catch (error) { disposeResources(views.flatMap(view => [view.keys, view.values])); throw error; }
+  }
+  signature(): string { return `kv:batched-quant:${this.bits}:${this.groupSize}`; }
+  get rowOffsets(): readonly number[] { return this.#planes[0]!.rowOffsets; }
+  get leftPad(): readonly number[] { return this.#planes[0]!.leftPad; }
+  get batchSize(): number | null { return this.#planes[0]!.batchSize; }
+  get offset(): number { return this.#planes[0]!.offset; }
+  get ropeOffsetArr(): MlxArray | undefined { return this.#planes[0]!.ropeOffsetArr; }
+  makeEmptyBatch(): BatchedQuantizedKVCache { return new BatchedQuantizedKVCache(this.groupSize, this.bits); }
+  bytesPerToken(): number { return this.#planes.reduce((n, plane) => n + plane.bytesPerToken(), 0); }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * tokens; }
+  makeMask(tokens: number, window: number | null): Mask { return this.#planes[0]!.makeMask(tokens, window); }
+  isTrimmable(): boolean { return true; }
+  trim(count: number): void { for (const plane of this.#planes) plane.trim(count); }
+  specRoundBegin(): void { for (const plane of this.#planes) plane.specRoundBegin(); }
+  specRoundCommit(): void { for (const plane of this.#planes) plane.specRoundCommit(); }
+  specRoundRollback(keep: number | readonly number[]): void {
+    for (const plane of this.#planes) plane.specRoundRollback(keep);
+  }
+  updateAndFetch(): [MlxArray, MlxArray] { throw new Error("quantized attention uses updateAndFetchQuantized"); }
+
+  updateAndFetchQuantized(k: MlxArray, v: MlxArray): [ops.QuantizedTensor, ops.QuantizedTensor] {
+    using inputs = new DisposableStack();
+    const key = ops.quantize(k, this.groupSize, this.bits);
+    for (const field of FIELDS) inputs.use(key[field]);
+    const value = ops.quantize(v, this.groupSize, this.bits);
+    for (const field of FIELDS) inputs.use(value[field]);
+    const held: MlxArray[] = [], keys: Partial<ops.QuantizedTensor> = {}, values: Partial<ops.QuantizedTensor> = {};
+    try {
+      for (const [index, field] of FIELDS.entries()) {
+        const [kk, vv] = this.#planes[index]!.updateAndFetch(key[field], value[field]);
+        held.push(kk, vv); keys[field] = kk; values[field] = vv;
+      }
+      return [keys as ops.QuantizedTensor, values as ops.QuantizedTensor];
+    } catch (error) { return cleanupFailure(error, () => disposeResources(held)); }
+  }
+
+  /** Borrow prepared request state. All planes stage replacements before
+   * releasing the current group, preserving donor ownership on failure. */
+  mergeRows(rows: readonly Cache[]): void {
+    const next: BatchedKVCache[] = [];
+    try {
+      for (const [index, field] of FIELDS.entries()) {
+        const held: KVCache[] = [], inputs: Cache[] = [];
+        try {
+          for (const row of rows) {
+            if (row instanceof BatchedQuantizedKVCache) inputs.push(row.#planes[index]!);
+            else if (isQuantizedKvCache(row)) {
+              const view = new KVCache(); held.push(view);
+              if (row.offset > 0) {
+                const key = ops.copyOf(row.keys![field]);
+                let value: MlxArray;
+                try { value = ops.copyOf(row.values![field]); }
+                catch (error) { key.dispose(); throw error; }
+                view.restoreState(key, value, row.offset);
+              }
+              inputs.push(view);
+            } else throw new Error(`quantized row layout cannot merge ${row.signature()}`);
+          }
+          const plane = new BatchedKVCache(); next.push(plane); plane.mergeRows(inputs);
+        } finally { disposeResources(held); }
+      }
+    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    const previous = this.#planes; this.#planes = next; disposeResources(previous);
+    this.#minimum = rows.flatMap(row => row instanceof BatchedQuantizedKVCache ? row.#minimum : [row.minimumReusableOffset ?? 0]);
+  }
+
+  alignRows(leftPad: readonly number[]): void {
+    if (leftPad.every((pad, row) => pad === this.leftPad[row])) return;
+    const aligned = new KvTensorRows();
+    try {
+    aligned.mergeRows([{ planes: this.#planes.flatMap(plane => plane.keys ? [plane.keys, plane.values!] : []),
+      rowOffsets: this.rowOffsets, leftPad: this.leftPad }]);
+    aligned.alignRows(leftPad);
+    for (const plane of this.#planes) {
+      plane.keys?.dispose(); plane.values?.dispose();
+      plane.keys = aligned.planes.shift() ?? null; plane.values = aligned.planes.shift() ?? null;
+      plane.leftPad = [...leftPad];
+    }
+    } finally { aligned.dispose(); }
+  }
+  filterRows(keep: readonly number[]): void {
+    for (const plane of this.#planes) plane.filterRows(keep);
+    this.#minimum = keep.map(row => this.#minimum[row]!);
+  }
+  state(): MlxArray[] {
+    if (!this.#planes[0]!.keys) return [];
+    return [...this.#planes.map(plane => plane.keys!), ...this.#planes.map(plane => plane.values!)];
+  }
+
+  /** Return the existing persisted representation, owning compact arrays. */
+  extractRow(row: number): QuantizedKVCache {
+    const offset = this.rowOffsets[row]!, result = new QuantizedKVCache(this.groupSize, this.bits);
+    if (this.#minimum[row]) result.minimumReusableOffset = this.#minimum[row];
+    if (offset === 0) return result;
+    const held: KVCache[] = [];
+    try {
+      for (const plane of this.#planes) held.push(plane.extractRow(row));
+      result.restoreState(
+        { packed: held[0]!.keys!, scales: held[1]!.keys!, biases: held[2]!.keys! },
+        { packed: held[0]!.values!, scales: held[1]!.values!, biases: held[2]!.values! }, offset);
+      for (const cache of held) cache.keys = cache.values = null;
+      return result;
+    } finally { disposeResources(held); }
+  }
+
+  dispose(): void { disposeResources(this.#planes); this.#minimum = []; }
+}

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

@@ -20,7 +20,8 @@
 // untouched. PERSISTENT (like BatchedRotatingCache, unlike the ephemeral
 // per-step BatchedQuantDecodeMaskCache): the scheduler passes it through
 // every step; it self-tracks per-row offset/leftPad and the scalar ring
-// state. N=1 decode updates only (rows solo-prefill serially, then merge).
+// state. Multi-token blocks use the same chronological storage operation as
+// the bf16 layout; one-token updates keep the existing in-place ring kernel.
 //
 // Gates: tests/unit/batched-rotating-quant.test.ts (model-free byte-identity vs
 // the serial oracle per row, through ring wrap, B=1 and B=2) + the gemma
@@ -29,11 +30,13 @@
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
 import { Dtype } from "../mlx/ffi";
-import { RotatingQuantizedKVCache, type Mask } from "./gemma4-base";
+import { RotatingQuantizedKVCache, type Mask, type PaddedPrefillCache, type PrefillPadding } from "./gemma4-base";
 import { buildBatchedRotatingMask } from "./batched-rotating";
 import type { QuantRow } from "./batched-quant";
-import { BatchedRotatingState } from "./batched-rotating-state";
+import { BatchedRotatingState, type RotatingPositionSnapshot } from "./batched-rotating-state";
 import {
+  appendRotatingStorage,
+  rowRollIndices,
   mergeStorageRows,
   quantizedRowStorage,
   temporalStorageView,
@@ -54,7 +57,7 @@ const mapTriple = (
  *  mechanics (see batched-rotating.ts) with RotatingQuantizedKVCache
  *  storage. Scalar ring state (`ringIdx`/`offset` reuse the base fields);
  *  per-row `offsetArr` (absolute positions → RoPE) and `leftPad`. */
-export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
+export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache implements PaddedPrefillCache {
   readonly #rows: BatchedRotatingState;
   /** Per-row RoPE positions. STABLE ACROSS A STEP — refreshed only at
    *  releaseRopeArr() (the scheduler's post-dispatch hook), never inside
@@ -72,7 +75,34 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
   ) {
     super(maxSize, groupSize, bits);
     this.#rows = new BatchedRotatingState(maxSize, leftPad, offsets);
-    this.ropeOffsetArr = MlxArray.fromInt32(Int32Array.from(offsets), [offsets.length]);
+    this.ropeOffsetArr = offsets.length ? MlxArray.fromInt32(Int32Array.from(offsets), [offsets.length]) : ops.zeros([0], Dtype.int32);
+  }
+
+  static empty(maxSize: number, groupSize: number, bits: number, leftPad: readonly number[]): BatchedRotatingQuantCache {
+    return new BatchedRotatingQuantCache(maxSize, groupSize, bits, [...leftPad], leftPad.map(pad => -pad));
+  }
+
+  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.packed.shape[2]!, shifts);
+      const keys = quantizedRowStorage.rollRows(this.keys, indices);
+      let values: ops.QuantizedTensor;
+      try { values = quantizedRowStorage.rollRows(this.values, indices); }
+      catch (error) { disposeTriple(keys); throw error; }
+      disposeTriple(this.keys); disposeTriple(this.values); this.keys = keys; this.values = values;
+    }
+    this.#rows.finalizePrefill(); this.releaseRopeArr();
+  }
+  /** Adopt owned tensors without changing their physical ring columns. */
+  static adoptPhysical(keys: ops.QuantizedTensor | null, values: ops.QuantizedTensor | null, groupSize: number, bits: number, position: RotatingPositionSnapshot): BatchedRotatingQuantCache {
+    const cache = new BatchedRotatingQuantCache(position.maxSize, groupSize, bits, [...position.leftPad], [...position.offsets]);
+    cache.#rows.restore(position); cache.keys = keys; cache.values = values;
+    cache.offset = position.totalOffset; cache.ringIdx = position.ringIndex;
+    return cache;
   }
 
   get offsetArr(): number[] { return this.#rows.offsets; }
@@ -93,7 +123,7 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
    *  old handle after this). */
   releaseRopeArr(): void {
     this.ropeOffsetArr.dispose();
-    this.ropeOffsetArr = MlxArray.fromInt32(Int32Array.from(this.offsetArr), [this.#B]);
+    this.ropeOffsetArr = this.#B ? MlxArray.fromInt32(Int32Array.from(this.offsetArr), [this.#B]) : ops.zeros([0], Dtype.int32);
   }
 
   #allocTriple(B: number, H: number, T: number, dim: number, dtype: Dtype): ops.QuantizedTensor {
@@ -111,7 +141,7 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
       mode: "array",
       arr: buildBatchedRotatingMask(
         this.#B, N, this.leftPad, this.maxSize, window,
-        this.#rows.ringIndex, this.#rows.totalOffset, this.#rows.rotated,
+        this.#rows.ringIndex, this.#rows.totalOffset, this.#rows.rotated, N > 1 || this.#rows.hasPendingPadding,
       ),
     };
   }
@@ -122,8 +152,7 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
   override updateAndFetchQuantized(k: MlxArray, v: MlxArray): [ops.QuantizedTensor, ops.QuantizedTensor] {
     const [B, H, S, D] = k.shape as [number, number, number, number];
     const vD = v.shape[3]!;
-    if (S !== 1)
-      throw new Error("BatchedRotatingQuantCache supports N=1 decode updates only (solo-prefill then merge)");
+    if (S !== 1 || this.#rows.hasPendingPadding) return this.#updateConcat(k, v);
     const prev = this.#rows.totalOffset;
 
     // Grow the buffer (in STEP chunks) until it reaches maxSize.
@@ -202,6 +231,29 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
     return [cut(this.keys!), cut(this.values!)];
   }
 
+  #updateConcat(k: MlxArray, v: MlxArray): [ops.QuantizedTensor, ops.QuantizedTensor] {
+    const kq = ops.quantize(k, this.groupSize, this.bits);
+    const vq = ops.quantize(v, this.groupSize, this.bits);
+    let keys: ops.QuantizedTensor | undefined;
+    let values: ops.QuantizedTensor;
+    try {
+      keys = appendRotatingStorage(quantizedRowStorage, this.keys, kq, this.#rows);
+      values = appendRotatingStorage(quantizedRowStorage, this.values, vq, this.#rows);
+    } catch (error) { if (keys) disposeTriple(keys); throw error; }
+    finally { disposeTriple(kq); disposeTriple(vq); }
+    const history = this.keys ? this.#rows.activeLength : 0;
+    if (this.keys) disposeTriple(this.keys);
+    if (this.values) disposeTriple(this.values);
+    this.keys = keys; this.values = values;
+    this.#rows.commitConcat(k.shape[2]!, history);
+    this.offset = this.#rows.totalOffset;
+    this.ringIdx = this.#rows.ringIndex;
+    return [
+      quantizedRowStorage.slice(keys, 0, this.#B, 0, this.ringIdx),
+      quantizedRowStorage.slice(values, 0, this.#B, 0, this.ringIdx),
+    ];
+  }
+
   #assign2(dst: MlxArray, src: MlxArray, S: number): MlxArray {
     const [B, H, , D] = dst.shape as [number, number, number, number];
     const out = ops.sliceUpdate(dst, src, [0, 0, this.ringIdx, 0], [B, H, this.ringIdx + S, D]);
@@ -214,8 +266,12 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
   override temporalView(): [ops.QuantizedTensor, ops.QuantizedTensor] {
     if (!this.keys || !this.values) throw new Error("cache is empty");
     return [
-      temporalStorageView(quantizedRowStorage, this.keys, this.#rows),
-      temporalStorageView(quantizedRowStorage, this.values, this.#rows),
+      temporalStorageView(quantizedRowStorage, this.keys, this.#rows, {
+        from: Math.max(0, this.#rows.activeLength - this.maxSize), to: this.#rows.activeLength,
+      }),
+      temporalStorageView(quantizedRowStorage, this.values, this.#rows, {
+        from: Math.max(0, this.#rows.activeLength - this.maxSize), to: this.#rows.activeLength,
+      }),
     ];
   }
 
@@ -228,15 +284,15 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
    *  is byte-safe (packing along HEAD_DIM — file header); per-row byte
    *  identity vs the serial oracle is the class invariant
    *  (tests/batched-rotating-quant), extraction is a pure slice+copy. */
-  extractRow(i: number): RotatingQuantizedKVCache | null {
+  extractRow(i: number, limit = Infinity): RotatingQuantizedKVCache | null {
     if (!this.keys || !this.values) return null;
-    const pad = Math.max(0, this.leftPad[i]!);
+    const pad = Math.max(0, this.leftPad[i]!, this.#rows.activeLength - limit);
     const c = new RotatingQuantizedKVCache(this.maxSize, this.groupSize, this.bits);
     const k = temporalStorageView(quantizedRowStorage, this.keys, this.#rows, {
-      row: i, from: pad, copy: true,
+      row: i, from: pad, to: this.#rows.activeLength, copy: true,
     });
     const v = temporalStorageView(quantizedRowStorage, 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.packed.shape[2]!);
     return c;
@@ -258,6 +314,11 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
 
   filterRows(keep: readonly number[]): void { this.filter([...keep]); }
 
+  override trim(count: number): void {
+    this.#rows.trim(count); this.offset = this.#rows.totalOffset; this.ringIdx = this.#rows.ringIndex;
+    this.releaseRopeArr();
+  }
+
   override isTrimmable(): boolean {
     return false; // batched rows never re-enter the prompt cache directly
   }