mlx-bun 0.0.4

package/src/mmap.ts

@@ -0,0 +1,90 @@
+// libc mmap via bun:ffi.
+//
+// Bun.mmap panics (SIGTRAP, Bun 1.3.3) on files larger than 4 GB because
+// JSC ArrayBuffers cap at 2^32 bytes — and weight shards routinely exceed
+// that. We don't need a whole-file ArrayBuffer anyway: tensors go to mlx
+// as raw pointers (base + offset), and only small ranges (headers) are
+// ever viewed from JS.
+
+import { dlopen, FFIType, toArrayBuffer } from "bun:ffi";
+import { closeSync, fstatSync, openSync } from "node:fs";
+
+const libc = dlopen("/usr/lib/libSystem.B.dylib", {
+  mmap: {
+    args: [FFIType.u64, FFIType.u64, FFIType.i32, FFIType.i32, FFIType.i32, FFIType.i64],
+    returns: FFIType.u64,
+  },
+  munmap: { args: [FFIType.u64, FFIType.u64], returns: FFIType.i32 },
+  madvise: { args: [FFIType.u64, FFIType.u64, FFIType.i32], returns: FFIType.i32 },
+}).symbols;
+
+const PROT_READ = 0x1;
+const PROT_WRITE = 0x2;
+const MAP_SHARED = 0x1;
+const MAP_PRIVATE = 0x2;
+const MAP_FAILED = 0xffffffffffffffffn;
+
+// madvise advice (macOS <sys/mman.h>). DONTNEED drops clean file-backed
+// pages (they re-fault from the file); FREE_REUSABLE is the Apple-specific
+// lazy-reclaim variant.
+export const MADV_DONTNEED = 4;
+export const MADV_FREE = 5;
+export const MADV_FREE_REUSABLE = 7;
+
+export class MmapFile {
+  readonly path: string;
+  readonly base: bigint;
+  readonly size: number;
+  #mapped = true;
+
+  private constructor(path: string, base: bigint, size: number) {
+    this.path = path;
+    this.base = base;
+    this.size = size;
+  }
+
+  /** Open read-only (MAP_SHARED) or copy-on-write (MAP_PRIVATE +
+   *  PROT_WRITE — writes go to private pages, never the file; used for
+   *  KV-cache reloads where mlx may donate buffers). */
+  static open(path: string, mode: "ro" | "cow" = "ro"): MmapFile {
+    const fd = openSync(path, "r");
+    try {
+      const size = fstatSync(fd).size;
+      if (size === 0) throw new Error(`${path}: empty file`);
+      const prot = mode === "cow" ? PROT_READ | PROT_WRITE : PROT_READ;
+      const flags = mode === "cow" ? MAP_PRIVATE : MAP_SHARED;
+      const base = libc.mmap(0n, BigInt(size), prot, flags, fd, 0n);
+      if (base === MAP_FAILED) throw new Error(`${path}: mmap failed`);
+      return new MmapFile(path, base, size);
+    } finally {
+      closeSync(fd); // mapping survives close
+    }
+  }
+
+  /** Raw pointer to byte `offset` (for FFI consumers like mlx). */
+  pointer(offset: number): number {
+    if (offset < 0 || offset > this.size) throw new RangeError(`offset ${offset} out of range`);
+    return Number(this.base + BigInt(offset));
+  }
+
+  /** JS view of a range. Only for small ranges (headers, spot checks);
+   *  the view aliases the mapping — do not retain past munmap. */
+  view(offset: number, length: number): Uint8Array {
+    if (offset + length > this.size) throw new RangeError("view out of range");
+    if (length >= 2 ** 32) throw new RangeError("JS views cap at 4 GB");
+    return new Uint8Array(toArrayBuffer(this.pointer(offset) as never, 0, length));
+  }
+
+  /** madvise a page-aligned sub-range (offset + length should be page
+   *  multiples). e.g. advise(off, len, MADV_DONTNEED) to release clean
+   *  expert pages back to the OS; they re-fault from the file on next read. */
+  advise(offset: number, length: number, advice: number): number {
+    return Number(libc.madvise(this.base + BigInt(offset), BigInt(length), advice));
+  }
+
+  unmap(): void {
+    if (!this.#mapped) return;
+    this.#mapped = false;
+    libc.munmap(this.base, BigInt(this.size));
+  }
+}

package/src/model/batched-mask.ts

@@ -0,0 +1,204 @@
+// Padding-aware attention masks for batched (B>1) serving — decode side.
+//
+// Companion to src/train/forward.ts's buildBatchedPadMask, which builds the
+// PREFILL mask ([B,1,L,L], offset 0, right-padded — the training case, already
+// parity-proven by tests/train-batch-e2e.test.ts). This module covers the
+// DECODE case for batched serving: left-padded rows sharing one growing
+// [B,H,S,D] KV buffer, where a step forwards N new query tokens at a nonzero
+// offset. (The two builders should be consolidated when batched prefill is
+// wired into the serving path — phase S1a, docs/design/parallel-slots.md.)
+//
+// Why left-padding for decode: with rows right-aligned in the KV buffer, every
+// row's next write lands in the SAME column, so one advancing offset serves
+// the whole batch. Each row's RoPE position is still per-row (handled via the
+// array-offset path, ropeOffsetArr / ops.ropeDynamic — not here); the only
+// thing the mask must encode is (a) causality, (b) each row's left padding,
+// and (c) the sliding window, if any.
+
+import { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { createCausalMask, type Cache, type Mask } from "./gemma4-base";
+
+/** Padding-aware key-validity mask for batched DECODE with left-padded rows.
+ *
+ *  The shared KV buffer has width S (== cache offset once this step's keys are
+ *  written). Row b has `leftPad[b]` padding columns on the left, then its real
+ *  keys. A step forwards N new query tokens whose absolute positions are
+ *  S-N .. S-1 (N == 1 in the steady decode state; N > 1 only for a batched
+ *  prefill chunk).
+ *
+ *  Returns [B, 1, N, S] bool, true where query i of row b may attend to key j:
+ *      allow[b,0,i,j] = (j <= S-N+i)            causal (query i at abs pos S-N+i)
+ *                       AND (j >= leftPad[b])    key past row b's left padding
+ *                       AND (S-N+i < j+window)   within the sliding window*
+ *      (*only when windowSize !== null)
+ *
+ *  Built from createCausalMask(N, S-N, windowSize) (the model's own [N, S] bool
+ *  matrix) broadcast over B, AND-ed with a per-row [B,1,1,S] key-validity mask
+ *  — exactly mirroring buildBatchedPadMask, but at a decode offset with
+ *  left-padding instead of offset-0 right-padding. Caller owns the result. */
+export function buildBatchedDecodeMask(
+  B: number,
+  N: number,
+  S: number,
+  leftPad: number[],
+  windowSize: number | null,
+): MlxArray {
+  // [N, S] causal(+window): queries at abs positions S-N..S-1, keys 0..S-1.
+  const causal2d = createCausalMask(N, S - N, windowSize); // [N, S] bool
+  const causal4d = ops.reshape(causal2d, [1, 1, N, S]);
+  causal2d.dispose();
+
+  // Per-row key validity: keyValid[b,0,0,j] = j >= leftPad[b]. Build as int32
+  // {0,1} then compare to bool, broadcasting over the query axis (i).
+  const keyValidData = new Int32Array(B * S);
+  for (let b = 0; b < B; b++) {
+    const pad = leftPad[b]!;
+    for (let j = 0; j < S; j++) keyValidData[b * S + j] = j >= pad ? 1 : 0;
+  }
+  const keyValidI = MlxArray.fromInt32(keyValidData, [B, 1, 1, S]);
+  const zero = MlxArray.fromInt32(new Int32Array([0]), []);
+  const keyValid = ops.less(zero, keyValidI); // 0 < x → bool, true where x==1
+  zero.dispose();
+  keyValidI.dispose();
+
+  const allow = ops.logicalAnd(causal4d, keyValid); // [1,1,N,S] & [B,1,1,S]
+  causal4d.dispose();
+  keyValid.dispose();
+  return allow; // [B, 1, N, S] bool
+}
+
+/** Cache wrapper for batched DECODE with left-padded rows. Delegates all KV
+ *  storage to a real, shape-[B,H,S,D] cache (KVCache et al. are already
+ *  B-generic), but overrides two things the stock single-sequence path gets
+ *  wrong for a padded batch:
+ *
+ *   - `makeMask` returns the per-row left-padding decode mask
+ *     (buildBatchedDecodeMask) instead of the stock mask, which is EMPTY at
+ *     N==1 and would let every row attend to its own (garbage) left padding.
+ *   - `ropeOffsetArr` exposes each row's REAL position (= pre-write offset −
+ *     leftPad[b]) as an int32 [B] array, so RoPE is applied per-row via the
+ *     array-offset path (mlx fast_rope_dynamic — verified to rotate each row
+ *     by its own offset). The scalar offset would mis-position every padded
+ *     row by its pad amount.
+ *
+ *  `leftPad` is fixed for the batch's lifetime (set when the left-padded decode
+ *  buffer is assembled from per-row prefills). The position array is recomputed
+ *  only when the inner offset advances (once per decode step), so the two reads
+ *  within a step (K before updateAndFetch, Q after — see Attention.forward)
+ *  return the SAME array; the attention graph holds it until eval, so it must
+ *  not be freed mid-step. */
+export class BatchedDecodeMaskCache implements Cache {
+  #ropeArr: MlxArray | null = null;
+  #ropeForOffset = -1;
+
+  constructor(
+    private readonly inner: Cache,
+    private readonly B: number,
+    private readonly leftPad: number[],
+    private readonly window: number | null,
+  ) {}
+
+  get offset(): number {
+    return this.inner.offset;
+  }
+
+  get ropeOffsetArr(): MlxArray {
+    const off = this.inner.offset;
+    if (this.#ropeArr && this.#ropeForOffset === off) return this.#ropeArr;
+    this.#ropeArr?.dispose(); // prior step's array; its graph is already evaluated
+    const data = new Int32Array(this.B);
+    for (let b = 0; b < this.B; b++) data[b] = off - this.leftPad[b]!;
+    this.#ropeArr = MlxArray.fromInt32(data, [this.B]);
+    this.#ropeForOffset = off;
+    return this.#ropeArr;
+  }
+
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    return this.inner.updateAndFetch(k, v);
+  }
+
+  makeMask(N: number, _windowSize: number | null): Mask {
+    // forwardLayers calls makeMask BEFORE this step's updateAndFetch, so the
+    // KV the attention then fetches spans inner.offset + N keys. Build a fresh
+    // mask each step and hand the model ownership — forwardLayers disposes
+    // mask.arr after building the layer graph (mlx refcounts it into the sdpa
+    // nodes, so post-build dispose is safe), so no view bookkeeping is needed.
+    const S = this.inner.offset + N;
+    return { mode: "array", arr: buildBatchedDecodeMask(this.B, N, S, this.leftPad, this.window) };
+  }
+
+  /** Free the per-step RoPE position array WITHOUT disposing the delegated KV
+   *  (inner). For callers that rebuild a fresh wrapper around a persistent inner
+   *  cache each decode step (the dynamic-B scheduler / parity harness) — the
+   *  inner outlives the wrapper, so wrapper.dispose() would be wrong. */
+  releaseRopeArr(): void {
+    this.#ropeArr?.dispose();
+    this.#ropeArr = null;
+    this.#ropeForOffset = -1;
+  }
+
+  state(): MlxArray[] {
+    return this.inner.state();
+  }
+  isTrimmable(): boolean {
+    return this.inner.isTrimmable();
+  }
+  trim(n: number): void {
+    this.inner.trim(n);
+  }
+  dispose(): void {
+    this.inner.dispose();
+    this.#ropeArr?.dispose();
+    this.#ropeArr = null;
+  }
+}
+
+// --- Dynamic-B cache ops (continuous batching) ---------------------------
+// Ports of mlx-lm's per-cache batch methods (merge / filter), the operations
+// the --batch N scheduler needs that our fixed-B verified forward doesn't have.
+// They run per layer on the raw [.,H,.,D] K/V; the scheduler applies them
+// across layers and wraps the result in a BatchedDecodeMaskCache. Pure array
+// surgery — unit-tested without a model (tests/batched-decode-mask.test.ts);
+// numerical correctness is gated separately vs an mlx-lm dynamic golden.
+
+/** mlx-lm `merge`: stack N single-row KV slices into one left-padded
+ *  [B,H,Smax,D] batch. Each row is its sequence's KV sliced to its true length
+ *  L_i; we left-pad each by Smax−L_i so all rows align on the right edge (the
+ *  layout decode + BatchedDecodeMaskCache expect). Generalizes
+ *  realBatchedGreedy's assemble to N rows. Caller owns the result; the input
+ *  row arrays are not disposed. */
+export function mergeKVRows(
+  rows: { keys: MlxArray; values: MlxArray }[],
+): { keys: MlxArray; values: MlxArray; leftPad: number[]; width: number } {
+  const lens = rows.map((r) => r.keys.shape[2]!);
+  const width = Math.max(...lens);
+  const leftPad = lens.map((l) => width - l);
+  const leftPadTo = (a: MlxArray, pad: number): MlxArray => {
+    if (pad === 0) return a.slice([0, 0, 0, 0], a.shape as number[]); // fresh full view
+    const [B, H, , D] = a.shape as [number, number, number, number];
+    const z = ops.zeros([B, H, pad, D], a.dtype);
+    const out = ops.concatAxis([z, a], 2);
+    z.dispose();
+    return out;
+  };
+  const ks = rows.map((r, i) => leftPadTo(r.keys, leftPad[i]!));
+  const vs = rows.map((r, i) => leftPadTo(r.values, leftPad[i]!));
+  const keys = ops.concatAxis(ks, 0);
+  const values = ops.concatAxis(vs, 0);
+  for (const a of [...ks, ...vs]) a.dispose();
+  return { keys, values, leftPad, width };
+}
+
+/** 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. */
+export function filterKVRows(
+  keys: MlxArray, values: MlxArray, keep: number[],
+): { 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);
+  idx.dispose();
+  return { keys: k, values: v };
+}

package/src/model/batched-rotating.ts

@@ -0,0 +1,300 @@
+// 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).
+//
+// 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/
+// 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
+// (tests/batched-rotating.test.ts) + end-to-end vs a long-context Gemma oracle.
+
+import { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import type { Cache, Mask } from "./gemma4-base";
+
+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,
+): 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 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;
+}
+
+/** Faithful port of mlx-lm BatchRotatingKVCache (decode-only — see file header). */
+export class BatchedRotatingCache implements Cache {
+  keys: MlxArray | null = null;
+  values: MlxArray | null = null;
+  /** Per-row absolute position (drives RoPE) and buffer padding. */
+  offsetArr: number[];
+  leftPad: number[];
+  #idx = 0;
+  #offset = 0; // scalar total tokens processed (mlx-lm _offset)
+  #rotated = false;
+  #ropeArr: MlxArray | null = null;
+  #ropeForOffset = -1;
+  readonly maxSize: number;
+
+  constructor(maxSize: number, leftPad: number[]) {
+    this.maxSize = maxSize;
+    this.leftPad = [...leftPad];
+    this.offsetArr = leftPad.map((l) => -l);
+  }
+
+  /** Current batch size — tracks filter() (which shrinks the per-row arrays). */
+  get #B(): number {
+    return this.leftPad.length;
+  }
+
+  /** 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.#offset;
+  }
+
+  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.#idx, this.#offset, this.#rotated,
+      ),
+    };
+  }
+
+  /** N=1 decode update — port of _update_in_place. */
+  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)");
+    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.#idx = 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.#idx = this.maxSize;
+      this.leftPad = this.leftPad.map((x) => x - trimSize);
+    }
+
+    // Rotate when the write head reaches the end of the ring.
+    if (this.#idx === this.maxSize) {
+      this.#rotated = true;
+      this.#idx = 0;
+    }
+    if (this.#rotated) this.leftPad = this.leftPad.map((x) => x - 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, this.#idx, 0], [B, H, this.#idx + S, DK]);
+    const v2 = ops.sliceUpdate(this.values!, v, [0, 0, this.#idx, 0], [B, H, this.#idx + S, vD]);
+    this.keys!.dispose();
+    this.values!.dispose();
+    this.keys = k2;
+    this.values = v2;
+    this.#offset += S;
+    this.offsetArr = this.offsetArr.map((x) => x + S);
+    this.#idx += 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]),
+    ];
+  }
+
+  /** 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");
+    const valid = Math.min(this.#offset, this.maxSize);
+    const order = (a: MlxArray): MlxArray => {
+      const [B, H, Sbuf, D] = a.shape as [number, number, number, number];
+      let t: MlxArray;
+      if (this.#idx === Sbuf) {
+        t = a.slice([0, 0, 0, 0], [B, H, Sbuf, D]);
+      } else if (this.#idx < this.#offset) {
+        const tail = a.slice([0, 0, this.#idx, 0], [B, H, Sbuf, D]);
+        const head = a.slice([0, 0, 0, 0], [B, H, this.#idx, D]);
+        t = ops.concatAxis([tail, head], 2);
+        tail.dispose();
+        head.dispose();
+      } else {
+        t = a.slice([0, 0, 0, 0], [B, H, this.#idx, D]);
+      }
+      const [, , , Dd] = t.shape as [number, number, number, number];
+      const cut = t.slice([0, 0, 0, 0], [B, H, valid, Dd]);
+      t.dispose();
+      return cut;
+    };
+    return [order(this.keys), order(this.values)];
+  }
+
+  /** Keep only `keep` rows along the batch axis (eviction). */
+  filter(keep: number[]): void {
+    if (this.keys && this.values) {
+      const idxArr = MlxArray.fromInt32(Int32Array.from(keep), [keep.length]);
+      const k = ops.takeAxis(this.keys, idxArr, 0);
+      const v = ops.takeAxis(this.values, idxArr, 0);
+      idxArr.dispose();
+      this.keys.dispose();
+      this.values.dispose();
+      this.keys = k;
+      this.values = v;
+    }
+    this.offsetArr = keep.map((i) => this.offsetArr[i]!);
+    this.leftPad = keep.map((i) => this.leftPad[i]!);
+    this.releaseRopeArr();
+  }
+
+  state(): MlxArray[] {
+    return this.keys && this.values ? [this.keys, this.values] : [];
+  }
+  isTrimmable(): boolean {
+    return this.#offset < this.maxSize;
+  }
+  trim(n: number): void {
+    const k = Math.min(this.#offset, n);
+    this.#offset -= k;
+    this.#idx -= k;
+    this.offsetArr = this.offsetArr.map((x) => x - k);
+  }
+  dispose(): void {
+    this.keys?.dispose();
+    this.values?.dispose();
+    this.keys = this.values = null;
+    this.#ropeArr?.dispose();
+    this.#ropeArr = null;
+  }
+
+  /** 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,
+  ): 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);
+    cache.offsetArr = [...offsets];
+    if (width === 0) return cache;
+
+    const padTo = (a: MlxArray, pad: number): MlxArray => {
+      if (pad === 0) return a.slice([0, 0, 0, 0], a.shape as number[]);
+      const [B, H, , D] = a.shape as [number, number, number, number];
+      const z = ops.zeros([B, H, pad, D], a.dtype);
+      const out = ops.concatAxis([z, a], 2);
+      z.dispose();
+      return out;
+    };
+    const ks = rows.map((r, i) => padTo(r.keys, leftPad[i]!));
+    const vs = rows.map((r, i) => padTo(r.values, leftPad[i]!));
+    cache.keys = ops.concatAxis(ks, 0);
+    cache.values = ops.concatAxis(vs, 0);
+    for (const a of [...ks, ...vs]) a.dispose();
+    cache.#idx = width;
+    cache.#offset = width;
+    return cache;
+  }
+}