@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/execution/batch-types.ts

@@ -0,0 +1,225 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { type MlxForwardWork,type MlxPreparationWork } from "../contracts/mlx/forward-work";
+import type { ExecutionContext } from "../contracts/portable/scheduling";
+import { type RuntimeConfig } from "../runtime/config";
+import type { PromptResponseTrace } from "../runtime/trace";
+import type { GrammarController } from "../sampling/grammar";
+import type { KvMaintenance } from "../contracts/mlx/cache";
+import type { RequestStateCache } from "../state/cache-service";
+import type { OrdinaryContinuation } from "./continuation-types";
+
+// NOTE — oMLX-style "adaptive burst decode" (several steps per event-loop
+// yield, omlx/engine_core.py _step_burst) was ported here 2026-07-02 and
+// REFUTED by measurement: on this runtime it REGRESSED cpm5 B=4 aggregate
+// 345→289 tok/s, batch-lane B=1 149→121, and TTFT ~+100 ms (first-token SSE
+// flush waits out the burst budget). Their win exists because each Python
+// step hand-off ping-pongs the GIL with asyncio/uvicorn (~1 ms/token); Bun's
+// setImmediate hop costs microseconds, so bursting here only delays socket
+// flushes. Don't re-add without new evidence — the per-yield step below is
+// the measured optimum (`02d723a:docs/design/batching.md` §2).
+
+/** A token sampler for one row: (logits [1,V], step) → token array [1] on
+ *  device. Greedy is `(l) => ops.argmaxAxis(l, -1)`; richer closures fold in
+ *  temperature / top-p / logits processors + history (built per request from
+ *  its sampler options, exactly like generate()'s sampleStep). Called once per
+ *  step 0..maxTokens-1 for this row (never beyond maxTokens). */
+export type RowSampler = (logits1V: MlxArray, step: number) => MlxArray;
+
+/** A method prepares and advances rows; the shared executor owns queueing. */
+export interface MlxGroupPreparation {
+  readonly rows: readonly Row[];
+  readonly tokenWeight?: number;
+  readonly canAdmit?: boolean;
+  readonly supportsMixedWork?: boolean;
+  admit?(row: Row): void;
+  advance(work?: MlxPreparationWork): Promise<boolean>;
+  dispose(): void;
+}
+
+export interface MlxGroupMethodHost {
+  readonly rows: readonly Row[];
+  readonly runtime: RuntimeConfig;
+  readonly prefillChunkSize: number;
+  /** The composition's `prefillTailSplit`: a prompt's last token prefills alone. */
+  readonly prefillTailSplit: boolean;
+  /** The composition's cache service: every row's state comes from it and returns to it. */
+  readonly cache: RequestStateCache;
+  /** The row-layout maintenance of the composition's KV scheme (where a row's
+   * composed caches join row layouts); absent under bf16. */
+  readonly kvMaintenance: KvMaintenance | undefined;
+  join(row: Row): void;
+  filterRows(keep: number[]): void;
+  publish(row: Row, token: number, logprobs?: import("../generation/index").TokenLogprobs): Promise<void | boolean>;
+  finish(row: Row, reason: "stop" | "length"): void;
+}
+
+export interface MlxGroupedMethod {
+  /** Maximum target tokens in the next iteration, including candidates. */
+  readonly runningTokens?: number;
+  prepare(row: Row): MlxGroupPreparation;
+  advance(work?: MlxForwardWork): Promise<void>;
+  filterRows(keep: readonly number[], discard: boolean): void;
+  dispose(): void;
+}
+
+export interface MlxGroupMethodRequest {
+  /** Equal keys declare compatible method state, independent of queue size. */
+  readonly key: string;
+  readonly data: unknown;
+  open(host: MlxGroupMethodHost): MlxGroupedMethod;
+}
+
+export type BatchRequest = BatchRequestFields & (
+  | { method: MlxGroupMethodRequest; sample?: RowSampler }
+  | { method?: undefined; sample: RowSampler }
+);
+
+export interface BatchRequestFields {
+  /** Borrowed prepared input; its owner retains it through preparation. */
+  promptInput?: import("../contracts/mlx/media").MlxPromptInput;
+  cacheSessionId?: string;
+  /** Optional ordinary continuation: its checkpoint key and sampler recovery;
+   * the cache service resumes and checkpoints it. */
+  continuation?: OrdinaryContinuation;
+  /** False for a request whose prefix has no identity (prepared media without
+   * one): it takes fresh state and its state is not kept. Rows that differ here
+   * do not share a group. */
+  reusesPrefix?: boolean;
+  context?: ExecutionContext;
+  /** Resolve loaded-state identity under the execution lease, after any
+   * queued adapter replacement has finished. */
+  cacheNamespace?: string | (() => string);
+  /** Resolved replay permission; standalone group callers retain auto mode. */
+  compiledDecode?: boolean;
+  promptIds: number[];
+  /** Request policy overrides the execution group's captured prefill default. */
+  prefillChunkSize?: number;
+  maxTokens: number;
+  eosTokenIds: number[];
+  /** Called per emitted (non-EOS) token, in order. Returning `false` halts this
+   *  row (a decoded-text stop sequence fired) — matches generate()'s onToken
+   *  contract. EOS terminates the row WITHOUT an onToken call. Throwing evicts
+   *  THIS row only (its submit promise rejects; siblings continue). May be
+   *  async; keep it cheap — it runs inline in the step loop. */
+  onToken: (token: number, logprobs?: import("../generation/index").TokenLogprobs) => void | boolean | Promise<void | boolean>;
+  /** Diagnostic-only request-local trace. */
+  trace?: PromptResponseTrace;
+  /** Closes the gateway's queue/admission span when this row leaves pending. */
+  onAdmitted?: () => void;
+  /** Client/request lifetime. An aborted pending row is removed before
+   *  admission; a row in prefill or decode is evicted at the next safe
+   *  scheduler boundary. */
+  signal?: AbortSignal;
+  /** The row's sampler is PLAIN GREEDY (temperature 0, no curve) with no
+   *  logits processors and no grammar — set by the gateway when true. Lets
+   *  the scheduler take the vectorized sampling fast path (ONE
+   *  log-softmax+argmax over [B,V] instead of B slice/sample/concat graphs)
+   *  when EVERY live row qualifies. Numerically identical per row
+   *  (row-independent ops, same per-row shapes); the per-row closure path
+   *  is the fallback and the composition's `vectorGreedySampling: false`. */
+  plainGreedy?: boolean;
+  /** Grammar controller for this row (B1: per-row matchers). When set, the
+   *  scheduler drives it: `accept()` after each emitted token (fires the async
+   *  bitmask fill), `await ready()` before the row's next sample (the sample
+   *  closure applies `applyMask`). Termination (grammar satisfied) is an
+   *  additional per-row stop source. The gateway OWNS disposal (finally around
+   *  submit()); the scheduler uses but never owns. Null on the degrade path. */
+  grammar?: GrammarController;
+  /** Stable cache boundary from the server's template probe (the serial
+   *  lane's options.snapshotAt): the prompt prefix the NEXT turn's render
+   *  preserves. The scheduler snapshots a trim-free prompt-cache entry at
+   *  min(this, promptIds.length-1) during the solo prefill — the oracle
+   *  invariant (mlx-lm insert_segments) that gives wrapped-ring models a
+   *  reuse path. Absent = still capped at promptIds.length-1. */
+  snapshotAt?: number;
+}
+
+export interface BatchStats {
+  spec?: import("../generation/index").GenerateStats["spec"];
+  fill?: import("../generation/index").GenerateStats["fill"];
+  promptTokens: number;
+  generatedTokens: number;
+  /** Prompt tokens served from the prompt cache: a joiner's prefill starts
+   *  from the longest usable cached prefix the cache service acquires. 0 on a
+   *  cold prefill. */
+  cachedTokens: number;
+  finishReason: "stop" | "length";
+  /** Wall-clock from admission (leaving the pending queue) to the first
+   *  emitted token — the row's prefill span on the batch lane. */
+  prefillMs: number;
+  /** Wall-clock from the first emitted token to finish. */
+  decodeMs: number;
+}
+
+export interface Row {
+  spec?: import("../generation/index").GenerateStats["spec"];
+  fill?: import("../generation/index").GenerateStats["fill"];
+  req: BatchRequest;
+  resolve: (s: BatchStats) => void;
+  reject: (e: unknown) => void;
+  current: number; // last emitted token, fed at the next step (pipeline cold)
+  generated: number; // tokens emitted so far (incl. a terminating EOS)
+  sampled: number; // sample() calls so far (leads `generated` by 1 in-pipeline)
+  promptTokens: number;
+  /** Prompt tokens restored from the prompt cache at admission (stats). */
+  cachedTokens: number;
+  /** performance.now() marks for BatchStats timing (0 = not reached). */
+  admittedAt: number;
+  firstTokenAt: number;
+  /** A method that computes its whole result before publishing reports that
+   *  work span; finish() then reports it as decode time, excluding admission
+   *  wait and output delivery. Denoising sets it. */
+  decodeSpan?: { start: number; end: number };
+  closeTokenZero?: () => void;
+  cacheNamespace?: string;
+  /** Generated tokens whose KV actually entered the cache (fed as a step
+   *  input) — serial generate()'s `forwarded` list. promptIds+fed is the
+   *  exact token coverage of the row's caches, the put() entry key.
+   *  PER-ROW EXACT for every row: pushes are gated on the pending
+   *  register's per-slot real/placeholder flags (#pendingReal), so a
+   *  doomed slot's placeholder is never recorded — and if one ever DID
+   *  feed a forward, `fedTainted` flips instead (see #step's read). */
+  fed: number[];
+  /** A placeholder slot value fed this row's KV (junk token in the cache):
+   *  the coverage key is no longer derivable — extraction refuses.
+   *  Structurally unreachable today (a doomed row always evicts at the
+   *  emit that follows its placeholder, filtering the slot before the
+   *  next forward; cold steps never contain doomed rows because a flush
+   *  emits every in-flight token) — this flag makes `fed` provably exact
+   *  instead of resting on that timing chain. */
+  fedTainted: boolean;
+  /** Ever merged into a multi-row batch (or poisoned by a reject): its KV
+   *  is interleaved in batched buffers — a finish put() must EXTRACT the
+   *  row (#extractRowCaches) rather than adopt the inners. */
+  merged: boolean;
+}
+
+/** Held while the batch is active; the serial fallback acquires the same lock. */
+export interface ExclusiveLock {
+  acquire(): Promise<() => void>;
+}
+
+export interface MlxBatchExecutionGroupOptions {
+  /** Scheduling work budget for extending an existing preparation cohort. */
+  prefillBatchTokenLimit?: number;
+  runtime?: RuntimeConfig;
+  maxQueued?: number;
+  /** Max rows in the running batch (mlx-lm `--decode-concurrency`). */
+  maxBatch: number;
+  lock?: ExclusiveLock;
+  /** Drain signal: while true, no NEW rows are admitted (running rows finish,
+   *  the lock is released) so a waiting serial-lane request can run — mlx-lm's
+   *  drain_batch. Pair with kick() when it flips back to false. */
+  admissionHeld?: () => boolean;
+  /** Joiner prefill chunk length (default 2048, the serial loop's constant). */
+  prefillChunkSize?: number;
+  /** The composition's cache service (`state/cache-service`): admission under
+   *  its KV budget (a joiner that would push the admitted rows' projected KV
+   *  over `--kv-budget` WAITS in the queue, FIFO and head-of-line, until rows
+   *  finish; one that can never fit is rejected), each row's state at
+   *  admission (a reused prefix, a resumed generation, or fresh state), and
+   *  the state rows offer back when they finish. Runs under the gateway mutex
+   *  domain. Default: a service over the graph's composition that reuses and
+   *  keeps nothing and admits without a budget. */
+  cache?: RequestStateCache;
+}

