mlx-bun 0.3.0 → 0.5.0

package/src/backends/mlx/projected-context-rows.ts

@@ -0,0 +1,72 @@
+import type { MlxArray } from "../../mlx/array";
+import { Dtype } from "../../mlx/ffi";
+import * as ops from "../../mlx/ops";
+import { BatchedKVCache } from "../../model/batched-kv";
+import type { Cache } from "../../model/gemma4-base";
+import { MlxStateRows } from "./state-rows";
+import { disposeResources } from "../../engine/resources";
+
+export interface ProjectedContext { k: MlxArray; v: MlxArray; }
+export interface ContextProjection {
+  project(hidden: MlxArray, positions: number | MlxArray): ProjectedContext[];
+}
+
+/** Incremental projected context reuses ordinary full-attention row storage.
+ * The projection graph owns numerics; this object owns positions and validity. */
+export class MlxProjectedContextRows {
+  readonly state: MlxStateRows<BatchedKVCache>;
+  constructor(layers: number, readonly graph: ContextProjection) {
+    this.state = new MlxStateRows(Array.from({ length: layers }, () => new BatchedKVCache()));
+  }
+  get positions(): readonly number[] { return this.state.caches[0]!.rowOffsets; }
+  prepareAppend(rows: readonly (readonly Cache[])[]) { return this.state.prepareAppendMany(rows); }
+  append(hidden: MlxArray, lengths: readonly number[]): void {
+    const width = Math.max(...lengths), B = lengths.length, positions = this.positions;
+    using context = hidden.slice([0, 0, 0], [B, width, hidden.shape[2]!]);
+    using offsets = positions.every(position => position === positions[0]) ? null : ops.fromInt32([...positions], [B]);
+    const projected = this.graph.project(context, offsets ?? positions[0]!);
+    try {
+      for (const [layer, cache] of this.state.caches.entries()) {
+        cache.specRoundBegin();
+        for (const array of cache.updateAndFetch(projected[layer]!.k, projected[layer]!.v)) array.dispose();
+        if (lengths.every(length => length === width)) cache.specRoundCommit();
+        else cache.specRoundRollback(lengths);
+      }
+    } finally { disposeResources(projected.flatMap(pair => [pair.k, pair.v])); }
+  }
+  materialize(): void { ops.evalAll(this.state.caches.flatMap(cache => cache.state())); }
+  readAttention() {
+    const views: import("../../model/gemma4-base").KvDonorRows[] = [];
+    try { for (const cache of this.state.caches) views.push(cache.captureDonorRows()); }
+    catch (error) { disposeResources(views.flatMap(view => [view.keys, view.values])); throw error; }
+    const masks = new Map<string, MlxArray | null>();
+    return {
+      attend(layer: number, query: MlxArray, blockKeys: MlxArray, blockValues: MlxArray, scale = 1): MlxArray {
+        const view = views[layer]!, width = view.keys.shape[2]!, block = blockKeys.shape[2]!, B = view.offsets.length;
+        // Membership and coverage advance together in every context layer.
+        const key = `${width}:${block}:${view.keys.dtype}`;
+        if (!masks.has(key)) {
+          let mask: MlxArray | null = null;
+          if (view.starts.some(start => start !== 0) || view.ends.some(end => end !== width)) {
+            using columns = ops.arange(0, width, 1, Dtype.int32);
+            using lower = ops.fromInt32([...view.starts], [B, 1, 1, 1]);
+            using upper = ops.fromInt32([...view.ends], [B, 1, 1, 1]);
+            using before = ops.less(columns, lower), after = ops.greaterEqual(columns, upper);
+            using invalid = ops.logicalOr(before, after);
+            using floats = invalid.astype(view.keys.dtype);
+            using prefix = ops.mulScalar(floats, -1e9);
+            using suffix = ops.zeros([B, 1, 1, block], view.keys.dtype);
+            mask = ops.concatAxis([prefix, suffix], 3);
+          }
+          masks.set(key, mask);
+        }
+        using keys = ops.concatAxis([view.keys, blockKeys], 2);
+        using values = ops.concatAxis([view.values, blockValues], 2);
+        const mask = masks.get(key)!;
+        return ops.sdpa(query, keys, values, scale, mask ? "array" : "", mask);
+      },
+      dispose() { disposeResources([...views.flatMap(view => [view.keys, view.values]), ...[...masks.values()].filter(value => value !== null)]); },
+    };
+  }
+  dispose(): void { this.state.dispose(); }
+}

