mlx-bun 0.2.0 → 0.4.0

package/src/spec/source.ts

@@ -1,6 +1,6 @@
 // DraftSource — the seam between the serve-time speculative verify loop
 // (src/spec/serve-loop.ts) and whatever produces draft tokens. Designed in
-// docs/design/mlx-lm-tool-parity-plan.md §7 so every drafter shares ONE
+// docs/reference/cli.md §7 so every drafter shares ONE
 // verify/accept executor:
 //   - TwoModelSource (src/spec/two-model.ts) — mlx-lm parity (L1 oracle:
 //     mlx_lm.server --draft-model), a full second model. Ignores `target`.
@@ -9,7 +9,7 @@
 //     Reads the target's donor K/V + anchor hidden each step.
 //   - DflashSource (src/spec/dflash-source.ts) — DSpark (L3, KL/quality-gated).
 //     Taps the target's multi-layer hiddens (prefill + verify) into a growing
-//     H_ctx (docs/design/dspark-speculative-decoding.md).
+//     H_ctx (docs/design/speculative-decoding.md).
 //   - NgramSource (src/spec/ngram-source.ts) — model-free prompt lookup.
 //   - Glm52NativeMtpSource (src/spec/glm52-mtp-source.ts) — the target
 //     artifact's native Colibri MTP row.
@@ -21,21 +21,60 @@
 // two-model). See [[dspark-seam-kv-borrowing]].
 
 import type { MlxArray } from "../mlx/array";
-import type { RuntimeModel } from "../model/factory";
+import type { Dtype } from "../mlx/ffi";
 import type { Cache } from "../model/gemma4";
+import type { CheckpointAttachment } from "../backends/mlx/checkpoint-state";
+import type { MlxDraftRows } from "../backends/mlx/speculative-round";
+import type { PreparedStateChange } from "../contracts/resources";
 
