mlx-bun 0.4.0 → 1.0.0

package/src/model/batched-rotating.ts

@@ -1,338 +0,0 @@
-// Batched sliding-window (rotating) KV cache for `--batch N` serving — the
-// hard half of dynamic-B that BatchedDecodeMaskCache (full-attention) doesn't
-// cover. Faithful port of mlx-lm's BatchRotatingKVCache (models/cache.py):
-// a shared ring buffer whose rotation state (_idx / _offset / rotated) is
-// SCALAR across the batch (all rows advance one token per step, so the write
-// column is the same for every row), with PER-ROW `offset` (absolute position,
-// drives RoPE) and `leftPad` (padding columns in the buffer, masked out).
-//
-// Includes merge, one-token writes, block prefill, padding finalization, masks
-// and row extraction. Prefill preparation is a storage capability; the serving
-// scheduler's use of whole prompt cohorts is a separate concern.
-// Positions (offset/leftPad/_idx/_offset/
-// rotated) are tracked on the HOST (small deterministic ints) so make_mask is
-// built in JS like buildBatchedDecodeMask; only K/V live on device.
-//
-// The genuinely tricky bit is make_mask once the ring wraps: the buffer is no
-// longer in temporal order, so the causal+window+padding mask is built in
-// temporal coordinates then ROLLED to the ring's physical layout — see
-// buildBatchedRotatingMask (a column-by-column port of mlx-lm make_mask,
-// including the `roll(shift=idx+1)`). Gated model-free against mlx-lm
-// (tests/unit/batched-rotating.test.ts) + end-to-end vs a long-context Gemma oracle.
-
-import { MlxArray } from "../mlx/array";
-import * as ops from "../mlx/ops";
-import { RotatingKVCache, type Cache, type Mask, type PaddedPrefillCache, type PrefillPadding } from "./gemma4-base";
-import { BatchedRotatingState, type RotatingPositionSnapshot } from "./batched-rotating-state";
-import {
-  appendRotatingStorage,
-  mergeStorageRows,
-  plainRowStorage,
-  temporalStorageView,
-  rowRollIndices,
-} from "./batched-row-storage";
-
-const STEP = 256;
-
-/** Port of BatchRotatingKVCache.make_mask for the DECODE case. Returns a
- *  [B,1,N,S] bool mask (S == min(maxSize-1, offsetScalar) + N) for the step
- *  ABOUT to be written: causal AND sliding-window AND per-row left-padding,
- *  rolled to the ring's physical layout once wrapped.
- *
- *  Inputs are the PRE-write scalar ring state (`idx`=_idx, `offsetScalar`=_offset,
- *  `rotated`) and per-row `leftPad`. The leftPad adjustments here are LOCAL to
- *  the mask (mlx-lm decrements left_padding both in make_mask AND in update; the
- *  update's decrement is the persistent one, applied separately by the cache). */
-export function buildBatchedRotatingMask(
-  B: number, N: number, leftPad: number[],
-  maxSize: number, window: number, idx: number, offsetScalar: number, rotated: boolean, block = N > 1,
-): MlxArray {
-  const off = Math.min(maxSize - 1, offsetScalar);
-  const S = off + N;
-
-  // Local (non-persisted) leftPad for the mask: trim + rotation shrink it.
-  const trimSize = idx - maxSize + (block ? 1 : 0);
-  const isRot = !block && (rotated || idx >= maxSize);
-  const lp = leftPad.map((x) => x - (trimSize > 0 ? trimSize : 0) - (isRot ? 1 : 0));
-
-  // roll(shift): physical column = (temporal column + shift) mod S.
-  const shift = isRot ? (idx >= maxSize ? 0 : idx) + 1 : 0;
-
-  const data = new Int32Array(B * N * S);
-  for (let b = 0; b < B; b++) {
-    const pad = lp[b]!;
-    for (let i = 0; i < N; i++) {
-      const lind = off + i; // query temporal position
-      for (let j = 0; j < S; j++) {
-        // temporal key position j: causal, within window, past this row's pad.
-        const allow = lind >= j && lind < j + window && j >= pad;
-        if (!allow) continue;
-        const col = ((j + shift) % S + S) % S;
-        data[(b * N + i) * S + col] = 1;
-      }
-    }
-  }
-  const intArr = MlxArray.fromInt32(data, [B, 1, N, S]);
-  const zero = MlxArray.fromInt32(new Int32Array([0]), []);
-  const mask = ops.less(zero, intArr); // 0 < x → bool
-  zero.dispose();
-  intArr.dispose();
-  return mask;
-}
-
-/** Port of mlx-lm BatchRotatingKVCache; padded singleton chunks use its
- * concat operation with the matching block mask until finalization. */
-export class BatchedRotatingCache implements Cache, PaddedPrefillCache {
-  keys: MlxArray | null = null;
-  values: MlxArray | null = null;
-  readonly #rows: BatchedRotatingState;
-  #ropeArr: MlxArray | null = null;
-  #ropeForOffset = -1;
-  readonly maxSize: number;
-
-  constructor(maxSize: number, leftPad: number[]) {
-    this.maxSize = maxSize;
-    this.#rows = new BatchedRotatingState(maxSize, leftPad);
-  }
-
-  /** Same signature as the serial RotatingKVCache — the scheduler's merge
-   *  guards recognize a batched ring as "rotating-plain + RowBatchCache"
-   *  (batch-scheduler #mergeJoiner prevRot branch); without this the next
-   *  join sees "unknown", skips the prevRot rows, and replaces the batch's
-   *  sliding-layer KV with a one-row cache (B≥3 collapse). Guard order
-   *  everywhere is isRowBatchCache FIRST, so the shared string never
-   *  misroutes a batched cache into a serial-only path. */
-  signature(): string { return "kv:rotating-plain"; }
-
-  get positionSnapshot(): RotatingPositionSnapshot { return this.#rows.snapshot(); }
-  validOffset(row: number): number { return this.#rows.validOffset(row); }
-  preparePrefill(padding: PrefillPadding, rightPaddedBatch?: boolean): void { this.#rows.preparePrefill(padding, rightPaddedBatch); this.releaseRopeArr(); }
-  finalizePrefill(): void {
-    const shifts = this.#rows.prefillRoll();
-    if (shifts?.some(shift => shift !== 0) && this.keys && this.values) {
-      using indices = rowRollIndices(this.keys.shape[2]!, shifts);
-      const keys = plainRowStorage.rollRows(this.keys, indices);
-      let values: MlxArray;
-      try { values = plainRowStorage.rollRows(this.values, indices); }
-      catch (error) { keys.dispose(); throw error; }
-      this.keys.dispose(); this.values.dispose(); this.keys = keys; this.values = values;
-    }
-    this.#rows.finalizePrefill(); this.releaseRopeArr();
-  }
-  /** Adopt owned tensors without changing their physical ring columns. */
-  static adoptPhysical(keys: MlxArray | null, values: MlxArray | null, position: RotatingPositionSnapshot): BatchedRotatingCache {
-    const cache = new BatchedRotatingCache(position.maxSize, [...position.leftPad]);
-    cache.#rows.restore(position); cache.keys = keys; cache.values = values;
-
-    return cache;
-  }
-
-  get offsetArr(): number[] { return this.#rows.offsets; }
-  get leftPad(): number[] { return this.#rows.leftPad; }
-  get batchSize(): number { return this.#rows.batchSize; }
-
-  /** Current batch size — tracks filter() (which shrinks the per-row arrays). */
-  get #B(): number {
-    return this.#rows.batchSize;
-  }
-
-  /** mlx-lm uses `cache.offset` (the per-row array) for the scalar interface
-   *  too; we expose the scalar total as `offset` and the per-row positions via
-   *  `ropeOffsetArr` (the model's per-row RoPE path). */
-  get offset(): number {
-    return this.#rows.totalOffset;
-  }
-
-  get ropeOffsetArr(): MlxArray {
-    if (this.#ropeArr && this.#ropeForOffset === this.offset) return this.#ropeArr;
-    this.#ropeArr?.dispose();
-    this.#ropeArr = MlxArray.fromInt32(Int32Array.from(this.offsetArr), [this.#B]);
-    this.#ropeForOffset = this.offset;
-    return this.#ropeArr;
-  }
-
-  makeMask(N: number, windowSize: number | null): Mask {
-    const window = windowSize ?? this.maxSize;
-    return {
-      mode: "array",
-      arr: buildBatchedRotatingMask(
-        this.#B, N, this.leftPad, this.maxSize, window,
-        this.#rows.ringIndex, this.offset, this.#rows.rotated, N > 1 || this.#rows.hasPendingPadding,
-      ),
-    };
-  }
-
-  /** N=1 decode update — port of _update_in_place. */
-  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
-    const [B, H, S, D] = k.shape as [number, number, number, number];
-    const vD = v.shape[3]!;
-    if (S !== 1 || this.#rows.hasPendingPadding) return this.#updateConcat(k, v);
-    const prev = this.offset;
-
-    // Grow the buffer (in STEP chunks) until it reaches maxSize.
-    if (!this.keys || (prev >= this.keys.shape[2]! && this.keys.shape[2]! < this.maxSize)) {
-      const newSize = Math.min(STEP, this.maxSize - prev);
-      const newK = ops.zeros([B, H, newSize, D], k.dtype);
-      const newV = ops.zeros([B, H, newSize, vD], v.dtype);
-      if (this.keys && this.values) {
-        const ck = ops.concatAxis([this.keys, newK], 2);
-        const cv = ops.concatAxis([this.values, newV], 2);
-        for (const a of [this.keys, this.values, newK, newV]) a.dispose();
-        this.keys = ck;
-        this.values = cv;
-      } else {
-        this.keys = newK;
-        this.values = newV;
-      }
-      this.#rows.markGrown(prev);
-    }
-
-    // Trim any overshoot past maxSize (decrements left padding persistently).
-    const trimSize = this.keys!.shape[2]! - this.maxSize;
-    if (trimSize > 0) {
-      const tk = this.keys!.slice([0, 0, trimSize, 0], this.keys!.shape as number[]);
-      const tv = this.values!.slice([0, 0, trimSize, 0], this.values!.shape as number[]);
-      this.keys!.dispose();
-      this.values!.dispose();
-      this.keys = tk;
-      this.values = tv;
-      this.#rows.trimOvershoot(trimSize);
-    }
-
-    // Rotate when the write head reaches the end of the ring.
-    const writeIndex = this.#rows.beginWrite(S);
-
-    // Write the new K/V at the ring head.
-    const [, , SK, DK] = this.keys!.shape as [number, number, number, number];
-    const k2 = ops.sliceUpdate(this.keys!, k, [0, 0, writeIndex, 0], [B, H, writeIndex + S, DK]);
-    const v2 = ops.sliceUpdate(this.values!, v, [0, 0, writeIndex, 0], [B, H, writeIndex + S, vD]);
-    this.keys!.dispose();
-    this.values!.dispose();
-    this.keys = k2;
-    this.values = v2;
-    this.#rows.commitWrite(S);
-
-    // Return the populated prefix (ring not yet full) or the whole buffer.
-    if (this.offset < this.maxSize) {
-      return [
-        this.keys.slice([0, 0, 0, 0], [B, H, this.offset, DK]),
-        this.values.slice([0, 0, 0, 0], [B, H, this.offset, vD]),
-      ];
-    }
-    return [
-      this.keys.slice([0, 0, 0, 0], [B, H, SK, DK]),
-      this.values.slice([0, 0, 0, 0], [B, H, SK, vD]),
-    ];
-  }
-
-  #updateConcat(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
-    const keys = appendRotatingStorage(plainRowStorage, this.keys, k, this.#rows);
-    let values: MlxArray;
-    try { values = appendRotatingStorage(plainRowStorage, this.values, v, this.#rows); }
-    catch (error) { keys.dispose(); throw error; }
-    const history = this.keys ? this.#rows.activeLength : 0;
-    this.keys?.dispose(); this.values?.dispose();
-    this.keys = keys; this.values = values;
-    this.#rows.commitConcat(k.shape[2]!, history);
-    return [
-      plainRowStorage.slice(keys, 0, this.#B, 0, this.#rows.ringIndex),
-      plainRowStorage.slice(values, 0, this.#B, 0, this.#rows.ringIndex),
-    ];
-  }
-
-  /** Free the per-step RoPE array without disposing KV (wrapper-rebuild path). */
-  releaseRopeArr(): void {
-    this.#ropeArr?.dispose();
-    this.#ropeArr = null;
-    this.#ropeForOffset = -1;
-  }
-
-  /** Ring contents in temporal order, cut to the valid length (extract). */
-  temporalView(): [MlxArray, MlxArray] {
-    if (!this.keys || !this.values) throw new Error("cache is empty");
-    return [
-      temporalStorageView(plainRowStorage, this.keys, this.#rows, {
-        from: Math.max(0, this.#rows.activeLength - this.maxSize), to: this.#rows.activeLength,
-      }),
-      temporalStorageView(plainRowStorage, this.values, this.#rows, {
-        from: Math.max(0, this.#rows.activeLength - this.maxSize), to: this.#rows.activeLength,
-      }),
-    ];
-  }
-
-  /** mlx-lm `BatchRotatingKVCache.extract` (models/cache.py:1417): row `i`
-   *  as a fresh SERIAL RotatingKVCache — de-rolled to temporal order (the
-   *  oracle's roll(-_idx) when rotated), left padding stripped
-   *  (`max(0, left_padding[idx])` — post-wrap pads go negative), OWNED
-   *  contiguous copies. `offset` = the row's absolute position (may exceed
-   *  the buffer: wrapped rings have evicted tokens); ring idx = the new
-   *  buffer length — temporal order ⇔ write head at the end, exactly the
-   *  oracle's `cache._idx = cache.keys.shape[2]`. Bit-exact vs a solo run:
-   *  merge/decode/filter keep each row's ring bytes identical to the serial
-   *  cache's (tests/batched-rotating) and this is a pure slice+copy. */
-  extractRow(i: number, limit?: number): RotatingKVCache | null {
-    if (!this.keys || !this.values) return null;
-    const pad = Math.max(0, this.leftPad[i]!, limit === undefined ? 0 : this.#rows.activeLength - limit);
-    const c = new RotatingKVCache(this.maxSize);
-    const k = temporalStorageView(plainRowStorage, this.keys, this.#rows, {
-      row: i, from: pad, to: this.#rows.activeLength, copy: true,
-    });
-    const v = temporalStorageView(plainRowStorage, this.values, this.#rows, {
-      row: i, from: pad, to: this.#rows.activeLength, copy: true,
-    });
-    c.restoreState(k, v, this.offsetArr[i]!, k.shape[2]!);
-    return c;
-  }
-
-  /** Keep only `keep` rows along the batch axis (eviction). */
-  filter(keep: number[]): void {
-    if (this.keys && this.values) {
-      const k = plainRowStorage.takeRows(this.keys, keep);
-      const v = plainRowStorage.takeRows(this.values, keep);
-      this.keys.dispose();
-      this.values.dispose();
-      this.keys = k;
-      this.values = v;
-    }
-    this.#rows.filter(keep);
-    this.releaseRopeArr();
-  }
-
-  filterRows(keep: readonly number[]): void { this.filter([...keep]); }
-
-  state(): MlxArray[] {
-    return this.keys && this.values ? [this.keys, this.values] : [];
-  }
-  isTrimmable(): boolean {
-    return this.#rows.trimmable;
-  }
-  trim(n: number): void {
-    this.#rows.trim(n);
-  }
-  dispose(): void {
-    this.keys?.dispose();
-    this.values?.dispose();
-    this.keys = this.values = null;
-    this.#ropeArr?.dispose();
-    this.#ropeArr = null;
-  }
-
-  /** Assemble a batch from per-row temporal KV slices (port of merge). Each row
-   *  is its solo cache's temporalView ([1,H,Li,D], Li ≤ maxSize, temporal
-   *  order) and `offsets[i]` its absolute position. Left-pads to the longest
-   *  row; the result is in temporal order (rotated=false, idx=offset=width). */
-  static merge(
-    rows: { keys: MlxArray; values: MlxArray }[], offsets: number[], maxSize: number,
-  ): BatchedRotatingCache {
-    const lens = rows.map((r) => r.keys.shape[2]!);
-    const width = Math.max(...lens, 0);
-    const leftPad = lens.map((l) => width - l);
-    const cache = new BatchedRotatingCache(maxSize, leftPad);
-    cache.#rows.restoreMerged(width, offsets);
-    if (width === 0) return cache;
-
-    cache.keys = mergeStorageRows(plainRowStorage, rows.map((row) => row.keys), leftPad);
-    cache.values = mergeStorageRows(plainRowStorage, rows.map((row) => row.values), leftPad);
-    return cache;
-  }
-}

package/src/model/batched-row-storage.ts

@@ -1,179 +0,0 @@
-import { MlxArray } from "../mlx/array";
-import * as ops from "../mlx/ops";
-import { BatchedRotatingState } from "./batched-rotating-state";
-import { Dtype } from "../mlx/ffi";
-
-export interface RowStorage<T> {
-  shape(value: T): readonly [number, number, number, number];
-  slice(value: T, batchFrom: number, batchTo: number, tokenFrom: number, tokenTo: number): T;
-  concatTokens(values: readonly T[]): T;
-  concatRows(values: readonly T[]): T;
-  padLeft(value: T, tokens: number): T;
-  takeRows(value: T, keep: readonly number[]): T;
-  copy(value: T): T;
-  rollRows(value: T, indices: MlxArray): T;
-  dispose(value: T): void;
-}
-
-export const plainRowStorage: RowStorage<MlxArray> = {
-  shape: (value) => value.shape as [number, number, number, number],
-  slice(value, batchFrom, batchTo, tokenFrom, tokenTo) {
-    const [, H, , D] = value.shape as [number, number, number, number];
-    return value.slice([batchFrom, 0, tokenFrom, 0], [batchTo, H, tokenTo, D]);
-  },
-  concatTokens: (values) => ops.concatAxis([...values], 2),
-  concatRows: (values) => ops.concatAxis([...values], 0),
-  padLeft(value, tokens) {
-    if (tokens === 0) return value.slice([0, 0, 0, 0], value.shape as number[]);
-    const [B, H, , D] = value.shape as [number, number, number, number];
-    const zeros = ops.zeros([B, H, tokens, D], value.dtype);
-    const out = ops.concatAxis([zeros, value], 2);
-    zeros.dispose();
-    return out;
-  },
-  takeRows(value, keep) {
-    const indices = MlxArray.fromInt32(Int32Array.from(keep), [keep.length]);
-    const out = ops.takeAxis(value, indices, 0);
-    indices.dispose();
-    return out;
-  },
-  copy: (value) => ops.copyOf(value),
-  rollRows: (value, indices) => ops.takeAlongAxis(value, indices, 2),
-  dispose: (value) => value.dispose(),
-};
-
-export const quantizedRowStorage: RowStorage<ops.QuantizedTensor> = {
-  shape: (value) => value.packed.shape as [number, number, number, number],
-  slice(value, batchFrom, batchTo, tokenFrom, tokenTo) {
-    const cut = (array: MlxArray): MlxArray => {
-      const [, H, , D] = array.shape as [number, number, number, number];
-      return array.slice([batchFrom, 0, tokenFrom, 0], [batchTo, H, tokenTo, D]);
-    };
-    return { packed: cut(value.packed), scales: cut(value.scales), biases: cut(value.biases) };
-  },
-  concatTokens(values) {
-    return {
-      packed: ops.concatAxis(values.map((value) => value.packed), 2),
-      scales: ops.concatAxis(values.map((value) => value.scales), 2),
-      biases: ops.concatAxis(values.map((value) => value.biases), 2),
-    };
-  },
-  concatRows(values) {
-    return {
-      packed: ops.concatAxis(values.map((value) => value.packed), 0),
-      scales: ops.concatAxis(values.map((value) => value.scales), 0),
-      biases: ops.concatAxis(values.map((value) => value.biases), 0),
-    };
-  },
-  padLeft(value, tokens) {
-    const pad = (array: MlxArray): MlxArray => plainRowStorage.padLeft(array, tokens);
-    return { packed: pad(value.packed), scales: pad(value.scales), biases: pad(value.biases) };
-  },
-  takeRows(value, keep) {
-    const take = (array: MlxArray): MlxArray => plainRowStorage.takeRows(array, keep);
-    return { packed: take(value.packed), scales: take(value.scales), biases: take(value.biases) };
-  },
-  copy(value) {
-    return {
-      packed: ops.copyOf(value.packed),
-      scales: ops.copyOf(value.scales),
-      biases: ops.copyOf(value.biases),
-    };
-  },
-  rollRows(value, indices) {
-    const packed = ops.takeAlongAxis(value.packed, indices, 2);
-    let scales: MlxArray | undefined;
-    try {
-      scales = ops.takeAlongAxis(value.scales, indices, 2);
-      return { packed, scales, biases: ops.takeAlongAxis(value.biases, indices, 2) };
-    } catch (error) { packed.dispose(); scales?.dispose(); throw error; }
-  },
-  dispose(value) {
-    value.packed.dispose();
-    value.scales.dispose();
-    value.biases.dispose();
-  },
-};
-
-/** Per-row cyclic shift, shared by plain and encoded storage planes.
- * This is mlx-lm's dynamic_roll indexing; only positions cross the FFI. */
-export function rowRollIndices(length: number, shifts: readonly number[]): MlxArray {
-  using columns = ops.arange(0, length, 1, Dtype.int32);
-  using positions = ops.reshape(columns, [1, 1, length, 1]);
-  using starts = ops.fromInt32(shifts.map(shift => ((-shift % length) + length) % length), [shifts.length, 1, 1, 1]);
-  using raw = ops.add(positions, starts);
-  using width = ops.fromInt32([length], []);
-  using wrapped = ops.sub(raw, width);
-  using over = ops.greaterEqual(raw, width);
-  return ops.where(over, wrapped, raw);
-}
-
-export function mergeStorageRows<T>(
-  storage: RowStorage<T>,
-  rows: readonly T[],
-  leftPad: readonly number[],
-): T {
-  const padded = rows.map((row, index) => storage.padLeft(row, leftPad[index]!));
-  try {
-    return storage.concatRows(padded);
-  } finally {
-    for (const row of padded) storage.dispose(row);
-  }
-}
-
-/** De-roll storage into temporal order, then cut one row or the whole batch. */
-export function temporalStorageView<T>(
-  storage: RowStorage<T>,
-  value: T,
-  state: BatchedRotatingState,
-  options: { row?: number; from?: number; to?: number; copy?: boolean } = {},
-): T {
-  const [batch, , tokens] = storage.shape(value);
-  const batchFrom = options.row ?? 0;
-  const batchTo = options.row === undefined ? batch : options.row + 1;
-  const parts = state.temporalRanges(tokens).map(([from, to]) =>
-    storage.slice(value, batchFrom, batchTo, from, to)
-  );
-  const ordered = parts.length === 1 ? parts[0]! : storage.concatTokens(parts);
-  if (parts.length > 1) for (const part of parts) storage.dispose(part);
-  const cut = storage.slice(
-    ordered,
-    0,
-    batchTo - batchFrom,
-    options.from ?? 0,
-    options.to ?? state.validLength,
-  );
-  storage.dispose(ordered);
-  if (!options.copy) return cut;
-  const owned = storage.copy(cut);
-  storage.dispose(cut);
-  return owned;
-}
-
-/** Borrow storage, return an owned chronological block without decoding the
- * representation. Port of mlx-lm's BatchRotatingKVCache._update_concat
- * (MIT): discard spare capacity, retain the full window for the first query. */
-export function appendRotatingStorage<T>(
-  storage: RowStorage<T>, previous: T | null, incoming: T, state: BatchedRotatingState,
-): T {
-  if (previous === null) return storage.slice(incoming, 0, storage.shape(incoming)[0], 0, storage.shape(incoming)[2]);
-  const active = state.activeLength;
-  const shifts = state.prefillRoll();
-  if (shifts?.some(shift => shift !== 0)) {
-    const chronological = temporalStorageView(storage, previous, state, { to: active });
-    let rolled: T | undefined, history: T | undefined;
-    try {
-      using indices = rowRollIndices(active, shifts);
-      rolled = storage.rollRows(chronological, indices);
-      history = storage.slice(rolled, 0, state.batchSize, Math.max(0, active - state.maxSize + 1), active);
-      return storage.concatTokens([history, incoming]);
-    } finally {
-      storage.dispose(chronological); if (rolled) storage.dispose(rolled); if (history) storage.dispose(history);
-    }
-  }
-  const history = temporalStorageView(storage, previous, state, {
-    from: Math.max(0, active - state.maxSize + 1), to: active,
-  });
-  try { return storage.concatTokens([history, incoming]); }
-  finally { storage.dispose(history); }
-}

package/src/model/batched-ssm.ts

@@ -1,52 +0,0 @@
-import { SsmPrefillPadding } from "./ssm-prefill-padding";
-import type { PaddedPrefillCache, PrefillPadding } from "./gemma4-base";
-import * as ops from "../mlx/ops";
-import type { MlxArray } from "../mlx/array";
-import type { BatchableCache, Cache } from "./gemma4-base";
-import { SSMCache } from "./qwen3-delta";
-
-/** Recurrent storage behind the same row-layout port as attention KV.
- * Inputs to merge are borrowed; request membership is owned by the caller. */
-export class BatchedSSMCache extends SSMCache implements BatchableCache, PaddedPrefillCache {
-  constructor() { super(); this.offsets = []; }
-  preparePrefill(padding: PrefillPadding): void {
-    if (!this.offsets!.length) this.offsets = padding.lengths.map(() => 0);
-    const padded = padding.leftPadding?.some(n => n > 0) || padding.rightPadding?.some(n => n > 0);
-    this.prefillPadding = padded ? new SsmPrefillPadding(padding) : null;
-  }
-  finalizePrefill(): void { this.prefillPadding = null; }
-  get rowOffsets(): readonly number[] { return this.offsets!; }
-  get leftPad(): readonly number[] { return this.rowOffsets.map(() => 0); }
-  makeEmptyBatch(): BatchedSSMCache { return new BatchedSSMCache(); }
-  projectedBytes(_tokens: number): number {
-    return this.batchSize ? this.state().reduce((bytes, array) => bytes + array.nbytes, 0) / this.batchSize : 0;
-  }
-
-  mergeRows(rows: readonly Cache[]): void {
-    if (!rows.length) { this.dispose(); return; }
-    const sources = rows as readonly SSMCache[];
-    const offsets = sources.flatMap(source => source.offsets ?? [source.offset]);
-    const merge = (field: "conv" | "recurrent"): MlxArray | null => {
-      const prototype = sources.find(source => source[field])?.[field];
-      if (!prototype) return null;
-      if (sources.length === 1) return ops.copyOf(prototype);
-      const empty: MlxArray[] = [];
-      try {
-        return ops.concatAxis(sources.map(source => {
-          if (source[field]) return source[field]!;
-          const zeros = ops.zeros([source.offsets?.length ?? 1, ...prototype.shape.slice(1)], prototype.dtype);
-          empty.push(zeros); return zeros;
-        }), 0);
-      } finally { for (const array of empty) array.dispose(); }
-    };
-    const conv = merge("conv");
-    let recurrent;
-    try { recurrent = merge("recurrent"); }
-    catch (error) { conv?.dispose(); throw error; }
-    this.dispose();
-    this.conv = conv; this.recurrent = recurrent;
-    this.offsets = offsets; this.offset = Math.max(0, ...offsets);
-  }
-
-  override dispose(): void { super.dispose(); this.offset = 0; this.offsets = []; }
-}

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

@@ -1,93 +0,0 @@
-import { decodedKvDonorAttention } from "./decoded-kv-donor";
-import type { MlxArray } from "../mlx/array";
-import { KvTensorRows, type KvTensorRowView } from "../backends/mlx/kv-tensor-rows";
-import { disposeResources } from "../engine/resources";
-import { TurboQuantKVCache, type BatchableCache, type Cache, type Mask, type RotatedValueAttentionState, type PaddedPrefillCache, type PrefillPadding } from "./gemma4-base";
-import { TurboQuantCodec, disposeTurboQuant, type TurboQuantTensor } from "./turboquant-codec";
-
-const fields = ["kIdx", "kScales", "kZeros", "vPacked", "vScales"] as const;
-function encoded(planes: readonly MlxArray[]): TurboQuantTensor {
-  return { kIdx: planes[0]!, kScales: planes[1]!, kZeros: planes[2]!, vPacked: planes[3]!, vScales: planes[4]! };
-}
-
-/** The existing TurboQuant codec over shared tensor-row storage. Row changes
- * copy encoded bytes without decoding, rotating or requantizing their values. */
-export class BatchedTurboQuantKVCache implements BatchableCache, PaddedPrefillCache {
-  readonly #storage = new KvTensorRows();
-  readonly #codec: TurboQuantCodec;
-  #reuseOffsets: number[] = [];
-  #headDim: number | null = null;
-  constructor(readonly kBits: number, readonly vBits: number) {
-    this.#codec = new TurboQuantCodec(kBits, vBits, process.env.MLX_BUN_TURBOQUANT_FUSED_DECODE === "1");
-  }
-  restorePrefillEnds(ends: readonly number[] | undefined): void { this.#storage.restorePrefillEnds(ends); }
-  preparePrefill(padding: PrefillPadding): void {
-    this.#storage.preparePrefill(padding);
-    if (!this.#reuseOffsets.length) this.#reuseOffsets = padding.lengths.map(() => 0);
-  }
-  finalizePrefill(): void { this.#storage.finalizePrefill(); }
-  captureDonorRows(): import("./gemma4-base").KvDonorRows {
-    if (!this.#storage.planes.length || this.#headDim === null) throw new Error("cache is empty");
-    const [keys, values] = this.#codec.decode(encoded(this.#storage.planes), this.offset, this.#headDim, false);
-    return { keys, values, ...this.#storage.captureDonorValidity() };
-  }
-  captureDonorAttention(): import("./gemma4-base").KvDonorAttention {
-    return decodedKvDonorAttention(this.captureDonorRows());
-  }
-  get rotatedValueAttention(): RotatedValueAttentionState { return this; }
-  signature(): string { return `kv:batched-turboquant:${this.kBits}:${this.vBits}`; }
-  get minimumReusableOffset(): number { return Math.max(0, ...this.#reuseOffsets); }
-  get headDim(): number | null { return this.#headDim; }
-  get rowOffsets(): readonly number[] { return this.#storage.rowOffsets; }
-  get leftPad(): readonly number[] { return this.#storage.leftPad; }
-  get batchSize(): number | null { return this.#storage.batchSize; }
-  get offset(): number { return this.#storage.offset; }
-  get ropeOffsetArr(): MlxArray | undefined { return this.#storage.ropeOffsetArr; }
-  makeEmptyBatch(): BatchedTurboQuantKVCache { return new BatchedTurboQuantKVCache(this.kBits, this.vBits); }
-  bytesPerToken(): number { return this.#storage.bytesPerToken(); }
-  projectedBytes(tokens: number): number { return this.bytesPerToken() * tokens; }
-  makeMask(tokens: number, window: number | null): Mask { return this.#storage.makeMask(tokens, window); }
-  isTrimmable(): boolean { return true; }
-  trim(count: number): void { this.#storage.trim(count); }
-  specRoundBegin(): void { this.#storage.specRoundBegin(); }
-  specRoundCommit(): void { this.#storage.specRoundCommit(); }
-  specRoundRollback(keep: number | readonly number[]): void { this.#storage.specRoundRollback(keep); }
-  #append(k: MlxArray, v: MlxArray, defer: boolean): [MlxArray, MlxArray] {
-    const packed = this.#codec.encode(k, v);
-    try {
-      this.#storage.append(fields.map(field => packed[field]));
-      this.#headDim = k.shape[3]!;
-      return this.#codec.decode(encoded(this.#storage.planes), this.offset, this.#headDim, defer);
-    } finally { disposeTurboQuant(packed); }
-  }
-  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.#append(k, v, false); }
-  updateAndFetchDeferredV(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.#append(k, v, true); }
-  mergeRows(rows: readonly Cache[]): void {
-    const views: KvTensorRowView[] = [], held: MlxArray[] = [];
-    const reuseOffsets: number[] = [];
-    let dim = this.#headDim;
-    try {
-      for (const row of rows) {
-        if (row instanceof BatchedTurboQuantKVCache) {
-          views.push(row.#storage); dim ??= row.headDim; reuseOffsets.push(...row.#reuseOffsets);
-        } else if (row instanceof TurboQuantKVCache) {
-          const planes = row.state(); held.push(...planes); reuseOffsets.push(row.minimumReusableOffset);
-          views.push({ planes, rowOffsets: [row.offset], leftPad: [0] }); dim ??= row.headDim;
-        } else throw new Error(`TurboQuant row layout cannot merge ${row.signature()}`);
-      }
-      this.#storage.mergeRows(views); this.#headDim = dim; this.#reuseOffsets = reuseOffsets;
-    } finally { disposeResources(held); }
-  }
-  alignRows(leftPad: readonly number[]): void { this.#storage.alignRows(leftPad); }
-  filterRows(keep: readonly number[]): void {
-    this.#storage.filterRows(keep); this.#reuseOffsets = keep.map(row => this.#reuseOffsets[row]!);
-  }
-  state(): MlxArray[] { return this.#storage.planes; }
-  extractRow(row: number): TurboQuantKVCache {
-    const result = new TurboQuantKVCache(this.kBits, this.vBits), count = this.rowOffsets[row]!;
-    if (count) result.restoreState(encoded(this.#storage.extractRow(row)), count, this.#headDim!);
-    result.minimumReusableOffset = this.#reuseOffsets[row] ?? 0;
-    return result;
-  }
-  dispose(): void { this.#storage.dispose(); this.#headDim = null; this.#reuseOffsets = []; }
-}