mlx-bun 0.0.8 → 0.0.10

package/src/serve/generation-gateway.ts

@@ -11,13 +11,22 @@
 // queue). One AsyncMutex enforces it: a serial run holds it for its duration;
 // the scheduler holds it for its whole active period (first admit → batch
 // empties). So batched requests run concurrently with EACH OTHER (the point),
-// but never alongside a serial-lane generation. A non-batchable request drains
-// the batch (waits for it to empty), then runs solo; queued batchable requests
-// wait for that solo run to finish. (mlx-lm's drain behavior.)
+// but never alongside a serial-lane generation. A non-batchable request DRAINS
+// the batch: while any serial-lane request is waiting, the scheduler stops
+// admitting new rows (admissionHeld), finishes the running ones, and releases
+// the mutex; the serial request runs solo, then kick() resumes admission.
+// (mlx-lm's drain_batch semantics — without this, sustained batchable traffic
+// starves the serial lane forever.) Non-generation GPU/model-state work (the
+// curve endpoints, adapter mount/unmount) goes through runExclusive() so there
+// is exactly ONE mutual-exclusion domain.
 //
-// v1 batchable gate (the rest → serial): batch>1 AND bf16 KV (no kv-quant — the
-// batched scheduler runs bf16; mixed-precision-KV batching is the novel-combo
-// L2 follow-up) AND no vision AND no LoRA adapters AND no repetition penalty
+// v1 batchable gate (the rest → serial): batch>1 AND every model cache is a
+// dynamic-B-capable type (KVCache | RotatingKVCache — mirrors mlx-lm
+// server.py's all-caches-have-merge check; Qwen3.5's SSMCache has no
+// merge/filter/temporalView, so hybrid models route serial until the ArraysCache
+// port — batching-v2-plan item h) AND bf16 KV (no kv-quant — the batched
+// scheduler runs bf16; mixed-precision-KV batching is the novel-combo L2
+// follow-up) AND no vision AND no LoRA adapters AND no repetition penalty
 // (per-row logits processors are a later refinement) AND no user-fixed seed
 // (reproducibility ⇒ solo, matching mlx-lm's _is_batchable). Temperature /
 // top-p / top-k DO batch (each row samples with its own seed). Full-attention
