mlx-bun 0.2.0 → 0.4.0

package/src/mlx/turboquant-ops.ts

@@ -1,8 +1,8 @@
-// TurboQuant KV-cache codec — pure mlx-op composition, ported op-for-op
+// TurboQuant KV-cache codec — reference operations ported op-for-op
 // from the vendored vllm-metal reference
 // (lab/repro/vllm-metal-turboquant/turboquant_reference.py). See
-// docs/design/turboquant-kv.md for the algorithm writeup and
-// goldens/turboquant.json (via tests/turboquant-ops.test.ts) for the
+// docs/design/turboquant.md for the algorithm writeup and
+// goldens/turboquant.json (via tests/research/turboquant-ops.test.ts) for the
 // bit-exactness oracle. Quantization groups run along the LAST axis
 // (head_dim) only — never the token axis, per the standing invariant.
 //
@@ -27,6 +27,10 @@ import {
   type TurboQuantHeadDim, type TurboQuantValueBits,
 } from "./turboquant-tables";
 
+import { tryJointDecodePackedKv, type PackedKvArrays } from "./turboquant-kv-decode";
+import { tryDecodePackedKvInverse256 } from "./turboquant-kv-inverse";
+import { isShapelessTracing } from "./compile";
+
 type S = MlxHandle;
 
 const BLOCK_SIZE = 32;
@@ -421,6 +425,32 @@ export function decodeValuesRotated(
   return outBf16;
 }
 
+/** Optional exact packed K/V decode. Inputs stay borrowed; the caller owns
+ * the returned pair. The eager path keeps the reference inverse rotation. */
+export function tryDecodePackedKv(
+  inputs: PackedKvArrays, kBits: number, vBits: number, headDim: number,
+  deferV = false, s: S = gpuStream,
+): [MlxArray, MlxArray] | null {
+  // Reject unsupported value widths before requesting their centroid table.
+  if (![2, 3, 4, 5, 8].includes(vBits)) return null;
+  if (!deferV && headDim === 256 && kBits === 8 && vBits === 3 &&
+      s === gpuStream && !isShapelessTracing() && inputs[0].shape[2]! >= 8192) {
+    const decoded = tryDecodePackedKvInverse256(
+      inputs, centroidsFor(vBits), signArray(256), kBits, vBits, headDim, s,
+    );
+    if (decoded) return decoded;
+  }
+  const decoded = tryJointDecodePackedKv(inputs, centroidsFor(vBits), kBits, vBits, headDim, deferV, s);
+  if (!decoded || deferV) return decoded;
+  const [key, rotated] = decoded;
+  try {
+    const value = fwht(rotated, false, s);
+    try { return [key, value.astype(Dtype.bfloat16, s)]; }
+    finally { value.dispose(); }
+  } catch (error) { key.dispose(); throw error; }
+  finally { rotated.dispose(); }
+}
+
 /** Undo the value rotation on an attention output computed against
  *  decodeValuesRotated values. f32 through the transform (matching the
  *  eager decode path's precision), returned in the input's dtype. */

package/src/mlx/turboquant-tables.ts

@@ -1,5 +1,5 @@
-// GENERATED by scripts/regen-turboquant-goldens.ts — DO NOT EDIT.
-// Regenerate:  bun scripts/regen-turboquant-goldens.ts
+// GENERATED by scripts/regen.ts turboquant — DO NOT EDIT.
+// Regenerate:  bun scripts/regen.ts turboquant
 //
 // Source: lab/repro/vllm-metal-turboquant/turboquant_reference.py (vendored
 // vllm-metal oracle, see PROVENANCE.md in that directory), executed under

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

@@ -2,11 +2,11 @@
 //
 // 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
+// parity-proven by tests/parity/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.)
+// wired into the serving path — phase S1a, docs/design/batching.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
@@ -15,6 +15,7 @@
 // thing the mask must encode is (a) causality, (b) each row's left padding,
 // and (c) the sliding window, if any.
 
+import { cacheSignature } from "./gemma4-base";
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
 import { createCausalMask, KVCache, type Cache, type Mask } from "./gemma4-base";
@@ -92,6 +93,11 @@ export class BatchedDecodeMaskCache implements Cache {
   #ropeArr: MlxArray | null = null;
   #ropeForOffset = -1;
 
+  /** A decorator, not a storage kind: never claims the inner cache's
+   *  identity (isPlainKvCache(wrapper) must stay false — it is not a
+   *  KVCache), but records what it wraps for diagnostics/persistence. */
+  signature(): string { return `kv:batched-mask:${cacheSignature(this.inner)}`; }
+
   constructor(
     private readonly inner: Cache,
     private readonly B: number,
@@ -253,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

@@ -17,7 +17,7 @@
 // (maybeQuantizeKv at chunk boundaries — see BatchScheduler.#quantizeSolo),
 // so a row's quantized bytes are bit-exact vs serial `--kv-quant config` by
 // construction; this module only re-arranges those bytes across the batch
-// axis. Gates: tests/batched-kv-quant-parity.test.ts.
+// axis. Gates: tests/parity/batched-kv-quant-parity.test.ts.
 
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
@@ -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 = []; }
+}