mlx-bun 0.0.13 → 0.2.0

package/src/serve/generation-gateway.ts

@@ -1,10 +1,11 @@
-// 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.
+// 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
@@ -34,16 +35,23 @@
 // back on finish (merged rows' KV is not extracted — their entries age out).
 
 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 {
+  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 { KvQuantSpec, TurboQuantScheme } from "../config";
-import { makeSampler, makeLogitsProcessors, toLogprobs } from "../sampler";
+import type { KvScheme } from "../kv-scheme";
+import { makeStepSampler } from "../sampler";
 import { BatchScheduler, type RowPromptCache } from "./batch-scheduler";
+import { runtimeValue } from "../runtime-config";
+import type { PromptResponseTrace } from "./prompt-response-trace";
 
 /** Async mutex: acquire() resolves to a release fn; releases run FIFO. */
 class AsyncMutex {
@@ -83,6 +91,9 @@ export type Vision = {
    *  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,17 +111,19 @@ export type SerialRun = (
   options: GenerateOptions & { stopSequences?: string[] },
   onToken: OnToken,
   vision?: Vision,
+  trace?: PromptResponseTrace,
 ) => Promise<GenerateStats>;
 
 /** What the batchable decision needs from a request (cheap to compute). */
 export interface RequestShape {
-  hasVision: boolean;
-  hasAdapters: boolean;
-  hasRepetitionPenalty: boolean;
+  readonly hasVision: boolean;
+  readonly hasAdapters: boolean;
+  readonly 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;
+  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
@@ -118,32 +131,42 @@ export interface RequestShape {
    *  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;
+  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
-   *  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
+   *  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. */
-  hasLogitsExtras: boolean;
+  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. */
-  hasGrammar: boolean;
+  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. */
-  wantsLogprobs: boolean;
+  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
-   *  modes by design (grammar-spec-batching-integration.md). */
-  hasDraft: boolean;
+   *  mechanisms by design (grammar-spec-batching-integration.md). */
+  readonly hasDraft: boolean;
+}
+
+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;
 }
 
 export class GenerationGateway {
@@ -156,12 +179,8 @@ export class GenerationGateway {
   #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,
     batch: number,
@@ -170,9 +189,9 @@ export class GenerationGateway {
       kvBudgetBytes?: number;
       /** 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
@@ -260,7 +279,7 @@ export class GenerationGateway {
           !a.maskArray && !a.layerTypes?.includes("sliding_attention");
         return this.#cacheBatchable;
       }
-      const ssmOk = process.env.MLX_BUN_BATCH_SSM !== "0";
+      const ssmOk = runtimeValue("MLX_BUN_BATCH_SSM") !== "0";
       const proto = this.model.makeCache(); // fresh caches hold no buffers
       this.#cacheBatchable = proto.every(
         (c) =>
@@ -284,37 +303,36 @@ export class GenerationGateway {
   #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.
+      if (scheme?.kind !== "affine-config") {
         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();
+        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;
   }
 
-  /** Decide whether a request joins the batch or runs serially. */
-  willBatch(shape: RequestShape): boolean {
+  /** 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: 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
+    // 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 (
@@ -336,20 +354,35 @@ export class GenerationGateway {
       // 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")
+      !(shape.hasGrammar && runtimeValue("MLX_BUN_GRAMMAR_BATCH") === "0")
     );
   }
 
+  place(shape: RequestShape): GenerationPlacement {
+    const frozenShape = Object.freeze(shape);
+    return Object.freeze({
+      shape: frozenShape,
+      mechanism: this.#supportsContinuous(frozenShape) ? "continuous" : "serial",
+    });
+  }
+
   /** Run `fn` with exclusive ownership of the GPU + shared model state (the
    *  serial lane's lock). THE single mutual-exclusion domain: the serial
    *  generation path, the curve /generate + /signal endpoints, and adapter
    *  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> {
+  async runExclusive<T>(
+    fn: () => Promise<T>,
+    trace?: PromptResponseTrace,
+  ): 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 {
@@ -369,8 +402,14 @@ export class GenerationGateway {
     onToken: OnToken,
     vision: Vision | undefined,
     shape: RequestShape,
+    placement: GenerationPlacement,
     signal?: AbortSignal,
+    trace?: PromptResponseTrace,
   ): Promise<GenerateStats> {
+    if (placement.shape !== shape) {
+      options.grammar?.dispose();
+      throw new Error("generation placement does not belong to this request shape");
+    }
     try {
       signal?.throwIfAborted();
     } catch (e) {
@@ -379,7 +418,7 @@ export class GenerationGateway {
       options.grammar?.dispose();
       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
@@ -404,75 +443,43 @@ export class GenerationGateway {
       };
       return this.runExclusive(async () => {
         signal?.throwIfAborted();
-        const stats = await this.serialRun(promptIds, options, hoppingOnToken, vision);
+        const stats = await this.serialRun(
+          promptIds, options, hoppingOnToken, vision, trace,
+        );
         signal?.throwIfAborted();
         return stats;
-      })
+      }, trace)
         // Covers a serial waiter aborted before serialRun takes ownership.
         // generate() also disposes defensively; the operation is idempotent.
         .finally(() => options.grammar?.dispose());
     }
 
-    // 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,
         maxTokens: options.maxTokens ?? 512,
         eosTokenIds: options.eosTokenIds ?? this.model.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 +490,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();
     }
@@ -507,7 +514,7 @@ export class GenerationGateway {
         // 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

package/src/serve/kv-budget.ts

@@ -0,0 +1,14 @@
+import type { ModelConfig } from "../config";
+import { kvBytesAt, KvScheme } from "../kv-scheme";
+
+/** Project one batch row's worst-case KV bytes using the same per-layer
+ * scheme that the server-wide admission ceiling uses. */
+export function batchRowKvBytes(
+  config: ModelConfig,
+  promptTokens: number,
+  maxTokens: number,
+  scheme?: KvScheme,
+): number {
+  const tokens = promptTokens + maxTokens;
+  return scheme ? scheme.bytesAt(config, tokens) : kvBytesAt(config, tokens);
+}