mlx-bun 0.0.13 → 0.3.0

package/src/mlx/turboquant-kv-decode.ts

@@ -0,0 +1,82 @@
+// Exact byte-packed TurboQuant K/V decoding. The caller owns output arrays.
+// Rotation and cache policy belong to the existing codec/cache layers.
+import { type MlxArray, gpuStream } from './array';
+import { Dtype, type MlxHandle } from './ffi';
+import { isShapelessTracing } from './compile';
+import { MetalKernel } from './metal-kernel';
+
+export type PackedKvArrays = readonly [
+  keyIndices: MlxArray, keyScales: MlxArray, keyZeros: MlxArray,
+  valueIndices: MlxArray, valueScales: MlxArray,
+];
+
+const source = `
+  uint col = thread_position_in_grid.x, pos = thread_position_in_grid.y;
+  uint bh = thread_position_in_grid.z;
+  if (col >= D || pos >= vp_shape[2] || bh >= vp_shape[0] * vp_shape[1]) return;
+  uint b = bh / vp_shape[1], h = bh % vp_shape[1];
+  uint ko = b * kp_strides[0] + h * kp_strides[1] + pos * kp_strides[2];
+  uint vpo = b * vp_strides[0] + h * vp_strides[1] + pos * vp_strides[2];
+  uint kso = b * ks_strides[0] + h * ks_strides[1] + pos * ks_strides[2] + (col / KG) * ks_strides[3];
+  uint kzo = b * kz_strides[0] + h * kz_strides[1] + pos * kz_strides[2] + (col / KG) * kz_strides[3];
+  uint vso = b * vs_strides[0] + h * vs_strides[1] + pos * vs_strides[2] + (col / VG) * vs_strides[3];
+  int key_code;
+  if (KBITS == 8) key_code = int(kp[ko + col * kp_strides[3]]);
+  else {
+    uint bit = col * KBITS, byte = bit / 8u, shift = bit % 8u;
+    uint word = uint(kp[ko + byte * kp_strides[3]]);
+    if (shift + KBITS > 8u) word |= uint(kp[ko + (byte + 1u) * kp_strides[3]]) << 8u;
+    key_code = int((word >> shift) & ((1u << KBITS) - 1u));
+  }
+  // Keep the incumbent f32 addition before multiplication.
+  float shifted = float(key_code) + float(kz[kzo]);
+  uint i = (bh * vp_shape[2] + pos) * D + col;
+  outK[i] = OutK(shifted * float(ks[kso]));
+  uint bit = col * VBITS, byte = bit / 8u, shift = bit % 8u;
+  uint word = uint(vp[vpo + byte * vp_strides[3]]);
+  if (shift + VBITS > 8u) word |= uint(vp[vpo + (byte + 1u) * vp_strides[3]]) << 8u;
+  uint code = (word >> shift) & ((1u << VBITS) - 1u);
+  outV[i] = OutV(centroids[code] * float(vs[vso]));
+`;
+
+let kernel: MetalKernel | undefined;
+function jointKernel(): MetalKernel {
+  return kernel ??= new MetalKernel({
+    name: 'turboquant_kv_decode',
+    inputNames: ['kp', 'ks', 'kz', 'vp', 'vs', 'centroids'],
+    outputNames: ['outK', 'outV'], source, ensureRowContiguous: false,
+  });
+}
+
+/** Return bf16 K and rotated V (f32 for eager inverse rotation, bf16 for
+ * deferred rotation). Null means the ordinary codec must handle this input.
+ * Input arrays and the shared centroid table remain borrowed. */
+export function tryJointDecodePackedKv(
+  inputs: PackedKvArrays, centroids: MlxArray, kBits: number, vBits: number,
+  headDim: number, deferV: boolean, stream: MlxHandle = gpuStream,
+): [MlxArray, MlxArray] | null {
+  if (stream !== gpuStream || isShapelessTracing() ||
+      ![64, 128, 256, 512].includes(headDim) ||
+      ![2, 4, 5, 8].includes(kBits) || ![2, 3, 4, 5, 8].includes(vBits)) return null;
+  const [kp, ks, kz, vp, vs] = inputs;
+  if (kp.dtype !== (kBits === 8 ? Dtype.int8 : Dtype.uint8) ||
+      vp.dtype !== Dtype.uint8 || centroids.dtype !== Dtype.float32 ||
+      centroids.ndim !== 1 || centroids.size !== (1 << vBits)) return null;
+  const shapes = inputs.map(a => a.shape);
+  if (shapes.some(shape => shape.length !== 4)) return null;
+  const [b, h, n] = shapes[3]!;
+  if (!b || !h || !n || b * h * n * headDim >= 2 ** 32 ||
+      shapes.some(shape => shape[0] !== b || shape[1] !== h || shape[2] !== n) ||
+      shapes[0]![3] !== headDim * kBits / 8 || shapes[3]![3] !== headDim * vBits / 8 ||
+      shapes[1]![3] !== shapes[2]![3]) return null;
+  const keyGroup = headDim / shapes[1]![3]!, valueGroup = headDim / shapes[4]![3]!;
+  if (![32, 64].includes(keyGroup) || ![32, 64].includes(valueGroup) ||
+      [ks, kz, vs].some(a => ![Dtype.float16, Dtype.bfloat16, Dtype.float32].includes(a.dtype))) return null;
+  const shape = [b, h, n, headDim], valueDtype = deferV ? Dtype.bfloat16 : Dtype.float32;
+  return jointKernel().apply([...inputs, centroids], {
+    templateDtypes: { OutK: Dtype.bfloat16, OutV: valueDtype },
+    templateInts: { D: headDim, KBITS: kBits, VBITS: vBits, KG: keyGroup, VG: valueGroup },
+    grid: [headDim, n, b * h], threadGroup: [Math.min(headDim, 256), Math.max(1, 256 / headDim), 1],
+    outputs: [{ shape, dtype: Dtype.bfloat16 }, { shape, dtype: valueDtype }], stream,
+  }) as [MlxArray, MlxArray];
+}

