mlx-bun 0.2.0 → 0.3.0

package/src/serve/completion-executor.ts

@@ -27,9 +27,10 @@ import {
   type RequestRejection,
 } from "./request-plan";
 import type { PromptResponseTrace } from "./prompt-response-trace";
+import { cleanupFailure } from "../engine/resources";
 
 export interface CompletionEngine {
-  place(shape: RequestShape): GenerationPlacement;
+  place(shape: RequestShape, options?: GenerateOptions): GenerationPlacement;
   run(
     promptIds: number[],
     options: GenerateOptions & { stopSequences?: string[] },
@@ -118,6 +119,8 @@ export interface CompletionUsage {
   completionTokens: number;
   totalTokens: number;
   speculation?: GenerateStats["spec"];
+  /** Token fast-forwarding telemetry (serial lane, MLX_BUN_FILL=strict). */
+  fill?: GenerateStats["fill"];
 }
 
 export interface CompletionSummary {
@@ -199,7 +202,7 @@ function finalLane(
   stats: GenerateStats,
 ): Lane {
   if (mechanism === "continuous") return "batched";
-  return stats.spec ? "serial+spec" : lane;
+  return stats.spec ? "serial+spec" : "serial";
 }
 
 function finishReason(
@@ -209,7 +212,7 @@ function finishReason(
 ): CompletionFinishReason {
   if (result.toolCalls.length) return "tool_calls";
   if (result.stopped) return "stop";
-  return stats.generatedTokens >= maxTokens ? "length" : "stop";
+  return stats.finishReason ?? (stats.generatedTokens >= maxTokens ? "length" : "stop");
 }
 
 function combineControl(
@@ -251,12 +254,14 @@ export class CompletionExecutor {
     try {
       if (usageProgress) control.onUsageProgress!(usageProgress);
       const closePlacement = control.trace?.begin("completion.placement");
-      const enginePlacement = this.engine.place(planned.shape);
+      const enginePlacement = this.engine.place(planned.shape, planned.options);
       closePlacement?.();
       if (enginePlacement.shape !== planned.shape)
         throw new Error("generation placement does not belong to this request shape");
       const mechanism = enginePlacement.mechanism;
-      const lane = initialLane(mechanism, planned.shape);
+      const lane = enginePlacement.execution
+        ? (mechanism === "continuous" ? "batched" : enginePlacement.execution.method === "speculative" ? "serial+spec" : "serial")
+        : initialLane(mechanism, planned.shape);
       recordLane(input.requestId, lane);
       const placement = { mechanism, lane, shape: planned.shape };
       input.onPlacement?.(placement);
@@ -314,6 +319,8 @@ export class CompletionExecutor {
         usageProgress.totalTokens = stats.promptTokens + stats.generatedTokens;
         if (stats.spec) usageProgress.speculation = stats.spec;
         else delete usageProgress.speculation;
+        if (stats.fill) usageProgress.fill = stats.fill;
+        else delete usageProgress.fill;
       }
       const summary = {
         content: result.content,
@@ -332,15 +339,16 @@ export class CompletionExecutor {
           completionTokens: stats.generatedTokens,
           totalTokens: stats.promptTokens + stats.generatedTokens,
           ...(stats.spec ? { speculation: stats.spec } : {}),
+          ...(stats.fill ? { fill: stats.fill } : {}),
         },
         logprobs: logprobs?.payload() ?? null,
       };
+      planned.dispose(); // retained preparation reservations outlive the native run
       closeCompletion?.();
       return summary;
     } catch (error) {
       closeCompletion?.();
-      planned.dispose();
-      throw error;
+      cleanupFailure(error, () => planned.dispose());
     }
   }
 }

package/src/serve/discovery-routes.ts

@@ -1,9 +1,8 @@
+import { modelServingBinding } from "../backends/mlx/model-serving";
 import pkgJson from "../../package.json" with { type: "json" };
-import type { ServerContext } from "../server";
+import type { ServingContext } from "./model-host";
 import type { GenerationGateway } from "./generation-gateway";
 import { fit } from "../fit";
-import { isEmbeddingModel } from "../embed";
-import { Glm52Model } from "../model/glm52";
 
 const pkgVersion = (pkgJson as { version: string }).version;
 
@@ -32,7 +31,7 @@ export interface DiscoveryRoutes {
 }
 
 export function createDiscoveryRoutes(
-  ctx: ServerContext,
+  ctx: ServingContext,
   gateway: Pick<GenerationGateway, "batchMode">,
   startedAt: number,
 ): DiscoveryRoutes {
@@ -46,7 +45,7 @@ export function createDiscoveryRoutes(
     async handle(url, request) {
       switch (matchDiscoveryRoute(request.method, url.pathname)) {
         case "library": {
-          if (!libraryCache || Date.now() - libraryCache.at > 30_000) {
+          if (url.searchParams.get("refresh") === "1" || !libraryCache || Date.now() - libraryCache.at > 30_000) {
             const { Registry, visionCapable, audioCapable } = await import("../registry");
             const { loadModelConfig } = await import("../config");
             const { supportTier } = await import("../model/support");
@@ -132,7 +131,7 @@ export function createDiscoveryRoutes(
             top_p: ctx.genDefaults.topP ?? null,
             top_k: ctx.genDefaults.topK ?? null,
           };
-          const isGlm52 = ctx.model instanceof Glm52Model;
+          const capabilities = modelServingBinding(ctx).discovery;
           const data: Array<Record<string, unknown>> = [{
             id: ctx.modelId,
             object: "model",
@@ -147,10 +146,10 @@ export function createDiscoveryRoutes(
             batch_mode: gateway.batchMode,
             tools: true,
             structured_output: true,
-            embeddings: isEmbeddingModel(ctx.model),
-            adapters: !isGlm52,
-            training: !isGlm52,
-            dsa: ctx.model instanceof Glm52Model && ctx.model.capabilities.dsa,
+            embeddings: capabilities.embeddings,
+            adapters: capabilities.adapters,
+            training: capabilities.training,
+            dsa: capabilities.dsa,
             mtp: ctx.draft?.provider.id === "glm52-native-mtp",
             capabilities: {
               chat_completions: true,
@@ -161,11 +160,11 @@ export function createDiscoveryRoutes(
               tools: true,
               structured_output: true,
               logprobs: true,
-              embeddings: isEmbeddingModel(ctx.model),
+              embeddings: capabilities.embeddings,
               vision: !!(ctx.vision || ctx.loadVision),
               audio: !!(ctx.audio || ctx.loadAudio),
-              adapters: !isGlm52,
-              training: !isGlm52,
+              adapters: capabilities.adapters,
+              training: capabilities.training,
             },
             gen_defaults: genDefaults,
           }];

package/src/serve/generation-gateway.ts

@@ -1,3 +1,6 @@
+import { AdmissionPool } from "../engine/admission";
+import { acquireReservation } from "./preparation";
+import { cleanupFailure, disposeResources } from "../engine/resources";
 // GenerationGateway is the scheduling seam between CompletionExecutor and
 // the inference runtime. It declares one mechanism per request: the preserved
 // strict/dedicated serial executor, or the continuous scheduler (`--batch N`).
@@ -32,25 +35,19 @@
 // AND sliding-window (Gemma) models both batch — the scheduler assembles each
 // layer's cache by type. Prompt-cache prefix reuse works on the batch lane
 // too (Phase 3.2): joiners take() at admission; never-merged lone rows put()
-// back on finish (merged rows' KV is not extracted — their entries age out).
+// back zero-copy on finish; merged rows with ≥256 prompt tokens are extracted
+// and put back by BatchScheduler.#extractAndPut (see docs/design/batching.md).
 
 import { MlxArray } from "../mlx/array";
 import type { RuntimeModel } from "../model/factory";
-import { DiffusionGemmaModel } from "../model/diffusion-gemma";
-import {
-  KVCache,
-  RotatingKVCache,
-  isBatchableCache,
-  isPlainKvCache,
-  isRotatingPlainCache,
-} from "../model/gemma4-base";
-import { SSMCache } from "../model/qwen3-delta";
-import { UniversalDenseModel } from "../model/universal/dense";
 import type { GenerateOptions, GenerateStats, TokenLogprobs } from "../generate";
 import type { KvScheme } from "../kv-scheme";
 import { makeStepSampler } from "../sampler";
-import { BatchScheduler, type RowPromptCache } from "./batch-scheduler";
-import { runtimeValue } from "../runtime-config";
+import type { RowPromptCache } from "../backends/mlx/batch-group";
+import type { MlxMemoryBudget } from "../backends/mlx/memory-guard";
+import { bindMlxGateway, type MlxGatewayBinding, type MlxBatchGroup } from "../backends/mlx/gateway-binding";
+import type { RuntimeConfig } from "../runtime-config";
+import type { ExecutionRequirements, ResolvedExecution } from "../contracts/execution";
 import type { PromptResponseTrace } from "./prompt-response-trace";
 
 /** Async mutex: acquire() resolves to a release fn; releases run FIFO. */
@@ -61,18 +58,38 @@ class AsyncMutex {
   get locked(): boolean {
     return this.#pending > 0;
   }
-  acquire(): Promise<() => void> {
+  acquire(signal?: AbortSignal): Promise<() => void> {
+    if (signal?.aborted) return Promise.reject(signal.reason);
     this.#pending++;
     let release!: () => void;
     const gate = new Promise<void>((r) => { release = r; });
     const wait = this.#tail;
     this.#tail = this.#tail.then(() => gate);
     let released = false;
-    return wait.then(() => () => {
+    const unlock = () => {
       if (released) return; // idempotent: a double release must not skew #pending
       released = true;
       this.#pending--;
       release();
+    };
+    return new Promise((resolve, reject) => {
+      let settled = false;
+      const abort = () => {
+        if (settled) return;
+        settled = true;
+        signal?.removeEventListener("abort", abort);
+        // Release this waiter's gate, not its predecessor's. Later acquirers
+        // still wait for the active holder, preserving mutual exclusion.
+        unlock();
+        reject(signal!.reason);
+      };
+      signal?.addEventListener("abort", abort, { once: true });
+      wait.then(() => {
+        if (settled) return;
+        settled = true;
+        signal?.removeEventListener("abort", abort);
+        resolve(unlock);
+      });
     });
   }
 }
@@ -85,7 +102,7 @@ export type Vision = {
   embeddings: MlxArray;
   /** bool [L] image-token mask for the bidirectional attention overlay.
    *  Absent when the prompt carries ANY audio — audio(-containing) prompts
-   *  run fully causal (audio-input-plan.md §3.3 Q1). */
+   *  run fully causal (docs/design/generic-model-support.md §3.3 Q1). */
   imageMask?: MlxArray;
   /** bool [L] union multimodal soft-token mask (image | audio) for
    *  per-layer-input id zeroing. Absent on the legacy vision-only shape,
@@ -112,53 +129,11 @@ export type SerialRun = (
   onToken: OnToken,
   vision?: Vision,
   trace?: PromptResponseTrace,
+  execution?: ResolvedExecution,
 ) => Promise<GenerateStats>;
 
 /** What the batchable decision needs from a request (cheap to compute). */
-export interface RequestShape {
-  readonly hasVision: boolean;
-  readonly hasAdapters: boolean;
-  readonly hasRepetitionPenalty: boolean;
-  /** The user explicitly set `seed` (reproducibility) — not the random default. */
-  readonly userSeed: boolean;
-  /** KV quantization is active. Per-layer configs may use the continuous
-   *  scheduler; uniform schemes use the serial mechanism. */
-  readonly kvQuant: boolean;
-  /** TurboQuant is active (docs/design/turboquant-kv.md). Solo-only in v1:
-   *  TurboQuantKVCache is a novel Cache implementation (not a KVCache/
-   *  RotatingKVCache subclass), so #modelCachesBatchable() already excludes
-   *  it automatically once a request's cache is converted — this flag is the
-   *  BELT on top of that automatic BRACES, an explicit refusal at the
-   *  request-shape level so a turbo request never reaches the scheduler even
-   *  before any cache conversion has happened. Both layers exist on purpose. */
-  readonly turboQuant: boolean;
-  /** Any of the mlx-lm sampler/processor extensions is active: min_p, XTC,
-   *  logit_bias, presence/frequency penalty. INFORMATIONAL ONLY: since the
-   *  continuous scheduler grew per-row samplers/logits processors, these all
-   *  compose. Placement deliberately does not gate on this field (see the
-   *  repetition-penalty note in #supportsContinuous). Kept so /stats and
-   *  scheduling traces can show what
-   *  a request carries. */
-  readonly hasLogitsExtras: boolean;
-  /** A grammar controller compiled for this request (response_format /
-   *  guided_*). Degrade-path requests (compile failed → prompt injection)
-   *  have NO controller and stay batchable — the injection already happened
-   *  at the prompt level. B0 routes grammar to serial; B1 makes it batchable
-   *  (per-row matchers), with MLX_BUN_GRAMMAR_BATCH=0 forcing the B0 serial
-   *  fallback as the A/B + kill switch. */
-  readonly hasGrammar: boolean;
-  /** The request asked for logprobs/top_logprobs capture. Batch-lane logprobs
-   *  are deferred (the scheduler's per-row sampler doesn't capture or read
-   *  back logprob arrays yet), so these route to the serial lane like the
-   *  other mlx-lm request extensions above. */
-  readonly wantsLogprobs: boolean;
-  /** A draft model is configured (serve --draft-model). Server-level and
-   *  upstream-parity: mlx_lm.server sets is_batchable = (draft is None), so
-   *  every request routes serial while a draft is mounted — speculation is a
-   *  B=1 latency optimization, batching a throughput one; they are different
-   *  mechanisms by design (grammar-spec-batching-integration.md). */
-  readonly hasDraft: boolean;
-}
+export interface RequestShape extends ExecutionRequirements {}
 
 export type GenerationMechanism = "serial" | "continuous";
 
@@ -167,12 +142,22 @@ export type GenerationMechanism = "serial" | "continuous";
 export interface GenerationPlacement {
   readonly shape: RequestShape;
   readonly mechanism: GenerationMechanism;
+  readonly execution?: ResolvedExecution;
+}
+
+export function disposeUnstartedRequest(options: GenerateOptions, vision?: Vision): void {
+  disposeResources([options.grammar, vision?.embeddings, vision?.imageMask,
+    vision?.multimodalMask, options.visionPixels].filter((resource) => resource != null));
 }
 
 export class GenerationGateway {
+  #closed = false;
+  readonly #runtime: RuntimeConfig;
+  readonly #binding: MlxGatewayBinding;
+  readonly #requests: AdmissionPool;
   readonly #mutex = new AsyncMutex();
   readonly #batch: number;
-  #scheduler: BatchScheduler | null = null;
+  #scheduler: MlxBatchGroup | null = null;
   #rowsSubmitted = 0;
   /** Serial-lane requests waiting for / holding the mutex. While > 0 the
    *  scheduler pauses admission (drain) so they can't be starved. */
@@ -182,11 +167,14 @@ export class GenerationGateway {
   /** Lazy, memoized: can the configured KV scheme convert every named cache? */
   #kvSchemeBatchable: boolean | null = null;
   constructor(
-    private readonly model: RuntimeModel,
+    model: RuntimeModel | MlxGatewayBinding,
     batch: number,
     private readonly serialRun: SerialRun,
     private readonly opts: {
       kvBudgetBytes?: number;
+      memoryBudget?: MlxMemoryBudget;
+      checkpoints?: boolean;
+      stateCodecs?: import("../kv-store").CacheCodecProvider;
       /** The server-wide KV scheme (server.ts kvScheme) — threaded to the
        *  scheduler at construction when batchable (Phase 3.1). turboQuant is
        *  never threaded to the scheduler (always solo-only — see placement);
@@ -197,9 +185,16 @@ export class GenerationGateway {
        *  back on finish. Safe to share with the serial lane — both use it
        *  only inside this gateway's mutual-exclusion domain. */
       promptCache?: RowPromptCache;
+      /** Finish cache writes before taking the generation lock. Writers
+       * need that same lock, so this hook must never run inside it. */
+      beforeSerial?: () => Promise<void>;
     } = {},
   ) {
+    this.#binding = "createBatchGroup" in model ? model : bindMlxGateway(model);
+    this.#runtime = this.#binding.runtime;
+    this.opts = Object.freeze({ ...opts });
     this.#batch = Math.max(1, Math.floor(batch));
+    this.#requests = new AdmissionPool(this.#batch);
   }
 
   /** True if `--batch N` (N>1) is on (batchability is then per-request). */
@@ -229,7 +224,15 @@ export class GenerationGateway {
    *  running/pending. activeRows/pendingRows alone MISS the serial lane —
    *  a serial generation holds the mutex but shows zero rows. */
   get busy(): boolean {
-    return this.#mutex.locked || this.activeRows > 0 || this.pendingRows > 0;
+    return this.#requests.active > 0 || this.#requests.queued > 0 ||
+      this.#mutex.locked || this.activeRows > 0 || this.pendingRows > 0;
+  }
+
+  async close(): Promise<void> {
+    this.#closed = true;
+    this.#requests.close();
+    await this.#scheduler?.close();
+    await this.onIdle();
   }
 
   /** Resolves once the engine is idle (poll-based, ~`pollMs` granularity —
@@ -257,113 +260,22 @@ export class GenerationGateway {
     };
   }
 
-  /** Cache-capability gate (mirrors mlx-lm server.py's all-caches-have-merge
-   *  check): the scheduler's dynamic-B ops (temporalView/merge/filter) exist
-   *  on KVCache, RotatingKVCache, and SSMCache (hybrid gated-DeltaNet models
-   *  — Qwen3.5 — batch via SSMCache.mergeRows/filter; MLX_BUN_BATCH_SSM=0 is
-   *  the kill switch back to serial routing). */
   #modelCachesBatchable(): boolean {
-    if (this.#cacheBatchable === null) {
-      // Tier-0 universal models: batchable for PLAIN full-attention archs
-      // only. The 2026-07-03 uneven-row bug (scalar-offset RoPE decoded
-      // padded rows at wrong positions) is fixed via
-      // UniversalRope.applyDynamic + ropeOffsetArr and GATED token-exact vs
-      // mlx-lm B=2 on Llama-3.2-3B (static + dynamic join/leave,
-      // tests/batched-decode-parity.test.ts "Llama 3B Tier-0"). The
-      // maskArray archs (gemma2-family: forwardLayers builds a pad-blind
-      // causal mask) and sliding-window universal archs remain UNVALIDATED
-      // cells → serial, per the per-model-cell discipline.
-      if (this.model instanceof UniversalDenseModel) {
-        const a = this.model.args;
-        this.#cacheBatchable =
-          !a.maskArray && !a.layerTypes?.includes("sliding_attention");
-        return this.#cacheBatchable;
-      }
-      const ssmOk = runtimeValue("MLX_BUN_BATCH_SSM") !== "0";
-      const proto = this.model.makeCache(); // fresh caches hold no buffers
-      this.#cacheBatchable = proto.every(
-        (c) =>
-          c instanceof KVCache ||
-          c instanceof RotatingKVCache ||
-          isBatchableCache(c) ||
-          (ssmOk && c instanceof SSMCache),
-      );
-      for (const c of proto) c.dispose();
-    }
-    return this.#cacheBatchable;
+    return this.#cacheBatchable ??= this.#binding.cachesBatchable();
   }
 
-  /** Phase 3.1 + milestone 2: the server's kv scheme batches iff it is a
-   *  per-layer kvConfig (the L2 mixed scheme) whose every configured layer
-   *  is a full-attention KVCache (BatchedQuantDecodeMaskCache + quantized
-   *  merge/extend/filter) or a rotating RotatingKVCache
-   *  (BatchedRotatingQuantCache — milestone 2, gemma's kv_config). Uniform
-   *  kvBits (quantizedKvStart threshold semantics via the serial no-byLayer
-   *  path) and configs naming SSM layers stay serial. */
   #kvBatchable(): boolean {
-    if (this.#kvSchemeBatchable === null) {
-      const scheme = this.opts.kvScheme;
-      if (scheme?.kind !== "affine-config") {
-        this.#kvSchemeBatchable = false;
-      } else {
-        const proto = this.model.makeCache();
-        try {
-          this.#kvSchemeBatchable = scheme.batchable(
-            this.model.config,
-            (layerIdx) =>
-              isPlainKvCache(proto[layerIdx]) || isRotatingPlainCache(proto[layerIdx]),
-          );
-        } finally {
-          for (const cache of proto) cache.dispose();
-        }
-      }
-    }
-    return this.#kvSchemeBatchable;
-  }
-
-  /** Whether the declared execution composition is implemented by the
-   *  continuous scheduler. This is a support check only: it never rewrites
-   *  MTP/drafting, KV, TurboQuant, grammar, adapters, or sampling. */
-  #supportsContinuous(shape: RequestShape): boolean {
-    // DiffusionGemma is non-autoregressive — the batch scheduler assumes the AR
-    // KV-cache decode path, so it always runs serially through generate().
-    if (this.model instanceof DiffusionGemmaModel) return false;
-    if (!this.batchingEnabled) return false;
-    if (!this.#modelCachesBatchable()) return false;
-    // repetitionPenalty / logits extras (min_p, XTC, logit_bias,
-    // presence/frequency penalties) batch through the same StepSampler and
-    // per-row device history as the serial lane. (Load-bearing beyond opt-in
-    // knobs: some models — Qwen3.5 — ship a default repetition penalty in
-    // generation_config.json, which used to route EVERY request serial.)
-    return (
-      !shape.hasVision &&
-      !shape.hasAdapters &&
-      !shape.wantsLogprobs &&
-      !shape.userSeed &&
-      // Phase 3.1: a kv-quant request batches when the server's scheme is
-      // the batchable kvConfig composition (the scheduler applies it);
-      // otherwise it routes serial exactly as before.
-      !(shape.kvQuant && !this.#kvBatchable()) &&
-      // TurboQuant: solo-only in v1, unconditionally — never partially
-      // batchable like kvConfig. The automatic instanceof exclusion in
-      // #modelCachesBatchable() already covers this once a cache converts;
-      // this is the explicit belt-and-braces refusal at request-shape time
-      // (mirrors the kvQuant pattern above — both layers exist on purpose).
-      !shape.turboQuant &&
-      !shape.hasDraft &&
-      // Grammar: B1 makes it batchable (per-row matchers) unless the kill
-      // switch forces serial. MLX_BUN_GRAMMAR_BATCH=0 = B0 behavior (serial),
-      // the A/B + kill lever for the new code, house style.
-      !(shape.hasGrammar && runtimeValue("MLX_BUN_GRAMMAR_BATCH") === "0")
-    );
+    return this.#kvSchemeBatchable ??= !!this.opts.kvScheme && this.#binding.kvBatchable(this.opts.kvScheme);
   }
 
-  place(shape: RequestShape): GenerationPlacement {
+  place(shape: RequestShape, options: GenerateOptions = {}): GenerationPlacement {
     const frozenShape = Object.freeze(shape);
-    return Object.freeze({
-      shape: frozenShape,
-      mechanism: this.#supportsContinuous(frozenShape) ? "continuous" : "serial",
+    const execution = this.#binding.plan(frozenShape, options, {
+      continuous: this.batchingEnabled,
+      quantizedBatch: shape.kvQuant && this.#kvBatchable(),
+      checkpoints: this.opts.checkpoints === true,
     });
+    return Object.freeze({ shape: frozenShape, mechanism: execution.mechanism, execution });
   }
 
   /** Run `fn` with exclusive ownership of the GPU + shared model state (the
@@ -372,26 +284,37 @@ export class GenerationGateway {
    *  mount/unmount all come through here, so nothing runs concurrently with
    *  batched decode steps (or with each other). Registers as a serial waiter
    *  so the scheduler drains (stops admitting) until `fn` has run. */
+  /** Shared by inference, preparation and managed GPU subprocesses. */
+  async acquireExecutionLease(
+    signal?: AbortSignal,
+    trace?: PromptResponseTrace,
+  ): Promise<import("../contracts/resources").DisposableResource> {
+    if (this.#closed) throw new Error("generation gateway is closed");
+    this.#serialWaiters++;
+    const closeAdmission = trace?.begin("engine.admission_wait", { mechanism: "serial" });
+    const resume = () => {
+      this.#serialWaiters--;
+      if (this.#serialWaiters === 0) this.#scheduler?.kick();
+    };
+    try {
+      const release = await this.#mutex.acquire(signal);
+      let released = false;
+      return { dispose: () => {
+        if (released) return;
+        released = true;
+        release(); resume();
+      } };
+    } catch (error) { resume(); throw error; }
+    finally { closeAdmission?.(); }
+  }
+
   async runExclusive<T>(
     fn: () => Promise<T>,
     trace?: PromptResponseTrace,
+    signal?: AbortSignal,
   ): Promise<T> {
-    this.#serialWaiters++;
-    const closeAdmission = trace?.begin("engine.admission_wait", {
-      mechanism: "serial",
-    });
-    try {
-      const release = await this.#mutex.acquire();
-      closeAdmission?.();
-      try {
-        return await fn();
-      } finally {
-        release();
-      }
-    } finally {
-      this.#serialWaiters--;
-      if (this.#serialWaiters === 0) this.#scheduler?.kick(); // resume admission
-    }
+    const lease = await this.acquireExecutionLease(signal, trace);
+    try { return await fn(); } finally { lease.dispose(); }
   }
 
   /** Run one generation on the appropriate lane. onToken is invoked per emitted
@@ -406,8 +329,26 @@ export class GenerationGateway {
     signal?: AbortSignal,
     trace?: PromptResponseTrace,
   ): Promise<GenerateStats> {
+    let reservation;
+    try { reservation = await acquireReservation(this.#requests, signal); }
+    catch (error) { cleanupFailure(error, () => disposeUnstartedRequest(options, vision)); }
+    try { return await this.#run(promptIds, options, onToken, vision, shape, placement, signal, trace); }
+    finally { reservation.dispose(); }
+  }
+
+  async #run(
+    promptIds: number[],
+    options: GenerateOptions & { stopSequences?: string[] },
+    onToken: OnToken,
+    vision: Vision | undefined,
+    shape: RequestShape,
+    placement: GenerationPlacement,
+    signal?: AbortSignal,
+    trace?: PromptResponseTrace,
+  ): Promise<GenerateStats> {
+    const disposeUnstarted = () => disposeUnstartedRequest(options, vision);
     if (placement.shape !== shape) {
-      options.grammar?.dispose();
+      disposeUnstarted();
       throw new Error("generation placement does not belong to this request shape");
     }
     try {
@@ -415,7 +356,7 @@ export class GenerationGateway {
     } catch (e) {
       // The gateway is the first component that accepts ownership of a
       // compiled controller. An already-aborted request reaches neither lane.
-      options.grammar?.dispose();
+      disposeUnstarted();
       throw e;
     }
     if (placement.mechanism === "serial") {
@@ -430,28 +371,40 @@ export class GenerationGateway {
       let lastHop = performance.now();
       const hop = (): Promise<void> => new Promise<void>((r) => setImmediate(r));
       const hoppingOnToken: OnToken = (token, lp) => {
-        if (signal?.aborted) return false;
+        signal?.throwIfAborted();
         const r = onToken(token, lp);
         if (r === false) return false;
         if (performance.now() - lastHop < 25) return r;
         lastHop = performance.now();
         if (r instanceof Promise)
-          return r.then((v) =>
-            v === false || signal?.aborted ? false : hop().then(() => signal?.aborted ? false : v),
-          );
-        return hop().then(() => signal?.aborted ? false : r);
+          return r.then(async (v) => {
+            if (v === false) return false;
+            signal?.throwIfAborted();
+            await hop();
+            signal?.throwIfAborted();
+            return v;
+          });
+        return hop().then(() => {
+          signal?.throwIfAborted();
+          return r;
+        });
       };
-      return this.runExclusive(async () => {
+      let started = false;
+      return Promise.resolve().then(() => this.opts.beforeSerial?.()).then(() => this.runExclusive(async () => {
         signal?.throwIfAborted();
+        started = true;
         const stats = await this.serialRun(
-          promptIds, options, hoppingOnToken, vision, trace,
+          promptIds, signal ? { ...options, signal } : options, hoppingOnToken, vision, trace, placement.execution,
         );
         signal?.throwIfAborted();
         return stats;
-      }, trace)
+      }, trace, signal))
         // Covers a serial waiter aborted before serialRun takes ownership.
         // generate() also disposes defensively; the operation is idempotent.
-        .finally(() => options.grammar?.dispose());
+        .finally(() => {
+          if (started) options.grammar?.dispose();
+          else disposeUnstarted();
+        });
     }
 
     // The scheduler owns grammar ready/accept sequencing. StepSampler owns the
@@ -473,8 +426,9 @@ export class GenerationGateway {
     try {
       st = await this.#ensureScheduler().submit({
         promptIds,
+        compiledDecode: placement.execution?.compiledDecode,
         maxTokens: options.maxTokens ?? 512,
-        eosTokenIds: options.eosTokenIds ?? this.model.config.eosTokenIds,
+        eosTokenIds: options.eosTokenIds ?? this.#binding.config.eosTokenIds,
         sample,
         plainGreedy: stepSampler.isPlainGreedy,
         onToken,
@@ -500,16 +454,23 @@ export class GenerationGateway {
       promptTokens: st.promptTokens,
       cachedTokens: st.cachedTokens,
       generatedTokens: st.generatedTokens,
-      prefillMs: 0, decodeMs: 0, prefillTps: 0, decodeTps: 0,
+      finishReason: st.finishReason,
+      prefillMs: st.prefillMs,
+      decodeMs: st.decodeMs,
+      prefillTps: st.prefillMs > 0 ? ((st.promptTokens - st.cachedTokens) / st.prefillMs) * 1000 : 0,
+      decodeTps: st.decodeMs > 0 && st.generatedTokens > 1 ? ((st.generatedTokens - 1) / st.decodeMs) * 1000 : 0,
       cacheTokens: [],
     };
   }
 
-  #ensureScheduler(): BatchScheduler {
+  #ensureScheduler(): MlxBatchGroup {
     if (!this.#scheduler)
-      this.#scheduler = new BatchScheduler(this.model, {
+      this.#scheduler = this.#binding.createBatchGroup({
+        runtime: this.#runtime,
         maxBatch: this.#batch,
+        stateCodecs: this.opts.stateCodecs,
         kvBudgetBytes: this.opts.kvBudgetBytes,
+        memoryBudget: this.opts.memoryBudget,
         // Phase 3.1: the batchable kvConfig composition is applied by the
         // scheduler (solo rows convert at serial chunk boundaries, then
         // merge as quantized triples). Only threaded when the scheme