package/src/backends/mlx/prompt-input.ts

@@ -0,0 +1,29 @@
+import type { MlxArray } from "../../mlx/array";
+import type { Cache } from "../../model/gemma4-base";
+
+/** Prepared native input with an indivisible prefill. The preparation owner
+ * retains its tensors; the model binding returns an owned final hidden row.
+ * Decode, sampling and row retirement use the ordinary execution group. */
+export interface MlxPromptInput {
+  forward(ids: MlxArray, caches: Cache[], start?: number): MlxArray;
+  readonly decodeState?: MlxDecodeState;
+}
+
+/** Additional model-owned forward state. Equal keys declare compatible row
+ * state; the backend supplies the current row order after joins/retirement. */
+export interface MlxDecodeState {
+  readonly key: string;
+  forward(ids: MlxArray, caches: Cache[], rows: readonly MlxDecodeState[]): MlxArray;
+}
+
+/** Preserve the ordinary generator's projection geometry for a prepared
+ * embedding sequence: process the whole prompt, then project only its tip. */
+export function bindEmbeddingsInput(
+  forward: (ids: MlxArray, caches: Cache[], start: number) => MlxArray,
+): MlxPromptInput {
+  return { forward(ids, caches, start = 0) {
+    using hidden = forward(ids, caches, start);
+    const [batch, length, width] = hidden.shape;
+    return hidden.slice([0, length! - 1, 0], [batch!, length!, width!]);
+  } };
+}

package/src/backends/mlx/qwen-prompt-input.ts

@@ -0,0 +1,40 @@
+import type { Qwen35Model } from "../../model/qwen3_5";
+import { mropePositionIds, type MropeRequestState } from "../../model/qwen3-mrope";
+import type { MlxArray } from "../../mlx/array";
+import type { Cache } from "../../model/gemma4-base";
+import * as ops from "../../mlx/ops";
+import { Dtype } from "../../mlx/ffi";
+import { bindEmbeddingsInput, type MlxDecodeState, type MlxPromptInput } from "./prompt-input";
+
+class QwenMediaDecodeState implements MlxDecodeState {
+  readonly key = "qwen-interleaved-mrope";
+  constructor(readonly model: Qwen35Model, readonly delta: number) {}
+
+  forward(ids: MlxArray, caches: Cache[], rows: readonly MlxDecodeState[]): MlxArray {
+    const B = rows.length, L = ids.shape[1]!;
+    const cache = caches[this.model.faIdx]!;
+    // The cache supplies logical positions, including left padding, on device.
+    // No offset readback or request-global position mutation is needed.
+    const offsets = (cache as Cache & { ropeOffsetArr?: MlxArray }).ropeOffsetArr;
+    using deltas = ops.fromInt32(rows.map(row => (row as QwenMediaDecodeState).delta), [B]);
+    using scalar = offsets ? null : ops.fromInt32([cache.offset], []);
+    using starts = ops.add(offsets ?? scalar!, deltas);
+    using startRows = ops.reshape(starts, [B, 1]);
+    using indices = ops.arange(0, L, 1, Dtype.int32);
+    using positions = ops.add(startRows, indices);
+    using axis = ops.reshape(positions, [1, B, L]);
+    using grid = ops.concatAxis([axis, axis, axis], 0);
+    return this.model.forwardHiddenAtPositions(ids, caches, grid);
+  }
+}
+
+/** The prompt owns its grid; only its continuation delta survives into decode. */
+export function bindQwenMediaInput(model: Qwen35Model, embeddings: MlxArray,
+  state: MropeRequestState): MlxPromptInput {
+  const input = bindEmbeddingsInput((ids, caches, start) => {
+    using positions = mropePositionIds(state, start, ids.shape[1]!);
+    return start > 0 ? model.forwardHiddenAtPositions(ids, caches, positions)
+      : model.forwardEmbeddingsAtPositions(embeddings, caches, positions);
+  });
+  return { ...input, decodeState: new QwenMediaDecodeState(model, state.delta) };
+}

