mlx-bun 0.2.0 → 0.4.0

package/src/kv-store.ts

@@ -1,5 +1,7 @@
+import { type CheckpointAttachment, attachmentBytes, disposeAttachments } from "./backends/mlx/checkpoint-state";
+import { cleanupFailure, disposeResources } from "./engine/resources";
 // KV-cache persistence: save prompt caches to disk, reload by streamed copy.
-// The serialization core of the SSD cold tier (docs/design/ssd-kv-cold-tier.md).
+// The serialization core of the SSD cold tier (docs/design/kv-cache.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):
@@ -85,6 +87,7 @@ interface TensorSlot {
 export interface CacheHeaderEntry {
   kind: CacheKind;
   offset: number;
+  minimumReusableOffset?: number;
   /** rotating variants: ring write index */
   idx?: number;
   /** rotating variants: window size */
@@ -113,6 +116,7 @@ export interface CacheHeaderEntry {
 }
 
 export interface KvSaveMeta {
+  attachments?: CheckpointAttachment[];
   modelId?: string;
   /** configFingerprint(config) — covers every graph-shaping field incl.
    *  the kv-quant scheme, so a scheme flip invalidates naturally. */
@@ -121,10 +125,28 @@ export interface KvSaveMeta {
   ns?: string;
   /** sha256 of tokenizer.json: identical ids must mean identical text. */
   tokenizerHash?: string;
+  /** In-flight generation state. The serialized caches cover `tokens`;
+   *  pendingToken was sampled from that state but has not been emitted. */
+  generationCheckpoint?: {
+    key: string;
+    cacheNs: string;
+    originalPromptTokens: number;
+    generatedTokens: number;
+    pendingToken: number;
+    /** Actual sampler seed, including a server-generated default. Required to
+     *  continue the random stream exactly after a restart. */
+    seed: number;
+    /** True when `seed` was part of the client request rather than generated
+     *  by the server. */
+    seedWasExplicit: boolean;
+  };
 }
 
-export interface KvFileHeader extends KvSaveMeta {
-  formatVersion: 3;
+export interface KvFileHeader extends Omit<KvSaveMeta, "attachments"> {
+  attachments?: Array<Omit<CheckpointAttachment, "tensors"> & { tensors: TensorSlot[] }>;
+  /** Absent in earlier v3 files, which used mlx-cache-v3. */
+  codecProvider?: string;
+  formatVersion: 3 | 4;
   createdAt: number;
   tokens: number[];
   caches: CacheHeaderEntry[];
@@ -143,7 +165,7 @@ const alignUp = (n: number) => Math.ceil(n / ALIGN) * ALIGN;
  *  materialization until rawBytes at write time). */
 interface TensorSource { arr: MlxArray; disposeAfter: boolean }
 
-interface SnapshotContext {
+export interface SnapshotContext {
   slots: TensorSlot[];
   push(a: MlxArray, disposeAfter: boolean): void;
   liveSlice(a: MlxArray, upTo: number): MlxArray;
@@ -151,21 +173,21 @@ interface SnapshotContext {
   pushTriple(t: ops.QuantizedTensor, upTo: number | null): void;
 }
 
-interface CloneContext {
+export interface CloneContext {
   view(a: MlxArray): MlxArray;
   liveView(a: MlxArray, upTo: number): MlxArray;
   mlaView(a: MlxArray, upTo: number): MlxArray;
   tripleView(t: ops.QuantizedTensor, upTo: number | null): ops.QuantizedTensor;
 }
 
-interface LoadContext {
+export interface LoadContext {
   path: string;
   arr(slot: TensorSlot): MlxArray;
   grownArr(slot: TensorSlot, offset: number): MlxArray;
   triple(slots: TensorSlot[], at: number): ops.QuantizedTensor;
 }
 
-interface CacheCodec {
+export interface CacheCodec {
   matches(cache: Cache): boolean;
   snapshot(cache: Cache, context: SnapshotContext): CacheHeaderEntry;
   clone(cache: Cache, context: CloneContext): Cache;
@@ -255,8 +277,8 @@ const CACHE_CODECS = {
     },
     clone: (cache, context) => {
       const c = cache as KVCache;
-      if (!c.keys || !c.values) throw new Error("cannot clone an empty cache");
       const clone = new KVCache();
+      if (!c.keys || !c.values) return clone;
       clone.restoreState(
         context.liveView(c.keys, c.offset),
         context.liveView(c.values, c.offset),
@@ -287,8 +309,8 @@ const CACHE_CODECS = {
     },
     clone: (cache, context) => {
       const c = cache as RotatingKVCache;
-      if (!c.keys || !c.values) throw new Error("cannot clone an empty cache");
       const clone = new RotatingKVCache(c.maxSize);
+      if (!c.keys || !c.values) return clone;
       clone.restoreState(context.view(c.keys), context.view(c.values), c.offset, c.ringIdx);
       return clone;
     },
@@ -314,8 +336,8 @@ const CACHE_CODECS = {
     },
     clone: (cache, context) => {
       const c = cache as QuantizedKVCache;
-      if (!c.keys || !c.values) throw new Error("cannot clone an empty cache");
       const clone = new QuantizedKVCache(c.groupSize, c.bits);
+      if (!c.keys || !c.values) return clone;
       clone.restoreState(
         context.tripleView(c.keys, c.offset),
         context.tripleView(c.values, c.offset),
@@ -345,8 +367,8 @@ const CACHE_CODECS = {
     },
     clone: (cache, context) => {
       const c = cache as RotatingQuantizedKVCache;
-      if (!c.keys || !c.values) throw new Error("cannot clone an empty cache");
       const clone = new RotatingQuantizedKVCache(c.maxSize, c.groupSize, c.bits);
+      if (!c.keys || !c.values) return clone;
       clone.restoreState(
         context.tripleView(c.keys, null), context.tripleView(c.values, null),
         c.offset, c.ringIdx,
@@ -458,16 +480,61 @@ const CACHE_CODECS = {
   },
 } satisfies Record<CacheKind, CacheCodec>;
 
-const cacheCodec = (cache: Cache): CacheCodec => {
-  const codec = Object.values(CACHE_CODECS).find((candidate) => candidate.matches(cache));
-  if (!codec) throw new Error(`unsupported cache signature ${cacheSignature(cache)}`);
-  return codec;
-};
+/** Codecs belong to the loaded backend binding. The persisted ID is checked
+ * before allocating state; codecs are engine code, never loaded from a file. */
+export interface CacheCodecProvider {
+  readonly id: string;
+  forCache(cache: Cache): CacheCodec;
+  forHeader(entry: CacheHeaderEntry): CacheCodec;
+}
 
-export const cacheHeadersTrimmable = (caches: readonly CacheHeaderEntry[]): boolean =>
-  caches.every((entry) => CACHE_CODECS[entry.kind].headerTrimmable(entry));
+export function createCacheCodecProvider(
+  id: string, entries: Readonly<Partial<Record<CacheKind, CacheCodec>>>,
+): CacheCodecProvider {
+  if (!id.trim()) throw new Error("cache codec provider needs an identity");
+  const table = new Map(Object.entries(entries).map(([kind, codec]) => [kind, Object.freeze({
+    matches: codec.matches.bind(codec), snapshot: codec.snapshot.bind(codec), clone: codec.clone.bind(codec),
+    load: codec.load.bind(codec), headerTrimmable: codec.headerTrimmable.bind(codec),
+  })]));
+  const candidates = [...table.values()];
+  const selected = new WeakMap<Cache, CacheCodec>();
+  return Object.freeze({
+    id,
+    forCache(cache: Cache) {
+      const bound = selected.get(cache);
+      if (bound) return bound;
+      const matches = candidates.filter((codec) => codec.matches(cache));
+      if (matches.length !== 1)
+        throw new Error(`cache codec provider ${id}: ${matches.length} matches for ${cacheSignature(cache)}`);
+      selected.set(cache, matches[0]!);
+      return matches[0]!;
+    },
+    forHeader(entry: CacheHeaderEntry) {
+      const codec = table.get(entry.kind);
+      if (!codec) throw new Error(`cache codec provider ${id}: unsupported cache kind ${entry.kind}`);
+      return codec;
+    },
+  });
+}
 
-function snapshotCache(c: Cache, dataOffset: number): { entry: CacheHeaderEntry; sources: TensorSource[]; next: number } {
+export const legacyCacheCodecs = createCacheCodecProvider("mlx-cache-v3", CACHE_CODECS);
+
+/** Old TQ files lack the conversion boundary. Keep them reusable at their
+ * stored offset, but never infer that an earlier precision state is valid. */
+export function cacheHeaderMinimumReusableOffset(entry: CacheHeaderEntry): number {
+  const offset = entry.minimumReusableOffset ?? (entry.kind === "turboquant" ? entry.offset : 0);
+  if (!Number.isSafeInteger(offset) || offset < 0 || offset > entry.offset)
+    throw new Error("invalid cache minimum reusable offset");
+  return offset;
+}
+export const cacheHeadersMinimumReusableOffset = (caches: readonly CacheHeaderEntry[]): number =>
+  Math.max(0, ...caches.map(cacheHeaderMinimumReusableOffset));
+
+export const cacheHeadersTrimmable = (
+  caches: readonly CacheHeaderEntry[], codecs: CacheCodecProvider = legacyCacheCodecs,
+): boolean => caches.every((entry) => codecs.forHeader(entry).headerTrimmable(entry));
+
+function snapshotCache(c: Cache, dataOffset: number, codecs: CacheCodecProvider): { entry: CacheHeaderEntry; sources: TensorSource[]; next: number } {
   const slots: TensorSlot[] = [];
   const sources: TensorSource[] = [];
   let cursor = dataOffset;
@@ -493,10 +560,14 @@ function snapshotCache(c: Cache, dataOffset: number): { entry: CacheHeaderEntry;
     }
   };
 
-  const entry = cacheCodec(c).snapshot(c, {
-    slots, push, liveSlice, liveMlaSlice, pushTriple,
-  });
-  return { entry, sources, next: cursor };
+  try {
+    const entry = codecs.forCache(c).snapshot(c, { slots, push, liveSlice, liveMlaSlice, pushTriple });
+    if (c.minimumReusableOffset !== undefined) entry.minimumReusableOffset = c.minimumReusableOffset;
+    return { entry, sources, next: cursor };
+  } catch (error) {
+    for (const source of sources) if (source.disposeAfter) source.arr.dispose();
+    throw error;
+  }
 }
 
 /** Persist `caches` (+ the exact token prefix they encode) to `path`.
@@ -513,15 +584,17 @@ function snapshotCache(c: Cache, dataOffset: number): { entry: CacheHeaderEntry;
  *  called while the caches hold EXACTLY the state to snapshot (post-wrap
  *  rings cannot be rewound later). Same per-kind dispatch as
  *  snapshotCache/loadKvCache. */
-export function cloneKvCaches(caches: Cache[]): Cache[] {
-  const view = (a: MlxArray): MlxArray => a.slice(a.shape.map(() => 0), [...a.shape]);
+export function cloneKvCaches(caches: Cache[], codecs: CacheCodecProvider = legacyCacheCodecs): Cache[] {
+  const pending: MlxArray[] = [];
+  const hold = (array: MlxArray): MlxArray => { pending.push(array); return array; };
+  const view = (a: MlxArray): MlxArray => hold(a.slice(a.shape.map(() => 0), [...a.shape]));
   const liveView = (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]);
+    return hold(a.slice([0, 0, 0, 0], [B, H, upTo, D]));
   };
   const mlaView = (a: MlxArray, upTo: number): MlxArray => {
     const [B, , D] = a.shape as [number, number, number];
-    return a.slice([0, 0, 0], [B, upTo, D]);
+    return hold(a.slice([0, 0, 0], [B, upTo, D]));
   };
   const tripleView = (t: ops.QuantizedTensor, upTo: number | null): ops.QuantizedTensor => ({
     packed: upTo === null ? view(t.packed) : liveView(t.packed, upTo),
@@ -531,11 +604,13 @@ export function cloneKvCaches(caches: Cache[]): Cache[] {
   const out: Cache[] = [];
   try {
     for (const c of caches) {
-      out.push(cacheCodec(c).clone(c, { view, liveView, mlaView, tripleView }));
+      const clone = codecs.forCache(c).clone(c, { view, liveView, mlaView, tripleView });
+      if (c.minimumReusableOffset !== undefined) clone.minimumReusableOffset = c.minimumReusableOffset;
+      out.push(clone);
+      pending.length = 0; // transferred to the returned cache
     }
   } catch (err) {
-    for (const c of out) c.dispose();
-    throw err;
+    cleanupFailure(err, () => disposeResources([...pending, ...out]));
   }
   return out;
 }
@@ -545,23 +620,37 @@ export function cloneKvCaches(caches: Cache[]): Cache[] {
  *  between yields so the EVENT LOOP KEEPS SERVING while a multi-hundred-MB
  *  entry flushes (the write-behind persistence contract: durability never
  *  blocks a request). try/finally still runs on early generator close. */
-function* saveKvCacheSteps(path: string, tokens: number[], caches: Cache[], meta: KvSaveMeta): Generator<void, void, void> {
+function* saveKvCacheSteps(path: string, tokens: number[], caches: Cache[], meta: KvSaveMeta, codecs: CacheCodecProvider): Generator<void, void, void> {
   // Plan pass: header entries + lazy tensor sources, NO bytes materialized.
   const entries: CacheHeaderEntry[] = [];
   const sources: TensorSource[] = [];
   let dataOffset = 0; // relative; rebased after header is sized
   try {
     for (const c of caches) {
-      const s = snapshotCache(c, dataOffset);
+      const s = snapshotCache(c, dataOffset, codecs);
       entries.push(s.entry);
       sources.push(...s.sources);
       dataOffset = s.next;
     }
 
+    const { attachments: liveAttachments, ...identity } = meta;
+    const attachments = liveAttachments?.map((attachment) => ({
+      schema: attachment.schema,
+      metadata: { ...attachment.metadata },
+      tensors: attachment.tensors.map((arr) => {
+        const slot: TensorSlot = { off: dataOffset, bytes: arr.nbytes,
+          shape: [...arr.shape], dtype: arr.dtype, hash: "0000000000000000" };
+        dataOffset = alignUp(dataOffset + slot.bytes);
+        sources.push({ arr, disposeAfter: false });
+        return slot;
+      }),
+    }));
+
     // Header size is final NOW: hashes are fixed-width placeholders that get
     // patched in place (same byte length) after the data pass computes them.
     const header: KvFileHeader = {
-      formatVersion: 3, createdAt: Date.now(), ...meta, tokens, caches: entries,
+      formatVersion: attachments?.length ? 4 : 3, createdAt: Date.now(), ...identity,
+      ...(attachments?.length ? { attachments } : {}), codecProvider: codecs.id, tokens, caches: entries,
     };
     const headerLen = new TextEncoder().encode(JSON.stringify(header)).length;
     const dataStart = alignUp(PREFIX_LEN + headerLen);
@@ -582,7 +671,7 @@ function* saveKvCacheSteps(path: string, tokens: number[], caches: Cache[], meta
       // pages, which makes them VISIBLE to ps RSS — accounting, not an
       // allocation (see bench-serve.ts' per-leg RSS note).
       let srcIdx = 0;
-      for (const e of entries) {
+      for (const e of [...entries, ...(attachments ?? [])]) {
         for (const slot of e.tensors) {
           const src = sources[srcIdx++]!;
           const c = ops.contiguous(src.arr);
@@ -626,38 +715,41 @@ function* saveKvCacheSteps(path: string, tokens: number[], caches: Cache[], meta
   }
 }
 
-export function saveKvCache(path: string, tokens: number[], caches: Cache[], meta: KvSaveMeta = {}): void {
-  for (const _ of saveKvCacheSteps(path, tokens, caches, meta)) { /* drain */ }
+export function saveKvCache(path: string, tokens: number[], caches: Cache[], meta: KvSaveMeta = {}, codecs: CacheCodecProvider = legacyCacheCodecs): void {
+  for (const _ of saveKvCacheSteps(path, tokens, caches, meta, codecs)) { /* drain */ }
 }
 
 /** Non-blocking variant: yields the event loop after every tensor write so
  *  serving interleaves with the flush. Caller owns `caches` lifetime for
  *  the duration (pass zero-copy clones, dispose after).
  *
- *  `waitTurn` (optional) is awaited BEFORE every step — including the first —
- *  so the caller can gate the flush's schedule (the server passes the
- *  gateway's onIdle: each tensor step is a blocking GPU sync on the decode
- *  stream + a synchronous multi-MB writeSync, and the old unconditional
- *  setImmediate pacing interleaved those slices exactly between decode
- *  tokens — the 2026-07-07 decode@ctx contamination. A request arriving
- *  MID-flush pauses the remaining tensors until idle again). Gate failures
- *  are swallowed: scheduling advice must never corrupt the write path
- *  (an early generator close inside the fd-open section would leak the fd). */
+ *  `runStep` (optional) runs each blocking tensor step inside the caller's
+ *  exclusion domain. Waiting for an idle observation is not sufficient: a
+ *  request can acquire the engine between that check and `steps.next()`,
+ *  putting rawBytesView's GPU sync on the decode stream during prefill. */
 export async function saveKvCacheAsync(
   path: string, tokens: number[], caches: Cache[], meta: KvSaveMeta = {},
-  waitTurn?: () => Promise<void>,
+  runStep?: <T>(step: () => T) => Promise<T>, codecs: CacheCodecProvider = legacyCacheCodecs,
 ): Promise<void> {
-  const steps = saveKvCacheSteps(path, tokens, caches, meta);
-  while (true) {
-    if (waitTurn) await waitTurn().catch(() => {});
-    if (steps.next().done) break;
-    await new Promise<void>((r) => setImmediate(r));
+  const steps = saveKvCacheSteps(path, tokens, caches, meta, codecs);
+  try {
+    while (true) {
+      const next = runStep ? await runStep(() => steps.next()) : steps.next();
+      if (next.done) break;
+      await new Promise<void>((r) => setImmediate(r));
+    }
+  } catch (error) {
+    // A rejected execution lease can leave the writer suspended with an
+    // open fd. Throw through it so its existing failure cleanup runs.
+    steps.throw(error);
+    throw error;
   }
 }
 
 /** One pending write-behind item: an entry's copied tokens + OWNED
  *  zero-copy cache clones (the queue disposes them on every exit path). */
 export interface SpillItem {
+  attachments?: CheckpointAttachment[];
   tokens: number[];
   caches: Cache[];
   ns: string;
@@ -666,7 +758,7 @@ export interface SpillItem {
 /** Bounded write-behind queue (2026-07-07 post-merge review fix).
  *
  *  Pending spill/snapshot clones pin their entries' GPU buffers until the
- *  idle-gated flush gets a turn; the old bare promise chain queued them
+ *  generation-locked flush gets a turn; the old bare promise chain queued them
  *  WITHOUT BOUND, so under sustained traffic (gate starved, evictions
  *  ongoing) resident memory = prompt-cache cap + every queued clone —
  *  allocator pressure exactly under the load that caused the evictions.
@@ -702,6 +794,8 @@ export class SpillQueue {
     readonly store: (item: SpillItem) => Promise<unknown>,
     /** Clone disposal — frees the pinned GPU memory. */
     readonly disposeClones: (caches: Cache[]) => void,
+    /** Zero-byte control barriers may need to survive capacity shedding. */
+    readonly canDrop: (item: SpillItem) => boolean = () => true,
   ) {}
 
   /** Bytes pinned by queued (not yet flushed) clones. */
@@ -717,7 +811,7 @@ export class SpillQueue {
     const result = new Promise<boolean>((resolve) => { settle = resolve; });
     const rec: SpillRec = {
       ...item,
-      bytes: this.bytesOf(item.caches),
+      bytes: this.bytesOf(item.caches) + attachmentBytes(item.attachments),
       dropped: false,
       settle,
     };
@@ -725,7 +819,7 @@ export class SpillQueue {
     this.#bytes += rec.bytes;
     while (this.#bytes > this.capBytes) {
       const victim = this.#queue.find(
-        (p) => !p.dropped && p !== this.#inFlight && p !== rec,
+        (p) => !p.dropped && p !== this.#inFlight && p !== rec && this.canDrop(p),
       );
       if (!victim) break; // only the new item (or in-flight) left — soft cap
       this.#drop(victim);
@@ -744,12 +838,19 @@ export class SpillQueue {
       } finally {
         this.#inFlight = null;
         this.#remove(rec);
-        this.disposeClones(rec.caches);
+        disposeResources([{ dispose: () => this.disposeClones(rec.caches) },
+          { dispose: () => disposeAttachments(rec.attachments) }]);
       }
     });
     return result;
   }
 
+  /** Cancel queued ownership, excluding the write already in flight. */
+  cancelWhere(matches: (item: SpillItem) => boolean): void {
+    for (const rec of [...this.#queue])
+      if (!rec.dropped && rec !== this.#inFlight && matches(rec)) this.#drop(rec);
+  }
+
   /** Settles when everything enqueued so far has flushed or dropped
    *  (tests and graceful shutdown). */
   drain(): Promise<void> { return this.#chain; }
@@ -759,7 +860,8 @@ export class SpillQueue {
     this.#dropped++;
     rec.settle(false);
     this.#remove(rec);
-    this.disposeClones(rec.caches);
+    disposeResources([{ dispose: () => this.disposeClones(rec.caches) },
+          { dispose: () => disposeAttachments(rec.attachments) }]);
   }
 
   #remove(rec: SpillRec): void {
@@ -789,7 +891,7 @@ export function readKvHeader(path: string): KvFileHeader & { dataStart: number }
     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 !== 3)
+    if (header.formatVersion !== 3 && header.formatVersion !== 4)
       throw new Error(`${path}: unsupported formatVersion ${header.formatVersion}`);
     return { ...header, dataStart };
   } finally {
@@ -798,6 +900,7 @@ export function readKvHeader(path: string): KvFileHeader & { dataStart: number }
 }
 
 export interface LoadedKvCache {
+  attachments?: CheckpointAttachment[];
   tokens: number[];
   header: KvFileHeader;
   caches: Cache[];
@@ -904,9 +1007,12 @@ function withStepCapacity(a: MlxArray, offset: number): MlxArray {
 export function loadKvCache(
   path: string,
   model: { makeCache(): Cache[] },
-  expect: KvLoadExpect = {},
+  expect: KvLoadExpect = {}, codecs: CacheCodecProvider = legacyCacheCodecs,
 ): LoadedKvCache {
   const header = readKvHeader(path);
+  if ((header.codecProvider ?? legacyCacheCodecs.id) !== codecs.id)
+    throw new Error(`${path}: cache codec provider mismatch`);
+  for (const entry of header.caches) codecs.forHeader(entry);
   for (const key of [
     "modelId",
     "configFingerprint",
@@ -964,22 +1070,30 @@ export function loadKvCache(
     ({ packed: arr(slots[at]!), scales: arr(slots[at + 1]!), biases: arr(slots[at + 2]!) });
 
   const caches: Cache[] = [];
+  const attachments: CheckpointAttachment[] = [];
   try {
     for (const e of header.caches) {
-      const codec = CACHE_CODECS[e.kind];
-      if (!codec)
-        throw new Error(`${path}: unknown cache kind ${(e as { kind: string }).kind}`);
-      caches.push(codec.load(e, { path, arr, grownArr, triple }));
+      const codec = codecs.forHeader(e);
+      const minimumOffset = cacheHeaderMinimumReusableOffset(e);
+      const cache = codec.load(e, { path, arr, grownArr, triple });
+      cache.minimumReusableOffset = minimumOffset;
+      caches.push(cache);
       // This entry's tensors are now owned by its cache — stop tracking them
       // (the catch must not double-free through both pending AND caches).
       pending.length = 0;
     }
+    for (const attachment of header.attachments ?? []) {
+      attachments.push({ schema: attachment.schema, metadata: { ...attachment.metadata },
+        tensors: attachment.tensors.map(arr) });
+      pending.length = 0;
+    }
   } catch (err) {
+    disposeAttachments(attachments);
     for (const a of pending) a.dispose(); // the mid-entry orphans
     for (const c of caches) c.dispose();
     throw err;
   } finally {
     mmap.unmap(); // every tensor was copied out — nothing aliases the file
   }
-  return { tokens: header.tokens, header, caches };
+  return { tokens: header.tokens, header, caches, ...(attachments.length ? { attachments } : {}) };
 }

package/src/lab/README.md

@@ -0,0 +1,18 @@
+# src/lab — quarantined, no-oracle code
+
+Everything here is a Lab-tier module (docs/design/unified-engine-frontier-plan.md):
+no external oracle, off by default, kept only while a measurement decides its
+fate (PLAN.md "Maintainability program", D6 darlings). Rules, gate-enforced by
+`scripts/check-hygiene.ts` (check 11):
+
+- nothing outside `src/lab/` may import from `src/lab/` except the edges listed
+  in `LAB_IMPORT_ALLOWLIST` — those are recorded debt, not permission;
+- a new production dependency on a Lab module is a decision, not a default:
+  either promote the module out of the lab with an oracle/measurement, or
+  route it through an option seam.
+
+| module | flag / entry | oracle |
+|---|---|---|
+| `curve/` (curve sampler + designer page) | `--curve-*`, `/curves` | none (Lab) |
+| `paged-kv/` | `--paged-kv` (serial-only) | parity vs plain KVCache (bit-exact), no measured win |
+| `expert-trace/` | `MLX_BUN_EXPERT_TRACE` hook in gemma4 | closed study (expert-offload) |

package/src/{assets → lab/curve}/curve-designer.html

@@ -88,7 +88,7 @@ button.ghost{background:var(--card);border:1px solid var(--bd);border-radius:8px
 
 <script>
 const ln=Math.log;
-const UMIN=ln(1e-6); // bottom-left / plot floor (matches src/curve-sampler.ts CURVE_UMIN)
+const UMIN=ln(1e-6); // bottom-left / plot floor (matches src/lab/curve/curve-sampler.ts CURVE_UMIN)
 const ENGINE_BASE = location.protocol.startsWith("http") ? location.origin : "http://localhost:8080";
 // ---- curve state: 5 movable control points, sorted ascending in p_in ----
 function identityPoints(){ return [{x_pct:1e-4,y_pct:1e-4},{x_pct:0.1,y_pct:0.1},{x_pct:1,y_pct:1},{x_pct:9,y_pct:9},{x_pct:100,y_pct:100}]; }
@@ -102,7 +102,7 @@ const PROMPTS=[
   "Describe the feeling of a city at 3am in one vivid sentence.",
 ];
 
-// ---- math: monotone cubic (PCHIP), IDENTICAL to src/curve-sampler.ts ----
+// ---- math: monotone cubic (PCHIP), IDENTICAL to src/lab/curve/curve-sampler.ts ----
 const Lp=pt=>[ln(Math.max(pt.x_pct,1e-9)/100),ln(Math.max(pt.y_pct,1e-9)/100)];
 function knotArrays(c){ const k=c.points.map(Lp).sort((a,b)=>a[0]-b[0]); return {X:k.map(q=>q[0]),Y:k.map(q=>q[1])}; }
 function secants(c){ const {X,Y}=knotArrays(c),s=[]; for(let i=X.length-1;i>0;i--)s.push((Y[i]-Y[i-1])/(X[i]-X[i-1])); return s; }

package/src/{curve-sampler.ts → lab/curve/curve-sampler.ts}

@@ -16,8 +16,8 @@
 // BELOW lowers them (renorm makes the literal number an intent); the SLOPE across
 // a region sets how evenly the tokens inside it spread (flat<1 even, steep>1 separated).
 
-import { MlxArray } from "./mlx/array";
-import * as ops from "./mlx/ops";
+import { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
 
 export interface CurvePoint { x_pct: number; y_pct: number; }
 export interface CurveParams {

package/src/{expert-trace.ts → lab/expert-trace/expert-trace.ts}

@@ -2,13 +2,13 @@
 //
 // Records every MoE router decision as JSONL — one record per router call:
 //   { c: callSeq, l: layerIdx, s: [B, L, k], i: [flattened uint32 expert ids] }
-// Analyse with scripts/analyze-expert-trace.ts (coverage curve, working-set
+// Analyse with analyze-expert-trace.ts (deleted 2026-08-23; git history) (coverage curve, working-set
 // size, within-task stability, cross-task specialisation).
 //
 // Two ways to drive it:
 //   - env: `MLX_BUN_EXPERT_TRACE=<path> mlx-bun serve ...` (auto-starts here)
 //   - programmatic: beginExpertTrace(path) / endExpertTrace() per domain
-//     (scripts/run-expert-trace.ts uses this to trace several domains from
+//     (run-expert-trace.ts (deleted 2026-08-23; git history) uses this to trace several domains from
 //     one model load).
 //
 // Recording adds a per-call GPU->host sync (it reads the routed indices back),
@@ -17,9 +17,9 @@
 // Router.forward is a single branch (zero overhead) — inert by default, same
 // discipline as the slots flag.
 
-import { Dtype } from "./mlx/ffi";
-import type { MlxArray } from "./mlx/array";
-import { runtimeValue } from "./runtime-config";
+import { Dtype } from "../../mlx/ffi";
+import type { MlxArray } from "../../mlx/array";
+import { runtimeValue } from "../../runtime-config";
 
 let sink: Bun.FileSink | null = null;
 let seq = 0;

package/src/lab/paged-kv/paged-kv-rows.ts

@@ -0,0 +1,83 @@
+import type { BatchableCache, Cache, Mask } from "../../model/gemma4-base";
+import { buildBatchedDecodeMask, mergeKVRows } from "../../model/batched-mask";
+import { MlxArray } from "../../mlx/array";
+import { disposeResources } from "../../engine/resources";
+import { PagedKVCache, poolBlocksFor } from "./paged-kv";
+
+/** Block storage owns row membership; attention consumes the same gathered
+ * K/V representation as the existing paged cache. No scheduler policy here. */
+export class PagedKvRows implements BatchableCache {
+  #rows: PagedKVCache[] = [];
+  #rope?: MlxArray;
+  #ropeOffset = -1;
+
+  constructor(readonly capacityTokens: number, readonly blockSize: number) {}
+  signature(): string { return "kv:paged-rows"; }
+  get batchSize(): number { return this.#rows.length; }
+  get offset(): number { return Math.max(0, ...this.#rows.map(row => row.offset)); }
+  get rowOffsets(): number[] { return this.#rows.map(row => row.offset); }
+  get leftPad(): number[] { const width = this.offset; return this.#rows.map(row => width - row.offset); }
+  get ropeOffsetArr(): MlxArray | undefined {
+    if (this.#rows.length < 2) return undefined;
+    if (!this.#rope || this.#ropeOffset !== this.offset) {
+      this.releaseRopeArr();
+      this.#rope = MlxArray.fromInt32(Int32Array.from(this.rowOffsets), [this.batchSize]);
+      this.#ropeOffset = this.offset;
+    }
+    return this.#rope;
+  }
+  releaseRopeArr(): void { this.#rope?.dispose(); this.#rope = undefined; this.#ropeOffset = -1; }
+  makeEmptyBatch(): PagedKvRows { return new PagedKvRows(this.capacityTokens, this.blockSize); }
+
+  mergeRows(rows: readonly Cache[]): void {
+    const next: PagedKVCache[] = [];
+    try {
+      for (const source of rows) {
+        if (source instanceof PagedKvRows) {
+          for (const row of source.#rows) next.push(row.clone());
+        } else if (source instanceof PagedKVCache) next.push(source.clone());
+        else throw new Error(`Paged rows cannot merge ${source.signature()}`);
+      }
+    } catch (error) { disposeResources(next); throw error; }
+    const previous = this.#rows; this.#rows = next;
+    this.releaseRopeArr(); disposeResources(previous);
+  }
+  extractRow(row: number): PagedKVCache { return this.#rows[row]!.clone(); }
+  filterRows(keep: readonly number[]): void {
+    const previous = this.#rows;
+    this.#rows = keep.map(row => previous[row]!);
+    const retained = new Set(keep);
+    this.releaseRopeArr();
+    disposeResources(previous.filter((_, row) => !retained.has(row)));
+  }
+
+  updateAndFetch(keys: MlxArray, values: MlxArray): [MlxArray, MlxArray] {
+    if (this.#rows.length === 1) return this.#rows[0]!.updateAndFetch(keys, values);
+    const fetched: Array<{ keys: MlxArray; values: MlxArray }> = [];
+    try {
+      for (const [row, storage] of this.#rows.entries()) {
+        using k = keys.slice([row, 0, 0, 0], [row + 1, ...keys.shape.slice(1)]);
+        using v = values.slice([row, 0, 0, 0], [row + 1, ...values.shape.slice(1)]);
+        const [nextK, nextV] = storage.updateAndFetch(k, v);
+        fetched.push({ keys: nextK, values: nextV });
+      }
+      const result = mergeKVRows(fetched);
+      return [result.keys, result.values];
+    } finally { disposeResources(fetched.flatMap(row => [row.keys, row.values])); }
+  }
+  makeMask(tokens: number, window: number | null): Mask {
+    if (this.#rows.length === 1) return this.#rows[0]!.makeMask(tokens, window);
+    return { mode: "array", arr: buildBatchedDecodeMask(this.batchSize, tokens,
+      this.offset + tokens, this.leftPad, window) };
+  }
+  state(): MlxArray[] { return this.#rows.flatMap(row => row.state()); }
+  isTrimmable(): boolean { return true; }
+  trim(tokens: number): void { for (const row of this.#rows) row.trim(tokens); this.releaseRopeArr(); }
+  projectedBytes(tokens: number): number {
+    const pool = this.#rows.find(row => row.pool)?.pool;
+    if (!pool) return 0;
+    const bytesPerBlock = (pool.keys.nbytes + pool.values.nbytes) / pool.numBlocks;
+    return poolBlocksFor(tokens, this.blockSize) * bytesPerBlock;
+  }
+  dispose(): void { const rows = this.#rows; this.#rows = []; this.releaseRopeArr(); disposeResources(rows); }
+}