package/src/mlx/turboquant-kv-inverse.ts

@@ -0,0 +1,122 @@
+// Joint packed K/V decode with eager inverse rotation in the measured D256 regime.
+// Hadamard butterfly/layout follows MLX 0.31.2,
+// mlx/backend/metal/kernels/hadamard.h, Copyright © 2024 Apple Inc., MIT.
+import { type MlxArray, gpuStream } from "./array";
+import { Dtype, type MlxHandle } from "./ffi";
+import { isShapelessTracing } from "./compile";
+import { MetalKernel } from "./metal-kernel";
+import { type PackedKvArrays } from "./turboquant-kv-decode";
+
+const header = `
+template <short R> inline void tq_inverse_radix(thread float* x) {
+  constexpr short logR = __builtin_ctz(R);
+  short h = 1;
+  #pragma unroll
+  for (short s = 0; s < logR; s++) {
+    #pragma unroll
+    for (short i = 0; i < R / 2; i++) {
+      short k = i & (h - 1), j = ((i - k) << 1) + k;
+      float a = x[j], b = x[j + h];
+      x[j] = a + b; x[j + h] = a - b;
+    }
+    h <<= 1;
+  }
+}
+`;
+const source = `
+  constexpr short D = 256, R = 16, NT = 16, LOG_R = 4;
+  uint pos = thread_position_in_grid.y, bh = thread_position_in_grid.z;
+  short i = thread_position_in_grid.x;
+  bool valid = pos < vp_shape[2];
+  uint b = bh / vp_shape[1], hidx = bh % vp_shape[1];
+  uint ko = b * kp_strides[0] + hidx * kp_strides[1] + pos * kp_strides[2];
+  uint vo = b * vp_strides[0] + hidx * vp_strides[1] + pos * vp_strides[2];
+  uint kso = b * ks_strides[0] + hidx * ks_strides[1] + pos * ks_strides[2];
+  uint kzo = b * kz_strides[0] + hidx * kz_strides[1] + pos * kz_strides[2];
+  uint vso = b * vs_strides[0] + hidx * vs_strides[1] + pos * vs_strides[2];
+  uint row = (bh * vp_shape[2] + pos) * D;
+  threadgroup float storage[D * 16];
+  threadgroup float* buf = storage + thread_position_in_threadgroup.y * D;
+  #pragma unroll
+  for (short j = 0; j < R / 4; j++) {
+    short index = j * 4 * NT + i * 4;
+    #pragma unroll
+    for (short r = 0; r < 4; r++) {
+      uint col = index + r;
+      if (!valid) { buf[col] = 0.0f; continue; }
+      int key_code = int(kp[ko + col * kp_strides[3]]);
+      float shifted = float(key_code) + float(kz[kzo + (col / 32) * kz_strides[3]]);
+      outK[row + col] = OutK(shifted * float(ks[kso + (col / 32) * ks_strides[3]]));
+      uint bit = col * 3u, byte = bit / 8u, shift = bit % 8u;
+      uint word = uint(vp[vo + byte * vp_strides[3]]);
+      if (shift + 3u > 8u) word |= uint(vp[vo + (byte + 1u) * vp_strides[3]]) << 8u;
+      uint code = (word >> shift) & 7u;
+      buf[col] = centroids[code * centroids_strides[0]] * float(vs[vso + (col / 32) * vs_strides[3]]);
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  float x[R]; short h = 1;
+  #pragma unroll
+  for (short s = 0; s < 2; s++) {
+    short k = i & (h - 1), j = ((i - k) << LOG_R) + k;
+    #pragma unroll
+    for (short r = 0; r < R; r++) x[r] = buf[j + h * r];
+    tq_inverse_radix<R>(x);
+    #pragma unroll
+    for (short r = 0; r < R; r++) buf[j + h * r] = x[r];
+    h <<= LOG_R;
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+  }
+  #pragma unroll
+  for (short j = 0; j < R / 4; j++) {
+    short index = j * 4 * NT + i * 4;
+    #pragma unroll
+    for (short r = 0; r < 4; r++) {
+      uint col = index + r;
+      float normalized = buf[col] * 0.0625f;
+      if (valid) outV[row + col] = OutV(normalized * signs[col * signs_strides[0]]);
+    }
+  }
+`;
+
+let kernel: MetalKernel | undefined;
+function inverseKernel(): MetalKernel {
+  return kernel ??= new MetalKernel({
+    name: "turboquant_kv_inverse256",
+    inputNames: ["kp", "ks", "kz", "vp", "vs", "centroids", "signs"],
+    outputNames: ["outK", "outV"], header, source, ensureRowContiguous: false,
+  });
+}
+
+/** Borrow packed inputs and codec tables; return owned bf16 K/V or null.
+ * Only the measured eager D256/k8v3/B1/H4/fp16/group32 regime is eligible.
+ * Other regimes continue through the existing joint decoder or original codec.
+ * N affects dispatch size but is never a shader template parameter. */
+export function tryDecodePackedKvInverse256(
+  inputs: PackedKvArrays, centroids: MlxArray, signs: MlxArray,
+  kBits: number, vBits: number, headDim: number, s: MlxHandle = gpuStream,
+): [MlxArray, MlxArray] | null {
+  if (s !== gpuStream || isShapelessTracing() || headDim !== 256 || kBits !== 8 || vBits !== 3)
+    return null;
+  const [kp, ks, kz, vp, vs] = inputs;
+  const shape = kp.shape;
+  if (shape.length !== 4 || shape[0] !== 1 || shape[1] !== 4 || shape[3] !== 256)
+    return null;
+  const n = shape[2]!;
+  if (n < 8192 || n * 1024 >= 2 ** 32) return null;
+  const samePrefix = (a: MlxArray, width: number) => {
+    const x = a.shape;
+    return x.length === 4 && x[0] === 1 && x[1] === 4 && x[2] === n && x[3] === width;
+  };
+  if (kp.dtype !== Dtype.int8 || vp.dtype !== Dtype.uint8 || !samePrefix(vp, 96)) return null;
+  if (![ks, kz, vs].every(a => a.dtype === Dtype.float16 && samePrefix(a, 8))) return null;
+  if (centroids.dtype !== Dtype.float32 || centroids.shape.length !== 1 || centroids.shape[0] !== 8)
+    return null;
+  if (signs.dtype !== Dtype.float32 || signs.shape.length !== 1 || signs.shape[0] !== 256)
+    return null;
+  return inverseKernel().apply([...inputs, centroids, signs], {
+    templateDtypes: { OutK: Dtype.bfloat16, OutV: Dtype.bfloat16 },
+    grid: [16, Math.ceil(n / 16) * 16, 4], threadGroup: [16, 16, 1],
+    outputs: [{ shape, dtype: Dtype.bfloat16 }, { shape, dtype: Dtype.bfloat16 }], stream: s,
+  }) as [MlxArray, MlxArray];
+}

