mlx-bun 0.3.0 → 0.4.0

package/src/backends/mlx/prefill-rows.ts

@@ -0,0 +1,200 @@
+import type { BatchableCache, Cache } from "../../model/gemma4-base";
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import { clearCache } from "../../mlx/ffi";
+import type { PrefillStep } from "../../inference/prefill";
+import { cloneKvCaches, type CacheCodecProvider } from "../../kv-store";
+import { disposeResources, cleanupFailure } from "../../engine/resources";
+import type { MlxGroupPreparation, Row } from "./batch-group";
+import { evalCacheState } from "./prefill";
+import { prefillCacheLayout } from "./cache-layout";
+import { MlxStateRows } from "./state-rows";
+import type { KvMaintenance } from "./kv-maintenance";
+
+export interface MlxPrefillStep extends PrefillStep {
+  /** Some methods maintain their final target chunk before sampling. */
+  readonly maintain?: boolean;
+}
+
+export interface MlxPrefillState {
+  row: Row;
+  solo: Cache[];
+  pos: number;
+  retain?: () => void;
+  planned?: MlxPrefillStep;
+}
+
+/** The method owns initialization, planning, companion state, snapshots and
+ * completion. This driver owns common target batching and work boundaries.
+ * Checkpoint factories return owned cache state; close/dispose are idempotent. */
+export interface MlxPrefillRowsHost<State extends MlxPrefillState> {
+  stateCodecs?: CacheCodecProvider;
+  maintain?: KvMaintenance;
+  open(row: Row): State;
+  plan(state: State): MlxPrefillStep;
+  /** Already prepared state can complete without any target forward. */
+  ready?(state: State): boolean;
+  forward(ids: MlxArray, caches: Cache[], states: readonly State[]): Promise<MlxArray>;
+  project(hidden: MlxArray, caches: Cache[], completed: readonly State[]): MlxArray | null;
+  /** Consume method context after target maintenance, before snapshots. */
+  afterForward?(ids: MlxArray, caches: Cache[], states: readonly State[], hidden: MlxArray): void | Promise<void>;
+  checkpoint(state: State, capture: () => Cache[], row: number): void;
+  complete(state: State, logits: MlxArray | null, row: number): Promise<void>;
+  reject(row: Row, error: unknown): void;
+  close(state: State): void;
+  filterRows?(keep: readonly number[]): void;
+  dispose?(): void;
+}
+
+/** One target preparation lifecycle at B1/B>1. Late joins retain each row's
+ * planned chunk endpoint, precision transition and checkpoint boundary. */
+export class MlxPrefillRows<State extends MlxPrefillState> implements MlxGroupPreparation {
+  #states: State[] = [];
+  #rows: Row[] = [];
+  #stateRows?: MlxStateRows;
+  constructor(readonly operations: MlxPrefillRowsHost<State>) {}
+  get rows(): readonly Row[] { return this.#rows; }
+  get canAdmit(): boolean { return this.#rows.length > 0; }
+  /** Keep a request's prefill weight until retirement so a nearly completed
+   * request does not repeatedly attract another long prefill into its tail. */
+  get tokenWeight(): number { return this.#states.reduce((sum, state) => sum + state.row.promptTokens - state.row.cachedTokens, 0); }
+
+  admit(row: Row): void {
+    this.#states.push(this.operations.open(row));
+    this.#rows.push(row);
+  }
+
+  #release(state: State): void {
+    const caches = state.solo; state.solo = [];
+    const retain = state.retain; state.retain = undefined;
+    disposeResources([...caches, { dispose: () => retain?.() }, { dispose: () => this.operations.close(state) }]);
+  }
+  #filter(keep: number[]): void {
+    if (keep.length) this.#stateRows?.filterRows(keep);
+    else { const rows = this.#stateRows; this.#stateRows = undefined; rows?.dispose(); }
+    this.operations.filterRows?.(keep);
+    this.#states = keep.map(row => this.#states[row]!);
+    this.#rows = keep.map(row => this.#rows[row]!);
+  }
+  /** Admission changes membership only between forwards. Existing planned
+   * endpoints survive a join, including partially consumed drain chunks. */
+  #synchronizeRows(): void {
+    if (!this.#stateRows && this.#states.length === 1) return;
+    if (!this.#stateRows) {
+      for (const state of this.#states) this.operations.maintain?.preparePrefill?.(state.solo);
+      const layouts: BatchableCache[] = [];
+      try { for (const cache of this.#states[0]!.solo) layouts.push(prefillCacheLayout(cache)); }
+      catch (error) { return cleanupFailure(error, () => disposeResources(layouts)); }
+      this.#stateRows = new MlxStateRows(layouts);
+      this.#stateRows.mergeRows(this.#states.map(state => state.solo));
+    } else {
+      for (const state of this.#states.slice(this.#stateRows.rowCount)) {
+        this.operations.maintain?.preparePrefill?.(state.solo);
+        this.#stateRows.append(state.solo);
+      }
+    }
+    // Resolve input-backed copies before releasing SSD restore leases.
+    if (this.#states.some(state => state.solo.length)) {
+      evalCacheState(this.#stateRows.caches);
+      for (const state of this.#states) {
+        const caches = state.solo; state.solo = [];
+        const retain = state.retain; state.retain = undefined;
+        disposeResources([...caches, { dispose: () => retain?.() }]);
+      }
+    }
+    for (const cache of this.#stateRows.caches) cache.prefillMaintenance?.beginPrefill();
+  }
+
+  async advance(): Promise<boolean> {
+    // Complete restored inputs before merging: their cache coverage must not
+    // introduce padding or a zero-width forward into unrelated cold rows.
+    if (this.operations.ready && this.#states.some(state => this.operations.ready!(state))) {
+      const remaining: State[] = [];
+      for (const state of this.#states) {
+        if (!this.operations.ready(state)) { remaining.push(state); continue; }
+        try {
+          state.row.req.signal?.throwIfAborted();
+          await this.operations.complete(state, null, -1);
+          state.solo = []; state.retain = undefined;
+        } catch (error) { this.operations.reject(state.row, error); this.#release(state); }
+        finally { this.operations.close(state); }
+      }
+      this.#states = remaining; this.#rows = remaining.map(state => state.row);
+      if (!remaining.length) { this.dispose(); return true; }
+    }
+    this.#synchronizeRows();
+    const keep: number[] = [];
+    for (const [index, state] of this.#states.entries()) {
+      if (!state.row.req.signal?.aborted) { keep.push(index); continue; }
+      this.#release(state); this.operations.reject(state.row, state.row.req.signal.reason);
+    }
+    if (keep.length !== this.#states.length) this.#filter(keep);
+    if (!this.#states.length) { this.dispose(); return true; }
+    while (this.#states.length) {
+      for (const state of this.#states) state.planned ??= this.operations.plan(state);
+      const count = Math.min(...this.#states.map(state => state.planned!.end - state.pos));
+      const closes = this.#states.map(state => state.row.req.trace?.begin("prefill.chunk", {
+        mechanism: "continuous", startToken: state.pos, batchSize: this.#states.length,
+      }));
+      const caches = this.#stateRows?.caches ?? this.#states[0]!.solo;
+      try {
+        using ids = count ? ops.fromInt32(this.#states.flatMap(state => state.row.req.promptIds.slice(state.pos, state.pos + count)),
+          [this.#states.length, count]) : null;
+        using hidden = ids ? await this.operations.forward(ids, caches, this.#states) : null;
+        const drains: number[] = [], finals: number[] = [];
+        let batchYield = false;
+        for (const [row, state] of this.#states.entries()) {
+          state.pos += count;
+          if (state.pos !== state.planned!.end) continue;
+          if (state.planned!.kind === "final") finals.push(row);
+          else drains.push(row);
+          batchYield ||= state.planned!.batchYield;
+        }
+        const maintained = [...drains, ...finals.filter(row => this.#states[row]!.planned!.maintain)];
+        if (count && (maintained.length || !finals.length)) {
+          evalCacheState(caches);
+          if (this.#stateRows) {
+            for (const cache of caches) cache.prefillMaintenance?.commitPrefill(maintained);
+          } else if (maintained.length) this.operations.maintain?.(this.#states[0]!.solo);
+          clearCache();
+        }
+        if (ids) {
+          const companion = this.operations.afterForward?.(ids, caches, this.#states, hidden!);
+          if (companion) await companion;
+        }
+        for (const row of [...drains, ...finals]) {
+          const state = this.#states[row]!;
+          if (state.planned!.snapshot) this.operations.checkpoint(state, () =>
+            this.#stateRows ? this.#stateRows.extractRow(row) : cloneKvCaches(state.solo, this.operations.stateCodecs), row);
+          if (!finals.includes(row)) state.planned = undefined;
+        }
+        using logits = hidden && finals.length ? this.operations.project(hidden, caches, finals.map(row => this.#states[row]!)) : null;
+        for (const row of finals) {
+          const state = this.#states[row]!;
+          if (this.#stateRows) state.solo = this.#stateRows.extractRow(row);
+          using part = logits ? logits.slice([row, logits.shape[1]! - 1, 0], [row + 1, logits.shape[1]!, logits.shape[2]!]) : null;
+          try {
+            state.row.req.signal?.throwIfAborted();
+            await this.operations.complete(state, part, row);
+            state.solo = []; state.retain = undefined;
+          } catch (error) {
+            this.operations.reject(state.row, error);
+            this.#release(state);
+          } finally { this.operations.close(state); }
+        }
+        for (const cache of caches) (cache as { releaseRopeArr?: () => void }).releaseRopeArr?.();
+        if (finals.length) this.#filter(this.#states.flatMap((_, row) => finals.includes(row) ? [] : [row]));
+        if (!this.#states.length) { this.dispose(); return true; }
+        if (batchYield || finals.length) return false;
+      } finally { for (const close of closes) close?.(); }
+    }
+    return true;
+  }
+
+  dispose(): void {
+    const states = this.#states; this.#states = []; this.#rows = [];
+    const rows = this.#stateRows; this.#stateRows = undefined;
+    disposeResources([...(rows ? [rows] : []), ...states.map(state => ({ dispose: () => this.#release(state) })),
+      { dispose: () => this.operations.dispose?.() }]);
+  }
+}

package/src/backends/mlx/prefill.ts

@@ -9,7 +9,7 @@ import { withResource } from "../../engine/resources";
 /** State views with explicit ownership are released even on deferred errors.
  * Owned views precede borrowed state, preserving the serial eval ordering. */
 export function evalCacheState(cache: Cache[]): void {
-  withResource(leaseCacheStates(cache), (state) => ops.evalAll([...state]));
+  withResource(leaseCacheStates(cache), (state) => { if (state.length) ops.evalAll([...state]); });
 }
 
 /** Forward once. Drains evaluate state, apply KV maintenance, then clear the

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/request-state-policy.ts

@@ -0,0 +1,41 @@
+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);
+    }
+  }
+}
+
+export function bindPagedRequestState(model: RuntimeModel, options: GenerateOptions,
+  capacityTokens: number): MlxRequestStatePolicy | undefined {
+  if (!options.pagedKv) return undefined;
+  const blockSize = options.pagedKv.blockSize ?? PagedKVCache.DEFAULT_BLOCK_SIZE;
+  return { key: `paged:${blockSize}`, create() {
+    const caches = model.makeCache();
+    try { maybePageKv(caches, { pagedKv: { blockSize } }, capacityTokens); return caches; }
+    catch (error) { disposeResources(caches); throw error; }
+  } };
+}

package/src/backends/mlx/rollback.ts

@@ -1,7 +1,36 @@
 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(): 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,
+    begin: () => cache.specRoundBegin(),
+    resolve(accepted: readonly number[], drafts: number) {
+      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 +53,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) => {
@@ -41,8 +79,8 @@ export function bindCacheRollback(caches: readonly Cache[]): SpeculativeTransact
       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,30 @@
+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 +35,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 +48,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 +66,17 @@ 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 () => {
+    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,11 +85,12 @@ 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.
@@ -111,32 +107,26 @@ export function createMlxSerialExecutor(binding: MlxSerialBinding, services: Mlx
       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);
       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 +150,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 +174,7 @@ export function createMlxSerialExecutor(binding: MlxSerialBinding, services: Mlx
           ? {
               initialPendingToken: checkpoint!.pendingToken,
               initialGeneratedTokens: checkpoint!.generatedTokens,
-              originalPromptTokens: checkpoint!.originalPromptTokens,
+              originalPromptTokens: promptIds.length,
             }
           : {}),
         ...(checkpointEligible
@@ -202,23 +186,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);
               },
             }
           : {}),
@@ -252,11 +222,10 @@ 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.