@@ -26,10 +35,13 @@
 // in v1 (cachedTokens = 0).
 
 import { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
 import type { RuntimeModel } from "../model/factory";
 import { DiffusionGemmaModel } from "../model/diffusion-gemma";
-import type { GenerateOptions, GenerateStats } from "../generate";
-import { makeSampler, toLogprobs } from "../sampler";
+import { KVCache, RotatingKVCache } from "../model/gemma4-base";
+import { SSMCache } from "../model/qwen3-delta";
+import type { GenerateOptions, GenerateStats, TokenLogprobs } from "../generate";
+import { makeSampler, makeLogitsProcessors, toLogprobs } from "../sampler";
 import { BatchScheduler } from "./batch-scheduler";
 
 /** Async mutex: acquire() resolves to a release fn; releases run FIFO. */
@@ -46,8 +58,13 @@ class AsyncMutex {
 
 export type Vision = { embeddings: MlxArray; imageMask: MlxArray };
 
-/** Per-token sink: returning `false` halts this generation (stop sequence). */
-export type OnToken = (token: number) => void | boolean | Promise<void | boolean>;
+/** Per-token sink: returning `false` halts this generation (stop sequence).
+ *  `logprobs` is only populated on the serial lane when the request asked for
+ *  logprobs capture (GenerateOptions.logprobs / topLogprobs). */
+export type OnToken = (
+  token: number,
+  logprobs?: TokenLogprobs,
+) => void | boolean | Promise<void | boolean>;
 
 /** The serial lane — exactly today's runGeneration (prompt-cache reuse + the
  *  generate() pipeline). The gateway calls it under the mutex. */
@@ -67,12 +84,29 @@ export interface RequestShape {
   userSeed: boolean;
   /** KV quantization is active (kvConfig/kvBits) — batched is bf16-only in v1. */
   kvQuant: boolean;
+  /** Any of the mlx-lm sampler/processor extensions is active: min_p, XTC,
+   *  logit_bias, presence/frequency penalty. Safe v1: they ALL route to the
+   *  serial lane alongside repetition penalty. min_p/XTC are per-row samplers
+   *  and could batch (the batched lane already builds a per-row sampler);
+   *  keeping them serial until the batched path grows per-row logits
+   *  processors keeps one gate for the whole family. */
+  hasLogitsExtras: 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;
 }
 
 export class GenerationGateway {
   readonly #mutex = new AsyncMutex();
   readonly #batch: number;
   #scheduler: BatchScheduler | null = null;
+  /** 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;
 
   constructor(
     private readonly model: RuntimeModel,
@@ -92,21 +126,69 @@ export class GenerationGateway {
     return this.#scheduler?.activeRows ?? 0;
   }
 
+  /** 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) {
+      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 ||
+          (ssmOk && c instanceof SSMCache),
+      );
+      for (const c of proto) c.dispose();
+    }
+    return this.#cacheBatchable;
+  }
+
   /** 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 (
-      this.batchingEnabled &&
       !shape.hasVision &&
       !shape.hasAdapters &&
-      !shape.hasRepetitionPenalty &&
+      !shape.wantsLogprobs &&
       !shape.userSeed &&
       !shape.kvQuant
     );
   }
 
+  /** 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> {
+    this.#serialWaiters++;
+    try {
+      const release = await this.#mutex.acquire();
+      try {
+        return await fn();
+      } finally {
+        release();
+      }
+    } finally {
+      this.#serialWaiters--;
+      if (this.#serialWaiters === 0) this.#scheduler?.kick(); // resume admission
+    }
+  }
+
   /** Run one generation on the appropriate lane. onToken is invoked per emitted
    *  token (its `false` halts); resolves with stats when the generation ends. */
   async run(
@@ -117,32 +199,74 @@ export class GenerationGateway {
     shape: RequestShape,
   ): Promise<GenerateStats> {
     if (!this.willBatch(shape)) {
-      const release = await this.#mutex.acquire();
-      try {
-        return await this.serialRun(promptIds, options, onToken, vision);
-      } finally {
-        release();
-      }
+      // 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
+      // stall until the request finishes (measured 2.5 s on a 512-token cpm5
+      // run). Rate-limited to ≥25 ms — the SSE flush hop's budget, ~0.1% of
+      // decode — and applied HERE so every serial caller (streaming or not,
+      // all four API surfaces) is covered. The batch lane needs nothing: its
+      // drive loop yields per step (measured /stats ~8 ms mid-generation).
+      let lastHop = performance.now();
+      const hop = (): Promise<void> => new Promise<void>((r) => setImmediate(r));
+      const hoppingOnToken: OnToken = (token, lp) => {
+        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 ? false : hop().then(() => v)));
+        return hop().then(() => r);
+      };
+      return this.runExclusive(() => this.serialRun(promptIds, options, hoppingOnToken, vision));
     }
 
-    // Per-row sampler, mirroring generate()'s sampleStep (no logits processors:
-    // repetition penalty was excluded by willBatch). Greedy (temperature 0) is
-    // argmax; temp>0 uses this request's own seed → independent per-row RNG.
+    // 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 => {
-      const lp = toLogprobs(logits1V);
+      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;
+      }
+      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;
     };
 
-    const st = await this.#ensureScheduler().submit({
-      promptIds,
-      maxTokens: options.maxTokens ?? 512,
-      eosTokenIds: options.eosTokenIds ?? this.model.config.eosTokenIds,
-      sample,
-      onToken,
-    });
+    let st;
+    try {
+      st = await this.#ensureScheduler().submit({
+        promptIds,
+        maxTokens: options.maxTokens ?? 512,
+        eosTokenIds: options.eosTokenIds ?? this.model.config.eosTokenIds,
+        sample,
+        onToken,
+      });
+    } finally {
+      history?.dispose();
+      history = null;
+    }
 
     return {
       promptTokens: st.promptTokens,
@@ -158,6 +282,9 @@ export class GenerationGateway {
       this.#scheduler = new BatchScheduler(this.model, {
         maxBatch: this.#batch,
         lock: { acquire: () => this.#mutex.acquire() },
+        // Drain: pause admission while any serial-lane request waits; the
+        // runExclusive finally-block kick()s the loop back awake.
+        admissionHeld: () => this.#serialWaiters > 0,
       });
     return this.#scheduler;
   }