mlx-bun 0.0.11 → 0.0.12

package/src/responses.ts

@@ -45,6 +45,7 @@ export class ResponseStore {
   constructor(
     readonly ttlMs: number = DEFAULT_TTL_MS,
     readonly maxBytes: number = DEFAULT_MAX_BYTES,
+    private readonly now: () => number = Date.now,
   ) {}
 
   get size(): number {
@@ -67,7 +68,7 @@ export class ResponseStore {
       input: [...entry.input],
       output: [...entry.output],
       instructions: entry.instructions,
-      createdAt: Date.now(),
+      createdAt: this.now(),
       size,
     });
     this.#bytes += size;
@@ -89,9 +90,11 @@ export class ResponseStore {
   }
 
   #evictExpired(): void {
-    const cutoff = Date.now() - this.ttlMs;
+    const cutoff = this.now() - this.ttlMs;
     for (const [id, entry] of this.#items) {
-      if (entry.createdAt >= cutoff) break; // oldest-first walk, like the oracle
+      // LRU touches reorder the map independently of TTL age, so every entry
+      // must be checked rather than stopping at the first fresh one.
+      if (entry.createdAt >= cutoff) continue;
       this.#bytes -= entry.size;
       this.#items.delete(id);
     }
@@ -733,8 +736,10 @@ export function translateOpenAiSseToResponses(
   const dec = new TextDecoder();
   let buf = "";
   let finalized = false;
+  let reader: ReadableStreamDefaultReader<Uint8Array> | null = null;
+  let cancelled = false;
   return new ReadableStream<Uint8Array>({
-    async start(controller) {
+    start(controller) {
       const emit = (frames: string[]) => {
         for (const f of frames) controller.enqueue(enc.encode(f));
       };
@@ -744,54 +749,62 @@ export function translateOpenAiSseToResponses(
         emit(translator.finalize());
         onComplete(translator.finalResponse());
       };
-      const reader = upstream.getReader();
-      try {
-        while (true) {
-          const { done, value } = await reader.read();
-          if (done) break;
-          buf += dec.decode(value, { stream: true });
-          let sep: number;
-          while ((sep = buf.indexOf("\n\n")) !== -1) {
-            const frame = buf.slice(0, sep);
-            buf = buf.slice(sep + 2);
-            const data = frame
-              .split("\n")
-              .find((l) => l.startsWith("data: "))
-              ?.slice(6);
-            if (data == null) continue;
-            if (data === "[DONE]" || data === '"[DONE]"') {
-              finish();
-              continue;
+      reader = upstream.getReader();
+      void (async () => {
+        try {
+          while (!cancelled) {
+            const { done, value } = await reader!.read();
+            if (done) break;
+            buf += dec.decode(value, { stream: true });
+            let sep: number;
+            while ((sep = buf.indexOf("\n\n")) !== -1) {
+              const frame = buf.slice(0, sep);
+              buf = buf.slice(sep + 2);
+              const data = frame
+                .split("\n")
+                .find((l) => l.startsWith("data: "))
+                ?.slice(6);
+              if (data == null) continue;
+              if (data === "[DONE]" || data === '"[DONE]"') {
+                finish();
+                continue;
+              }
+              const parsed = JSON.parse(data);
+              if (parsed?.error) {
+                emit([
+                  sse("error", {
+                    type: "error",
+                    code: "server_error",
+                    message: parsed.error.message ?? "generation failed",
+                    param: null,
+                  }),
+                ]);
+                continue;
+              }
+              emit(translator.addChunk(parsed));
             }
-            const parsed = JSON.parse(data);
-            if (parsed?.error) {
-              emit([
-                sse("error", {
-                  type: "error",
-                  code: "server_error",
-                  message: parsed.error.message ?? "generation failed",
-                  param: null,
-                }),
-              ]);
-              continue;
-            }
-            emit(translator.addChunk(parsed));
           }
+          if (!cancelled) finish();
+        } catch (e) {
+          if (!cancelled)
+            emit([
+              sse("error", {
+                type: "error",
+                code: "server_error",
+                message: (e as Error).message,
+                param: null,
+              }),
+            ]);
+        } finally {
+          reader?.releaseLock();
+          reader = null;
+          if (!cancelled) controller.close();
         }
-        finish();
-      } catch (e) {
-        emit([
-          sse("error", {
-            type: "error",
-            code: "server_error",
-            message: (e as Error).message,
-            param: null,
-          }),
-        ]);
-      } finally {
-        reader.releaseLock();
-        controller.close();
-      }
+      })();
+    },
+    async cancel(reason) {
+      cancelled = true;
+      await reader?.cancel(reason);
     },
   });
 }

package/src/serve/batch-scheduler.ts

@@ -61,7 +61,8 @@ import { flagOn } from "../flags";
 import { CompiledDecode } from "../model/compiled-decode";
 import type { Gemma4Model } from "../model/gemma4";
 import {
-  KVCache, QuantizedKVCache, RotatingKVCache, RotatingQuantizedKVCache, type Cache,
+  KVCache, QuantizedKVCache, RotatingKVCache, RotatingQuantizedKVCache,
+  isBatchableCache, type BatchableCache, type Cache,
 } from "../model/gemma4-base";
 import { BatchedDecodeMaskCache, mergeKVRows, extendKVRows, extractKVRow, filterKVRows } from "../model/batched-mask";
 import {
@@ -124,6 +125,10 @@ export interface BatchRequest {
    *  THIS row only (its submit promise rejects; siblings continue). May be
    *  async; keep it cheap — it runs inline in the step loop. */
   onToken: (token: number) => void | boolean | Promise<void | boolean>;
+  /** Client/request lifetime. An aborted pending row is removed before
+   *  admission; a row in prefill or decode is evicted at the next safe
+   *  scheduler boundary. */
+  signal?: AbortSignal;
   /** The row's sampler is PLAIN GREEDY (temperature 0, no curve) with no
    *  logits processors and no grammar — set by the gateway when true. Lets
    *  the scheduler take the vectorized sampling fast path (ONE
@@ -269,7 +274,8 @@ type LayerInner =
   | BatchedRotatingQuantCache // rot layer under a kv_config scheme (milestone 2)
   | RotatingKVCache // adopted lone-row state only (see #mergeJoiner adopt)
   | RotatingQuantizedKVCache // adopted lone-row state, quantized rot layer
-  | SSMCache;
+  | SSMCache
+  | BatchableCache;
 type Row1 = { keys: MlxArray; values: MlxArray };
 
 export class BatchScheduler {
@@ -300,8 +306,10 @@ export class BatchScheduler {
    *  only a lone adopted row holds un-copied entry caches). Runs after the
    *  adopted caches are disposed, or transfers back on put(). */
   #adoptedRetain: (() => void) | null = null;
-  readonly #kinds: ("full" | "rot" | "ssm")[]; // per-layer attention type
+  readonly #kinds: ("full" | "rot" | "ssm" | "owned-batch")[];
   readonly #rotMaxSize: number[]; // per-layer sliding window (rot layers only)
+  readonly #compressedProjectors: Array<(tokens: number) => number> | null;
+  readonly #batchCacheMaxTokens: number | null;
   /** layerIdx → mixed-precision spec (Phase 3.1); null = bf16 batch (v1). */
   readonly #kvByLayer: Map<number, KvQuantSpec> | null;
   /** Compiled decode runner for the B=1 serial-class case (Phase 3.2) —
@@ -320,8 +328,20 @@ export class BatchScheduler {
     this.#promptCache = opts.promptCache;
     const proto = model.makeCache(); // fresh caches hold no buffers
     this.#kinds = proto.map((c) =>
-      c instanceof RotatingKVCache ? "rot" : c instanceof SSMCache ? "ssm" : "full",
+      isBatchableCache(c)
+        ? "owned-batch"
+        : c instanceof RotatingKVCache
+          ? "rot"
+          : c instanceof SSMCache
+            ? "ssm"
+            : "full",
     );
+    this.#compressedProjectors = proto.every(isBatchableCache)
+      ? proto.map((cache) => (tokens: number) => cache.projectedBytes(tokens))
+      : null;
+    this.#batchCacheMaxTokens = proto.every(isBatchableCache)
+      ? Math.min(...proto.map((cache) => cache.maxTokens ?? Number.MAX_SAFE_INTEGER))
+      : null;
     this.#kvByLayer = opts.kvConfig?.length
       ? new Map(opts.kvConfig.map((e) => [e.layerIdx, e]))
       : null;
@@ -346,7 +366,19 @@ export class BatchScheduler {
   /** Projected KV bytes of one row at its worst case (full prompt + full
    *  completion; the sliding-window term is window-capped by kvBytesAt). */
   #rowKvBytes(row: Row): number {
-    return kvBytesAt(this.model.config, row.promptTokens + row.req.maxTokens);
+    const tokens = row.promptTokens + row.req.maxTokens;
+    return this.#compressedProjectors
+      ? this.#compressedProjectors.reduce((sum, project) => sum + project(tokens), 0)
+      : kvBytesAt(this.model.config, tokens);
+  }
+
+  async #forwardHidden(ids: MlxArray, cache: Cache[]): Promise<MlxArray> {
+    const asyncModel = this.model as RuntimeModel & {
+      forwardHiddenAsync?: (ids: MlxArray, cache: Cache[]) => Promise<MlxArray>;
+    };
+    return typeof asyncModel.forwardHiddenAsync === "function"
+      ? await asyncModel.forwardHiddenAsync(ids, cache)
+      : this.model.forwardHidden(ids, cache);
   }
 
   /** Projected aggregate KV of everything admitted (running + mid-prefill). */
@@ -364,6 +396,18 @@ export class BatchScheduler {
    *  that cannot fit even alone is rejected here (never deadlocks the
    *  queue); one that fits alone but not alongside the current batch waits. */
   #kvAdmits(candidate: Row): boolean {
+    const requestedTokens = candidate.promptTokens + candidate.req.maxTokens;
+    if (
+      this.#batchCacheMaxTokens !== null &&
+      requestedTokens > this.#batchCacheMaxTokens
+    ) {
+      this.#pending.shift();
+      candidate.reject(new RangeError(
+        `context limit: prompt ${candidate.promptTokens} + max_tokens ` +
+        `${candidate.req.maxTokens} exceeds ${this.#batchCacheMaxTokens}`,
+      ));
+      return false;
+    }
     if (this.#kvBudgetBytes === undefined) return true;
     const need = this.#rowKvBytes(candidate);
     if (need > this.#kvBudgetBytes && this.#running.length === 0 && !this.#prefill) {
@@ -383,12 +427,24 @@ export class BatchScheduler {
 
   /** Submit a request; resolves when its row finishes (EOS, stop, or length). */
   submit(req: BatchRequest): Promise<BatchStats> {
+    if (req.signal?.aborted) return Promise.reject(req.signal.reason);
     return new Promise<BatchStats>((resolve, reject) => {
+      let abortListener: (() => void) | null = null;
+      const cleanup = () => {
+        if (abortListener) req.signal?.removeEventListener("abort", abortListener);
+        abortListener = null;
+      };
       this.#pending.push({
-        req, resolve, reject,
+        req,
+        resolve: (stats) => { cleanup(); resolve(stats); },
+        reject: (error) => { cleanup(); reject(error); },
         current: 0, generated: 0, sampled: 0, promptTokens: req.promptIds.length,
         cachedTokens: 0, fed: [], fedTainted: false, merged: false,
       });
+      if (req.signal) {
+        abortListener = () => this.kick();
+        req.signal.addEventListener("abort", abortListener, { once: true });
+      }
       this.#ensureLoop();
     });
   }
@@ -409,6 +465,12 @@ export class BatchScheduler {
     let release: (() => void) | null = null;
     try {
       while (true) {
+        for (let i = this.#pending.length - 1; i >= 0; i--) {
+          const row = this.#pending[i]!;
+          if (!row.req.signal?.aborted) continue;
+          this.#pending.splice(i, 1);
+          row.reject(row.req.signal.reason);
+        }
         const held = this.#admissionHeld?.() === true;
         // Work = running rows, an in-flight prefill (finish it even under
         // drain — the row is already half-admitted), or admissible pending.
@@ -543,6 +605,7 @@ export class BatchScheduler {
    *  The final chunk samples token 0, emits it, and — if the row survives —
    *  merges it into the running batch; returns true (admission complete). */
   async #prefillChunk(p: PrefillState): Promise<boolean> {
+    p.row.req.signal?.throwIfAborted();
     const prompt = p.row.req.promptIds;
     // Boundary-snapshot chunking (mirrors generate.ts's snapshotAt split):
     // chunk edges land EXACTLY on snapAt so the clone is taken while the
@@ -554,7 +617,7 @@ export class BatchScheduler {
       const end = Math.min(p.pos + this.#prefillChunkSize, p.snapAt);
       const chunk = prompt.slice(p.pos, end);
       const ids = ops.fromInt32(chunk, [1, chunk.length]);
-      const h = this.model.forwardHidden(ids, p.solo);
+      const h = await this.#forwardHidden(ids, p.solo);
       ids.dispose();
       h.dispose();
       ops.evalAll(p.solo.flatMap((c) => c.state()));
@@ -584,7 +647,7 @@ export class BatchScheduler {
         this.#prefillChunkSize) {
       const chunk = prompt.slice(p.pos, p.pos + this.#prefillChunkSize);
       const ids = ops.fromInt32(chunk, [1, chunk.length]);
-      const h = this.model.forwardHidden(ids, p.solo);
+      const h = await this.#forwardHidden(ids, p.solo);
       ids.dispose();
       h.dispose(); // logits never computed for non-final chunks
       ops.evalAll(p.solo.flatMap((c) => c.state()));
@@ -602,7 +665,7 @@ export class BatchScheduler {
     if (tailSplit && p.pos < prompt.length - 1) {
       const head = prompt.slice(p.pos, prompt.length - 1);
       const ids = ops.fromInt32(head, [1, head.length]);
-      const h = this.model.forwardHidden(ids, p.solo);
+      const h = await this.#forwardHidden(ids, p.solo);
       ids.dispose();
       h.dispose();
       ops.evalAll(p.solo.flatMap((c) => c.state()));
@@ -618,7 +681,7 @@ export class BatchScheduler {
     // the last prompt token (L=1).
     const chunk = prompt.slice(p.pos);
     const ids = ops.fromInt32(chunk, [1, chunk.length]);
-    const h = this.model.forwardHidden(ids, p.solo);
+    const h = await this.#forwardHidden(ids, p.solo);
     ids.dispose();
     const [, Lc, H] = h.shape as [number, number, number];
     const hLast = h.slice([0, Lc - 1, 0], [1, Lc, H]);
@@ -700,6 +763,17 @@ export class BatchScheduler {
     const newInners: LayerInner[] = [];
     let newFullPad = this.#fullLeftPad;
     for (let layer = 0; layer < this.#kinds.length; layer++) {
+      if (this.#kinds[layer] === "owned-batch") {
+        const solo = p.solo[layer]!;
+        if (!isBatchableCache(solo))
+          throw new Error(`batch-capable layer ${layer} lost its cache capability`);
+        const merged = solo.makeEmptyBatch();
+        const previous = prev?.[layer];
+        merged.mergeRows(previous ? [previous, solo] : [solo]);
+        newFullPad = [...merged.leftPad];
+        newInners.push(merged);
+        continue;
+      }
       if (this.#kinds[layer] === "ssm") {
         // No temporal axis, no left-pad: B-axis concat of the state slots.
         // mergeRows steals the solo arrays when the batch starts cold, so the
@@ -1005,7 +1079,7 @@ export class BatchScheduler {
         if (!lg) {
           fwd = inners.map((c) =>
             c instanceof BatchedRotatingCache || c instanceof BatchedRotatingQuantCache ||
-            c instanceof SSMCache || unpadded
+            c instanceof SSMCache || isBatchableCache(c) || unpadded
               ? c
               : c instanceof QuantizedKVCache
                 ? new BatchedQuantDecodeMaskCache(c, B, this.#fullLeftPad)
@@ -1014,7 +1088,7 @@ export class BatchScheduler {
           const ids = this.#pendingToks
             ? ops.reshape(this.#pendingToks, [B, 1]) // feed the unread tokens
             : ops.fromInt32(rows.map((r) => r.current), [B, 1]); // pipeline cold
-          const h = this.model.forwardHidden(ids, fwd);
+          const h = await this.#forwardHidden(ids, fwd);
           ids.dispose();
           lg = this.model.logitsFromHidden(h); // [B,1,V]
           h.dispose();
@@ -1180,7 +1254,7 @@ export class BatchScheduler {
       const unpadded = this.#fullLeftPad.every((p) => p === 0);
       const fwd: Cache[] = inners.map((c) =>
         c instanceof BatchedRotatingCache || c instanceof BatchedRotatingQuantCache ||
-            c instanceof SSMCache || unpadded
+            c instanceof SSMCache || isBatchableCache(c) || unpadded
           ? c
           : c instanceof QuantizedKVCache
             ? new BatchedQuantDecodeMaskCache(c, B, this.#fullLeftPad)
@@ -1191,7 +1265,7 @@ export class BatchScheduler {
         const ids = prev
           ? ops.reshape(prev, [B, 1]) // feed the unread tokens (device array)
           : ops.fromInt32(rows.map((r) => r.current), [B, 1]); // pipeline cold
-        const h = this.model.forwardHidden(ids, fwd);
+        const h = await this.#forwardHidden(ids, fwd);
         ids.dispose();
         const lg = this.model.logitsFromHidden(h); // [B,1,V]
         h.dispose();
@@ -1294,6 +1368,7 @@ export class BatchScheduler {
    *  WITHOUT an onToken call; otherwise onToken(token) runs and `false` halts;
    *  reaching maxTokens ends with "length". Advances row.current on continue. */
   async #emit(row: Row, token: number): Promise<"continue" | "stop" | "length"> {
+    row.req.signal?.throwIfAborted();
     if (row.req.eosTokenIds.includes(token)) return "stop";
     const cont = await row.req.onToken(token);
     if (cont === false) return "stop";
@@ -1366,7 +1441,8 @@ export class BatchScheduler {
     const out: Cache[] = [];
     for (const inner of this.#inners!) {
       let c: Cache | null;
-      if (inner instanceof BatchedRotatingQuantCache) c = inner.extractRow(b);
+      if (isBatchableCache(inner)) c = inner.extractRow(b);
+      else if (inner instanceof BatchedRotatingQuantCache) c = inner.extractRow(b);
       else if (inner instanceof BatchedRotatingCache) c = inner.extractRow(b);
       else if (inner instanceof SSMCache)
         // Coverage gate: recurrent state is UNTRIMMABLE, so an entry is only
@@ -1427,7 +1503,10 @@ export class BatchScheduler {
     }
     const out: LayerInner[] = [];
     for (const inner of inners) {
-      if (inner instanceof BatchedRotatingCache) {
+      if (isBatchableCache(inner)) {
+        inner.filterRows(keep);
+        out.push(inner);
+      } else if (inner instanceof BatchedRotatingCache) {
         inner.filter(keep); // in-place (mutates + drops rows + reduces padding)
         out.push(inner);
       } else if (inner instanceof BatchedRotatingQuantCache) {

package/src/serve/generation-gateway.ts

@@ -20,14 +20,12 @@
 // 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 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
+// Batchable gate (the rest → serial): batch>1 AND every model cache is a
+// dynamic-B-capable built-in type or implements BatchableCache (mirrors mlx-lm
+// 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
 // (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
@@ -39,7 +37,7 @@ 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 } from "../model/gemma4-base";
+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";
@@ -187,7 +185,13 @@ export class GenerationGateway {
 
   /** True if `--batch N` (N>1) is on (batchability is then per-request). */
   get batchingEnabled(): boolean {
-    return this.#batch > 1;
+    return this.batchMode === "batch";
+  }
+
+  /** Truthful configured/model-capability mode for stats and discovery. */
+  get batchMode(): "off" | "serial" | "batch" {
+    if (this.#batch <= 1) return "off";
+    return this.#modelCachesBatchable() ? "batch" : "serial";
   }
 
   /** Rows currently decoding in the batch (0 if no scheduler / idle). */
@@ -262,6 +266,7 @@ export class GenerationGateway {
         (c) =>
           c instanceof KVCache ||
           c instanceof RotatingKVCache ||
+          isBatchableCache(c) ||
           (ssmOk && c instanceof SSMCache),
       );
       for (const c of proto) c.dispose();
@@ -364,7 +369,16 @@ export class GenerationGateway {
     onToken: OnToken,
     vision: Vision | undefined,
     shape: RequestShape,
+    signal?: AbortSignal,
   ): Promise<GenerateStats> {
+    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();
+      throw e;
+    }
     if (!this.willBatch(shape)) {
       // Serial decode is an unbroken microtask chain (FFI + generator
       // resumes): without a periodic macrotask hop the event loop starves for
@@ -377,14 +391,26 @@ 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;
         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);
+        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 this.runExclusive(() => this.serialRun(promptIds, options, hoppingOnToken, vision));
+      return this.runExclusive(async () => {
+        signal?.throwIfAborted();
+        const stats = await this.serialRun(promptIds, options, hoppingOnToken, vision);
+        signal?.throwIfAborted();
+        return stats;
+      })
+        // 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
@@ -447,6 +473,7 @@ export class GenerationGateway {
           processors.length === 0 &&
           !options.grammar,
         onToken,
+        ...(signal ? { signal } : {}),
         // B1: pass the per-row grammar controller through. The scheduler drives
         // accept/ready/terminate; this gateway OWNS disposal (finally below)
         // across resolve, reject, eviction, and the whole-batch-drop error path.
@@ -458,9 +485,7 @@ export class GenerationGateway {
     } finally {
       history?.dispose();
       history = null;
-      // The controller's WASM matcher/compiled-grammar are per-request; free
-      // them on every exit path. (generate()'s finally does this on the serial
-      // lane; the batch lane owns it here since the scheduler never does.)
+      // The scheduler never owns per-row grammar state.
       options.grammar?.dispose();
     }