mlx-bun 0.3.0 → 0.5.0

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
   }

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 = []; }
+}