mlx-bun 0.4.0 → 0.5.0

package/src/serve/generation-gateway.ts

@@ -94,11 +94,12 @@ class AsyncMutex {
   }
 }
 
-/** Embeddings-prefill payload for media prompts (vision and/or audio — the
- *  name predates audio). Any request carrying one routes to the serial lane
- *  (shape.hasVision) and prefills through generate()'s promptEmbeddings
- *  path, bypassing the prompt cache. */
+/** Prepared embeddings for vision/audio prompts. Qualified models consume
+ * this through shared preparation and decode; other bindings retain serial
+ * execution. Token-only prefix caching remains disabled for media. */
 export type Vision = {
+  /** Producer identity includes all media and preceding rendered tokens. */
+  prefixIdentity?: string;
   embeddings: MlxArray;
   /** bool [L] image-token mask for the bidirectional attention overlay.
    *  Absent when the prompt carries ANY audio — audio(-containing) prompts
@@ -109,7 +110,7 @@ export type Vision = {
    *  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). */
+   *  owned by shared forward input, or scoped to the explicit serial run. */
   mrope?: import("../model/qwen3-mrope").MropeRequestState;
 };
 
@@ -318,6 +319,19 @@ export class GenerationGateway {
     try { return await fn(); } finally { lease.dispose(); }
   }
 