package/src/execution/cancellation.ts

@@ -0,0 +1,38 @@
+import type { CancelReason, Cancellation } from "../contracts/portable/generation";
+
+export class GenerationCancelled extends Error {
+  constructor(readonly reason: CancelReason) {
+    super(`generation cancelled: ${reason}`);
+    this.name = "GenerationCancelled";
+  }
+}
+
+export function throwIfCancelled(cancellation: Cancellation): void {
+  if (cancellation.reason !== undefined) throw new GenerationCancelled(cancellation.reason);
+}
+
+export class CancellationSource implements Cancellation {
+  #reason: CancelReason | undefined;
+  #listeners = new Set<(reason: CancelReason) => void>();
+  get reason(): CancelReason | undefined { return this.#reason; }
+
+  subscribe(listener: (reason: CancelReason) => void): () => void {
+    if (this.#reason !== undefined) {
+      listener(this.#reason);
+      return () => {};
+    }
+    this.#listeners.add(listener);
+    return () => { this.#listeners.delete(listener); };
+  }
+
+  cancel(reason: CancelReason): void {
+    if (this.#reason !== undefined) return;
+    this.#reason = reason;
+    const listeners = [...this.#listeners];
+    this.#listeners.clear();
+    // A faulty observer cannot prevent the remaining owners from cancelling.
+    for (const listener of listeners) {
+      try { listener(reason); } catch { /* Observers do not own settlement. */ }
+    }
+  }
+}

package/src/execution/continuation-request.ts

@@ -0,0 +1,45 @@
+import { snapshotGenerationPolicy } from "./request-policy";
+import type { GenerateOptions } from "../generation/index";
+import type { ResolvedExecution } from "../contracts/portable/execution";
+import { generationCheckpointKey } from "../generation/checkpoint-identity";
+import { makeStepSampler } from "../sampling/index";
+import { createRowSampling } from "./row-sampling";
+import { type OrdinaryContinuation } from "./continuation-types";
+
+/** Optional policy and sampler recovery for the existing ordinary driver. The
+ * cache service resumes and checkpoints the generation under the key this
+ * request names; KV layout and delayed conversion stay with the caches. */
+export function createOrdinaryContinuationRequest(input: {
+  options: GenerateOptions; prompt: number[]; execution: ResolvedExecution;
+  /** Checkpoint every this many generated tokens (the cache service's `checkpointEveryTokens`). */
+  interval: number;
+  onToken: Parameters<typeof createRowSampling>[1];
+}) {
+  const options: GenerateOptions = snapshotGenerationPolicy(input.options);
+  if (!Number.isSafeInteger(input.interval) || input.interval < 1) throw new Error("continuation interval must be positive");
+  if (options.grammar || options.logprobs || options.topLogprobs || options.fill ||
+      input.execution.method !== "autoregressive")
+    throw new Error("ordinary continuation requires replayable ordinary sampling");
+  const make = (seed = options.seed, history: readonly number[] = input.prompt) => createRowSampling(makeStepSampler(
+    { ...options, seed }, { tokenRepresentation: "device", grammarWait: "external", historyUpdate: "after-sample",
+      initialHistory: [...history] }), input.onToken);
+  let sampling = make();
+  const continuation: OrdinaryContinuation = {
+    interval: input.interval,
+    key: (namespace, identity) => generationCheckpointKey(input.prompt, options, namespace, input.execution, identity),
+    seed: options.seed ?? 0,
+    seedWasExplicit: options.seedWasExplicit === true,
+    resumeSampling(state) {
+      const next = make(state.seed, [...state.cacheTokens, state.pendingToken]);
+      const previous = sampling; sampling = next; previous.dispose();
+    },
+  };
+  return {
+    continuation,
+    sample: ((...args) => sampling.sample(...args)) as ReturnType<typeof createRowSampling>["sample"],
+    // History-free greedy sampling needs no recovery work in the fast path.
+    plainGreedy: sampling.plainGreedy,
+    onToken: input.onToken,
+    dispose: () => { sampling.dispose(); },
+  };
+}

package/src/execution/continuation-types.ts

@@ -0,0 +1,10 @@
+import type { ContinuationRequest, ContinuationState } from "../state/cache-service";
+
+/** A replayable ordinary request: what the cache service checkpoints it under
+ * (`ContinuationRequest`), how often the scheduler offers a checkpoint, and
+ * the sampler recovery a resumed generation runs. No store knowledge enters
+ * the driver; the cache service resumes and checkpoints it. */
+export interface OrdinaryContinuation extends ContinuationRequest {
+  readonly interval: number;
+  resumeSampling(state: Omit<ContinuationState, "caches">): void;
+}

package/src/execution/coordinator.ts

@@ -0,0 +1,70 @@
+import type { Cancellation } from "../contracts/portable/generation";
+import type { DisposableResource } from "../contracts/portable/resources";
+import { AdmissionRejected } from "./admission";
+import { throwIfCancelled } from "./cancellation";
+
+type Mode = "shared" | "exclusive";
+interface Waiter {
+  mode: Mode;
+  resolve(lease: DisposableResource): void;
+  reject(error: unknown): void;
+  unsubscribe(): void;
+}
+
+/** Shared worker activity; exclusive managed jobs. FIFO writers cannot starve
+ * behind newly arriving inference. This coordinates execution, not memory size. */
+export class ExecutionCoordinator {
+  #readers = 0;
+  #writer = false;
+  #closed = false;
+  #queue: Waiter[] = [];
+  constructor(readonly maxQueued = 64) {
+    if (!Number.isSafeInteger(maxQueued) || maxQueued < 0) throw new Error("invalid coordination queue bound");
+  }
+  acquire(mode: Mode, cancellation?: Cancellation): Promise<DisposableResource> {
+    if (this.#closed) return Promise.reject(new Error("execution coordinator is closed"));
+    try { if (cancellation) throwIfCancelled(cancellation); }
+    catch (error) { return Promise.reject(error); }
+    if (!this.#queue.length && !this.#writer && (mode === "shared" || this.#readers === 0))
+      return Promise.resolve(this.#lease(mode));
+    if (this.#queue.length >= this.maxQueued) return Promise.reject(new AdmissionRejected());
+    return new Promise((resolve, reject) => {
+      const waiter: Waiter = { mode, resolve, reject, unsubscribe() {} };
+      this.#queue.push(waiter);
+      waiter.unsubscribe = cancellation?.subscribe((reason) => {
+        const index = this.#queue.indexOf(waiter);
+        if (index < 0) return;
+        this.#queue.splice(index, 1);
+        waiter.unsubscribe();
+        reject(new Error(`execution admission cancelled: ${reason}`));
+        this.#drain();
+      }) ?? (() => {});
+      if (!this.#queue.includes(waiter)) waiter.unsubscribe();
+    });
+  }
+  #lease(mode: Mode): DisposableResource {
+    if (mode === "shared") this.#readers++;
+    else this.#writer = true;
+    let disposed = false;
+    return { dispose: () => {
+      if (disposed) return;
+      disposed = true;
+      if (mode === "shared") this.#readers--;
+      else this.#writer = false;
+      this.#drain();
+    } };
+  }
+  #drain(): void {
+    while (!this.#writer && this.#queue.length) {
+      const next = this.#queue[0]!;
+      if (next.mode === "exclusive" && this.#readers > 0) return;
+      this.#queue.shift(); next.unsubscribe(); next.resolve(this.#lease(next.mode));
+    }
+  }
+  close(): void {
+    this.#closed = true;
+    for (const next of this.#queue.splice(0)) {
+      next.unsubscribe(); next.reject(new Error("execution coordinator is closed"));
+    }
+  }
+}

package/src/execution/denoising-group.ts

@@ -0,0 +1,180 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { GenerateOptions } from "../generation/index";
+import { denoisingRequestOptions } from "../generation/denoising";
+import { openDenoisingRun, type DenoisingRun, type DiffusionGenOptions, type DiffusionGenResult } from "../generation/diffusion";
+import type { MlxDenoisingBinding } from "../generation/bindings/denoising";
+import { cleanupFailure, disposeResources } from "../runtime/resources";
+import { snapshotGenerationPolicy } from "./request-policy";
+import type { MlxGroupedMethod, MlxGroupMethodHost, MlxGroupMethodRequest, MlxGroupPreparation, Row } from "./batch-types";
+
+interface RequestState {
+  readonly run: DenoisingRun;
+  readonly adapters: string[];
+  readonly maxTokens: number;
+  result?: DiffusionGenResult;
+  /** performance.now() when the final unit returned the result. */
+  computedAt?: number;
+}
+
+/** Interleaved shared denoising: each row keeps its own canvas, feedback,
+ * encoder state and key sequence, and advances by one bounded unit at a time
+ * in round-robin order. Graph calls stay single-row; canvases are not stacked. */
+export function bindDenoisingGroupRequests<State>(
+  binding: MlxDenoisingBinding<State>,
+  policy: (options: GenerateOptions) => DiffusionGenOptions = denoisingRequestOptions,
+) {
+  return (input: GenerateOptions): MlxGroupMethodRequest => {
+    const options = snapshotGenerationPolicy(input);
+    return {
+      key: JSON.stringify(["denoising", binding.graph.descriptor.id]),
+      data: options,
+      open: host => new DenoisingGroup(host, binding, policy),
+    };
+  };
+}
+
+class DenoisingGroup<State> implements MlxGroupedMethod {
+  readonly #requests = new Map<Row, RequestState>();
+  /** One immutable base-weight table shared by this group's rows. Rows borrow
+   * it; the group releases it after the last row's run has closed. */
+  #table: MlxArray | null = null;
+  #tableUsers = 0;
+  #next = 0;
+
+  constructor(readonly host: MlxGroupMethodHost, readonly binding: MlxDenoisingBinding<State>,
+    readonly policy: (options: GenerateOptions) => DiffusionGenOptions) {}
+
+  prepare(first: Row): MlxGroupPreparation {
+    let row: Row | null = first;
+    const preparation: MlxGroupPreparation = {
+      get rows() { return row ? [row] : []; },
+      canAdmit: false,
+      supportsMixedWork: false,
+      get tokenWeight() { return row?.promptTokens ?? 0; },
+      // The whole prompt prefill and first canvas draw form one unit.
+      advance: async () => {
+        const current = row!;
+        if (current.req.signal?.aborted) throw current.req.signal.reason;
+        this.#advanceRun(this.#requests.get(current)!);
+        row = null;
+        this.host.join(current);
+        return true;
+      },
+      dispose: () => {
+        if (row && !this.host.rows.includes(row)) this.#close(row);
+        row = null;
+      },
+    };
+    try { this.#open(first); return preparation; }
+    catch (error) { return cleanupFailure(error, () => preparation.dispose()); }
+  }
+
+  async advance(): Promise<void> {
+    const rows = this.host.rows;
+    const live = rows.flatMap((row, index) => row.req.signal?.aborted ? [] : [index]);
+    if (live.length !== rows.length) {
+      for (const row of rows) if (row.req.signal?.aborted) row.reject(row.req.signal.reason);
+      this.host.filterRows(live);
+    }
+    if (!this.host.rows.length) return;
+    const index = this.#next % this.host.rows.length;
+    const row = this.host.rows[index]!;
+    const state = this.#requests.get(row)!;
+    try {
+      if (!state.result) this.#advanceRun(state);
+    } catch (error) {
+      this.#retire(index);
+      row.reject(error);
+      return;
+    }
+    this.#next = index + 1;
+    if (state.result) await this.#publish(row, index, state);
+  }
+
+  filterRows(keep: readonly number[], _discard: boolean): void {
+    const rows = this.host.rows;
+    const retained = new Set(keep.map(index => rows[index]!));
+    for (const row of rows) if (!retained.has(row)) this.#close(row);
+    this.#next = keep.filter(index => index < this.#next).length;
+  }
+
+  dispose(): void {
+    disposeResources([...this.#requests.keys()].map(row => ({ dispose: () => this.#close(row) })));
+  }
+
+  /** Image pixels stay borrowed from the request owner. Only the row's first
+   * unit reads them; the owner releases them after the row has settled. */
+  #open(row: Row): void {
+    const options = row.req.method!.data as GenerateOptions;
+    const denoise = this.policy(options);
+    const table = this.#acquireTable();
+    let run: DenoisingRun;
+    try { run = openDenoisingRun(this.binding.graph, row.req.promptIds, { ...denoise, embedWeight: table }); }
+    catch (error) { return cleanupFailure(error, () => this.#releaseTable()); }
+    this.#requests.set(row, { run, maxTokens: denoise.maxTokens,
+      adapters: options.adapters?.length ? [...options.adapters] : [] });
+  }
+
+  #advanceRun(state: RequestState): void {
+    const result = this.#unit(state);
+    if (!result) return;
+    state.result = result;
+    state.computedAt = performance.now();
+  }
+
+  /** One synchronous unit with only this row's adapters active. Units never
+   * await, so no other row can observe them. */
+  #unit(state: RequestState): DiffusionGenResult | undefined {
+    const adapters = this.binding.adapters;
+    if (!adapters) return state.run.advance();
+    const previous = adapters.active;
+    adapters.active = state.adapters;
+    try { return state.run.advance(); }
+    finally { adapters.active = previous; }
+  }
+
+  /** Publish the finished result only; no canvas token leaves before it. */
+  async #publish(row: Row, index: number, state: RequestState): Promise<void> {
+    try {
+      for (const token of state.result!.tokens) {
+        row.generated++; row.sampled++;
+        if (await this.host.publish(row, token) === false) break;
+      }
+    } catch (error) {
+      this.#retire(index);
+      row.reject(error);
+      return;
+    }
+    // The whole result exists before its first token is published. Report the
+    // work from admission to the final unit as decode time, as the serial
+    // method did, excluding output delivery.
+    row.decodeSpan = { start: row.admittedAt, end: state.computedAt! };
+    this.host.finish(row, row.generated >= state.maxTokens ? "length" : "stop");
+    this.#retire(index);
+  }
+
+  #retire(index: number): void {
+    this.host.filterRows(this.host.rows.flatMap((_row, position) => position === index ? [] : [position]));
+  }
+
+  /** The run's cleanup completes before the shared table can be released. */
+  #close(row: Row): void {
+    const state = this.#requests.get(row);
+    if (!state) return;
+    this.#requests.delete(row);
+    disposeResources([{ dispose: () => state.run.close() }, { dispose: () => this.#releaseTable() }]);
+  }
+
+  #acquireTable(): MlxArray {
+    this.#table ??= this.binding.graph.dequantEmbedWeight();
+    this.#tableUsers++;
+    return this.#table;
+  }
+
+  #releaseTable(): void {
+    if (--this.#tableUsers > 0) return;
+    const table = this.#table;
+    this.#table = null;
+    table?.dispose();
+  }
+}

package/src/execution/engine.ts

@@ -0,0 +1,56 @@
+import type { ExecutionPlanner, GenerationSession, InferenceEngine, RunControl, Timer } from "../contracts/portable/generation";
+import { EngineSession, type SessionResources } from "./session";
+
+export interface EngineOptions {
+  readonly timer: Timer;
+  /** Deadline for an unread stream or a producer blocked by an idle consumer. */
+  readonly idleTimeoutMs?: number;
+  readonly maxQueuedEvents?: number;
+  readonly maxQueuedTokens?: number;
+  readonly maxTopLogprobs?: number;
+  /** Aggregate output buffers for active collecting sessions, in tokens (4 bytes each).
+   * Once terminal, output storage belongs to the caller. */
+  readonly maxCollectTokens?: number;
+}
+
+export function createInferenceEngine<Request, Metrics>(
+  planner: ExecutionPlanner<Request, Metrics>, options: EngineOptions,
+): InferenceEngine<Request, Metrics> {
+  const idleTimeoutMs = options.idleTimeoutMs ?? 30_000;
+  const maxQueuedEvents = options.maxQueuedEvents ?? 32;
+  const maxQueuedTokens = options.maxQueuedTokens ?? 256;
+  const maxCollectTokens = options.maxCollectTokens ?? 65_536;
+  const maxTopLogprobs = options.maxTopLogprobs ?? 20;
+  for (const [name, value] of Object.entries({ idleTimeoutMs, maxQueuedEvents, maxQueuedTokens, maxCollectTokens, maxTopLogprobs }))
+    if (!Number.isSafeInteger(value) || value <= 0) throw new Error(`${name} must be a positive safe integer`);
+  const sessions = new Set<GenerationSession<Metrics>>();
+  let closed = false;
+  let closing: Promise<void> | undefined;
+  let nextId = 0;
+  let reserved = 0;
+  const resources: SessionResources = {
+    timer: options.timer, idleTimeoutMs, maxQueuedEvents, maxQueuedTokens, maxTopLogprobs,
+    reserveOutput(tokens) {
+      if (tokens > maxCollectTokens - reserved) throw new Error("collection output capacity exceeded");
+      const buffer = new Uint32Array(tokens);
+      reserved += tokens;
+      let released = false;
+      return { buffer, release() { if (!released) { released = true; reserved -= tokens; } } };
+    },
+  };
+  return {
+    async open(request: Request, control: RunControl) {
+      if (closed) throw new Error("inference engine is closed");
+      let session: EngineSession<Metrics> | undefined;
+      session = new EngineSession(String(++nextId), (cancellation) => planner.plan(request, cancellation),
+        { ...control }, resources,
+        () => { if (session) sessions.delete(session); });
+      if (session.state !== "terminal") sessions.add(session);
+      return session;
+    },
+    close() {
+      closed = true;
+      return closing ??= Promise.all([...sessions].map((session) => session.cancel("engine_closed"))).then(() => {});
+    },
+  };
+}