mlx-bun 0.3.0 → 0.5.0

package/src/mlx/ops.ts

@@ -472,14 +472,12 @@ export function where(cond: MlxArray, x: MlxArray, y: MlxArray, s: S = gpuStream
   );
 }
 
-// arange is built host-side and uploaded. The original mlx_arange binding
-// broke once the calling path got JIT-optimized; a standalone repro
-// (lab/repro/bun-ffi-f64/) traced the root cause to a Bun bug: after DFG
-// tier-up, typed-array reads following a bun:ffi call return stale values
-// (the JIT eliminates the load across the native call). Not f64 marshaling
-// — args reach C intact. See PLAN.md Phase 4 findings. Reuse large ranges
-// across attention layers and sampling, but bound their retained storage:
-// causal masks request a new range on every round once context exceeds 64K.
+// The old native arange workaround concerned stale typed-array out-pointer
+// reads after Bun JIT tier-up, not f64 argument marshaling. outArray uses
+// native pointer reads now. Large nonnegative unit-step int32 ranges avoid
+// host construction; small/general integer ranges retain JS conversion,
+// including its wrapping behavior outside int32. Cache storage stays bounded
+// as causal masks request new lengths on successive long-context rounds.
 const arangeCache = new Map<string, { array: MlxArray; bytes: number }>();
 const ARANGE_CACHE_MIN = 65536;
 const ARANGE_CACHE_BYTES = 8 * 1024 * 1024;
@@ -498,8 +496,10 @@ export function arange(start: number, stop: number, step: number, dtype: Dtype,
   }
 
   let arr: MlxArray;
-  if (dtype === Dtype.int32) {
-    // host-built int32 leaf (the Bun-f64-FFI-safe path — see note above).
+  const nativePositions = n === 0 ||
+    (n >= ARANGE_CACHE_MIN && step === 1 && start >= 0 && stop <= 2147483647);
+  if (dtype === Dtype.int32 && !nativePositions) {
+    // Preserve JS conversion for small/general integer ranges.
     const data = new Int32Array(n);
     for (let i = 0; i < n; i++) data[i] = start + i * step;
     arr = MlxArray.fromInt32(data, [n]);
@@ -1019,3 +1019,90 @@ export function leftShift(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray
 export function rightShift(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
   return new MlxArray(outArray("right_shift", (o) => C.mlx_right_shift(o, a.handle, b.handle, s)));
 }
+
+// --- Whisper additions (fft / strided framing / pad / reductions) ----------
+
+/** mx.fft.rfft(a, n, axis) with numpy "backward" normalization. Real input
+ *  → complex64 output of length n/2+1 along `axis`. */
+export function rfft(a: MlxArray, n: number, axis = -1, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("fft_rfft", (o) => C.mlx_fft_rfft(o, a.handle, n, axis, 0, s)));
+}
+
+/** mx.as_strided(a, shape, strides, offset) — element strides. */
+export function asStrided(
+  a: MlxArray, shape: number[], strides: number[], offset = 0, s: S = gpuStream,
+): MlxArray {
+  const shp = new Int32Array(shape);
+  const str = new BigInt64Array(strides.map((v) => BigInt(v)));
+  return new MlxArray(
+    outArray("as_strided", (o) =>
+      C.mlx_as_strided(o, a.handle, ptr(shp), BigInt(shape.length), ptr(str), BigInt(strides.length), BigInt(offset), s),
+    ),
+  );
+}
+
+/** mx.pad(a, [(lo, hi), …]) with constant zero fill (all axes given). */
+export function pad(a: MlxArray, widths: [number, number][], s: S = gpuStream): MlxArray {
+  const axes = new Int32Array(widths.map((_, i) => i));
+  const lo = new Int32Array(widths.map((w) => w[0]));
+  const hi = new Int32Array(widths.map((w) => w[1]));
+  const zero = scalarLike(0, a);
+  try {
+    return new MlxArray(
+      outArray("pad", (o) =>
+        C.mlx_pad(
+          o, a.handle, ptr(axes), BigInt(widths.length), ptr(lo), BigInt(widths.length),
+          ptr(hi), BigInt(widths.length), zero.handle, ptr(cstr("constant")), s,
+        ),
+      ),
+    );
+  } finally {
+    zero.dispose();
+  }
+}
+
+export function log10(a: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("log10", (o) => C.mlx_log10(o, a.handle, s)));
+}
+
+/** Reduce-max over all elements → 0-d array (keepdims=false). */
+export function maxAll(a: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("max", (o) => C.mlx_max(o, a.handle, false, s)));
+}
+
+export function broadcastTo(a: MlxArray, shape: number[], s: S = gpuStream): MlxArray {
+  const buf = new Int32Array(shape);
+  return new MlxArray(
+    outArray("broadcast_to", (o) => C.mlx_broadcast_to(o, a.handle, ptr(buf), BigInt(shape.length), s)),
+  );
+}
+
+/** mx.topk(a, k, axis) — the k largest values along axis (unsorted). */
+export function topkAxis(a: MlxArray, k: number, axis: number, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("topk_axis", (o) => C.mlx_topk_axis(o, a.handle, k, axis, s)));
+}
+
+export function varAxis(a: MlxArray, axis: number, keepdims: boolean, ddof = 0, s: S = gpuStream): MlxArray {
+  return new MlxArray(
+    outArray("var_axis", (o) => C.mlx_var_axis(o, a.handle, axis, keepdims, ddof, s)),
+  );
+}
+
+export function stackAxis(arrays: MlxArray[], axis: number, s: S = gpuStream): MlxArray {
+  const handles = new BigUint64Array(arrays.map((a) => a.handle));
+  const vec = C.mlx_vector_array_new_data(ptr(handles), BigInt(arrays.length));
+  try {
+    return new MlxArray(outArray("stack_axis", (o) => C.mlx_stack_axis(o, vec, axis, s)));
+  } finally {
+    C.mlx_vector_array_free(vec);
+  }
+}
+
+export function minimum(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("minimum", (o) => C.mlx_minimum(o, a.handle, b.handle, s)));
+}
+
+/** Flat mx.take(a, indices) — gathers from the flattened array. */
+export function takeFlat(a: MlxArray, indices: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("take", (o) => C.mlx_take(o, a.handle, indices.handle, s)));
+}

