@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/execution/fit.ts

@@ -0,0 +1,184 @@
+// Memory contracts — `mlx-bun fit`. Every term is deterministic:
+//   weights        — safetensors byte sizes (registry)
+//   KV bytes/token — config: layers × kv_heads × head_dim × bytes;
+//                    sliding-window layers saturate at the window
+//   prefill transient — chunk size we choose × calibrated bytes/token
+//   machine        — RAM (queried) + Metal wired ceiling fraction
+//
+// Calibration constants come from measured runs on the reference M4 Pro.
+
+import { totalmem } from "node:os";
+import type { ModelConfig } from "../artifacts/config";
+import type { MemoryPlan } from "../contracts/portable/memory-plan";
+import { moeOf } from "../models/support";
+import { kvBytesAt, sdpaFallbackBytes, type KvSchemeOptions } from "../state/kv-scheme";
+
+/** Decode-efficiency vs theoretical bandwidth ceiling, measured on the
+ *  reference machine (24.9 tok/s vs 30.3 ceiling @600 ctx). */
+export const DECODE_EFFICIENCY = 0.82;
+/** MoE decode efficiency vs the active-bytes ceiling. RECALIBRATED from
+ *  a cleared-machine matrix: 26B-A4B measured 54.5 tok/s
+ *  (python 55.7 — parity) vs ~71 tok/s raw ceiling → 0.76. The earlier
+ *  0.42 came from a session where BOTH stacks were memory-degraded
+ *  (32.3/33.0 "parity" — equally wrong, mutually consistent). */
+export const MOE_DECODE_EFFICIENCY = 0.76;
+/** Prefill transient bytes per chunk token (measured: ~1.1 GB @ 2048). */
+export const TRANSIENT_PER_TOKEN = 0.55e6;
+/** Fraction of unified RAM usable as GPU working set (Metal's
+ *  recommendedMaxWorkingSetSize is ~75% on consumer SKUs). */
+export const WIRED_FRACTION = 0.75;
+export const DEFAULT_CHUNK = 2048;
+
+export interface MachineSpec {
+  name: string;
+  ramBytes: number;
+  bandwidthGBs: number;
+}
+
+/** Representative Apple Silicon SKUs (memory bandwidth GB/s). */
+export const APPLE_SKUS: { chip: string; bandwidthGBs: number; ramOptions: number[] }[] = [
+  { chip: "M1", bandwidthGBs: 68, ramOptions: [8, 16] },
+  { chip: "M1 Pro", bandwidthGBs: 200, ramOptions: [16, 32] },
+  { chip: "M1 Max", bandwidthGBs: 400, ramOptions: [32, 64] },
+  { chip: "M1 Ultra", bandwidthGBs: 800, ramOptions: [64, 128] },
+  { chip: "M2", bandwidthGBs: 100, ramOptions: [8, 16, 24] },
+  { chip: "M2 Pro", bandwidthGBs: 200, ramOptions: [16, 32] },
+  { chip: "M2 Max", bandwidthGBs: 400, ramOptions: [32, 64, 96] },
+  { chip: "M2 Ultra", bandwidthGBs: 800, ramOptions: [64, 128, 192] },
+  { chip: "M3", bandwidthGBs: 100, ramOptions: [8, 16, 24] },
+  { chip: "M3 Pro", bandwidthGBs: 150, ramOptions: [18, 36] },
+  { chip: "M3 Max", bandwidthGBs: 400, ramOptions: [36, 48, 64, 96, 128] },
+  { chip: "M3 Ultra", bandwidthGBs: 819, ramOptions: [96, 256, 512] },
+  { chip: "M4", bandwidthGBs: 120, ramOptions: [16, 24, 32] },
+  { chip: "M4 Pro", bandwidthGBs: 273, ramOptions: [24, 48] },
+  { chip: "M4 Max", bandwidthGBs: 546, ramOptions: [36, 48, 64, 128] },
+];
+
+/** Chip name from sysctl ("M1 Max") + its bandwidth from APPLE_SKUS. */
+export function detectChip(): { name: string | null; bandwidthGBs: number | null } {
+  try {
+    const proc = Bun.spawnSync(["sysctl", "-n", "machdep.cpu.brand_string"]);
+    const name = proc.stdout.toString().trim().replace(/^Apple\s+/, "");
+    if (!name) return { name: null, bandwidthGBs: null };
+    const sku = APPLE_SKUS.find((s) => s.chip === name);
+    return { name, bandwidthGBs: sku?.bandwidthGBs ?? null };
+  } catch {
+    return { name: null, bandwidthGBs: null };
+  }
+}
+
+export function thisMachine(bandwidthGBs?: number): MachineSpec {
+  // Default bandwidth: the detected chip's table entry; 273 (M4 Pro, the
+  // original dev machine) only as the last resort.
+  const bw = bandwidthGBs ?? detectChip().bandwidthGBs ?? 273;
+  return { name: "this machine", ramBytes: totalmem(), bandwidthGBs: bw };
+}
+
+
+export interface FitReport extends MemoryPlan {
+  readonly strategy: "generic-kv";
+  readonly predictedDecodeTps: number;
+}
+
+export function fit(
+  config: ModelConfig,
+  weightsBytes: number,
+  ctx: number,
+  machine: MachineSpec = thisMachine(),
+  chunk: number = DEFAULT_CHUNK,
+  /** Bytes of `.experts.` tensors (registry). MoE decode reads only
+   *  top_k/num_experts of them per token; residency still needs all. */
+  expertsBytes = 0,
+  /** Explicit memory budget in bytes (admission control). When set it
+   *  replaces the machine-derived usable ceiling (ram × WIRED_FRACTION)
+   *  outright — the budget IS the usable envelope. */
+  usableBytes?: number,
+  /** Active KV-quant scheme; a quantized cache holds more context in the
+   *  same budget, so the solved ceiling (and admission) must bill it. */
+  kvScheme?: KvSchemeOptions,
+): FitReport {
+  const usable = usableBytes ?? machine.ramBytes * WIRED_FRACTION;
+  // Prefill transient = the calibrated per-chunk-token constant + the
+  // SDPA-fallback scores tensor (heads × min(chunk, ctx) × ctx × 2 B) for
+  // attention layers whose head dim MLX cannot fuse (kv-scheme.ts). The
+  // second term is what grows with context on Qwen3.x (head_dim 256) and
+  // Gemma 4 (256/512) and was unbilled before 2026-09-08.
+  const transientAt = (n: number) =>
+    Math.min(chunk, n) * TRANSIENT_PER_TOKEN + sdpaFallbackBytes(config, chunk, n);
+  const transient = transientAt(ctx);
+  const kv = kvBytesAt(config, ctx, kvScheme);
+  const total = weightsBytes + kv + transient;
+
+  // solve max context: weights + kv(n) + transient(n) ≤ usable. The same
+  // formula `fits` uses, so the solved ceiling always fits. It is monotone
+  // in n (KV linear, sliding saturating at the window, scores quadratic
+  // below the chunk then linear), so bisect over [0, maxPositionEmbeddings].
+  const totalAt = (n: number) => weightsBytes + kvBytesAt(config, n, kvScheme) + transientAt(n);
+  let maxCtx = 0;
+  {
+    let lo = 0;
+    let hi = Math.max(0, Math.floor(config.text.maxPositionEmbeddings));
+    if (totalAt(hi) <= usable) lo = hi;
+    else {
+      // invariant: totalAt(lo) ≤ usable (lo = 0 may itself not fit → 0)
+      if (totalAt(0) > usable) hi = 0;
+      while (hi - lo > 1) {
+        const mid = Math.floor((lo + hi) / 2);
+        if (totalAt(mid) <= usable) lo = mid;
+        else hi = mid;
+      }
+    }
+    maxCtx = lo;
+  }
+
+  // decode reads all weights + the KV cache once per token — except MoE
+  // expert weights, where only top_k of num_experts are touched per token
+  const moe = moeOf(config);
+  const isMoe = moe !== null;
+  const expertsSkipped = moe ? expertsBytes * (1 - moe.topK / moe.experts) : 0;
+  const bytesPerToken = weightsBytes - expertsSkipped + kv;
+  const predictedDecodeTps =
+    ((machine.bandwidthGBs * 1e9) / bytesPerToken) *
+    (isMoe ? MOE_DECODE_EFFICIENCY : DECODE_EFFICIENCY);
+
+  return {
+    schemaVersion: 1,
+    strategy: "generic-kv",
+    fits: total <= usable,
+    contextTokens: ctx,
+    weightsBytes,
+    kvBytes: kv,
+    transientBytes: transient,
+    reserveBytes: 0,
+    totalBytes: total,
+    usableBytes: usable,
+    maxSafeContext: maxCtx,
+    predictedDecodeTps,
+  };
+}
+
+/** The SKU matrix: which Apple Silicon configs run this model at `ctx`. */
+export function skuMatrix(
+  config: ModelConfig, weightsBytes: number, ctx: number, expertsBytes = 0,
+  kvScheme?: KvSchemeOptions,
+): { sku: string; ramGB: number; fits: boolean; maxContext: number; decodeTps: number }[] {
+  const rows: ReturnType<typeof skuMatrix> = [];
+  for (const sku of APPLE_SKUS) {
+    for (const ram of sku.ramOptions) {
+      const m: MachineSpec = {
+        name: `${sku.chip} ${ram}GB`,
+        ramBytes: ram * 2 ** 30,
+        bandwidthGBs: sku.bandwidthGBs,
+      };
+      const r = fit(config, weightsBytes, ctx, m, DEFAULT_CHUNK, expertsBytes, undefined, kvScheme);
+      rows.push({
+        sku: sku.chip,
+        ramGB: ram,
+        fits: r.fits,
+        maxContext: r.maxSafeContext,
+        decodeTps: r.predictedDecodeTps,
+      });
+    }
+  }
+  return rows;
+}

