mlx-bun 0.2.0 → 0.3.0

package/src/jobs/runner.ts

@@ -1,3 +1,4 @@
+import { CancellationSource } from "../engine/cancellation";
 // Job runner: turns a registered `JobRunner` into a tracked, logged job.
 // Two execution modes — both create a row, both stream NDJSON events to the
 // log, both own terminal status:
@@ -14,6 +15,8 @@
 // Port of optiq lab jobs.py `submit` + `_job_main`, split across this file
 // and job-entry.ts.
 
+import type { DisposableResource } from "../contracts/resources";
+import { createJobTaskClient } from "./task-client";
 import { appendFileSync } from "node:fs";
 import { fileURLToPath } from "node:url";
 import type { Emit, JobEvent, JobRunner } from "./types";
@@ -104,18 +107,28 @@ export interface SubmitResult {
  *  on success sets 'done'/progress 1/ended_at + emits 'done', on throw sets
  *  'failed'/error/ended_at + emits 'failed'. Returns once the task is
  *  scheduled (does not await completion — poll the row / tail the log). */
+interface InProcessWork { cancellation: CancellationSource; done: Promise<void>; }
+const inProcess = new WeakMap<JobStore, Set<InProcessWork>>();
+
 export function submitInProcess(
   store: JobStore,
   kind: string,
   config: Record<string, unknown>,
   outputPath?: string,
 ): SubmitResult {
+  if (closedStores.has(store)) throw new Error("job host is closed");
   const row = store.create(kind, config, outputPath);
   const emit = makeEmit(store, row.id, row.log_path);
   store.setStatus(row.id, "running");
   emit({ type: "started", ts: Date.now() });
 
   const runner = getRunner(kind);
+  const cancellation = new CancellationSource();
+  const running = inProcess.get(store) ?? new Set<InProcessWork>();
+  inProcess.set(store, running);
+  let finish!: () => void;
+  const work = { cancellation, done: new Promise<void>((resolve) => { finish = resolve; }) };
+  running.add(work);
   void (async () => {
     if (!runner) {
       const error = `no runner registered for kind "${kind}"`;
@@ -124,7 +137,7 @@ export function submitInProcess(
       return;
     }
     try {
-      const result = await runner(emit, JSON.parse(row.config_json));
+      const result = await createJobTaskClient(runner).run(JSON.parse(row.config_json), emit, cancellation);
       const out = result?.outputPath;
       if (out) store.setOutputPath(row.id, out);
       store.setProgress(row.id, 1);
@@ -135,7 +148,7 @@ export function submitInProcess(
       store.setStatus(row.id, "failed", { error, endedAt: nowIso() });
       emit({ type: "failed", error, ts: Date.now() });
     }
-  })();
+  })().finally(() => { running.delete(work); finish(); });
 
   return { jobId: row.id, outputPath };
 }
@@ -145,16 +158,25 @@ export function submitInProcess(
 // ---------------------------------------------------------------------------
 
 interface QueuedSpawn {
+  acquire?: (signal: AbortSignal) => Promise<DisposableResource>;
+  lease?: DisposableResource;
+  abort: AbortController;
+  proc?: Bun.Subprocess;
+  finished: Promise<void>;
+  finish(): void;
   store: JobStore;
   jobId: string;
   entry: string;
   bin: string;
+  compiled: boolean;
   spawn: typeof Bun.spawn;
   onComplete?: (jobId: string, code: number) => void;
 }
 
 let gpuLeaseHolder: string | null = null;
 const spawnQueue: QueuedSpawn[] = [];
+let activeSpawn: QueuedSpawn | undefined;
+const closedStores = new WeakSet<JobStore>();
 
 /** True while a GPU (subprocess) job is running — the server gates token
  *  generation on this so quantize/finetune don't fight live inference. */
@@ -168,6 +190,8 @@ export function currentGpuJob(): string | null {
 }
 
 export interface SubprocessOpts {
+  /** Host execution lease, held until child exit (including crashes). */
+  acquire?: (signal: AbortSignal) => Promise<DisposableResource>;
   /** Override the entry script (tests). Defaults to JOB_ENTRY_PATH. */
   entry?: string;
   /** Override the runtime binary (tests). Defaults to "bun". */
@@ -189,14 +213,20 @@ export function submitSubprocess(
   outputPath?: string,
   opts: SubprocessOpts = {},
 ): SubmitResult {
+  if (closedStores.has(store)) throw new Error("job host is closed");
   const row = store.create(kind, config, outputPath);
+  let finish!: () => void;
+  const finished = new Promise<void>((resolve) => { finish = resolve; });
   const item: QueuedSpawn = {
+    abort: new AbortController(), finished, finish,
     store,
     jobId: row.id,
     entry: opts.entry ?? JOB_ENTRY_PATH,
+    compiled: opts.entry === undefined && opts.bin === undefined && JOB_ENTRY_PATH.includes("$bunfs"),
     bin: opts.bin ?? "bun",
     spawn: opts.spawn ?? Bun.spawn,
     onComplete: opts.onComplete,
+    acquire: opts.acquire,
   };
   spawnQueue.push(item);
   drainQueue();
@@ -215,7 +245,7 @@ function spawnNow(item: QueuedSpawn): void {
   // a file-backed DB.
   let proc: Bun.Subprocess<"ignore", "pipe", "pipe">;
   try {
-    proc = spawn([bin, entry, jobId], {
+    proc = spawn(item.compiled ? [process.execPath, "__job", jobId] : [bin, entry, jobId], {
       stdout: "pipe",
       stderr: "pipe",
       env: {
@@ -230,14 +260,16 @@ function spawnNow(item: QueuedSpawn): void {
         error: errString(e),
         endedAt: nowIso(),
       });
+    } catch (error) {
+      console.error(`[jobs] failed to record spawn failure: ${errString(error)}`);
     } finally {
-      releaseLease(jobId);
+      releaseLease(item);
     }
     return;
   }
 
-  const row = store.get(jobId);
-  const logPath = row?.log_path;
+  item.proc = proc;
+  const logPath = store.get(jobId)?.log_path;
   const logLine = (line: string) => {
     if (!line || !logPath) return;
     const ev: JobEvent = { type: "log", line };
@@ -252,7 +284,7 @@ function spawnNow(item: QueuedSpawn): void {
     // code 0 ⇒ trust the child's terminal status (it set done/failed itself).
     // non-zero ⇒ if the row never reached terminal (crash before the wrapper
     // could write), force it failed.
-    if (code !== 0) {
+    try { if (code !== 0) {
       const cur = store.get(jobId);
       if (cur && (cur.status === "queued" || cur.status === "running")) {
         store.setStatus(jobId, "failed", {
@@ -260,15 +292,25 @@ function spawnNow(item: QueuedSpawn): void {
           endedAt: nowIso(),
         });
       }
+    } } catch (error) {
+      console.error(`[jobs] failed to record child exit: ${errString(error)}`);
+    } finally {
+      releaseLease(item);
+      try { item.onComplete?.(jobId, code); } catch {}
+      drainQueue();
     }
-    releaseLease(jobId);
-    try { item.onComplete?.(jobId, code); } catch {}
-    drainQueue();
   })();
 }
 
-function releaseLease(jobId: string): void {
-  if (gpuLeaseHolder === jobId) gpuLeaseHolder = null;
+function releaseLease(item: QueuedSpawn): void {
+  try { item.lease?.dispose(); }
+  catch (error) { console.error(`[jobs] execution lease cleanup failed: ${errString(error)}`); }
+  finally {
+    item.lease = undefined;
+    if (gpuLeaseHolder === item.jobId) gpuLeaseHolder = null;
+    if (activeSpawn === item) activeSpawn = undefined;
+    item.finish();
+  }
 }
 
 /** Run the next queued subprocess job if the lease is free. */
@@ -279,8 +321,51 @@ export function drainQueue(): void {
     // A synchronous spawn failure releases the lease. Keep draining
     // iteratively so a run of bad queue entries cannot recurse or wedge the
     // first startable job behind them.
-    spawnNow(next);
+    activeSpawn = next;
+    if (next.acquire) {
+      gpuLeaseHolder = next.jobId; // reserve FIFO order while inference drains
+      void Promise.resolve().then(() => next.acquire!(next.abort.signal)).then((lease) => {
+        next.lease = lease;
+        if (next.abort.signal.aborted) throw next.abort.signal.reason;
+        spawnNow(next);
+      }).catch((error) => {
+        try {
+          next.store.setStatus(next.jobId, "failed", { error: errString(error), endedAt: nowIso() });
+        } catch (recordError) {
+          console.error(`[jobs] failed to record admission failure: ${errString(recordError)}`);
+        } finally { releaseLease(next); }
+      }).finally(drainQueue);
+    } else spawnNow(next);
+  }
+}
+
+/** Dataset work stays in-process, but belongs to the application host too. */
+export async function closeInProcessJobs(store: JobStore): Promise<void> {
+  closedStores.add(store);
+  const active = [...inProcess.get(store) ?? []];
+  for (const work of active) work.cancellation.cancel("engine_closed");
+  await Promise.all(active.map((work) => work.done));
+}
+
+/** Stop only this host's managed GPU jobs. Await process death before releasing
+ * the execution lease; queued and admission-waiting jobs never spawn. */
+export async function closeSubprocessJobs(store: JobStore): Promise<void> {
+  closedStores.add(store);
+  for (let i = spawnQueue.length - 1; i >= 0; i--) {
+    const item = spawnQueue[i]!;
+    if (item.store !== store) continue;
+    spawnQueue.splice(i, 1);
+    store.setStatus(item.jobId, "failed", { error: "job host closed", endedAt: nowIso() });
+    item.finish();
   }
+  const active = activeSpawn;
+  if (!active || active.store !== store) return;
+  active.abort.abort(new Error("job host closed"));
+  const proc = active.proc;
+  if (proc) proc.kill("SIGTERM");
+  const force = proc ? setTimeout(() => { if (proc.exitCode === null) proc.kill("SIGKILL"); }, 3000) : undefined;
+  try { await active.finished; }
+  finally { if (force) clearTimeout(force); }
 }
 
 /** Read a child stream line-by-line, buffering partial trailing lines, and

package/src/jobs/task-client.ts

@@ -0,0 +1,19 @@
+import type { TaskClient } from "../contracts/completion";
+import { throwIfCancelled } from "../engine/cancellation";
+import type { JobEvent, JobRunner } from "./types";
+
+/** Numerical runners retain their artifact policy; the client owns delivery
+ * and cooperative cancellation at progress boundaries. No event buffering. */
+export function createJobTaskClient(runner: JobRunner): TaskClient<
+  Record<string, unknown>, JobEvent, { outputPath?: string } | void
+> {
+  return {
+    async run(request, report, cancellation) {
+      const check = () => { if (cancellation) throwIfCancelled(cancellation); };
+      check();
+      const result = await runner((event) => { check(); report(event); check(); }, request);
+      check();
+      return result;
+    },
+  };
+}

package/src/jobs/types.ts

@@ -1,58 +1,2 @@
-// Shared job-system contract. Both the job runner (src/jobs/*) and every
-// job producer (src/dataset/*, src/quantize/*, src/train/*) speak this
-// vocabulary. Mirrors the optiq lab SSE event grammar so the web wizards
-// port over mechanically (started → log/stage/metric → done|failed).
-
-/** A single streamed job event, serialized one-per-line into the job's
- *  NDJSON log and re-emitted to the browser as SSE `data:` frames. */
-export type JobEvent =
-  | { type: "started"; ts: number }
-  | { type: "log"; line: string }
-  | {
-      type: "stage";
-      stage: string;
-      progress?: number; // 0..1
-      message?: string;
-      [k: string]: unknown;
-    }
-  | {
-      type: "metric";
-      kind: "train" | "val";
-      step: number;
-      loss: number;
-      [k: string]: unknown; // grad_norm, learning_rate, tokens_per_sec, accuracy, margin, progress, message
-    }
-  | { type: "done"; ts: number; output_dir?: string; summary?: unknown; [k: string]: unknown }
-  | { type: "failed"; error: string; ts: number };
-
-/** Sink a job calls to report progress. Implementations append to the log
- *  file and (for stage/metric events carrying `progress`/`message`) update
- *  the SQLite row. Never throws — a logging failure must not kill a job. */
-export type Emit = (e: JobEvent) => void;
-
-/** The unit of work. Returns an optional output path recorded on the row.
- *  Receives the parsed `config` from the submit request. */
-export type JobRunner = (
-  emit: Emit,
-  config: Record<string, unknown>,
-) => Promise<{ outputPath?: string } | void>;
-
-/** Job kinds in the system. */
-export type JobKind = "quantize" | "finetune" | "dataset" | "noop" | "crash";
-
-export type JobStatus = "queued" | "running" | "done" | "failed" | "zombie";
-
-/** A persisted job row (shape returned by the HTTP `GET /api/jobs/:id`). */
-export interface JobRow {
-  id: string;
-  kind: string;
-  status: JobStatus;
-  config_json: string;
-  progress: number;
-  message: string | null;
-  log_path: string;
-  output_path: string | null;
-  error: string | null;
-  started_at: string;
-  ended_at: string | null;
-}
+/** Compatibility exports; producer and browser share the portable wire contract. */
+export type * from "../contracts/jobs";

package/src/kv-store.ts

@@ -1,5 +1,6 @@
+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):
@@ -121,9 +122,26 @@ 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 {
+  /** Absent in earlier v3 files, which used mlx-cache-v3. */
+  codecProvider?: string;
   formatVersion: 3;
   createdAt: number;
   tokens: number[];
@@ -143,7 +161,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 +169,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;
@@ -458,16 +476,50 @@ 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 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;
+    },
+  });
+}
+
+export const legacyCacheCodecs = createCacheCodecProvider("mlx-cache-v3", CACHE_CODECS);
 
-export const cacheHeadersTrimmable = (caches: readonly CacheHeaderEntry[]): boolean =>
-  caches.every((entry) => CACHE_CODECS[entry.kind].headerTrimmable(entry));
+export const cacheHeadersTrimmable = (
+  caches: readonly CacheHeaderEntry[], codecs: CacheCodecProvider = legacyCacheCodecs,
+): boolean => caches.every((entry) => codecs.forHeader(entry).headerTrimmable(entry));
 
-function snapshotCache(c: Cache, dataOffset: number): { entry: CacheHeaderEntry; sources: TensorSource[]; next: number } {
+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 +545,13 @@ 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 });
+    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 +568,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 +588,11 @@ 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 }));
+      out.push(codecs.forCache(c).clone(c, { view, liveView, mlaView, tripleView }));
+      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,14 +602,14 @@ 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;
@@ -561,7 +618,7 @@ function* saveKvCacheSteps(path: string, tokens: number[], caches: Cache[], meta
     // 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: 3, createdAt: Date.now(), ...meta, codecProvider: codecs.id, tokens, caches: entries,
     };
     const headerLen = new TextEncoder().encode(JSON.stringify(header)).length;
     const dataStart = alignUp(PREFIX_LEN + headerLen);
@@ -626,32 +683,34 @@ 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;
   }
 }
 
@@ -666,7 +725,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.
@@ -904,9 +963,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",
@@ -966,9 +1028,7 @@ export function loadKvCache(
   const caches: Cache[] = [];
   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}`);
+      const codec = codecs.forHeader(e);
       caches.push(codec.load(e, { path, arr, grownArr, triple }));
       // This entry's tensors are now owned by its cache — stop tracking them
       // (the catch must not double-free through both pending AND caches).

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 {