@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/state/batched-rotating.ts

@@ -0,0 +1,430 @@
+// Batched sliding-window (rotating) KV cache for `--batch N` serving — the
+// hard half of dynamic-B that BatchedDecodeMaskCache (full-attention) doesn't
+// cover. Faithful port of mlx-lm's BatchRotatingKVCache (models/cache.py):
+// a shared ring buffer whose rotation state (_idx / _offset / rotated) is
+// SCALAR across the batch (all rows advance one token per step, so the write
+// column is the same for every row), with PER-ROW `offset` (absolute position,
+// drives RoPE) and `leftPad` (padding columns in the buffer, masked out).
+//
+// 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.
+//
+// The genuinely tricky bit is make_mask once the ring wraps: the buffer is no
+// longer in temporal order, so the causal+window+padding mask is built in
+// temporal coordinates then ROLLED to the ring's physical layout — see
+// buildBatchedRotatingMask (a column-by-column port of mlx-lm make_mask,
+// including the `roll(shift=idx+1)`). Gated model-free against mlx-lm
+// (`02d723a:tests/unit/batched-rotating.test.ts`) + end-to-end vs a long-context Gemma oracle.
+
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { RotatingKVCache } from "./rotating-kv";
+import { type AttentionCache, type AttentionRead, type BatchableCache, type Cache, type Mask, type PaddedPrefillCache, type PrefillPadding } from "../contracts/mlx/cache";
+import { AttentionMasks, sdpaRead, withLease, type MaskLease } from "./attention-read";
+import { isRotatingPlainCache } from "./capabilities";
+import { BatchedRotatingState, type RotatingPositionSnapshot } from "./batched-rotating-state";
+import {
+  appendRotatingStorage,
+  mergeStorageRows,
+  plainRowStorage,
+  temporalStorageView,
+  rowRollIndices,
+} from "./batched-row-storage";
+
+const STEP = 256;
+
+/** Port of BatchRotatingKVCache.make_mask for the DECODE case. Returns a
+ *  [B,1,N,S] bool mask (S == min(maxSize-1, offsetScalar) + N) for the step
+ *  ABOUT to be written: causal AND sliding-window AND per-row left-padding,
+ *  rolled to the ring's physical layout once wrapped.
+ *
+ *  Inputs are the PRE-write scalar ring state (`idx`=_idx, `offsetScalar`=_offset,
+ *  `rotated`) and per-row `leftPad`. The leftPad adjustments here are LOCAL to
+ *  the mask (mlx-lm decrements left_padding both in make_mask AND in update; the
+ *  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, 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. A
+  // block write first puts the ring in temporal order, so its history is the
+  // ring's active length (the whole ring once rotated), not the write head.
+  const history = block && rotated ? maxSize : idx;
+  const trimSize = history - 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.
+  const shift = isRot ? (idx >= maxSize ? 0 : idx) + 1 : 0;
+
+  const data = new Int32Array(B * N * S);
+  for (let b = 0; b < B; b++) {
+    const pad = lp[b]!;
+    for (let i = 0; i < N; i++) {
+      const lind = off + i; // query temporal position
+      for (let j = 0; j < S; j++) {
+        // temporal key position j: causal, within window, past this row's pad.
+        const allow = lind >= j && lind < j + window && j >= pad;
+        if (!allow) continue;
+        const col = ((j + shift) % S + S) % S;
+        data[(b * N + i) * S + col] = 1;
+      }
+    }
+  }
+  const intArr = MlxArray.fromInt32(data, [B, 1, N, S]);
+  const zero = MlxArray.fromInt32(new Int32Array([0]), []);
+  const mask = ops.less(zero, intArr); // 0 < x → bool
+  zero.dispose();
+  intArr.dispose();
+  return mask;
+}
+
+/** Port of mlx-lm BatchRotatingKVCache; padded singleton chunks use its
+ * concat operation with the matching block mask until finalization.
+ *
+ * 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 of any length. Each attends with the stock fused SDPA under mlx-lm's
+ * batched rotating mask (`buildBatchedRotatingMask`) for its write, built once
+ * per forward through `masks`. The deprecated `updateAndFetch` 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 BatchedRotatingCache implements AttentionCache, BatchableCache, PaddedPrefillCache {
+  keys: MlxArray | null = null;
+  values: MlxArray | null = null;
+  #rows: BatchedRotatingState;
+  #ropeArr: MlxArray | null = null;
+  #ropeForOffset = -1;
+  readonly maxSize: number;
+
+  constructor(
+    maxSize: number,
+    leftPad: number[],
+    /** The per-forward masks of the model this layout serves. */
+    readonly masks = new AttentionMasks(),
+  ) {
+    this.maxSize = maxSize;
+    this.#rows = new BatchedRotatingState(maxSize, leftPad);
+  }
+
+  appendDecode(k: MlxArray, v: MlxArray): AttentionRead {
+    return withLease(this.#mask(1, false), mask => {
+      const [keys, values] = this.#updateInPlace(k, v);
+      return sdpaRead(keys, values, mask);
+    });
+  }
+
+  appendWindow(k: MlxArray, v: MlxArray): AttentionRead {
+    return withLease(this.#mask(k.shape[2]!, true), mask => {
+      const [keys, values] = this.#updateConcat(k, v);
+      return sdpaRead(keys, values, mask);
+    });
+  }
+
+  /** The mask of an `N`-position write before it: `block` for concatenation,
+   * not for an in-place write. */
+  #mask(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),
+    }));
+  }
+
+  /** Same signature as the serial RotatingKVCache — the scheduler's merge
+   *  guards recognize a batched ring as "rotating-plain + RowBatchCache"
+   *  (packages/inference/src/execution/batch-group.ts #mergeJoiner prevRot branch); without this the next
+   *  join sees "unknown", skips the prevRot rows, and replaces the batch's
+   *  sliding-layer KV with a one-row cache (B≥3 collapse). Guard order
+   *  everywhere is isRowBatchCache FIRST, so the shared string never
+   *  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,
+    masks = new AttentionMasks()): BatchedRotatingCache {
+    const cache = new BatchedRotatingCache(position.maxSize, [...position.leftPad], masks);
+    cache.#rows.restore(position); cache.keys = keys; cache.values = values;
+
+    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(): BatchedRotatingCache { return new BatchedRotatingCache(this.maxSize, [], this.masks); }
+  /** Bytes one row holds for `tokens` more positions, capped at the window. */
+  projectedBytes(tokens: number): number {
+    if (!this.keys || !this.values || !this.batchSize) return 0;
+    return (this.keys.nbytes + this.values.nbytes) / (this.batchSize * this.keys.shape[2]!) * Math.min(tokens, this.maxSize);
+  }
+
+  /** Current batch size — tracks filter() (which shrinks the per-row arrays). */
+  get #B(): number {
+    return this.#rows.batchSize;
+  }
+
+  /** mlx-lm uses `cache.offset` (the per-row array) for the scalar interface
+   *  too; we expose the scalar total as `offset` and the per-row positions via
+   *  `ropeOffsetArr` (the model's per-row RoPE path). */
+  get offset(): number {
+    return this.#rows.totalOffset;
+  }
+
+  get ropeOffsetArr(): MlxArray {
+    if (this.#ropeArr && this.#ropeForOffset === this.offset) return this.#ropeArr;
+    this.#ropeArr?.dispose();
+    this.#ropeArr = MlxArray.fromInt32(Int32Array.from(this.offsetArr), [this.#B]);
+    this.#ropeForOffset = this.offset;
+    return this.#ropeArr;
+  }
+
+  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.offset, this.#rows.rotated, N > 1 || this.#rows.hasPendingPadding,
+      ),
+    };
+  }
+
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    if (k.shape[2] !== 1 || this.#rows.hasPendingPadding) return this.#updateConcat(k, v);
+    return this.#updateInPlace(k, v);
+  }
+
+  /** N=1 decode update — port of _update_in_place. */
+  #updateInPlace(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    const [B, H, S, D] = k.shape as [number, number, number, number];
+    const vD = v.shape[3]!;
+    const prev = this.offset;
+
+    // Grow the buffer (in STEP chunks) until it reaches maxSize.
+    if (!this.keys || (prev >= this.keys.shape[2]! && this.keys.shape[2]! < this.maxSize)) {
+      const newSize = Math.min(STEP, this.maxSize - prev);
+      const newK = ops.zeros([B, H, newSize, D], k.dtype);
+      const newV = ops.zeros([B, H, newSize, vD], v.dtype);
+      if (this.keys && this.values) {
+        const ck = ops.concatAxis([this.keys, newK], 2);
+        const cv = ops.concatAxis([this.values, newV], 2);
+        for (const a of [this.keys, this.values, newK, newV]) a.dispose();
+        this.keys = ck;
+        this.values = cv;
+      } else {
+        this.keys = newK;
+        this.values = newV;
+      }
+      this.#rows.markGrown(prev);
+    }
+
+    // Trim any overshoot past maxSize (decrements left padding persistently).
+    const trimSize = this.keys!.shape[2]! - this.maxSize;
+    if (trimSize > 0) {
+      const tk = this.keys!.slice([0, 0, trimSize, 0], this.keys!.shape as number[]);
+      const tv = this.values!.slice([0, 0, trimSize, 0], this.values!.shape as number[]);
+      this.keys!.dispose();
+      this.values!.dispose();
+      this.keys = tk;
+      this.values = tv;
+      this.#rows.trimOvershoot(trimSize);
+    }
+
+    // Rotate when the write head reaches the end of the ring.
+    const writeIndex = this.#rows.beginWrite(S);
+
+    // Write the new K/V at the ring head.
+    const [, , SK, DK] = this.keys!.shape as [number, number, number, number];
+    const k2 = ops.sliceUpdate(this.keys!, k, [0, 0, writeIndex, 0], [B, H, writeIndex + S, DK]);
+    const v2 = ops.sliceUpdate(this.values!, v, [0, 0, writeIndex, 0], [B, H, writeIndex + S, vD]);
+    this.keys!.dispose();
+    this.values!.dispose();
+    this.keys = k2;
+    this.values = v2;
+    this.#rows.commitWrite(S);
+
+    // Return the populated prefix (ring not yet full) or the whole buffer.
+    if (this.offset < this.maxSize) {
+      return [
+        this.keys.slice([0, 0, 0, 0], [B, H, this.offset, DK]),
+        this.values.slice([0, 0, 0, 0], [B, H, this.offset, vD]),
+      ];
+    }
+    return [
+      this.keys.slice([0, 0, 0, 0], [B, H, SK, DK]),
+      this.values.slice([0, 0, 0, 0], [B, H, SK, vD]),
+    ];
+  }
+
+  #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();
+    this.#ropeArr = null;
+    this.#ropeForOffset = -1;
+  }
+
+  /** Ring contents in temporal order, cut to the valid length (extract). */
+  temporalView(): [MlxArray, MlxArray] {
+    if (!this.keys || !this.values) throw new Error("cache is empty");
+    return [
+      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,
+      }),
+    ];
+  }
+
+  /** mlx-lm `BatchRotatingKVCache.extract` (models/cache.py:1417): row `i`
+   *  as a fresh SERIAL RotatingKVCache — de-rolled to temporal order (the
+   *  oracle's roll(-_idx) when rotated), left padding stripped
+   *  (`max(0, left_padding[idx])` — post-wrap pads go negative), OWNED
+   *  contiguous copies. `offset` = the row's absolute position (may exceed
+   *  the buffer: wrapped rings have evicted tokens); ring idx = the new
+   *  buffer length — temporal order ⇔ write head at the end, exactly the
+   *  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 (`02d723a:tests/unit/batched-rotating.test.ts`) and this is a pure
+   *  slice+copy. */
+  extractRow(i: number, limit?: number): RotatingKVCache {
+    if (!this.keys || !this.values) return new RotatingKVCache(this.maxSize);
+    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, to: this.#rows.activeLength, copy: true,
+    });
+    const v = temporalStorageView(plainRowStorage, this.values, this.#rows, {
+      row: i, from: pad, to: this.#rows.activeLength, copy: true,
+    });
+    c.restoreState(k, v, this.offsetArr[i]!, k.shape[2]!);
+    return c;
+  }
+
+  /** Keep only `keep` rows along the batch axis (eviction). */
+  filter(keep: number[]): void {
+    if (this.keys && this.values) {
+      const k = plainRowStorage.takeRows(this.keys, keep);
+      const v = plainRowStorage.takeRows(this.values, keep);
+      this.keys.dispose();
+      this.values.dispose();
+      this.keys = k;
+      this.values = v;
+    }
+    this.#rows.filter(keep);
+    this.releaseRopeArr();
+  }
+
+  filterRows(keep: readonly number[]): void { this.filter([...keep]); }
+
+  state(): MlxArray[] {
+    return this.keys && this.values ? [this.keys, this.values] : [];
+  }
+  isTrimmable(): boolean {
+    return this.#rows.trimmable;
+  }
+  trim(n: number): void {
+    this.#rows.trim(n);
+  }
+  dispose(): void {
+    this.keys?.dispose();
+    this.values?.dispose();
+    this.keys = this.values = null;
+    this.#ropeArr?.dispose();
+    this.#ropeArr = null;
+    this.masks.clear();
+  }
+
+  /** Join rows into this empty ring: sources are running rings, whose rows are
+   *  cut out in temporal order, or serial RotatingKVCaches (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("BatchedRotatingCache.mergeRows needs an empty batch");
+    const parts: { keys: MlxArray; values: MlxArray }[] = [], offsets: number[] = [], held: MlxArray[] = [];
+    const own = (array: MlxArray) => { held.push(array); return array; };
+    try {
+      for (const source of rows) {
+        if (source instanceof BatchedRotatingCache) {
+          const [k0, v0] = source.temporalView(); own(k0); own(v0);
+          const [, H, valid, D] = k0.shape as [number, number, number, number];
+          const vD = v0.shape[3]!;
+          for (let b = 0; b < source.batchSize; b++) {
+            const pad = Math.max(0, source.leftPad[b]!);
+            parts.push({ keys: own(k0.slice([b, 0, pad, 0], [b + 1, H, valid, D])),
+              values: own(v0.slice([b, 0, pad, 0], [b + 1, H, valid, vD])) });
+            offsets.push(source.offsetArr[b]!);
+          }
+        } else if (isRotatingPlainCache(source)) {
+          const [keys, values] = (source as RotatingKVCache).temporalView();
+          parts.push({ keys: own(keys), values: own(values) }); offsets.push(source.offset);
+        } else throw new Error(`sliding-window batch cannot merge ${source.signature()}`);
+      }
+      const merged = BatchedRotatingCache.merge(parts, offsets, this.maxSize);
+      this.#rows = merged.#rows; this.keys = merged.keys; this.values = merged.values;
+      merged.keys = merged.values = null;
+    } finally { for (const array of held) array.dispose(); }
+  }
+
+  /** Assemble a batch from per-row temporal KV slices (port of merge). Each row
+   *  is its solo cache's temporalView ([1,H,Li,D], Li ≤ maxSize, temporal
+   *  order) and `offsets[i]` its absolute position. Left-pads to the longest
+   *  row; the result is in temporal order (rotated=false, idx=offset=width). */
+  static merge(
+    rows: { keys: MlxArray; values: MlxArray }[], offsets: number[], maxSize: number,
+    masks = new AttentionMasks(),
+  ): BatchedRotatingCache {
+    const lens = rows.map((r) => r.keys.shape[2]!);
+    const width = Math.max(...lens, 0);
+    const leftPad = lens.map((l) => width - l);
+    const cache = new BatchedRotatingCache(maxSize, leftPad, masks);
+    cache.#rows.restoreMerged(width, offsets);
+    if (width === 0) return cache;
+
+    cache.keys = mergeStorageRows(plainRowStorage, rows.map((row) => row.keys), leftPad);
+    cache.values = mergeStorageRows(plainRowStorage, rows.map((row) => row.values), leftPad);
+    return cache;
+  }
+}

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

