@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/execution/methods.ts

@@ -0,0 +1,75 @@
+import { runtimeConfig, withRuntimeConfig, type RuntimeConfig } from "../runtime/config";
+import { snapshotGenerationPolicy } from "./request-policy";
+import type { GenerationOutput, InferenceMethod, MethodResult } from "../contracts/portable/generation";
+import { throwIfCancelled } from "./cancellation";
+import { generateAutoregressive, generateDenoising, type GenerateOptions, type GenerateStats } from "../generation/index";
+import type { MlxAutoregressiveBinding } from "../generation/bindings/autoregressive";
+import type { MlxDenoisingBinding } from "../generation/bindings/denoising";
+
+/** The native algorithm owns cleanup in its finally. This host adapter owns
+ * cancellation subscriptions and waits for the algorithm before closing. */
+function nativeMethod<Metrics>(
+  id: string,
+  execute: (output: GenerationOutput, signal: AbortSignal) => Promise<MethodResult<Metrics>>,
+  runtime: RuntimeConfig,
+): InferenceMethod<Metrics> {
+  return {
+    id,
+    async createRun(cancellation) {
+      throwIfCancelled(cancellation);
+      const abort = new AbortController();
+      const unsubscribe = cancellation.subscribe((reason) => abort.abort(
+        new DOMException(`generation cancelled: ${reason}`, "AbortError"),
+      ));
+      let work: Promise<MethodResult<Metrics>> | undefined;
+      let closing: Promise<void> | undefined;
+      return {
+        execute(output) {
+          if (work || closing) return Promise.reject(new Error("method run can execute only once"));
+          work = withRuntimeConfig(runtime, () => execute(output, abort.signal));
+          return work;
+        },
+        close() {
+          return closing ??= (async () => {
+            abort.abort(new DOMException("method closed", "AbortError"));
+            try { await work; } catch { /* execute owns the algorithm's error */ }
+            finally { unsubscribe(); }
+          })();
+        },
+      };
+    },
+  };
+}
+
+/** Each plan supplies its prepared per-run resources. Caches/media in options
+ * retain the same ownership contract as generateAutoregressive. */
+export function createAutoregressiveMethod(
+  binding: MlxAutoregressiveBinding, promptIds: readonly number[], options: GenerateOptions,
+): InferenceMethod<GenerateStats> {
+  const prompt = [...promptIds];
+  options = snapshotGenerationPolicy(options);
+  return nativeMethod("mlx-autoregressive", async (output, signal) => {
+    const generation = generateAutoregressive(binding, prompt, { ...options, signal });
+    for await (const item of generation)
+      if (await output.commit([item.token], item.logprobs ? [item.logprobs] : undefined) === false) break;
+    const metrics = generation.stats;
+    if (!metrics) throw new Error("AR method did not settle its metrics");
+    return { finishReason: metrics.finishReason ?? (metrics.generatedTokens >= (options.maxTokens ?? 512) ? "length" : "stop"), metrics };
+  }, binding.runtime ?? runtimeConfig());
+}
+
+/** The canvas stays private. This method publishes only the finished result. */
+export function createDenoisingMethod<State>(
+  binding: MlxDenoisingBinding<State>, promptIds: readonly number[], options: GenerateOptions,
+): InferenceMethod<GenerateStats> {
+  const prompt = [...promptIds];
+  options = snapshotGenerationPolicy(options);
+  return nativeMethod("mlx-denoising", async (output, signal) => {
+    const generation = generateDenoising(binding, prompt, { ...options, signal });
+    for await (const item of generation)
+      if (await output.commit([item.token]) === false) break;
+    const metrics = generation.stats;
+    if (!metrics) throw new Error("denoising method did not settle its metrics");
+    return { finishReason: metrics.generatedTokens >= (options.maxTokens ?? 256) ? "length" : "stop", metrics };
+  }, binding.runtime ?? runtimeConfig());
+}

package/src/execution/mixed-iteration.ts

