mlx-bun 0.0.13 → 0.2.0

package/src/media-fetch.ts

@@ -21,6 +21,7 @@
 // runs model-free.
 
 import { lookup } from "node:dns/promises";
+import { runtimeValue } from "./runtime-config";
 
 export interface MediaFetchPolicy {
   /** Permit private/loopback/link-local destinations (LAN hosts). */
@@ -71,7 +72,7 @@ export interface RestrictedFetchResult {
 
 export function defaultMediaFetchPolicy(): MediaFetchPolicy {
   return {
-    allowPrivate: process.env.MLX_BUN_ALLOW_PRIVATE_MEDIA === "1",
+    allowPrivate: runtimeValue("MLX_BUN_ALLOW_PRIVATE_MEDIA") === "1",
     // 10 s / 64 MB: well above what a 30 s audio clip (~5 MB of 16 kHz
     // WAV, less compressed) or any supported image needs.
     timeoutMs: 10_000,
@@ -405,10 +406,19 @@ export async function fetchRestrictedHttpBytes(
  *  redirect hop re-validated, one wall-clock timeout, and a streaming
  *  size cap. Throws Error with a client-presentable message — callers
  *  surface it as a 400 (the prompt-build catch in server.ts). */
+/** Video clips are legitimately larger than images/audio: same SSRF guard
+ *  and timeout discipline, quadruple the body cap (a 30 s 1080p H.264 clip
+ *  runs tens of MB; frame sampling truncates long clips server-side). */
+export function videoMediaFetchPolicy(): MediaFetchPolicy {
+  return { ...defaultMediaFetchPolicy(), maxBytes: 256 * 1024 * 1024 };
+}
+
 export async function fetchMediaBytes(
   url: string,
-  kind: "image" | "audio",
-  policy: MediaFetchPolicy = defaultMediaFetchPolicy(),
+  kind: "image" | "audio" | "video",
+  policy: MediaFetchPolicy = kind === "video"
+    ? videoMediaFetchPolicy()
+    : defaultMediaFetchPolicy(),
 ): Promise<Uint8Array> {
   if (url.startsWith("data:")) {
     const comma = url.indexOf(",");
@@ -416,6 +426,12 @@ export async function fetchMediaBytes(
     const meta = url.slice(0, comma);
     const body = url.slice(comma + 1);
     if (!meta.includes("base64")) throw new Error("data: URL must be base64");
+    // The same body cap as the HTTP path — inline payloads must not become
+    // the uncapped route (base64 inflates 4/3, so check pre-decode too).
+    if (body.length * 0.75 > policy.maxBytes)
+      throw new Error(
+        `${kind} data: URL exceeds the ${Math.round(policy.maxBytes / 1024 / 1024)} MB cap`,
+      );
     return Uint8Array.from(Buffer.from(body, "base64"));
   }
   return (

package/src/memory/model.ts

@@ -30,6 +30,7 @@ import type { LoadedTokenizer } from "../tokenizer";
 import { BatchScheduler, type RowSampler } from "../serve/batch-scheduler";
 import * as ops from "../mlx/ops";
 import type { MlxArray } from "../mlx/array";
+import { runtimeValue } from "../runtime-config";
 
 const HF_HUB = `${process.env.HOME}/.cache/huggingface/hub`;
 const E4B_REPO = "models--mlx-community--gemma-4-e4b-it-OptiQ-4bit";
@@ -166,7 +167,7 @@ interface MemoryRuntime {
  *  Opt back in with MLX_BUN_MEMORY_BATCH=8 if a length-bucketed scheduler
  *  lands. (Decision: Josh, 2026-07-01.) */
 export function memoryBatchSize(): number {
-  const n = Number(process.env.MLX_BUN_MEMORY_BATCH ?? "1");
+  const n = Number(runtimeValue("MLX_BUN_MEMORY_BATCH") ?? "1");
   return Number.isFinite(n) && n >= 1 ? Math.floor(n) : 1;
 }
 

package/src/memory/vault.ts

@@ -19,12 +19,13 @@ import { homedir } from "node:os";
 import { isAbsolute, join } from "node:path";
 import { spawn } from "node:child_process";
 import { isFenceLine } from "./article";
+import { runtimeValue } from "../runtime-config";
 
 // ---- paths -----------------------------------------------------------
 
 /** Root of the memory vault: ~/.mlx-bun/wiki (override with MLX_BUN_WIKI). */
 export function vaultRoot(): string {
-  const override = process.env.MLX_BUN_WIKI?.trim();
+  const override = runtimeValue("MLX_BUN_WIKI")?.trim();
   if (override) return override.replace(/^~(?=$|[/\\])/, homedir());
   return join(homedir(), ".mlx-bun", "wiki");
 }

package/src/memory-plan.ts

@@ -0,0 +1,18 @@
+/** Common admission contract for every runtime. Representation-specific
+ * planners may expose richer fields, but serving consumes only this shape. */
+export interface MemoryPlan {
+  readonly schemaVersion: 1;
+  readonly strategy: "generic-kv" | "glm52-colibri";
+  readonly fits: boolean;
+  readonly contextTokens: number;
+  readonly maxSafeContext: number;
+  readonly weightsBytes: number;
+  readonly kvBytes: number;
+  readonly transientBytes: number;
+  readonly reserveBytes: number;
+  readonly totalBytes: number;
+  readonly usableBytes: number;
+  readonly predictedDecodeTps: number | null;
+  /** mlx allocator-cache limit, distinct from the process admission limit. */
+  readonly allocatorLimitBytes?: number;
+}

package/src/mlx/ffi.ts

@@ -14,6 +14,7 @@ import { dlopen, FFIType, ptr, read } from "bun:ffi";
 import { existsSync } from "node:fs";
 import { dirname, join } from "node:path";
 import { nativePackDir } from "../native-pack";
+import { runtimeValue } from "../runtime-config";
 
 const { ptr: P, i32, u64, f32, f64, cstring } = FFIType;
 
@@ -24,7 +25,7 @@ const { ptr: P, i32, u64, f32, f64, cstring } = FFIType;
  *  (src/native-pack.ts; populated on first run by the CLI) → homebrew
  *  (arm64, then Intel prefix). Keep in sync with nativeRuntimeDir(). */
 function resolveLibmlxc(): string {
-  const env = process.env.MLX_BUN_LIBMLXC;
+  const env = runtimeValue("MLX_BUN_LIBMLXC");
   if (env) return env;
   const beside = join(dirname(process.execPath), "libmlxc.dylib");
   if (existsSync(beside)) return beside;
@@ -99,6 +100,7 @@ export const C = dlopen(LIBMLXC_PATH, {
   mlx_fast_layer_norm: { args: [P, u64, u64, u64, f32, u64], returns: i32 },
   mlx_matmul: { args: [P, u64, u64, u64], returns: i32 },
   mlx_contiguous: { args: [P, u64, FFIType.bool, u64], returns: i32 },
+  mlx_copy: { args: [P, u64, u64], returns: i32 },
   // (res: mlx_vector_array*, w, opt group_size, opt bits, mode, global_scale, stream)
   mlx_quantize: { args: [P, u64, u64, u64, cstring, u64, u64], returns: i32 },
   mlx_vector_array_get: { args: [P, u64, u64], returns: i32 },
@@ -173,16 +175,11 @@ export const C = dlopen(LIBMLXC_PATH, {
   // (res, input, weight, stride, padding, dilation, groups, stream) — depthwise
   // causal conv1d for Qwen3.5 gated-DeltaNet (groups == channels, padding 0).
   mlx_conv1d: { args: [P, u64, u64, i32, i32, i32, i32, u64], returns: i32 },
-  // C signature: (res, input, weight, stride_0, stride_1, padding_0,
-  // padding_1, dilation_0, dilation_1, groups, stream) — NHWC conv2d for the
-  // gemma-4 audio SSCP subsampler.
-  // HAZARD (lab/repro/bun-ffi-stack-args): Apple arm64 packs sub-8-byte STACK
-  // args at natural size (dilation_1@sp+0, groups@sp+4, stream@sp+8) but
-  // bun:ffi 1.3.14 writes one 8-byte slot per arg — the callee reads shifted
-  // garbage and mlx segfaults on a bogus stream. Workaround: declare the two
-  // stack ints as ONE u64 (dilation_1 | groups<<32) so Bun's slot layout
-  // byte-matches the Apple ABI. ops.conv2d does the packing.
-  mlx_conv2d: { args: [P, u64, u64, i32, i32, i32, i32, i32, u64, u64], returns: i32 },
+  // Natural mlx-c signatures. Bun 1.4 fixed the macOS arm64 stack packing
+  // bug recorded in lab/repro/bun-ffi-stack-args, so these must match the
+  // header exactly. mlx-bun now requires Bun 1.4 or newer.
+  mlx_conv2d: { args: [P, u64, u64, i32, i32, i32, i32, i32, i32, i32, u64], returns: i32 },
+  mlx_conv3d: { args: [P, u64, u64, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, u64], returns: i32 },
   // --- training: autograd (value_and_grad) — proven in lab/spikes/phase-train-vag.ts.
   // mlx_value_and_grad(res: mlx_closure_value_and_grad*, fun: mlx_closure,
   //   argnums: int*, n) → builds a value+grad closure differentiating the

package/src/mlx/ops.ts

@@ -113,6 +113,15 @@ export function matmul(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
   return new MlxArray(outArray("matmul", (o) => C.mlx_matmul(o, a.handle, b.handle, s)));
 }
 
+/** Bit-identical copy through a FRESH primitive node (mx.copy). Unlike
+ *  `contiguous`, this always creates a new graph node even for an already-
+ *  contiguous input — the tool for DETACHING an array from a multi-output
+ *  primitive's sibling set (a retained sibling output pins every co-output's
+ *  buffer natively; see gatedDeltaUpdate's state detach, 2026-08-20). */
+export function copyOf(a: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("copy", (o) => C.mlx_copy(o, a.handle, s)));
+}
+
 /** Materialize a row-major contiguous copy (needed before rawBytes on views). */
 export function contiguous(a: MlxArray, s: S = gpuStream): MlxArray {
   return new MlxArray(
@@ -880,20 +889,36 @@ export function conv1d(
 
 /** nn.Conv2d: input [B, H, W, C_in] (NHWC), weight [C_out, kH, kW, C_in/groups]
  *  → [B, H', W', C_out]. mlx weight layout — the oracle transposes PyTorch's
- *  [C_out, C_in, kH, kW] at load. For the gemma-4 audio SSCP subsampler.
- *  dilation_1 and groups ride in one packed u64 — see the mlx_conv2d hazard
- *  note in ffi.ts (bun:ffi stack-arg ABI bug). */
+ *  [C_out, C_in, kH, kW] at load. For the gemma-4 audio SSCP subsampler. */
 export function conv2d(
   input: MlxArray, weight: MlxArray, stride: [number, number] = [1, 1],
   padding: [number, number] = [0, 0], dilation: [number, number] = [1, 1],
   groups = 1, s: S = gpuStream,
 ): MlxArray {
-  const dilation1AndGroups =
-    (BigInt(dilation[1] >>> 0) & 0xffffffffn) | (BigInt(groups >>> 0) << 32n);
   return new MlxArray(
     outArray("conv2d", (o) =>
       C.mlx_conv2d(o, input.handle, weight.handle, stride[0], stride[1],
-        padding[0], padding[1], dilation[0], dilation1AndGroups, s),
+        padding[0], padding[1], dilation[0], dilation[1], groups, s),
+    ),
+  );
+}
+
+/** nn.Conv3d: input [B, D, H, W, C_in] (NDHWC), weight [C_out, kD, kH, kW,
+ *  C_in/groups] → [B, D', H', W', C_out]. The Qwen3-VL vision patch embed
+ *  (kernel == stride) — the conv KERNEL's accumulation order is the parity-
+ *  relevant arithmetic; a value-equivalent GEMM differs by bf16 ulp. */
+export function conv3d(
+  input: MlxArray, weight: MlxArray,
+  stride: [number, number, number] = [1, 1, 1],
+  padding: [number, number, number] = [0, 0, 0],
+  dilation: [number, number, number] = [1, 1, 1],
+  groups = 1, s: S = gpuStream,
+): MlxArray {
+  return new MlxArray(
+    outArray("conv3d", (o) =>
+      C.mlx_conv3d(o, input.handle, weight.handle, stride[0], stride[1],
+        stride[2], padding[0], padding[1], padding[2], dilation[0],
+        dilation[1], dilation[2], groups, s),
     ),
   );
 }

package/src/model/batched-mask.ts

@@ -242,7 +242,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

@@ -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

@@ -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]];
+  }
+}