-/** Read-only handle to the target's live decode state, handed to every source
- *  at open(). Two-model ignores it; KV-borrowing sources read the target's
- *  model (embed / lm-head / the draft module) and its LIVE cache array (donor
- *  K/V views + offset — the SAME caches the serve loop drives). */
+/** Numerical ports are backend-specific; returned arrays are caller-owned. */
+export interface DraftProjection {
+  readonly embed: {
+    encode(ids: MlxArray): MlxArray;
+    readonly scales: { readonly dtype: Dtype };
+  };
+  logitsFromHidden(hidden: MlxArray): MlxArray;
+}
+
+/** Read-only target state and embedding, independent of draft scheduling. */
+export interface AssistantRowsTarget {
+  readonly hiddenSize: number;
+  embed(ids: MlxArray): MlxArray;
+  readDonors(): import("./drafter").AssistantDonors & {
+    readonly positions: readonly number[];
+    dispose(): void;
+  };
+}
+
+export interface QwenMtpTarget {
+  readonly hiddenSize: number;
+  readonly layerCount: number;
+  embed(ids: MlxArray): MlxArray;
+  logitsFromHidden(hidden: MlxArray): MlxArray;
+}
+
+/** Graph-declared extensions over the target's live state. Sources request
+ * only the ports they consume; they never inspect a concrete model or cache.
+ * Absence refuses an unsupported pairing before draft-side allocation. */
 export interface TargetView {
-  readonly model: RuntimeModel;
-  readonly caches: Cache[];
+  /** Opaque identity for providers borrowing weights from one exact target. */
+  readonly identity: object;
+  readonly assistantRows?: AssistantRowsTarget;
+  readonly gemmaTaps?: { readonly layerCount: number; readonly projection: DraftProjection };
+  readonly qwenMtp?: QwenMtpTarget;
 }
 
 /** A per-request draft-token producer. Created per generation (owns its own
  *  draft-side state), disposed by the serve loop's finally. */
 export interface DraftSource {
+  /** Method-owned companion state, independent of retention and storage.
+   * Capture returns owned immutable views. Restore borrows its attachment;
+   * the source retains the views it needs before returning. */
+  readonly checkpoint?: {
+    readonly namespace: string;
+    capture(processedTokens: number): import("../backends/mlx/checkpoint-state").CheckpointAttachment;
+    restore(processedTokens: number, attachment: import("../backends/mlx/checkpoint-state").CheckpointAttachment): void;
+  };
   /** Target prefill shape required by this source's oracle. Most mlx-lm
    *  sources leave the final prompt token pending; native Colibri MTP starts
    *  from a full-prompt target forward. */
@@ -54,8 +93,9 @@ export interface DraftSource {
   /** Process the prompt (two-model: prefill the draft model's cache;
    *  assistant/dflash: read the target's state — mostly a no-op, but DSpark
    *  seeds H_ctx from `ctxML`, the tapped prefill context [1,Lp,m*H], present
-   *  iff tapLayers is set). */
-  prefill(promptIds: number[], ctxML?: MlxArray): void | Promise<void>;
+   *  iff tapLayers is set). A composed prefill may supply its exact processed
+   *  token count, so a history source need not infer the target's boundary. */
+  prefill(promptIds: number[], ctxML?: MlxArray, processedTokens?: number): void | Promise<void>;
 
   /** Propose 0..n tokens (RETURN LENGTH IS AUTHORITATIVE — a source may
    *  return fewer than n, e.g. DSpark's confidence-scheduled draft-length
@@ -100,18 +140,75 @@ export interface DraftSource {
   dispose(): void;
 }
 
+export interface DraftRowSampling {
+  /** Borrow [B,V] log-probabilities; return owned [B] IDs on device. */
+  sample(logprobs: MlxArray, steps: readonly number[]): MlxArray;
+}
+
+export interface DraftRowCheckpoint {
+  readonly processedTokens: number;
+  readonly attachment: CheckpointAttachment;
+}
+
+/** Method-owned state membership, independent of scheduling and cache tiers.
+ * Open/append borrow checkpoints; capture returns owned immutable state.
+ * Membership changes and capture occur only between committed rounds.
+ * The caller retains the provider's residency lease for the group's lifetime. */
+export interface DraftRowGroup extends MlxDraftRows {
+  readonly namespace: string;
+  readonly tapLayers: readonly number[];
+  readonly rowCount: number;
+  append(checkpoints: readonly DraftRowCheckpoint[]): void;
+  prepareAppend(checkpoints: readonly DraftRowCheckpoint[]): PreparedStateChange;
+  filterRows(keep: readonly number[]): void;
+  capture(row: number): DraftRowCheckpoint;
+  dispose(): void;
+}
+
+/** Draft preparation consumes each target chunk at the same batch geometry.
+ * Inputs are borrowed and contain only new tokens/context, with equal chunk
+ * lengths across rows. A null checkpoint appends a cold request. The provider
+ * owns alignment and retained hidden state; scheduling and persistence do not.
+ * Membership and capture occur between prefill calls. Capture returns owned
+ * state after at least one target token has been consumed for that row. */
+export interface DraftPrefillGroup {
+  readonly namespace: string;
+  readonly prefillMode: NonNullable<DraftSource["prefillMode"]>;
+  readonly tapLayers: readonly number[];
+  readonly rowCount: number;
+  append(checkpoints: readonly (DraftRowCheckpoint | null)[]): void;
+  prepareAppend(checkpoints: readonly (DraftRowCheckpoint | null)[]): PreparedStateChange;
+  prefill(tokens: MlxArray, context?: MlxArray): void | Promise<void>;
+  /** Resolve borrowed backing before its external restore lease is released. */
+  materialize(): void;
+  filterRows(keep: readonly number[]): void;
+  capture(row: number): DraftRowCheckpoint;
+  dispose(): void;
+}
+
+export interface GroupedDraftProvider {
+  /** Draft state remains valid when target forwards run under a mounted
+   * adapter context. Unqualified learned draft graphs leave this absent. */
+  readonly supportsTargetAdapters?: boolean;
+  openPrefill(options: { target: TargetView;
+    checkpoints: readonly (DraftRowCheckpoint | null)[] }): DraftPrefillGroup;
+  open(options: { target: TargetView; sampling: DraftRowSampling;
+    checkpoints: readonly DraftRowCheckpoint[] }): DraftRowGroup;
+}
+
 /** Server-lifetime owner of the draft machinery (the loaded draft model);
- *  open() mints a per-request DraftSource. */
+ * open() mints a per-request DraftSource; grouped opens shared row state. */
 export interface DraftProvider {
   /** Human-readable id (registry id / path tail) for logs + cache namespacing. */
   readonly id: string;
   readonly weightsBytes: number;
+  readonly grouped?: GroupedDraftProvider;
   open(opts: {
     /** The request's sampler over logprobs [1,V] → token array [1] (the SAME
      *  sampler as the target — mlx-lm parity; greedy drafting under a
      *  temperature>0 request is NOT parity). */
     sampler: (logprobs: MlxArray, step: number) => MlxArray;
-    /** The target model + its live caches (donor views / offset). */
+    /** Bound target capabilities, with explicit donor-view ownership. */
     target: TargetView;
   }): DraftSource;
   dispose(): void;

package/src/spec/two-model-rows.ts

@@ -0,0 +1,61 @@
+import type { MlxAutoregressiveBinding } from "../backends/mlx/autoregressive";
+import { MlxAutoregressiveDraftRows, type MlxDraftLayout } from "../backends/mlx/autoregressive-draft-rows";
+import { prefillCacheLayout } from "../backends/mlx/cache-layout";
+import { MlxStateRows } from "../backends/mlx/state-rows";
+import { captureCacheAttachment, restoreCacheAttachment } from "../backends/mlx/cache-attachment";
+import { applyStateChanges, cleanupFailure, disposeResources } from "../engine/resources";
+import type { Cache } from "../model/gemma4-base";
+import type { DraftPrefillGroup, DraftRowCheckpoint, DraftRowGroup, DraftRowSampling, GroupedDraftProvider } from "./source";
+import type { PreparedStateChange } from "../contracts/resources";
+import type { MlxArray } from "../mlx/array";
+
+export interface StandaloneDraftRows extends DraftPrefillGroup {
+  draft: DraftRowGroup["draft"];
+  commit(accepted: readonly number[], context?: MlxArray): Promise<void>;
+}
+
+/** The standalone model supplies graph/state capabilities once. Target
+ * verification, scheduling and cache tiers use the existing provider ports. */
+export function openStandaloneDraftRows(binding: MlxAutoregressiveBinding, namespace: string,
+  sampling: DraftRowSampling | null, checkpoints: readonly (DraftRowCheckpoint | null)[]): StandaloneDraftRows {
+    const empty = binding.makeCache(), layouts: MlxDraftLayout[] = [];
+    try { for (const cache of empty) layouts.push(prefillCacheLayout(cache) as MlxDraftLayout); }
+    catch (error) { disposeResources(layouts); throw error; }
+    finally { disposeResources(empty); }
+    const state = new MlxStateRows(layouts), rows = new MlxAutoregressiveDraftRows(binding.graph, state, sampling);
+    const prepareAppend = (checkpoints: readonly (DraftRowCheckpoint | null)[]): PreparedStateChange => {
+      const restored: Cache[][] = [], pending: Array<number | null> = [];
+      try {
+        for (const checkpoint of checkpoints) {
+          restored.push(checkpoint ? restoreCacheAttachment(checkpoint.attachment, binding.makeCache) : binding.makeCache());
+          pending.push(checkpoint && checkpoint.attachment.metadata.pendingLast !== false ? Number(checkpoint.attachment.metadata.pendingLast) : null);
+        }
+        return rows.prepareAppend(restored, pending);
+      } finally { disposeResources(restored.flat()); }
+    };
+    const append = (checkpoints: readonly (DraftRowCheckpoint | null)[]) => applyStateChanges([() => prepareAppend(checkpoints)]);
+    try { append(checkpoints); }
+    catch (error) { return cleanupFailure(error, () => rows.dispose()); }
+    return {
+      namespace, prefillMode: "tail-split", tapLayers: [],
+      get rowCount() { return state.rowCount; },
+      append, prepareAppend, prefill: tokens => rows.prefill(tokens), materialize: () => rows.materialize(),
+      filterRows: keep => rows.filterRows(keep), draft: rows.draft.bind(rows), commit: accepted => rows.commit(accepted),
+      capture(row): DraftRowCheckpoint {
+        const caches = state.extractRow(row), pendingLast = rows.pendingLast[row] ?? null;
+        try {
+          const processedTokens = caches[0]!.offset + Number(pendingLast !== null);
+          return { processedTokens, attachment: captureCacheAttachment("standalone-draft-v1", caches,
+            { processedTokens, pendingLast: pendingLast ?? false }) };
+        } finally { disposeResources(caches); }
+      },
+      dispose: () => rows.dispose(),
+    };
+}
+
+export function standaloneDraftGroups(binding: MlxAutoregressiveBinding, namespace: string): GroupedDraftProvider {
+  return {
+    open: options => openStandaloneDraftRows(binding, namespace, options.sampling, options.checkpoints),
+    openPrefill: options => openStandaloneDraftRows(binding, namespace, null, options.checkpoints),
+  };
+}

package/src/spec/two-model.ts

@@ -16,13 +16,15 @@
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
 import { clearCache } from "../mlx/ffi";
-import { flagOn } from "../flags";
+import { flagOn } from "../runtime-config";
 import { loadModelConfig } from "../config";
 import { Weights } from "../weights";
 import { createModel, type RuntimeModel } from "../model/factory";
-import type { Cache } from "../model/gemma4";
-import { toLogprobs } from "../sampler";
-import type { DraftProvider, DraftSource } from "./source";
+import type { DraftProvider, DraftSource, GroupedDraftProvider } from "./source";
+import { standaloneDraftGroups, openStandaloneDraftRows } from "./two-model-rows";
+import { bindLegacyAutoregressiveModel } from "../backends/mlx/autoregressive";
+import { artifactIdentity } from "../model/artifact-identity";
+import { configFingerprint } from "../model/fingerprint";
 
 /** mlx-lm's prefill_step_size — the draft drain chunks at the same stride
  *  as the target's (serve-loop.ts); chunking is numerically exact for
@@ -32,13 +34,16 @@ const PREFILL_CHUNK = 2048;
 export class TwoModelProvider implements DraftProvider {
   readonly id: string;
   readonly weightsBytes: number;
+  readonly grouped: GroupedDraftProvider;
 
   private constructor(
     readonly model: RuntimeModel,
     private readonly weights: Weights,
     id: string,
+    private readonly namespace: string,
   ) {
     this.id = id;
+    this.grouped = standaloneDraftGroups(bindLegacyAutoregressiveModel(model), namespace);
     this.weightsBytes = [...weights.shards.files.values()].reduce(
       (a, f) => a + f.mmap.size,
       0,
@@ -51,20 +56,25 @@ export class TwoModelProvider implements DraftProvider {
   static async load(modelDir: string, targetVocabSize?: number): Promise<TwoModelProvider> {
     const config = await loadModelConfig(modelDir);
     const weights = await Weights.open(modelDir);
-    const model = createModel(weights, config);
-    if (
-      targetVocabSize !== undefined &&
-      config.text.vocabSize !== targetVocabSize
-    )
-      console.warn(
-        `draft model vocab size ${config.text.vocabSize} != target ${targetVocabSize} — ` +
-          `speculation will accept nothing useful (mlx-lm warns identically)`,
-      );
-    return new TwoModelProvider(model, weights, modelDir.split("/").filter(Boolean).at(-1)!);
+    try {
+      const model = createModel(weights, config);
+      if (
+        targetVocabSize !== undefined &&
+        config.text.vocabSize !== targetVocabSize
+      )
+        console.warn(
+          `draft model vocab size ${config.text.vocabSize} != target ${targetVocabSize} — ` +
+            `speculation will accept nothing useful (mlx-lm warns identically)`,
+        );
+      const identity = await artifactIdentity(configFingerprint(config),
+        [...weights.shards.files].map(([name, shard]) => ({ name, path: shard.path })));
+      return new TwoModelProvider(model, weights, modelDir.split("/").filter(Boolean).at(-1)!,
+        `standalone-draft-v1:${identity}`);
+    } catch (error) { weights.dispose(); throw error; }
   }
 
   open(opts: Parameters<DraftProvider["open"]>[0]): DraftSource {
-    return new TwoModelSource(this.model, opts.sampler);
+    return new TwoModelSource(this.model, opts.sampler, this.namespace);
   }
 
   dispose(): void {
@@ -72,100 +82,32 @@ export class TwoModelProvider implements DraftProvider {
   }
 }
 
+/** Compatibility adapter: the existing request loop also uses the shared
+ * draft graph at B=1. Proposal generation and state recovery have one owner. */
 class TwoModelSource implements DraftSource {
-  readonly weightsBytes = 0; // provider-owned weights; per-request adds caches only
-  private caches: Cache[];
-
-  constructor(
-    private readonly model: RuntimeModel,
-    private readonly sampler: (lp: MlxArray, step: number) => MlxArray,
-  ) {
-    this.caches = this.model.makeCache();
+  readonly weightsBytes = 0;
+  private readonly rows: ReturnType<typeof openStandaloneDraftRows>;
+  private processed = 0;
+  constructor(model: RuntimeModel, sampler: (lp: MlxArray, step: number) => MlxArray, namespace: string) {
+    this.rows = openStandaloneDraftRows(bindLegacyAutoregressiveModel(model), namespace,
+      { sample: (logprobs, steps) => sampler(logprobs, steps[0]!) }, [null]);
   }
-
-  prefill(promptIds: number[]): void {
-    // Oracle convention (mlx-lm speculative_generate_step._prefill, re-anchored
-    // 2026-07-07): the DRAFT model also drains only to len-1 (`while y.size >
-    // 1`) — the last prompt token stays unprocessed and is consumed by the
-    // FIRST draft() round's feed ([pending] = that same token under the serve
-    // loop's oracle shape), mirroring _draft_generate's first _step. Draft-side
-    // full-prefill was the residual knife-edge flipper in the 2026-07-07 live
-    // oracle gate (γ=2 haiku cell). MLX_BUN_PREFILL_TAIL_SPLIT=0 reverts.
-    // (mlx-lm chunks its drain at prefill_step_size=2048; chunking is
-    // numerically exact for causal attention — same KV, same positions —
-    // so we chunk too. Pre-2026-07-07 this was a SINGLE forward of the
-    // whole head: nothing enforced the "<2048 serve regime" the old
-    // comment assumed, and a 32k prompt through the spec lane ran one
-    // 32k-position draft forward — a large activation/mask transient on a
-    // 24 GB box, and untested numerics vs the chunked oracle.)
+  async prefill(promptIds: number[]): Promise<void> {
     const tailSplit = flagOn("MLX_BUN_PREFILL_TAIL_SPLIT", true);
     const upTo = tailSplit && promptIds.length > 1 ? promptIds.length - 1 : promptIds.length;
-    for (let at = 0; at < upTo; at += PREFILL_CHUNK) {
+    for (let at = this.processed; at < upTo; at += PREFILL_CHUNK) {
       const end = Math.min(at + PREFILL_CHUNK, upTo);
-      const ids = ops.fromInt32(promptIds.slice(at, end), [1, end - at]);
-      const h = this.model.forwardHidden(ids, this.caches);
-      ids.dispose();
-      // The head's last-position logits are not needed: the serve loop's
-      // first `feed` is consumed at the top of draft().
-      h.dispose();
-      clearCache();
-    }
-  }
-
-  /** One forward over `feed`, sample the first draft from its last position,
-   *  then n-1 single-token steps — mirrors _draft_generate's autoregressive
-   *  chain (with async_eval pipelining upstream; eager here, the draft is
-   *  small and the serve loop overlaps nothing with it in v1). */
-  draft(feed: number[], n: number, stepBase: number): number[] {
-    const drafts: number[] = [];
-    let input = feed;
-    for (let k = 0; k < n; k++) {
-      const ids = ops.fromInt32(input, [1, input.length]);
-      const h = this.model.forwardHidden(ids, this.caches);
-      ids.dispose();
-      const L = h.shape[1]!;
-      const H = h.shape[2]!;
-      const hLast = h.slice([0, L - 1, 0], [1, L, H]);
-      h.dispose();
-      const logits = this.model.logitsFromHidden(hLast);
-      hLast.dispose();
-      const V = logits.shape[logits.shape.length - 1]!;
-      const flat = ops.reshape(logits, [1, V]);
-      logits.dispose();
-      const lp = toLogprobs(flat);
-      flat.dispose();
-      const tokArr = this.sampler(lp, stepBase + k);
-      lp.dispose();
-      const tok = ops.itemUint32(tokArr);
-      tokArr.dispose();
-      drafts.push(tok);
-      input = [tok];
+      using ids = ops.fromInt32(promptIds.slice(at, end), [1, end - at]);
+      await this.rows.prefill(ids); clearCache();
     }
-    return drafts;
+    this.processed = upTo;
   }
-
-  commit(n: number, kAccept: number): void {
-    const rewind = Math.max(n - kAccept - 1, 0);
-    if (rewind > 0) {
-      for (const c of this.caches) {
-        if (!c.isTrimmable()) {
-          // The serve loop stops speculating on the TARGET-side trim check
-          // before this can matter; a wrapped draft ring here just means the
-          // source is done being useful. Recreate is the safe recovery.
-          this.resetCaches();
-          return;
-        }
-        c.trim(rewind);
-      }
-    }
-  }
-
-  private resetCaches(): void {
-    for (const c of this.caches) c.dispose();
-    this.caches = this.model.makeCache();
+  async draft(feed: number[], count: number, step: number): Promise<number[]> {
+    return (await this.rows.draft([feed.at(-1)!], count, [step]))[0]!;
   }
-
-  dispose(): void {
-    for (const c of this.caches) c.dispose();
+  async commit(_drafted: number, accepted: number): Promise<void> {
+    await this.rows.commit([accepted]);
+    this.processed += accepted + 1;
   }
+  dispose(): void { this.rows.dispose(); }
 }

package/src/ssd-cache.ts

@@ -1,4 +1,5 @@
-// SSD cold tier for the prompt/KV cache (docs/design/ssd-kv-cold-tier.md).
+import type { CheckpointAttachment } from "./backends/mlx/checkpoint-state";
+// SSD cold tier for the prompt/KV cache (docs/design/kv-cache.md).
 //
 // Files ARE the database: no sidecar index. Layout
 //   <dir>/<configFingerprint>/<nsHash>/<uuid>.mlxkv
@@ -22,10 +23,11 @@ import { randomUUID } from "node:crypto";
 import { commonPrefixLength } from "./prompt-cache";
 import {
   saveKvCache, saveKvCacheAsync, loadKvCache, readKvHeader,
-  cacheHeadersTrimmable,
+  cacheHeadersTrimmable, cacheHeadersMinimumReusableOffset, legacyCacheCodecs, type CacheCodecProvider,
   type KvSaveMeta, type KvLoadExpect, type LoadedKvCache,
 } from "./kv-store";
 import type { Cache } from "./model/gemma4-base";
+import { minimumReusableOffset } from "./backends/mlx/state-views";
 
 export interface SsdIndexEntry {
   path: string;
@@ -39,9 +41,12 @@ export interface SsdIndexEntry {
    *  over a shorter usable one (the 2026-07-06 restart-0 defect: the
    *  [prompt+gen] file always won, always got rejected after restore). */
   trimmable: boolean;
+  minimumReusableOffset?: number;
+  generationCheckpoint?: NonNullable<KvSaveMeta["generationCheckpoint"]>;
 }
 
 export interface SsdStoreOptions {
+  codecs?: CacheCodecProvider;
   dir: string;
   maxBytes: number;
   /** Server identity — written into every file, enforced on scan/load. */
@@ -57,6 +62,7 @@ const nsHash = (ns: string): string =>
 
 export class SsdCacheStore {
   readonly #opts: SsdStoreOptions;
+  readonly #codecs: CacheCodecProvider;
   readonly #root: string; // <dir>/<configFingerprint>
   #index: SsdIndexEntry[] = [];
   #warnedWriteFailure = false;
@@ -64,6 +70,7 @@ export class SsdCacheStore {
 
   constructor(opts: SsdStoreOptions) {
     this.#opts = opts;
+    this.#codecs = opts.codecs ?? legacyCacheCodecs;
     this.#root = join(opts.dir, opts.configFingerprint);
   }
 
@@ -88,7 +95,7 @@ export class SsdCacheStore {
    * trimmable descendant is equivalent to an exact boundary snapshot. */
   hasDurablePrefix(tokens: number[], ns = ""): boolean {
     return this.#index.some((entry) => {
-      if (entry.ns !== ns || entry.tokens.length < tokens.length) return false;
+      if (entry.ns !== ns || entry.tokens.length < tokens.length || tokens.length < (entry.minimumReusableOffset ?? 0)) return false;
       if (commonPrefixLength(entry.tokens, tokens) !== tokens.length) return false;
       return entry.tokens.length === tokens.length || entry.trimmable;
     });
@@ -121,10 +128,15 @@ export class SsdCacheStore {
             h.tokenizerHash !== this.#opts.tokenizerHash
           )
             throw new Error("metadata mismatch");
+          if ((h.codecProvider ?? legacyCacheCodecs.id) !== this.#codecs.id) continue;
           const st = statSync(path);
           this.#index.push({
             path, ns: h.ns ?? "", tokens: h.tokens, bytes: st.size, mtimeMs: st.mtimeMs,
-            trimmable: cacheHeadersTrimmable(h.caches),
+            trimmable: !h.attachments?.length && cacheHeadersTrimmable(h.caches, this.#codecs),
+            minimumReusableOffset: cacheHeadersMinimumReusableOffset(h.caches),
+            ...(h.generationCheckpoint
+              ? { generationCheckpoint: h.generationCheckpoint }
+              : {}),
           });
         } catch {
           try { rmSync(path, { force: true }); } catch {}
@@ -149,7 +161,7 @@ export class SsdCacheStore {
     for (const e of this.#index) {
       if (e.ns !== ns) continue;
       const p = Math.min(commonPrefixLength(e.tokens, prompt), prompt.length - 1);
-      if (p === 0 || p <= bestLen) continue;
+      if (p === 0 || p <= bestLen || p < (e.minimumReusableOffset ?? 0)) continue;
       if (e.tokens.length - p > 0 && !e.trimmable) continue;
       bestLen = p;
       best = e;
@@ -171,7 +183,7 @@ export class SsdCacheStore {
         ns: entry.ns,
         verify: this.#opts.verify,
       };
-      const loaded = loadKvCache(entry.path, model, expect);
+      const loaded = loadKvCache(entry.path, model, expect, this.#codecs);
       const now = new Date();
       try { utimesSync(entry.path, now, now); entry.mtimeMs = now.getTime(); } catch {}
       this.stats.restores++;
@@ -193,13 +205,13 @@ export class SsdCacheStore {
    *  the same ns whose tokens are a PREFIX of the new entry are superseded
    *  (the agent-conversation pattern: one growing entry, not N generations
    *  of it). Returns true if stored. */
-  store(tokens: number[], caches: Cache[], ns = ""): boolean {
+  store(tokens: number[], caches: Cache[], ns = "", attachments?: CheckpointAttachment[]): boolean {
     const dir = join(this.#root, nsHash(ns));
     const path = join(dir, `${randomUUID()}.mlxkv`);
     try {
       mkdirSync(dir, { recursive: true });
-      saveKvCache(path, tokens, caches, this.#meta(ns));
-      return this.#indexStored(path, tokens, caches, ns);
+      saveKvCache(path, tokens, caches, { ...this.#meta(ns), attachments }, this.#codecs);
+      return this.#indexStored(path, tokens, caches, ns, undefined, attachments);
     } catch (err) {
       return this.#storeFailed(path, err);
     }
@@ -210,21 +222,81 @@ export class SsdCacheStore {
    *  yields the event loop between tensors so serving interleaves.
    *  Caller passes zero-copy CLONES it owns (a consistent snapshot no
    *  matter what the live entry does meanwhile) and disposes them after.
-   *  `waitTurn` gates every per-tensor step (see saveKvCacheAsync) — the
-   *  server passes the gateway's onIdle so the flush only progresses while
-   *  the engine is idle, never mid-decode. */
-  async storeAsync(tokens: number[], caches: Cache[], ns = "", waitTurn?: () => Promise<void>): Promise<boolean> {
+   *  `runStep` owns every per-tensor step (see saveKvCacheAsync) — the
+   *  server passes the gateway's exclusive runner so a request cannot start
+   *  between an idle check and the blocking MLX readback. */
+  async storeAsync(
+    tokens: number[], caches: Cache[], ns = "",
+    runStep?: <T>(step: () => T) => Promise<T>,
+    attachments?: CheckpointAttachment[],
+  ): Promise<boolean> {
     const dir = join(this.#root, nsHash(ns));
     const path = join(dir, `${randomUUID()}.mlxkv`);
     try {
       mkdirSync(dir, { recursive: true });
-      await saveKvCacheAsync(path, tokens, caches, this.#meta(ns), waitTurn);
-      return this.#indexStored(path, tokens, caches, ns);
+      await saveKvCacheAsync(path, tokens, caches, { ...this.#meta(ns), attachments }, runStep, this.#codecs);
+      return this.#indexStored(path, tokens, caches, ns, undefined, attachments);
     } catch (err) {
       return this.#storeFailed(path, err);
     }
   }
 
+  /** Longest durable in-flight generation for an identical request. The
+   *  checkpoint namespace is separate from the ordinary prompt-cache tier,
+   *  so incomplete assistant output can never seed an unrelated prefill. */
+  findGenerationCheckpoint(
+    prompt: number[], key: string, cacheNs = "",
+  ): SsdIndexEntry | null {
+    let best: SsdIndexEntry | null = null;
+    for (const entry of this.#index) {
+      const checkpoint = entry.generationCheckpoint;
+      if (!checkpoint || checkpoint.key !== key || checkpoint.cacheNs !== cacheNs)
+        continue;
+      if (checkpoint.originalPromptTokens !== prompt.length) continue;
+      if (entry.tokens.length <= prompt.length) continue;
+      if (commonPrefixLength(entry.tokens, prompt) !== prompt.length) continue;
+      if (!best || entry.tokens.length > best.tokens.length) best = entry;
+    }
+    return best;
+  }
+
+  /** Persist one in-flight generation atomically. A completed newer
+   *  checkpoint supersedes older checkpoints for the same request only after
+   *  its rename succeeds, so a crash during a write leaves the prior point. */
+  async storeGenerationCheckpoint(
+    tokens: number[], caches: Cache[], checkpoint: NonNullable<KvSaveMeta["generationCheckpoint"]>,
+    runStep?: <T>(step: () => T) => Promise<T>,
+  ): Promise<boolean> {
+    const ns = `__generation_checkpoint__:${checkpoint.key}`;
+    const dir = join(this.#root, nsHash(ns));
+    const path = join(dir, `${randomUUID()}.mlxkv`);
+    try {
+      if (runStep) await runStep(() => mkdirSync(dir, { recursive: true }));
+      else mkdirSync(dir, { recursive: true });
+      await saveKvCacheAsync(path, tokens, caches, {
+        ...this.#meta(ns), generationCheckpoint: checkpoint,
+      }, runStep, this.#codecs);
+      const index = () => {
+        const stored = this.#indexStored(path, tokens, caches, ns, checkpoint);
+        if (stored) {
+          for (const entry of [...this.#index]) {
+            if (entry.path === path) continue;
+            if (entry.generationCheckpoint?.key === checkpoint.key) this.remove(entry.path);
+          }
+        }
+        return stored;
+      };
+      return runStep ? await runStep(index) : index();
+    } catch (err) {
+      return runStep ? await runStep(() => this.#storeFailed(path, err)) : this.#storeFailed(path, err);
+    }
+  }
+
+  removeGenerationCheckpoints(key: string): void {
+    for (const entry of [...this.#index])
+      if (entry.generationCheckpoint?.key === key) this.remove(entry.path);
+  }
+
   #meta(ns: string): KvSaveMeta {
     return {
       modelId: this.#opts.modelId,
@@ -234,14 +306,18 @@ export class SsdCacheStore {
     };
   }
 
-  #indexStored(path: string, tokens: number[], caches: Cache[], ns: string): boolean {
+  #indexStored(
+    path: string, tokens: number[], caches: Cache[], ns: string,
+    generationCheckpoint?: NonNullable<KvSaveMeta["generationCheckpoint"]>,
+    attachments?: CheckpointAttachment[],
+  ): boolean {
     const st = statSync(path);
     if (st.size > this.#opts.maxBytes) {
       rmSync(path, { force: true });
       console.warn(`[ssd-cache] entry not stored: ${st.size} bytes exceeds the ${this.#opts.maxBytes}-byte cap`);
       return false;
     }
-    const trimmable = caches.every((c) => c.isTrimmable());
+    const trimmable = !attachments?.length && caches.every((c) => c.isTrimmable());
     // Exact duplicates (same tokens) are replaced regardless of
     // trimmability — the new file serves exactly the old one's matches.
     for (const e of [...this.#index]) {
@@ -256,11 +332,15 @@ export class SsdCacheStore {
     // unrestorable one landed.
     if (trimmable) {
       for (const e of [...this.#index]) {
-        if (e.ns !== ns || e.tokens.length >= tokens.length) continue;
+        if (e.ns !== ns || e.tokens.length >= tokens.length || e.tokens.length < minimumReusableOffset(caches)) continue;
         if (commonPrefixLength(e.tokens, tokens) === e.tokens.length) this.remove(e.path);
       }
     }
-    this.#index.push({ path, ns, tokens, bytes: st.size, mtimeMs: Date.now(), trimmable });
+    this.#index.push({
+      path, ns, tokens, bytes: st.size, mtimeMs: Date.now(), trimmable,
+      minimumReusableOffset: minimumReusableOffset(caches),
+      ...(generationCheckpoint ? { generationCheckpoint } : {}),
+    });
     this.stats.spills++;
     this.evictToCap();
     return true;