@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/state/batched-quant.ts

@@ -0,0 +1,389 @@
+// Batched QUANTIZED KV for continuous batching — the quantized twins of
+// batched-mask.ts's dynamic-B ops, over (packed, scales, biases) triples.
+//
+// Why this is safe surgery: mlx quantization packs along HEAD_DIM (the last
+// axis; packed last-dim = D/(32/bits), scales/biases last-dim = D/groupSize),
+// so the TOKEN axis (axis 2) stays token-granular in all three tensors —
+// left-padding, concat, and row-filter at token granularity never touch a
+// quantization group. Pad columns hold zeros; they are never attended (the
+// padding mask covers them) and never dequantized into anything a real row
+// reads. QuantizedKVCache.updateAndFetchQuantized already handles arbitrary
+// starting widths (it trims to `offset` and regrows in STEP chunks), so a
+// merged exact-width buffer needs no STEP alignment here.
+//
+// Composition contract (the L2 oracle binds per ROW): each joiner prefills
+// alone through the per-layer conversion its bf16-first lego runs (the
+// conversion at the scheme's start), so a row's quantized bytes are bit-exact
+// vs the same request run alone under `--kv-quant config` by construction;
+// this module only re-arranges those bytes across the batch axis. Gates: `02d723a:tests/parity/batched-kv-quant-parity.test.ts`.
+
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { QuantizedKVCache } from "./quantized-kv";
+import { type AttentionCache, type AttentionRead, type BatchableCache, type Cache, type Mask } from "../contracts/mlx/cache";
+import type { QuantizedAttentionState } from "./kv-storage";
+import { buildBatchedDecodeMask } from "./batched-mask";
+import { isQuantizedKvCache } from "./capabilities";
+import { runtimeValue } from "../runtime/config";
+import { AttentionMasks, withLease } from "./attention-read";
+import { affineRead, rowLayoutKernels, type AffineKernel, type AffineKernels } from "./affine-attention";
+
+export interface QuantRow {
+  keys: ops.QuantizedTensor;
+  values: ops.QuantizedTensor;
+}
+
+const disposeTriple = (t: ops.QuantizedTensor): void => {
+  t.packed.dispose();
+  t.scales.dispose();
+  t.biases.dispose();
+};
+
+/** Apply an MlxArray op to each member of a triple. Each tensor keeps its
+ *  OWN last-dim (packed vs scales/biases differ) — helpers must never share
+ *  one `D` across the triple (recon hazard, 2026-07-05). */
+const tripleMap = (
+  t: ops.QuantizedTensor, f: (a: MlxArray) => MlxArray,
+): ops.QuantizedTensor => ({ packed: f(t.packed), scales: f(t.scales), biases: f(t.biases) });
+
+/** Left-pad one tensor of a triple by `pad` tokens (axis 2) with zeros of
+ *  its own dtype/last-dim. pad=0 returns a fresh full view (mirrors
+ *  batched-mask.ts leftPadTo so ownership is uniform). */
+const leftPadTo = (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;
+};
+
+/** Quantized twin of mergeKVRows: stack N single-row quantized KV slices
+ *  into one left-padded [B,H,Smax,*] triple pair. Caller owns the result;
+ *  input rows are not disposed. */
+export function mergeQuantRows(rows: QuantRow[]): {
+  keys: ops.QuantizedTensor; values: ops.QuantizedTensor; leftPad: number[]; width: number;
+} {
+  const lens = rows.map((r) => r.keys.packed.shape[2]!);
+  const width = Math.max(...lens);
+  const leftPad = lens.map((l) => width - l);
+  const padded = rows.map((r, i) => ({
+    keys: tripleMap(r.keys, (a) => leftPadTo(a, leftPad[i]!)),
+    values: tripleMap(r.values, (a) => leftPadTo(a, leftPad[i]!)),
+  }));
+  const cat = (pick: (r: QuantRow) => ops.QuantizedTensor): ops.QuantizedTensor => ({
+    packed: ops.concatAxis(padded.map((r) => pick(r).packed), 0),
+    scales: ops.concatAxis(padded.map((r) => pick(r).scales), 0),
+    biases: ops.concatAxis(padded.map((r) => pick(r).biases), 0),
+  });
+  const keys = cat((r) => r.keys);
+  const values = cat((r) => r.values);
+  for (const r of padded) { disposeTriple(r.keys); disposeTriple(r.values); }
+  return { keys, values, leftPad, width };
+}
+
+/** Quantized twin of extendKVRows: append ONE right-justified quantized row
+ *  to an existing batched triple pair in one pad + one B-axis concat.
+ *  Existing pads grow, never shrink (mlx-lm extend semantics). Caller owns
+ *  the result; inputs are not disposed. */
+export function extendQuantRows(
+  keys: ops.QuantizedTensor, values: ops.QuantizedTensor, leftPad: number[], row: QuantRow,
+): { keys: ops.QuantizedTensor; values: ops.QuantizedTensor; leftPad: number[]; width: number } {
+  const S = keys.packed.shape[2]!;
+  const Lr = row.keys.packed.shape[2]!;
+  const width = Math.max(S, Lr);
+  const bk = tripleMap(keys, (a) => leftPadTo(a, width - S));
+  const bv = tripleMap(values, (a) => leftPadTo(a, width - S));
+  const rk = tripleMap(row.keys, (a) => leftPadTo(a, width - Lr));
+  const rv = tripleMap(row.values, (a) => leftPadTo(a, width - Lr));
+  const cat = (b: ops.QuantizedTensor, r: ops.QuantizedTensor): ops.QuantizedTensor => ({
+    packed: ops.concatAxis([b.packed, r.packed], 0),
+    scales: ops.concatAxis([b.scales, r.scales], 0),
+    biases: ops.concatAxis([b.biases, r.biases], 0),
+  });
+  const outK = cat(bk, rk);
+  const outV = cat(bv, rv);
+  for (const t of [bk, bv, rk, rv]) disposeTriple(t);
+  return {
+    keys: outK, values: outV,
+    leftPad: [...leftPad.map((p) => p + (width - S)), width - Lr],
+    width,
+  };
+}
+
+/** Quantized twin of extractKVRow (mlx-lm BatchKVCache.extract,
+ *  models/cache.py:1080, over triples): pull row `i` past its left padding
+ *  into a fresh SERIAL QuantizedKVCache with OWNED contiguous copies.
+ *  Token-axis slicing is byte-safe because quantization packs along
+ *  HEAD_DIM (file header) — a row's quantized bytes come out exactly as
+ *  the serial oracle wrote them. */
+export function extractQuantRow(
+  cache: QuantizedKVCache, leftPad: number, i: number, kernels: AffineKernels = cache.kernels,
+): QuantizedKVCache {
+  if (!cache.keys || !cache.values) throw new Error("extractQuantRow: empty cache");
+  const S = cache.offset;
+  const cut = (t: ops.QuantizedTensor): ops.QuantizedTensor =>
+    tripleMap(t, (a) => {
+      const [, H, , D] = a.shape as [number, number, number, number];
+      const view = a.slice([i, 0, leftPad, 0], [i + 1, H, S, D]);
+      const own = ops.copyOf(view); // TRUE copy: contiguous(view) is a no-op VIEW when already contiguous — pins the source buffer (2026-08-20 DeltaNet conv leak class)
+      view.dispose();
+      return own;
+    });
+  const out = new QuantizedKVCache(cache.groupSize, cache.bits, kernels);
+  out.restoreState(cut(cache.keys), cut(cache.values), S - leftPad);
+  return out;
+}
+
+/** 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[], 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) => {
+      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 };
+}
+
+/** Per-step cache wrapper for batched QUANTIZED decode with left-padded
+ *  rows — the quantized twin of BatchedDecodeMaskCache. Subclasses
+ *  QuantizedKVCache so attention's `instanceof` dispatch routes to
+ *  quantizedSdpa exactly as for a serial quantized cache; delegates KV
+ *  storage to the persistent batched inner and overrides the two things a
+ *  padded batch needs: the padding-aware decode mask and per-row RoPE
+ *  positions. Ephemeral — the scheduler builds a fresh wrapper each step
+ *  and calls releaseRopeArr() after (never dispose(), which would free the
+ *  inner). Rope is captured ONCE per step before updateAndFetch (attention
+ *  reads it pre-write; see packages/inference/src/models/minicpm5/model.ts
+ *  LlamaAttention.forward), so an
+ *  eagerly-built [B] position array is exact.
+ *
+ *  Reads: append to the inner storage, then the inner cache's kernel for the
+ *  read under the padding-aware mask (`buildBatchedDecodeMask`) for the query
+ *  count, built once per forward through the inner cache's `masks`. That mask
+ *  is an array, so a window takes the masked-window kernel. */
+export class BatchedQuantDecodeMaskCache extends QuantizedKVCache {
+  override readonly ropeOffsetArr: MlxArray;
+
+  constructor(
+    private readonly inner: QuantizedKVCache,
+    private readonly B: number,
+    private readonly leftPad: number[],
+  ) {
+    super(inner.groupSize, inner.bits, inner.kernels, inner.masks);
+    this.offset = inner.offset;
+    const data = new Int32Array(B);
+    for (let b = 0; b < B; b++) data[b] = inner.offset - leftPad[b]!;
+    this.ropeOffsetArr = MlxArray.fromInt32(data, [B]);
+  }
+
+  override appendDecode(k: MlxArray, v: MlxArray): AttentionRead { return this.#read(k, v, this.kernels.decode); }
+
+  override appendWindow(k: MlxArray, v: MlxArray): AttentionRead { return this.#read(k, v, this.kernels.maskedWindow); }
+
+  #read(k: MlxArray, v: MlxArray, kernel: AffineKernel): AttentionRead {
+    const N = k.shape[2]!, S = this.inner.offset + N;
+    const mask = this.masks.lease("padded-kv", `${N}|${S}|${this.leftPad.join(",")}`,
+      () => ({ mode: "array", arr: buildBatchedDecodeMask(this.B, N, S, this.leftPad, null) }));
+    return withLease(mask, held => {
+      const [keys, values] = this.updateAndFetchQuantized(k, v);
+      return affineRead(keys, values, held, kernel);
+    });
+  }
+
+  override updateAndFetchQuantized(k: MlxArray, v: MlxArray): [ops.QuantizedTensor, ops.QuantizedTensor] {
+    const r = this.inner.updateAndFetchQuantized(k, v);
+    this.offset = this.inner.offset;
+    return r;
+  }
+
+  override makeMask(N: number, _windowSize: number | null): Mask {
+    // Same convention as BatchedDecodeMaskCache: built pre-write, S spans
+    // offset + N; quantizedSdpaUnfused expands it for the 5-D GQA scores.
+    const S = this.inner.offset + N;
+    return { mode: "array", arr: buildBatchedDecodeMask(this.B, N, S, this.leftPad, null) };
+  }
+
+  /** Free the per-step RoPE array WITHOUT touching the persistent inner —
+   *  name and signature must match the scheduler's structural
+   *  `(c as { releaseRopeArr?: () => void })` cast. */
+  releaseRopeArr(): void {
+    this.ropeOffsetArr.dispose();
+  }
+
+  override state(): MlxArray[] {
+    return this.inner.state();
+  }
+  override isTrimmable(): boolean {
+    return this.inner.isTrimmable();
+  }
+  override trim(n: number): void {
+    this.inner.trim(n);
+  }
+  override dispose(): void {
+    this.inner.dispose();
+  }
+}
+
+/** The quantized twin of PaddedKVRows: a running batch's affine-quantized
+ *  full-attention KV as (packed, scales, biases) triples with right-aligned rows.
+ *  Its rows were converted by the serial per-layer maintenance before they
+ *  joined, so their quantized bytes are the serial cache's; this layout only
+ *  re-arranges them along the batch axis (`extendQuantRows`, or the re-merge
+ *  with MLX_BUN_BATCH_EXTEND=0) and, while any row is padded, supplies the
+ *  padding-aware mask and per-row RoPE positions. Sources are borrowed: a
+ *  running layout or one serial QuantizedKVCache, then one per joining row.
+ *
+ *  Reads: while no row is padded, the serial cache's reads; while one is, the
+ *  step wrapper's (`BatchedQuantDecodeMaskCache`): the same kernels under the
+ *  padding-aware mask, with the masked-window kernel for a window. The kernels
+ *  are `kernels` while the layout holds one unpadded row and its `multiRow` set
+ *  otherwise (`rowLayoutKernels`), chosen when the rows change; rows it
+ *  extracts keep `kernels`. */
+export class PaddedQuantKVRows implements AttentionCache, BatchableCache {
+  #inner: QuantizedKVCache;
+  #leftPad: number[] = [];
+  #padded = false;
+  #view: BatchedQuantDecodeMaskCache | null = null;
+
+  constructor(
+    readonly groupSize: number,
+    readonly bits: number,
+    /** The kernel of each read; empty batches and extracted rows keep them. */
+    readonly kernels: AffineKernels,
+    /** The per-forward masks of the model this layout serves. */
+    readonly masks = new AttentionMasks(),
+  ) { this.#inner = this.#storage([]); }
+
+  appendDecode(k: MlxArray, v: MlxArray): AttentionRead {
+    return (this.#padded ? this.#step() : this.#inner).appendDecode(k, v);
+  }
+
+  appendWindow(k: MlxArray, v: MlxArray): AttentionRead {
+    return (this.#padded ? this.#step() : this.#inner).appendWindow(k, v);
+  }
+
+  /** Empty serial storage that reads with this layout's masks and the kernels for rows at `leftPad`. */
+  #storage(leftPad: readonly number[]): QuantizedKVCache {
+    return new QuantizedKVCache(this.groupSize, this.bits, rowLayoutKernels(this.kernels, leftPad.length <= 1 && !leftPad[0]), this.masks);
+  }
+
+  signature(): string { return `kv:padded-rows-quant:${this.bits}:${this.groupSize}`; }
+  get batchSize(): number | null { return this.#leftPad.length || null; }
+  get leftPad(): readonly number[] { return this.#leftPad; }
+  get rowOffsets(): readonly number[] { return this.#leftPad.map((pad) => this.#inner.offset - pad); }
+  get offset(): number { return this.#inner.offset; }
+  /** Attention reads and appends the quantized planes directly. */
+  get quantizedAttention(): QuantizedAttentionState { return this.#inner; }
+  /** Present only while a row is padded: each row's own position. */
+  get ropeOffsetArr(): MlxArray | undefined { return this.#padded ? this.#step().ropeOffsetArr : undefined; }
+  makeEmptyBatch(): PaddedQuantKVRows { return new PaddedQuantKVRows(this.groupSize, this.bits, this.kernels, this.masks); }
+
+  bytesPerToken(): number {
+    return this.#leftPad.length ? this.#inner.bytesPerToken() / this.#leftPad.length : 0;
+  }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * tokens; }
+
+  /** One wrapper per step: its RoPE positions are fixed at construction, before the step's writes. */
+  #step(): BatchedQuantDecodeMaskCache {
+    return this.#view ??= new BatchedQuantDecodeMaskCache(this.#inner, this.#leftPad.length, [...this.#leftPad]);
+  }
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.#inner.updateAndFetch(k, v); }
+  makeMask(N: number, windowSize: number | null): Mask {
+    return this.#padded ? this.#step().makeMask(N, windowSize) : this.#inner.makeMask(N, windowSize);
+  }
+  /** Free this step's RoPE positions after the step's graph is built; the KV stays. */
+  releaseRopeArr(): void { this.#view?.releaseRopeArr(); this.#view = null; }
+  state(): MlxArray[] { return this.#inner.state(); }
+  isTrimmable(): boolean { return this.#inner.isTrimmable(); }
+  trim(n: number): void { this.#inner.trim(n); }
+
+  #set(inner: QuantizedKVCache, leftPad: number[]): void {
+    this.releaseRopeArr();
+    this.#inner = inner; this.#leftPad = leftPad; this.#padded = leftPad.some((pad) => pad !== 0);
+  }
+
+  mergeRows(rows: readonly Cache[]): void {
+    if (this.#leftPad.length) throw new Error("PaddedQuantKVRows.mergeRows needs an empty batch");
+    if (!rows.length) return;
+    const held: ops.QuantizedTensor[] = [];
+    const own = (triple: ops.QuantizedTensor): ops.QuantizedTensor => { held.push(triple); return triple; };
+    type Part = { row: QuantRow; pads: number[] };
+    const partOf = (cache: Cache): Part => {
+      if (cache instanceof PaddedQuantKVRows) {
+        const [keys, values] = cache.#inner.temporalView();
+        return { row: { keys: own(keys), values: own(values) }, pads: [...cache.#leftPad] };
+      }
+      if (!isQuantizedKvCache(cache)) throw new Error(`quantized full-attention batch cannot merge ${cache.signature()}`);
+      const [keys, values] = cache.temporalView();
+      return { row: { keys: own(keys), values: own(values) }, pads: [0] };
+    };
+    const cut = (triple: ops.QuantizedTensor, b: number, pad: number): ops.QuantizedTensor => own(tripleMap(triple, (a) =>
+      a.slice([b, 0, pad, 0], [b + 1, a.shape[1]!, a.shape[2]!, a.shape[3]!])));
+    const extend = runtimeValue("MLX_BUN_BATCH_EXTEND") !== "0";
+    let final: { keys: ops.QuantizedTensor; values: ops.QuantizedTensor } | undefined;
+    try {
+      let acc = partOf(rows[0]!);
+      if (rows.length === 1) {
+        const single = mergeQuantRows(acc.pads.map((pad, b) => ({ keys: cut(acc.row.keys, b, pad), values: cut(acc.row.values, b, pad) })));
+        acc = { row: { keys: own(single.keys), values: own(single.values) }, pads: single.leftPad };
+      }
+      for (const source of rows.slice(1)) {
+        const next = partOf(source);
+        if (next.pads.length !== 1) throw new Error("PaddedQuantKVRows joins one solo row at a time");
+        let joined: { keys: ops.QuantizedTensor; values: ops.QuantizedTensor; leftPad: number[] };
+        if (extend) {
+          joined = extendQuantRows(acc.row.keys, acc.row.values, acc.pads, next.row);
+        } else {
+          const rowsOut = acc.pads.map((pad, b) => ({ keys: cut(acc.row.keys, b, pad), values: cut(acc.row.values, b, pad) }));
+          joined = mergeQuantRows([...rowsOut, next.row]);
+        }
+        acc = { row: { keys: own(joined.keys), values: own(joined.values) }, pads: joined.leftPad };
+      }
+      final = acc.row;
+      held.splice(held.indexOf(final.keys), 1); held.splice(held.indexOf(final.values), 1);
+      const cache = this.#storage(acc.pads);
+      cache.restoreState(final.keys, final.values, final.keys.packed.shape[2]!);
+      final = undefined;
+      this.#set(cache, acc.pads);
+    } finally {
+      for (const triple of held) disposeTriple(triple);
+      if (final) { disposeTriple(final.keys); disposeTriple(final.values); }
+    }
+  }
+
+  extractRow(row: number): Cache {
+    return this.#inner.keys ? extractQuantRow(this.#inner, this.#leftPad[row]!, row, this.kernels)
+      : new QuantizedKVCache(this.groupSize, this.bits, this.kernels);
+  }
+
+  filterRows(keep: readonly number[]): void {
+    if (!keep.length) { this.dispose(); return; }
+    const kept = keep.map((row) => this.#leftPad[row]!);
+    const shared = Math.min(...kept); // padding every survivor shares
+    const [keys, values] = this.#inner.temporalView();
+    let filtered: { keys: ops.QuantizedTensor; values: ops.QuantizedTensor };
+    try { filtered = filterQuantRows(keys, values, [...keep], shared); }
+    finally { disposeTriple(keys); disposeTriple(values); }
+    const pads = kept.map((pad) => pad - shared);
+    const cache = this.#storage(pads);
+    cache.restoreState(filtered.keys, filtered.values, this.#inner.offset - shared);
+    this.#inner.dispose();
+    this.#set(cache, pads);
+  }
+
+  dispose(): void {
+    this.#inner.dispose();
+    this.#set(this.#storage([]), []);
+  }
+}

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

@@ -0,0 +1,209 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { cleanupFailure, disposeResources } from "../runtime/resources";
+import { BatchedKVCache } from "./batched-kv";
+import { KVCache } from "./kv";
+import { QuantizedKVCache } from "./quantized-kv";
+import { isQuantizedKvCache } from "./capabilities";
+import { type AttentionRead, type BatchableCache, type Cache, type CommittedAttentionCache, type Mask, type PaddedPrefillCache, type PrefillPadding } from "../contracts/mlx/cache";
+import type { QuantizedAttentionState } from "./kv-storage";
+
+import { KvTensorRows } from "./kv-tensor-rows";
+import { quantizedDonorAttention } from "./kv-attention-view";
+import { dequantizedFetch } from "./quantized-tensor";
+import { AttentionMasks, unmaskedLease, withLease, type MaskLease } from "./attention-read";
+import { affineCausalLease, affineRead, committedAffineRead, rowLayoutKernels, type AffineKernel, type AffineKernels } from "./affine-attention";
+
+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.
+ *
+ * Reads: append through `updateAndFetchQuantized`, then the kernel `kernels`
+ * fixed for the read at construction. While every row is unpadded and at one
+ * position, decode passes no mask and a window the bottom-right causal matrix,
+ * as `makeMask` leads the graphs to today; otherwise both pass the per-row mask
+ * (causal, past each row's left padding) for their query count, built once per
+ * forward through `masks`, and a window takes the masked-window kernel. The
+ * kernels are `kernels` while the layout holds one unpadded row and its
+ * `multiRow` set otherwise (`rowLayoutKernels`); rows it extracts keep
+ * `kernels`. The committed read is `quantizedAppendAttention`, for one row. */
+export class BatchedQuantizedKVCache implements BatchableCache, QuantizedAttentionState, PaddedPrefillCache, CommittedAttentionCache {
+  #planes = FIELDS.map(() => new BatchedKVCache());
+  #minimum: number[] = [];
+  get minimumReusableOffset(): number { return Math.max(0, ...this.#minimum); }
+  constructor(
+    readonly groupSize: number,
+    readonly bits: number,
+    /** The kernel of each read; empty batches and extracted rows keep them. */
+    readonly kernels: AffineKernels,
+    /** The per-forward masks of the model this layout serves. */
+    readonly masks = new AttentionMasks(),
+  ) {}
+
+  appendDecode(k: MlxArray, v: MlxArray): AttentionRead {
+    return this.#read(k, v, this.#planes[0]!.aligned ? unmaskedLease : this.#rowMask(1), this.#kernels().decode);
+  }
+
+  appendWindow(k: MlxArray, v: MlxArray): AttentionRead {
+    const L = k.shape[2]!;
+    return this.#planes[0]!.aligned
+      ? this.#read(k, v, affineCausalLease(this.masks, L, this.offset + L), this.#kernels().window)
+      : this.#read(k, v, this.#rowMask(L), this.#kernels().maskedWindow);
+  }
+
+  /** `kernels` while this layout holds one unpadded row, its `multiRow` set otherwise. */
+  #kernels(): AffineKernels {
+    return rowLayoutKernels(this.kernels, this.batchSize === 1 && this.#planes[0]!.aligned);
+  }
+
+  /** Fill spans run on this layout (the target layout of a full-attention
+   * layer): the graph's independent-rows path. */
+  appendCommitted(k: MlxArray, v: MlxArray): AttentionRead {
+    const [keys, values] = this.updateAndFetchQuantized(k, v);
+    return committedAffineRead(keys, values, this.groupSize, this.bits);
+  }
+
+  #read(k: MlxArray, v: MlxArray, mask: MaskLease, kernel: AffineKernel): AttentionRead {
+    return withLease(mask, held => {
+      const [keys, values] = this.updateAndFetchQuantized(k, v);
+      return affineRead(keys, values, held, kernel);
+    });
+  }
+
+  /** The per-row mask of an `N`-position append, before it: the bf16 layout's
+   * (`BatchedKVCache`), whose positions every plane shares. */
+  #rowMask(N: number): MaskLease {
+    return this.masks.lease("batched-kv", `${N}|${this.rowOffsets.join(",")}|${this.leftPad.join(",")}`,
+      () => this.#planes[0]!.makeMask(N, null));
+  }
+  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("../contracts/mlx/cache").KvDonorAttention {
+    const views: import("../contracts/mlx/cache").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]!;
+      // The kernels this layout reads with for its rows (a one-sequence set
+      // reads several rows with its `multiRow` set).
+      return quantizedDonorAttention(field("keys"), field("values"), { offsets, starts, ends }, this.#kernels());
+    } 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, this.kernels, this.masks); }
+  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);
+  }
+  /** The dense fetch: appends, then the rows' stored windows dequantized. */
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    return dequantizedFetch(this.updateAndFetchQuantized(k, v), this.groupSize, this.bits);
+  }
+
+  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, this.kernels);
+    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 = []; this.masks.clear(); }
+}