mlx-bun 0.0.9 → 0.0.10

package/src/kv-store.ts

@@ -1,23 +1,38 @@
 // KV-cache persistence: save prompt caches to disk, reload zero-copy.
+// The serialization core of the SSD cold tier (docs/design/ssd-kv-cold-tier.md).
 //
 // File layout (every tensor PAGE-ALIGNED — the Phase 1 corollary: files
 // we write can be mmap'd and handed to the GPU without copies):
-//   [magic "MLXBUNKV1\n"][u32 LE header length][JSON header][padding]
+//   [magic "MLXBUNKV2\n"][u32 LE header length][u32 LE dataStart]
+//   [u64 LE header hash][JSON header][padding]
 //   [tensor data at 16 KiB-aligned offsets]
-// Header: { tokens, caches: [{ kind, offset, idx?, maxSize?, dtype,
-//           kShape, vShape, kOffset, vOffset }] }
+// Header: { formatVersion, modelId, configFingerprint, ns, tokenizerHash,
+//           createdAt, tokens, caches: [{ kind, offset, idx?, maxSize?,
+//           groupSize?, bits?, tensors: [{ off, bytes, shape, dtype, hash }] }] }
+//
+// v2 over v1: quantized cache kinds (the serving DEFAULT is kv_config
+// quantization, which v1 could not persist), SSM kind (Qwen3.5 hybrid),
+// invalidation metadata (configFingerprint covers the kv-quant scheme;
+// tokenizerHash guards vocab drift; ns = adapter spec), per-tensor hashes
+// (verified opt-in — eager verification would defeat lazy page fault-in),
+// a header hash (always verified — cheap), and atomic tmp+fsync+rename
+// writes. v1 files are not migrated: nothing shipped wrote them (the
+// format was test/experiment-only) — they read as "not a v2 file" and the
+// SSD tier deletes-and-regenerates.
 //
 // Reload mmaps copy-on-write (MAP_PRIVATE): if mlx ever donates one of
 // these buffers, writes hit private pages, never the file.
 