@@ -0,0 +1,96 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { MlxPhaseGroup, MlxMixedRange, MlxSeparateMixed } from "../contracts/mlx/graph";
+import type { MlxForwardStep, MlxForwardWork } from "../contracts/mlx/forward-work";
+
+/** Join two method-owned work producers at their forward boundary: the first
+ * forward preparation offers and the first forward decode makes run as one
+ * mixed forward (`mixed`, the composition's), each group in its own phase;
+ * every other forward runs on its own (`step.own`). Decode finishes
+ * sampling/publication/retirement before preparation can join rows. The model
+ * owns packing; this coordinator never merges caches or samples. A
+ * restored/cancelled preparation need not produce a forward at all. */
+export async function runMixedTokenIteration(options: {
+  prepare(forward: MlxForwardWork): Promise<void>;
+  decode(forward: MlxForwardWork): Promise<void>;
+  mixed(groups: readonly MlxPhaseGroup[], rows: number): MlxArray[];
+}): Promise<boolean> {
+  let notify!: () => void;
+  const offered = new Promise<void>(resolve => { notify = resolve; });
+  let pending: { group: MlxPhaseGroup; step: MlxForwardStep; resolve(hidden: MlxArray): void; reject(error: unknown): void } | undefined;
+  let captured = false;
+  let preparationError: unknown;
+  let preparationFailed = false;
+  const groupOf = (ids: MlxArray, state: MlxPhaseGroup["state"], step: MlxForwardStep): MlxPhaseGroup =>
+    ({ phase: step.phase, rows: step.rows, ids, state, capture: step.capture });
+  const preparing = options.prepare(async (ids, cache, step) => {
+    if (captured) return step.own();
+    captured = true;
+    return new Promise<MlxArray>((resolve, reject) => {
+      pending = { group: groupOf(ids, cache, step), step, resolve, reject }; notify();
+    });
+  }).catch(error => { preparationFailed = true; preparationError = error; }).finally(notify);
+  await offered;
+  if (!pending) {
+    await preparing;
+    if (preparationFailed) throw preparationError;
+    return false;
+  }
+  const held = pending;
+  let preparedHidden: MlxArray | undefined;
+  let mixed = false;
+  try {
+    await options.decode(async (ids, cache, step) => {
+      if (mixed) return step.own();
+      const outputs = options.mixed([groupOf(ids, cache, step), held.group], step.rows + held.step.rows);
+      mixed = true; preparedHidden = outputs[1]!;
+      return outputs[0]!;
+    });
+    // A final unread token can retire decode without another model call.
+    preparedHidden ??= await held.step.own();
+    held.resolve(preparedHidden); preparedHidden = undefined;
+    await preparing;
+    if (preparationFailed) throw preparationError;
+    return true;
+  } catch (error) {
+    preparedHidden?.dispose(); held.reject(error);
+    await preparing;
+    throw error;
+  }
+}
+
+export { type MlxForwardWork, type MlxPreparationWork } from "../contracts/mlx/forward-work";
+
+
+/** The scheduler routes declared packed work and preserves each separate group's geometry. */
+export function dispatchMixedGroups(groups: readonly MlxPhaseGroup[], rows: number,
+  packed: readonly MlxMixedRange[], separate?: MlxSeparateMixed): MlxArray[] {
+  const packedOperation = (ranges: readonly MlxMixedRange[], packedRows: number) => {
+    const operation = ranges.find(range => packedRows >= range.from && packedRows <= range.to
+      && rows >= (range.totalFrom ?? 1) && rows <= (range.totalTo ?? Infinity))?.operation;
+    if (!operation) throw new Error(`graph has no mixed operation for ${packedRows} packed rows (${rows} total rows)`);
+    return operation;
+  };
+  if (!separate || (packed.length > 0 && !groups.some(group => group.phase === "verify")))
+    return packedOperation(packed, rows)(groups, rows);
+  const outputs: Array<MlxArray | undefined> = Array(groups.length);
+  try {
+    const candidates = separate.packed?.length
+      ? groups.flatMap((group, index) => group.phase !== "verify" ? [index] : []) : [];
+    if (candidates.length > 1) {
+      const joined = candidates.map(index => groups[index]!);
+      const packedRows = joined.reduce((sum, group) => sum + group.rows, 0);
+      const returned = packedOperation(separate.packed!, packedRows)(joined, rows);
+      for (const [position, index] of candidates.entries()) outputs[index] = returned[position]!;
+    }
+    for (const [index, group] of groups.entries()) {
+      if (outputs[index]) continue;
+      const operation = separate[group.phase].find(range => group.rows >= range.from && group.rows <= range.to)?.operation;
+      if (!operation) throw new Error(`graph has no separate ${group.phase} operation for ${group.rows} rows`);
+      outputs[index] = operation(group);
+    }
+    return outputs as MlxArray[];
+  } catch (error) {
+    for (const output of outputs) output?.dispose();
+    throw error;
+  }
+}

package/src/execution/plan.ts

