mlx-bun 0.2.0 → 0.4.0

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

@@ -0,0 +1,242 @@
+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 { 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 { bindContinuationPolicy } from "./continuation";
+import { runtimeConfig, withRuntimeConfig, type RuntimeConfig } from "../../runtime-config";
+import { bindLegacySpeculativeModel } from "./speculative";
+
+/** 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?: (
+    promptIds: Parameters<SerialRun>[0], options: Parameters<SerialRun>[1],
+    onToken: Parameters<SerialRun>[2], services?: SpecRunServices,
+  ) => ReturnType<SerialRun>;
+  makeCache(): Cache[];
+  enterMedia?(vision?: Vision): () => void;
+}
+
+/** Family-specific context is bound once at the compatibility boundary. */
+export function bindLegacySerialModel(
+  model: RuntimeModel,
+  draft?: { provider: DraftProvider; numDraftTokens: number },
+): MlxSerialBinding {
+  const speculative = draft ? bindLegacySpeculativeModel(model, draft.provider) : undefined;
+  return {
+    runtime: runtimeConfig(),
+    generate: bindGeneration(model),
+    makeCache: model.makeCache.bind(model),
+    ...(speculative && draft ? { speculate: (async (prompt, options, onToken, services) => {
+      const { specRun } = await import("../../spec/serve-loop");
+      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;
+      return () => { model.mrope = previous; };
+    } } : {}),
+  };
+}
+
+export interface MlxSerialServices {
+  readonly promptCache: MlxPrefixCache & { readonly maxBytes?: number };
+  readonly checkpoints: Pick<SsdCacheStore, "findGenerationCheckpoint" | "restore" |
+    "storeGenerationCheckpoint" | "removeGenerationCheckpoints"> | null;
+  readonly checkpointEveryTokens?: number;
+  readonly checkpointPersistence?: import("./continuation-persistence").ContinuationPersistence;
+  /** Artifact, implementation, state ABI and codec identity captured at load. */
+  readonly identity: unknown;
+  adapterNamespace(adapters: string[]): string;
+  cloneState(caches: Cache[]): Cache[];
+}
+
+/** Cache lookup, replay, prompt snapshots and checkpoint persistence belong to
+ * native execution. HTTP only composes this service and frames its output.
+ * The caller transfers prepared native inputs; caches created/taken here are
+ * owned until returned to the prefix store or disposed after execution. */
+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;
+    const cleanup = ownResource(null, () => disposeResources([
+      { dispose() {
+        disposeResources(caches);
+        retain?.(); // backing release requires successful cache disposal
+      } },
+      { dispose: () => closeMedia?.() },
+      { dispose: () => releaseContinuation?.() },
+      ...[options.grammar, vision?.embeddings, vision?.imageMask,
+        vision?.multimodalMask, options.visionPixels].filter((value) => value != null),
+    ]));
+    try {
+      options.signal?.throwIfAborted();
+      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, 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) : "";
+      // 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.
+      const skipPromptCache = !execution.promptCache;
+      const checkpointEvery = services.checkpointEveryTokens;
+      const checkpointEligible = execution.checkpoint;
+      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 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 (!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
+      // window (wrapped rings) and under quantized KV (mid-group), so any
+      // decode→encode roundtrip drift in the reply the client sends back
+      // turns the next turn into a total miss (measured: 12B turn-2 TTFT
+      // 8.9 s instead of ~0.2 s). A prompt-ONLY entry is always an exact
+      // prefix of the next turn's rendering regardless of reply drift.
+      // Zero-copy (cloneKvCaches = slice views); only for substantial cold
+      // prefills, where the re-prefill it saves is worth an extra entry.
+      // The oracle invariant (mlx-lm insert_segments): a trim-free STRICT
+      // prefix of the prompt exists for EVERY substantial request — cap the
+      // boundary at len-1 so even a stableLen == len prompt (e4b: the
+      // template tail survives the probe render) snapshots prompt[:-1]. An
+      // exact repeat then matches with trimNeeded == 0, bypassing
+      // isTrimmable() entirely — the only reuse path a wrapped ring has.
+      const boundary = Math.min(options.snapshotAt ?? promptIds.length, promptIds.length - 1);
+      // Re-snapshot on EVERY substantial request whose stable boundary extends
+      // past the cached prefix; the clone is zero-copy views, so re-putting
+      // is ~free.
+      const snapshotBoundary =
+        !skipPromptCache && !resuming && boundary >= 256 &&
+        boundary > (entry?.tokens.length ?? 0);
+      closeMedia = binding.enterMedia?.(vision);
+      if (resuming) {
+        const replay = generationPromptIds.slice(promptIds.length);
+        console.log(
+          `[generation-checkpoint] resuming ${checkpointKey} at ` +
+          `${replay.length} emitted tokens`,
+        );
+        for (const token of replay) {
+          options.signal?.throwIfAborted();
+          if ((await onToken(token)) === false)
+            throw new Error("saved generation prefix triggered a terminal stop while replaying");
+        }
+      }
+      const gen = binding.generate(generationPromptIds, {
+        ...options,
+        decodePolicy: execution,
+        ...(resuming ? { seed: checkpoint!.seed } : {}),
+        fill: execution.fill ? options.fill : undefined,
+        pagedKv, // request-scoped (undefined strips the server-wide flag)
+        cache: caches,
+        ...(resuming
+          ? {
+              initialPendingToken: checkpoint!.pendingToken,
+              initialGeneratedTokens: checkpoint!.generatedTokens,
+              originalPromptTokens: promptIds.length,
+            }
+          : {}),
+        ...(checkpointEligible
+          ? {
+              checkpointEveryTokens: checkpointEvery,
+              onDecodeCheckpoint: async (state: {
+                cacheTokens: number[];
+                caches: Cache[];
+                generatedTokens: number;
+                pendingToken: number;
+              }) => {
+                if (services.checkpointPersistence)
+                  continuation!.captureOwned({ ...state, caches: services.cloneState(state.caches) });
+                else await continuation!.capture(state);
+              },
+            }
+          : {}),
+        ...(snapshotBoundary
+          ? {
+              // snapshotAt MUST travel with the hook: generate() splits the
+              // prefill at exactly this many tokens and fires the hook while
+              // the caches hold exactly that prefix — putting boundary
+              // tokens against caches at any other offset is silent KV
+              // corruption.
+              snapshotAt: boundary,
+              onPrefillDone: () => {
+                try {
+                  const snapshot = ownResource(services.cloneState(caches), disposeResources);
+                  try {
+                    promptCache.put(promptIds.slice(0, boundary), snapshot.borrow(), cacheNs);
+                    snapshot.transfer();
+                  } finally { snapshot.close(); }
+                } catch (e) {
+                  console.warn(`prompt-boundary snapshot skipped: ${(e as Error).message}`);
+                }
+              },
+            }
+          : {}),
+        ...(vision
+          ? {
+              promptEmbeddings: vision.embeddings,
+              ...(vision.imageMask ? { imageMask: vision.imageMask } : {}),
+              ...(vision.multimodalMask ? { multimodalMask: vision.multimodalMask } : {}),
+            }
+          : {}),
+      }, { trace, mechanism: "serial" });
+      for await (const t of gen) {
+        if ((await onToken(t.token, t.logprobs)) === false) break;
+      }
+      const s = gen.stats!; // set on completion AND on early break
+      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);
+        caches = []; retain = undefined; // ownership returned to the prefix store
+      }
+      return s;
+    } catch (error) {
+      cleanupFailure(error, () => cleanup.close());
+    } finally {
+      cleanup.close();
+    }
+  });
+}