package/src/backends/mlx/request-state-policy.ts

@@ -0,0 +1,54 @@
+import { namespacedCache } from "../../engine/namespaced-cache";
+import { runtimeConfig, runtimeFlag, withRuntimeConfig, type RuntimeConfig } from "../../runtime-config";
+import type { GenerateOptions } from "../../generate";
+import type { RuntimeModel } from "../../model/factory";
+import { KVCache, type Cache } from "../../model/gemma4-base";
+import { PagedKVCache } from "../../lab/paged-kv/paged-kv";
+import type { RowPromptCache } from "./batch-group";
+import { disposeResources } from "../../engine/resources";
+
+/** Request state construction and reusable-prefix policy. Equal keys declare
+ * compatible state layouts; the executor need not inspect their settings. */
+export interface MlxRequestStatePolicy {
+  readonly key: string;
+  readonly promptCache?: RowPromptCache;
+  create(): Cache[];
+}
+
+/** Replace fresh full-attention storage before prefill. Sliding layers retain
+ * their layout; pre-warmed library caches retain their existing storage. */
+export function maybePageKv(
+  cache: Cache[], options: GenerateOptions, capacityTokens: number,
+): void {
+  if (!options.pagedKv) return;
+  if (cache.some((c) => c.offset > 0)) return;
+  const blockSize = options.pagedKv.blockSize ?? PagedKVCache.DEFAULT_BLOCK_SIZE;
+  for (let i = 0; i < cache.length; i++) {
+    if (cache[i] instanceof KVCache) {
+      cache[i]!.dispose(); // fresh (offset 0) — nothing stored yet
+      cache[i] = new PagedKVCache(capacityTokens, blockSize, undefined,
+        options.kvBits === 4 || options.kvBits === 8 ? { bits: options.kvBits, groupSize: options.kvGroupSize ?? 64 } : undefined);
+    }
+  }
+}
+
+export function pagedPrefixNamespace(options: GenerateOptions, base: string,
+  direct = runtimeFlag("MLX_BUN_PAGED_ATTN", false)): string {
+  return JSON.stringify(["paged-v1", options.pagedKv?.blockSize ?? PagedKVCache.DEFAULT_BLOCK_SIZE,
+    options.kvBits ?? 0, options.kvGroupSize ?? 64, direct, base]);
+}
+
+export function bindPagedRequestState(model: RuntimeModel, options: GenerateOptions,
+  capacityTokens: number, cache?: RowPromptCache, runtime: RuntimeConfig = runtimeConfig()): MlxRequestStatePolicy | undefined {
+  if (!options.pagedKv) return undefined;
+  const blockSize = options.pagedKv.blockSize ?? PagedKVCache.DEFAULT_BLOCK_SIZE;
+  const stateOptions = { pagedKv: { blockSize }, kvBits: options.kvBits, kvGroupSize: options.kvGroupSize };
+  const direct = runtime.flag("MLX_BUN_PAGED_ATTN", false);
+  const namespace = (base: string) => pagedPrefixNamespace(stateOptions, base, direct);
+  return { key: namespace(""),
+    promptCache: cache ? namespacedCache(cache, namespace) : undefined, create: () => withRuntimeConfig(runtime, () => {
+    const caches = model.makeCache();
+    try { maybePageKv(caches, stateOptions, capacityTokens); return caches; }
+    catch (error) { disposeResources(caches); throw error; }
+  }) };
+}

package/src/backends/mlx/rollback.ts