@@ -0,0 +1,65 @@
+import type { ExecutionCapabilities, ExecutionFeatures, ExecutionReason, ExecutionRequirements, ResolvedExecution } from "../contracts/portable/execution";
+
+/** Select the request method, then place it using executor capabilities. The
+ * composition's caches serve every method; nothing here reads how they store
+ * keys and values. */
+export function resolveExecution(
+  request: ExecutionRequirements, capabilities: ExecutionCapabilities,
+  features: ExecutionFeatures = { fill: false },
+): ResolvedExecution {
+  const reasons: ExecutionReason[] = [];
+  const sharedRequestExclusions = [
+    [request.hasVision && !capabilities.mediaBatch, "media-batch-unsupported"],
+    [request.hasGrammar && !capabilities.grammarBatch, "grammar-batch-unsupported"],
+  ] as const;
+  const excluded = sharedRequestExclusions.some(([exclusion]) => exclusion);
+  const sharedAdapterMethod = capabilities.sharedSpeculativeAdapters === true && !excluded &&
+    capabilities.groupedMethods?.includes("speculative") === true;
+  const grammarProposals = features.grammarJump === true && request.hasGrammar && !request.hasDraft && !request.hasVision &&
+    (!request.wantsLogprobs || capabilities.speculativeLogprobs === true) && capabilities.method === "autoregressive" &&
+    capabilities.sharedGrammarProposals === true && capabilities.groupedMethods?.includes("speculative") === true &&
+    (!request.hasAdapters || sharedAdapterMethod) && !excluded;
+  const speculative = grammarProposals || (capabilities.method === "autoregressive" && request.hasDraft &&
+    !request.hasVision && (!request.hasAdapters || sharedAdapterMethod) &&
+    (!request.wantsLogprobs || capabilities.speculativeLogprobs === true));
+  const method = speculative ? "speculative" : capabilities.method;
+  // Token-level request features act on sampled tokens. A method that does not
+  // sample tokens refuses them instead of silently discarding them.
+  const tokenMethod = capabilities.method === "autoregressive";
+  const continuousExclusions = [
+    [!(capabilities.groupedMethods ?? ["autoregressive"]).includes(method), "method-batch-unsupported"],
+    [request.hasDraft && !tokenMethod, "draft-method-unsupported"],
+    [request.wantsLogprobs && !tokenMethod, "logprobs-method-unsupported"],
+    [request.hasRepetitionPenalty && !tokenMethod, "repetition-penalty-method-unsupported"],
+    [request.hasLogitsExtras && !tokenMethod, "logits-extras-method-unsupported"],
+    [features.fill && !tokenMethod, "fill-method-unsupported"],
+    ...sharedRequestExclusions,
+  ] as const;
+  for (const [exclusion, reason] of continuousExclusions) if (exclusion) reasons.push(reason);
+  const mechanism = reasons.length ? "unsupported" : "continuous";
+  if (request.hasDraft && !speculative && tokenMethod) reasons.push("draft-incompatible-with-request");
+  // Prefix reuse restores autoregressive token caches only.
+  const promptCache = method === "autoregressive" && (!request.hasVision ||
+    (mechanism === "continuous" && request.hasPreparedPrefixIdentity === true && capabilities.mediaPrefixCache === true));
+  const sharedFill = mechanism === "continuous" && capabilities.sharedFill === true;
+  const speculativeEcho = method === "speculative" && mechanism === "continuous" &&
+    capabilities.sharedSpeculativeEcho === true;
+  const fill = features.fill && !request.hasVision && !request.hasGrammar && !request.wantsLogprobs &&
+    (speculativeEcho || (method === "autoregressive" && sharedFill && !request.hasDraft));
+  if (features.fill && !fill && tokenMethod) reasons.push("fill-incompatible-with-request");
+  const checkpoint = capabilities.checkpoints && method === "autoregressive" && !request.hasVision &&
+    capabilities.sharedCheckpoints === true && promptCache && !request.hasGrammar &&
+    !features.fill && !request.wantsLogprobs;
+  const compiledDecode = features.compiledDecode === true && capabilities.compiledDecode === true &&
+    method === "autoregressive" && !request.hasAdapters && !(sharedFill && fill);
+  if (features.compiledDecode && !compiledDecode) reasons.push("compiled-decode-unavailable-for-request");
+  const directGrammarJump = features.grammarJump === true && request.hasGrammar &&
+    !request.wantsLogprobs && method === "autoregressive" && mechanism === "continuous" &&
+    capabilities.sharedGrammarJump === true;
+  const grammarJump = grammarProposals || directGrammarJump;
+  if (features.grammarJump && request.hasGrammar && !grammarJump)
+    reasons.push("grammar-jump-incompatible-with-request");
+  return Object.freeze({ method, mechanism, promptCache, checkpoint, fill,
+    compiledDecode, grammarJump,
+    reasons: Object.freeze(reasons) });
+}