package/src/mlx/turboquant-ops.ts

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

package/src/mlx/turboquant-tables.ts

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

package/src/model/batched-mask.ts

@@ -2,11 +2,11 @@
 //
 // Companion to src/train/forward.ts's buildBatchedPadMask, which builds the
 // PREFILL mask ([B,1,L,L], offset 0, right-padded — the training case, already
-// parity-proven by tests/train-batch-e2e.test.ts). This module covers the
+// parity-proven by tests/parity/train-batch-e2e.test.ts). This module covers the
 // DECODE case for batched serving: left-padded rows sharing one growing
 // [B,H,S,D] KV buffer, where a step forwards N new query tokens at a nonzero
 // offset. (The two builders should be consolidated when batched prefill is
-// wired into the serving path — phase S1a, docs/design/parallel-slots.md.)
+// wired into the serving path — phase S1a, docs/design/batching.md.)
 //
 // Why left-padding for decode: with rows right-aligned in the KV buffer, every
 // row's next write lands in the SAME column, so one advancing offset serves
@@ -15,6 +15,7 @@
 // thing the mask must encode is (a) causality, (b) each row's left padding,
 // and (c) the sliding window, if any.
 
+import { cacheSignature } from "./gemma4-base";
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
 import { createCausalMask, KVCache, type Cache, type Mask } from "./gemma4-base";