@@ -1,7 +1,39 @@
 import type { SpeculativeTransaction } from "../../inference/rollback";
 import type { Cache } from "../../model/gemma4-base";
 
-/** Bind the full legacy recurrent transaction or the trim implementation once.
+interface TransactionLayer<Acceptance> {
+  fits(drafts: number): boolean;
+  begin(drafts: number): void;
+  resolve(accepted: Acceptance, drafts: number): void;
+}
+
+/** State storage that can retain a different verified prefix for each row.
+ * Both attention and recurrent layouts implement this same transaction. */
+export interface RowSpeculativeState {
+  specRoundBegin(): void;
+  specRoundCommit(): void;
+  specRoundRollback(keep: number | readonly number[]): void;
+}
+
+export function bindRowCacheRollback(
+  caches: readonly RowSpeculativeState[], rows: number,
+): SpeculativeTransaction<readonly number[]> {
+  return bindTransaction(caches.map(cache => ({
+    fits: () => true,
+    // A zero-candidate round consumes only the committed pending token.
+    // It cannot reject any input, so retaining rollback state is unnecessary.
+    begin: drafts => { if (drafts > 0) cache.specRoundBegin(); },
+    resolve(accepted: readonly number[], drafts: number) {
+      if (drafts === 0) return;
+      if (accepted.every(count => count === drafts)) cache.specRoundCommit();
+      else cache.specRoundRollback(accepted.every(count => count === accepted[0])
+        ? accepted[0]! + 1 : accepted.map(count => count + 1));
+    },
+  })), (accepted, drafts) => accepted.length === rows &&
+    accepted.every(count => Number.isSafeInteger(count) && count >= 0 && count <= drafts));
+}
+
+/** Bind the recurrent transaction or the trim implementation once.
  * Ring-capacity checks still happen at each safe boundary before any writes. */
 export function bindCacheRollback(caches: readonly Cache[]): SpeculativeTransaction {
   const layers = caches.map((cache) => {
@@ -24,6 +56,15 @@ export function bindCacheRollback(caches: readonly Cache[]): SpeculativeTransact
       },
     };
   });
