mlx-bun 0.0.13 → 0.3.0

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):
@@ -46,7 +47,7 @@ import { MlxArray } from "./mlx/array";
 import type { Dtype } from "./mlx/ffi";
 import * as ops from "./mlx/ops";
 import {
-  type Cache, KVCache, RotatingKVCache,
+  type Cache, cacheSignature, KVCache, RotatingKVCache,
   QuantizedKVCache, RotatingQuantizedKVCache,
   TurboQuantKVCache, type TurboQuantTensor,
 } from "./model/gemma4-base";
@@ -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[];
@@ -142,7 +160,366 @@ const alignUp = (n: number) => Math.ceil(n / ALIGN) * ALIGN;
  *  as-laid-out with ringIdx. Slice handles are lazy graph nodes (no GPU
  *  materialization until rawBytes at write time). */
 interface TensorSource { arr: MlxArray; disposeAfter: boolean }
-function snapshotCache(c: Cache, dataOffset: number): { entry: CacheHeaderEntry; sources: TensorSource[]; next: number } {
+
+export interface SnapshotContext {
+  slots: TensorSlot[];
+  push(a: MlxArray, disposeAfter: boolean): void;
+  liveSlice(a: MlxArray, upTo: number): MlxArray;
+  liveMlaSlice(a: MlxArray, upTo: number): MlxArray;
+  pushTriple(t: ops.QuantizedTensor, upTo: number | null): void;
+}
+
+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;
+}
+
+export interface LoadContext {
+  path: string;
+  arr(slot: TensorSlot): MlxArray;
+  grownArr(slot: TensorSlot, offset: number): MlxArray;
+  triple(slots: TensorSlot[], at: number): ops.QuantizedTensor;
+}
+
+export interface CacheCodec {
+  matches(cache: Cache): boolean;
+  snapshot(cache: Cache, context: SnapshotContext): CacheHeaderEntry;
+  clone(cache: Cache, context: CloneContext): Cache;
+  load(entry: CacheHeaderEntry, context: LoadContext): Cache;
+  headerTrimmable(entry: CacheHeaderEntry): boolean;
+}
+
+const requireMlaState = (cache: Glm52Cache, operation: string): void => {
+  if (!cache.latent || !cache.rope || cache.batchSize === null || cache.offset === 0)
+    throw new Error(`cannot ${operation} an empty GLM compressed cache`);
+  if (cache.dsa && (!cache.dsa.data || cache.dsa.offset !== cache.offset))
+    throw new Error(`cannot ${operation} misaligned GLM DSA state`);
+};
+
+const snapshotMla = (
+  kind: Extract<CacheKind, "mla" | "mla-dsa" | "mtp-mla">,
+  cache: Cache,
+  context: SnapshotContext,
+): CacheHeaderEntry => {
+  const c = cache as Glm52Cache;
+  requireMlaState(c, "persist");
+  context.push(context.liveMlaSlice(c.latent!, c.offset), true);
+  context.push(context.liveMlaSlice(c.rope!, c.offset), true);
+  if (c.dsa) context.push(context.liveMlaSlice(c.dsa.data!, c.offset), true);
+  return {
+    kind,
+    offset: c.offset,
+    kvLoraRank: c.kvLoraRank,
+    ropeHeadDim: c.ropeHeadDim,
+    ...(c.dsa ? { dsaHeadDim: c.dsa.headDim } : {}),
+    maxTokens: c.maxTokens,
+    tensors: context.slots,
+  };
+};
+
+const cloneMla = (cache: Cache, context: CloneContext): Cache => {
+  const c = cache as Glm52Cache;
+  requireMlaState(c, "clone");
+  const clone = new MLACache({
+    kvLoraRank: c.kvLoraRank,
+    ropeHeadDim: c.ropeHeadDim,
+    ...(c.dsa ? { dsa: { headDim: c.dsa.headDim } } : {}),
+    maxTokens: c.maxTokens,
+    role: c.role,
+  });
+  clone.restoreCompressedState(
+    context.mlaView(c.latent!, c.offset),
+    context.mlaView(c.rope!, c.offset),
+    c.dsa ? context.mlaView(c.dsa.data!, c.offset) : null,
+    c.offset,
+  );
+  return clone;
+};
+
+const loadMla = (entry: CacheHeaderEntry, context: LoadContext): Cache => {
+  const hasDsa = entry.kind === "mla-dsa";
+  if (entry.kind === "mtp-mla" && entry.dsaHeadDim !== undefined)
+    throw new Error(`${context.path}: native MTP cache cannot contain DSA state`);
+  const cache = new MLACache({
+    kvLoraRank: entry.kvLoraRank!,
+    ropeHeadDim: entry.ropeHeadDim!,
+    ...(hasDsa ? { dsa: { headDim: entry.dsaHeadDim! } } : {}),
+    maxTokens: entry.maxTokens,
+    role: entry.kind === "mtp-mla" ? "mtp" : "target",
+  });
+  cache.restoreCompressedState(
+    context.arr(entry.tensors[0]!),
+    context.arr(entry.tensors[1]!),
+    hasDsa ? context.arr(entry.tensors[2]!) : null,
+    entry.offset,
+  );
+  return cache;
+};
+
+/** One authoritative persistence algebra. Adding a cache storage kind must
+ *  define all four operations here; save/clone/load and SSD eligibility
+ *  therefore cannot grow independent instanceof/switch chains. */
+const CACHE_CODECS = {
+  kv: {
+    matches: (cache) => cacheSignature(cache) === "kv:plain",
+    snapshot: (cache, context) => {
+      const c = cache as KVCache;
+      if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
+      context.push(context.liveSlice(c.keys, c.offset), true);
+      context.push(context.liveSlice(c.values, c.offset), true);
+      return { kind: "kv", offset: c.offset, tensors: context.slots };
+    },
+    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();
+      clone.restoreState(
+        context.liveView(c.keys, c.offset),
+        context.liveView(c.values, c.offset),
+        c.offset,
+      );
+      return clone;
+    },
+    load: (entry, context) => {
+      const cache = new KVCache();
+      cache.restoreState(
+        context.grownArr(entry.tensors[0]!, entry.offset),
+        context.grownArr(entry.tensors[1]!, entry.offset),
+        entry.offset,
+      );
+      return cache;
+    },
+    headerTrimmable: () => true,
+  },
+  rotating: {
+    matches: (cache) => cacheSignature(cache) === "kv:rotating-plain",
+    snapshot: (cache, context) => {
+      const c = cache as RotatingKVCache;
+      if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
+      context.push(c.keys, false);
+      context.push(c.values, false);
+      return { kind: "rotating", offset: c.offset, idx: c.ringIdx,
+        maxSize: c.maxSize, tensors: context.slots };
+    },
+    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);
+      clone.restoreState(context.view(c.keys), context.view(c.values), c.offset, c.ringIdx);
+      return clone;
+    },
+    load: (entry, context) => {
+      const cache = new RotatingKVCache(entry.maxSize!);
+      cache.restoreState(
+        context.arr(entry.tensors[0]!), context.arr(entry.tensors[1]!),
+        entry.offset, entry.idx!,
+      );
+      return cache;
+    },
+    headerTrimmable: (entry) => entry.offset < (entry.maxSize ?? 0),
+  },
+  qkv: {
+    matches: (cache) => cacheSignature(cache).startsWith("kv:quant:"),
+    snapshot: (cache, context) => {
+      const c = cache as QuantizedKVCache;
+      if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
+      context.pushTriple(c.keys, c.offset);
+      context.pushTriple(c.values, c.offset);
+      return { kind: "qkv", offset: c.offset, groupSize: c.groupSize,
+        bits: c.bits, tensors: context.slots };
+    },
+    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);
+      clone.restoreState(
+        context.tripleView(c.keys, c.offset),
+        context.tripleView(c.values, c.offset),
+        c.offset,
+      );
+      return clone;
+    },
+    load: (entry, context) => {
+      const cache = new QuantizedKVCache(entry.groupSize!, entry.bits!);
+      cache.restoreState(
+        context.triple(entry.tensors, 0), context.triple(entry.tensors, 3), entry.offset,
+      );
+      return cache;
+    },
+    headerTrimmable: () => true,
+  },
+  "rotating-qkv": {
+    matches: (cache) => cacheSignature(cache).startsWith("kv:rotating-quant:"),
+    snapshot: (cache, context) => {
+      const c = cache as RotatingQuantizedKVCache;
+      if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
+      context.pushTriple(c.keys, null);
+      context.pushTriple(c.values, null);
+      return { kind: "rotating-qkv", offset: c.offset, idx: c.ringIdx,
+        maxSize: c.maxSize, groupSize: c.groupSize, bits: c.bits,
+        tensors: context.slots };
+    },
+    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);
+      clone.restoreState(
+        context.tripleView(c.keys, null), context.tripleView(c.values, null),
+        c.offset, c.ringIdx,
+      );
+      return clone;
+    },
+    load: (entry, context) => {
+      const cache = new RotatingQuantizedKVCache(
+        entry.maxSize!, entry.groupSize!, entry.bits!,
+      );
+      cache.restoreState(
+        context.triple(entry.tensors, 0), context.triple(entry.tensors, 3),
+        entry.offset, entry.idx!,
+      );
+      return cache;
+    },
+    headerTrimmable: (entry) => entry.offset < (entry.maxSize ?? 0),
+  },
+  ssm: {
+    matches: (cache) => cacheSignature(cache) === "ssm",
+    snapshot: (cache, context) => {
+      const c = cache as SSMCache;
+      if (!c.conv || !c.recurrent) throw new Error("cannot persist an empty cache");
+      context.push(c.conv, false);
+      context.push(c.recurrent, false);
+      return { kind: "ssm", offset: c.offset, tensors: context.slots };
+    },
+    clone: (cache, context) => {
+      const c = cache as SSMCache;
+      if (!c.conv || !c.recurrent) throw new Error("cannot clone an empty cache");
+      const clone = new SSMCache();
+      clone.conv = context.view(c.conv);
+      clone.recurrent = context.view(c.recurrent);
+      clone.offset = c.offset;
+      return clone;
+    },
+    load: (entry, context) => {
+      const cache = new SSMCache();
+      cache.conv = context.arr(entry.tensors[0]!);
+      cache.recurrent = context.arr(entry.tensors[1]!);
+      cache.offset = entry.offset;
+      return cache;
+    },
+    headerTrimmable: () => false,
+  },
+  turboquant: {
+    matches: (cache) => cacheSignature(cache).startsWith("kv:turboquant:"),
+    snapshot: (cache, context) => {
+      const c = cache as TurboQuantKVCache;
+      const state = c.state();
+      if (state.length === 0) throw new Error("cannot persist an empty cache");
+      for (const array of state) context.push(array, true);
+      return { kind: "turboquant", offset: c.offset, kBits: c.kBits,
+        vBits: c.vBits, headDim: c.headDim ?? 0, tensors: context.slots };
+    },
+    clone: (cache, context) => {
+      const c = cache as TurboQuantKVCache;
+      const state = c.state();
+      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);
+        clone.restoreState({
+          kIdx: context.view(kIdx!),
+          kScales: context.view(kScales!),
+          kZeros: context.view(kZeros!),
+          vPacked: context.view(vPacked!),
+          vScales: context.view(vScales!),
+        }, c.offset, c.headDim!);
+        return clone;
+      } finally {
+        for (const array of state) array.dispose();
+      }
+    },
+    load: (entry, context) => {
+      const cache = new TurboQuantKVCache(entry.kBits!, entry.vBits!);
+      const tensors: TurboQuantTensor = {
+        kIdx: context.arr(entry.tensors[0]!),
+        kScales: context.arr(entry.tensors[1]!),
+        kZeros: context.arr(entry.tensors[2]!),
+        vPacked: context.arr(entry.tensors[3]!),
+        vScales: context.arr(entry.tensors[4]!),
+      };
+      cache.restoreState(tensors, entry.offset, entry.headDim!);
+      return cache;
+    },
+    headerTrimmable: () => true,
+  },
+  mla: {
+    matches: (cache) => cacheSignature(cache) === "kv:mla:target",
+    snapshot: (cache, context) => snapshotMla("mla", cache, context),
+    clone: cloneMla,
+    load: loadMla,
+    headerTrimmable: () => true,
+  },
+  "mla-dsa": {
+    matches: (cache) => cacheSignature(cache) === "kv:mla:target:dsa",
+    snapshot: (cache, context) => snapshotMla("mla-dsa", cache, context),
+    clone: cloneMla,
+    load: loadMla,
+    headerTrimmable: () => true,
+  },
+  "mtp-mla": {
+    matches: (cache) => cacheSignature(cache) === "kv:mla:mtp",
+    snapshot: (cache, context) => snapshotMla("mtp-mla", cache, context),
+    clone: cloneMla,
+    load: loadMla,
+    headerTrimmable: () => true,
+  },
+} satisfies Record<CacheKind, CacheCodec>;
+
+/** 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[], codecs: CacheCodecProvider = legacyCacheCodecs,
+): boolean => caches.every((entry) => codecs.forHeader(entry).headerTrimmable(entry));
+
+function snapshotCache(c: Cache, dataOffset: number, codecs: CacheCodecProvider): { entry: CacheHeaderEntry; sources: TensorSource[]; next: number } {
   const slots: TensorSlot[] = [];
   const sources: TensorSource[] = [];
   let cursor = dataOffset;
@@ -168,65 +545,13 @@ function snapshotCache(c: Cache, dataOffset: number): { entry: CacheHeaderEntry;
     }
   };
 
-  let entry: CacheHeaderEntry;
-  if (c instanceof Glm52Cache) {
-    if (!c.latent || !c.rope || c.batchSize === null || c.offset === 0)
-      throw new Error("cannot persist an empty GLM compressed cache");
-    push(liveMlaSlice(c.latent, c.offset), true);
-    push(liveMlaSlice(c.rope, c.offset), true);
-    if (c.dsa) {
-      if (!c.dsa.data || c.dsa.offset !== c.offset)
-        throw new Error("cannot persist misaligned GLM DSA state");
-      push(liveMlaSlice(c.dsa.data, c.offset), true);
-    }
-    entry = {
-      kind: c.role === "mtp" ? "mtp-mla" : c.dsa ? "mla-dsa" : "mla",
-      offset: c.offset,
-      kvLoraRank: c.kvLoraRank,
-      ropeHeadDim: c.ropeHeadDim,
-      ...(c.dsa ? { dsaHeadDim: c.dsa.headDim } : {}),
-      maxTokens: c.maxTokens,
-      tensors: slots,
-    };
-  } else if (c instanceof TurboQuantKVCache) {
-    const t = c.state();
-    if (t.length === 0) throw new Error("cannot persist an empty cache");
-    const [kIdx, kScales, kZeros, vPacked, vScales] = t as [MlxArray, MlxArray, MlxArray, MlxArray, MlxArray];
-    // state() already trims to offset — dispose these lazy slice views
-    // once written (same disposeAfter contract as liveSlice elsewhere).
-    for (const a of [kIdx, kScales, kZeros, vPacked, vScales]) push(a, true);
-    entry = { kind: "turboquant", offset: c.offset, kBits: c.kBits, vBits: c.vBits,
-      headDim: c.headDim ?? 0, tensors: slots };
-  } else if (c instanceof RotatingQuantizedKVCache) {
-    if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
-    pushTriple(c.keys, null);
-    pushTriple(c.values, null);
-    entry = { kind: "rotating-qkv", offset: c.offset, idx: c.ringIdx, maxSize: c.maxSize,
-      groupSize: c.groupSize, bits: c.bits, tensors: slots };
-  } else if (c instanceof QuantizedKVCache) {
-    if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
-    pushTriple(c.keys, c.offset);
-    pushTriple(c.values, c.offset);
-    entry = { kind: "qkv", offset: c.offset, groupSize: c.groupSize, bits: c.bits, tensors: slots };
-  } else if (c instanceof RotatingKVCache) {
-    if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
-    push(c.keys, false);
-    push(c.values, false);
-    entry = { kind: "rotating", offset: c.offset, idx: c.ringIdx, maxSize: c.maxSize, tensors: slots };
-  } else if (c instanceof SSMCache) {
-    if (!c.conv || !c.recurrent) throw new Error("cannot persist an empty cache");
-    push(c.conv, false);
-    push(c.recurrent, false);
-    entry = { kind: "ssm", offset: c.offset, tensors: slots };
-  } else if (c instanceof KVCache) {
-    if (!c.keys || !c.values) throw new Error("cannot persist an empty cache");
-    push(liveSlice(c.keys, c.offset), true);
-    push(liveSlice(c.values, c.offset), true);
-    entry = { kind: "kv", offset: c.offset, tensors: slots };
-  } else {
-    throw new Error("unknown cache type");
+  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;
   }
-  return { entry, sources, next: cursor };
 }
 
 /** Persist `caches` (+ the exact token prefix they encode) to `path`.
@@ -243,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),
@@ -261,69 +588,11 @@ export function cloneKvCaches(caches: Cache[]): Cache[] {
   const out: Cache[] = [];
   try {
     for (const c of caches) {
-      if (c instanceof Glm52Cache) {
-        if (!c.latent || !c.rope || c.batchSize === null || c.offset === 0)
-          throw new Error("cannot clone an empty GLM compressed cache");
-        const n = new MLACache({
-          kvLoraRank: c.kvLoraRank,
-          ropeHeadDim: c.ropeHeadDim,
-          ...(c.dsa ? { dsa: { headDim: c.dsa.headDim } } : {}),
-          maxTokens: c.maxTokens,
-          role: c.role,
-        });
-        n.restoreCompressedState(
-          mlaView(c.latent, c.offset),
-          mlaView(c.rope, c.offset),
-          c.dsa ? mlaView(c.dsa.data!, c.offset) : null,
-          c.offset,
-        );
-        out.push(n);
-      } else if (c instanceof TurboQuantKVCache) {
-        const state = c.state();
-        if (state.length === 0) throw new Error("cannot clone an empty cache");
-        const [kIdx, kScales, kZeros, vPacked, vScales] = state;
-        const n = new TurboQuantKVCache(c.kBits, c.vBits);
-        n.restoreState(
-          { kIdx: view(kIdx!), kScales: view(kScales!), kZeros: view(kZeros!),
-            vPacked: view(vPacked!), vScales: view(vScales!) },
-          c.offset, c.headDim!,
-        );
-        for (const a of state) a.dispose();
-        out.push(n);
-      } else if (c instanceof RotatingQuantizedKVCache) {
-        if (!c.keys || !c.values) throw new Error("cannot clone an empty cache");
-        const n = new RotatingQuantizedKVCache(c.maxSize, c.groupSize, c.bits);
-        n.restoreState(tripleView(c.keys, null), tripleView(c.values, null), c.offset, c.ringIdx);
-        out.push(n);
-      } else if (c instanceof QuantizedKVCache) {
-        if (!c.keys || !c.values) throw new Error("cannot clone an empty cache");
-        const n = new QuantizedKVCache(c.groupSize, c.bits);
-        n.restoreState(tripleView(c.keys, c.offset), tripleView(c.values, c.offset), c.offset);
-        out.push(n);
-      } else if (c instanceof RotatingKVCache) {
-        if (!c.keys || !c.values) throw new Error("cannot clone an empty cache");
-        const n = new RotatingKVCache(c.maxSize);
-        n.restoreState(view(c.keys), view(c.values), c.offset, c.ringIdx);
-        out.push(n);
-      } else if (c instanceof SSMCache) {
-        if (!c.conv || !c.recurrent) throw new Error("cannot clone an empty cache");
-        const n = new SSMCache();
-        n.conv = view(c.conv);
-        n.recurrent = view(c.recurrent);
-        n.offset = c.offset;
-        out.push(n);
-      } else if (c instanceof KVCache) {
-        if (!c.keys || !c.values) throw new Error("cannot clone an empty cache");
-        const n = new KVCache();
-        n.restoreState(liveView(c.keys, c.offset), liveView(c.values, c.offset), c.offset);
-        out.push(n);
-      } else {
-        throw new Error("cloneKvCaches: unknown cache type");
-      }
+      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;
 }
@@ -333,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;
@@ -349,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);
@@ -414,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;
   }
 }
 
@@ -454,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.
@@ -468,12 +739,17 @@ export interface SpillItem {
  *  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. */