@@ -92,6 +93,11 @@ export class BatchedDecodeMaskCache implements Cache {
   #ropeArr: MlxArray | null = null;
   #ropeForOffset = -1;
 
+  /** A decorator, not a storage kind: never claims the inner cache's
+   *  identity (isPlainKvCache(wrapper) must stay false — it is not a
+   *  KVCache), but records what it wraps for diagnostics/persistence. */
+  signature(): string { return `kv:batched-mask:${cacheSignature(this.inner)}`; }
+
   constructor(
     private readonly inner: Cache,
     private readonly B: number,
@@ -242,7 +248,7 @@ export function extractKVRow(cache: KVCache, leftPad: number, i: number): KVCach
   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.contiguous(view);
+    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;
   };

package/src/model/batched-quant.ts

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

package/src/model/batched-rotating-quant.ts

@@ -22,9 +22,9 @@
 // every step; it self-tracks per-row offset/leftPad and the scalar ring
 // state. N=1 decode updates only (rows solo-prefill serially, then merge).
 //
-// Gates: tests/batched-rotating-quant.test.ts (model-free byte-identity vs
+// Gates: tests/unit/batched-rotating-quant.test.ts (model-free byte-identity vs
 // the serial oracle per row, through ring wrap, B=1 and B=2) + the gemma
-// scheduler gate in tests/batched-kv-quant-parity.test.ts.
+// scheduler gate in tests/parity/batched-kv-quant-parity.test.ts.
 
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
@@ -32,6 +32,12 @@ import { Dtype } from "../mlx/ffi";
 import { RotatingQuantizedKVCache, type Mask } from "./gemma4-base";
 import { buildBatchedRotatingMask } from "./batched-rotating";
 import type { QuantRow } from "./batched-quant";
+import { BatchedRotatingState } from "./batched-rotating-state";
+import {
+  mergeStorageRows,
+  quantizedRowStorage,
+  temporalStorageView,
+} from "./batched-row-storage";
 
 const STEP = 256; // shared with RotatingQuantizedKVCache.STEP and the bf16 twin
 
@@ -49,9 +55,7 @@ const mapTriple = (
  *  storage. Scalar ring state (`ringIdx`/`offset` reuse the base fields);
  *  per-row `offsetArr` (absolute positions → RoPE) and `leftPad`. */
 export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
-  offsetArr: number[];
-  leftPad: number[];
-  #rotated = false;
+  readonly #rows: BatchedRotatingState;
   /** Per-row RoPE positions. STABLE ACROSS A STEP — refreshed only at
    *  releaseRopeArr() (the scheduler's post-dispatch hook), never inside
    *  the update: the monolith CAPTURES it pre-update and ropes Q after
@@ -67,13 +71,16 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
     maxSize: number, groupSize: number, bits: number, leftPad: number[], offsets: number[],
   ) {
     super(maxSize, groupSize, bits);
-    this.leftPad = [...leftPad];
-    this.offsetArr = [...offsets];
+    this.#rows = new BatchedRotatingState(maxSize, leftPad, offsets);
     this.ropeOffsetArr = MlxArray.fromInt32(Int32Array.from(offsets), [offsets.length]);
   }
 
+  get offsetArr(): number[] { return this.#rows.offsets; }
+  get leftPad(): number[] { return this.#rows.leftPad; }
+  get batchSize(): number { return this.#rows.batchSize; }
+
   get #B(): number {
-    return this.leftPad.length;
+    return this.#rows.batchSize;
   }
 
   #seqLen2(): number {
@@ -104,7 +111,7 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
       mode: "array",
       arr: buildBatchedRotatingMask(
         this.#B, N, this.leftPad, this.maxSize, window,
-        this.ringIdx, this.offset, this.#rotated,
+        this.#rows.ringIndex, this.#rows.totalOffset, this.#rows.rotated,
       ),
     };
   }
@@ -117,7 +124,7 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
     const vD = v.shape[3]!;
     if (S !== 1)
       throw new Error("BatchedRotatingQuantCache supports N=1 decode updates only (solo-prefill then merge)");
-    const prev = this.offset;
+    const prev = this.#rows.totalOffset;
 
     // Grow the buffer (in STEP chunks) until it reaches maxSize.
     if (!this.keys || (prev >= this.#seqLen2() && this.#seqLen2() < this.maxSize)) {
@@ -144,7 +151,8 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
         this.keys = newK;
         this.values = newV;
       }
-      this.ringIdx = prev;
+      this.#rows.markGrown(prev);
+      this.ringIdx = this.#rows.ringIndex;
     }
 
     // Trim any overshoot past maxSize (decrements left padding persistently).
@@ -159,16 +167,12 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
         });
       this.keys = cutFrom(this.keys!);
       this.values = cutFrom(this.values!);
-      this.ringIdx = this.maxSize;
-      this.leftPad = this.leftPad.map((x) => x - trimSize);
+      this.#rows.trimOvershoot(trimSize);
+      this.ringIdx = this.#rows.ringIndex;
     }
 
     // Rotate when the write head reaches the end of the ring.
-    if (this.ringIdx === this.maxSize) {
-      this.#rotated = true;
-      this.ringIdx = 0;
-    }
-    if (this.#rotated) this.leftPad = this.leftPad.map((x) => x - S);
+    this.ringIdx = this.#rows.beginWrite(S);
 
     // Quantize + write the new K/V at the ring head (six sliceUpdates).
     const kq = ops.quantize(k, this.groupSize, this.bits);
@@ -184,9 +188,9 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
     disposeTriple(kq);
     disposeTriple(vq);
 
-    this.offset += S;
-    this.offsetArr = this.offsetArr.map((x) => x + S);
-    this.ringIdx += S;
+    this.#rows.commitWrite(S);
+    this.offset = this.#rows.totalOffset;
+    this.ringIdx = this.#rows.ringIndex;
 
     // Return the populated prefix (ring not yet full) or the whole buffer.
     const upTo = this.offset < this.maxSize ? this.offset : this.#seqLen2();
@@ -209,27 +213,10 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
    *  overrides the base to use the BATCHED ring state. */
   override temporalView(): [ops.QuantizedTensor, ops.QuantizedTensor] {
     if (!this.keys || !this.values) throw new Error("cache is empty");
-    const valid = Math.min(this.offset, this.maxSize);
-    const order = (t: ops.QuantizedTensor): ops.QuantizedTensor =>
-      mapTriple(t, (a) => {
-        const [B, H, Sbuf, D] = a.shape as [number, number, number, number];
-        let tt: MlxArray;
-        if (this.ringIdx === Sbuf) {
-          tt = a.slice([0, 0, 0, 0], [B, H, Sbuf, D]);
-        } else if (this.ringIdx < this.offset) {
-          const tail = a.slice([0, 0, this.ringIdx, 0], [B, H, Sbuf, D]);
-          const head = a.slice([0, 0, 0, 0], [B, H, this.ringIdx, D]);
-          tt = ops.concatAxis([tail, head], 2);
-          tail.dispose();
-          head.dispose();
-        } else {
-          tt = a.slice([0, 0, 0, 0], [B, H, this.ringIdx, D]);
-        }
-        const cut = tt.slice([0, 0, 0, 0], [B, H, valid, D]);
-        tt.dispose();
-        return cut;
-      });
-    return [order(this.keys), order(this.values)];
+    return [
+      temporalStorageView(quantizedRowStorage, this.keys, this.#rows),
+      temporalStorageView(quantizedRowStorage, this.values, this.#rows),
+    ];
   }
 
   /** Quantized twin of BatchedRotatingCache.extractRow (mlx-lm
@@ -244,34 +231,13 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
   extractRow(i: number): RotatingQuantizedKVCache | null {
     if (!this.keys || !this.values) return null;
     const pad = Math.max(0, this.leftPad[i]!);
-    const valid = Math.min(this.offset, this.maxSize);
-    const cut = (t: ops.QuantizedTensor): ops.QuantizedTensor =>
-      mapTriple(t, (a) => {
-        const [, H, Sbuf, D] = a.shape as [number, number, number, number];
-        const row = a.slice([i, 0, 0, 0], [i + 1, H, Sbuf, D]);
-        let tt: MlxArray;
-        if (this.ringIdx === Sbuf) {
-          tt = row;
-        } else if (this.ringIdx < this.offset) {
-          const tail = row.slice([0, 0, this.ringIdx, 0], [1, H, Sbuf, D]);
-          const head = row.slice([0, 0, 0, 0], [1, H, this.ringIdx, D]);
-          tt = ops.concatAxis([tail, head], 2);
-          tail.dispose();
-          head.dispose();
-          row.dispose();
-        } else {
-          tt = row.slice([0, 0, 0, 0], [1, H, this.ringIdx, D]);
-          row.dispose();
-        }
-        const cutV = tt.slice([0, 0, pad, 0], [1, H, valid, D]);
-        tt.dispose();
-        const own = ops.contiguous(cutV);
-        cutV.dispose();
-        return own;
-      });
     const c = new RotatingQuantizedKVCache(this.maxSize, this.groupSize, this.bits);
-    const k = cut(this.keys);
-    const v = cut(this.values);
+    const k = temporalStorageView(quantizedRowStorage, this.keys, this.#rows, {
+      row: i, from: pad, copy: true,
+    });
+    const v = temporalStorageView(quantizedRowStorage, this.values, this.#rows, {
+      row: i, from: pad, copy: true,
+    });
     c.restoreState(k, v, this.offsetArr[i]!, k.packed.shape[2]!);
     return c;
   }
@@ -279,22 +245,19 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
   /** Keep only `keep` rows along the batch axis (eviction), in place. */
   filter(keep: number[]): void {
     if (this.keys && this.values) {
-      const idxArr = MlxArray.fromInt32(Int32Array.from(keep), [keep.length]);
-      const take = (t: ops.QuantizedTensor): ops.QuantizedTensor =>
-        mapTriple(t, (a) => ops.takeAxis(a, idxArr, 0));
-      const k = take(this.keys);
-      const v = take(this.values);
-      idxArr.dispose();
+      const k = quantizedRowStorage.takeRows(this.keys, keep);
+      const v = quantizedRowStorage.takeRows(this.values, keep);
       disposeTriple(this.keys);
       disposeTriple(this.values);
       this.keys = k;
       this.values = v;
     }
-    this.offsetArr = keep.map((i) => this.offsetArr[i]!);
-    this.leftPad = keep.map((i) => this.leftPad[i]!);
+    this.#rows.filter(keep);
     this.releaseRopeArr(); // safe here: filter runs between steps
   }
 
+  filterRows(keep: readonly number[]): void { this.filter([...keep]); }
+
   override isTrimmable(): boolean {
     return false; // batched rows never re-enter the prompt cache directly
   }
@@ -317,30 +280,13 @@ export class BatchedRotatingQuantCache extends RotatingQuantizedKVCache {
     const width = Math.max(...lens, 0);
     const leftPad = lens.map((l) => width - l);
     const cache = new BatchedRotatingQuantCache(maxSize, groupSize, bits, leftPad, offsets);
+    cache.#rows.restoreMerged(width, offsets);
+    cache.offset = cache.#rows.totalOffset;
+    cache.ringIdx = cache.#rows.ringIndex;
     if (width === 0) return cache;
 
-    const padTo = (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;
-    };
-    const padded = rows.map((r, i) => ({
-      keys: mapTriple(r.keys, (a) => padTo(a, leftPad[i]!)),
-      values: mapTriple(r.values, (a) => padTo(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),
-    });
-    cache.keys = cat((r) => r.keys);
-    cache.values = cat((r) => r.values);
-    for (const r of padded) { disposeTriple(r.keys); disposeTriple(r.values); }
-    cache.ringIdx = width;
-    cache.offset = width;
+    cache.keys = mergeStorageRows(quantizedRowStorage, rows.map((row) => row.keys), leftPad);
+    cache.values = mergeStorageRows(quantizedRowStorage, rows.map((row) => row.values), leftPad);
     return cache;
   }
 }

package/src/model/batched-rotating-state.ts

@@ -0,0 +1,73 @@
+/** Host-side row and ring bookkeeping shared by plain and quantized storage. */
+export class BatchedRotatingState {
+  readonly maxSize: number;
+  offsets: number[];
+  leftPad: number[];
+  ringIndex = 0;
+  totalOffset = 0;
+  rotated = false;
+
+  constructor(maxSize: number, leftPad: readonly number[], offsets?: readonly number[]) {
+    this.maxSize = maxSize;
+    this.leftPad = [...leftPad];
+    this.offsets = offsets ? [...offsets] : leftPad.map((pad) => -pad);
+  }
+
+  get batchSize(): number { return this.leftPad.length; }
+  get validLength(): number { return Math.min(this.totalOffset, this.maxSize); }
+  get trimmable(): boolean { return this.totalOffset < this.maxSize; }
+
+  markGrown(previousOffset: number): void {
+    this.ringIndex = previousOffset;
+  }
+
+  trimOvershoot(tokens: number): void {
+    if (tokens <= 0) return;
+    this.ringIndex = this.maxSize;
+    this.leftPad = this.leftPad.map((pad) => pad - tokens);
+  }
+
+  /** Prepare an N-token write and return its physical ring column. */
+  beginWrite(tokens: number): number {
+    if (this.ringIndex === this.maxSize) {
+      this.rotated = true;
+      this.ringIndex = 0;
+    }
+    if (this.rotated) this.leftPad = this.leftPad.map((pad) => pad - tokens);
+    return this.ringIndex;
+  }
+
+  commitWrite(tokens: number): void {
+    this.totalOffset += tokens;
+    this.offsets = this.offsets.map((offset) => offset + tokens);
+    this.ringIndex += tokens;
+  }
+
+  filter(keep: readonly number[]): void {
+    this.offsets = keep.map((row) => this.offsets[row]!);
+    this.leftPad = keep.map((row) => this.leftPad[row]!);
+  }
+
+  trim(tokens: number): number {
+    const amount = Math.min(this.totalOffset, tokens);
+    this.totalOffset -= amount;
+    this.ringIndex -= amount;
+    this.offsets = this.offsets.map((offset) => offset - amount);
+    return amount;
+  }
+
+  restoreMerged(width: number, offsets: readonly number[]): void {
+    this.totalOffset = width;
+    this.ringIndex = width;
+    this.offsets = [...offsets];
+    this.rotated = false;
+  }
+
+  /** Physical ranges that expose a ring buffer in temporal order. */
+  temporalRanges(bufferLength: number): readonly [number, number][] {
+    if (this.ringIndex === bufferLength) return [[0, bufferLength]];
+    if (this.ringIndex < this.totalOffset)
+      return [[this.ringIndex, bufferLength], [0, this.ringIndex]];
+    return [[0, this.ringIndex]];
+  }
+}