@mlx-bun/inference 0.0.0-stage → 1.0.0

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

@@ -0,0 +1,457 @@
+// Batched sliding-window QUANTIZED KV — the composition that lets a
+// per-layer kv_config naming ROTATING layers (gemma) batch. No stack ships
+// this (mlx-lm's batched path is bf16-only; optiq's mixed-KV hook is
+// serial-only), so per the composition rule the oracle is the SERIAL
+// RotatingQuantizedKVCache per row: each joiner's prefill alone converts through
+// its bf16-first lego, and this class only
+// re-arranges those bytes across the batch axis + runs the mlx-lm
+// BatchRotatingKVCache ring mechanics (batched-rotating.ts) over
+// (packed, scales, biases) triples instead of bf16 tensors.
+//
+// Same safety argument as batched-quant.ts: quantization packs along
+// HEAD_DIM, so the token axis stays token-granular in all three tensors —
+// left-pad columns are zeros, masked by buildBatchedRotatingMask (token-axis
+// only; shared verbatim with the bf16 twin), never dequantized into anything
+// a real row reads.
+//
+// Subclasses RotatingQuantizedKVCache so Attention.forward's `instanceof`
+// dispatch (packages/inference/src/models/gemma4/model.ts) routes to the quantized path with the L1 model file
+// 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. Multi-token blocks use the same chronological storage operation as
+// the bf16 layout; one-token updates keep the existing in-place ring kernel.
+//
+// Gates: packages/inference/tests/state/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
+// scheduler gate in `02d723a:tests/parity/batched-kv-quant-parity.test.ts`.
+
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import { RotatingQuantizedKVCache } from "./rotating-quantized-kv";
+import { type AttentionRead, type BatchableCache, type Cache, type Mask, type PaddedPrefillCache, type PrefillPadding } from "../contracts/mlx/cache";
+import { buildBatchedRotatingMask } from "./batched-rotating";
+import { AttentionMasks, withLease, type MaskLease } from "./attention-read";
+import { affineRead, type AffineKernel, type AffineKernels } from "./affine-attention";
+import { isRotatingQuantizedCache } from "./capabilities";
+import type { QuantRow } from "./batched-quant";
+import { BatchedRotatingState, type RotatingPositionSnapshot } from "./batched-rotating-state";
+import {
+  appendRotatingStorage,
+  rowRollIndices,
+  mergeStorageRows,
+  quantizedRowStorage,
+  temporalStorageView,
+} from "./batched-row-storage";
+
+const STEP = 256; // shared with RotatingQuantizedKVCache.STEP and the bf16 twin
+
+const disposeTriple = (t: ops.QuantizedTensor): void => {
+  t.packed.dispose();
+  t.scales.dispose();
+  t.biases.dispose();
+};
+const mapTriple = (
+  t: ops.QuantizedTensor, f: (a: MlxArray) => MlxArray,
+): ops.QuantizedTensor => ({ packed: f(t.packed), scales: f(t.scales), biases: f(t.biases) });
+
+/** Batched ring over quantized triples — mlx-lm BatchRotatingKVCache
+ *  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`.
+ *
+ *  Reads: the sliding window is `maxSize`. Decode, which runs with no prefill
+ *  padding pending (finalize first), writes one position in place at the ring
+ *  head; a window writes by concatenation, which also serves padded prefill
+ *  chunks. Each attends under mlx-lm's batched rotating mask
+ *  (`buildBatchedRotatingMask`) for its write, built once per forward through
+ *  `masks`, with the decode and masked-window kernels of `kernels`: that mask
+ *  is always an array, which `quantizedSdpa` never tiles, so the unfused
+ *  kernel both composed sets hold there is what every configuration computes
+ *  today. A row it extracts attends alone with the same `kernels`. The
+ *  deprecated `updateAndFetchQuantized` picks the write from the length
+ *  instead, so a one-row window without pending padding wrote in place there
+ *  and concatenates here: the same bits until the ring wraps, a different key
+ *  order after. */
+export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache implements BatchableCache, PaddedPrefillCache {
+  #rows: BatchedRotatingState;
+  /** Per-row RoPE positions. STABLE ACROSS A STEP — refreshed only at
+   *  releaseRopeArr() (the scheduler's post-dispatch hook), never inside
+   *  the update: a graph that captures it before the update ropes Q after
+   *  (an in-update dispose would be a use-after-dispose), and one that
+   *  re-reads it for Q after the update would see K and Q at different
+   *  positions after an in-update refresh. Widens the base's readonly to
+   *  mutable. */
+  override ropeOffsetArr: MlxArray;
+
+  private constructor(
+    maxSize: number, groupSize: number, bits: number, kernels: AffineKernels, leftPad: number[], offsets: number[],
+    masks: AttentionMasks,
+  ) {
+    super(maxSize, groupSize, bits, kernels, masks);
+    this.#rows = new BatchedRotatingState(maxSize, leftPad, offsets);
+    this.ropeOffsetArr = offsets.length ? MlxArray.fromInt32(Int32Array.from(offsets), [offsets.length]) : ops.zeros([0], Dtype.int32);
+  }
+
+  /** `kernels`: the composition's kernel set for these bits and group size.
+   * `masks`: the per-forward masks of the model this layout serves. */
+  static empty(maxSize: number, groupSize: number, bits: number, kernels: AffineKernels, leftPad: readonly number[],
+    masks = new AttentionMasks()): BatchedRotatingQuantCache {
+    return new BatchedRotatingQuantCache(maxSize, groupSize, bits, kernels, [...leftPad], leftPad.map(pad => -pad), masks);
+  }
+
+  override appendDecode(k: MlxArray, v: MlxArray): AttentionRead {
+    return withLease(this.#ringMask(1, false), mask => {
+      const [keys, values] = this.#updateInPlace(k, v);
+      return affineRead(keys, values, mask, this.kernels.decode);
+    });
+  }
+
+  override appendWindow(k: MlxArray, v: MlxArray): AttentionRead {
+    return this.appendWindowUnder(k, v, this.#ringMask(k.shape[2]!, true), this.kernels.maskedWindow);
+  }
+
+  /** The window write (concatenation) attended with `kernel` under `mask`,
+   * for the layout that owns this ring and knows its mask is plain causal
+   * (`SpeculativeRotatingAffineLayout`). Takes ownership of `mask`. */
+  appendWindowUnder(k: MlxArray, v: MlxArray, mask: MaskLease, kernel: AffineKernel): AttentionRead {
+    return withLease(mask, held => {
+      const [keys, values] = this.#updateConcat(k, v);
+      return affineRead(keys, values, held, kernel);
+    });
+  }
+
+  /** The mask of an `N`-position write before it: `block` for concatenation,
+   * not for an in-place write. The bf16 ring's family and builder
+   * (`BatchedRotatingCache`). */
+  #ringMask(N: number, block: boolean): MaskLease {
+    const rows = this.#rows;
+    const key = `${N}|${block}|${rows.leftPad.join(",")}|${rows.ringIndex}|${rows.totalOffset}|${rows.rotated}`;
+    return this.masks.lease(`batched-rotating:${this.maxSize}`, key, () => ({
+      mode: "array",
+      arr: buildBatchedRotatingMask(rows.batchSize, N, rows.leftPad, this.maxSize, this.maxSize,
+        rows.ringIndex, rows.totalOffset, rows.rotated, block),
+    }));
+  }
+
+  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,
+    kernels: AffineKernels, position: RotatingPositionSnapshot, masks = new AttentionMasks()): BatchedRotatingQuantCache {
+    const cache = new BatchedRotatingQuantCache(position.maxSize, groupSize, bits, kernels, [...position.leftPad], [...position.offsets], masks);
+    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; }
+  get rowOffsets(): readonly number[] { return this.#rows.offsets; }
+  get leftPad(): number[] { return this.#rows.leftPad; }
+  get batchSize(): number { return this.#rows.batchSize; }
+  makeEmptyBatch(): BatchedRotatingQuantCache {
+    return BatchedRotatingQuantCache.empty(this.maxSize, this.groupSize, this.bits, this.kernels, [], this.masks);
+  }
+  /** Bytes one row holds for `tokens` more positions, capped at the window. */
+  projectedBytes(tokens: number): number {
+    const perRow = this.keys && this.values && this.batchSize
+      ? [this.keys.packed, this.keys.scales, this.keys.biases, this.values.packed, this.values.scales, this.values.biases]
+        .reduce((bytes, array) => bytes + array.nbytes, 0) / (this.batchSize * this.keys.packed.shape[2]!) : 0;
+    return perRow * Math.min(tokens, this.maxSize);
+  }
+
+  get #B(): number {
+    return this.#rows.batchSize;
+  }
+
+  #seqLen2(): number {
+    return this.keys ? this.keys.packed.shape[2]! : 0;
+  }
+
+  /** Scheduler hook, called on every fwd cache AFTER the step's graph is
+   *  built — the ONE safe point to swap the rope array for the advanced
+   *  positions (the step's graph holds its own reference; no JS reads the
+   *  old handle after this). */
+  releaseRopeArr(): void {
+    this.ropeOffsetArr.dispose();
+    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 {
+    const elPerInt = 32 / this.bits;
+    return {
+      packed: ops.zeros([B, H, T, dim / elPerInt], Dtype.uint32),
+      scales: ops.zeros([B, H, T, dim / this.groupSize], dtype),
+      biases: ops.zeros([B, H, T, dim / this.groupSize], dtype),
+    };
+  }
+
+  override makeMask(N: number, windowSize: number | null): Mask {
+    const window = windowSize ?? this.maxSize;
+    return {
+      mode: "array",
+      arr: buildBatchedRotatingMask(
+        this.#B, N, this.leftPad, this.maxSize, window,
+        this.#rows.ringIndex, this.#rows.totalOffset, this.#rows.rotated, N > 1 || this.#rows.hasPendingPadding,
+      ),
+    };
+  }
+
+  override updateAndFetchQuantized(k: MlxArray, v: MlxArray): [ops.QuantizedTensor, ops.QuantizedTensor] {
+    if (k.shape[2] !== 1 || this.#rows.hasPendingPadding) return this.#updateConcat(k, v);
+    return this.#updateInPlace(k, v);
+  }
+
+  /** N=1 decode update — the bf16 twin's updateAndFetch (mlx-lm
+   *  _update_in_place) over triples, with quantize-on-write from the
+   *  serial oracle. */
+  #updateInPlace(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]!;
+    const prev = this.#rows.totalOffset;
+
+    // Grow the buffer (in STEP chunks) until it reaches maxSize.
+    if (!this.keys || (prev >= this.#seqLen2() && this.#seqLen2() < this.maxSize)) {
+      const newSize = Math.min(STEP, this.maxSize - prev);
+      const newK = this.#allocTriple(B, H, newSize, D, k.dtype);
+      const newV = this.#allocTriple(B, H, newSize, vD, v.dtype);
+      if (this.keys && this.values) {
+        const grow = (old: ops.QuantizedTensor, add: ops.QuantizedTensor): ops.QuantizedTensor => {
+          const cat = (a: MlxArray, b: MlxArray): MlxArray => {
+            const out = ops.concatAxis([a, b], 2);
+            a.dispose();
+            b.dispose();
+            return out;
+          };
+          return {
+            packed: cat(old.packed, add.packed),
+            scales: cat(old.scales, add.scales),
+            biases: cat(old.biases, add.biases),
+          };
+        };
+        this.keys = grow(this.keys, newK);
+        this.values = grow(this.values, newV);
+      } else {
+        this.keys = newK;
+        this.values = newV;
+      }
+      this.#rows.markGrown(prev);
+      this.ringIdx = this.#rows.ringIndex;
+    }
+
+    // Trim any overshoot past maxSize (decrements left padding persistently).
+    const trimSize = this.#seqLen2() - this.maxSize;
+    if (trimSize > 0) {
+      const cutFrom = (t: ops.QuantizedTensor): ops.QuantizedTensor =>
+        mapTriple(t, (a) => {
+          const [b2, h2, s2, d2] = a.shape as [number, number, number, number];
+          const s = a.slice([0, 0, trimSize, 0], [b2, h2, s2, d2]);
+          a.dispose();
+          return s;
+        });
+      this.keys = cutFrom(this.keys!);
+      this.values = cutFrom(this.values!);
+      this.#rows.trimOvershoot(trimSize);
+      this.ringIdx = this.#rows.ringIndex;
+    }
+
+    // Rotate when the write head reaches the end of the ring.
+    this.ringIdx = this.#rows.beginWrite(S);
+
+    // Quantize + write the new K/V at the ring head (six sliceUpdates).
+    const kq = ops.quantize(k, this.groupSize, this.bits);
+    const vq = ops.quantize(v, this.groupSize, this.bits);
+    const writeAt = (dst: ops.QuantizedTensor, src: ops.QuantizedTensor): ops.QuantizedTensor =>
+      ({
+        packed: this.#assign2(dst.packed, src.packed, S),
+        scales: this.#assign2(dst.scales, src.scales, S),
+        biases: this.#assign2(dst.biases, src.biases, S),
+      });
+    this.keys = writeAt(this.keys!, kq);
+    this.values = writeAt(this.values!, vq);
+    disposeTriple(kq);
+    disposeTriple(vq);
+
+    this.#rows.commitWrite(S);
+    this.offset = this.#rows.totalOffset;
+    this.ringIdx = this.#rows.ringIndex;
+
+    // Return the populated prefix (ring not yet full) or the whole buffer.
+    const upTo = this.offset < this.maxSize ? this.offset : this.#seqLen2();
+    const cut = (t: ops.QuantizedTensor): ops.QuantizedTensor =>
+      mapTriple(t, (a) => {
+        const [b2, h2, , d2] = a.shape as [number, number, number, number];
+        return a.slice([0, 0, 0, 0], [b2, h2, upTo, d2]);
+      });
+    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]);
+    dst.dispose();
+    return out;
+  }
+
+  /** Ring contents in temporal order, cut to the valid length (extract) —
+   *  overrides the base to use the BATCHED ring state. */
+  override temporalView(): [ops.QuantizedTensor, ops.QuantizedTensor] {
+    if (!this.keys || !this.values) throw new Error("cache is empty");
+    return [
+      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,
+      }),
+    ];
+  }
+
+  /** Quantized twin of BatchedRotatingCache.extractRow (mlx-lm
+   *  BatchRotatingKVCache.extract, models/cache.py:1417, over triples):
+   *  row `i` as a fresh SERIAL RotatingQuantizedKVCache — de-rolled to
+   *  temporal order, left padding stripped, OWNED contiguous copies;
+   *  offset = the row's absolute position, ringIdx = the new buffer length
+   *  (the oracle's `cache._idx = cache.keys.shape[2]`). Token-axis slicing
+   *  is byte-safe (packing along HEAD_DIM — file header); per-row byte
+   *  identity vs the serial oracle is the class invariant
+   *  (packages/inference/tests/state/batched-rotating-quant.test.ts), extraction
+   *  is a pure slice+copy. */
+  extractRow(i: number, limit = Infinity): RotatingQuantizedKVCache {
+    if (!this.keys || !this.values) return new RotatingQuantizedKVCache(this.maxSize, this.groupSize, this.bits, this.kernels);
+    const pad = Math.max(0, this.leftPad[i]!, this.#rows.activeLength - limit);
+    const c = new RotatingQuantizedKVCache(this.maxSize, this.groupSize, this.bits, this.kernels);
+    const k = temporalStorageView(quantizedRowStorage, this.keys, this.#rows, {
+      row: i, from: pad, to: this.#rows.activeLength, copy: true,
+    });
+    const v = temporalStorageView(quantizedRowStorage, this.values, this.#rows, {
+      row: i, from: pad, to: this.#rows.activeLength, copy: true,
+    });
+    c.restoreState(k, v, this.offsetArr[i]!, k.packed.shape[2]!);
+    return c;
+  }
+
+  /** Keep only `keep` rows along the batch axis (eviction), in place. */
+  filter(keep: number[]): void {
+    if (this.keys && this.values) {
+      const k = quantizedRowStorage.takeRows(this.keys, keep);
+      const v = quantizedRowStorage.takeRows(this.values, keep);
+      disposeTriple(this.keys);
+      disposeTriple(this.values);
+      this.keys = k;
+      this.values = v;
+    }
+    this.#rows.filter(keep);
+    this.releaseRopeArr(); // safe here: filter runs between steps
+  }
+
+  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
+  }
+
+  override dispose(): void {
+    super.dispose();
+    this.releaseRopeArr();
+    this.ropeOffsetArr?.dispose();
+  }
+
+  /** Join rows into this empty ring: sources are running rings, whose rows are
+   *  cut out in temporal order, or serial RotatingQuantizedKVCaches (an adopted
+   *  lone row or a joiner), in row order. Sources are borrowed. */
+  mergeRows(rows: readonly Cache[]): void {
+    if (this.#rows.batchSize) throw new Error("BatchedRotatingQuantCache.mergeRows needs an empty batch");
+    const parts: QuantRow[] = [], offsets: number[] = [], held: ops.QuantizedTensor[] = [];
+    const own = (triple: ops.QuantizedTensor) => { held.push(triple); return triple; };
+    try {
+      for (const source of rows) {
+        if (source instanceof BatchedRotatingQuantCache) {
+          const [k0, v0] = source.temporalView(); own(k0); own(v0);
+          const valid = k0.packed.shape[2]!;
+          for (let b = 0; b < source.batchSize; b++) {
+            const pad = Math.max(0, source.leftPad[b]!);
+            const cutRow = (t: ops.QuantizedTensor) => own(mapTriple(t, (a) =>
+              a.slice([b, 0, pad, 0], [b + 1, a.shape[1]!, valid, a.shape[3]!])));
+            parts.push({ keys: cutRow(k0), values: cutRow(v0) });
+            offsets.push(source.offsetArr[b]!);
+          }
+        } else if (isRotatingQuantizedCache(source)) {
+          const [keys, values] = (source as RotatingQuantizedKVCache).temporalView();
+          parts.push({ keys: own(keys), values: own(values) }); offsets.push(source.offset);
+        } else throw new Error(`quantized sliding-window batch cannot merge ${source.signature()}`);
+      }
+      const merged = BatchedRotatingQuantCache.merge(parts, offsets, this.maxSize, this.groupSize, this.bits, this.kernels);
+      this.#rows = merged.#rows; this.keys = merged.keys; this.values = merged.values;
+      this.offset = merged.offset; this.ringIdx = merged.ringIdx;
+      merged.keys = merged.values = null;
+      merged.dispose();
+      this.releaseRopeArr();
+    } finally { for (const triple of held) disposeTriple(triple); }
+  }
+
+  /** Assemble a batch from per-row temporal quantized slices (the bf16
+   *  twin's merge over triples). Each row is a solo/adopted/batched cache's
+   *  temporalView ([1,H,Li,*], Li ≤ maxSize, temporal order) and
+   *  `offsets[i]` its absolute position. Result is in temporal order
+   *  (rotated=false, ringIdx=offset=width). */
+  static merge(
+    rows: QuantRow[], offsets: number[], maxSize: number, groupSize: number, bits: number, kernels: AffineKernels,
+    masks = new AttentionMasks(),
+  ): BatchedRotatingQuantCache {
+    const lens = rows.map((r) => r.keys.packed.shape[2]!);
+    const width = Math.max(...lens, 0);
+    const leftPad = lens.map((l) => width - l);
+    const cache = new BatchedRotatingQuantCache(maxSize, groupSize, bits, kernels, leftPad, offsets, masks);
+    cache.#rows.restoreMerged(width, offsets);
+    cache.offset = cache.#rows.totalOffset;
+    cache.ringIdx = cache.#rows.ringIndex;
+    if (width === 0) return cache;
+
+    cache.keys = mergeStorageRows(quantizedRowStorage, rows.map((row) => row.keys), leftPad);
+    cache.values = mergeStorageRows(quantizedRowStorage, rows.map((row) => row.values), leftPad);
+    return cache;
+  }
+}

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