+  /** Preparation reads model weights without replacing active model state.
+   * A compatible backend can perform it between existing decode iterations. */
+  runPreparation<T>(work: () => Promise<T>, signal?: AbortSignal): Promise<T> {
+    if (this.#closed) return Promise.reject(new Error("generation gateway is closed"));
+    if (signal?.aborted) return Promise.reject(signal.reason);
+    const group = this.batchingEnabled ? this.#ensureScheduler() : undefined;
+    return group?.runPreparation
+      ? group.runPreparation(work, signal)
+      : this.runExclusive(work, undefined, signal);
+  }
+
+  get mediaBatchingEnabled(): boolean { return this.batchingEnabled && !!this.#binding.mediaInput; }
+
   /** 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(
@@ -331,10 +345,32 @@ export class GenerationGateway {
     trace?: PromptResponseTrace,
   ): Promise<GenerateStats> {
     let reservation;
+    const closeRequestWait = trace?.begin("engine.request_wait", { capacity: this.#batch });
     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(); }
+    finally { closeRequestWait?.(); }
+    let releasePrefix: (() => void) | undefined;
+    try {
+      try {
+        const adapters = options.adapters?.length
+          ? this.opts.adapterNamespace?.(options.adapters) ?? JSON.stringify(options.adapters) : "";
+        const baseNamespace = this.#binding.prefixNamespace?.(placement.execution, options, adapters) ??
+          (placement.execution?.method === "speculative" ? null : adapters);
+        const namespace = baseNamespace !== null && vision?.prefixIdentity
+          ? JSON.stringify(["prepared-prefix-v1", vision.prefixIdentity, baseNamespace]) : baseNamespace;
+        if ((!vision || placement.execution?.promptCache) && namespace !== null) {
+          const closePrefetch = this.opts.promptCache?.prefetch ? trace?.begin("cache.prefetch") : undefined;
+          try { releasePrefix = await this.opts.promptCache?.prefetch?.(promptIds, namespace, options.cacheSessionId); }
+          finally { closePrefetch?.(); }
+        }
+      } catch (error) { cleanupFailure(error, () => disposeUnstartedRequest(options,
+        placement.mechanism === "continuous" ? undefined : vision)); }
+      return await this.#run(promptIds, options, onToken, vision, shape, placement, signal, trace);
+    } finally {
+      disposeResources([{ dispose: () => releasePrefix?.() }, reservation,
+        ...(placement.mechanism === "continuous" && vision
+          ? [vision.embeddings, vision.imageMask, vision.multimodalMask].filter(value => value != null) : [])]);
+    }
   }
 
   async #run(
@@ -347,7 +383,8 @@ export class GenerationGateway {
     signal?: AbortSignal,
     trace?: PromptResponseTrace,
   ): Promise<GenerateStats> {
-    const disposeUnstarted = () => disposeUnstartedRequest(options, vision);
+    const disposeUnstarted = () => disposeUnstartedRequest(options,
+      placement.mechanism === "continuous" ? undefined : vision);
     if (placement.shape !== shape) {
       disposeUnstarted();
       throw new Error("generation placement does not belong to this request shape");
@@ -410,8 +447,10 @@ export class GenerationGateway {
 
     const adapters = options.adapters?.length ? [...options.adapters] : undefined;
     const adapterKey = adapters ? JSON.stringify(adapters) : "";
-    const cacheNamespace = adapters
-      ? () => this.opts.adapterNamespace?.(adapters) ?? adapterKey : "";
+    const adapterNamespace = () => adapters ? this.opts.adapterNamespace?.(adapters) ?? adapterKey : "";
+    const cacheNamespace = vision?.prefixIdentity
+      ? () => JSON.stringify(["prepared-prefix-v1", vision.prefixIdentity, adapterNamespace()])
+      : adapters ? adapterNamespace : "";
     const context = adapters
       ? this.#binding.bindAdapterContext?.(adapters, `adapters:${adapterKey}`) : undefined;
 
@@ -433,7 +472,8 @@ export class GenerationGateway {
     });
     try {
       st = await this.#ensureScheduler().submit({
-        promptIds, context, cacheNamespace, prefillChunkSize: options.prefillChunkSize,
+        promptIds, context, cacheNamespace, cacheSessionId: options.cacheSessionId, prefillChunkSize: options.prefillChunkSize,
+        ...(vision ? { promptInput: this.#binding.mediaInput!(vision) } : {}),
         continuation: continuation?.continuation,
         statePolicy: this.#binding.statePolicy?.(placement.execution, options, promptIds.length + (options.maxTokens ?? 512)),
         compiledDecode: placement.execution?.compiledDecode,
@@ -470,6 +510,7 @@ export class GenerationGateway {
       decodeTps: st.decodeMs > 0 && st.generatedTokens > 1 ? ((st.generatedTokens - 1) / st.decodeMs) * 1000 : 0,
       cacheTokens: [],
       ...(st.spec ? { spec: st.spec } : {}),
+      ...(st.fill ? { fill: st.fill } : {}),
     };
   }
 

package/src/serve/model-binding.ts

@@ -21,9 +21,14 @@ export interface BuiltPrompt {
   diffusionPixels: MlxArray | null;
 }
 
+/** Enter the model's execution domain for native preparation. Fetching and
+ * container transcoding stay outside this callback. */
+export type PromptNativeWork = <T>(work: () => Promise<T>) => Promise<T>;
+
 export interface ModelPromptBuilder {
   (body: ChatRequestParams, tools: ChatRequestParams["tools"] | null,
-    ownership: RequestOwnership, prep: RequestPrep): Promise<BuiltPrompt>;
+    ownership: RequestOwnership, prep: RequestPrep, nativeWork?: PromptNativeWork,
+    objects?: import("../contracts/object-cache").ObjectCache<import("../backends/mlx/checkpoint-state").CheckpointAttachment[]>): Promise<BuiltPrompt>;
 }
 
 /** Model-owned native services. A replacement supplies this one binding;
@@ -35,6 +40,7 @@ export interface ModelServingBinding {
   createSerial(services: MlxSerialServices): SerialRun;
   readonly buildPrompt: ModelPromptBuilder;
   restore(store: SsdCacheStore, entry: SsdIndexEntry): ReturnType<SsdCacheStore["restore"]>;
+  restoreAsync?(store: SsdCacheStore, entry: SsdIndexEntry): ReturnType<SsdCacheStore["restoreAsync"]>;
   signal(promptIds: number[], bins: number, minimum: number): Promise<{ bins: number[]; vocab: number }>;
   diagnostics(): Record<string, unknown>;
   readonly discovery: { readonly adapters: boolean; readonly training: boolean;

package/src/serve/preparation.ts

@@ -17,11 +17,11 @@ export async function acquireReservation(pool: Pick<AdmissionPool, "acquire">, s
   finally { signal?.removeEventListener("abort", abort); }
 }
 
-/** One retained media preparation at a time; grammar rows may form a batch. */
+/** Bound retained preparations independently from the execution queue. */
 export function createPreparationExecutor(
-  run: PreparationExecutor["run"], grammarCapacity: number,
+  run: PreparationExecutor["run"], grammarCapacity: number, mediaCapacity = 1,
 ): PreparationExecutor & { close(): void } {
-  const media = new AdmissionPool(1);
+  const media = new AdmissionPool(Math.max(1, mediaCapacity));
   const constraints = new AdmissionPool(Math.max(1, grammarCapacity));
   return {
     run,

package/src/serve/prompt-response-trace.ts

@@ -10,16 +10,27 @@ export type P2RTracePhase =
   | "request.prompt_prepare"
   | "completion.total"
   | "completion.placement"
+  | "engine.request_wait"
   | "engine.admission_wait"
+  | "engine.mixed_forward"
   | "cache.lookup_restore"
+  | "cache.prefetch"
   | "prefill.total"
   | "prefill.batch_setup"
   | "prefill.chunk"
+  | "prefill.row_sync"
+  | "prefill.forward"
+  | "prefill.evaluate"
+  | "prefill.checkpoint"
+  | "prefill.project"
+  | "prefill.complete"
+  | "prefill.companion"
   | "prefill.kv_maintenance"
   | "token_zero.total"
   | "token_zero.forward"
   | "token_zero.head"
   | "token_zero.sample"
+  | "response.token_route"
   | "response.first_write"
   | "response.final_write";
 
@@ -42,6 +53,8 @@ export interface P2RTraceRecord {
   route: string;
   clock: "monotonic-ms";
   outcome: P2RTraceOutcome;
+  /** Process-local monotonic origin; aligns concurrent requests in one server. */
+  startedAtMs?: number;
   totalMs: number;
   events: P2RTraceEvent[];
   attributes?: P2RTraceAttributes;
@@ -129,6 +142,7 @@ export class PromptResponseTrace {
       route: this.route,
       clock: "monotonic-ms",
       outcome,
+      startedAtMs: this.#startedAt,
       totalMs,
       events: [...this.#events].sort((a, b) => a.startMs - b.startMs),
       ...(attributes ? { attributes } : {}),

package/src/serve/request-plan.ts

@@ -39,6 +39,7 @@ export interface RequestPlanInput {
   topLogprobs: number;
   adapterIds: string[];
   hasVision: boolean;
+  hasPreparedPrefixIdentity?: boolean;
   userSeed: boolean;
   hasGrammar: boolean;
   hasDraft: boolean;
@@ -137,6 +138,7 @@ export function planRequest(input: RequestPlanInput): PlanRequestResult {
 
   const shape: RequestShape = {
     hasVision: input.hasVision,
+    hasPreparedPrefixIdentity: input.hasPreparedPrefixIdentity,
     hasAdapters: !!options.adapters?.length,
     hasRepetitionPenalty: !!options.repetitionPenalty,
     hasLogitsExtras: !!(

package/src/serve/request-prep.ts

@@ -5,7 +5,7 @@
 // tool-call stream router. Resolved ONCE per server (createRequestPrep) and
 // captured by the route handlers. Extracted from src/server.ts (repo-taming
 // Phase 4).
-import type { ToolDefinition } from "../chat-template";
+import type { ChatMessage, ToolDefinition } from "../chat-template";
 import {
   fillEchoConfig,
   fillMaxSpan,
@@ -27,6 +27,7 @@ import {
   type ChatRequestParams,
 } from "./chat-request";
 import type { ServingContext } from "./model-host";
+import type { PromptTokenHistory } from "./generated-token-history";
 import { selectToolStreamMode, ToolAwareStream, type ToolStreamMode } from "./token-streams";
 
 /** The ServerOptions fields request preparation reads (server-wide
@@ -49,6 +50,7 @@ export function createRequestPrep(input: {
   /** `--max-tokens` / the GLM memory plan's generation cap; undefined = the
    *  surface's own default (chat 65 536, raw completion 512). */
   defaultGeneratedTokens: number | undefined;
+  tokenHistory?: PromptTokenHistory;
 }) {
   const { ctx, serverOptions, kvScheme, defaultGeneratedTokens } = input;
 
@@ -92,6 +94,8 @@ export function createRequestPrep(input: {
       ? Math.min(genTemp, NO_THINK_TEMPERATURE)
       : genTemp;
     return {
+      cacheSessionId: typeof req.session_id === "string" ? req.session_id :
+        typeof req.prompt_cache_key === "string" ? req.prompt_cache_key : undefined,
       // Unset means generate until EOS/stop. Only a client cap or an explicit
       // operator limit constrains completion length; fit estimates do not.
       maxTokens: req.max_completion_tokens ?? req.max_tokens ??
@@ -191,7 +195,8 @@ export function createRequestPrep(input: {
     const ids = ctx.tokenizer.encode(rendered);
     // template includes <bos>; tokenizer post-processor also prepends one
     const trimmed = ids[0] === ids[1] && ids[0] === ctx.tokenizer.bosTokenId ? ids.slice(1) : ids;
-    return { ids: trimmed, startInThinking: promptEndsInOpenThink(rendered) };
+    return { ids: input.tokenHistory?.resolve(rendered, trimmed) ?? trimmed,
+      startInThinking: promptEndsInOpenThink(rendered) };
   };
 
   /** STABLE cache boundary: the prompt's tail is a generation primer (e.g.
@@ -219,18 +224,32 @@ export function createRequestPrep(input: {
     if (primer !== undefined) return Math.max(0, trimmed.length - primer);
     let stableLen = trimmed.length;
     try {
-      const probe = ctx.tokenizer.encode(
-        ctx.template.render(
-          [...normalizeMessages(req.messages), { role: "assistant", content: "x" }],
-          opts,
-        ),
-      );
-      const probeTrimmed =
-        probe[0] === probe[1] && probe[0] === ctx.tokenizer.bosTokenId ? probe.slice(1) : probe;
-      let i = 0;
-      const n = Math.min(trimmed.length, probeTrimmed.length);
-      while (i < n && trimmed[i] === probeTrimmed[i]) i++;
-      stableLen = i;
+      const canonical = ctx.tokenizer.encode(ctx.template.render(normalizeMessages(req.messages), opts));
+      const current = canonical[0] === canonical[1] && canonical[0] === ctx.tokenizer.bosTokenId
+        ? canonical.slice(1) : canonical;
+      const replies: ChatMessage[] = [{ role: "assistant", content: "x" }];
+      // Some templates retain an empty think primer for text replies but
+      // omit it before a tool call. Both reply forms must preserve the
+      // boundary retained by an untrimmable recurrent cache.
+      if (tools?.length) replies.push({ role: "assistant", content: null, tool_calls: [{
+        id: "cache-boundary-probe", type: "function",
+        function: { name: tools[0]!.function.name, arguments: {} },
+      }] });
+      let i = current.length;
+      for (const reply of replies) {
+        const probe = ctx.tokenizer.encode(ctx.template.render([
+          ...normalizeMessages(req.messages), reply, { role: "user", content: "x" },
+        ], opts));
+        const probeTrimmed = probe[0] === probe[1] && probe[0] === ctx.tokenizer.bosTokenId ? probe.slice(1) : probe;
+        let common = 0;
+        const n = Math.min(i, probeTrimmed.length);
+        while (common < n && current[common] === probeTrimmed[common]) common++;
+        i = common;
+      }
+      // Generated-history provenance may use a different BPE segmentation
+      // inside old turns. Measure the primer canonically, then subtract it
+      // from the actual request IDs that inference will receive.
+      stableLen = Math.max(0, trimmed.length - (current.length - i));
       // Memoize the mode's primer length (sanity-capped: a "primer" longer
       // than 64 tokens means the probe diverged for content reasons —
       // don't generalize that).
@@ -271,10 +290,10 @@ export function createRequestPrep(input: {
    *  when the feature is off or the request's shape refuses it.
    *
    *  Refusals owned here (body-level); the ones only ChatStage can see —
-   *  a compiled grammar, media prompts, a mounted draft model, a quantized KV
-   *  scheme — are applied there. Continuous (batch) placement needs no
-   *  refusal: generate() is the only fill site, so a batched request simply
-   *  does not fill.
+   *  a compiled grammar, media prompts and a mounted draft model — are applied
+   *  there. The execution binding owns cache-format and method support. Strict
+   *  shared execution keeps the sampler's absolute token positions, including
+   *  known tokens; an explicit seed therefore remains replayable.
    *
    *  MLX_BUN_FILL=echo additionally arms the echo index (K3c, Lab tier): the
    *  strict rows keep policy "assert", and copied spans ride the same
@@ -287,9 +306,7 @@ export function createRequestPrep(input: {
     const mode = resolveFillMode();
     if (mode === "off") return null;
     if (!tools?.length) return null;
-    // A user-fixed seed means reproducibility: the sampler's step index would
-    // skip the injected positions, so an identical request could not replay.
-    if (req.seed !== undefined) return null;
+    // Seed support is selected by the execution binding after placement.
     // Injected tokens are never sampled — they have no logprob row.
     if (req.logprobs === true) return null;
     if (typeof req.top_logprobs === "number" && req.top_logprobs > 0) return null;

package/src/serve/transcription-server.ts

@@ -0,0 +1,64 @@
+// Whisper-only server: `mlx-bun serve <whisper checkpoint>` starts this
+// instead of the chat server. Routes: the audio routes, /v1/models,
+// /v1, /health, /stats. No chat machinery, no prompt cache — the process
+// idles at a few tens of MB with the weights paged out (see
+// TranscriptionService). This is the footprint sotto wants from a
+// dictation backend.
+
+import type { Server } from "bun";
+import pkgJson from "../../package.json" with { type: "json" };
+import { handleAudioRoute } from "./audio-routes";
+import { TranscriptionService } from "./transcription-service";
+
+const pkgVersion = (pkgJson as { version: string }).version;
+
+export interface TranscriptionServerOptions {
+  modelDir: string;
+  modelId: string;
+  port?: number;
+  hostname?: string;
+  unixSocket?: string;
+  idleUnloadSec?: number;
+  resident?: boolean;
+  /** Load the weights before serving (default: on first request). */
+  preload?: boolean;
+}
+
+export async function createTranscriptionServer(opts: TranscriptionServerOptions): Promise<{ server: Server<unknown>; service: TranscriptionService }> {
+  const service = new TranscriptionService({
+    modelDir: opts.modelDir, modelId: opts.modelId, idleUnloadSec: opts.idleUnloadSec, resident: opts.resident,
+  });
+  if (opts.preload) await service.ensureLoaded();
+  const startedAt = Date.now();
+  const server = Bun.serve({
+    ...(opts.unixSocket
+      ? ({ unix: opts.unixSocket } as unknown as Record<string, never>)
+      : { port: opts.port ?? 0, ...(opts.hostname ? { hostname: opts.hostname } : {}) }),
+    idleTimeout: 0,
+    async fetch(request) {
+      const url = new URL(request.url);
+      const audio = await handleAudioRoute(url, request, { service: async () => service });
+      if (audio) return audio;
+      if (request.method === "GET" && url.pathname === "/health")
+        return Response.json({ status: "ok", transcription: { ...service.stats, sessions: service.sessionCount } });
+      if (request.method === "GET" && url.pathname === "/stats")
+        return Response.json({ model: opts.modelId, transcription: service.stats, uptime_s: (Date.now() - startedAt) / 1000 });
+      if (request.method === "GET" && url.pathname === "/v1")
+        return Response.json({
+          name: "mlx-bun", version: pkgVersion, model: opts.modelId, mode: "transcription",
+          endpoints: ["POST /v1/audio/transcriptions", "POST /v1/audio/translations", "POST /v1/audio/sessions", "POST /v1/audio/sessions/:id/audio", "POST /v1/audio/sessions/:id/finish", "DELETE /v1/audio/sessions/:id", "POST /admin/transcription/unload", "GET /v1/models", "GET /health", "GET /stats"],
+        });
+      if (request.method === "GET" && (url.pathname === "/v1/models" || url.pathname.startsWith("/v1/models/")))
+        return Response.json({
+          object: "list",
+          data: [{
+            id: opts.modelId, object: "model", created: Math.floor(startedAt / 1000), owned_by: "mlx-bun",
+            transcription: true, resident: service.resident,
+            capabilities: { transcription: true, translation: true, chat_completions: false },
+          }],
+        });
+      return Response.json({ error: { message: "not found" } }, { status: 404 });
+    },
+  });
+  return { server, service };
+}