package/src/execution/prefill-cohort.ts

@@ -0,0 +1,100 @@
+import type { Cache } from "../contracts/mlx/cache";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { nextPrefillStep, prefillPhase } from "../contracts/portable/prefill";
+import { disposeResources, cleanupFailure } from "../runtime/resources";
+import type { Row } from "./batch-types";
+import type { KvMaintenance } from "../contracts/mlx/cache";
+import type { ContinuationState, RequestStateCache } from "../state/cache-service";
+import { MlxPrefillRows, type MlxPrefillForward, type MlxPrefillState } from "./prefill-rows";
+
+import { promptWorkStep } from "../generation/phase-dispatch";
+
+export interface PrefillState extends MlxPrefillState {
+  snapAt: number | null;
+  continuation?: ContinuationState;
+  closePrefill?: () => void;
+}
+
+interface PrefillHost {
+  chunkSize: number;
+  tailSplit: boolean;
+  /** Each row's state comes from the cache service, and boundary snapshots return to it. */
+  cache: RequestStateCache;
+  maintain?: KvMaintenance;
+  /** One cohort forward of `ids` [sequences, tokens] in the phase its planner named. */
+  forward(ids: MlxArray, caches: Cache[], step: MlxPrefillForward): Promise<MlxArray>;
+  /** Target projection preserves the forward batch geometry before sampling. */
+  project(hidden: MlxArray, caches: Cache[], completed: readonly PrefillState[], step: MlxPrefillForward): MlxArray;
+  /** Borrows [1,1,V] logits; takes ownership of the completed request caches. */
+  complete(state: PrefillState, logits: MlxArray): Promise<void>;
+  resume?(state: PrefillState): Promise<void>;
+  reject(row: Row, error: unknown): void;
+}
+
+/** Ordinary sampling/cache policy over the shared target preparation driver. */
+export class MlxPrefillCohort extends MlxPrefillRows<PrefillState> {
+  override get canAdmit(): boolean {
+    return super.canAdmit && this.supportsMixedWork;
+  }
+  get supportsMixedWork(): boolean { return !this.rows.some(row => row.req.promptInput); }
+  constructor(readonly host: PrefillHost) {
+    super({
+      codecs: host.cache.codecs, maintain: host.maintain,
+      open(row) {
+        let owned: { caches: Cache[]; retain?: () => void } | undefined;
+        let closeCache: (() => void) | undefined;
+        let closePrefill: (() => void) | undefined;
+        try {
+          closeCache = row.req.trace?.begin("cache.lookup_restore", { mechanism: "continuous" });
+          // The row's resumed generation, a reused prefix composed again as the
+          // graph's caches, or fresh state: each converts inside its own appends.
+          const acquired = owned = host.cache.acquire(row.req, row.cacheNamespace ?? "");
+          if (acquired.continuation)
+            return { row, solo: acquired.caches, pos: row.req.promptIds.length,
+              snapAt: null, continuation: acquired.continuation };
+          closeCache?.(); closeCache = undefined;
+          row.cachedTokens = acquired.reusedTokens;
+          const length = row.req.promptIds.length;
+          const boundary = Math.min(row.req.snapshotAt ?? length, length - 1);
+          const snapAt = host.cache.keeps(row.req) && boundary >= 256 && boundary > acquired.reusedTokens ? boundary : null;
+          closePrefill = row.req.trace?.begin("prefill.total", {
+            mechanism: "continuous", promptTokens: row.promptTokens, cachedTokens: row.cachedTokens,
+          });
+          return { row, solo: acquired.caches, pos: acquired.reusedTokens, retain: acquired.retain, snapAt, closePrefill };
+        } catch (error) {
+          return cleanupFailure(error, () => disposeResources([...(owned?.caches ?? []),
+            { dispose: () => owned?.retain?.() }, { dispose: () => closePrefill?.() }]));
+        } finally { closeCache?.(); }
+      },
+      ready: state => !!state.continuation,
+      plan(state) {
+        if (state.row.req.promptInput) return { kind: "final", start: state.pos,
+          end: state.row.promptTokens, snapshot: false, batchYield: false, atomic: true,
+          phase: prefillPhase(state.pos, state.row.promptTokens, false) };
+        return nextPrefillStep({ length: state.row.promptTokens, position: state.pos,
+        chunkSize: state.row.req.prefillChunkSize ?? host.chunkSize,
+        tailSplit: host.tailSplit, snapshotAt: state.snapAt }); },
+      forward: (ids, caches, states, step, work) => states[0]?.row.req.promptInput
+        ? Promise.resolve(states[0].row.req.promptInput.prefill(ids, caches, states[0].pos,
+          step.phase === "decode" ? "decode" : "tail", (step.tokens)))
+        : work ? work(ids, caches, promptWorkStep(step, null, () => host.forward(ids, caches, step))) : host.forward(ids, caches, step),
+      // A media prompt input returns its prompt's last position: one row per sequence to project.
+      project: (hidden, caches, completed, step) => host.project(hidden, caches, completed,
+        completed.some(state => state.row.req.promptInput) ? { ...step, tokens: 1 } : step),
+      complete: (state, logits) => state.continuation ? host.resume!(state) : host.complete(state, logits!), reject: host.reject.bind(host),
+      checkpoint(state, capture) {
+        let snapshot: Cache[] | undefined;
+        try {
+          snapshot = capture();
+          host.cache.publish(state.row.req, state.row.cacheNamespace ?? "", state.row.req.promptIds.slice(0, state.pos), { caches: snapshot });
+          snapshot = undefined;
+        } catch (error) {
+          if (snapshot) disposeResources(snapshot);
+          console.warn(`batch-lane boundary snapshot skipped: ${(error as Error).message}`);
+        }
+        state.snapAt = null;
+      },
+      close(state) { const close = state.closePrefill; state.closePrefill = undefined; close?.(); },
+    });
+  }
+}