-import { openSync, writeSync, closeSync } from "node:fs";
+import { openSync, writeSync, readSync, closeSync, fsyncSync, renameSync, rmSync } from "node:fs";
 import { MmapFile } from "./mmap";
 import { MlxArray } from "./mlx/array";
 import type { Dtype } from "./mlx/ffi";
 import * as ops from "./mlx/ops";
 import {
-  type Cache, KVCache, RotatingKVCache, Gemma4Model,
-} from "./model/gemma4";
+  type Cache, KVCache, RotatingKVCache,
+  QuantizedKVCache, RotatingQuantizedKVCache,
+} from "./model/gemma4-base";
+import { SSMCache } from "./model/qwen3-delta";
 
 /** Contiguous raw bytes of (possibly a view of) an array. */
 function contiguousBytes(a: MlxArray): Uint8Array {
@@ -27,120 +42,184 @@ function contiguousBytes(a: MlxArray): Uint8Array {
   return bytes;
 }
 
-const MAGIC = "MLXBUNKV1\n";
+const MAGIC = "MLXBUNKV2\n";
 const ALIGN = 16384;
+/** magic + u32 headerLen + u32 dataStart + u64 headerHash */
+const PREFIX_LEN = MAGIC.length + 4 + 4 + 8;
+
+const hash64 = (bytes: Uint8Array): string => Bun.hash(bytes).toString(16);
 
-interface CacheHeaderEntry {
-  kind: "kv" | "rotating";
+export type CacheKind = "kv" | "rotating" | "qkv" | "rotating-qkv" | "ssm";
+
+interface TensorSlot {
+  off: number;
+  bytes: number;
+  shape: number[];
+  dtype: number;
+  hash: string;
+}
+
+export interface CacheHeaderEntry {
+  kind: CacheKind;
   offset: number;
+  /** rotating variants: ring write index */
   idx?: number;
+  /** rotating variants: window size */
   maxSize?: number;
-  dtype: number;
-  kShape: number[];
-  vShape: number[];
-  kOffset: number;
-  vOffset: number;
-  kBytes: number;
-  vBytes: number;
+  /** quantized variants */
+  groupSize?: number;
+  bits?: number;
+  /** kv/rotating: [k, v] · qkv/rotating-qkv: [kPacked, kScales, kBiases,
+   *  vPacked, vScales, vBiases] · ssm: [conv, recurrent] */
+  tensors: TensorSlot[];
 }
 
-export interface KvFileHeader {
+export interface KvSaveMeta {
+  modelId?: string;
+  /** configFingerprint(config) — covers every graph-shaping field incl.
+   *  the kv-quant scheme, so a scheme flip invalidates naturally. */
+  configFingerprint?: string;
+  /** Adapter namespace (PromptCache ns — adapters joined with "+"). */
+  ns?: string;
+  /** sha256 of tokenizer.json: identical ids must mean identical text. */
+  tokenizerHash?: string;
+}
+
+export interface KvFileHeader extends KvSaveMeta {
+  formatVersion: 2;
+  createdAt: number;
   tokens: number[];
   caches: CacheHeaderEntry[];
 }
 
 const alignUp = (n: number) => Math.ceil(n / ALIGN) * ALIGN;
 
-export function saveKvCache(path: string, tokens: number[], caches: Cache[]): void {
-  const entries: CacheHeaderEntry[] = [];
+/** Snapshot one cache into header entry + blobs. Quantized full caches are
+ *  sliced to the live [.., :offset, :] region (token axis — quantization
+ *  groups run along the FEATURE dim, so a sequence-axis cut never splits a
+ *  group); rotating caches persist the whole ring as-laid-out with ringIdx. */
+function snapshotCache(c: Cache, dataOffset: number): { entry: CacheHeaderEntry; blobs: Uint8Array[]; next: number } {
+  const slots: TensorSlot[] = [];
   const blobs: Uint8Array[] = [];
+  let cursor = dataOffset;
+  const push = (a: MlxArray, disposeAfter: boolean): void => {
+    const bytes = contiguousBytes(a);
+    const off = alignUp(cursor);
+    slots.push({ off, bytes: bytes.length, shape: a.shape, dtype: a.dtype, hash: hash64(bytes) });
+    blobs.push(bytes);
+    cursor = off + bytes.length;
+    if (disposeAfter) a.dispose();
+  };
+  const liveSlice = (a: MlxArray, upTo: number): MlxArray => {
+    const [B, H, , D] = a.shape as [number, number, number, number];
+    return a.slice([0, 0, 0, 0], [B, H, upTo, D]);
+  };
+  const pushTriple = (t: ops.QuantizedTensor, upTo: number | null): void => {
+    for (const a of [t.packed, t.scales, t.biases]) {
+      if (upTo === null) push(a, false);
+      else push(liveSlice(a, upTo), true);
+    }
+  };
+
+  let entry: CacheHeaderEntry;
+  if (c instanceof RotatingQuantizedKVCache) {
+    if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
+    pushTriple(c.keys, null);
+    pushTriple(c.values, null);
+    entry = { kind: "rotating-qkv", offset: c.offset, idx: c.ringIdx, maxSize: c.maxSize,
+      groupSize: c.groupSize, bits: c.bits, tensors: slots };
+  } else if (c instanceof QuantizedKVCache) {
+    if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
+    pushTriple(c.keys, c.offset);
+    pushTriple(c.values, c.offset);
+    entry = { kind: "qkv", offset: c.offset, groupSize: c.groupSize, bits: c.bits, tensors: slots };
+  } else if (c instanceof RotatingKVCache) {
+    if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
+    push(c.keys, false);
+    push(c.values, false);
+    entry = { kind: "rotating", offset: c.offset, idx: c.ringIdx, maxSize: c.maxSize, tensors: slots };
+  } else if (c instanceof SSMCache) {
+    if (!c.conv || !c.recurrent) throw new Error("cannot persist an empty cache");
+    push(c.conv, false);
+    push(c.recurrent, false);
+    entry = { kind: "ssm", offset: c.offset, tensors: slots };
+  } else if (c instanceof KVCache) {
+    if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
+    push(liveSlice(c.keys, c.offset), true);
+    push(liveSlice(c.values, c.offset), true);
+    entry = { kind: "kv", offset: c.offset, tensors: slots };
+  } else {
+    throw new Error("unknown cache type");
+  }
+  return { entry, blobs, next: cursor };
+}
 
-  // compact each cache to its live region and snapshot bytes
+/** Persist `caches` (+ the exact token prefix they encode) to `path`.
+ *  ATOMIC: written to `<path>.tmp`, fsync'd, renamed — a crash mid-write
+ *  leaves only a .tmp orphan (ignored + reaped by SsdCacheStore.scan). */
+export function saveKvCache(path: string, tokens: number[], caches: Cache[], meta: KvSaveMeta = {}): void {
+  const entries: CacheHeaderEntry[] = [];
+  const blobs: Uint8Array[] = [];
   let dataOffset = 0; // relative; rebased after header is sized
   for (const c of caches) {
-    let keys: MlxArray;
-    let values: MlxArray;
-    let kind: "kv" | "rotating";
-    let idx: number | undefined;
-    let maxSize: number | undefined;
-    if (c instanceof RotatingKVCache) {
-      kind = "rotating";
-      idx = c.ringIdx;
-      maxSize = c.maxSize;
-      if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
-      keys = c.keys;
-      values = c.values;
-    } else if (c instanceof KVCache) {
-      kind = "kv";
-      if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
-      // store only the live [.., :offset, :] region
-      const [B, H, , D] = c.keys.shape as [number, number, number, number];
-      const vD = c.values.shape[3]!;
-      keys = c.keys.slice([0, 0, 0, 0], [B, H, c.offset, D]);
-      values = c.values.slice([0, 0, 0, 0], [B, H, c.offset, vD]);
-    } else {
-      throw new Error("unknown cache type");
-    }
-
-    const kBytesArr = contiguousBytes(keys);
-    const vBytesArr = contiguousBytes(values);
-    const kOffset = alignUp(dataOffset);
-    const vOffset = alignUp(kOffset + kBytesArr.length);
-    dataOffset = vOffset + vBytesArr.length;
-
-    entries.push({
-      kind, offset: c.offset, idx, maxSize,
-      dtype: keys.dtype,
-      kShape: keys.shape, vShape: values.shape,
-      kOffset, vOffset,
-      kBytes: kBytesArr.length, vBytes: vBytesArr.length,
-    });
-    blobs.push(kBytesArr, vBytesArr);
-    if (c instanceof KVCache) {
-      keys.dispose();
-      values.dispose();
-    }
+    const s = snapshotCache(c, dataOffset);
+    entries.push(s.entry);
+    blobs.push(...s.blobs);
+    dataOffset = s.next;
   }
 
-  const header: KvFileHeader = { tokens, caches: entries };
+  const header: KvFileHeader = {
+    formatVersion: 2, createdAt: Date.now(), ...meta, tokens, caches: entries,
+  };
   const headerJson = new TextEncoder().encode(JSON.stringify(header));
-  const dataStart = alignUp(MAGIC.length + 8 + headerJson.length);
+  const dataStart = alignUp(PREFIX_LEN + headerJson.length);
 
-  const fd = openSync(path, "w");
+  const tmp = `${path}.tmp`;
+  const fd = openSync(tmp, "w");
   try {
     const pre = new Uint8Array(dataStart);
     pre.set(new TextEncoder().encode(MAGIC), 0);
     const dv = new DataView(pre.buffer);
     dv.setUint32(MAGIC.length, headerJson.length, true);
     dv.setUint32(MAGIC.length + 4, dataStart, true);
-    pre.set(headerJson, MAGIC.length + 8);
+    dv.setBigUint64(MAGIC.length + 8, BigInt(Bun.hash(headerJson)), true);
+    pre.set(headerJson, PREFIX_LEN);
     writeSync(fd, pre, 0, pre.length, 0);
 
     let blobIdx = 0;
-    for (const e of entries) {
-      writeSync(fd, blobs[blobIdx]!, 0, e.kBytes, dataStart + e.kOffset);
-      writeSync(fd, blobs[blobIdx + 1]!, 0, e.vBytes, dataStart + e.vOffset);
-      blobIdx += 2;
-    }
-  } finally {
+    for (const e of entries)
+      for (const slot of e.tensors)
+        writeSync(fd, blobs[blobIdx++]!, 0, slot.bytes, dataStart + slot.off);
+    fsyncSync(fd);
+  } catch (err) {
     closeSync(fd);
+    try { rmSync(tmp, { force: true }); } catch {}
+    throw err;
   }
+  closeSync(fd);
+  renameSync(tmp, path);
 }
 
-/** Read only the header (cheap — for prefix matching across many files). */
+/** Read only the header (cheap — for prefix matching across many files).
+ *  Always verifies the header hash; throws on any structural mismatch. */
 export function readKvHeader(path: string): KvFileHeader & { dataStart: number } {
   const fd = openSync(path, "r");
   try {
-    const head = new Uint8Array(MAGIC.length + 8);
-    require("node:fs").readSync(fd, head, 0, head.length, 0);
+    const head = new Uint8Array(PREFIX_LEN);
+    readSync(fd, head, 0, head.length, 0);
     if (new TextDecoder().decode(head.subarray(0, MAGIC.length)) !== MAGIC)
-      throw new Error(`${path}: not an mlx-bun KV cache file`);
+      throw new Error(`${path}: not an mlx-bun KV v2 cache file`);
     const dv = new DataView(head.buffer);
     const len = dv.getUint32(MAGIC.length, true);
     const dataStart = dv.getUint32(MAGIC.length + 4, true);
+    const expectHash = dv.getBigUint64(MAGIC.length + 8, true);
     const body = new Uint8Array(len);
-    require("node:fs").readSync(fd, body, 0, len, MAGIC.length + 8);
+    readSync(fd, body, 0, len, PREFIX_LEN);
+    if (BigInt(Bun.hash(body)) !== expectHash)
+      throw new Error(`${path}: header hash mismatch (truncated or corrupt)`);
     const header = JSON.parse(new TextDecoder().decode(body)) as KvFileHeader;
+    if (header.formatVersion !== 2)
+      throw new Error(`${path}: unsupported formatVersion ${header.formatVersion}`);
     return { ...header, dataStart };
   } finally {
     closeSync(fd);
@@ -149,37 +228,100 @@ export function readKvHeader(path: string): KvFileHeader & { dataStart: number }
 
 export interface LoadedKvCache {
   tokens: number[];
+  header: KvFileHeader;
   caches: Cache[];
   /** Keep referenced as long as the caches live. */
   mmap: MmapFile;
 }
 
-export function loadKvCache(path: string, model: Gemma4Model): LoadedKvCache {
+export interface KvLoadExpect {
+  /** Reject on metadata mismatch (pass what the server is running). */
+  configFingerprint?: string;
+  tokenizerHash?: string;
+  ns?: string;
+  /** Verify every tensor hash (reads all bytes — defeats lazy fault-in;
+   *  off by default, `--ssd-cache-verify`). */
+  verify?: boolean;
+}
+
+/** Reload zero-copy. `model` is anything with makeCache() — the entry count
+ *  is validated against the DONOR cache list (model.layers.length was wrong
+ *  for KV-shared models like e4b, whose makeCache() returns donors only). */
+export function loadKvCache(
+  path: string,
+  model: { makeCache(): Cache[] },
+  expect: KvLoadExpect = {},
+): LoadedKvCache {
   const header = readKvHeader(path);
+  for (const key of ["configFingerprint", "tokenizerHash", "ns"] as const) {
+    if (expect[key] !== undefined && header[key] !== expect[key])
+      throw new Error(`${path}: ${key} mismatch (file ${header[key]}, expected ${expect[key]})`);
+  }
+  const proto = model.makeCache();
+  const cacheCount = proto.length;
+  for (const c of proto) c.dispose();
+  if (header.caches.length !== cacheCount)
+    throw new Error(`${path}: ${header.caches.length} cached layers but model has ${cacheCount}`);
+
   const mmap = MmapFile.open(path, "cow");
   const dataStart = header.dataStart;
+  const arr = (slot: TensorSlot): MlxArray => {
+    if (expect.verify) {
+      const view = mmap.view(dataStart + slot.off, slot.bytes);
+      if (hash64(view) !== slot.hash) {
+        mmap.unmap();
+        throw new Error(`${path}: tensor hash mismatch at offset ${slot.off}`);
+      }
+    }
+    return MlxArray.fromPointer(mmap.pointer(dataStart + slot.off), slot.shape, slot.dtype as Dtype);
+  };
+  const triple = (slots: TensorSlot[], at: number): ops.QuantizedTensor =>
+    ({ packed: arr(slots[at]!), scales: arr(slots[at + 1]!), biases: arr(slots[at + 2]!) });
 
   const caches: Cache[] = [];
-  for (const e of header.caches) {
-    const keys = MlxArray.fromPointer(
-      mmap.pointer(dataStart + e.kOffset), e.kShape, e.dtype as Dtype,
-    );
-    const values = MlxArray.fromPointer(
-      mmap.pointer(dataStart + e.vOffset), e.vShape, e.dtype as Dtype,
-    );
-    if (e.kind === "rotating") {
-      const c = new RotatingKVCache(e.maxSize!);
-      c.restoreState(keys, values, e.offset, e.idx!);
-      caches.push(c);
-    } else {
-      const c = new KVCache();
-      c.restoreState(keys, values, e.offset);
-      caches.push(c);
+  try {
+    for (const e of header.caches) {
+      const t = e.tensors;
+      switch (e.kind) {
+        case "kv": {
+          const c = new KVCache();
+          c.restoreState(arr(t[0]!), arr(t[1]!), e.offset);
+          caches.push(c);
+          break;
+        }
+        case "rotating": {
+          const c = new RotatingKVCache(e.maxSize!);
+          c.restoreState(arr(t[0]!), arr(t[1]!), e.offset, e.idx!);
+          caches.push(c);
+          break;
+        }
+        case "qkv": {
+          const c = new QuantizedKVCache(e.groupSize!, e.bits!);
+          c.restoreState(triple(t, 0), triple(t, 3), e.offset);
+          caches.push(c);
+          break;
+        }
+        case "rotating-qkv": {
+          const c = new RotatingQuantizedKVCache(e.maxSize!, e.groupSize!, e.bits!);
+          c.restoreState(triple(t, 0), triple(t, 3), e.offset, e.idx!);
+          caches.push(c);
+          break;
+        }
+        case "ssm": {
+          const c = new SSMCache();
+          c.conv = arr(t[0]!);
+          c.recurrent = arr(t[1]!);
+          c.offset = e.offset;
+          caches.push(c);
+          break;
+        }
+        default:
+          throw new Error(`${path}: unknown cache kind ${(e as { kind: string }).kind}`);
+      }
     }
+  } catch (err) {
+    for (const c of caches) c.dispose();
+    throw err;
   }
-  if (caches.length !== model.layers.length)
-    throw new Error(
-      `${path}: ${caches.length} cached layers but model has ${model.layers.length}`,
-    );
-  return { tokens: header.tokens, caches, mmap };
+  return { tokens: header.tokens, header, caches, mmap };
 }

package/src/mlx/array.ts

@@ -108,6 +108,27 @@ export class MlxArray {
     );
   }
 
+  /** Graph-free LEAF with this array's exact bytes — ONE copy, straight from
+   *  the (evaluated) mlx buffer into a fresh mlx-owned array, no JS-heap
+   *  intermediate (`fromBytesCopy(rawBytes())` is TWO copies: device→JS heap
+   *  →new leaf). mlx_array_new_data copies synchronously, so the source only
+   *  has to outlive this call. The buffer is read LINEARLY — the caller must
+   *  pass a row-major contiguous array (ops.contiguous first for views), same
+   *  contract as rawBytes. Segmented-backward boundary detach hot path. */
+  detachCopy(): MlxArray {
+    this.eval();
+    const dt = this.dtype;
+    const p =
+      dt === Dtype.float32 ? C.mlx_array_data_float32(this.handle)
+      : dt === Dtype.float16 ? C.mlx_array_data_float16(this.handle)
+      : dt === Dtype.bfloat16 ? C.mlx_array_data_bfloat16(this.handle)
+      : dt === Dtype.uint32 ? C.mlx_array_data_uint32(this.handle)
+      : null;
+    if (p === null) throw new Error(`detachCopy: unsupported dtype ${this.dtypeName}`);
+    const sb = shapeBuf(this.shape);
+    return new MlxArray(C.mlx_array_new_data(p as never, ptr(sb), this.ndim, dt));
+  }
+
   get handle(): MlxHandle {
     if (this.#disposed) throw new Error("MlxArray used after dispose");
     return this.#handle;

package/src/mlx/ops.ts

@@ -172,6 +172,32 @@ export function quantizedMatmul(
   x: MlxArray, w: MlxArray, scales: MlxArray, biases: MlxArray | null,
   spec: QuantSpec, transpose = true, s: S = gpuStream,
 ): MlxArray {
+  // UPSTREAM MLX BUG WORKAROUND (found 2026-07-02, reproduced on stock mlx
+  // 0.31.2): quantized_matmul with transpose=false returns WRONG VALUES when
+  // the row count is 2 or 3 (relnorm ~0.3-1.1 vs the dequantized matmul);
+  // M=1 and M>=4 are correct, transpose=true is correct at every M. Pad the
+  // row axis to 4 with zeros and slice the result back. NOTE this wrapper
+  // cannot protect mlx's INTERNAL vjp of a transpose=true qmm (C++ calls the
+  // primitive directly) — training heads pad the span through the head
+  // instead (loss.ts logitsFromHiddenPadM).
+  const rowAxis = x.shape.length - 2;
+  const rows = x.shape[rowAxis]!;
+  if (!transpose && (rows === 2 || rows === 3)) {
+    const padShape = [...x.shape];
+    padShape[rowAxis] = 4 - rows;
+    const pad = zeros(padShape, x.dtype, s);
+    const xPad = concatAxis([x, pad], rowAxis, s);
+    pad.dispose();
+    const yPad = quantizedMatmul(xPad, w, scales, biases, spec, transpose, s);
+    xPad.dispose();
+    const start = MlxArray.fromInt32(new Int32Array([0]), [1]);
+    const outShape = [...yPad.shape];
+    outShape[rowAxis] = rows;
+    const y = sliceDynamic(yPad, start, [rowAxis], outShape, s);
+    start.dispose();
+    yPad.dispose();
+    return y;
+  }
   return new MlxArray(
     outArray("quantized_matmul", (o) =>
       C.mlx_quantized_matmul(

package/src/model/gemma4-base.ts

@@ -590,6 +590,15 @@ export class QuantizedKVCache implements Cache {
     this.offset = Math.max(0, this.offset - n);
   }
 
+  /** Adopt persisted state (takes ownership of the triples' arrays) —
+   *  the quantized twin of KVCache.restoreState (kv-store persistence). */
+  restoreState(keys: ops.QuantizedTensor, values: ops.QuantizedTensor, offset: number): void {
+    this.dispose();
+    this.keys = keys;
+    this.values = values;
+    this.offset = offset;
+  }
+
   dispose(): void {
     for (const a of this.state()) a.dispose();
     this.keys = this.values = null;
@@ -1162,6 +1171,16 @@ export class RotatingQuantizedKVCache implements Cache {
     }
   }
 
+  /** Adopt persisted state (takes ownership of the triples' arrays) —
+   *  ring order as-laid-out, ringIdx carried with it (kv-store persistence). */
+  restoreState(keys: ops.QuantizedTensor, values: ops.QuantizedTensor, offset: number, idx: number): void {
+    this.dispose();
+    this.keys = keys;
+    this.values = values;
+    this.offset = offset;
+    this.ringIdx = idx;
+  }
+
   /** Oracle: to_quantized on an already-quantized rotating cache is
    *  idempotent. */
   toQuantized(): RotatingQuantizedKVCache {
@@ -1572,8 +1591,22 @@ export function quantizedSdpa(
   )
     return fusedDecodeSdpa(q, kq, vq, groupSize, bits);
   const tile = q.shape[2]! > 1 || process.env.MLX_BUN_FUSED_DECODE === "1";
-  if (tile && fusedSdpaSupported(q, mask, groupSize, bits))
+  if (tile && fusedSdpaSupported(q, mask, groupSize, bits)) {
+    // The N-tiled loop CANNOT be traced: its tile count and slice extents
+    // bake at trace-time N, so under shapeless replay a growing (concat-
+    // phase) quantized cache's newest rows are silently never attended —
+    // reproduced on e4b (whole-graph form, FUSED_DECODE=1: diverges within
+    // 3 tokens, then repetition-loops). generate() refuses to compile when
+    // MLX_BUN_FUSED_DECODE=1; this throw is the backstop for any other
+    // compiled caller — it lands in the trace's transactional fallback
+    // instead of producing silently wrong attention.
+    if (isCompiledTrace())
+      throw new Error(
+        "quantizedSdpaTiled inside a compiled trace: the tile loop bakes trace-time N " +
+          "(growing KV rows would be silently dropped) — run this generation uncompiled",
+      );
     return quantizedSdpaTiled(q, kq, vq, scale, mask, groupSize, bits);
+  }
   return quantizedSdpaUnfused(q, kq, vq, scale, mask, groupSize, bits);
 }
 

package/src/model/qwen3-delta.ts

@@ -192,7 +192,10 @@ export class SSMCache implements Cache {
   }
 
   makeMask(_N: number, _windowSize: number | null): Mask {
-    // B=1 single-stream: ssm_mask is None (no left-pad / variable lengths).
+    // ssm_mask is None both single-stream and in the batch lane: rows
+    // solo-prefill UNPADDED at B=1 (state never sees pad tokens), and
+    // batched decode feeds one real token per row — so unlike mlx-lm's
+    // left-padded batch-prefill there is never a pad position to mask.
     return { mode: "", arr: null };
   }
 
@@ -221,4 +224,48 @@ export class SSMCache implements Cache {
     this.conv = null;
     this.recurrent = null;
   }
+
+  // --- batch-lane dynamic-B ops (BatchScheduler only) ----------------------
+  // Both state slots are plain [B, ...] tensors with no temporal axis and no
+  // left-padding, so the batched twins of mergeKVRows/filterKVRows collapse
+  // to B-axis concat/take. `offset` is bookkeeping only for SSM layers (RoPE
+  // reads the FULL-attention caches; the linear-attn forward never reads it),
+  // so merged rows carrying different token counts share the max.
+
+  /** Merge a fully-prefilled solo row into a running batched cache (`prev`
+   *  null ⇒ the solo row becomes the batch). The returned cache owns fresh
+   *  arrays — or, when there is nothing to concat, arrays STOLEN from `solo`
+   *  (nulled there so the scheduler's unconditional dispose of the solo
+   *  caches stays safe). `prev`/`solo` disposal remains the caller's job. */
+  static mergeRows(prev: SSMCache | null, solo: SSMCache): SSMCache {
+    if (!solo.conv || !solo.recurrent)
+      throw new Error("SSMCache.mergeRows: solo row has no state (prefill first)");
+    const out = new SSMCache();
+    if (prev) {
+      if (!prev.conv || !prev.recurrent)
+        throw new Error("SSMCache.mergeRows: batched cache has no state");
+      out.conv = ops.concatAxis([prev.conv, solo.conv], 0);
+      out.recurrent = ops.concatAxis([prev.recurrent, solo.recurrent], 0);
+    } else {
+      out.conv = solo.conv;
+      out.recurrent = solo.recurrent;
+      solo.conv = null;
+      solo.recurrent = null;
+    }
+    out.offset = Math.max(prev?.offset ?? 0, solo.offset);
+    return out;
+  }
+
+  /** Evict rows not in `keep` (ascending row indices), in place. */
+  filter(keep: number[]): void {
+    if (!this.conv || !this.recurrent) return;
+    const idx = ops.fromInt32(keep, [keep.length]);
+    const conv = ops.takeAxis(this.conv, idx, 0);
+    const recurrent = ops.takeAxis(this.recurrent, idx, 0);
+    idx.dispose();
+    this.conv.dispose();
+    this.recurrent.dispose();
+    this.conv = conv;
+    this.recurrent = recurrent;
+  }
 }

package/src/model/qwen3_5.ts

@@ -212,8 +212,17 @@ class Qwen3Attention {
     v = disposing(v, ops.reshape(v, [B, L, this.nKvHeads, this.headDim]));
     v = disposing(v, ops.transposeAxes(v, [0, 2, 1, 3]));
 
-    q = disposing(q, ops.rope(q, this.ropeDims, this.ropeBase, cache.offset, null));
-    k = disposing(k, ops.rope(k, this.ropeDims, this.ropeBase, cache.offset, null));
+    // Batched decode: the scheduler's mask wrapper exposes each row's REAL
+    // position as ropeOffsetArr (rows have different prompt lengths); the
+    // dynamic-offset kernel is the same fast::rope, bit-exact vs the static
+    // form (tests/compile.test.ts). Serial lane: scalar offset, unchanged.
+    const offArr = (cache as { ropeOffsetArr?: MlxArray }).ropeOffsetArr;
+    q = disposing(q, offArr
+      ? ops.ropeDynamic(q, this.ropeDims, this.ropeBase, offArr, null)
+      : ops.rope(q, this.ropeDims, this.ropeBase, cache.offset, null));
+    k = disposing(k, offArr
+      ? ops.ropeDynamic(k, this.ropeDims, this.ropeBase, offArr, null)
+      : ops.rope(k, this.ropeDims, this.ropeBase, cache.offset, null));
 
     let attn: MlxArray;
     if (cache instanceof QuantizedKVCache) {