mlx-bun 0.3.0 → 0.5.0

package/src/kv-store.ts

@@ -1,3 +1,4 @@
+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/kv-cache.md).
@@ -42,8 +43,14 @@ import { cleanupFailure, disposeResources } from "./engine/resources";
 // whole entry again.
 
 import { openSync, writeSync, readSync, closeSync, fsyncSync, renameSync, rmSync } from "node:fs";
+import { pagedCacheCodec } from "./backends/mlx/paged-cache-codec";
+import { HostBuffer } from "./storage/host-buffer";
+import { readFileSync } from "node:fs";
+import { join, dirname } from "node:path";
+import { createHash } from "node:crypto";
 import { MmapFile, MADV_DONTNEED } from "./mmap";
 import { MlxArray } from "./mlx/array";
+import { kvWriter, kvReader, type KvWriteRequest } from "./storage/kv-writer";
 import type { Dtype } from "./mlx/ffi";
 import * as ops from "./mlx/ops";
 import {
@@ -66,6 +73,7 @@ const hash64 = (bytes: Uint8Array): string => Bun.hash(bytes).toString(16).padSt
 
 export type CacheKind =
   | "kv"
+  | "paged"
   | "rotating"
   | "qkv"
   | "rotating-qkv"
@@ -75,7 +83,8 @@ export type CacheKind =
   | "mla-dsa"
   | "mtp-mla";
 
-interface TensorSlot {
+export interface TensorSlot {
+  blocks?: Array<{ hash: string; bytes: number }>;
   off: number;
   bytes: number;
   shape: number[];
@@ -84,8 +93,11 @@ interface TensorSlot {
 }
 
 export interface CacheHeaderEntry {
+  paged?: { capacityTokens: number; blockSize: number; blockTable: number[]; numBlocks: number;
+    headDim: number; vHeadDim: number; dtype: Dtype; direct: boolean };
   kind: CacheKind;
   offset: number;
+  minimumReusableOffset?: number;
   /** rotating variants: ring write index */
   idx?: number;
   /** rotating variants: window size */
@@ -114,6 +126,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. */
@@ -139,10 +152,11 @@ export interface KvSaveMeta {
   };
 }
 
-export interface KvFileHeader extends KvSaveMeta {
+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;
+  formatVersion: 3 | 4 | 5;
   createdAt: number;
   tokens: number[];
   caches: CacheHeaderEntry[];
@@ -273,8 +287,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),
@@ -305,8 +319,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;
     },
@@ -332,8 +346,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),
@@ -363,8 +377,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,
@@ -426,7 +440,7 @@ const CACHE_CODECS = {
       if (state.length === 0) throw new Error("cannot clone an empty cache");
       try {
         const [kIdx, kScales, kZeros, vPacked, vScales] = state;
-        const clone = new TurboQuantKVCache(c.kBits, c.vBits);
+        const clone = new TurboQuantKVCache(c.kBits, c.vBits, c.fusedDecode);
         clone.restoreState({
           kIdx: context.view(kIdx!),
           kScales: context.view(kScales!),
@@ -453,6 +467,7 @@ const CACHE_CODECS = {
     },
     headerTrimmable: () => true,
   },
+  paged: pagedCacheCodec,
   mla: {
     matches: (cache) => cacheSignature(cache) === "kv:mla:target",
     snapshot: (cache, context) => snapshotMla("mla", cache, context),
@@ -515,6 +530,17 @@ export function createCacheCodecProvider(
 
 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));
@@ -547,6 +573,7 @@ function snapshotCache(c: Cache, dataOffset: number, codecs: CacheCodecProvider)
 
   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();
@@ -588,7 +615,9 @@ export function cloneKvCaches(caches: Cache[], codecs: CacheCodecProvider = lega
   const out: Cache[] = [];
   try {
     for (const c of caches) {
-      out.push(codecs.forCache(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) {
@@ -597,12 +626,8 @@ export function cloneKvCaches(caches: Cache[], codecs: CacheCodecProvider = lega
   return out;
 }
 
-/** Shared v3 writer core, one `yield` after each tensor write. The sync
- *  wrapper drains it in a tight loop; the async wrapper awaits a macrotask
- *  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, codecs: CacheCodecProvider): Generator<void, void, void> {
+/** Shared file layout and codec planning for synchronous and worker writes. */
+function planKvCache(tokens: number[], caches: Cache[], meta: KvSaveMeta, codecs: CacheCodecProvider): { header: KvFileHeader; sources: TensorSource[] } {
   // Plan pass: header entries + lazy tensor sources, NO bytes materialized.
   const entries: CacheHeaderEntry[] = [];
   const sources: TensorSource[] = [];
@@ -615,11 +640,36 @@ function* saveKvCacheSteps(path: string, tokens: number[], caches: Cache[], meta
       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, codecProvider: codecs.id, tokens, caches: entries,
+      formatVersion: attachments?.length ? 4 : 3, createdAt: Date.now(), ...identity,
+      ...(attachments?.length ? { attachments } : {}), codecProvider: codecs.id, tokens, caches: entries,
     };
+    return { header, sources };
+  } catch (error) {
+    for (const source of sources) if (source.disposeAfter) source.arr.dispose();
+    throw error;
+  }
+}
+
+function* saveKvCacheSteps(path: string, tokens: number[], caches: Cache[], meta: KvSaveMeta, codecs: CacheCodecProvider): Generator<void, void, void> {
+  const { header, sources } = planKvCache(tokens, caches, meta, codecs);
+  const entries = header.caches, attachments = header.attachments;
+  try {
     const headerLen = new TextEncoder().encode(JSON.stringify(header)).length;
     const dataStart = alignUp(PREFIX_LEN + headerLen);
 
@@ -639,7 +689,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);
@@ -687,58 +737,38 @@ export function saveKvCache(path: string, tokens: number[], caches: Cache[], met
   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).
- *
- *  `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. */
+/** Prepare immutable native storage on the owner thread, then pack, hash and
+ * write it on the CPU worker. The optional runner covers preparation only.
+ * Callers retain cache and attachment ownership until this promise settles. */
 export async function saveKvCacheAsync(
   path: string, tokens: number[], caches: Cache[], meta: KvSaveMeta = {},
   runStep?: <T>(step: () => T) => Promise<T>, codecs: CacheCodecProvider = legacyCacheCodecs,
+  storage: Pick<KvWriteRequest, "layout" | "blockBytes" | "segmented"> = {},
 ): Promise<void> {
-  const steps = saveKvCacheSteps(path, tokens, caches, meta, codecs);
+  const { header, sources } = planKvCache(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;
+    // Snapshot preparation belongs to the MLX owner. After publication the
+    // worker sees only immutable bytes and never evaluates or calls MLX.
+    const prepare = () => sources.map(source => source.arr.storageView());
+    const tensors = runStep ? await runStep(prepare) : prepare();
+    await kvWriter.write({ path, header, tensors, ...storage });
+  } finally {
+    for (const source of sources) if (source.disposeAfter) source.arr.dispose();
   }
 }
 
-/** One pending write-behind item: an entry's copied tokens + OWNED
- *  zero-copy cache clones (the queue disposes them on every exit path). */
+/** Queued immutable snapshot; the queue releases its shared views on settle. */
 export interface SpillItem {
+  attachments?: CheckpointAttachment[];
   tokens: number[];
   caches: Cache[];
   ns: string;
 }
 
-/** Bounded write-behind queue (2026-07-07 post-merge review fix).
- *
- *  Pending spill/snapshot clones pin their entries' GPU buffers until the
- *  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.
- *
- *  Policy: queued bytes are capped. Over cap, the OLDEST not-in-flight
- *  item is dropped and its clones disposed immediately — cache semantics
- *  (a dropped spill is a future cache miss, never a wrong result; the
- *  contention-free alternative to letting the flush cut into decode).
- *  The item being enqueued is never its own victim (soft cap: one
- *  oversized entry may exceed the cap alone rather than never spilling).
- *  Items run strictly serially in enqueue order via an internal chain;
- *  store failures are swallowed (cold tier is best-effort) but clones
- *  are disposed on every settle path. */
+/** Serial snapshot queue with optional capacity shedding. TieredPromptCache
+ * disables shedding and owns residency itself. Interrupted-generation
+ * persistence uses the cap for supersedable intervals. Queue ownership ends
+ * after each write settles; the cache can retain independent shared views. */
 type SpillRec = SpillItem & {
   bytes: number;
   dropped: boolean;
@@ -757,10 +787,12 @@ export class SpillQueue {
     readonly capBytes: number,
     /** Byte size of a clone set (prompt-cache's cacheBytes). */
     readonly bytesOf: (caches: Cache[]) => number,
-    /** The actual write (server passes storeAsync + the idle gate). */
+    /** The asynchronous storage operation. */
     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. */
@@ -776,7 +808,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,
     };
@@ -784,7 +816,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);
@@ -803,12 +835,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; }
@@ -818,7 +857,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 {
@@ -848,7 +888,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 && header.formatVersion !== 5)
       throw new Error(`${path}: unsupported formatVersion ${header.formatVersion}`);
     return { ...header, dataStart };
   } finally {
@@ -857,6 +897,7 @@ export function readKvHeader(path: string): KvFileHeader & { dataStart: number }
 }
 
 export interface LoadedKvCache {
+  attachments?: CheckpointAttachment[];
   tokens: number[];
   header: KvFileHeader;
   caches: Cache[];
@@ -953,19 +994,8 @@ function withStepCapacity(a: MlxArray, offset: number): MlxArray {
   }
 }
 
-/** Reload by STREAMED COPY (see the header note): every tensor is copied
- *  into an mlx-owned leaf; the mapping is read-only, its pages dropped
- *  per-tensor, and unmapped before returning — the caches own their bytes
- *  outright (no dtor contract, no pinned mapping, nothing to retain).
- *  `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 = {}, codecs: CacheCodecProvider = legacyCacheCodecs,
-): LoadedKvCache {
-  const header = readKvHeader(path);
+function validateKvHeader(path: string, header: KvFileHeader, model: { makeCache(): Cache[] },
+  expect: KvLoadExpect, codecs: CacheCodecProvider): void {
   if ((header.codecProvider ?? legacyCacheCodecs.id) !== codecs.id)
     throw new Error(`${path}: cache codec provider mismatch`);
   for (const entry of header.caches) codecs.forHeader(entry);
@@ -992,8 +1022,11 @@ export function loadKvCache(
     for (const c of proto) c.dispose();
   }
 
-  const mmap = MmapFile.open(path, "ro");
-  const dataStart = header.dataStart;
+}
+
+/** Decode cache layouts independently of the transport supplying arrays. */
+function restoreKvArrays(path: string, header: KvFileHeader, codecs: CacheCodecProvider,
+  readArray: (slot: TensorSlot) => MlxArray): LoadedKvCache {
   // Every tensor materialized for the CURRENT entry, drained on any throw
   // (2026-07-07 review: a hash mismatch under --ssd-cache-verify — or any
   // per-tensor failure — orphaned the entry's already-built tensors: the
@@ -1002,19 +1035,7 @@ export function loadKvCache(
   // is safe (dispose is idempotent).
   const pending: MlxArray[] = [];
   const arr = (slot: TensorSlot): MlxArray => {
-    const view = mmap.view(dataStart + slot.off, slot.bytes);
-    if (expect.verify && hash64(view) !== slot.hash)
-      throw new Error(`${path}: tensor hash mismatch at offset ${slot.off}`);
-    // mlx_array_new_data COPIES synchronously — the view only has to
-    // outlive this call, so the mapping can be advised/unmapped freely.
-    const a = MlxArray.fromBytesCopy(view, slot.shape, slot.dtype as Dtype);
-    pending.push(a);
-    // Drop the just-read clean pages so the load's transient stays at
-    // live entry + one tensor (ALIGN = 16 KiB == the arm64 page size, so
-    // every tensor offset is page-aligned by construction).
-    const len = Math.min(alignUp(slot.bytes), mmap.size - (dataStart + slot.off));
-    mmap.advise(dataStart + slot.off, len, MADV_DONTNEED);
-    return a;
+    const a = readArray(slot); pending.push(a); return a;
   };
   /** withStepCapacity whose (possibly fresh) result is pending-tracked. */
   const grownArr = (slot: TensorSlot, offset: number): MlxArray => {
@@ -1026,20 +1047,89 @@ 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 = codecs.forHeader(e);
-      caches.push(codec.load(e, { path, arr, grownArr, triple }));
+      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 } : {}) };
+}
+
+/** CPU-only asynchronous reads into unpublished allocations; no tensor copy
+ * crosses the event loop. Native array construction uses the shared codecs. */
+export async function loadKvCacheAsync(path: string, model: { makeCache(): Cache[] },
+  expect: KvLoadExpect = {}, codecs: CacheCodecProvider = legacyCacheCodecs): Promise<LoadedKvCache> {
+  const header = readKvHeader(path);
+  validateKvHeader(path, header, model, expect, codecs);
+  const slots = [...header.caches, ...(header.attachments ?? [])].flatMap(e => e.tensors);
+  const buffers = new Map<TensorSlot, HostBuffer>();
+  try {
+    for (const slot of slots) buffers.set(slot, new HostBuffer(slot.bytes));
+    await kvReader.read({ path, header, pointers: slots.map(slot => buffers.get(slot)!.pointer), verify: expect.verify });
+    return restoreKvArrays(path, header, codecs, slot =>
+      MlxArray.adoptHostBuffer(buffers.get(slot)!, slot.shape, slot.dtype as Dtype));
+  } finally { for (const buffer of buffers.values()) buffer.dispose(); }
+}
+
+/** Reload by STREAMED COPY (see the header note): every tensor is copied
+ *  into an mlx-owned leaf; the mapping is read-only, its pages dropped
+ *  per-tensor, and unmapped before returning — the caches own their bytes
+ *  outright (no dtor contract, no pinned mapping, nothing to retain).
+ *  `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 = {}, codecs: CacheCodecProvider = legacyCacheCodecs,
+): LoadedKvCache {
+  const header = readKvHeader(path);
+  validateKvHeader(path, header, model, expect, codecs);
+  if (header.formatVersion === 5) {
+    return restoreKvArrays(path, header, codecs, slot => {
+      const bytes = new Uint8Array(slot.bytes);
+      let offset = 0;
+      for (const block of slot.blocks ?? []) {
+        if (!/^[a-f0-9]{64}$/.test(block.hash) || block.bytes <= 0 || offset + block.bytes > bytes.length)
+          throw new Error(`${path}: invalid KV block reference`);
+        const part = readFileSync(join(dirname(path), "blocks", block.hash));
+        if (part.length !== block.bytes || (expect.verify && createHash("sha256").update(part).digest("hex") !== block.hash))
+          throw new Error(`${path}: KV block mismatch`);
+        bytes.set(part, offset); offset += part.length;
+      }
+      if (offset !== bytes.length)
+        throw new Error(`${path}: KV tensor mismatch`);
+      return MlxArray.fromBytesCopy(bytes, slot.shape, slot.dtype as Dtype);
+    });
+  }
+  const mmap = MmapFile.open(path, "ro"), dataStart = header.dataStart;
+  try {
+    return restoreKvArrays(path, header, codecs, slot => {
+      const view = mmap.view(dataStart + slot.off, slot.bytes);
+      if (expect.verify && hash64(view) !== slot.hash)
+        throw new Error(`${path}: tensor hash mismatch at offset ${slot.off}`);
+      const array = MlxArray.fromBytesCopy(view, slot.shape, slot.dtype as Dtype);
+      const len = Math.min(alignUp(slot.bytes), mmap.size - (dataStart + slot.off));
+      mmap.advise(dataStart + slot.off, len, MADV_DONTNEED);
+      return array;
+    });
+  } finally { mmap.unmap(); }
 }

package/src/lab/README.md

@@ -14,5 +14,9 @@ fate (PLAN.md "Maintainability program", D6 darlings). Rules, gate-enforced by
 | 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 |
+| `paged-kv/` | `--paged-kv`; direct read via `MLX_BUN_PAGED_ATTN` | gathered storage is bit-exact; direct Metal is tolerance-tested, with paired M1/M4 measurements in benchmarks.md |
 | `expert-trace/` | `MLX_BUN_EXPERT_TRACE` hook in gemma4 | closed study (expert-offload) |
+
+The optional paged layout has a separate backend persistence codec. Its
+allowlisted import keeps `kv-store` independent of row-execution modules and
+preserves RAM/SSD reuse without promoting direct-attention numerics to L1.

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

@@ -0,0 +1,117 @@
+import { MlxArray } from "../../mlx/array";
+import { Dtype } from "../../mlx/ffi";
+import { MetalKernel } from "../../mlx/metal-kernel";
+import * as ops from "../../mlx/ops";
+import type { KvAttentionView, Mask } from "../../model/gemma4-base";
+import { disposeResources } from "../../engine/resources";
+import type { BlockPool } from "./paged-kv";
+
+// Each SIMD group reduces one query/head over a bounded sequence partition.
+// Only partial output vectors and normalizers materialize, never gathered KV.
+const partialSpec = { name: "paged_attention_partial", ensureRowContiguous: true,
+  inputNames: ["q", "keys", "values", "table", "scale"], outputNames: ["out", "normal"], source: `
+  uint lane = thread_position_in_threadgroup.x;
+  uint hq = threadgroup_position_in_grid.y % HQ;
+  uint query = threadgroup_position_in_grid.y / HQ;
+  uint part = threadgroup_position_in_grid.z;
+  uint hk = hq / (HQ / HK);
+  float accum[(DV + 31) / 32];
+  for (uint j = 0; j < (DV + 31) / 32; ++j) accum[j] = 0;
+  float maximum = -INFINITY, total = 0;
+  uint length = uint(scale[1]);
+  uint count = (length + PART - 1) / PART;
+  uint end = min(uint(length - N + query + 1), (part + 1) * PART);
+  for (uint token = part * PART; token < end; ++token) {
+    uint physical = table[token / BS];
+    uint keyBase = ((physical * HK + hk) * BS + token % BS) * DK;
+    float score = 0;
+    for (uint d = lane; d < DK; d += 32)
+      score += float(q[(hq * N + query) * DK + d]) * float(keys[keyBase + d]);
+    score = simd_sum(score) * scale[0];
+    float next = max(maximum, score);
+    float previous = exp(maximum - next), weight = exp(score - next);
+    uint valueBase = ((physical * HK + hk) * BS + token % BS) * DV;
+    for (uint j = 0; j < (DV + 31) / 32; ++j) {
+      uint d = j * 32 + lane;
+      if (d < DV) accum[j] = accum[j] * previous + weight * float(values[valueBase + d]);
+    }
+    total = total * previous + weight; maximum = next;
+  }
+  uint base = ((hq * N + query) * count + part);
+  for (uint j = 0; j < (DV + 31) / 32; ++j) {
+    uint d = j * 32 + lane;
+    if (d < DV) out[base * DV + d] = accum[j];
+  }
+  if (lane == 0) { normal[base * 2] = maximum; normal[base * 2 + 1] = total; }
+` };
+const partial = new MetalKernel(partialSpec);
+const quantizedPartial = new MetalKernel({ ...partialSpec, name: "paged_attention_affine_partial",
+  inputNames: [...partialSpec.inputNames, "ks", "kb", "vs", "vb"],
+  source: partialSpec.source
+    .replace("float(keys[keyBase + d])", "float(InT(float(ks[(keyBase+d)/GROUP]) * float((keys[(keyBase+d)/PACK] >> (((keyBase+d)%PACK)*BITS)) & ((1u<<BITS)-1u)) + float(kb[(keyBase+d)/GROUP])))")
+    .replace("float(values[valueBase + d])", "float(InT(float(vs[(valueBase+d)/GROUP]) * float((values[(valueBase+d)/PACK] >> (((valueBase+d)%PACK)*BITS)) & ((1u<<BITS)-1u)) + float(vb[(valueBase+d)/GROUP])))"),
+});
+const merge = new MetalKernel({ name: "paged_attention_merge", inputNames: ["parts", "normal"], outputNames: ["out"], source: `
+  uint d = thread_position_in_grid.x;
+  uint row = thread_position_in_grid.y;
+  if (d >= DV) return;
+  uint count = normal_shape[2];
+  float maximum = -INFINITY;
+  for (uint part = 0; part < count; ++part) maximum = max(maximum, normal[(row * count + part) * 2]);
+  float sum = 0, value = 0;
+  for (uint part = 0; part < count; ++part) {
+    uint base = row * count + part;
+    float weight = exp(normal[base * 2] - maximum);
+    sum += normal[base * 2 + 1] * weight;
+    value += parts[base * DV + d] * weight;
+  }
+  out[row * DV + d] = OutT(value / sum);
+` });
+
+/** Independent immutable view: later appends or row retirement cannot alter it. */
+export function pagedAttentionView(pool: BlockPool, blockTable: number[], length: number,
+  direct: boolean): KvAttentionView {
+  const keys = ops.contiguous(pool.keys), values = ops.contiguous(pool.values);
+  const table = MlxArray.fromInt32(Int32Array.from(blockTable), [blockTable.length]);
+  const bs = pool.blockSize, quantization = pool.quantization;
+  const extra = quantization ? [pool.keyScales!, pool.keyBiases!, pool.valueScales!, pool.valueBiases!].map(a => ops.contiguous(a)) : [];
+  const dkStored = pool.headDim, dvStored = pool.vHeadDim;
+  const dtype = pool.dtype;
+  return {
+    attend(q: MlxArray, scale: number, mask: Mask): MlxArray {
+      const [, hq, n, dk] = q.shape as [number, number, number, number];
+      const hk = keys.shape[1]!, dv = dvStored;
+      // Long prefill and explicit masks use the established attention oracle.
+      // This is shape dispatch within one attention interface, not scheduling.
+      if (!direct || n > 8 || mask.mode === "array") {
+        const gather = (tensor: MlxArray) => {
+          using picked = ops.takeAxis(tensor, table, 0);
+          using transposed = ops.transposeAxes(picked, [1, 0, 2, 3]);
+          using flat = ops.reshape(transposed, [1, hk, blockTable.length * bs, tensor.shape[3]!]);
+          return flat.slice([0, 0, 0, 0], [1, hk, length, tensor.shape[3]!]);
+        };
+        using k = gather(keys), v = gather(values);
+        if (!quantization) return ops.sdpa(q, k, v, scale, mask.mode, mask.arr);
+        using ks = gather(extra[0]!), kb = gather(extra[1]!), vs = gather(extra[2]!), vb = gather(extra[3]!);
+        using kd = ops.dequantize(k, ks, kb, { ...quantization, mode: "affine" });
+        using vd = ops.dequantize(v, vs, vb, { ...quantization, mode: "affine" });
+        return ops.sdpa(q, kd, vd, scale, mask.mode, mask.arr);
+      }
+      const part = 64, parts = Math.ceil(length / part);
+      using multiplier = MlxArray.fromFloat32(new Float32Array([scale, length]), [2]);
+      const [out, normal] = (quantization ? quantizedPartial : partial).apply([q, keys, values, table, multiplier, ...extra], {
+        outputs: [{ shape: [hq, n, parts, dv], dtype: Dtype.float32 }, { shape: [hq, n, parts, 2], dtype: Dtype.float32 }],
+        grid: [32, hq * n, parts], threadGroup: [32, 1, 1],
+        templateDtypes: quantization ? { InT: dtype } : {},
+        templateInts: { HQ: hq, HK: hk, N: n, DK: dkStored, DV: dv, BS: bs, PART: part,
+          ...(quantization ? { BITS: quantization.bits, GROUP: quantization.groupSize, PACK: 32 / quantization.bits } : {}) },
+      });
+      try {
+        return merge.apply([out!, normal!], { outputs: [{ shape: [1, hq, n, dv], dtype: q.dtype }],
+          grid: [Math.ceil(dv / 32) * 32, hq * n, 1], threadGroup: [32, 1, 1],
+          templateDtypes: { OutT: q.dtype }, templateInts: { DV: dv } })[0]!;
+      } finally { disposeResources([out!, normal!]); }
+    },
+    dispose() { disposeResources([keys, values, table, ...extra]); },
+  };
+}