mlx-bun 0.2.0 → 0.4.0

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`).
@@ -26,31 +29,25 @@
 // server.py's all-caches-have-merge check). GLM's capability owns compressed
 // MLA/DSA merge/extract/filter and byte projection; Qwen3.5's SSM cache owns
 // its row state. The remaining request-shape exclusions include vision, LoRA,
-// serial-only KV modes, and user-fixed seed
+// serial-only KV modes
 // (reproducibility ⇒ solo, matching mlx-lm's _is_batchable). Temperature /
 // top-p / top-k DO batch (each row samples with its own seed). Full-attention
 // 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 { createRowSampling } from "../backends/mlx/row-sampling";
 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 { 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,
@@ -97,7 +114,7 @@ export type Vision = {
 };
 
 /** Per-token sink: returning `false` halts this generation (stop sequence).
- *  `logprobs` is only populated on the serial lane when the request asked for
+ *  `logprobs` is populated on either lane when the request asked for
  *  logprobs capture (GenerateOptions.logprobs / topLogprobs). */
 export type OnToken = (
   token: number,
@@ -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,13 @@ 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;
+      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 +184,14 @@ 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;
+      adapterNamespace?: (adapters: string[]) => string;
     } = {},
   ) {
+    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,19 +221,31 @@ 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;
   }
 
-  /** Resolves once the engine is idle (poll-based, ~`pollMs` granularity —
-   *  cheap vs the multi-hundred-ms flushes it paces; resolves without a
-   *  timer when already idle). The SSD write-behind gates each per-tensor
-   *  flush step on this so durability work never steals blocking
-   *  GPU-sync + writeSync slices from an active decode (the 2026-07-07
-   *  decode@ctx contamination — see server.ts's write-behind block). */
+  async close(): Promise<void> {
+    this.#closed = true;
+    this.#requests.close();
+    await this.#scheduler?.close();
+    await this.onIdle();
+  }
+
+  /** Observe idleness. Call runWhenIdle for background work that must also
+   * acquire ownership atomically with the idle check. */
   async onIdle(pollMs = 20): Promise<void> {
     while (this.busy) await new Promise<void>((r) => setTimeout(r, pollMs));
   }
 
+  /** Background work waits without requesting a batch drain. The final busy
+   * check and mutex acquisition have no await between them, so a foreground
+   * request cannot enter between the check and ownership registration. */
+  async runWhenIdle<T>(fn: () => Promise<T>, pollMs = 20): Promise<T> {
+    while (this.busy) await new Promise<void>(resolve => setTimeout(resolve, pollMs));
+    return this.runExclusive(fn);
+  }
+
   /** Cumulative rows routed to the batch lane since server start. The serial
    *  lane never advances this — /stats' race-free lane-routing observable
    *  (active/pending are instantaneous and read 0 once a request finishes). */
@@ -257,113 +261,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;
+    return this.#kvSchemeBatchable ??= !!this.opts.kvScheme && this.#binding.kvBatchable(this.opts.kvScheme);
   }
 
-  /** 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")
-    );
-  }
-
-  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 || shape.turboQuant) && 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 +285,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 +330,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 +357,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,40 +372,59 @@ 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;
+        });
       };
+      let started = false;
       return 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();
+        });
     }
 
+    const adapters = options.adapters?.length ? [...options.adapters] : undefined;
+    const adapterKey = adapters ? JSON.stringify(adapters) : "";
+    const cacheNamespace = adapters
+      ? () => this.opts.adapterNamespace?.(adapters) ?? adapterKey : "";
+    const context = adapters
+      ? this.#binding.bindAdapterContext?.(adapters, `adapters:${adapterKey}`) : undefined;
+
     // The scheduler owns grammar ready/accept sequencing. StepSampler owns the
     // shared processors -> mask -> logprobs -> sample -> history contract.
-    const stepSampler = makeStepSampler(options, {
-      tokenRepresentation: "device",
-      grammarWait: "external",
-      historyUpdate: "after-sample",
-      initialHistory: promptIds,
-    });
-    const sample = (logits1V: MlxArray, step: number): MlxArray =>
-      stepSampler.sample(logits1V, step).token;
+    const method = this.#binding.methodRequest?.(placement.execution, options);
+    const continuation = method ? undefined : this.#binding.continuationRequest?.(placement.execution, options, promptIds, onToken);
+    const sampling = method ? undefined : continuation ?? createRowSampling(makeStepSampler(options, {
+      tokenRepresentation: "device", grammarWait: "external",
+      historyUpdate: "after-sample", initialHistory: promptIds,
+      captureSelectedLogprob: options.logprobs === true,
+      captureTopLogprobs: options.topLogprobs,
+    }), onToken);
 
     let st;
     this.#rowsSubmitted++;
@@ -472,12 +433,14 @@ export class GenerationGateway {
     });
     try {
       st = await this.#ensureScheduler().submit({
-        promptIds,
+        promptIds, context, cacheNamespace, prefillChunkSize: options.prefillChunkSize,
+        continuation: continuation?.continuation,
+        statePolicy: this.#binding.statePolicy?.(placement.execution, options, promptIds.length + (options.maxTokens ?? 512)),
+        compiledDecode: placement.execution?.compiledDecode,
         maxTokens: options.maxTokens ?? 512,
-        eosTokenIds: options.eosTokenIds ?? this.model.config.eosTokenIds,
-        sample,
-        plainGreedy: stepSampler.isPlainGreedy,
-        onToken,
+        eosTokenIds: options.eosTokenIds ?? this.#binding.config.eosTokenIds,
+        ...(method ? { method, onToken } : { sample: sampling!.sample,
+          plainGreedy: sampling!.plainGreedy, onToken: sampling!.onToken }),
         onAdmitted: closeAdmission,
         trace,
         ...(signal ? { signal } : {}),
@@ -491,7 +454,7 @@ export class GenerationGateway {
       });
     } finally {
       closeAdmission?.();
-      stepSampler.dispose();
+      sampling?.dispose();
       // The scheduler never owns per-row grammar state.
       options.grammar?.dispose();
     }
@@ -500,15 +463,22 @@ 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: [],
+      ...(st.spec ? { spec: st.spec } : {}),
     };
   }
 
-  #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,
         // Phase 3.1: the batchable kvConfig composition is applied by the
         // scheduler (solo rows convert at serial chunk boundaries, then