mlx-bun 0.3.0 → 0.5.0

package/src/sampler.ts

@@ -9,8 +9,11 @@
 // and the XTC-off path is unchanged.)
 
 import { Dtype } from "./mlx/ffi";
+import { normalizedArgmax } from "./mlx/normalized-argmax";
 import { MlxArray } from "./mlx/array";
 import * as ops from "./mlx/ops";
+import type { TokenLogprobs } from "./contracts/generation";
+import type { SampledToken, SamplingSession } from "./contracts/sampling";
 import { applyCurve, type CurveParams } from "./lab/curve/curve-sampler";
 
 /** Resolved HLG sampling config (the user-facing knobs). The mid gain is NOT
@@ -120,10 +123,7 @@ export interface StepExtras {
   topVals: MlxArray | null;
 }
 
-export interface StepSample<T> {
-  token: T;
-  extras: StepExtras | null;
-}
+export type StepSample<T> = SampledToken<T, StepExtras>;
 
 interface StepGrammar {
   readonly isTerminated: boolean;
@@ -161,25 +161,8 @@ export interface NumberStepSamplerConfig extends StepSamplerConfigBase {
   grammarWait: "before-sample";
 }
 
-export interface DeviceStepSampler {
-  readonly isPlainGreedy: boolean;
-  readonly needsHistory: boolean;
-  sample(logits: MlxArray, step: number): StepSample<MlxArray>;
-  seedHistory(tokens: readonly number[]): void;
-  commitDevice(token: MlxArray): void;
-  commitNumbers(tokens: readonly number[]): void;
-  dispose(): void;
-}
-
-export interface NumberStepSampler {
-  readonly isPlainGreedy: boolean;
-  readonly needsHistory: boolean;
-  sample(logits: MlxArray, step: number): Promise<StepSample<number>>;
-  seedHistory(tokens: readonly number[]): void;
-  commitDevice(token: MlxArray): void;
-  commitNumbers(tokens: readonly number[]): void;
-  dispose(): void;
-}
+export type DeviceStepSampler = SamplingSession<MlxArray, MlxArray, StepSample<MlxArray>>;
+export type NumberStepSampler = SamplingSession<MlxArray, MlxArray, Promise<StepSample<number>>>;
 
 const GOLDEN = 0x9e3779b97f4a7c15n;
 
@@ -786,7 +769,7 @@ export function makeLogitsProcessors(opts: LogitsProcessorOptions = {}): LogitsP
   return out;
 }
 
-/** Whether a whole [B,V] row can use the scheduler's vectorized argmax path. */
+/** Whether scores can use stateless greedy sampling across independent rows. */
 export function isPlainGreedy(
   opts: StepSamplerOptions,
   hasProcessors = makeLogitsProcessors(opts).length > 0,
@@ -809,6 +792,23 @@ export function disposeStepExtras(extras: StepExtras | null): void {
   extras.topVals?.dispose();
 }
 
+/** Read captured probabilities with their token, without adding casts behind
+ * the next decode step. Consumes the capture on success and failure. */
+export function readStepExtras(extras: StepExtras | null): TokenLogprobs | undefined {
+  if (!extras) return undefined;
+  try {
+    const out: TokenLogprobs = {};
+    if (extras.sel) out.logprob = extras.sel.toFloat32Host()[0]!;
+    if (extras.topIdx && extras.topVals) {
+      const ids = extras.topIdx.toIntTokens();
+      const vals = extras.topVals.toFloat32Host();
+      out.top = Array.from(ids, (id, i) => ({ id, logprob: vals[i]! }))
+        .sort((a, b) => b.logprob - a.logprob);
+    }
+    return out;
+  } finally { disposeStepExtras(extras); }
+}
+
 export function makeStepSampler(
   options: StepSamplerOptions,
   config: DeviceStepSamplerConfig,
@@ -827,6 +827,7 @@ export function makeStepSampler(
   const captureSelected = config.captureSelectedLogprob === true;
   const captureTop = Math.max(0, config.captureTopLogprobs ?? 0);
   const capture = captureSelected || captureTop > 0;
+  const greedyWithoutMetadata = (options.temperature ?? 0) === 0 && !options.curve && !capture && !config.sampler;
   let history: MlxArray | null = null;
 
   const seedHistory = (tokens: readonly number[]): void => {
@@ -889,6 +890,8 @@ export function makeStepSampler(
         current = next;
         ownsCurrent = current !== input;
       }
+      if (greedyWithoutMetadata)
+        return { token: normalizedArgmax(current), extras: null };
       logprobs = toLogprobs(current);
       if (ownsCurrent) current.dispose();
       current = input;
@@ -941,7 +944,10 @@ export function makeStepSampler(
   if (config.initialHistory) seedHistory(config.initialHistory);
 
   const common = {
+    independent: isPlainGreedy(options, processors.length > 0) && !capture && !config.sampler
+      ? independentGreedySampling : undefined,
     isPlainGreedy: isPlainGreedy(options, processors.length > 0),
+    capturesLogprobs: capture,
     needsHistory: processors.length > 0,
     seedHistory,
     commitDevice,
@@ -989,7 +995,13 @@ export function makeStepSampler(
   };
 }
 
-/** logits [1, V] → logprobs [1, V] (logits - logsumexp). */
+/** Stateless sampling across rows or verification positions. Preserve the
+ * normalized-score argmax, including rounding-created ties. */
+export const independentGreedySampling = Object.freeze({
+  sample: normalizedArgmax,
+});
+
+/** logits [..., V] → logprobs [..., V] (logits - logsumexp). */
 export function toLogprobs(logits: MlxArray): MlxArray {
   const lse = ops.logsumexpAxis(logits, -1, true);
   const out = ops.sub(logits, lse);

package/src/serve/audio-routes.ts

@@ -0,0 +1,323 @@
+// OpenAI-compatible speech-to-text routes:
+//   POST /v1/audio/transcriptions   (task transcribe)
+//   POST /v1/audio/translations     (task translate → English)
+//   POST /admin/transcription/unload (explicit page-out of the Whisper weights)
+//   POST /v1/audio/sessions · POST /v1/audio/sessions/:id/audio ·
+//   POST /v1/audio/sessions/:id/finish · DELETE /v1/audio/sessions/:id
+//   (streaming dictation: chunks arrive during capture, each finished 30 s
+//   window transcribes immediately, finish returns the whole transcript)
+// Bodies: multipart/form-data (OpenAI: file, model, language, prompt,
+// response_format, temperature, stream, timestamp_granularities[]) or JSON
+// ({ file: <base64 | data: URL>, ...same fields }). Non-standard fields:
+// beam_size, condition_on_previous_text, no_speech_threshold, vocabulary
+// (array of hint terms fitted into the prompt budget; usage reported),
+// faithful (oracle graph instead of the fast path), audio_ctx (Lab), vad /
+// vad_threshold / vad_min_speech_ms / vad_trim (Silero gate: silence returns
+// an empty transcript without running Whisper).
+// Streaming (`stream=true`): SSE `transcript.text.delta` per segment,
+// then `transcript.text.done` carrying the final body.
+
+import { formatTranscription, TRANSCRIPTION_FORMATS, toVerboseJson, type TranscriptionFormat } from "../audio/whisper-format";
+import type { WhisperSegment } from "../audio/whisper-transcribe";
+import { TranscriptionError, TranscriptionService, type TranscriptionOutcome, type TranscriptionParams } from "./transcription-service";
+
+export type AudioRoute =
+  | { kind: "transcriptions"; task: "transcribe" | "translate" }
+  | { kind: "unload" }
+  | { kind: "session-create" }
+  | { kind: "session-audio"; id: string }
+  | { kind: "session-finish"; id: string }
+  | { kind: "session-delete"; id: string };
+
+export function matchAudioRoute(method: string, pathname: string): AudioRoute | null {
+  switch (`${method} ${pathname}`) {
+    case "POST /v1/audio/transcriptions": return { kind: "transcriptions", task: "transcribe" };
+    case "POST /v1/audio/translations": return { kind: "transcriptions", task: "translate" };
+    case "POST /admin/transcription/unload": return { kind: "unload" };
+    case "POST /v1/audio/sessions": return { kind: "session-create" };
+  }
+  if (pathname.startsWith("/v1/audio/sessions/")) {
+    const rest = pathname.slice("/v1/audio/sessions/".length).split("/");
+    const id = decodeURIComponent(rest[0] ?? "");
+    if (!id) return null;
+    if (method === "POST" && rest[1] === "audio" && rest.length === 2) return { kind: "session-audio", id };
+    if (method === "POST" && rest[1] === "finish" && rest.length === 2) return { kind: "session-finish", id };
+    if (method === "DELETE" && rest.length === 1) return { kind: "session-delete", id };
+  }
+  return null;
+}
+
+/** Everything the route needs from the host: a service (or a reason there is none). */
+export interface AudioRouteHost {
+  service(): Promise<TranscriptionService | null>;
+  /** Message for the 503 when no Whisper model is available. */
+  unavailableMessage?: string;
+}
+
+export interface ParsedAudioRequest {
+  audio: Uint8Array;
+  format: TranscriptionFormat;
+  stream: boolean;
+  params: TranscriptionParams;
+}
+
+const sessionFormats = new Map<string, TranscriptionFormat>();
+
+const segmentJson = (s: WhisperSegment) => ({ id: s.id, seek: s.seek, start: s.start, end: s.end, text: s.text, tokens: s.tokens, temperature: s.temperature, avg_logprob: s.avgLogprob, compression_ratio: s.compressionRatio, no_speech_prob: s.noSpeechProb, ...(s.words ? { words: s.words } : {}) });
+
+function errorJson(message: string, status: number, type = "invalid_request_error"): Response {
+  return Response.json({ error: { message, type } }, { status });
+}
+
+const B64_MAX = 200 * 1024 * 1024;
+
+function decodeBase64Audio(v: unknown): Uint8Array | null {
+  if (typeof v !== "string" || v.length === 0 || v.length > B64_MAX * 1.4) return null;
+  const data = v.startsWith("data:") ? v.slice(v.indexOf(",") + 1) : v;
+  try {
+    return new Uint8Array(Buffer.from(data, "base64"));
+  } catch {
+    return null;
+  }
+}
+
+function num(v: unknown): number | null {
+  if (v === null || v === undefined || v === "") return null;
+  const n = typeof v === "number" ? v : Number(v);
+  return Number.isFinite(n) ? n : null;
+}
+
+function bool(v: unknown): boolean | null {
+  if (typeof v === "boolean") return v;
+  if (v === "true" || v === "1") return true;
+  if (v === "false" || v === "0") return false;
+  return null;
+}
+
+/** Decoding params shared by one-shot requests and sessions (no audio). */
+function parseParams(fields: Record<string, unknown>, task: "transcribe" | "translate"): TranscriptionParams | Response {
+  const language = typeof fields.language === "string" && fields.language.trim() ? fields.language.trim() : null;
+  const temperature = num(fields.temperature);
+  if (temperature !== null && (temperature < 0 || temperature > 1)) return errorJson("temperature must be in [0, 1]", 400);
+  const beamSize = num(fields.beam_size);
+  if (beamSize !== null && (!Number.isInteger(beamSize) || beamSize < 1 || beamSize > 16))
+    return errorJson("beam_size must be an integer in [1, 16]", 400);
+  if (temperature !== null && temperature > 0 && beamSize !== null && beamSize > 1)
+    return errorJson("beam_size requires temperature 0", 400);
+  const vocabulary = Array.isArray(fields.vocabulary)
+    ? (fields.vocabulary as unknown[]).map(String)
+    : typeof fields.vocabulary === "string" && fields.vocabulary.trim()
+      ? fields.vocabulary.split(",").map((s) => s.trim()).filter(Boolean)
+      : undefined;
+  if (vocabulary && (vocabulary.length > 8192 || vocabulary.some((t) => t.length > 16384)))
+    return errorJson("vocabulary: at most 8192 terms of at most 16 KiB each", 400);
+  const noSpeech = num(fields.no_speech_threshold);
+  const vadOn = bool(fields.vad) ?? false;
+  const vadThreshold = num(fields.vad_threshold);
+  if (vadThreshold !== null && (vadThreshold <= 0 || vadThreshold >= 1)) return errorJson("vad_threshold must be in (0, 1)", 400);
+  const vadMinSpeech = num(fields.vad_min_speech_ms);
+  const audioCtx = num(fields.audio_ctx);
+  if (audioCtx !== null && (!Number.isInteger(audioCtx) || audioCtx < 64 || audioCtx > 1500))
+    return errorJson("audio_ctx must be an integer in [64, 1500]", 400);
+  const granularities = Array.isArray(fields.timestamp_granularities) ? (fields.timestamp_granularities as string[]) : [];
+  return {
+    task, language,
+    prompt: typeof fields.prompt === "string" && fields.prompt.trim() ? fields.prompt : null,
+    vocabulary, temperature, beamSize,
+    conditionOnPreviousText: bool(fields.condition_on_previous_text) ?? true,
+    ...(noSpeech !== null ? { noSpeechThreshold: noSpeech } : {}),
+    withoutTimestamps: bool(fields.without_timestamps) ?? false,
+    vad: vadOn ? { threshold: vadThreshold ?? undefined, minSpeechMs: vadMinSpeech ?? undefined, trim: bool(fields.vad_trim) ?? false } : null,
+    wordTimestamps: granularities.includes("word"),
+    faithful: bool(fields.faithful) ?? false,
+    audioCtx,
+  };
+}
+
+/** Parse a multipart or JSON transcription request into audio bytes + params. */
+export async function parseAudioRequest(request: Request, task: "transcribe" | "translate"): Promise<ParsedAudioRequest | Response> {
+  const ct = request.headers.get("content-type") ?? "";
+  let fields: Record<string, unknown> = {};
+  let audio: Uint8Array | null = null;
+  if (ct.includes("multipart/form-data")) {
+    let form: FormData;
+    try { form = await request.formData(); }
+    catch (e) { return errorJson(`invalid multipart body: ${(e as Error).message}`, 400); }
+    const file = form.get("file");
+    if (!(file instanceof Blob)) return errorJson("`file` (audio upload) is required", 400);
+    audio = new Uint8Array(await file.arrayBuffer());
+    for (const [k, v] of form.entries()) {
+      if (k === "file") continue;
+      if (k.endsWith("[]")) {
+        const key = k.slice(0, -2);
+        (fields[key] ??= [] as string[]);
+        (fields[key] as string[]).push(String(v));
+      } else fields[k] = typeof v === "string" ? v : String(v);
+    }
+  } else if (ct.includes("application/json")) {
+    try { fields = (await request.json()) as Record<string, unknown>; }
+    catch { return errorJson("invalid JSON body", 400); }
+    const raw = fields.file ?? (fields.input_audio as { data?: unknown } | undefined)?.data ?? fields.audio;
+    audio = decodeBase64Audio(raw);
+    if (!audio) return errorJson("`file` must be base64 audio (optionally a data: URL)", 400);
+  } else {
+    return errorJson("content-type must be multipart/form-data or application/json", 415);
+  }
+  if (audio.length === 0) return errorJson("empty audio upload", 400);
+  const format = (fields.response_format as string | undefined) ?? "json";
+  if (!TRANSCRIPTION_FORMATS.includes(format as TranscriptionFormat))
+    return errorJson(`response_format must be one of ${TRANSCRIPTION_FORMATS.join(", ")}`, 400);
+  const params = parseParams(fields, task);
+  if (params instanceof Response) return params;
+  const stream = bool(fields.stream) ?? false;
+  if (params.wordTimestamps && format !== "verbose_json" && !stream)
+    return errorJson("timestamp_granularities[]=word requires response_format=verbose_json", 400);
+  return { audio, format: format as TranscriptionFormat, stream, params };
+}
+
+/** Body for a finished transcription in the requested format. */
+export function transcriptionResponse(outcome: TranscriptionOutcome, format: TranscriptionFormat, task: string): Response {
+  const extra = {
+    model: outcome.modelId,
+    language: outcome.result.language,
+    duration: outcome.durationSeconds,
+    timings: outcome.timings,
+    ...(outcome.vocabulary ? { vocabulary: outcome.vocabulary } : {}),
+    ...(outcome.vad ? { vad: outcome.vad } : {}),
+  };
+  const headers: Record<string, string> = {
+    "x-mlx-bun-language": outcome.result.language,
+    "x-mlx-bun-timings": JSON.stringify(outcome.timings),
+  };
+  if (format === "json")
+    return Response.json({ text: outcome.result.text.trim(), mlx_bun: extra }, { headers });
+  if (format === "verbose_json")
+    return Response.json({ ...toVerboseJson(outcome.result, outcome.durationSeconds, task), mlx_bun: extra }, { headers });
+  const body = formatTranscription(outcome.result, format, outcome.durationSeconds, task);
+  const type = format === "srt" ? "application/x-subrip" : format === "vtt" ? "text/vtt" : "text/plain";
+  return new Response(body, { headers: { ...headers, "content-type": `${type}; charset=utf-8` } });
+}
+
+export async function handleAudioRoute(url: URL, request: Request, host: AudioRouteHost): Promise<Response | null> {
+  const route = matchAudioRoute(request.method, url.pathname);
+  if (!route) return null;
+  const service = await host.service();
+  if (!service)
+    return errorJson(
+      host.unavailableMessage ??
+        "no Whisper model is available — download one (mlx-bun get mlx-community/whisper-large-v3-turbo) or pass --whisper-model",
+      503, "model_unavailable",
+    );
+  if (route.kind === "unload") {
+    const unloaded = service.unload();
+    return Response.json({ unloaded, ...service.stats });
+  }
+  if (route.kind === "session-create") {
+    let fields: Record<string, unknown> = {};
+    if ((request.headers.get("content-type") ?? "").includes("application/json")) {
+      try { fields = (await request.json()) as Record<string, unknown>; }
+      catch { return errorJson("invalid JSON body", 400); }
+    }
+    const params = parseParams(fields, fields.task === "translate" ? "translate" : "transcribe");
+    if (params instanceof Response) return params;
+    const format = (fields.response_format as string | undefined) ?? "json";
+    if (!["json", "verbose_json"].includes(format)) return errorJson("sessions return json or verbose_json", 400);
+    try {
+      const session = await service.createSession(params);
+      sessionFormats.set(session.id, format as TranscriptionFormat);
+      return Response.json({ id: session.id, model: service.modelId, vocabulary: session.vocabulary ?? null });
+    } catch (e) {
+      return transcriptionFailure(e);
+    }
+  }
+  if (route.kind === "session-audio" || route.kind === "session-finish" || route.kind === "session-delete") {
+    const session = service.session(route.id);
+    if (!session) return errorJson("unknown transcription session", 404);
+    if (route.kind === "session-delete") {
+      session.close();
+      sessionFormats.delete(route.id);
+      return new Response(null, { status: 204 });
+    }
+    if (route.kind === "session-audio") {
+      const ct = request.headers.get("content-type") ?? "";
+      const bytes = new Uint8Array(await request.arrayBuffer());
+      if (bytes.length === 0) return errorJson("empty audio chunk", 400);
+      let pcm: Float32Array;
+      try {
+        if (ct.startsWith("audio/pcm") || ct.startsWith("audio/l32") || ct.includes("float32")) {
+          if (bytes.length % 4 !== 0) return errorJson("float32 PCM chunk length must be a multiple of 4", 400);
+          pcm = new Float32Array(bytes.buffer.slice(bytes.byteOffset, bytes.byteOffset + bytes.byteLength));
+          const rate = Number(/rate=(\d+)/.exec(ct)?.[1] ?? 16000);
+          if (rate !== 16000) return errorJson("float32 PCM chunks must be 16 kHz mono", 400);
+        } else pcm = await (service.constructor as typeof TranscriptionService).decodeAudio(bytes);
+      } catch (e) {
+        return transcriptionFailure(e);
+      }
+      try {
+        const r = await session.append(pcm);
+        return Response.json({ id: session.id, samples: session.samples, duration: session.durationSeconds, speech: r.speech, segments: r.segments.map(segmentJson) });
+      } catch (e) {
+        return transcriptionFailure(e);
+      }
+    }
+    try {
+      const outcome = await session.finish();
+      sessionFormats.delete(route.id);
+      return transcriptionResponse(outcome, sessionFormats.get(route.id) ?? "json", session.params.task ?? "transcribe");
+    } catch (e) {
+      sessionFormats.delete(route.id);
+      return transcriptionFailure(e);
+    }
+  }
+  const parsed = await parseAudioRequest(request, route.task);
+  if (parsed instanceof Response) return parsed;
+  const taskName = route.task;
+  if (!parsed.stream) {
+    try {
+      const outcome = await service.transcribe(parsed.audio, { ...parsed.params, signal: request.signal });
+      return transcriptionResponse(outcome, parsed.format, taskName);
+    } catch (e) {
+      return transcriptionFailure(e);
+    }
+  }
+  // Streaming: one SSE event per finalized segment, then the final body.
+  const encoder = new TextEncoder();
+  const stream = new ReadableStream<Uint8Array>({
+    start: async (controller) => {
+      const send = (event: string, data: unknown) =>
+        controller.enqueue(encoder.encode(`event: ${event}\ndata: ${JSON.stringify(data)}\n\n`));
+      try {
+        const outcome = await service.transcribe(parsed.audio, {
+          ...parsed.params,
+          signal: request.signal,
+          onSegment: (s: WhisperSegment) => {
+            if (s.text) send("transcript.text.delta", { delta: s.text, segment: { id: s.id, start: s.start, end: s.end } });
+          },
+          onProgress: (done, total) => send("transcript.progress", { done, total }),
+        });
+        const body = parsed.format === "verbose_json"
+          ? toVerboseJson(outcome.result, outcome.durationSeconds, taskName)
+          : { text: outcome.result.text.trim() };
+        send("transcript.text.done", {
+          ...body,
+          mlx_bun: { model: outcome.modelId, language: outcome.result.language, duration: outcome.durationSeconds, timings: outcome.timings, ...(outcome.vocabulary ? { vocabulary: outcome.vocabulary } : {}) },
+        });
+      } catch (e) {
+        send("error", { message: (e as Error).message });
+      } finally {
+        controller.close();
+      }
+    },
+  });
+  return new Response(stream, {
+    headers: { "content-type": "text/event-stream", "cache-control": "no-cache", connection: "keep-alive" },
+  });
+}
+
+function transcriptionFailure(e: unknown): Response {
+  if (e instanceof TranscriptionError) return errorJson(e.message, e.status);
+  const message = (e as Error).message ?? String(e);
+  if (/cancelled/.test(message)) return errorJson("request cancelled", 499);
+  if (/WAV|afconvert|transcode|audio|RIFF/i.test(message)) return errorJson(`undecodable audio: ${message}`, 400);
+  return errorJson(message, 500, "server_error");
+}

package/src/serve/chat-request.ts

@@ -7,6 +7,8 @@ import type { HlgConfig } from "../sampler";
 import { RequestError } from "./pipeline";
 
 export interface ChatRequestParams {
+  session_id?: string;
+  prompt_cache_key?: string;
   messages: ChatMessage[];
   stream?: boolean;
   max_tokens?: number;

package/src/serve/chat-stage.ts

@@ -12,7 +12,7 @@ import {
 } from "./chat-request";
 import type { InferenceRequest } from "./inference-request";
 import type { ServingContext } from "./model-host";
-import type { BuiltPrompt, ModelPromptBuilder } from "./model-binding";
+import type { BuiltPrompt, ModelPromptBuilder, PromptNativeWork } from "./model-binding";
 import { modelPromptBuilder } from "../backends/mlx/model-serving";
 import { RequestError } from "./pipeline";
 import { RequestOwnership } from "./request-plan";
@@ -23,9 +23,9 @@ export class ChatStage {
   constructor(
     private readonly ctx: ServingContext,
     private readonly prep: RequestPrep,
-    private readonly promptCache: Pick<PromptCache, "peekPrefixLen">,
+    private readonly promptCache: Pick<PromptCache, "peekPrefixLen"> & Partial<Pick<PromptCache, "objects">>,
     /** Admission ceiling (fit() / the GLM memory plan), resolved once. */
-    private readonly maxSafeContext: number,
+    private readonly contextLimit: number | null,
     /** `serve --adapter <dir>` startup default; a request's explicit
      *  `adapter` (incl. "none") wins. */
     private readonly defaultAdapter?: string,
@@ -40,9 +40,14 @@ export class ChatStage {
         message.content.some((part) => part.type !== "text"));
     const native = media || body.response_format != null || !!body.guided_grammar ||
       !!body.guided_regex || !!body.guided_choice?.length || body.structured_outputs != null;
+    const nativeWork: PromptNativeWork = async (work) => {
+      signal?.throwIfAborted();
+      const entered = async () => { signal?.throwIfAborted(); return work(); };
+      return this.preparation ? this.preparation.run(entered, signal) : entered();
+    };
     const build = async () => {
       signal?.throwIfAborted();
-      const result = await this.runPrepared(request, requestId);
+      const result = await this.runPrepared(request, requestId, nativeWork);
       if (signal?.aborted) {
         result.plan.ownership.dispose();
         signal.throwIfAborted();
@@ -52,14 +57,15 @@ export class ChatStage {
     if (!native || !this.preparation) return build();
     let reservation = await this.preparation.reserve?.(media ? "media" : "constraint", signal);
     try {
-      const result = await this.preparation.run(build, signal);
+      const result = await build();
       if (reservation) result.plan.ownership.retain(reservation);
       reservation = undefined;
       return result;
     } finally { reservation?.dispose(); }
   }
 
-  private async runPrepared(request: ChatRequest, requestId?: string): Promise<InferenceRequest> {
+  private async runPrepared(request: ChatRequest, requestId: string | undefined,
+    nativeWork: PromptNativeWork): Promise<InferenceRequest> {
     const { ctx, prep } = this;
     let body = request.params;
     const id = requestId ?? `chatcmpl-${crypto.randomUUID()}`;
@@ -83,7 +89,7 @@ export class ChatStage {
       !!body.guided_regex || !!body.guided_choice?.length ||
       body.structured_outputs != null;
     if (grammarReq) {
-      const g = await prep.compileGrammarForRequest(body);
+      const g = await nativeWork(() => prep.compileGrammarForRequest(body));
       grammarCtrl = ownership.own(g.controller);
       if (!g.controller && g.degradeHint) {
         const degraded = applyGrammarDegrade(body, g.degradeHint);
@@ -94,7 +100,7 @@ export class ChatStage {
 
     let built: BuiltPrompt;
     try {
-      built = await this.buildPrompt(body, tools, ownership, prep);
+      built = await this.buildPrompt(body, tools, ownership, prep, nativeWork, this.promptCache.objects);
     } catch (e) {
       if (e instanceof RequestError) return reject(e.status, e.message);
       return reject(400, `prompt build failed: ${(e as Error).message}`);
@@ -131,13 +137,11 @@ export class ChatStage {
     // Token fast-forwarding (K3, MLX_BUN_FILL=strict): the determined-span
     // table for this request's tools. Refused for the shapes only this stage
     // can see — a compiled grammar (forced tokens are its job), media prompts
-    // (embeddings prefill / mRoPE), a mounted draft model (the spec loop is a
-    // different executor and never reads options.fill), and quantized-KV
-    // schemes (post-conversion multi-token append is unvalidated). The
+    // (embeddings prefill / mRoPE). Provider and cache-format support
+    // are declared by the engine's method binding. The
     // body-level refusals live in prep.fillPlanFor.
     let fillSession: import("../fill/fill-session").FillSession | null = null;
-    if (!grammarCtrl && !vision && !diffusionPixels && !ctx.draft &&
-        !options.kvBits && !options.kvConfig?.length && !options.turboQuant) {
+    if (!grammarCtrl && !vision && !diffusionPixels) {
       fillSession = prep.fillPlanFor(body, tools, promptIds);
       if (fillSession) options.fill = fillSession;
     }
@@ -165,12 +169,13 @@ export class ChatStage {
         promptIds,
         options,
         requestedMaxTokens,
-        maxSafeContext: this.maxSafeContext,
+        contextLimit: this.contextLimit,
         stream: request.stream,
         wantLogprobs,
         topLogprobs,
         adapterIds,
         hasVision: !!vision,
+        hasPreparedPrefixIdentity: !!vision?.prefixIdentity,
         userSeed: body.seed !== undefined,
         hasGrammar: !!grammarCtrl,
         hasDraft: !!ctx.draft,

package/src/serve/completion-executor.ts

@@ -280,9 +280,17 @@ export class CompletionExecutor {
       control.signal?.throwIfAborted();
       planned.transferOwnership();
 
+      // Bounded diagnostic of hidden markers versus visible events. No token
+      // readbacks are added: these IDs already crossed the sampling interface.
+      let traceTokens = control.trace ? 8 : 0;
       const consumeToken: OnToken = (token, tokenLogprobs) => {
         logprobs?.push(token, tokenLogprobs);
-        const pushed = sink.push(token);
+        const route = traceTokens ? { token, index: 8 - traceTokens, events: 0 } : undefined;
+        const closeRoute = route ? control.trace!.begin("response.token_route", route) : undefined;
+        let pushed: ReturnType<CompletionSink["push"]>;
+        try { pushed = sink.push(token); if (route) route.events = pushed.events.length; }
+        finally { closeRoute?.(); }
+        if (traceTokens) traceTokens = pushed.events.length ? 0 : traceTokens - 1;
         const observed = pushed.events.length
           ? control.onEvents?.(pushed.events)
           : undefined;

package/src/serve/discovery-routes.ts

@@ -30,10 +30,16 @@ export interface DiscoveryRoutes {
   invalidateLibrary(): void;
 }
 
+export interface TranscriptionInfo {
+  id: string;
+  resident: boolean;
+}
+
 export function createDiscoveryRoutes(
   ctx: ServingContext,
   gateway: Pick<GenerationGateway, "batchMode">,
   startedAt: number,
+  transcription: () => Promise<TranscriptionInfo | null> = async () => null,
 ): DiscoveryRoutes {
   let libraryCache: { at: number; rows: unknown[] } | null = null;
 
@@ -106,6 +112,8 @@ export function createDiscoveryRoutes(
               "POST /v1/messages",
               "POST /v1/responses",
               "POST /v1/embeddings",
+              "POST /v1/audio/transcriptions",
+              "POST /v1/audio/translations",
               "GET /v1/models",
               "GET/POST/DELETE /v1/adapters",
               "GET /health",
@@ -132,6 +140,7 @@ export function createDiscoveryRoutes(
             top_k: ctx.genDefaults.topK ?? null,
           };
           const capabilities = modelServingBinding(ctx).discovery;
+          const stt = await transcription();
           const data: Array<Record<string, unknown>> = [{
             id: ctx.modelId,
             object: "model",
@@ -165,9 +174,16 @@ export function createDiscoveryRoutes(
               audio: !!(ctx.audio || ctx.loadAudio),
               adapters: capabilities.adapters,
               training: capabilities.training,
+              transcription: stt !== null,
             },
             gen_defaults: genDefaults,
           }];
+          if (stt)
+            data.push({
+              id: stt.id, object: "model", created, owned_by: "mlx-bun",
+              transcription: true, resident: stt.resident,
+              capabilities: { transcription: true, translation: true, chat_completions: false },
+            });
           try {
             const { Registry, visionCapable } = await import("../registry");
             const { supportTier } = await import("../model/support");
@@ -175,7 +191,7 @@ export function createDiscoveryRoutes(
             try {
               if (registry.list().length === 0) await registry.scan();
               for (const model of registry.listCanonical()) {
-                if (model.repoId === ctx.modelId) continue;
+                if (model.repoId === ctx.modelId || model.repoId === stt?.id) continue;
                 const tier = supportTier(model.modelType, model.repoId);
                 if (tier === null) continue;
                 data.push({