package/src/execution/gateway-binding.ts

@@ -0,0 +1,178 @@
+import { pagedDecodeCompiles } from "../state/layout";
+import { speculativePrefixNamespace, captureSpeculativeOptions } from "../generation/speculative/cache-identity";
+import { createOrdinaryContinuationRequest } from "./continuation-request";
+import type { ExecutionContext } from "../contracts/portable/scheduling";
+import type { ModelConfig } from "../artifacts/config";
+import type { MlxTokenGraph } from "../models/graph";
+import { declaredGraph, graphCapabilitiesOf } from "../models/capabilities";
+import { BF16_KV, COMPOSITION_DEFAULTS, type KvStorageScheme } from "../contracts/portable/composition";
+import { runtimeConfig, type RuntimeConfig } from "../runtime/config";
+import type { CacheServiceGraph } from "../state/cache-service";
+import { disposeResources } from "../runtime/resources";
+import { MlxBatchExecutionGroup } from "./batch-group";
+import { type MlxBatchExecutionGroupOptions, type MlxGroupMethodRequest } from "./batch-types";
+import type { DraftProvider } from "../generation/speculative/source";
+import { constraintDraftProvider } from "../generation/speculative/sources/ngram-source";
+import { bindSpeculativeGroupRequests } from "./speculative-group";
+import { bindGrammarGroupRequests } from "./grammar-group";
+import { bindFillGroupRequests } from "./fill-group";
+import { bindDenoisingGroupRequests } from "./denoising-group";
+import { bindLegacyDenoisingModel } from "../generation/bindings/denoising";
+import type { GenerateOptions } from "../generation/index";
+import type { ExecutionRequirements, ResolvedExecution } from "../contracts/portable/execution";
+import type { GraphCapabilities } from "../contracts/portable/graph";
+import { resolveExecution } from "./plan";
+import type { MlxPromptInput, Vision } from "../contracts/mlx/media";
+import type { AdaptiveDraftGateSettings } from "../contracts/mlx/graph";
+
+export interface MlxBatchGroup extends Pick<MlxBatchExecutionGroup,
+  "activeRows" | "pendingRows" | "submit" | "kick" | "close"> {
+  runPreparation?<T>(work: () => Promise<T>, signal?: AbortSignal): Promise<T>;
+}
+
+/** Planning consumes the graph's declared capabilities (`declaredGraph`), the
+ * request, and dynamic state; it never inspects a model's class, type, or name,
+ * nor how the composition's caches store keys and values: every cache lego
+ * serves every method. Request state, prefix reuse and persistence are the
+ * cache service's (`state/cache-service`), handed to `createBatchGroup`. */
+export interface MlxGatewayBinding {
+  mediaInput?(input: Vision): MlxPromptInput;
+  /** A replayable ordinary request when the plan qualified one: its sampler
+   * and the key the cache service checkpoints it under every `interval` tokens
+   * (the service's `checkpointEveryTokens`). */
+  continuationRequest?(execution: ResolvedExecution | undefined, options: GenerateOptions, prompt: number[],
+    onToken: Parameters<typeof createOrdinaryContinuationRequest>[0]["onToken"],
+    interval: number | undefined): ReturnType<typeof createOrdinaryContinuationRequest> | undefined;
+  readonly config: ModelConfig;
+  readonly runtime: RuntimeConfig;
+  /** What the composition's cache service builds request state from: the
+   * bound graph's caches, config and composition. */
+  readonly graph: CacheServiceGraph;
+  /** What the bound graph declares, as resolved when it was bound. */
+  readonly capabilities: GraphCapabilities;
+  /** `scheduling.checkpoints`: the cache service keeps generation checkpoints. */
+  plan(request: ExecutionRequirements, options: GenerateOptions, scheduling: { checkpoints: boolean }): ResolvedExecution;
+  bindAdapterContext?(adapters: string[], key: string): ExecutionContext;
+  /** Whether a request's prefix may be reused and its state kept (the batch
+   * request's `reusesPrefix`): media token IDs alone do not identify prepared
+   * embeddings, so an unidentified media prefix takes fresh state. */
+  reusesPrefix(execution: ResolvedExecution | undefined): boolean;
+  /** The prefix namespace of a request whose own namespace is `base` (its
+   * adapters, a prepared media prefix): its method's over it. The cache service
+   * adds the composition's paging. Null: the method keeps no prefixes. */
+  prefixNamespace?(execution: ResolvedExecution | undefined, options: GenerateOptions, base: string): string | null;
+  methodRequest?(execution: ResolvedExecution | undefined, options: GenerateOptions): MlxGroupMethodRequest | undefined;
+  createBatchGroup(options: MlxBatchExecutionGroupOptions): MlxBatchGroup;
+}
+
+/** `adaptiveDepth`: each round drafts up to `numDraftTokens`, chosen from the
+ *  decayed acceptance rate and the graph's measured `verifyRoundCosts` (the
+ *  graph must declare them) by a gate with the settings `gate` (the graph's
+ *  `adaptiveDraftGates` for the drafter's kind); otherwise every round drafts
+ *  `numDraftTokens`.
+ *  The graph's composition names its KV scheme and scheduler switches (a
+ *  graph without one runs the defaults): every request's, so the binding
+ *  reads them once, here. */
+export function bindMlxGateway(model: MlxTokenGraph,
+  draft?: { provider: DraftProvider; numDraftTokens: number; adaptiveDepth?: boolean; gate?: AdaptiveDraftGateSettings }): MlxGatewayBinding {
+  const kv: KvStorageScheme = model.composition?.kv ?? BF16_KV;
+  // The composition facts this binding plans with: the graph's record, else the defaults.
+  const facts = model.composition ?? COMPOSITION_DEFAULTS;
+  const runtime = runtimeConfig();
+  // The declarations are resolved once, when the graph is bound; planning reads
+  // only these and the request.
+  const graph = declaredGraph(model);
+  const caps = graphCapabilitiesOf(graph);
+  const verification = caps.speculation;
+  // Denoising rows interleave through their own grouped method. Token-level
+  // methods (speculation, grammar proposals, fill) never bind to this graph.
+  const denoising = caps.method === "denoising" ? bindDenoisingGroupRequests(bindLegacyDenoisingModel(model)) : undefined;
+  // Grouped speculation binds any provider whose taps the graph's forward
+  // captures; one that cannot bind is refused by placement rather than served
+  // ordinarily.
+  const speculative = !denoising && draft?.provider.grouped
+    ? bindSpeculativeGroupRequests(model, draft.provider, draft.numDraftTokens, draft.adaptiveDepth ?? false, draft.gate) : undefined;
+  // Grammar jump: verified proposals where the graph's verifier qualifies them,
+  // committed spans otherwise.
+  const grammarProvider = caps.method === "autoregressive" && verification.grammarProposals && facts.grammarJump
+    ? constraintDraftProvider() : undefined;
+  const grammarProposals = grammarProvider ? bindSpeculativeGroupRequests(model, grammarProvider,
+    Math.max(1, Math.trunc(runtime.number("MLX_BUN_GRAMMAR_DRAFT_TOKENS", 3)))) : undefined;
+  const grammarSpans = caps.method === "autoregressive" && !verification.grammarProposals
+    ? bindGrammarGroupRequests(model) : undefined;
+  const adapterState = model.loraState;
+  const fillRequests = !denoising ? bindFillGroupRequests(model) : undefined;
+  const mediaInput = caps.media && caps.media.input !== "pixels"
+    ? (input: Vision): MlxPromptInput => graph.bindMediaInput!(input) : undefined;
+  return {
+    mediaInput,
+    config: model.config, runtime, graph: model, capabilities: caps,
+    continuationRequest(execution, options, prompt, onToken, interval) {
+      if (!execution?.checkpoint || execution.mechanism !== "continuous") return undefined;
+      if (interval === undefined) throw new Error("qualified continuation requires bound persistence services");
+      // The cache service keys checkpoints in the request state's namespace, as its prefixes are.
+      return createOrdinaryContinuationRequest({ options, prompt, onToken, execution, interval });
+    },
+    // Media token IDs alone do not identify the prepared embeddings, so an
+    // unidentified media prefix takes fresh state and reuses none.
+    reusesPrefix: execution => !(execution?.method === "autoregressive" && !execution.promptCache),
+    // The method's namespace: the request's base, or the speculative group's over it.
+    prefixNamespace: (execution, options, base) => {
+      if (execution?.method !== "speculative") return base;
+      const namespace = (execution.grammarJump ? grammarProvider : draft?.provider)?.grouped?.checkpointNamespace?.();
+      return namespace === undefined ? null : speculativePrefixNamespace(namespace, base, captureSpeculativeOptions(options));
+    },
+    methodRequest: (execution, options) => execution?.method === "denoising" ? denoising?.(options)
+      : execution?.method === "speculative"
+      ? (execution.grammarJump ? grammarProposals : speculative)?.(
+        options.fill && !execution.fill ? { ...options, fill: undefined } : options)
+      : execution?.grammarJump ? grammarSpans?.(options)
+      : execution?.fill ? fillRequests?.(options) : undefined,
+    ...(adapterState ? { bindAdapterContext(adapters: string[], key: string): ExecutionContext {
+      // Denoising applies each row's adapters around that row's own units, so
+      // its rows share one neutral group context whatever their adapters.
+      if (denoising) return { key: "", enter: () => () => {} };
+      const selected = [...adapters];
+      return { key, enter() {
+        const previous = adapterState.active;
+        adapterState.active = selected;
+        return () => { adapterState.active = previous; };
+      } };
+    } } : {}),
+    plan(request, options, scheduling) {
+      const sharedMethod = request.hasDraft ? speculative : grammarProposals;
+      const provider = request.hasDraft ? draft?.provider : grammarProvider;
+      // A graph whose verifier does not qualify adapters, or a provider that
+      // cannot serve the target's adapters, serves adapter requests ordinarily
+      // even when a draft is configured.
+      const sharedSpeculativeAdapters = verification.adapters && !!sharedMethod &&
+        provider?.grouped?.supportsTargetAdapters === true;
+      // Such a drafted request opens no draft provider and keeps the ordinary checkpoint rules.
+      const ignoredDraft = request.hasDraft && request.hasAdapters && !sharedSpeculativeAdapters;
+      return resolveExecution(request, {
+        checkpoints: scheduling.checkpoints,
+        sharedCheckpoints: scheduling.checkpoints &&
+          (!request.hasDraft || ignoredDraft) && !request.hasVision && !request.hasGrammar &&
+          !request.wantsLogprobs && !options.fill,
+        mediaBatch: !!mediaInput,
+        mediaPrefixCache: runtime.flag("MLX_BUN_MEDIA_PREFIX_CACHE", true),
+        groupedMethods: denoising ? ["denoising"] : sharedMethod ? ["autoregressive", "speculative"] : ["autoregressive"],
+        sharedGrammarProposals: !!grammarProposals,
+        // Committed spans take ordinary grammar rows; prepared media keeps masking.
+        sharedGrammarJump: !!grammarSpans && !request.hasVision,
+        sharedFill: !!fillRequests && !!options.fill,
+        sharedSpeculativeEcho: verification.externalTokens && !!options.fill?.plan.echo &&
+          provider?.grouped?.supportsExternalTokens === true,
+        speculativeLogprobs: !!sharedMethod && verification.logprobs,
+        sharedSpeculativeAdapters,
+        method: caps.method,
+        compiledDecode: caps.compiledDecode,
+        // The batch group coordinates per-row grammar for every token method;
+        // masks come from the shared sampler, so no graph qualifies or declines
+        // them. Denoising samples canvases, which an AR token mask cannot apply to.
+        grammarBatch: !denoising,
+      }, { fill: !!options.fill, compiledDecode: (!model.composition?.pagedKv || pagedDecodeCompiles(model.composition.pagedKv, model.composition.kv)), grammarJump: facts.grammarJump });
+    },
+    createBatchGroup: (options) => new MlxBatchExecutionGroup(model, options),
+  };
+}

