mlx-bun 0.0.13 → 0.3.0

package/src/serve/generation-gateway.ts

@@ -1,10 +1,14 @@
-// Generation gateway — the seam between the server's request handlers and the
-// two execution lanes: the serial single-queue path (today's default) and the
-// continuous-batching scheduler (`--batch N`). The handler builds an onToken
-// closure (its own StopMatcher + tool router + SSE stream) and calls run();
-// the gateway routes it to the right lane. Per-row SSE fan-out falls out for
-// free — each request keeps its own onToken/stream and the scheduler just
-// invokes the right row's onToken.
+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`).
+// The scheduler itself chooses its B=1 fast path or B=N step from its active
+// row count; placement does not predict a batch size. The executor owns each
+// request's semantic token sink and passes its callback plus one immutable
+// placement into run(). Each request keeps its own callback, so the scheduler
+// can fan tokens out to the matching response stream.
 //
 // The two lanes are MUTUALLY EXCLUSIVE on the GPU (and on shared model state
 // like loraState, which generate() mutates per-generation assuming a serialized
@@ -31,19 +35,20 @@
 // 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 * as ops from "../mlx/ops";
 import type { RuntimeModel } from "../model/factory";
-import { DiffusionGemmaModel } from "../model/diffusion-gemma";
-import { KVCache, RotatingKVCache, isBatchableCache } 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 { KvQuantSpec, TurboQuantScheme } from "../config";
-import { makeSampler, makeLogitsProcessors, toLogprobs } from "../sampler";
-import { BatchScheduler, type RowPromptCache } from "./batch-scheduler";
+import type { KvScheme } from "../kv-scheme";
+import { makeStepSampler } from "../sampler";
+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. */
 class AsyncMutex {
@@ -53,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);
+      });
     });
   }
 }
@@ -77,12 +102,15 @@ 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,
    *  where zeroing falls back to imageMask. */
   multimodalMask?: MlxArray;
+  /** Qwen3.5/3.8 vision: the request's mRoPE positions + decode delta,
+   *  installed on the model for exactly this serial run (server-side). */
+  mrope?: import("../model/qwen3-mrope").MropeRequestState;
 };
 
 /** Per-token sink: returning `false` halts this generation (stop sequence).
@@ -100,87 +128,73 @@ export type SerialRun = (
   options: GenerateOptions & { stopSequences?: string[] },
   onToken: OnToken,
   vision?: Vision,
+  trace?: PromptResponseTrace,
+  execution?: ResolvedExecution,
 ) => Promise<GenerateStats>;
 
 /** What the batchable decision needs from a request (cheap to compute). */
-export interface RequestShape {
-  hasVision: boolean;
-  hasAdapters: boolean;
-  hasRepetitionPenalty: boolean;
-  /** The user explicitly set `seed` (reproducibility) — not the random default. */
-  userSeed: boolean;
-  /** KV quantization is active (kvConfig/kvBits) — batched is bf16-only in v1. */
-  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. */
-  turboQuant: boolean;
-  /** Any of the mlx-lm sampler/processor extensions is active: min_p, XTC,
-   *  logit_bias, presence/frequency penalty. INFORMATIONAL ONLY: since the
-   *  batched lane grew per-row samplers/logits processors these all batch —
-   *  willBatch() deliberately does not gate on this field (see its
-   *  repetitionPenalty note). Kept so /stats and lane tracing can show what
-   *  a request carries. */
-  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. */
-  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. */
-  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
-   *  modes by design (grammar-spec-batching-integration.md). */
-  hasDraft: boolean;
+export interface RequestShape extends ExecutionRequirements {}
+
+export type GenerationMechanism = "serial" | "continuous";
+
+/** One immutable scheduling declaration for one exact request shape.
+ *  `continuous` means scheduler admission, not that another row is present. */
+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. */
   #serialWaiters = 0;
   /** Lazy, memoized: can this model's caches do the dynamic-B ops? */
   #cacheBatchable: boolean | null = null;
-  /** Lazy, memoized: can the server's kv scheme run under batching?
-   *  (Phase 3.1: kvConfig whose every configured layerIdx is a plain
-   *  full-attention KVCache — uniform kvBits and rotating-layer configs
-   *  stay serial. The scheme is a server-lifetime constant, so one memo.) */
+  /** 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 willBatch);
+       *  never threaded to the scheduler (always solo-only — see placement);
        *  it's here only so the gateway can see it's active. */
-      kvScheme?: { kvBits?: number; kvConfig?: KvQuantSpec[]; turboQuant?: TurboQuantScheme };
+      kvScheme?: KvScheme;
       /** The server's prompt cache (Phase 3.2): batch-lane joiners take()
        *  the longest usable prefix at admission; never-merged rows put()
        *  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). */
@@ -210,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 —
@@ -238,106 +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 = process.env.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 || (!scheme.kvBits && !scheme.kvConfig?.length)) {
-        // Only consulted when shape.kvQuant is SET: a kv-quant request with
-        // no scheme threaded to the gateway cannot be applied by the
-        // scheduler — batching it would SILENTLY DROP the quantization (the
-        // exact composition bug optiq serve ships). Route serial.
-        this.#kvSchemeBatchable = false;
-      } else if (scheme.kvBits || !scheme.kvConfig?.length) {
-        this.#kvSchemeBatchable = false; // uniform bits: serial
-      } else {
-        const proto = this.model.makeCache();
-        this.#kvSchemeBatchable = scheme.kvConfig.every(
-          (e) => proto[e.layerIdx] instanceof KVCache ||
-                 proto[e.layerIdx] instanceof RotatingKVCache,
-        );
-        for (const c of proto) c.dispose();
-      }
-    }
-    return this.#kvSchemeBatchable;
+    return this.#kvSchemeBatchable ??= !!this.opts.kvScheme && this.#binding.kvBatchable(this.opts.kvScheme);
   }
 
-  /** Decide whether a request joins the batch or runs serially. */
-  willBatch(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: the per-row sampler folds in the
-    // same makeSampler/makeLogitsProcessors closures the serial lane runs,
-    // over a per-row device-side token history. (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 && process.env.MLX_BUN_GRAMMAR_BATCH === "0")
-    );
+  place(shape: RequestShape, options: GenerateOptions = {}): GenerationPlacement {
+    const frozenShape = Object.freeze(shape);
+    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
@@ -346,19 +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. */
-  async runExclusive<T>(fn: () => Promise<T>): Promise<T> {
+  /** 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++;
-    try {
-      const release = await this.#mutex.acquire();
-      try {
-        return await fn();
-      } finally {
-        release();
-      }
-    } finally {
+    const closeAdmission = trace?.begin("engine.admission_wait", { mechanism: "serial" });
+    const resume = () => {
       this.#serialWaiters--;
-      if (this.#serialWaiters === 0) this.#scheduler?.kick(); // resume admission
-    }
+      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> {
+    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
@@ -369,17 +325,41 @@ export class GenerationGateway {
     onToken: OnToken,
     vision: Vision | undefined,
     shape: RequestShape,
+    placement: GenerationPlacement,
     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) {
+      disposeUnstarted();
+      throw new Error("generation placement does not belong to this request shape");
+    }
     try {
       signal?.throwIfAborted();
     } 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 (!this.willBatch(shape)) {
+    if (placement.mechanism === "serial") {
       // Serial decode is an unbroken microtask chain (FFI + generator
       // resumes): without a periodic macrotask hop the event loop starves for
       // the WHOLE generation — /stats, /health, and new-connection accepts
@@ -391,88 +371,69 @@ 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();
-        const stats = await this.serialRun(promptIds, options, hoppingOnToken, vision);
+        started = true;
+        const stats = await this.serialRun(
+          promptIds, signal ? { ...options, signal } : options, hoppingOnToken, vision, trace, placement.execution,
+        );
         signal?.throwIfAborted();
         return stats;
-      })
+      }, 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();
+        });
     }
 
-    // Per-row sampler, mirroring generate()'s sampleStep: logits processors
-    // (logit_bias / repetition / presence / frequency penalties — the same
-    // device-side closures the serial lane runs) fold in over a per-row token
-    // history seeded with the prompt, then the sampler. Greedy (temperature 0)
-    // is argmax; temp>0 uses this request's own seed → independent per-row RNG.
-    const sampler = makeSampler(options);
-    const processors = makeLogitsProcessors(options);
-    // Device-side history (prompt + generated), maintained only when a
-    // processor needs it — generate()'s pushHistory, one per row. The sample
-    // closure is called once per sampled token in order, so extending it
-    // inside the closure keeps history exact.
-    let history: MlxArray | null = null;
-    if (processors.length > 0) history = ops.fromInt32(promptIds, [promptIds.length]);
-    const sample = (logits1V: MlxArray, step: number): MlxArray => {
-      let cur = logits1V; // caller-owned; only intermediates are disposed here
-      for (const p of processors) {
-        const next = p(history, cur);
-        if (cur !== logits1V) cur.dispose();
-        cur = next;
-      }
-      // Grammar mask (B1): after the standard processors, before the sampler —
-      // same "grammar has the final say" ordering as serial sampleStep.
-      // Precondition: the scheduler has awaited ready() before invoking this
-      // closure (the mask for this step is materialized). applyMask is sync.
-      if (options.grammar) {
-        const masked = options.grammar.applyMask(cur);
-        if (cur !== logits1V) cur.dispose();
-        cur = masked;
-      }
-      const lp = toLogprobs(cur);
-      if (cur !== logits1V) cur.dispose();
-      const tok = sampler(lp, step);
-      lp.dispose();
-      if (history) {
-        const t1 = ops.reshape(tok, [1]);
-        const prev = history;
-        history = ops.concatAxis([prev, t1], 0);
-        prev.dispose();
-        t1.dispose();
-      }
-      return tok;
-    };
+    // 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;
 
     let st;
     this.#rowsSubmitted++;
+    const closeAdmission = trace?.begin("engine.admission_wait", {
+      mechanism: "continuous",
+    });
     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,
-        // Vectorized-sampling eligibility: mirrors makeSampler's greedy branch
-        // (temperature 0, no curve override) with nothing per-row in the way
-        // (no processors, no grammar mask).
-        plainGreedy:
-          (options.temperature ?? 0) === 0 &&
-          !options.curve &&
-          processors.length === 0 &&
-          !options.grammar,
+        plainGreedy: stepSampler.isPlainGreedy,
         onToken,
+        onAdmitted: closeAdmission,
+        trace,
         ...(signal ? { signal } : {}),
         // B1: pass the per-row grammar controller through. The scheduler drives
         // accept/ready/terminate; this gateway OWNS disposal (finally below)
@@ -483,8 +444,8 @@ export class GenerationGateway {
         ...(options.snapshotAt !== undefined ? { snapshotAt: options.snapshotAt } : {}),
       });
     } finally {
-      history?.dispose();
-      history = null;
+      closeAdmission?.();
+      stepSampler.dispose();
       // The scheduler never owns per-row grammar state.
       options.grammar?.dispose();
     }
@@ -493,21 +454,28 @@ 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
         // passed #kvBatchable — otherwise those requests never reach here.
-        kvConfig: this.#kvBatchable() ? this.opts.kvScheme?.kvConfig : undefined,
+        kvScheme: this.#kvBatchable() ? this.opts.kvScheme : undefined,
         promptCache: this.opts.promptCache,
         lock: { acquire: () => this.#mutex.acquire() },
         // Drain: pause admission while any serial-lane request waits; the