+  return bindTransaction(layers, (accepted, drafts) =>
+    Number.isSafeInteger(accepted) && accepted >= 0 && accepted <= drafts);
+}
+
+/** One transaction lifecycle for scalar and row-wise state bindings. */
+function bindTransaction<Acceptance>(
+  layers: readonly TransactionLayer<Acceptance>[],
+  accepts: (accepted: Acceptance, drafts: number) => boolean,
+): SpeculativeTransaction<Acceptance> {
   let active: number | undefined;
   let invalid = false;
   const validate = (drafts: number) => {
@@ -37,12 +78,12 @@ export function bindCacheRollback(caches: readonly Cache[]): SpeculativeTransact
       if (active !== undefined || !layers.every((layer) => layer.fits(drafts)))
         throw new Error("speculative state cannot begin this round");
       active = drafts;
-      try { for (const layer of layers) layer.begin(); }
+      try { for (const layer of layers) layer.begin(drafts); }
       catch (error) { invalid = true; throw error; }
     },
     resolve(accepted) {
-      validate(accepted);
-      if (active === undefined || accepted > active) throw new Error("invalid speculative acceptance count");
+      if (invalid) throw new Error("speculative state is invalid and must be discarded");
+      if (active === undefined || !accepts(accepted, active)) throw new Error("invalid speculative acceptance count");
       try { for (const layer of layers) layer.resolve(accepted, active); }
       catch (error) { invalid = true; throw error; }
       active = undefined;

package/src/backends/mlx/row-sampling.ts

@@ -0,0 +1,49 @@
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import type { TokenLogprobs } from "../../contracts/generation";
+import { readStepExtras, stepExtrasArrays, disposeStepExtras,
+  type DeviceStepSampler, type StepExtras } from "../../sampler";
+import { disposeResources, cleanupFailure } from "../../engine/resources";
+
+type TokenSink = (token: number, logprobs?: TokenLogprobs) => void | boolean | Promise<void | boolean>;
+
+/** Request-owned sampling and output metadata. Execution groups receive token
+ * tensors and an output callback; they need no logprob or retention policy. */
+export function createRowSampling(
+  sampler: DeviceStepSampler, onToken: TokenSink,
+) {
+  if (!sampler.capturesLogprobs) return {
+    sample: (logits: MlxArray, step: number) => sampler.sample(logits, step).token,
+    plainGreedy: sampler.isPlainGreedy, onToken, dispose: () => sampler.dispose(),
+  };
+
+  // A pipelined group can sample the next token before emitting the previous
+  // one. Keep captures by request-local step, independently of row position.
+  const pending = new Map<number, StepExtras>();
+  let emitted = 0;
+  return {
+    sample(logits: MlxArray, step: number): MlxArray {
+      const result = sampler.sample(logits, step);
+      try {
+        ops.asyncEvalAll([result.token, ...stepExtrasArrays(result.extras)]);
+        if (result.extras) pending.set(step, result.extras);
+        return result.token;
+      } catch (error) {
+        return cleanupFailure(error, () => disposeResources([result.token,
+          { dispose: () => disposeStepExtras(result.extras) }]));
+      }
+    },
+    // The vectorized token-only shortcut bypasses the per-request sampler.
+    plainGreedy: false,
+    onToken(token: number) {
+      const extras = pending.get(emitted) ?? null;
+      pending.delete(emitted++);
+      return onToken(token, readStepExtras(extras));
+    },
+    dispose() {
+      const extras = [...pending.values()];
+      pending.clear();
+      disposeResources([sampler, ...extras.map(value => ({ dispose: () => disposeStepExtras(value) }))]);
+    },
+  };
+}

package/src/backends/mlx/serial-executor.ts

@@ -1,32 +1,31 @@
+import { pagedPrefixNamespace } from "./request-state-policy";
+import type { SpecRunServices } from "../../spec/serve-loop";
 import { cleanupFailure, disposeResources, ownResource } from "../../engine/resources";
 import { bindGeneration } from "../../generate";
 import type { RuntimeModel } from "../../model/factory";
 import { Qwen35Model } from "../../model/qwen3_5";
 import type { Cache } from "../../model/gemma4";
-import type { PromptCache } from "../../prompt-cache";
+import type { MlxPrefixCache } from "./checkpoint-state";
 import type { SsdCacheStore } from "../../ssd-cache";
 import type { DraftProvider } from "../../spec/source";
 import type { SerialRun, Vision } from "../../serve/generation-gateway";
-import { generationCheckpointKey } from "../../serve/checkpoint-identity";
+import { bindContinuationPolicy } from "./continuation";
 import { runtimeConfig, withRuntimeConfig, type RuntimeConfig } from "../../runtime-config";
 import { bindLegacySpeculativeModel } from "./speculative";
-import { enterMlxMemoryGuard, type MlxMemoryBudget } from "./memory-guard";
 
 /** Native serial execution depends on this bound port, never a model union.
  * Weights remain borrowed. The gateway supplies the exclusive runtime lease. */
 export interface MlxSerialBinding {
   readonly runtime: RuntimeConfig;
   readonly generate: ReturnType<typeof bindGeneration>;
-  readonly speculate?: SerialRun;
+  readonly speculate?: (
+    promptIds: Parameters<SerialRun>[0], options: Parameters<SerialRun>[1],
+    onToken: Parameters<SerialRun>[2], services?: SpecRunServices,
+  ) => ReturnType<SerialRun>;
   makeCache(): Cache[];
   enterMedia?(vision?: Vision): () => void;
-  enterMemoryGuard?(budget: SerialMemoryBudget, promptTokens: number,
-    prefillChunkSize?: number): { check(): void; close(): void };
 }
 
-/** Admission's ceiling and the cache that can release memory under the GPU lease. */
-export type SerialMemoryBudget = MlxMemoryBudget;
-
 /** Family-specific context is bound once at the compatibility boundary. */
 export function bindLegacySerialModel(
   model: RuntimeModel,
@@ -37,13 +36,10 @@ export function bindLegacySerialModel(
     runtime: runtimeConfig(),
     generate: bindGeneration(model),
     makeCache: model.makeCache.bind(model),
-    enterMemoryGuard(budget, promptTokens, prefillChunkSize) {
-      return enterMlxMemoryGuard(model, budget, promptTokens, prefillChunkSize);
-    },
-    ...(speculative && draft ? { speculate: (async (prompt, options, onToken) => {
+    ...(speculative && draft ? { speculate: (async (prompt, options, onToken, services) => {
       const { specRun } = await import("../../spec/serve-loop");
-      return specRun(speculative, draft.numDraftTokens, prompt, options, onToken);
-    }) satisfies SerialRun } : {}),
+      return specRun(speculative, draft.numDraftTokens, prompt, options, onToken, services);
+    }) satisfies NonNullable<MlxSerialBinding["speculate"]> } : {}),
     ...(model instanceof Qwen35Model ? { enterMedia(vision?: Vision) {
       const previous = model.mrope;
       if (vision?.mrope) model.mrope = vision.mrope;
@@ -53,11 +49,11 @@ export function bindLegacySerialModel(
 }
 
 export interface MlxSerialServices {
-  readonly promptCache: Pick<PromptCache, "take" | "put"> & Partial<Pick<PromptCache, "maxBytes">>;
+  readonly promptCache: MlxPrefixCache & { readonly maxBytes?: number };
   readonly checkpoints: Pick<SsdCacheStore, "findGenerationCheckpoint" | "restore" |
     "storeGenerationCheckpoint" | "removeGenerationCheckpoints"> | null;
   readonly checkpointEveryTokens?: number;
-  readonly memoryBudget?: SerialMemoryBudget;
+  readonly checkpointPersistence?: import("./continuation-persistence").ContinuationPersistence;
   /** Artifact, implementation, state ABI and codec identity captured at load. */
   readonly identity: unknown;
   adapterNamespace(adapters: string[]): string;
@@ -71,17 +67,25 @@ export interface MlxSerialServices {
 export function createMlxSerialExecutor(binding: MlxSerialBinding, services: MlxSerialServices): SerialRun {
   const { promptCache, checkpoints: ssdStore } = services;
   return (promptIds, options, onToken, vision, trace, execution) => withRuntimeConfig(binding.runtime, async () => {
+    if (promptCache.reclaim) {
+      const publish = onToken;
+      let steps = 0;
+      onToken = (token, logprobs) => {
+        if (steps++ % 256 === 0) promptCache.reclaim!();
+        return publish(token, logprobs);
+      };
+    }
+    let releaseContinuation: (() => void) | undefined;
     let caches: Cache[] = [];
     let retain: (() => void) | undefined;
     let closeMedia: (() => void) | undefined;
-    let memoryGuard: ReturnType<NonNullable<MlxSerialBinding["enterMemoryGuard"]>> | undefined;
     const cleanup = ownResource(null, () => disposeResources([
       { dispose() {
         disposeResources(caches);
         retain?.(); // backing release requires successful cache disposal
       } },
       { dispose: () => closeMedia?.() },
-      { dispose: () => memoryGuard?.close() },
+      { dispose: () => releaseContinuation?.() },
       ...[options.grammar, vision?.embeddings, vision?.imageMask,
         vision?.multimodalMask, options.visionPixels].filter((value) => value != null),
     ]));
@@ -90,53 +94,48 @@ export function createMlxSerialExecutor(binding: MlxSerialBinding, services: Mlx
       if (!execution) throw new Error("serial execution requires a resolved plan");
       if (execution.method === "speculative") {
         if (!binding.speculate) throw new Error("resolved speculation requires a bound verifier");
-        return await binding.speculate(promptIds, {
-          ...options,
-          speculativeCacheBytes: binding.runtime.flag("MLX_BUN_MTP_PROMPT_CACHE", false)
-            ? services.promptCache.maxBytes ?? 0 : 0,
-        }, onToken);
+        return await binding.speculate(promptIds, options, onToken, {
+          ...(binding.runtime.flag("MLX_BUN_MTP_PROMPT_CACHE", true) && services.promptCache.maxBytes !== 0
+            ? { prefixCache: services.promptCache } : {}),
+          cloneState: services.cloneState,
+          cacheNamespace: options.adapters?.length ? services.adapterNamespace(options.adapters) : "",
+        });
       }
       // Cache entries are adapter-specific: KV computed under one adapter
       // must never seed another's (or the base's) prefill.
-      const cacheNs = options.adapters?.length ? services.adapterNamespace(options.adapters) : "";
+      const baseNs = options.adapters?.length ? services.adapterNamespace(options.adapters) : "";
+      const cacheNs = execution.pagedKv ? pagedPrefixNamespace(options, baseNs) : baseNs;
       // Paged-KV request scope (docs/design/kv-cache.md): media
       // prompts (bidir overlay) and LoRA-adapter requests are v1 non-goals —
       // they run the PLAIN cache path even under --paged-kv (scope the flag
       // per request, never 400). Effective value computed ONCE so the
       // prompt-cache bypass below and the generate() options can't disagree.
       const pagedKv = execution.pagedKv ? options.pagedKv : undefined;
-      // Paged requests bypass the prompt cache entirely (v1 non-goal:
-      // PagedKVCache has no cloneKvCaches/restore path — the vision
-      // precedent). Fresh caches per request, disposed on completion.
+      // The resolved policy controls reuse; paged layouts have their own
+      // numerical namespace and use the same RAM/SSD cache interface.
       const skipPromptCache = !execution.promptCache;
       const checkpointEvery = services.checkpointEveryTokens;
       const checkpointEligible = execution.checkpoint;
-      const checkpointKey = checkpointEligible
-        ? generationCheckpointKey(promptIds, options, cacheNs, execution, services.identity)
-        : null;
+      const continuation = checkpointEligible ? bindContinuationPolicy({ store: ssdStore!,
+        restore: entry => ssdStore!.restore(entry, binding), prompt: promptIds,
+        options, execution, identity: services.identity, namespace: cacheNs, persistence: services.checkpointPersistence }) : null;
+      releaseContinuation = continuation?.release;
+      const checkpointKey = continuation?.key;
       // Both tiers in one call (Layer 0): take() prefers a strictly-longer
       // SSD prefix, restores it zero-copy, and trims — see PromptCache.take.
       const closeCacheLookup = trace?.begin("cache.lookup_restore", {
         mechanism: "serial",
         bypassed: skipPromptCache,
       });
-      const checkpointEntry = checkpointKey
-        ? ssdStore!.findGenerationCheckpoint(promptIds, checkpointKey, cacheNs)
-        : null;
-      const restoredCheckpoint = checkpointEntry
-        ? ssdStore!.restore(checkpointEntry, binding)
-        : null;
-      caches = restoredCheckpoint?.caches ?? [];
-      const checkpoint = restoredCheckpoint?.header.generationCheckpoint;
-      if (restoredCheckpoint && !checkpoint)
-        throw new Error("restored generation checkpoint has no continuation metadata");
-      const resuming = Boolean(restoredCheckpoint && checkpoint);
-      const generationPromptIds = resuming ? restoredCheckpoint!.tokens : promptIds;
+      const checkpoint = continuation?.restore();
+      caches = checkpoint?.caches ?? [];
+      const resuming = Boolean(checkpoint);
+      const generationPromptIds = checkpoint?.cacheTokens ?? promptIds;
       const entry = skipPromptCache || resuming
         ? null
-        : promptCache.take(promptIds, cacheNs);
+        : promptCache.take(promptIds, cacheNs, options.cacheSessionId);
       if (entry) { caches = entry.caches; retain = entry.retain; }
-      if (!restoredCheckpoint && !entry) caches = binding.makeCache();
+      if (!checkpoint && !entry) caches = binding.makeCache();
       closeCacheLookup?.();
       // Prompt-boundary snapshot (the multi-turn agent fix, 2026-07-04): the
       // prompt+gen entry put() below is UNTRIMMABLE at context > sliding
@@ -160,12 +159,6 @@ export function createMlxSerialExecutor(binding: MlxSerialBinding, services: Mlx
       const snapshotBoundary =
         !skipPromptCache && !resuming && boundary >= 256 &&
         boundary > (entry?.tokens.length ?? 0);
-      if (services.memoryBudget) {
-        memoryGuard = binding.enterMemoryGuard?.(
-          services.memoryBudget, generationPromptIds.length, options.prefillChunkSize,
-        );
-        memoryGuard?.check();
-      }
       closeMedia = binding.enterMedia?.(vision);
       if (resuming) {
         const replay = generationPromptIds.slice(promptIds.length);
@@ -190,7 +183,7 @@ export function createMlxSerialExecutor(binding: MlxSerialBinding, services: Mlx
           ? {
               initialPendingToken: checkpoint!.pendingToken,
               initialGeneratedTokens: checkpoint!.generatedTokens,
-              originalPromptTokens: checkpoint!.originalPromptTokens,
+              originalPromptTokens: promptIds.length,
             }
           : {}),
         ...(checkpointEligible
@@ -202,23 +195,9 @@ export function createMlxSerialExecutor(binding: MlxSerialBinding, services: Mlx
                 generatedTokens: number;
                 pendingToken: number;
               }) => {
-                const stored = await ssdStore!.storeGenerationCheckpoint(
-                  state.cacheTokens,
-                  state.caches,
-                  {
-                    key: checkpointKey!,
-                    cacheNs,
-                    originalPromptTokens: promptIds.length,
-                    generatedTokens: state.generatedTokens,
-                    pendingToken: state.pendingToken,
-                    seed: options.seed ?? 0,
-                    seedWasExplicit: options.seedWasExplicit === true,
-                  },
-                );
-                if (stored)
-                  console.log(
-                    `[generation-checkpoint] saved ${state.generatedTokens} emitted tokens`,
-                  );
+                if (services.checkpointPersistence)
+                  continuation!.captureOwned({ ...state, caches: services.cloneState(state.caches) });
+                else await continuation!.capture(state);
               },
             }
           : {}),
@@ -234,7 +213,7 @@ export function createMlxSerialExecutor(binding: MlxSerialBinding, services: Mlx
                 try {
                   const snapshot = ownResource(services.cloneState(caches), disposeResources);
                   try {
-                    promptCache.put(promptIds.slice(0, boundary), snapshot.borrow(), cacheNs);
+                    promptCache.put(promptIds.slice(0, boundary), snapshot.borrow(), cacheNs, undefined, undefined, options.cacheSessionId);
                     snapshot.transfer();
                   } finally { snapshot.close(); }
                 } catch (e) {
@@ -252,15 +231,14 @@ export function createMlxSerialExecutor(binding: MlxSerialBinding, services: Mlx
           : {}),
       }, { trace, mechanism: "serial" });
       for await (const t of gen) {
-        if (t.index % 256 === 0) memoryGuard?.check();
         if ((await onToken(t.token, t.logprobs)) === false) break;
       }
       const s = gen.stats!; // set on completion AND on early break
-      if (checkpointKey) ssdStore!.removeGenerationCheckpoints(checkpointKey);
+      continuation?.complete();
       if (!skipPromptCache) {
         // put() fires onPut → the debounced write-behind SSD snapshot
         // (wired below), covering the batch lane's puts too.
-        promptCache.put(s.cacheTokens, caches, cacheNs, retain);
+        promptCache.put(s.cacheTokens, caches, cacheNs, retain, undefined, options.cacheSessionId);
         caches = []; retain = undefined; // ownership returned to the prefix store
       }
       return s;