mlx-bun 0.4.0 → 0.5.0

package/src/fill/fill-session.ts

@@ -20,7 +20,7 @@
 // sources in priority order, clamps what they return, and counts. It never
 // touches MLX, the tokenizer, or the template. Row COMPILATION is
 // ./schema-rows.ts.
-import { flagOn, runtimeNumber, runtimeValue } from "../runtime-config";
+import { runtimeFlag, runtimeNumber, runtimeValue } from "../runtime-config";
 import { EchoSource, type EchoConfig } from "./echo-index";
 import type {
   Proposal, ProposalOrigin, ProposalPolicy, ProposalSource, TokenView,
@@ -153,7 +153,7 @@ export function fillEchoConfig(): FillEchoConfig {
  *  the engine's fill path. Off by default (the invariant it checks is the
  *  easiest bug in the feature, and the check costs a JS compare). */
 export function fillTraceEnabled(): boolean {
-  return flagOn("MLX_BUN_FILL_TRACE", false) || fillTracePath() !== null;
+  return runtimeFlag("MLX_BUN_FILL_TRACE", false) || fillTracePath() !== null;
 }
 
 /** MLX_BUN_FILL_TRACE=<file.jsonl> — additionally append one record per
@@ -319,12 +319,23 @@ export class FillSession {
   /** Feed one SAMPLED token, then ask the sources what follows. `budget` is
    *  the caller's remaining allowance (max_tokens − generated). */
   push(tokenId: number, budget: number = Number.POSITIVE_INFINITY): Proposal | null {
+    this.observe(tokenId);
+    return this.propose(budget);
+  }
+
+  /** Record sampled output independently of when the method requests proposals. */
+  observe(tokenId: number): void {
     this.stats.decodeSteps++;
     this.#append([tokenId]);
+  }
+
+  /** Select a continuation after a committed output boundary. Methods may
+   * request only verified proposals while retaining ordinary sampled policy. */
+  propose(budget: number = Number.POSITIVE_INFINITY, policy?: ProposalPolicy): Proposal | null {
     this.#budget = budget;
     for (const source of this.#sources) {
       const proposal = source.propose(this);
-      if (!proposal) continue;
+      if (!proposal || (policy && proposal.policy !== policy)) continue;
       if (proposal.branchStop) this.stats.branchStops++;
       const clamped = this.#clamp(proposal.ids, budget, proposal.origin === "echo");
       if (!clamped) continue;

package/src/generate.ts

@@ -27,7 +27,7 @@ import { denoiseAsync } from "./diffusion/diffusion-generate";
 import { bindLegacyDenoisingModel, type MlxDenoisingBinding } from "./backends/mlx/diffusion";
 import type { RuntimeModel } from "./model/factory";
 import {
-  assertMlxAutoregressiveBinding, bindLegacyAutoregressiveModel,
+  assertMlxAutoregressiveBinding, bindLegacyAutoregressiveModel, supportsCommittedAppendCache,
   type MlxAutoregressiveBinding, type MlxModelMemory, type MlxDecodeStep, type MlxTokenAppend,
 } from "./backends/mlx/autoregressive";
 import { createKvMaintenance } from "./backends/mlx/kv-maintenance";
@@ -61,6 +61,8 @@ import {
 } from "./fill/fill-session";
 
 export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions {
+  /** Application session affinity; independent of numerical cache identity. */
+  cacheSessionId?: string;
   /** Resolved host policy. Direct compatibility calls resolve from their
    * captured binding when this is absent. */
   decodePolicy?: Readonly<Pick<import("./contracts/execution").ResolvedExecution, "compiledDecode" | "grammarJump">>;
@@ -201,17 +203,16 @@ export function shouldUseGrammarJump(
 }
 
 /** Composition gate for token fast-forwarding, enforced inside the engine
- *  (the serve layer refuses more shapes it alone can see — a user-fixed seed,
- *  a mounted draft model, continuous placement: src/serve/request-prep.ts,
- *  src/serve/chat-stage.ts). Throws when MLX_BUN_FILL names the unbuilt echo
- *  tier — an operator who asked for K3c must not get K3b silently.
+ *  The serving planner owns request-level eligibility, including explicit
+ *  seeds in serial execution and mounted draft models. Shared strict fill
+ *  has its own grouped method and uses the same append-format declaration.
  *   - grammar: forced-token production is the grammar's job; two producers in
  *     one iteration is forbidden (asserted in the loop).
  *   - logprobs/top_logprobs: injected tokens are never sampled, so they have
  *     no distribution row (same rule as shouldUseGrammarJump).
  *   - promptEmbeddings: vision/audio prompts (and mRoPE) — untested shape.
- *   - kvBits/kvConfig/turboQuant: post-conversion multi-token append is
- *     L-generic but unvalidated; refuse in v1 rather than pay it silently. */
+ *   - affine KV: the model's append binding declares its supported formats.
+ *     Verify-policy quantized appends remain a separate gate. */
 export function shouldUseFill(
   options: Pick<
     GenerateOptions,
@@ -219,6 +220,7 @@ export function shouldUseFill(
     | "kvBits" | "kvConfig" | "turboQuant"
   >,
   runtime: RuntimeConfig = runtimeConfig(),
+  append?: Pick<MlxTokenAppend, "affineKvBits" | "turboQuantFormats"> | null,
 ): boolean {
   if (!options.fill) return false;
   if (resolveFillMode(runtime.value("MLX_BUN_FILL") ?? "off") === "off") return false;
@@ -226,9 +228,7 @@ export function shouldUseFill(
     !options.logprobs &&
     !(options.topLogprobs && options.topLogprobs > 0) &&
     options.promptEmbeddings === undefined &&
-    !options.kvBits &&
-    !options.kvConfig?.length &&
-    options.turboQuant === undefined;
+    supportsCommittedAppendCache(append, options);
 }
 
 /** Compatibility entry point for callers that perform a single conversion.
@@ -599,7 +599,7 @@ async function* generateInner(
   const {
     maxTokens = 512,
     eosTokenIds = binding.eosTokenIds,
-    prefillChunkSize = 2048,
+    prefillChunkSize = binding.prefillPolicy?.chunkSize(promptTokens.length) ?? 2048,
   } = options;
 
   try {
@@ -641,6 +641,7 @@ async function* generateInner(
   const ownsCache = !options.cache;
   const resuming = options.initialPendingToken !== undefined;
   let cache: Cache[] = options.cache ?? [];
+  let appender: MlxTokenAppend | null = null;
   try {
     if (ownsCache) cache = binding.makeCache();
     if (!cache.length) throw new Error("AR binding returned an empty cache");
@@ -655,6 +656,9 @@ async function* generateInner(
     // Replace fresh full-attention caches before any forward. The deepest
     // write is prompt + maxTokens - 1, so capacity covers every decode step.
     maybePageKv(cache, options, promptTokens.length + maxTokens);
+    appender = options.fill ? binding.createAppend?.({
+      hasAdapters: !!options.adapters?.length, pagedKv: !!options.pagedKv,
+    }) ?? null : null;
   } catch (error) {
     if (ownsCache) for (const state of cache) state.dispose();
     stepSampler.dispose();
@@ -668,7 +672,7 @@ async function* generateInner(
   // window) layers append multi-token writes through #updateConcat — O(window)
   // per append rather than the O(L) a plain ring pays — so v1 warns and skips
   // the whole feature for those models rather than paying it silently.
-  let fillOn = shouldUseFill(options, runtime);
+  let fillOn = shouldUseFill(options, runtime, appender);
   if (fillOn && cache.some((c) => c instanceof RotatingKVCache)) {
     fillOn = false;
     if (!warnedFillRotating) {
@@ -688,7 +692,7 @@ async function* generateInner(
   // restore a pre-round snapshot and bit-exactly replay the accepted prefix.
   // A model whose caches can do neither still gets assert-policy fills; verify
   // proposals are dropped and counted (stats.verifyUnsupported).
-  const verifyCapable = fillOn && cache.every(
+  const verifyCapable = fillOn && !options.kvBits && !options.kvConfig?.length && !options.turboQuant && cache.every(
     (c) => c.isTrimmable() || typeof rewindable(c).specRoundRollback === "function",
   );
   closeBatchSetup?.();
@@ -720,7 +724,6 @@ async function* generateInner(
   let pendingExtras: StepExtras | null = null;
   let nextExtras: StepExtras | null = null;
   let decoder: MlxDecodeStep | null = null;
-  let appender: MlxTokenAppend | null = null;
   let finished = false;
   let threw = false;
   let executionError: unknown;
@@ -823,8 +826,7 @@ async function* generateInner(
     nextPending = null;
     disposeStepExtras(nextExtras);
     nextExtras = null;
-    // No-op under the fill exclusions (kv quant refuses fill); kept so the
-    // boundary stays correct if those exclusions ever loosen.
+    // Conversion runs at the ordinary committed-token boundary as well.
     maintainKv(cache);
 
     let checkpointMs = 0;
@@ -840,11 +842,18 @@ async function* generateInner(
     let logitsAll: MlxArray | null = null;
     try {
       const chunkSize = (state: readonly Cache[]) => {
-        const maxChunk = appender?.maxChunkSize(state) ?? 1;
+        const maxChunk = Math.min(appender?.maxChunkSize(state) ?? 1,
+          maintainKv.maxAppendTokens?.(state) ?? Number.POSITIVE_INFINITY);
         return fill.appendChunkSize > 0 ? Math.min(fill.appendChunkSize, maxChunk) : maxChunk;
       };
+      const committedGraph = appender ? appender.forwardHidden.bind(appender) : graph.forwardHidden.bind(graph);
+      const forwardCommitted = maintainKv.maxAppendTokens ? async (ids: MlxArray, state: Cache[]) => {
+        const hidden = await committedGraph(ids, state);
+        try { maintainKv(state); return hidden; }
+        catch (error) { hidden.dispose(); throw error; }
+      } : committedGraph;
       hf = await appendFillHidden(
-        !verify && appender ? appender.forwardHidden.bind(appender) : graph.forwardHidden.bind(graph),
+        verify ? graph.forwardHidden.bind(graph) : forwardCommitted,
         cache, ids, verify ? () => ids.length : chunkSize,
       );
       const [, Lf, Hf] = hf.shape as [number, number, number];
@@ -1032,9 +1041,6 @@ async function* generateInner(
     decoder = options.decodePolicy?.compiledDecode === false ? null : binding.createDecode?.({
       hasAdapters: !!options.adapters?.length, pagedKv: !!options.pagedKv,
     }) ?? null;
-    appender = fillOn ? binding.createAppend?.({
-      hasAdapters: !!options.adapters?.length, pagedKv: !!options.pagedKv,
-    }) ?? null : null;
     let stop = false;
     /** Token id read eagerly at the top of the loop for grammar (reused for
      *  the yield, avoiding a second readback). -1 when grammar is off — the
@@ -1050,7 +1056,7 @@ async function* generateInner(
     const grammarJump = options.decodePolicy?.grammarJump ?? shouldUseGrammarJump(options, runtime);
     // Yield token zero before its own decode forward. The captured runtime
     // keeps this scheduling choice stable across consumer awaits.
-    const earlyFirstToken = runtime.value("MLX_BUN_EARLY_FIRST_TOKEN") === "1" &&
+    const earlyFirstToken = runtime.flag("MLX_BUN_EARLY_FIRST_TOKEN", true) &&
       maxTokens > 1 && !resuming && !fillOn && !options.grammar;
     while (!stop) {
       options.signal?.throwIfAborted();

package/src/grammar.ts

@@ -314,6 +314,19 @@ export class GrammarController {
     return this.terminated;
   }
 
+  /** Propose the forced string without advancing the matcher. Tokenization
+   * can differ at its first boundary; verification and the normal grammar
+   * sampler accept or reject the candidates before anything is committed. */
+  async proposeTokens(maxTokens: number): Promise<number[]> {
+    await this.ready();
+    if (this.terminated || !this.encodeRaw) return [];
+    return wasmQueue(async () => {
+      if (this.disposed) return [];
+      const text = this.matcher.findJumpForwardString();
+      return text ? this.encodeRaw!(text).slice(0, maxTokens) : [];
+    });
+  }
+
   /** Jump-forward decoding (SGLang's technique via xgrammar's
    *  findJumpForwardString; opt-in — MLX_BUN_GRAMMAR_JUMP=1): when the grammar
    *  admits exactly one continuation string (JSON keys/punctuation, enum

package/src/inference/prefill.ts

@@ -1,3 +1,9 @@
+/** Request-level chunk policy. Execution consumes the selected size; model
+ * backends can account for their temporary workspace without changing admission. */
+export interface PrefillPolicy {
+  chunkSize(promptTokens: number): number;
+}
+
 export interface PrefillPosition {
   readonly length: number;
   readonly position: number;

package/src/kv-scheme.ts

@@ -172,6 +172,28 @@ export class KvScheme {
     return "bf16";
   }
 
+  /** Configured attention storage, not live allocator occupancy. Recurrent
+   * layers retain their own state and never acquire an attention KV codec. */
+  describe(config: ModelConfig) {
+    const layers: Record<string, number> = {};
+    const attention = { global: 0, sliding_window: 0 };
+    let recurrent_layers = 0;
+    const perLayer = new Map(this.options.kvConfig?.map(entry => [entry.layerIdx, entry.bits]));
+    for (let layer = 0; layer < config.text.numHiddenLayers; layer++) {
+      const type = config.text.layerTypes[layer] ?? "full_attention";
+      if (type === "linear_attention") { recurrent_layers++; continue; }
+      const sliding = type === "sliding_attention";
+      attention[sliding ? "sliding_window" : "global"]++;
+      const bits = this.kind === "affine-config" ? perLayer.get(layer)
+        : this.kind === "affine-uniform" ? this.options.kvBits : undefined;
+      const codec = this.kind === "turbo" && !sliding
+        ? `turbo-k${this.options.turboQuant!.kBits}v${this.options.turboQuant!.vBits}`
+        : bits ? `kv${bits}` : "bf16";
+      layers[codec] = (layers[codec] ?? 0) + 1;
+    }
+    return { mode: this.label, layers, attention, recurrent_layers };
+  }
+
   bytesAt(config: ModelConfig, tokens: number): number {
     // TurboQuant stays conservatively billed as bf16 until its packed layout
     // exposes a stable projector.

package/src/kv-store.ts

@@ -43,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 {
@@ -67,6 +73,7 @@ const hash64 = (bytes: Uint8Array): string => Bun.hash(bytes).toString(16).padSt
 
 export type CacheKind =
   | "kv"
+  | "paged"
   | "rotating"
   | "qkv"
   | "rotating-qkv"
@@ -76,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[];
@@ -85,6 +93,8 @@ 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;
@@ -146,7 +156,7 @@ 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;
+  formatVersion: 3 | 4 | 5;
   createdAt: number;
   tokens: number[];
   caches: CacheHeaderEntry[];
@@ -430,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!),
@@ -457,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),
@@ -615,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[] = [];
@@ -652,6 +659,17 @@ function* saveKvCacheSteps(path: string, tokens: number[], caches: Cache[], meta
       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);
 
@@ -719,35 +737,27 @@ 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[];
@@ -755,23 +765,10 @@ export interface SpillItem {
   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;
@@ -790,7 +787,7 @@ 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,
@@ -891,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 && header.formatVersion !== 4)
+    if (header.formatVersion !== 3 && header.formatVersion !== 4 && header.formatVersion !== 5)
       throw new Error(`${path}: unsupported formatVersion ${header.formatVersion}`);
     return { ...header, dataStart };
   } finally {
@@ -997,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);
@@ -1036,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
@@ -1046,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 => {
@@ -1092,8 +1069,67 @@ export function loadKvCache(
     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, ...(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.