package/src/backends/mlx/speculative-group.ts

@@ -0,0 +1,359 @@
+import {captureSpeculativeOptions, speculativePrefixNamespace} from "../../spec/cache-identity";
+import type { RuntimeModel } from "../../model/factory";
+import type { GenerateOptions } from "../../generate";
+import type { DraftProvider, DraftRowGroup, DraftPrefillGroup, DraftRowCheckpoint } from "../../spec/source";
+import { makeSampler, makeStepSampler, readStepExtras, type NumberStepSampler, type Sampler } from "../../sampler";
+import { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import { clearCache } from "../../mlx/ffi";
+import { applyStateChanges, cleanupFailure, disposeResources } from "../../engine/resources";
+import { cloneKvCaches } from "../../kv-store";
+import type { Cache } from "../../model/gemma4-base";
+import { targetCacheLayout } from "./cache-layout";
+import { MlxStateRows } from "./state-rows";
+import { bindLegacySpeculativeModel, type MlxSpeculativeBinding } from "./speculative";
+import { bindLegacyDraftTarget } from "./draft-target";
+import { bindRowCacheRollback } from "./rollback";
+import { createKvMaintenance } from "./kv-maintenance";
+import { disposeAttachments, type CheckpointAttachment } from "./checkpoint-state";
+import { advanceSpeculativeOutputs } from "./speculative-round";
+import { MlxPrefillRows, type MlxPrefillState } from "./prefill-rows";
+import type { MlxGroupMethodHost, MlxGroupMethodRequest, MlxGroupPreparation, MlxGroupedMethod, Row } from "./batch-group";
+
+interface RequestState {
+  sampling: NumberStepSampler;
+  draftSampling: Sampler;
+  retain?: () => void;
+  namespace: string;
+  prefixLength: number;
+  /** Generated inputs retained by committed target/draft rounds. The pending
+   * correction is excluded until a later round actually processes it. */
+  processed?: number[];
+}
+
+/** Binding owns graph/layout selection. The executor receives only the method
+ * key and lifecycle; sampling, checkpoints and numerical state remain ports. */
+export function bindSpeculativeGroupRequests(model: RuntimeModel, provider: DraftProvider, depth: number) {
+  const binding = bindLegacySpeculativeModel(model, provider);
+  return (input: GenerateOptions): MlxGroupMethodRequest => {
+    const options = captureSpeculativeOptions(input);
+    return {
+      key: JSON.stringify(["speculative", provider.id, depth, options.kvBits ?? null,
+        options.kvGroupSize ?? 64, options.quantizedKvStart ?? null, options.turboQuant ?? null,
+        ...(options.kvConfig?.length ? [options.kvConfig] : [])]),
+      data: options,
+      open: host => new SpeculativeGroup(host, model, provider, binding, depth),
+    };
+  };
+}
+
+class SpeculativeGroup implements MlxGroupedMethod {
+  #target: MlxStateRows<ReturnType<typeof targetCacheLayout>> | null = null;
+  #draft: DraftRowGroup | null = null;
+  readonly #requests = new Map<Row, RequestState>();
+  #steps = 0;
+
+  constructor(readonly host: MlxGroupMethodHost, readonly model: RuntimeModel,
+    readonly provider: DraftProvider, readonly binding: MlxSpeculativeBinding, readonly depth: number) {}
+
+  prepare(first: Row): MlxGroupPreparation {
+    const method = this;
+    const options = first.req.method!.data as GenerateOptions;
+    const maintain = createKvMaintenance(options);
+    interface Preparation extends MlxPrefillState {
+      request: RequestState;
+      end: number;
+      chunkSize: number;
+      pendingPrompt: boolean;
+      boundary: number;
+      transferred: boolean;
+      closed: boolean;
+    }
+    interface Ready {
+      row: Row;
+      request: RequestState;
+      caches: Cache[];
+      draft: DraftRowCheckpoint;
+      retain?: () => void;
+      transferred: boolean;
+    }
+    let prefix: DraftPrefillGroup | undefined;
+    let context: MlxArray | null = null;
+    const ready: Ready[] = [];
+    const closeReady = (state: Ready) => {
+      const caches = state.caches.splice(0), attachments = [state.draft.attachment];
+      const release = state.retain; state.retain = undefined;
+      disposeResources([...caches, { dispose: () => disposeAttachments(attachments) },
+        { dispose: () => release?.() }, ...(!state.transferred ? [state.request.sampling] : [])]);
+    };
+    const capture = (state: Preparation, snapshot: () => Cache[], index: number) => {
+      let caches: Cache[] = [], attachments: CheckpointAttachment[] = [];
+      try {
+        caches = snapshot(); attachments.push(prefix!.capture(index).attachment);
+        method.host.promptCache!.put(state.row.req.promptIds.slice(0, state.pos), caches,
+          state.request.namespace, undefined, attachments);
+        caches = []; attachments = [];
+      } finally { disposeResources([...caches, { dispose: () => disposeAttachments(attachments) }]); }
+    };
+    const target = new MlxPrefillRows<Preparation>({
+      maintain,
+      open(row) {
+        const options = row.req.method!.data as GenerateOptions;
+        row.spec = { drafted: 0, accepted: 0, targetCalls: 0, rejected: 0, rounds: 0,
+          acceptanceLengths: [], tokensPerForward: 0, forwardsSaved: 0, draftedByPos: [], acceptedByPos: [] };
+        const sampler = makeSampler(options);
+        const request: RequestState = { draftSampling: sampler, namespace: "", prefixLength: 0,
+          processed: method.host.promptCache ? [] : undefined,
+          sampling: makeStepSampler(options, { tokenRepresentation: "number", grammarWait: "before-sample",
+            historyUpdate: "after-sample", initialHistory: row.req.promptIds, acceptGrammar: true,
+            eosTokenIds: row.req.eosTokenIds, captureSelectedLogprob: options.logprobs === true,
+            captureTopLogprobs: options.topLogprobs }) };
+        let caches: Cache[] = [], retained: (() => void) | undefined;
+        let appended = false;
+        const previous = prefix?.rowCount ?? 0;
+        try {
+          caches = method.binding.makeCache();
+          prefix ??= method.provider.grouped!.openPrefill({ target: bindLegacyDraftTarget(method.model, caches), checkpoints: [] });
+          const prompt = row.req.promptIds;
+          const pendingPrompt = prefix.prefillMode !== "full" && method.binding.prefillTailSplit && prompt.length > 1;
+          request.prefixLength = prompt.length - Number(pendingPrompt);
+          request.namespace = speculativePrefixNamespace(prefix.namespace, row.cacheNamespace ?? "", options);
+          const hit = method.host.promptCache?.take(prompt, request.namespace);
+          if (hit) {
+            try {
+              disposeResources(caches); caches = hit.caches.splice(0);
+              retained = hit.retain; hit.retain = undefined;
+              prefix.append([{ processedTokens: hit.tokens.length, attachment: hit.attachments![0]! }]);
+              appended = true;
+              // Target membership later releases backing leases after its own
+              // copies resolve. Settle companion copies under the same lease.
+              if (retained) prefix.materialize();
+              row.cachedTokens = hit.tokens.length;
+            } finally { disposeResources([...hit.caches, { dispose: () => disposeAttachments(hit.attachments) },
+              { dispose: () => hit.retain?.() }]); }
+          } else { prefix.append([null]); appended = true; }
+          return { row, request, solo: caches, retain: retained, pos: row.cachedTokens,
+            end: request.prefixLength, pendingPrompt, transferred: false, closed: false,
+            chunkSize: options.prefillChunkSize ?? method.host.prefillChunkSize,
+            boundary: method.host.promptCache ? Math.min(row.req.snapshotAt ?? prompt.length, prompt.length - 1) : -1 };
+        } catch (error) {
+          return cleanupFailure(error, () => disposeResources([
+            { dispose: () => { if (appended) prefix!.filterRows(Array.from({ length: previous }, (_, i) => i)); } },
+            ...caches, request.sampling, { dispose: () => retained?.() },
+          ]));
+        }
+      },
+      plan(state) {
+        const end = Math.min(state.pos + state.chunkSize, state.end, state.boundary > state.pos ? state.boundary : Infinity);
+        return { start: state.pos, end, kind: end === state.end ? "final" : "drain",
+          snapshot: end === state.boundary && end > state.row.cachedTokens, batchYield: true, maintain: true };
+      },
+      async forward(ids, caches) {
+        const taps = prefix!.tapLayers;
+        const output = await method.binding.forward(ids, caches, taps.length ? [...taps] : undefined);
+        context = output.ctxML;
+        return output.hidden;
+      },
+      afterForward(ids, _caches, _states, hidden) {
+        // Tapped providers consume their captured layers; assistant drafting
+        // consumes the true target hidden. Other providers ignore context.
+        const captured = context; context = null;
+        try {
+          const pending = prefix!.prefill(ids, captured ?? hidden);
+          if (pending) return pending.finally(() => captured?.dispose());
+        } catch (error) { captured?.dispose(); throw error; }
+        captured?.dispose();
+      },
+      project(hidden, _caches, completed) {
+        if (completed.every(state => state.pendingPrompt)) return null;
+        using last = hidden.slice([0, hidden.shape[1]! - 1, 0], [hidden.shape[0]!, hidden.shape[1]!, hidden.shape[2]!]);
+        return method.binding.projectLogits(last);
+      },
+      checkpoint: capture,
+      async complete(state, logits, index) {
+        const row = state.row;
+        row.spec!.targetCalls++;
+        if (state.pendingPrompt) { row.current = row.req.promptIds[state.end]!; row.generated = row.sampled = 0; }
+        else {
+          using scores = ops.reshape(logits!, [1, logits!.shape.at(-1)!]);
+          const sampled = await state.request.sampling.sample(scores, 0);
+          row.current = sampled.token; row.generated = row.sampled = 1;
+          if (row.req.eosTokenIds.includes(row.current)) {
+            if (method.host.promptCache) capture(state, () => cloneKvCaches(state.solo), index);
+            disposeResources(state.solo.splice(0)); state.retain?.(); state.retain = undefined;
+            method.host.finish(row, "stop"); return;
+          }
+          const more = await method.host.publish(row, row.current, readStepExtras(sampled.extras));
+          if (more === false || row.req.grammar?.isTerminated || row.generated >= row.req.maxTokens) {
+            if (method.host.promptCache && !row.req.signal?.aborted) capture(state, () => cloneKvCaches(state.solo), index);
+            disposeResources(state.solo.splice(0)); state.retain?.(); state.retain = undefined;
+            method.host.finish(row, more === false || row.req.grammar?.isTerminated ? "stop" : "length"); return;
+          }
+        }
+        ready.push({ row, request: state.request, caches: state.solo, retain: state.retain,
+          draft: prefix!.capture(index), transferred: false });
+        state.transferred = true;
+      },
+      reject(row, error) { row.reject(error); },
+      close(state) {
+        if (state.closed) return;
+        state.closed = true;
+        if (!state.transferred) state.request.sampling.dispose();
+      },
+      filterRows(keep) { prefix!.filterRows(keep); },
+      dispose() {
+        const held = [prefix, context].filter(value => value != null);
+        prefix = undefined; context = null; disposeResources(held);
+      },
+    });
+    const preparation: MlxGroupPreparation = {
+      get rows() { return [...target.rows, ...ready.map(state => state.row)]; },
+      get canAdmit() { return target.canAdmit; },
+      get tokenWeight() { return target.tokenWeight + ready.reduce((sum, state) => sum + state.row.promptTokens - state.row.cachedTokens, 0); },
+      admit: row => target.admit(row),
+      advance: async () => {
+        // Publication returns to scheduling before prepared state joins active
+        // decode. This preserves the first-output flush boundary at every B.
+        if (ready.length) {
+          while (ready.length) {
+            const state = ready[0]!;
+            try {
+              if (state.row.req.signal?.aborted) {
+                state.row.reject(state.row.req.signal.reason);
+                continue;
+              }
+              maintain.prepareBatch?.(state.caches);
+              this.#target ??= new MlxStateRows(state.caches.map(targetCacheLayout));
+              this.#draft ??= this.provider.grouped!.open({ target: bindLegacyDraftTarget(this.model, this.#target.caches),
+                checkpoints: [], sampling: { sample: (lp, steps) => this.#sampleDraftRows(lp, steps) } });
+              applyStateChanges([() => this.#target!.prepareAppend(state.caches),
+                () => this.#draft!.prepareAppend([state.draft]), () => ({ commit: () => {
+                  state.request.retain = state.retain; state.retain = undefined;
+                  this.#requests.set(state.row, state.request); state.transferred = true;
+                  this.host.join(state.row);
+                }, dispose() {} })]);
+            } finally { ready.shift(); closeReady(state); }
+          }
+          return target.rows.length === 0;
+        }
+        const done = await target.advance();
+        return done && ready.length === 0;
+      },
+      dispose() {
+        const held = ready.splice(0);
+        disposeResources([target, ...held.map(state => ({ dispose: () => closeReady(state) }))]);
+      },
+    };
+    try { target.admit(first); return preparation; }
+    catch (error) { return cleanupFailure(error, () => preparation.dispose()); }
+  }
+
+  #sampleDraftRows(logprobs: MlxArray, steps: readonly number[]): MlxArray {
+    const tokens: MlxArray[] = [];
+    try {
+      for (let row = 0; row < this.host.rows.length; row++) {
+        using scores = logprobs.slice([row, 0], [row + 1, logprobs.shape[1]!]);
+        tokens.push(this.#requests.get(this.host.rows[row]!)!.draftSampling(scores, steps[row]!));
+      }
+      return ops.concatAxis(tokens, 0);
+    } finally { disposeResources(tokens); }
+  }
+
+  async advance(): Promise<void> {
+    const active = [...this.host.rows];
+    const live = active.flatMap((row, index) => row.req.signal?.aborted ? [] : [index]);
+    if (live.length !== active.length) {
+      for (const row of active) if (row.req.signal?.aborted) row.reject(row.req.signal.reason);
+      this.host.filterRows(live);
+    }
+    const rows = [...this.host.rows];
+    if (!rows.length) return;
+    const depth = Math.min(this.depth, Math.max(...rows.map(row => row.req.maxTokens - row.generated)));
+    const halted = new Set<Row>();
+    const pin = this.binding.pinVerify?.();
+    try {
+      const completed = await advanceSpeculativeOutputs(rows.map(row => ({ pending: row.current,
+        step: row.sampled, remaining: row.req.maxTokens - row.generated, eosTokenIds: row.req.eosTokenIds,
+        sampling: this.#requests.get(row)!.sampling, grammarDone: () => row.req.grammar?.isTerminated === true,
+        output: { commit: async (ids, metadata) => {
+          row.generated++;
+          const more = await this.host.publish(row, ids[0]!, metadata?.[0]);
+          if (more === false) { halted.add(row); return false; }
+        } },
+      })), depth, this.#draft!, {
+        transaction: bindRowCacheRollback(this.#target!.caches, rows.length),
+        forward: async ids => {
+          const taps = this.#draft!.tapLayers;
+          const result = await this.binding.forward(ids, this.#target!.caches, taps.length ? [...taps] : undefined);
+          const context = result.ctxML ?? result.hidden;
+          try { return { logits: this.binding.projectLogits(result.hidden), context }; }
+          catch (error) { context.dispose(); throw error; }
+          finally { if (result.ctxML) result.hidden.dispose(); }
+        },
+      });
+      const keep: number[] = [];
+      for (let index = 0; index < rows.length; index++) {
+        const row = rows[index]!, output = completed.outputs[index]!, round = completed.rounds[index]!;
+        const request = this.#requests.get(row)!;
+        if (output.kind !== "failed" && output.kind !== "cancelled")
+          request.processed?.push(row.current, ...round.drafts.slice(0, output.accepted));
+        const stats = row.spec!;
+        stats.drafted += round.drafts.length;
+        stats.accepted += round.acceptance.accepted;
+        stats.rejected = stats.drafted - stats.accepted;
+        stats.rounds!++; stats.targetCalls++;
+        stats.acceptanceLengths!.push(round.acceptance.accepted);
+        for (let p = 0; p < round.drafts.length; p++) stats.draftedByPos![p] = (stats.draftedByPos![p] ?? 0) + 1;
+        for (let p = 0; p < round.acceptance.accepted; p++) stats.acceptedByPos![p] = (stats.acceptedByPos![p] ?? 0) + 1;
+        row.sampled += round.acceptance.emitted.length + Number(round.acceptance.sawEos);
+        // A consumer may stop inside a verified burst before its terminal EOS.
+        if (output.kind === "stop" && !halted.has(row) && round.acceptance.sawEos &&
+          output.generated === round.acceptance.emitted.length && row.generated < row.req.maxTokens)
+          row.generated++;
+        stats.tokensPerForward = row.generated / stats.rounds!;
+        stats.forwardsSaved = Math.max(0, row.generated - 1 - stats.rounds!);
+        if (output.kind === "continue") { row.current = output.pending; keep.push(index); }
+        else if (output.kind === "failed") row.reject(output.error);
+        else if (output.kind === "cancelled") row.reject(row.req.signal?.reason ?? new Error(output.reason));
+        else {
+          if (!row.req.signal?.aborted) this.#publishCheckpoint(row, index, request);
+          this.host.finish(row, output.kind);
+        }
+      }
+      if (keep.length !== rows.length) this.host.filterRows(keep);
+      if (++this.#steps % 256 === 0) clearCache();
+    } finally { pin?.close(); }
+  }
+
+  /** The method aligns token coverage and captures its state. The cache owns
+   * RAM retention and eventual SSD persistence; no storage work is awaited. */
+  #publishCheckpoint(row: Row, index: number, request: RequestState): void {
+    if (!this.host.promptCache || !request.processed) return;
+    let target: Cache[] = [], attachments: CheckpointAttachment[] = [];
+    try {
+      target = this.#target!.extractRow(index);
+      attachments.push(this.#draft!.capture(index).attachment);
+      this.host.promptCache.put([...row.req.promptIds.slice(0, request.prefixLength), ...request.processed], target,
+        request.namespace, undefined, attachments);
+      target = []; attachments = [];
+    } finally { disposeResources([...target, { dispose: () => disposeAttachments(attachments) }]); }
+  }
+
+  filterRows(keep: readonly number[], _discard: boolean): void {
+    this.#target?.filterRows(keep); this.#draft?.filterRows(keep);
+    const retained = new Set(keep.map(index => this.host.rows[index]!));
+    const releases: { dispose(): void }[] = [];
+    for (const [row, request] of this.#requests) if (!retained.has(row)) {
+      this.#requests.delete(row);
+      releases.push(request.sampling, { dispose: () => request.retain?.() });
+    }
+    disposeResources(releases);
+  }
+
+  dispose(): void {
+    const resources = [...(this.#target ? [this.#target] : []), ...(this.#draft ? [this.#draft] : []),
+      ...[...this.#requests.values()].flatMap(request => [request.sampling, { dispose: () => request.retain?.() }])];
+    this.#target = null; this.#draft = null; this.#requests.clear();
+    disposeResources(resources);
+  }
+}