-type SpillRec = SpillItem & { bytes: number; dropped: boolean };
+type SpillRec = SpillItem & {
+  bytes: number;
+  dropped: boolean;
+  settle: (stored: boolean) => void;
+};
 
 export class SpillQueue {
   #queue: SpillRec[] = [];
   #bytes = 0;
   #dropped = 0;
+  #failed = 0;
   #inFlight: object | null = null;
   #chain: Promise<void> = Promise.resolve();
 
@@ -492,9 +768,18 @@ export class SpillQueue {
   get pendingCount(): number { return this.#queue.length; }
   /** Spills dropped by the cap since start (each = one future cache miss). */
   get droppedCount(): number { return this.#dropped; }
+  /** Store calls that threw or returned false. */
+  get failedCount(): number { return this.#failed; }
 
-  enqueue(item: SpillItem): void {
-    const rec = { ...item, bytes: this.bytesOf(item.caches), dropped: false };
+  enqueue(item: SpillItem): Promise<boolean> {
+    let settle!: (stored: boolean) => void;
+    const result = new Promise<boolean>((resolve) => { settle = resolve; });
+    const rec: SpillRec = {
+      ...item,
+      bytes: this.bytesOf(item.caches),
+      dropped: false,
+      settle,
+    };
     this.#queue.push(rec);
     this.#bytes += rec.bytes;
     while (this.#bytes > this.capBytes) {
@@ -508,24 +793,30 @@ export class SpillQueue {
       if (rec.dropped) return; // clones already disposed at drop time
       this.#inFlight = rec;
       try {
-        await this.store(rec);
+        const stored = await this.store(rec);
+        if (stored === false) this.#failed++;
+        rec.settle(stored !== false);
       } catch {
         // best-effort tier — the entry simply won't restore
+        this.#failed++;
+        rec.settle(false);
       } finally {
         this.#inFlight = null;
         this.#remove(rec);
         this.disposeClones(rec.caches);
       }
     });
+    return result;
   }
 
   /** Settles when everything enqueued so far has flushed or dropped
-   *  (tests; a future shutdown hook would await this). */
+   *  (tests and graceful shutdown). */
   drain(): Promise<void> { return this.#chain; }
 
   #drop(rec: SpillRec): void {
     rec.dropped = true;
     this.#dropped++;
+    rec.settle(false);
     this.#remove(rec);
     this.disposeClones(rec.caches);
   }
@@ -592,24 +883,26 @@ function validateGlm52Prototype(
 ): void {
   const glmKind = entry.kind === "mla" || entry.kind === "mla-dsa" ||
     entry.kind === "mtp-mla";
-  if (!(prototype instanceof Glm52Cache)) {
+  const prototypeSignature = cacheSignature(prototype);
+  if (!prototypeSignature.startsWith("kv:mla:")) {
     if (glmKind)
       throw new Error(`${path}: GLM cache kind ${entry.kind} does not match model cache`);
     return;
   }
-  const expectedKind: CacheKind = prototype.role === "mtp"
+  const glmPrototype = prototype as Glm52Cache;
+  const expectedKind: CacheKind = glmPrototype.role === "mtp"
     ? "mtp-mla"
-    : prototype.dsa ? "mla-dsa" : "mla";
+    : glmPrototype.dsa ? "mla-dsa" : "mla";
   if (entry.kind !== expectedKind) {
     throw new Error(
       `${path}: cache kind ${entry.kind} != model cache kind ${expectedKind}`,
     );
   }
   for (const [field, actual, expected] of [
-    ["kvLoraRank", entry.kvLoraRank, prototype.kvLoraRank],
-    ["ropeHeadDim", entry.ropeHeadDim, prototype.ropeHeadDim],
-    ["dsaHeadDim", entry.dsaHeadDim, prototype.dsa?.headDim],
-    ["maxTokens", entry.maxTokens, prototype.maxTokens],
+    ["kvLoraRank", entry.kvLoraRank, glmPrototype.kvLoraRank],
+    ["ropeHeadDim", entry.ropeHeadDim, glmPrototype.ropeHeadDim],
+    ["dsaHeadDim", entry.dsaHeadDim, glmPrototype.dsa?.headDim],
+    ["maxTokens", entry.maxTokens, glmPrototype.maxTokens],
   ] as const) {
     if (actual !== expected) {
       throw new Error(
@@ -617,7 +910,7 @@ function validateGlm52Prototype(
       );
     }
   }
-  const expectedTensors = prototype.dsa ? 3 : 2;
+  const expectedTensors = glmPrototype.dsa ? 3 : 2;
   if (entry.tensors.length !== expectedTensors) {
     throw new Error(
       `${path}: ${entry.kind} has ${entry.tensors.length} tensors, ` +
@@ -670,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",
@@ -732,79 +1028,8 @@ export function loadKvCache(
   const caches: Cache[] = [];
   try {
     for (const e of header.caches) {
-      const t = e.tensors;
-      switch (e.kind) {
-        case "kv": {
-          const c = new KVCache();
-          c.restoreState(
-            grownArr(t[0]!, e.offset),
-            grownArr(t[1]!, e.offset),
-            e.offset,
-          );
-          caches.push(c);
-          break;
-        }
-        case "rotating": {
-          const c = new RotatingKVCache(e.maxSize!);
-          c.restoreState(arr(t[0]!), arr(t[1]!), e.offset, e.idx!);
-          caches.push(c);
-          break;
-        }
-        case "qkv": {
-          const c = new QuantizedKVCache(e.groupSize!, e.bits!);
-          c.restoreState(triple(t, 0), triple(t, 3), e.offset);
-          caches.push(c);
-          break;
-        }
-        case "rotating-qkv": {
-          const c = new RotatingQuantizedKVCache(e.maxSize!, e.groupSize!, e.bits!);
-          c.restoreState(triple(t, 0), triple(t, 3), e.offset, e.idx!);
-          caches.push(c);
-          break;
-        }
-        case "ssm": {
-          const c = new SSMCache();
-          c.conv = arr(t[0]!);
-          c.recurrent = arr(t[1]!);
-          c.offset = e.offset;
-          caches.push(c);
-          break;
-        }
-        case "turboquant": {
-          const c = new TurboQuantKVCache(e.kBits!, e.vBits!);
-          const kv: TurboQuantTensor = {
-            kIdx: arr(t[0]!), kScales: arr(t[1]!), kZeros: arr(t[2]!),
-            vPacked: arr(t[3]!), vScales: arr(t[4]!),
-          };
-          c.restoreState(kv, e.offset, e.headDim!);
-          caches.push(c);
-          break;
-        }
-        case "mla":
-        case "mla-dsa":
-        case "mtp-mla": {
-          const hasDsa = e.kind === "mla-dsa";
-          if (e.kind === "mtp-mla" && e.dsaHeadDim !== undefined)
-            throw new Error(`${path}: native MTP cache cannot contain DSA state`);
-          const c = new MLACache({
-            kvLoraRank: e.kvLoraRank!,
-            ropeHeadDim: e.ropeHeadDim!,
-            ...(hasDsa ? { dsa: { headDim: e.dsaHeadDim! } } : {}),
-            maxTokens: e.maxTokens,
-            role: e.kind === "mtp-mla" ? "mtp" : "target",
-          });
-          c.restoreCompressedState(
-            arr(t[0]!),
-            arr(t[1]!),
-            hasDsa ? arr(t[2]!) : null,
-            e.offset,
-          );
-          caches.push(c);
-          break;
-        }
-        default:
-          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).
       pending.length = 0;