package/src/model/artifact-identity.ts

@@ -0,0 +1,73 @@
+import { createHash, randomUUID } from "node:crypto";
+import { mkdir, readFile, realpath, rename, rm, stat, writeFile } from "node:fs/promises";
+import { homedir } from "node:os";
+import { join } from "node:path";
+
+/** Identity calculation owns bytes; this interface owns memo persistence. */
+export interface ArtifactIdentityStore {
+  get(key: string): Promise<string | undefined>;
+  put(key: string, digest: string): void;
+}
+
+export class FileArtifactIdentityStore implements ArtifactIdentityStore {
+  readonly #memory = new Map<string, string>();
+  readonly #pending = new Set<Promise<void>>();
+  constructor(readonly directory: string) {}
+
+  async get(key: string): Promise<string | undefined> {
+    const memory = this.#memory.get(key);
+    if (memory) return memory;
+    try {
+      const digest = await readFile(join(this.directory, key), "utf8");
+      if (!/^[0-9a-f]{64}$/.test(digest)) return undefined;
+      this.#memory.set(key, digest);
+      return digest;
+    } catch { return undefined; }
+  }
+
+  put(key: string, digest: string): void {
+    this.#memory.set(key, digest);
+    const temporary = join(this.directory, `${key}.${process.pid}.${randomUUID()}.tmp`);
+    const write = async () => {
+      try {
+        await mkdir(this.directory, { recursive: true });
+        await writeFile(temporary, digest);
+        await rename(temporary, join(this.directory, key));
+      } catch { await rm(temporary, { force: true }).catch(() => {}); }
+    };
+    const pending = write().finally(() => this.#pending.delete(pending));
+    this.#pending.add(pending);
+  }
+
+  async flush(): Promise<void> { await Promise.all(this.#pending); }
+}
+
+const defaultStore = new FileArtifactIdentityStore(join(homedir(), ".cache", "mlx-bun", "artifact-identities"));
+export interface ArtifactIdentityFile { name: string; path: string; }
+
+async function revision(files: readonly ArtifactIdentityFile[]): Promise<string> {
+  return JSON.stringify(await Promise.all(files.map(async file => {
+    const path = await realpath(file.path), info = await stat(path, { bigint: true });
+    return [file.name, path, info.dev, info.ino, info.size, info.mtimeNs, info.ctimeNs].map(String);
+  })));
+}
+
+/** Preserve the existing seed/name/byte digest, reusing it while file revisions
+ * are unchanged. ctime catches same-size edits even when mtime is restored.
+ * A memo miss computes the full identity; persistence never delays its caller. */
+export async function artifactIdentity(seed: string, input: readonly ArtifactIdentityFile[],
+  store: ArtifactIdentityStore = defaultStore): Promise<string> {
+  const files = [...input].sort((a, b) => a.name < b.name ? -1 : a.name > b.name ? 1 : 0);
+  const before = await revision(files);
+  const key = createHash("sha256").update(JSON.stringify([seed, before])).digest("hex");
+  const cached = await store.get(key);
+  if (cached) return cached;
+  const digest = createHash("sha256").update(seed);
+  for (const file of files) {
+    digest.update(file.name);
+    for await (const chunk of Bun.file(file.path).stream()) digest.update(chunk);
+  }
+  const value = digest.digest("hex");
+  if (before === await revision(files)) store.put(key, value);
+  return value;
+}

package/src/model/batched-kv.ts

@@ -0,0 +1,235 @@
+import { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import * as ops from "../mlx/ops";
+import { materializeCopy } from "../mlx/materialize";
+import { KVCache, isPlainKvCache, type BatchableCache, type Cache, type Mask, type PaddedPrefillCache, type PrefillPadding } from "./gemma4-base";
+import { FullPrefillPadding } from "./full-prefill-padding";
+
+/** 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. */
+export class BatchedKVCache implements BatchableCache, PaddedPrefillCache {
+  keys: MlxArray | null = null;
+  values: MlxArray | null = null;
+  rowOffsets: number[] = [];
+  leftPad: number[] = [];
+  #rope: MlxArray | undefined;
+  #beforeRound: number[] | undefined;
+  readonly #padding = new FullPrefillPadding();
+
+  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("./gemma4-base").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(); }
+  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 {
+    const B = this.rowOffsets.length, ends = this.#ends(), S = this.offset + N;
+    if (window === null && this.leftPad.every(pad => pad === 0) && ends.every(end => end === ends[0]))
+      return { mode: N === 1 ? "" : "causal", arr: null };
+    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 { mode: "array", arr: mask };
+    using windowSize = ops.fromInt32([window], []);
+    using limit = ops.add(keys, windowSize);
+    using within = ops.less(position, limit);
+    try { return { mode: "array", arr: 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();
+  }
+}

package/src/model/batched-mask.ts

@@ -259,13 +259,21 @@ export function extractKVRow(cache: KVCache, leftPad: number, i: number): KVCach
 
 /** 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. */
+ *  `trimLeft` removes padding shared by all survivors (the caller updates
+ *  offsets and per-row padding). Caller owns the result; inputs stay owned. */
 export function filterKVRows(
-  keys: MlxArray, values: MlxArray, keep: number[],
+  keys: MlxArray, values: MlxArray, keep: number[], trimLeft = 0,
 ): { 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);
+  const take = (a: MlxArray) => {
+    const rows = ops.takeAxis(a, idx, 0);
+    if (trimLeft === 0) return rows;
+    const trimmed = rows.slice([0, 0, trimLeft, 0], rows.shape);
+    rows.dispose();
+    return trimmed;
+  };
+  const k = take(keys);
+  const v = take(values);
   idx.dispose();
   return { keys: k, values: v };
 }

package/src/model/batched-quant.ts

@@ -134,11 +134,17 @@ export function extractQuantRow(
 /** Quantized twin of filterKVRows: keep `keep` rows along the batch axis.
  *  Caller owns the result; inputs are not disposed. */
 export function filterQuantRows(
-  keys: ops.QuantizedTensor, values: ops.QuantizedTensor, keep: number[],
+  keys: ops.QuantizedTensor, values: ops.QuantizedTensor, keep: number[], trimLeft = 0,
 ): { keys: ops.QuantizedTensor; values: ops.QuantizedTensor } {
   const idx = MlxArray.fromInt32(Int32Array.from(keep), [keep.length]);
   const take = (t: ops.QuantizedTensor): ops.QuantizedTensor =>
-    tripleMap(t, (a) => ops.takeAxis(a, idx, 0));
+    tripleMap(t, (a) => {
+      const rows = ops.takeAxis(a, idx, 0);
+      if (trimLeft === 0) return rows;
+      const trimmed = rows.slice([0, 0, trimLeft, 0], rows.shape);
+      rows.dispose();
+      return trimmed;
+    });
   const k = take(keys);
   const v = take(values);
   idx.dispose();

package/src/model/batched-quantized-kv.ts

@@ -0,0 +1,144 @@
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { cleanupFailure, disposeResources } from "../engine/resources";
+import { BatchedKVCache } from "./batched-kv";
+import { KVCache, QuantizedKVCache, isQuantizedKvCache,
+  type BatchableCache, type Cache, type Mask, type QuantizedAttentionState, type PaddedPrefillCache, type PrefillPadding } from "./gemma4-base";
+
+import { KvTensorRows } from "../backends/mlx/kv-tensor-rows";
+import { quantizedDonorAttention } from "./kv-attention-view";
+
+const FIELDS = ["packed", "scales", "biases"] as const;
+
+/** Affine KV uses the same row-position storage operations for its packed
+ * values, scales and biases. No dequantization or requantization is needed
+ * for merge, rollback, retirement or extraction. Inputs share this codec. */
+export class BatchedQuantizedKVCache implements BatchableCache, QuantizedAttentionState, PaddedPrefillCache {
+  #planes = FIELDS.map(() => new BatchedKVCache());
+  #minimum: number[] = [];
+  get minimumReusableOffset(): number { return Math.max(0, ...this.#minimum); }
+  constructor(readonly groupSize: number, readonly bits: number) {}
+  restorePrefillEnds(ends: readonly number[] | undefined): void { for (const plane of this.#planes) plane.restorePrefillEnds(ends); }
+  preparePrefill(padding: PrefillPadding): void {
+    for (const plane of this.#planes) plane.preparePrefill(padding);
+    if (!this.#minimum.length) this.#minimum = padding.lengths.map(() => 0);
+  }
+  finalizePrefill(): void { for (const plane of this.#planes) plane.finalizePrefill(); }
+  get quantizedAttention(): QuantizedAttentionState { return this; }
+  captureDonorAttention(): import("./gemma4-base").KvDonorAttention {
+    const views: import("./gemma4-base").KvDonorRows[] = [];
+    try {
+      for (const plane of this.#planes) views.push(plane.captureDonorRows());
+      const field = (name: "keys" | "values") => ({
+        packed: views[0]![name], scales: views[1]![name], biases: views[2]![name],
+      });
+      const { offsets, starts, ends } = views[0]!;
+      return quantizedDonorAttention(field("keys"), field("values"), { offsets, starts, ends }, this.groupSize, this.bits);
+    } catch (error) { disposeResources(views.flatMap(view => [view.keys, view.values])); throw error; }
+  }
+  signature(): string { return `kv:batched-quant:${this.bits}:${this.groupSize}`; }
+  get rowOffsets(): readonly number[] { return this.#planes[0]!.rowOffsets; }
+  get leftPad(): readonly number[] { return this.#planes[0]!.leftPad; }
+  get batchSize(): number | null { return this.#planes[0]!.batchSize; }
+  get offset(): number { return this.#planes[0]!.offset; }
+  get ropeOffsetArr(): MlxArray | undefined { return this.#planes[0]!.ropeOffsetArr; }
+  makeEmptyBatch(): BatchedQuantizedKVCache { return new BatchedQuantizedKVCache(this.groupSize, this.bits); }
+  bytesPerToken(): number { return this.#planes.reduce((n, plane) => n + plane.bytesPerToken(), 0); }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * tokens; }
+  makeMask(tokens: number, window: number | null): Mask { return this.#planes[0]!.makeMask(tokens, window); }
+  isTrimmable(): boolean { return true; }
+  trim(count: number): void { for (const plane of this.#planes) plane.trim(count); }
+  specRoundBegin(): void { for (const plane of this.#planes) plane.specRoundBegin(); }
+  specRoundCommit(): void { for (const plane of this.#planes) plane.specRoundCommit(); }
+  specRoundRollback(keep: number | readonly number[]): void {
+    for (const plane of this.#planes) plane.specRoundRollback(keep);
+  }
+  updateAndFetch(): [MlxArray, MlxArray] { throw new Error("quantized attention uses updateAndFetchQuantized"); }
+
+  updateAndFetchQuantized(k: MlxArray, v: MlxArray): [ops.QuantizedTensor, ops.QuantizedTensor] {
+    using inputs = new DisposableStack();
+    const key = ops.quantize(k, this.groupSize, this.bits);
+    for (const field of FIELDS) inputs.use(key[field]);
+    const value = ops.quantize(v, this.groupSize, this.bits);
+    for (const field of FIELDS) inputs.use(value[field]);
+    const held: MlxArray[] = [], keys: Partial<ops.QuantizedTensor> = {}, values: Partial<ops.QuantizedTensor> = {};
+    try {
+      for (const [index, field] of FIELDS.entries()) {
+        const [kk, vv] = this.#planes[index]!.updateAndFetch(key[field], value[field]);
+        held.push(kk, vv); keys[field] = kk; values[field] = vv;
+      }
+      return [keys as ops.QuantizedTensor, values as ops.QuantizedTensor];
+    } catch (error) { return cleanupFailure(error, () => disposeResources(held)); }
+  }
+
+  /** Borrow prepared request state. All planes stage replacements before
+   * releasing the current group, preserving donor ownership on failure. */
+  mergeRows(rows: readonly Cache[]): void {
+    const next: BatchedKVCache[] = [];
+    try {
+      for (const [index, field] of FIELDS.entries()) {
+        const held: KVCache[] = [], inputs: Cache[] = [];
+        try {
+          for (const row of rows) {
+            if (row instanceof BatchedQuantizedKVCache) inputs.push(row.#planes[index]!);
+            else if (isQuantizedKvCache(row)) {
+              const view = new KVCache(); held.push(view);
+              if (row.offset > 0) {
+                const key = ops.copyOf(row.keys![field]);
+                let value: MlxArray;
+                try { value = ops.copyOf(row.values![field]); }
+                catch (error) { key.dispose(); throw error; }
+                view.restoreState(key, value, row.offset);
+              }
+              inputs.push(view);
+            } else throw new Error(`quantized row layout cannot merge ${row.signature()}`);
+          }
+          const plane = new BatchedKVCache(); next.push(plane); plane.mergeRows(inputs);
+        } finally { disposeResources(held); }
+      }
+    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    const previous = this.#planes; this.#planes = next; disposeResources(previous);
+    this.#minimum = rows.flatMap(row => row instanceof BatchedQuantizedKVCache ? row.#minimum : [row.minimumReusableOffset ?? 0]);
+  }
+
+  alignRows(leftPad: readonly number[]): void {
+    if (leftPad.every((pad, row) => pad === this.leftPad[row])) return;
+    const aligned = new KvTensorRows();
+    try {
+    aligned.mergeRows([{ planes: this.#planes.flatMap(plane => plane.keys ? [plane.keys, plane.values!] : []),
+      rowOffsets: this.rowOffsets, leftPad: this.leftPad }]);
+    aligned.alignRows(leftPad);
+    for (const plane of this.#planes) {
+      plane.keys?.dispose(); plane.values?.dispose();
+      plane.keys = aligned.planes.shift() ?? null; plane.values = aligned.planes.shift() ?? null;
+      plane.leftPad = [...leftPad];
+    }
+    } finally { aligned.dispose(); }
+  }
+  filterRows(keep: readonly number[]): void {
+    for (const plane of this.#planes) plane.filterRows(keep);
+    this.#minimum = keep.map(row => this.#minimum[row]!);
+  }
+  state(): MlxArray[] {
+    if (!this.#planes[0]!.keys) return [];
+    return [...this.#planes.map(plane => plane.keys!), ...this.#planes.map(plane => plane.values!)];
+  }
+
+  /** Return the existing persisted representation, owning compact arrays. */
+  extractRow(row: number): QuantizedKVCache {
+    const offset = this.rowOffsets[row]!, result = new QuantizedKVCache(this.groupSize, this.bits);
+    if (this.#minimum[row]) result.minimumReusableOffset = this.#minimum[row];
+    if (offset === 0) return result;
+    const held: KVCache[] = [];
+    try {
+      for (const plane of this.#planes) held.push(plane.extractRow(row));
+      result.restoreState(
+        { packed: held[0]!.keys!, scales: held[1]!.keys!, biases: held[2]!.keys! },
+        { packed: held[0]!.values!, scales: held[1]!.values!, biases: held[2]!.values! }, offset);
+      for (const cache of held) cache.keys = cache.values = null;
+      return result;
+    } finally { disposeResources(held); }
+  }
+
+  dispose(): void { disposeResources(this.#planes); this.#minimum = []; }
+}