package/src/execution/grammar-group.ts

@@ -0,0 +1,224 @@
+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 { Cache } from "../contracts/mlx/cache";
+import type { MlxTokenGraph } from "../models/graph";
+import type { GenerateOptions } from "../generation/index";
+import { MlxPhaseDispatch } from "../generation/phase-dispatch";
+import { tokenPhasesOf } from "../models/capabilities";
+import { appendGrammarSpan } from "../generation/grammar-step";
+import { makeStepSampler, type DeviceStepSampler } from "../sampling/index";
+import { cleanupFailure, disposeResources } from "../runtime/resources";
+import { snapshotGenerationPolicy } from "./request-policy";
+import { MlxPrefillCohort } from "./prefill-cohort";
+import type { MlxGroupedMethod, MlxGroupMethodHost, MlxGroupMethodRequest, MlxGroupPreparation, Row } from "./batch-types";
+
+interface RequestState {
+  sampling: DeviceStepSampler;
+  caches: Cache[];
+  retain?: () => void;
+  pending: MlxArray | null;
+}
+
+/** Preserve committed grammar spans while the shared scheduler interleaves
+ * active requests. Graph calls remain B1: spans of different lengths are not
+ * padded, split, or verified. The gateway owns each borrowed grammar matcher. */
+export function bindGrammarGroupRequests(model: MlxTokenGraph) {
+  return (input: GenerateOptions): MlxGroupMethodRequest => ({
+    key: "grammar-forced-span", data: snapshotGenerationPolicy(input),
+    open: host => new GrammarGroup(host, model),
+  });
+}
+
+class GrammarGroup implements MlxGroupedMethod {
+  readonly #dispatch: MlxPhaseDispatch;
+  readonly #requests = new Map<Row, RequestState>();
+  #next = 0;
+
+  constructor(readonly host: MlxGroupMethodHost, readonly model: MlxTokenGraph) {
+    this.#dispatch = new MlxPhaseDispatch(tokenPhasesOf(model));
+  }
+
+  prepare(row: Row): MlxGroupPreparation {
+    const options = row.req.method!.data as GenerateOptions;
+    if (!row.req.grammar || options.logprobs || (options.topLogprobs ?? 0) > 0 ||
+        row.req.continuation || row.req.promptInput)
+      throw new Error("forced grammar span requires an ordinary grammar request without logprobs, media, or continuation");
+    const sampling = makeStepSampler(options, { tokenRepresentation: "device", grammarWait: "external",
+      historyUpdate: "manual", initialHistory: row.req.promptIds });
+    this.#requests.set(row, { sampling, caches: [], pending: null });
+    const target = new MlxPrefillCohort({
+      chunkSize: this.host.prefillChunkSize,
+      tailSplit: this.host.prefillTailSplit,
+      cache: this.host.cache, maintain: this.host.kvMaintenance,
+      forward: async (ids, caches, step) => this.#dispatch.prompt(step.phase, step.sequences, step.tokens)(ids, caches, null),
+      project: (hidden, _caches, _completed, step) => {
+        const [, length, width] = hidden.shape as [number, number, number];
+        using last = hidden.slice([0, length - 1, 0], [1, length, width]);
+        return this.#dispatch.logits(step.sequences)(last);
+      },
+      complete: async (state, logits) => {
+        const request = this.#requests.get(row)!;
+        await row.req.grammar!.ready();
+        row.req.signal?.throwIfAborted();
+        request.pending = request.sampling.sample(logits, 0).token;
+        ops.asyncEvalAll([request.pending]);
+        // No publication here: the first complete grammar span must enter KV
+        // before even token zero becomes visible to an output callback.
+        request.caches = state.solo; state.solo = [];
+        request.retain = state.retain; state.retain = undefined;
+        row.sampled = 1;
+        this.host.join(row);
+      },
+      reject: (failed, error) => {
+        try { this.#close(failed); }
+        catch (cleanup) { error = new AggregateError([error, cleanup], "grammar preparation cleanup failed"); }
+        failed.reject(error);
+      },
+    });
+    const preparation: MlxGroupPreparation = {
+      get rows() { return target.rows; },
+      get tokenWeight() { return target.tokenWeight; },
+      canAdmit: false, supportsMixedWork: false,
+      advance: () => target.advance(),
+      dispose: () => disposeResources([target, { dispose: () => {
+        if (!this.host.rows.includes(row)) this.#close(row);
+      } }]),
+    };
+    try { target.admit(row); return preparation; }
+    catch (error) { return cleanupFailure(error, () => preparation.dispose()); }
+  }
+
+  async advance(): Promise<void> {
+    const live = this.host.rows.flatMap((row, index) => row.req.signal?.aborted ? [] : [index]);
+    if (live.length !== this.host.rows.length) {
+      const cancelled = this.host.rows.filter(row => row.req.signal?.aborted).map(row => {
+        let error = row.req.signal!.reason;
+        try { this.#close(row); }
+        catch (cleanup) { error = new AggregateError([error, cleanup], "cancelled grammar cleanup failed"); }
+        return { row, error };
+      });
+      // Filter membership only after every cancelled owner has closed. A
+      // request must not settle while its tensors are still request-owned.
+      this.host.filterRows(live);
+      for (const { row, error } of cancelled) row.reject(error);
+    }
+    if (!this.host.rows.length) return;
+    const index = this.#next % this.host.rows.length;
+    const row = this.host.rows[index]!;
+    this.#next = index + 1;
+    try { await this.#step(row); }
+    catch (error) {
+      try { this.#retire(row); }
+      catch (cleanup) { error = new AggregateError([error, cleanup], "grammar step cleanup failed"); }
+      row.reject(error);
+    }
+  }
+
+  async #step(row: Row): Promise<void> {
+    const state = this.#requests.get(row)!;
+    const grammar = row.req.grammar!;
+    row.req.signal?.throwIfAborted();
+    const current = state.pending!;
+    const token = ops.itemUint32(current);
+    const step = row.generated;
+    let forced: number[] | null = null;
+    if (step + 1 < row.req.maxTokens) {
+      grammar.accept(token);
+      await grammar.ready();
+      if (!grammar.isTerminated) {
+        forced = grammar.jumpForward(row.req.maxTokens - (step + 1));
+        if (forced) await grammar.ready();
+      }
+    }
+    let next: MlxArray | null = null;
+    try {
+      if (forced) {
+        state.sampling.commitDevice(current);
+        state.sampling.commitNumbers(forced);
+        const after = step + 1 + forced.length;
+        await appendGrammarSpan(this.#dispatch, state.caches, [token, ...forced],
+          ids => { row.fed.push(...ids); },
+          after < row.req.maxTokens && !grammar.isTerminated
+            ? logits => { next = state.sampling.sample(logits, after).token; } : undefined);
+      } else if (step + 1 < row.req.maxTokens && !grammar.isTerminated) {
+        state.sampling.commitDevice(current);
+        using ids = ops.reshape(current, [1, 1]);
+        using hidden = await this.#dispatch.decode(1)(ids, state.caches, null);
+        using logits = this.#dispatch.logits(1)(hidden);
+        next = state.sampling.sample(logits, step + 1).token;
+        row.fed.push(token);
+      }
+      if (next) ops.asyncEvalAll([next]);
+      row.req.signal?.throwIfAborted();
+      current.dispose(); state.pending = null;
+      row.generated++;
+      let stop = row.req.eosTokenIds.includes(token);
+      if (!stop) {
+        stop = await this.host.publish(row, token) === false;
+        if (!stop) {
+          row.req.signal?.throwIfAborted();
+          if ((row.generated - 1) % 256 === 0) clearCache();
+          for (const id of forced ?? []) {
+            row.req.signal?.throwIfAborted();
+            row.generated++;
+            if (await this.host.publish(row, id) === false) { stop = true; break; }
+            if ((row.generated - 1) % 256 === 0) clearCache();
+          }
+        }
+      }
+      row.req.signal?.throwIfAborted();
+      row.sampled = row.generated + Number(next !== null);
+      if (stop || !next) {
+        const unused = next; next = null; unused?.dispose();
+        this.#complete(row, stop || row.generated < row.req.maxTokens ? "stop" : "length");
+      } else { state.pending = next; next = null; }
+    } finally { next?.dispose(); }
+  }
+
+  #complete(row: Row, reason: "stop" | "length"): void {
+    const state = this.#requests.get(row)!;
+    let caches = state.caches; state.caches = [];
+    let retain = state.retain; state.retain = undefined;
+    const release = () => {
+      const owned = caches; caches = [];
+      const lease = retain; retain = undefined;
+      disposeResources([...owned, { dispose: () => lease?.() }]);
+    };
+    try {
+      // Match the direct generator: pending/sampler cleanup must succeed
+      // before a final prefix can become visible to another request.
+      this.#retire(row);
+      row.req.signal?.throwIfAborted();
+      if (this.host.cache.keeps(row.req) && !row.fedTainted) {
+        this.host.cache.publish(row.req, row.cacheNamespace ?? "", [...row.req.promptIds, ...row.fed], { caches, retain });
+        caches = []; retain = undefined;
+      }
+    } catch (error) { return cleanupFailure(error, release); }
+    release();
+    this.host.finish(row, reason);
+  }
+
+  #retire(row: Row): void {
+    this.host.filterRows(this.host.rows.flatMap((other, index) => other === row ? [] : [index]));
+  }
+
+  filterRows(keep: readonly number[], _discard: boolean): void {
+    const retained = new Set(keep.map(index => this.host.rows[index]!));
+    const removed = this.host.rows.filter(row => !retained.has(row));
+    this.#next = keep.filter(index => index < this.#next).length;
+    disposeResources(removed.map(row => ({ dispose: () => this.#close(row) })));
+  }
+
+  #close(row: Row): void {
+    const state = this.#requests.get(row);
+    if (!state) return;
+    this.#requests.delete(row);
+    disposeResources([...(state.pending ? [state.pending] : []), state.sampling,
+      ...state.caches, { dispose: () => state.retain?.() }]);
+  }
+
+  dispose(): void {
+    disposeResources([...this.#requests.keys()].map(row => ({ dispose: () => this.#close(row) })));
+  }
+}

package/src/execution/index.ts

@@ -0,0 +1,18 @@
+export * from "./engine";
+export * from "./session";
+export * from "./cancellation";
+export * from "./admission";
+export * from "./scheduler";
+export * from "./coordinator";
+export * from "./tasks";
+export * from "./plan";
+export * from "./methods";
+export * from "./batch-group";
+export * from "./gateway-binding";
+export * from "./row-sampling";
+export * from "./prefill-cohort";
+export * from "./prefill-rows";
+export * from "./continuation-types";
+export * from "./continuation-request";
+export * from "../runtime/resources";
+export * from "../runtime/config";