mlx-bun 0.4.0 → 1.0.0

package/src/model/batched-mask.ts

@@ -1,279 +0,0 @@
-// Padding-aware attention masks for batched (B>1) serving — decode side.
-//
-// Companion to src/train/forward.ts's buildBatchedPadMask, which builds the
-// PREFILL mask ([B,1,L,L], offset 0, right-padded — the training case, already
-// parity-proven by tests/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/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
-// the whole batch. Each row's RoPE position is still per-row (handled via the
-// array-offset path, ropeOffsetArr / ops.ropeDynamic — not here); the only
-// thing the mask must encode is (a) causality, (b) each row's left padding,
-// and (c) the sliding window, if any.
-
-import { 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";
-
-/** Padding-aware key-validity mask for batched DECODE with left-padded rows.
- *
- *  The shared KV buffer has width S (== cache offset once this step's keys are
- *  written). Row b has `leftPad[b]` padding columns on the left, then its real
- *  keys. A step forwards N new query tokens whose absolute positions are
- *  S-N .. S-1 (N == 1 in the steady decode state; N > 1 only for a batched
- *  prefill chunk).
- *
- *  Returns [B, 1, N, S] bool, true where query i of row b may attend to key j:
- *      allow[b,0,i,j] = (j <= S-N+i)            causal (query i at abs pos S-N+i)
- *                       AND (j >= leftPad[b])    key past row b's left padding
- *                       AND (S-N+i < j+window)   within the sliding window*
- *      (*only when windowSize !== null)
- *
- *  Built from createCausalMask(N, S-N, windowSize) (the model's own [N, S] bool
- *  matrix) broadcast over B, AND-ed with a per-row [B,1,1,S] key-validity mask
- *  — exactly mirroring buildBatchedPadMask, but at a decode offset with
- *  left-padding instead of offset-0 right-padding. Caller owns the result. */
-export function buildBatchedDecodeMask(
-  B: number,
-  N: number,
-  S: number,
-  leftPad: number[],
-  windowSize: number | null,
-): MlxArray {
-  // [N, S] causal(+window): queries at abs positions S-N..S-1, keys 0..S-1.
-  const causal2d = createCausalMask(N, S - N, windowSize); // [N, S] bool
-  const causal4d = ops.reshape(causal2d, [1, 1, N, S]);
-  causal2d.dispose();
-
-  // Per-row key validity: keyValid[b,0,0,j] = j >= leftPad[b]. Build as int32
-  // {0,1} then compare to bool, broadcasting over the query axis (i).
-  const keyValidData = new Int32Array(B * S);
-  for (let b = 0; b < B; b++) {
-    const pad = leftPad[b]!;
-    for (let j = 0; j < S; j++) keyValidData[b * S + j] = j >= pad ? 1 : 0;
-  }
-  const keyValidI = MlxArray.fromInt32(keyValidData, [B, 1, 1, S]);
-  const zero = MlxArray.fromInt32(new Int32Array([0]), []);
-  const keyValid = ops.less(zero, keyValidI); // 0 < x → bool, true where x==1
-  zero.dispose();
-  keyValidI.dispose();
-
-  const allow = ops.logicalAnd(causal4d, keyValid); // [1,1,N,S] & [B,1,1,S]
-  causal4d.dispose();
-  keyValid.dispose();
-  return allow; // [B, 1, N, S] bool
-}
-
-/** Cache wrapper for batched DECODE with left-padded rows. Delegates all KV
- *  storage to a real, shape-[B,H,S,D] cache (KVCache et al. are already
- *  B-generic), but overrides two things the stock single-sequence path gets
- *  wrong for a padded batch:
- *
- *   - `makeMask` returns the per-row left-padding decode mask
- *     (buildBatchedDecodeMask) instead of the stock mask, which is EMPTY at
- *     N==1 and would let every row attend to its own (garbage) left padding.
- *   - `ropeOffsetArr` exposes each row's REAL position (= pre-write offset −
- *     leftPad[b]) as an int32 [B] array, so RoPE is applied per-row via the
- *     array-offset path (mlx fast_rope_dynamic — verified to rotate each row
- *     by its own offset). The scalar offset would mis-position every padded
- *     row by its pad amount.
- *
- *  `leftPad` is fixed for the batch's lifetime (set when the left-padded decode
- *  buffer is assembled from per-row prefills). The position array is recomputed
- *  only when the inner offset advances (once per decode step), so the two reads
- *  within a step (K before updateAndFetch, Q after — see Attention.forward)
- *  return the SAME array; the attention graph holds it until eval, so it must
- *  not be freed mid-step. */
-export class BatchedDecodeMaskCache implements Cache {
-  #ropeArr: MlxArray | null = null;
-  #ropeForOffset = -1;
-
-  /** 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,
-    private readonly leftPad: number[],
-    private readonly window: number | null,
-  ) {}
-
-  get offset(): number {
-    return this.inner.offset;
-  }
-
-  get ropeOffsetArr(): MlxArray {
-    const off = this.inner.offset;
-    if (this.#ropeArr && this.#ropeForOffset === off) return this.#ropeArr;
-    this.#ropeArr?.dispose(); // prior step's array; its graph is already evaluated
-    const data = new Int32Array(this.B);
-    for (let b = 0; b < this.B; b++) data[b] = off - this.leftPad[b]!;
-    this.#ropeArr = MlxArray.fromInt32(data, [this.B]);
-    this.#ropeForOffset = off;
-    return this.#ropeArr;
-  }
-
-  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
-    return this.inner.updateAndFetch(k, v);
-  }
-
-  makeMask(N: number, _windowSize: number | null): Mask {
-    // forwardLayers calls makeMask BEFORE this step's updateAndFetch, so the
-    // KV the attention then fetches spans inner.offset + N keys. Build a fresh
-    // mask each step and hand the model ownership — forwardLayers disposes
-    // mask.arr after building the layer graph (mlx refcounts it into the sdpa
-    // nodes, so post-build dispose is safe), so no view bookkeeping is needed.
-    const S = this.inner.offset + N;
-    return { mode: "array", arr: buildBatchedDecodeMask(this.B, N, S, this.leftPad, this.window) };
-  }
-
-  /** Free the per-step RoPE position array WITHOUT disposing the delegated KV
-   *  (inner). For callers that rebuild a fresh wrapper around a persistent inner
-   *  cache each decode step (the dynamic-B scheduler / parity harness) — the
-   *  inner outlives the wrapper, so wrapper.dispose() would be wrong. */
-  releaseRopeArr(): void {
-    this.#ropeArr?.dispose();
-    this.#ropeArr = null;
-    this.#ropeForOffset = -1;
-  }
-
-  state(): MlxArray[] {
-    return this.inner.state();
-  }
-  isTrimmable(): boolean {
-    return this.inner.isTrimmable();
-  }
-  trim(n: number): void {
-    this.inner.trim(n);
-  }
-  dispose(): void {
-    this.inner.dispose();
-    this.#ropeArr?.dispose();
-    this.#ropeArr = null;
-  }
-}
-
-// --- Dynamic-B cache ops (continuous batching) ---------------------------
-// Ports of mlx-lm's per-cache batch methods (merge / filter), the operations
-// the --batch N scheduler needs that our fixed-B verified forward doesn't have.
-// They run per layer on the raw [.,H,.,D] K/V; the scheduler applies them
-// across layers and wraps the result in a BatchedDecodeMaskCache. Pure array
-// surgery — unit-tested without a model (tests/batched-decode-mask.test.ts);
-// numerical correctness is gated separately vs an mlx-lm dynamic golden.
-
-/** mlx-lm `merge`: stack N single-row KV slices into one left-padded
- *  [B,H,Smax,D] batch. Each row is its sequence's KV sliced to its true length
- *  L_i; we left-pad each by Smax−L_i so all rows align on the right edge (the
- *  layout decode + BatchedDecodeMaskCache expect). Generalizes
- *  realBatchedGreedy's assemble to N rows. Caller owns the result; the input
- *  row arrays are not disposed. */
-export function mergeKVRows(
-  rows: { keys: MlxArray; values: MlxArray }[],
-): { keys: MlxArray; values: MlxArray; leftPad: number[]; width: number } {
-  const lens = rows.map((r) => r.keys.shape[2]!);
-  const width = Math.max(...lens);
-  const leftPad = lens.map((l) => width - l);
-  const leftPadTo = (a: MlxArray, pad: number): MlxArray => {
-    if (pad === 0) return a.slice([0, 0, 0, 0], a.shape as number[]); // fresh full view
-    const [B, H, , D] = a.shape as [number, number, number, number];
-    const z = ops.zeros([B, H, pad, D], a.dtype);
-    const out = ops.concatAxis([z, a], 2);
-    z.dispose();
-    return out;
-  };
-  const ks = rows.map((r, i) => leftPadTo(r.keys, leftPad[i]!));
-  const vs = rows.map((r, i) => leftPadTo(r.values, leftPad[i]!));
-  const keys = ops.concatAxis(ks, 0);
-  const values = ops.concatAxis(vs, 0);
-  for (const a of [...ks, ...vs]) a.dispose();
-  return { keys, values, leftPad, width };
-}
-
-/** mlx-lm `BatchKVCache.extend`: append ONE new right-justified row to an
- *  existing batched [B,H,S,D] buffer WITHOUT re-extracting the running rows
- *  (the O(B·S) slice storm mergeKVRows pays on a join). Pad whichever side
- *  is narrower on the LEFT to the common width, then a single B-axis concat.
- *  Existing rows' left-pads grow by the batch-side pad (they never shrink —
- *  mlx-lm extend semantics; the decode mask covers pads either way).
- *  Caller owns the result; inputs are not disposed. */
-export function extendKVRows(
-  keys: MlxArray, values: MlxArray, leftPad: number[],
-  row: { keys: MlxArray; values: MlxArray },
-): { keys: MlxArray; values: MlxArray; leftPad: number[]; width: number } {
-  const S = keys.shape[2]!;
-  const Lr = row.keys.shape[2]!;
-  const width = Math.max(S, Lr);
-  const padTo = (a: MlxArray, pad: number): MlxArray => {
-    if (pad === 0) return a.slice([0, 0, 0, 0], a.shape as number[]); // fresh full view
-    const [B, H, , D] = a.shape as [number, number, number, number];
-    const z = ops.zeros([B, H, pad, D], a.dtype);
-    const out = ops.concatAxis([z, a], 2);
-    z.dispose();
-    return out;
-  };
-  const bk = padTo(keys, width - S);
-  const bv = padTo(values, width - S);
-  const rk = padTo(row.keys, width - Lr);
-  const rv = padTo(row.values, width - Lr);
-  const outK = ops.concatAxis([bk, rk], 0);
-  const outV = ops.concatAxis([bv, rv], 0);
-  for (const a of [bk, bv, rk, rv]) a.dispose();
-  return {
-    keys: outK,
-    values: outV,
-    leftPad: [...leftPad.map((p) => p + (width - S)), width - Lr],
-    width,
-  };
-}
-
-/** mlx-lm `BatchKVCache.extract` (models/cache.py:1080): pull row `i` out of
- *  a batched full-attention cache into a fresh SERIAL KVCache — the row slice
- *  past its left padding, materialized as an OWNED contiguous copy (extraction
- *  outlives the batch; the source buffers are about to be filter()ed or
- *  disposed). `leftPad` is the scheduler's per-row padding for this row (full
- *  layers share one offset; the scheduler owns the pad list). Bit-exact vs a
- *  solo run by construction: merge/extend/filter are byte-preserving per row
- *  (tests/batched-decode-parity, tests/batched-extract) and this is a pure
- *  slice+copy of those bytes. Off the hot decode path. */
-export function extractKVRow(cache: KVCache, leftPad: number, i: number): KVCache {
-  if (!cache.keys || !cache.values) throw new Error("extractKVRow: empty cache");
-  const S = cache.offset; // valid columns (buffer may be STEP-padded past it)
-  const cut = (a: MlxArray): MlxArray => {
-    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 KVCache();
-  out.restoreState(cut(cache.keys), cut(cache.values), S - leftPad);
-  return out;
-}
-
-/** 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.
- *  `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[], trimLeft = 0,
-): { keys: MlxArray; values: MlxArray } {
-  const idx = MlxArray.fromInt32(Int32Array.from(keep), [keep.length]);
-  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

@@ -1,212 +0,0 @@
-// Batched QUANTIZED KV for continuous batching (Phase 3.1 of
-// docs/design/unified-engine-frontier-plan.md) — 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 solo-
-// prefills with the SAME per-layer conversion the serial path runs
-// (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/parity/batched-kv-quant-parity.test.ts.
-
-import { MlxArray } from "../mlx/array";
-import * as ops from "../mlx/ops";
-import { QuantizedKVCache, type Mask } from "./gemma4-base";
-import { buildBatchedDecodeMask } from "./batched-mask";
-
-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,
-): 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);
-  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 minicpm5.ts LlamaAttention.forward), so an
- *  eagerly-built [B] position array is exact. */
-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);
-    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 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();
-  }
-}

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

@@ -1,144 +0,0 @@
-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 = []; }
-}