@@ -0,0 +1,179 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { BatchedRotatingState } from "./batched-rotating-state";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+
+export interface RowStorage<T> {
+  shape(value: T): readonly [number, number, number, number];
+  slice(value: T, batchFrom: number, batchTo: number, tokenFrom: number, tokenTo: number): T;
+  concatTokens(values: readonly T[]): T;
+  concatRows(values: readonly T[]): 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;
+}
+
+export const plainRowStorage: RowStorage<MlxArray> = {
+  shape: (value) => value.shape as [number, number, number, number],
+  slice(value, batchFrom, batchTo, tokenFrom, tokenTo) {
+    const [, H, , D] = value.shape as [number, number, number, number];
+    return value.slice([batchFrom, 0, tokenFrom, 0], [batchTo, H, tokenTo, D]);
+  },
+  concatTokens: (values) => ops.concatAxis([...values], 2),
+  concatRows: (values) => ops.concatAxis([...values], 0),
+  padLeft(value, tokens) {
+    if (tokens === 0) return value.slice([0, 0, 0, 0], value.shape as number[]);
+    const [B, H, , D] = value.shape as [number, number, number, number];
+    const zeros = ops.zeros([B, H, tokens, D], value.dtype);
+    const out = ops.concatAxis([zeros, value], 2);
+    zeros.dispose();
+    return out;
+  },
+  takeRows(value, keep) {
+    const indices = MlxArray.fromInt32(Int32Array.from(keep), [keep.length]);
+    const out = ops.takeAxis(value, indices, 0);
+    indices.dispose();
+    return out;
+  },
+  copy: (value) => ops.copyOf(value),
+  rollRows: (value, indices) => ops.takeAlongAxis(value, indices, 2),
+  dispose: (value) => value.dispose(),
+};
+
+export const quantizedRowStorage: RowStorage<ops.QuantizedTensor> = {
+  shape: (value) => value.packed.shape as [number, number, number, number],
+  slice(value, batchFrom, batchTo, tokenFrom, tokenTo) {
+    const cut = (array: MlxArray): MlxArray => {
+      const [, H, , D] = array.shape as [number, number, number, number];
+      return array.slice([batchFrom, 0, tokenFrom, 0], [batchTo, H, tokenTo, D]);
+    };
+    return { packed: cut(value.packed), scales: cut(value.scales), biases: cut(value.biases) };
+  },
+  concatTokens(values) {
+    return {
+      packed: ops.concatAxis(values.map((value) => value.packed), 2),
+      scales: ops.concatAxis(values.map((value) => value.scales), 2),
+      biases: ops.concatAxis(values.map((value) => value.biases), 2),
+    };
+  },
+  concatRows(values) {
+    return {
+      packed: ops.concatAxis(values.map((value) => value.packed), 0),
+      scales: ops.concatAxis(values.map((value) => value.scales), 0),
+      biases: ops.concatAxis(values.map((value) => value.biases), 0),
+    };
+  },
+  padLeft(value, tokens) {
+    const pad = (array: MlxArray): MlxArray => plainRowStorage.padLeft(array, tokens);
+    return { packed: pad(value.packed), scales: pad(value.scales), biases: pad(value.biases) };
+  },
+  takeRows(value, keep) {
+    const take = (array: MlxArray): MlxArray => plainRowStorage.takeRows(array, keep);
+    return { packed: take(value.packed), scales: take(value.scales), biases: take(value.biases) };
+  },
+  copy(value) {
+    return {
+      packed: ops.copyOf(value.packed),
+      scales: ops.copyOf(value.scales),
+      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();
+    value.biases.dispose();
+  },
+};
+
+/** 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[],
+  leftPad: readonly number[],
+): T {
+  const padded = rows.map((row, index) => storage.padLeft(row, leftPad[index]!));
+  try {
+    return storage.concatRows(padded);
+  } finally {
+    for (const row of padded) storage.dispose(row);
+  }
+}
+
+/** De-roll storage into temporal order, then cut one row or the whole batch. */
+export function temporalStorageView<T>(
+  storage: RowStorage<T>,
+  value: T,
+  state: BatchedRotatingState,
+  options: { row?: number; from?: number; to?: number; copy?: boolean } = {},
+): T {
+  const [batch, , tokens] = storage.shape(value);
+  const batchFrom = options.row ?? 0;
+  const batchTo = options.row === undefined ? batch : options.row + 1;
+  const parts = state.temporalRanges(tokens).map(([from, to]) =>
+    storage.slice(value, batchFrom, batchTo, from, to)
+  );
+  const ordered = parts.length === 1 ? parts[0]! : storage.concatTokens(parts);
+  if (parts.length > 1) for (const part of parts) storage.dispose(part);
+  const cut = storage.slice(
+    ordered,
+    0,
+    batchTo - batchFrom,
+    options.from ?? 0,
+    options.to ?? state.validLength,
+  );
+  storage.dispose(ordered);
+  if (!options.copy) return cut;
+  const owned = storage.copy(cut);
+  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/state/batched-ssm.ts

@@ -0,0 +1,54 @@
+import { SsmPrefillPadding } from "./ssm-prefill-padding";
+import type { PaddedPrefillCache, PrefillPadding } from "../contracts/mlx/cache";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { BatchableCache, Cache } from "../contracts/mlx/cache";
+import { SSMCache } from "./ssm";
+
+/** 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(); }
+  /** Recurrent state cannot be trimmed: only a row whose own advance count is the key. */
+  canPublishRow(row: number, tokens: number): boolean { return !!this.conv && this.rowOffset(row) === tokens; }
+  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 = []; }
+}