@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/state/attention-read.ts

@@ -0,0 +1,215 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { AttentionCache, AttentionRead, Mask } from "../contracts/mlx/cache";
+import { bidirMask } from "../kernels/attention/masks";
+import { foldedSdpa, foldedSdpaEligible } from "../layers/folded-attention";
+import { disposeResources } from "../runtime/resources";
+
+/** A mask held by one attention read until that read is disposed. */
+export interface MaskLease {
+  readonly mask: Mask;
+  /** Drops this read's hold on the mask. Call exactly once. */
+  release(): void;
+}
+
+/** No mask: one query over every stored position. Nothing to release. */
+export const unmaskedLease: MaskLease = { mask: { mode: "", arr: null }, release() {} };
+/** The fused kernel's bottom-right-aligned causal mask. Nothing to release. */
+export const causalLease: MaskLease = { mask: { mode: "causal", arr: null }, release() {} };
+
+interface MaskSlot {
+  readonly key: string;
+  readonly source: object | null;
+  readonly mask: Mask;
+  /** This memo's hold plus one per unreleased lease. */
+  holders: number;
+}
+
+/**
+ * The attention masks of one forward, built once and shared by every layer
+ * cache composed with this instance.
+ *
+ * A cache leases its mask here inside each named read, under a family and a
+ * key. The family names how the mask is built and is fixed when the cache is
+ * constructed: the cache's mask builder and its sliding window. The key holds
+ * every input that builder reads from the cache before the append: the query
+ * count and the row offsets, padding and ring position, plus `source`, compared
+ * by identity, for a mask built from a caller's tensor. The first lease of a key
+ * builds the mask. The other layers of the same forward stand at the same
+ * positions, so they lease the same mask. The next forward's appends have moved
+ * the offsets, so its first lease builds a new mask and drops this memo's hold
+ * on the old one. Each family keeps one mask, so a model's sliding and full
+ * layers do not evict each other. Equal keys always mean equal masks, so sharing
+ * an instance can only save builds, never change a result; caches at different
+ * positions that share one and alternate within a forward rebuild on each turn.
+ *
+ * A lease keeps its mask alive until the read holding it is disposed, so a read
+ * stays valid after a later forward replaces the mask. `clear()` drops only this
+ * memo's holds; a cache calls it from `dispose()`.
+ *
+ * Neither the graph nor the scheduler sees this. A model composes one instance
+ * into the caches its `makeCache()` makes (each cache constructor takes it; a
+ * cache built without one keeps its own); empty batches and the row layouts
+ * made from a cache (`state/layout`) take their source's instance. Not a module
+ * singleton: caches share masks only through an instance they were given.
+ */
+export class AttentionMasks {
+  readonly #slots = new Map<string, MaskSlot>();
+  #builds = 0;
+
+  /** Masks built so far (diagnostic; tests count it). */
+  get builds(): number { return this.#builds; }
+
+  /** The mask `build` makes for `family` at `key` (and `source`), built on the
+   * first lease of that key. The caller owns the lease, not the mask: release
+   * the lease, never dispose `lease.mask.arr`. */
+  lease(family: string, key: string, build: () => Mask, source: object | null = null): MaskLease {
+    let slot = this.#slots.get(family);
+    if (!slot || slot.key !== key || slot.source !== source) {
+      const mask = build();
+      this.#builds++;
+      if (slot) drop(slot);
+      slot = { key, source, mask, holders: 1 };
+      this.#slots.set(family, slot);
+    }
+    const held = slot;
+    held.holders++;
+    return { mask: held.mask, release: () => drop(held) };
+  }
+
+  /** Drop this memo's holds. Leased masks stay valid until released. */
+  clear(): void {
+    for (const slot of this.#slots.values()) drop(slot);
+    this.#slots.clear();
+  }
+}
+
+function drop(slot: MaskSlot): void {
+  if (--slot.holders === 0) slot.mask.arr?.dispose();
+}
+
+/** The encoder window's mask over an `L`-position first window
+ * (`EncoderAttentionCache.appendEncoderWindow`): causal, OR'd with the overlay
+ * of positions both flagged in `bidirectional`, with no sliding window, as one
+ * [1, 1, L, L] bool array (the shape the DiffusionGemma encoder built on the
+ * host). No cache's window enters it, so every cache of a model leases it
+ * under one family and a forward builds it once. */
+export function encoderWindowMask(masks: AttentionMasks, L: number, bidirectional: MlxArray): MaskLease {
+  return masks.lease("encoder-window", `${L}`, () => {
+    const { arr } = bidirMask(L, null, bidirectional);
+    try { return { mode: "array", arr: ops.reshape(arr!, [1, 1, L, L]) }; }
+    finally { arr!.dispose(); }
+  }, bidirectional);
+}
+
+/** DiffusionGemma's sliding-layer decoder mask (its `#makeDecoderMasks`) for a
+ * block read over a ring of `maxSize` holding `offset` positions:
+ * `storedLength` stored positions in temporal order, of which the newest
+ * `maxSize - 1` that hold a valid position are visible, then all `N` block
+ * positions, as one bool [1, 1, 1, storedLength + N] row every query shares.
+ * The rings of a model lease it under one family per window. */
+export function rotatingBlockMask(masks: AttentionMasks, maxSize: number, offset: number, N: number,
+  storedLength: number): MaskLease {
+  const valid = Math.min(offset, storedLength);
+  return masks.lease(`rotating-block:${maxSize}`, `${N}:${storedLength}:${valid}`, () => {
+    const start = Math.max(0, valid - Math.max(maxSize - 1, 0));
+    const positions = ops.arange(0, storedLength, 1, Dtype.int32);
+    const startArr = ops.scalarLike(start, positions);
+    const validArr = ops.scalarLike(valid, positions);
+    const geStart = ops.greaterEqual(positions, startArr);
+    const ltValid = ops.less(positions, validArr);
+    const stored = ops.logicalAnd(geStart, ltValid);
+    positions.dispose();
+    startArr.dispose();
+    validArr.dispose();
+    geStart.dispose();
+    ltValid.dispose();
+    using ones = ops.fromInt32(new Array(N).fill(1), [N]);
+    const block = ones.astype(stored.dtype);
+    const row = ops.concatAxis([stored, block], 0);
+    stored.dispose();
+    block.dispose();
+    const reshaped = ops.reshape(row, [1, 1, 1, storedLength + N]);
+    row.dispose();
+    return { mode: "array", arr: reshaped };
+  });
+}
+
+/** A read over dense stored keys and values under a leased mask; it takes
+ * ownership of `keys`, `values` and `mask`. A bf16 storage reads its windows
+ * with the one its composition names (`KVCache.windowRead`). */
+export type DenseRead = (keys: MlxArray, values: MlxArray, mask: MaskLease) => AttentionRead;
+
+/** The read over dense stored keys and values: the stock fused SDPA kernel
+ * (`ops.sdpa`) under the leased mask. Takes ownership of `keys`, `values` and
+ * `mask`. */
+export function sdpaRead(keys: MlxArray, values: MlxArray, mask: MaskLease): AttentionRead {
+  return {
+    attend: (q, scale) => ops.sdpa(q, keys, values, scale, mask.mask.mode, mask.mask.arr),
+    dispose() { keys.dispose(); values.dispose(); mask.release(); },
+  };
+}
+
+/** `sdpaRead`, reading each KV head's keys and values once for a causal window
+ * MLX's fused SDPA hands to its unfused fallback (up to 8 queries per row, more
+ * than 32 per KV head: `foldedSdpaEligible`), where that fallback reads them
+ * once per query head: such a window folds GQA into rows (`foldedSdpa`: the
+ * fallback's bits on `applegpu_g13s`, where it was measured, bf16 precision
+ * elsewhere); every other read is the stock fused SDPA. Takes ownership of
+ * `keys`, `values` and `mask`. */
+export function foldedWindowRead(keys: MlxArray, values: MlxArray, mask: MaskLease): AttentionRead {
+  return {
+    attend: (q, scale) => mask.mask.mode === "causal" && foldedSdpaEligible(q, keys)
+      ? foldedSdpa(q, keys, values, scale) : ops.sdpa(q, keys, values, scale, mask.mask.mode, mask.mask.arr),
+    dispose() { keys.dispose(); values.dispose(); mask.release(); },
+  };
+}
+
+/** A committed span (`CommittedAttentionCache.appendCommitted`) read one
+ * position at a time: each of the span's `L` positions appends through
+ * `cache.appendDecode` and keeps that one-position read, so query `i` attends
+ * its own causal prefix with the decode arithmetic and the storage holds what
+ * `L` decodes leave. The committed read of ring storage, whose window read is
+ * not the one-position arithmetic (a ring writes a window by concatenation and,
+ * once it wraps, reorders its keys). `k` and `v` are borrowed, as for every
+ * append. */
+export function decodeByPosition(cache: AttentionCache, k: MlxArray, v: MlxArray): AttentionRead {
+  const L = k.shape[2]!;
+  const at = (a: MlxArray, i: number) => a.slice([0, 0, i, 0], [a.shape[0]!, a.shape[1]!, i + 1, a.shape[3]!]);
+  const reads: AttentionRead[] = [];
+  try {
+    for (let i = 0; i < L; i++) {
+      using key = at(k, i), value = at(v, i);
+      reads.push(cache.appendDecode(key, value));
+    }
+  } catch (error) { disposeResources(reads); throw error; }
+  return {
+    attend(q, scale) {
+      const outputs: MlxArray[] = [];
+      try {
+        for (const [i, read] of reads.entries()) {
+          using query = at(q, i);
+          outputs.push(read.attend(query, scale));
+        }
+        return outputs.length === 1 ? ops.contiguous(outputs[0]!) : ops.concatAxis(outputs, 2);
+      } finally { disposeResources(outputs); }
+    },
+    dispose() { disposeResources(reads); },
+  };
+}
+
+/** Run a read's storage step under a mask leased before it, releasing the
+ * lease if the step throws. */
+export function withLease(mask: MaskLease, read: (mask: MaskLease) => AttentionRead): AttentionRead {
+  try { return read(mask); }
+  catch (error) { mask.release(); throw error; }
+}
+
+/** Retain an independent array handle for positions used across cache writes.
+ * The contiguous view preserves the integer values and shares their storage;
+ * callers release it with the fetched attention state. */
+export function captureRopeOffsets(cache: { readonly ropeOffsetArr?: MlxArray }): MlxArray | undefined {
+  const offsets = cache.ropeOffsetArr;
+  return offsets ? ops.contiguous(offsets) : undefined;
+}

package/src/state/batched-kv.ts

@@ -0,0 +1,294 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import { materializeCopy } from "@mlx-bun/mlx/materialize";
+import { KVCache } from "./kv";
+import { isPlainKvCache } from "./capabilities";
+import { type AttentionRead, type BatchableCache, type Cache, type CommittedAttentionCache, type Mask, type PaddedPrefillCache, type PrefillPadding } from "../contracts/mlx/cache";
+import { FullPrefillPadding } from "./full-prefill-padding";
+import { AttentionMasks, causalLease, sdpaRead, unmaskedLease, withLease, type DenseRead, type MaskLease } from "./attention-read";
+
+/** Full-attention storage with independently advancing row positions.
+ * Speculative rollback changes coverage, without moving the retained KV.
+ * Subsequent appends overwrite each row's rejected suffix in place when MLX
+ * can donate the buffer. Queue membership is outside this layout.
+ *
+ * Reads: append through `updateAndFetch`, then the stock fused SDPA, a window
+ * through `windowRead` (its source storage's). While every row is unpadded and
+ * at one position, decode passes no mask and a window the fused causal mask;
+ * otherwise both pass the per-row mask (causal, past each row's left padding)
+ * for their query count, built once per forward through `masks`. */
+export class BatchedKVCache implements BatchableCache, PaddedPrefillCache, CommittedAttentionCache {
+  keys: MlxArray | null = null;
+  values: MlxArray | null = null;
+  rowOffsets: number[] = [];
+  leftPad: number[] = [];
+  #rope: MlxArray | undefined;
+  #beforeRound: number[] | undefined;
+  readonly #padding = new FullPrefillPadding();
+
+  constructor(
+    /** The per-forward masks of the model this layout serves. */
+    readonly masks = new AttentionMasks(),
+    /** The window read of the storage it was made from (`KVCache.windowRead`). */
+    readonly windowRead: DenseRead = sdpaRead,
+  ) {}
+
+  appendDecode(k: MlxArray, v: MlxArray): AttentionRead {
+    return this.#read(k, v, this.#mask(1, unmaskedLease), sdpaRead);
+  }
+
+  appendWindow(k: MlxArray, v: MlxArray): AttentionRead {
+    return this.#read(k, v, this.#mask(k.shape[2]!, causalLease), this.windowRead);
+  }
+
+  /** Fill spans run on this layout (the target layout of a full-attention
+   * layer). A bf16 committed span reads exactly as a window does. */
+  appendCommitted(k: MlxArray, v: MlxArray): AttentionRead {
+    return this.appendWindow(k, v);
+  }
+
+  #read(k: MlxArray, v: MlxArray, mask: MaskLease, read: DenseRead): AttentionRead {
+    return withLease(mask, held => {
+      const [keys, values] = this.updateAndFetch(k, v);
+      return read(keys, values, held);
+    });
+  }
+
+  /** `aligned` while every row is unpadded at one position; otherwise the
+   * per-row mask of an `N`-position append, before it. */
+  #mask(N: number, aligned: MaskLease): MaskLease {
+    if (this.#aligned()) return aligned;
+    return this.masks.lease("batched-kv", `${N}|${this.rowOffsets.join(",")}|${this.leftPad.join(",")}`,
+      () => ({ mode: "array", arr: this.#rowMask(N, null) }));
+  }
+
+  restorePrefillEnds(ends: readonly number[] | undefined): void { this.#padding.restoreEnds(ends); }
+  preparePrefill(padding: PrefillPadding): void {
+    this.#padding.prepare(this, padding); this.#positionsChanged();
+  }
+  finalizePrefill(): void {
+    const next = this.#padding.finalize(this.state(), this);
+    if (!next) return;
+    this.keys?.dispose(); this.values?.dispose();
+    this.keys = next[0] ?? null; this.values = next[1] ?? null;
+    this.#positionsChanged();
+  }
+
+  captureDonorRows(): import("../contracts/mlx/cache").KvDonorRows {
+    const B = this.rowOffsets.length, width = this.offset;
+    const keys = this.keys!.slice([0,0,0,0], [B,this.keys!.shape[1]!,width,this.keys!.shape[3]!]);
+    let values: MlxArray;
+    try { values = this.values!.slice([0,0,0,0], [B,this.values!.shape[1]!,width,this.values!.shape[3]!]); }
+    catch (error) { keys.dispose(); throw error; }
+    return { keys, values, offsets: [...this.rowOffsets], starts: [...this.leftPad], ends: this.#ends() };
+  }
+
+  signature(): string { return "kv:batched-positions"; }
+  get batchSize(): number | null { return this.rowOffsets.length || null; }
+  get offset(): number { return Math.max(0, ...this.#ends()); }
+  #ends(): number[] { return this.rowOffsets.map((offset, row) => offset + this.leftPad[row]!); }
+  get ropeOffsetArr(): MlxArray | undefined {
+    // A common unpadded position uses the existing scalar RoPE operation.
+    if (this.leftPad.every(pad => pad === 0) && this.rowOffsets.every(offset => offset === this.rowOffsets[0])) return undefined;
+    return this.#rope ??= ops.fromInt32(this.rowOffsets, [this.rowOffsets.length]);
+  }
+  #positionsChanged(): void { this.#rope?.dispose(); this.#rope = undefined; }
+  makeEmptyBatch(): BatchedKVCache { return new BatchedKVCache(this.masks, this.windowRead); }
+  bytesPerToken(): number {
+    if (!this.keys || !this.values) return 0;
+    return (this.keys.nbytes + this.values.nbytes) / (this.rowOffsets.length * this.keys.shape[2]!);
+  }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * tokens; }
+  state(): MlxArray[] { return this.keys && this.values ? [this.keys, this.values] : []; }
+  isTrimmable(): boolean { return true; }
+  trim(count: number): void {
+    this.rowOffsets = this.rowOffsets.map(offset => offset - count);
+    this.#positionsChanged();
+  }
+  specRoundBegin(): void { this.#beforeRound = [...this.rowOffsets]; }
+  specRoundCommit(): void { this.#beforeRound = undefined; }
+  specRoundRollback(keep: number | readonly number[]): void {
+    this.rowOffsets = this.#beforeRound!.map((offset, row) => offset + (typeof keep === "number" ? keep : keep[row]!));
+    this.#beforeRound = undefined;
+    this.#positionsChanged();
+  }
+
+  makeMask(N: number, window: number | null): Mask {
+    if (window === null && this.#aligned()) return { mode: N === 1 ? "" : "causal", arr: null };
+    return { mode: "array", arr: this.#rowMask(N, window) };
+  }
+
+  /** Every row unpadded and at one position: a windowless read needs no
+   * per-row mask. Layouts that keep their positions in this one read it. */
+  get aligned(): boolean { return this.#aligned(); }
+
+  #aligned(): boolean {
+    const ends = this.#ends();
+    return this.leftPad.every(pad => pad === 0) && ends.every(end => end === ends[0]);
+  }
+
+  /** Causal, past each row's left padding, within `window`: [B, 1, N, S]. */
+  #rowMask(N: number, window: number | null): MlxArray {
+    const B = this.rowOffsets.length, ends = this.#ends(), S = this.offset + N;
+    using starts = ops.fromInt32(ends, [B, 1, 1, 1]);
+    using relative = ops.arange(0, N, 1, Dtype.int32);
+    using queries = ops.reshape(relative, [1, 1, N, 1]);
+    using position = ops.add(starts, queries);
+    using indices = ops.arange(0, S, 1, Dtype.int32);
+    using keys = ops.reshape(indices, [1, 1, 1, S]);
+    using pads = ops.fromInt32(this.leftPad, [B, 1, 1, 1]);
+    using causal = ops.lessEqual(keys, position);
+    using valid = ops.greaterEqual(keys, pads);
+    const mask = ops.logicalAnd(causal, valid);
+    if (window === null) return mask;
+    using windowSize = ops.fromInt32([window], []);
+    using limit = ops.add(keys, windowSize);
+    using within = ops.less(position, limit);
+    try { return ops.logicalAnd(mask, within); }
+    finally { mask.dispose(); }
+  }
+
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    const B = this.rowOffsets.length, N = k.shape[2]!, ends = this.#ends();
+    const needed = Math.max(...ends) + N;
+    const grow = (array: MlxArray) => {
+      const [rows, heads, capacity, width] = array.shape as [number, number, number, number];
+      if (capacity >= needed) return array;
+      using padding = ops.zeros([rows, heads, Math.ceil(needed / 256) * 256 - capacity, width], array.dtype);
+      const grown = ops.concatAxis([array, padding], 2);
+      array.dispose();
+      return grown;
+    };
+    const allocate = (input: MlxArray) => ops.zeros(
+      [B, input.shape[1]!, Math.ceil(needed / 256) * 256, input.shape[3]!], input.dtype);
+    this.keys = this.keys ? grow(this.keys) : allocate(k);
+    this.values = this.values ? grow(this.values) : allocate(v);
+    const common = ends.every(end => end === ends[0]);
+    using positions = common ? null : ops.fromInt32(ends.flatMap(end =>
+      Array.from({ length: N }, (_, position) => end + position)), [B, 1, N, 1]);
+    const append = (array: MlxArray, values: MlxArray) => common
+      ? ops.sliceUpdate(array, values, [0, 0, ends[0]!, 0], [B, array.shape[1]!, ends[0]! + N, array.shape[3]!])
+      : ops.putAlongAxis(array, positions!, values, 2);
+    const keys = append(this.keys, k);
+    let values: MlxArray;
+    try { values = append(this.values, v); }
+    catch (error) { keys.dispose(); throw error; }
+    this.keys.dispose(); this.values.dispose();
+    this.keys = keys; this.values = values;
+    this.rowOffsets = this.rowOffsets.map(offset => offset + N);
+    this.#positionsChanged();
+    return [keys.slice([0, 0, 0, 0], [B, keys.shape[1]!, needed, keys.shape[3]!]),
+      values.slice([0, 0, 0, 0], [B, values.shape[1]!, needed, values.shape[3]!])];
+  }
+
+  mergeRows(rows: readonly Cache[]): void {
+    const only = rows.length === 1 ? rows[0] : undefined;
+    if (only && (only instanceof BatchedKVCache || isPlainKvCache(only))) {
+      // Preserve allocated capacity when a request starts alone. Immutable
+      // aliases transfer ownership without copying its prefill storage.
+      const keys = only.keys ? ops.copyOf(only.keys) : null;
+      let values: MlxArray | null;
+      try { values = only.values ? ops.copyOf(only.values) : null; }
+      catch (error) { keys?.dispose(); throw error; }
+      const offsets = only instanceof BatchedKVCache ? [...only.rowOffsets] : [only.offset];
+      const pads = only instanceof BatchedKVCache ? [...only.leftPad] : [0];
+      this.dispose(); this.keys = keys; this.values = values;
+      this.rowOffsets = offsets; this.leftPad = pads;
+      return;
+    }
+    const chunks: Array<{ keys: MlxArray | null; values: MlxArray | null; offsets: number[]; pads: number[]; width: number }> = [];
+    const held: MlxArray[] = [];
+    try {
+      for (const cache of rows) {
+        let keys: MlxArray | null, values: MlxArray | null, offsets: number[], pads: number[];
+        if (cache instanceof BatchedKVCache && cache.offset === 0) {
+          keys = values = null; offsets = [...cache.rowOffsets]; pads = [...cache.leftPad];
+        } else if (cache instanceof BatchedKVCache) {
+          const B = cache.rowOffsets.length, width = cache.offset;
+          keys = cache.keys!.slice([0, 0, 0, 0], [B, cache.keys!.shape[1]!, width, cache.keys!.shape[3]!]);
+          held.push(keys);
+          values = cache.values!.slice([0, 0, 0, 0], [B, cache.values!.shape[1]!, width, cache.values!.shape[3]!]);
+          held.push(values);
+          offsets = [...cache.rowOffsets]; pads = [...cache.leftPad];
+        } else if (isPlainKvCache(cache)) {
+          if (cache.offset === 0) keys = values = null;
+          else { [keys, values] = cache.temporalView(); held.push(keys, values); }
+          offsets = [cache.offset]; pads = [0];
+        } else throw new Error(`full-attention batch cannot merge ${cache.signature()}`);
+        chunks.push({ keys, values, offsets, pads, width: keys?.shape[2] ?? 0 });
+      }
+      const width = Math.max(0, ...chunks.map(chunk => chunk.width));
+      const padded = (array: MlxArray, count: number) => {
+        if (!count) return array;
+        using padding = ops.zeros([array.shape[0]!, array.shape[1]!, count, array.shape[3]!], array.dtype);
+        const result = ops.concatAxis([padding, array], 2); held.push(result); return result;
+      };
+      const merge = (field: "keys" | "values") => {
+        if (!width) return null; // Geometry arrives with the first append.
+        const prototype = chunks.find(chunk => chunk[field])![field]!;
+        return ops.concatAxis(chunks.map(chunk => {
+          const array = chunk[field];
+          if (array) return padded(array, width - chunk.width);
+          const empty = ops.zeros([chunk.offsets.length, prototype.shape[1]!, width, prototype.shape[3]!], prototype.dtype);
+          held.push(empty); return empty;
+        }), 0);
+      };
+      const keys = merge("keys");
+      let values: MlxArray | null;
+      try { values = merge("values"); }
+      catch (error) { keys?.dispose(); throw error; }
+      this.dispose();
+      this.keys = keys; this.values = values;
+      this.rowOffsets = chunks.flatMap(chunk => chunk.offsets);
+      this.leftPad = chunks.flatMap(chunk => chunk.pads.map(pad => pad + width - chunk.width));
+    } finally { for (const array of held) array.dispose(); }
+  }
+
+  filterRows(rows: readonly number[]): void {
+    if (!rows.length) { this.dispose(); return; }
+    const pads = rows.map(row => this.leftPad[row]!);
+    const sharedPadding = Math.min(...pads);
+    using indices = this.keys ? ops.fromInt32([...rows], [rows.length]) : null;
+    const select = (array: MlxArray | null): MlxArray | null => {
+      if (!array) return null;
+      const selected = ops.takeAxis(array, indices!, 0);
+      if (!sharedPadding) return selected;
+      // mlx-lm BatchKVCache.filter removes padding common to all survivors.
+      // Keeping it changes attention reduction geometry after retirement.
+      try { return selected.slice([0, 0, sharedPadding, 0], [...selected.shape]); }
+      finally { selected.dispose(); }
+    };
+    const keys = select(this.keys);
+    let values: MlxArray | null;
+    try { values = select(this.values); }
+    catch (error) { keys?.dispose(); throw error; }
+    this.keys?.dispose(); this.values?.dispose();
+    this.keys = keys; this.values = values;
+    this.rowOffsets = rows.map(row => this.rowOffsets[row]!);
+    this.leftPad = pads.map(pad => pad - sharedPadding);
+    this.#padding.filter(rows);
+    this.#positionsChanged();
+  }
+
+  extractRow(row: number): KVCache {
+    const start = this.leftPad[row]!, count = this.rowOffsets[row]!;
+    if (count === 0) return new KVCache();
+    const copy = (array: MlxArray) => {
+      using view = array.slice([row, 0, start, 0], [row + 1, array.shape[1]!, start + count, array.shape[3]!]);
+      return materializeCopy(view);
+    };
+    const keys = copy(this.keys!);
+    let values: MlxArray;
+    try { values = copy(this.values!); }
+    catch (error) { keys.dispose(); throw error; }
+    const cache = new KVCache(); cache.restoreState(keys, values, count); return cache;
+  }
+
+  dispose(): void {
+    this.keys?.dispose(); this.values?.dispose(); this.#positionsChanged();
+    this.keys = this.values = null;
+    this.rowOffsets = []; this.leftPad = []; this.#beforeRound = undefined; this.#padding.clear();
+    this.masks.clear();
+  }
+}