package/src/execution/prefill-rows.ts

@@ -0,0 +1,263 @@
+import type { BatchableCache, Cache } from "../contracts/mlx/cache";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { clearCache } from "@mlx-bun/mlx/ffi";
+import type { PrefillStep } from "../contracts/portable/prefill";
+import type { MlxPromptStep } from "../generation/phase-dispatch";
+import { cloneKvCaches } from "../state/persistence";
+import { type CacheCodecProvider } from "../state/persistence-types";
+import { disposeResources, cleanupFailure } from "../runtime/resources";
+import type { MlxGroupPreparation, Row } from "./batch-types";
+import { evalCacheState } from "../generation/prefill";
+import { prefillCacheLayout } from "../state/layout";
+import { MlxStateRows } from "../state/rows";
+import type { P2RTracePhase, P2RTraceAttributes } from "../runtime/trace";
+import type { MlxForwardWork, MlxPreparationWork } from "../contracts/mlx/forward-work";
+import type { KvMaintenance } from "../contracts/mlx/cache";
+
+// Shared work is recorded on every participating request. workId identifies
+// duplicates when constructing a process timeline; row spans are not additive.
+let nextWorkId = 0;
+function traceRows(states: readonly MlxPrefillState[], phase: P2RTracePhase, attributes: P2RTraceAttributes) {
+  if (!states.some(state => state.row.req.trace)) return null;
+  const closes = states.map(state => state.row.req.trace?.begin(phase, attributes));
+  return { [Symbol.dispose]() { for (const close of closes) close?.(); } };
+}
+
+export interface MlxPrefillStep extends PrefillStep {
+  /** A non-causal input cannot be split at a scheduler token budget.
+   * Its preparation must decline joins until this step completes. */
+  readonly atomic?: boolean;
+  /** 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;
+}
+
+/** One cohort forward: the rows' common phase and their positions per row. */
+export type MlxPrefillForward = MlxPromptStep;
+
+/** 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> {
+  /** The snapshot form of the composition's caches (the cache service's), for a single row's boundary snapshot. */
+  codecs?: CacheCodecProvider;
+  maintain?: KvMaintenance;
+  open(row: Row): State;
+  plan(state: State): MlxPrefillStep;
+  /** Already prepared state can complete without any target forward. */
+  ready?(state: State): boolean;
+  /** Run one cohort forward of `ids` [sequences, tokens] in `step.phase`. */
+  forward(ids: MlxArray, caches: Cache[], states: readonly State[], step: MlxPrefillForward,
+    work?: MlxForwardWork): Promise<MlxArray>;
+  /** Project the cohort forward `step` produced (its rows: sequences times tokens). */
+  project(hidden: MlxArray, caches: Cache[], completed: readonly State[], step: MlxPrefillForward): 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(workLimit?: MlxPreparationWork): Promise<boolean> {
+    let remaining = workLimit?.maxTokens ?? Infinity;
+    // 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; }
+    }
+    {
+      using span = traceRows(this.#states, "prefill.row_sync", { workId: ++nextWorkId, batchSize: this.#states.length });
+      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.some(state => state.planned!.atomic) ? Infinity
+        : Math.max(1, Math.floor(remaining / this.#states.length)),
+        ...this.#states.map(state => state.planned!.end - state.pos));
+      // A planned step stays the unit of maintenance: a row whose step is longer
+      // than this forward (a sibling's step is shorter, or the budget is) runs
+      // `count` of its positions now and is maintained at the step's planned end,
+      // so its caches convert where its own plan put the boundary, whatever rows
+      // it shares the forward with. Only every row's whole planned chunk is a
+      // chunk forward; any other piece of more than one position is a tail.
+      const whole = this.#states.every(state => state.planned!.phase === "chunk" &&
+        state.pos === state.planned!.start && state.pos + count === state.planned!.end);
+      const step: MlxPrefillForward = { sequences: this.#states.length, tokens: count,
+        phase: count === 1 ? "decode" : whole ? "chunk" : "tail" };
+      const work = { workId: ++nextWorkId, batchSize: this.#states.length, tokensPerRow: count };
+      const closes = this.#states.map(state => state.row.req.trace?.begin("prefill.chunk", {
+        mechanism: "continuous", startToken: state.pos, ...work,
+      }));
+      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;
+        let forwarded: MlxArray | null;
+        {
+          using span = traceRows(this.#states, "prefill.forward", work);
+          forwarded = ids ? await this.operations.forward(ids, caches, this.#states, step, workLimit?.forward) : null;
+        }
+        using hidden = forwarded;
+        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)) {
+          {
+            using span = traceRows(this.#states, "prefill.evaluate", work);
+            evalCacheState(caches);
+          }
+          using span = traceRows(this.#states, "prefill.kv_maintenance", work);
+          // Cohort rows convert at their planned boundaries; a single row's
+          // composed caches converted inside their own appends.
+          if (this.#stateRows) for (const cache of caches) cache.prefillMaintenance?.commitPrefill(maintained);
+          // A scheduler budget can split one planned chunk into many pieces.
+          // Resolve state at each yield, but retain allocator reuse until the
+          // method's maintenance boundary instead of purging it every piece.
+          if (maintained.length) clearCache();
+        }
+        if (ids && this.operations.afterForward) {
+          using span = traceRows(this.#states, "prefill.companion", work);
+          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) {
+            using span = traceRows([state], "prefill.checkpoint", { ...work, row });
+            this.operations.checkpoint(state, () =>
+              this.#stateRows ? this.#stateRows.extractRow(row) : cloneKvCaches(state.solo, this.operations.codecs), row);
+          }
+          if (!finals.includes(row)) state.planned = undefined;
+        }
+        let projected: MlxArray | null = null;
+        if (hidden && finals.length) {
+          using span = traceRows(this.#states, "prefill.project", work);
+          projected = this.operations.project(hidden, caches, finals.map(row => this.#states[row]!), step);
+        }
+        using logits = projected;
+        for (const row of finals) {
+          const state = this.#states[row]!;
+          using span = traceRows([state], "prefill.complete", { ...work, 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; }
+        remaining -= count * work.batchSize;
+        if (batchYield || finals.length || remaining < this.#states.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/execution/request-policy.ts

@@ -0,0 +1,19 @@
+import type { GenerateOptions } from "../generation/index";
+
+/** Copy policy data without cloning native caches, controllers or callbacks. */
+export function snapshotGenerationPolicy<T extends GenerateOptions & { stopSequences?: readonly string[] }>(options: T): T {
+  const snapshot = { ...options };
+  const freeze = (value: unknown): void => {
+    if (!value || typeof value !== "object") return;
+    for (const item of Object.values(value)) freeze(item);
+    Object.freeze(value);
+  };
+  for (const key of ["stopSequences", "eosTokenIds", "adapters", "xtcSpecialTokens", "logitBias",
+    "hlg", "curve", "decodePolicy"] as const) {
+    if (snapshot[key] === undefined) continue;
+    const value = structuredClone(snapshot[key]);
+    freeze(value);
+    Object.assign(snapshot, { [key]: value });
+  }
+  return Object.freeze(snapshot);
+}