@@ -0,0 +1,149 @@
+import type { PrefillPadding } from "../contracts/mlx/cache";
+
+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;
+  offsets: number[];
+  leftPad: number[];
+  ringIndex = 0;
+  totalOffset = 0;
+  rotated = false;
+  #prefillEnds?: number[];
+
+  constructor(maxSize: number, leftPad: readonly number[], offsets?: readonly number[]) {
+    this.maxSize = maxSize;
+    this.leftPad = [...leftPad];
+    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 {
+    this.ringIndex = previousOffset;
+  }
+
+  trimOvershoot(tokens: number): void {
+    if (tokens <= 0) return;
+    this.ringIndex = this.maxSize;
+    this.leftPad = this.leftPad.map((pad) => pad - tokens);
+  }
+
+  /** Prepare an N-token write and return its physical ring column. */
+  beginWrite(tokens: number): number {
+    if (this.ringIndex === this.maxSize) {
+      this.rotated = true;
+      this.ringIndex = 0;
+    }
+    if (this.rotated) this.leftPad = this.leftPad.map((pad) => pad - tokens);
+    return this.ringIndex;
+  }
+
+  commitWrite(tokens: number): void {
+    this.totalOffset += tokens;
+    this.offsets = this.offsets.map((offset) => offset + tokens);
+    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 {
+    const amount = Math.min(this.totalOffset, tokens);
+    this.totalOffset -= amount;
+    this.ringIndex -= amount;
+    this.offsets = this.offsets.map((offset) => offset - amount);
+    return amount;
+  }
+
+  restoreMerged(width: number, offsets: readonly number[]): void {
+    this.totalOffset = width;
+    this.ringIndex = width;
+    this.offsets = [...offsets];
+    this.rotated = false;
+    this.#prefillEnds = undefined;
+  }
+
+  /** Physical ranges that expose a ring buffer in temporal order. */
+  temporalRanges(bufferLength: number): readonly [number, number][] {
+    if (this.ringIndex === bufferLength) return [[0, bufferLength]];
+    if (this.ringIndex < this.totalOffset)
+      return [[this.ringIndex, bufferLength], [0, this.ringIndex]];
+    return [[0, this.ringIndex]];
+  }
+}