mlx-bun 0.0.7 → 0.0.9

package/src/server.ts

@@ -23,9 +23,11 @@ const pkgVersion = (pkgJson as { version: string }).version;
 import { loadModelConfig, type KvQuantSpec, type ModelConfig } from "./config";
 import { Weights } from "./weights";
 import { Gemma4Model } from "./model/gemma4";
+import { DiffusionGemmaModel } from "./model/diffusion-gemma";
+import { spliceImageTokens } from "./vision/diffusion-vision";
 import { createModel, type RuntimeModel } from "./model/factory";
-import { isMiniCPM5Config, isQwen35Config, isSupportedModelRecord } from "./model/support";
-import { generate, type GenerateOptions } from "./generate";
+import { isMiniCPM5Config, isSupportedModelRecord } from "./model/support";
+import { generate, type GenerateOptions, type TokenLogprobs } from "./generate";
 import type { HlgConfig } from "./sampler";
 import { isMonotone, CURVE_UMIN, type CurveParams } from "./curve-sampler";
 const CURVE_PAGE = curveDesignerHtml as unknown as string;
@@ -49,6 +51,7 @@ import {
   type ResponsesRequest,
 } from "./responses";
 import { AdapterManager, listAvailableAdapters } from "./lora";
+import { embedMany, isEmbeddingModel } from "./embed";
 import { fit } from "./fit";
 import { setMemoryLimit } from "./mlx/ffi";
 import { VisionTower } from "./vision/embedder";
@@ -83,8 +86,10 @@ export interface ServerOptions {
    *  it). Exposed at /stats.server so other mlx-bun processes can
    *  warn before attaching to a server that may vanish. */
   owner?: "serve" | "pi-session" | "embedded";
-  /** Interface to bind (Bun.serve hostname). Default: Bun's default
-   *  (all interfaces); pass "127.0.0.1" for loopback-only. */
+  /** Interface to bind (Bun.serve hostname). Unset ⇒ Bun's default
+   *  (all interfaces) — embedded/library use. The CLI always passes one:
+   *  "127.0.0.1" (loopback, mlx_lm.server parity) unless --host says
+   *  otherwise. */
   hostname?: string;
   /** Server-wide default for the chat template's `enable_thinking`
    *  variable (MiniCPM5/CPM and other hybrid-reasoning models). A
@@ -99,6 +104,10 @@ export interface ServerOptions {
   defaultTemperature?: number;
   defaultTopP?: number;
   defaultTopK?: number;
+  /** Completion cap when the request omits max_tokens (`--max-tokens`).
+   *  mlx_lm.server's flag; its default there is 512 — ours stays 65,536 so
+   *  thinking traces never truncate. `--max-tokens 512` = mlx-lm behavior. */
+  defaultMaxTokens?: number;
   /** Max concurrent requests batched through the mlx-lm-parity engine
    *  (`--batch N`). Default 1 = today's serialized single-queue path. >1 opts
    *  the WHOLE server into the batched engine (continuous batching, B floats
@@ -109,6 +118,11 @@ export interface ServerOptions {
   /** HLG tone-curve sampling default (set via --hlg-sampling on + sub-knobs).
    *  A per-request `hlg` object overrides it field-by-field. Off when unset. */
   hlg?: HlgConfig;
+  /** Adapter id mounted at startup via `serve --adapter <dir>` (mlx-lm's
+   *  `--adapter-path`). Used when a request sends no `adapter` field; a
+   *  request's explicit `adapter` (including "none") always wins. Hot-swap
+   *  via /v1/adapters is unchanged. */
+  defaultAdapter?: string;
 }
 
 export interface ServerContext {
@@ -260,6 +274,34 @@ interface ChatRequest {
   top_k?: number;
   seed?: number;
   repetition_penalty?: number;
+  /** mlx-lm extension: recent-token window for repetition_penalty
+   *  (default 20; 0 = whole history, Python `[-0:]` semantics). */
+  repetition_context_size?: number;
+  /** min-p sampling (mlx_lm.server's `min_p`): keep tokens whose probability
+   *  is ≥ min_p · p(top token). 0 = off. */
+  min_p?: number;
+  /** XTC sampling (mlx_lm.server names): with probability `xtc_probability`
+   *  per step, remove every token above `xtc_threshold` except the least
+   *  likely of them. EOS + the newline token are always exempt (the server
+   *  injects them as xtc special tokens, matching mlx_lm.server). */
+  xtc_probability?: number;
+  xtc_threshold?: number;
+  /** OpenAI logit_bias: {tokenId: additive bias}. JSON object keys arrive as
+   *  strings; coerced to int keys / float values like mlx-lm (400 on failure). */
+  logit_bias?: Record<string, number>;
+  /** OpenAI presence/frequency penalties + mlx-lm's context-size extensions
+   *  (window of recent tokens the penalty looks at; default 20). */
+  presence_penalty?: number;
+  presence_context_size?: number;
+  frequency_penalty?: number;
+  frequency_context_size?: number;
+  /** mlx_lm.server logprobs: `logprobs` is a BOOL (even on /v1/completions —
+   *  not OpenAI's legacy int), `top_logprobs` an int in [0, 11] or the -1
+   *  "unset" sentinel (server.py validates exactly that; OpenAI's cap is 20,
+   *  mlx-lm's is 11 — we copy the reference). Non-stream responses carry
+   *  mlx-lm's logprobs block; stream chunks never do (reference behavior). */
+  logprobs?: boolean;
+  top_logprobs?: number;
   /** OpenAI stop sequences: plain string or array (spec allows up to 4).
    *  Matched on DECODED text, not token ids — see StopMatcher. */
   stop?: string | string[];
@@ -288,6 +330,15 @@ interface ChatRequest {
   };
 }
 
+/** POST /v1/completions body (mlx_lm.server's raw text completion — no chat
+ *  template). Sampling/penalty/stop fields are the same names as ChatRequest;
+ *  `prompt` replaces `messages`. mlx_lm.server accepts only a STRING prompt
+ *  (it calls `tokenizer.encode(request.prompt)` directly) — token-array
+ *  prompts are rejected there too, so we match. No `echo` (mlx-lm has none). */
+type TextCompletionRequest = Omit<ChatRequest, "messages" | "tools" | "tool_choice"> & {
+  prompt?: unknown;
+};
+
 /** Per-field default HLG knobs when enabling without specifying them. */
 const HLG_DEFAULTS = { width: 4, shoulder: 4, toe: 6, pivotOffset: 6 } as const;
 
@@ -317,7 +368,23 @@ interface OpenAIToolCall {
   function: { name: string; arguments: string };
 }
 
-type ToolStreamMode = "gemma-sentinel" | "plain" | "buffered-text";
+export type ToolStreamMode = "gemma-sentinel" | "plain" | "buffered-text";
+
+/** Pick the stream router for a model family. The token-id sentinel router is
+ *  Gemma-4-ONLY: ids 48/49 (<|tool_call>/<tool_call|>) and 100/101
+ *  (<|channel>/<channel|>) are special tokens of that tokenizer family
+ *  (src/tool-call.ts). On every other tokenizer — MiniCPM5, Qwen3/3.5, and
+ *  the Tier-0 generics (llama, qwen2, phi3, …) — those ids are ordinary
+ *  low-id vocab, so the sentinel router would silently swallow output into a
+ *  phantom tool/reasoning segment. Everyone else parses tool calls from
+ *  decoded text (buffered-text; parseGeneratedToolCalls covers the
+ *  OpenAI-JSON <tool_call>, Qwen <function=…>, and MiniCPM5 <function name=…>
+ *  shapes) when tools are present, and streams plain otherwise. Models whose
+ *  markup isn't covered fail soft: the markup stays in content. */
+export function selectToolStreamMode(modelType: string, hasTools: boolean): ToolStreamMode {
+  if (modelType.startsWith("gemma4")) return "gemma-sentinel";
+  return hasTools ? "buffered-text" : "plain";
+}
 
 /** Routes generated tokens. Gemma uses family-specific sentinel token ids;
  *  MiniCPM5 and other text-template models use decoded-text parsing so
@@ -727,6 +794,129 @@ export function normalizeMessages(messages: ChatMessage[]): ChatMessage[] {
   });
 }
 
+/** Coerce a wire `logit_bias` ({"tokenId": bias}) to numeric keys/values —
+ *  mlx_lm.server's `{int(k): float(v)}` coercion; throws its exact error
+ *  message on anything non-numeric (surfaced as a 400). JSON object keys
+ *  always arrive as strings, hence the coercion. */
+export function parseLogitBias(
+  raw: Record<string, number> | undefined | null,
+): Record<number, number> | undefined {
+  if (raw == null) return undefined;
+  if (typeof raw !== "object" || Array.isArray(raw))
+    throw new Error("logit_bias must be a dict of int to float");
+  const out: Record<number, number> = {};
+  for (const [k, v] of Object.entries(raw)) {
+    const id = Number(k);
+    const bias = Number(v);
+    if (!Number.isInteger(id) || !Number.isFinite(bias))
+      throw new Error("logit_bias must be a dict of int to float");
+    out[id] = bias;
+  }
+  return Object.keys(out).length ? out : undefined;
+}
+
+/** Default seed for a request that didn't pin one. `Date.now()` alone is NOT
+ *  request-unique: under `--batch N` the batch lane serves ONLY default-seed
+ *  requests (explicit-seed requests route serial, see GenerationGateway.willBatch),
+ *  so two identical prompts arriving in the same millisecond would share a seed
+ *  and — with per-row RNG keyed as stepKey(seed, generatedCount) — produce
+ *  byte-identical completions, silently collapsing best-of-N diversity. Mix a
+ *  per-process Weyl counter (golden-ratio increment, period 2^32) into the
+ *  timestamp so every call yields a distinct uint32 within any given ms.
+ *  Determinism contract unchanged: reproducibility is only promised for an
+ *  EXPLICIT request seed (`req.seed ?? nextDefaultSeed()` — explicit wins,
+ *  byte-identical to before); a default seed is fresh entropy per request. */
+let seedWeyl = 0;
+export function nextDefaultSeed(): number {
+  seedWeyl = (seedWeyl + 0x9e3779b9) >>> 0;
+  return ((Date.now() & 0xffffffff) ^ seedWeyl) >>> 0;
+}
+
+/** mlx_lm.server's logprobs request validation, copied exactly (server.py
+ *  APIHandler.validate_model_parameters: `_validate("logprobs", bool)` and
+ *  `_validate("top_logprobs", int, min_val=0, max_val=11, whitelist=[-1])`
+ *  with defaults logprobs=False / top_logprobs=-1). Returns the reference's
+ *  error message (→ 400) or null when valid. Note mlx-lm caps top_logprobs
+ *  at 11, not OpenAI's 20 — we mirror the reference. */
+export function validateLogprobsParams(body: {
+  logprobs?: unknown;
+  top_logprobs?: unknown;
+}): string | null {
+  const lp = body.logprobs ?? false;
+  if (typeof lp !== "boolean") return "logprobs must be of type bool";
+  const tl = body.top_logprobs ?? -1;
+  if (typeof tl !== "number" || !Number.isInteger(tl))
+    return "top_logprobs must be of type int";
+  if (tl === -1) return null; // the "unset" whitelist sentinel
+  if (tl < 0) return "top_logprobs must be at least 0";
+  if (tl > 11) return "top_logprobs must be at most 11";
+  return null;
+}
+
+/** Collects per-token logprob info and shapes mlx_lm.server's response block
+ *  (server.py generate_response L1317-1327). NOT OpenAI's shape — entries
+ *  carry token *ids* (and raw vocab token strings in the top-k form), and the
+ *  SAME block is attached under choices[0].logprobs for chat AND text
+ *  completions:
+ *  - top_logprobs > 0 → {content: [{id, token, logprob,
+ *      top_logprobs: [{id, token, logprob}, …]}, …]} — each entry is the
+ *      top-1 candidate merged with its own top-k list (`dict(i[0],
+ *      top_logprobs=i)`); mlx-lm leaves argpartition order unspecified, we
+ *      sort descending, so the entry is deterministically the argmax.
+ *  - logprobs=true (and top_logprobs ≤ 0) → {content: [{id, logprob}, …]}
+ *      with the SAMPLED token's logprob.
+ *  When both are set, only the top_logprobs form is emitted (the reference's
+ *  if/elif). Stream chunks never carry logprobs — mlx-lm's streaming
+ *  generate_response calls pass no token_logprobs/top_tokens. */
+class LogprobsCollector {
+  readonly #tokens: number[] = [];
+  readonly #tokenLogprobs: number[] = [];
+  readonly #topTokens: { id: number; token: string; logprob: number }[][] = [];
+
+  constructor(
+    private readonly wantLogprobs: boolean,
+    private readonly topK: number,
+    private readonly idToToken: (id: number) => string,
+  ) {}
+
+  get active(): boolean {
+    return this.wantLogprobs || this.topK > 0;
+  }
+
+  push(token: number, info?: TokenLogprobs): void {
+    if (!this.active) return;
+    this.#tokens.push(token);
+    if (this.wantLogprobs) this.#tokenLogprobs.push(info?.logprob ?? NaN);
+    if (this.topK > 0)
+      this.#topTokens.push(
+        (info?.top ?? []).map((t) => ({
+          id: t.id,
+          token: this.idToToken(t.id),
+          logprob: t.logprob,
+        })),
+      );
+  }
+
+  /** choices[0].logprobs value, or null when nothing to attach (mirrors the
+   *  reference: the key is omitted entirely for zero collected tokens). */
+  payload(): { content: Record<string, unknown>[] } | null {
+    if (this.#topTokens.length)
+      return {
+        content: this.#topTokens.map((t) =>
+          t.length ? { ...t[0]!, top_logprobs: t } : {},
+        ),
+      };
+    if (this.#tokenLogprobs.length)
+      return {
+        content: this.#tokens.map((id, i) => ({
+          id,
+          logprob: this.#tokenLogprobs[i]!,
+        })),
+      };
+    return null;
+  }
+}
+
 /** Incremental detokenizer: emits the longest stable decoded prefix. */
 class StreamDecoder {
   #ids: number[] = [];
@@ -786,13 +976,6 @@ export function createServer(
   // lane (see GenerationGateway.willBatch).
   const batch = Math.max(1, Math.floor(serverOptions.batch ?? 1));
 
-  let queue: Promise<unknown> = Promise.resolve();
-  const enqueue = <T>(fn: () => Promise<T>): Promise<T> => {
-    const run = queue.then(fn, fn);
-    queue = run.catch(() => {});
-    return run;
-  };
-
   const promptCache = new PromptCache(serverOptions.promptCacheBytes ?? 2e9);
   // Responses-API store for previous_response_id resumption (Phase 11):
   // TTL + byte-capped LRU, port of optiq/response_store.py. Pairs with
@@ -800,8 +983,11 @@ export function createServer(
   // so its KV prefill is already cached.
   const responseStore = new ResponseStore();
 
-  /** Run one generation with prompt-cache reuse. Must be called inside
-   *  the queue. onToken returning `false` halts generation early (stop
+  /** Run one generation with prompt-cache reuse. Must be called under the
+   *  gateway's exclusive lock — the gateway invokes it on the serial lane;
+   *  the curve /generate endpoint wraps it in gateway.runExclusive. (The old
+   *  `enqueue` promise queue is gone; the gateway's AsyncMutex is THE lock.)
+   *  onToken returning `false` halts generation early (stop
    *  sequence fired); the cache snapshot stays valid and is still kept.
    *  Vision requests bypass the prompt cache: image tokens are
    *  identical placeholder ids, so prefix matching across different
@@ -809,7 +995,7 @@ export function createServer(
   const runGeneration = async (
     promptIds: number[],
     options: GenerateOptions,
-    onToken: (token: number) => void | boolean | Promise<void | boolean>,
+    onToken: (token: number, logprobs?: TokenLogprobs) => void | boolean | Promise<void | boolean>,
     vision?: { embeddings: import("./mlx/array").MlxArray; imageMask: import("./mlx/array").MlxArray },
   ) => {
     // Cache entries are adapter-specific: KV computed under one adapter
@@ -824,7 +1010,7 @@ export function createServer(
         ...(vision ? { promptEmbeddings: vision.embeddings, imageMask: vision.imageMask } : {}),
       });
       for await (const t of gen) {
-        if ((await onToken(t.token)) === false) break;
+        if ((await onToken(t.token, t.logprobs)) === false) break;
       }
       const s = gen.stats!; // set on completion AND on early break
       if (vision) {
@@ -839,6 +1025,7 @@ export function createServer(
     } finally {
       vision?.embeddings.dispose();
       vision?.imageMask.dispose();
+      options.visionPixels?.dispose();
     }
   };
 
@@ -911,6 +1098,17 @@ export function createServer(
         `Omit --kv-quant to batch in bf16. (docs/design/parallel-slots.md)`,
     );
 
+  // XTC never removes EOS or the newline token: mlx_lm.server passes
+  // [tokenizer.eos_token_id, tokenizer.encode("\n")] as xtc_special_tokens.
+  // We pass the flat equivalent — ALL configured EOS ids plus the bare newline
+  // id(s), encoded WITHOUT special tokens (mlx-lm's encode("\n") can drag a
+  // BOS along on some tokenizers; protecting BOS from XTC is a no-op, so the
+  // flat/no-BOS form is behaviorally identical and cleaner).
+  const xtcSpecialTokens = [
+    ...ctx.model.config.eosTokenIds,
+    ...ctx.tokenizer.encode("\n", false),
+  ];
+
   // Effective enable_thinking for a request, with the same precedence the chat
   // template uses (extracted so sampling and template rendering can't disagree):
   // explicit chat_template_kwargs.enable_thinking → reasoning_effort ("none" =
@@ -944,12 +1142,25 @@ export function createServer(
       // truncates the visible reply. Default very generously (the model's
       // context is far larger); only an explicit max_tokens narrows it. The
       // model still stops at its eos_token well before this in normal replies.
-      maxTokens: req.max_completion_tokens ?? req.max_tokens ?? 65_536,
+      maxTokens: req.max_completion_tokens ?? req.max_tokens ?? serverOptions.defaultMaxTokens ?? 65_536,
       temperature: req.temperature ?? serverOptions.defaultTemperature ?? defaultTemp,
       topP: req.top_p ?? serverOptions.defaultTopP ?? ctx.genDefaults.topP ?? 0,
       topK: req.top_k ?? serverOptions.defaultTopK ?? ctx.genDefaults.topK ?? 0,
-      seed: req.seed ?? (Date.now() & 0xffffffff),
+      seed: req.seed ?? nextDefaultSeed(),
       repetitionPenalty: req.repetition_penalty ?? ctx.genDefaults.repetitionPenalty,
+      // mlx_lm.server sampling extensions (defaults mirror server.py: min_p /
+      // xtc off, context windows 20). parseLogitBias throws on non-numeric
+      // input — callers surface it as a 400 (mlx-lm's coercion error).
+      minP: req.min_p ?? 0,
+      xtcProbability: req.xtc_probability ?? 0,
+      xtcThreshold: req.xtc_threshold ?? 0,
+      ...(req.xtc_probability ? { xtcSpecialTokens } : {}),
+      logitBias: parseLogitBias(req.logit_bias),
+      repetitionContextSize: req.repetition_context_size ?? 20,
+      presencePenalty: req.presence_penalty ?? 0,
+      presenceContextSize: req.presence_context_size ?? 20,
+      frequencyPenalty: req.frequency_penalty ?? 0,
+      frequencyContextSize: req.frequency_context_size ?? 20,
       hlg: resolveHlg(req.hlg, serverOptions.hlg),
       // generate() yields token ids; stop sequences match on decoded text
       // (they can span token boundaries), so the StopMatcher sits at the
@@ -977,15 +1188,11 @@ export function createServer(
     return { ids: trimmed, startInThinking: promptEndsInOpenThink(rendered) };
   };
 
-  const toolStreamMode = (tools: ToolDefinition[] | null): ToolStreamMode => {
-    // MiniCPM5 and Qwen3.5 emit tool calls as DECODED TEXT (Qwen:
-    // <tool_call><function=name><parameter=…>; both handled by
-    // parseGeneratedToolCalls), not Gemma's token sentinels — so they take the
-    // buffered-text path. Gemma keeps the token-level sentinel path.
-    if (isMiniCPM5Config(ctx.model.config) || isQwen35Config(ctx.model.config))
-      return tools?.length ? "buffered-text" : "plain";
-    return "gemma-sentinel";
-  };
+  const toolStreamMode = (tools: ToolDefinition[] | null): ToolStreamMode =>
+    // Gemma-4 keeps the token-level sentinel path; every other family
+    // (MiniCPM5, Qwen3/3.5, Tier-0 generics) emits tool calls as DECODED TEXT
+    // — see selectToolStreamMode for why sentinel ids must stay family-gated.
+    selectToolStreamMode(ctx.model.config.modelType, !!tools?.length);
 
   const toolRouter = (tools: ToolDefinition[] | null): ToolAwareStream =>
     new ToolAwareStream(ctx.tokenizer, toolStreamMode(tools), tools);
@@ -1064,9 +1271,14 @@ export function createServer(
           // surface on their own. (scan() is INSERT-OR-REPLACE + prunes deleted,
           // so it's idempotent and cheap for a local cache.)
           await reg.scan();
+          const { visionCapable } = await import("./registry");
+          const { supportTier } = await import("./model/support");
           const rows = [];
-          for (const m of reg.list()) {
-            const supported = isSupportedModelRecord(m.modelType, m.repoId);
+          // listCanonical: one row per repo (refs/main) — duplicate snapshots
+          // from upstream re-pushes stay visible only in `ls --all-revisions`.
+          for (const m of reg.listCanonical()) {
+            const tier = supportTier(m.modelType, m.repoId);
+            const supported = tier !== null;
             let assessment = null;
             try {
               const config = await loadModelConfig(m.path);
@@ -1080,7 +1292,7 @@ export function createServer(
             rows.push({
               repo_id: m.repoId, model_type: m.modelType,
               size_bytes: m.sizeBytes, quant_bits: m.quantBits,
-              vision: m.hasVisionSidecar, supported,
+              vision: visionCapable(m), supported, support_tier: tier,
               serving: m.repoId === ctx.modelId,
               assessment,
             });
@@ -1095,6 +1307,39 @@ export function createServer(
         return Response.json({ downloads: downloadsSnapshot() });
       }
 
+      // Memory synthesis progress (P8-T5). The same DAG the nightly launchd job
+      // runs (`mlx-bun memory synthesize`), streamed as Server-Sent Events so the
+      // status page can show live stage/log/done progress. `?dry=1` plans the DAG
+      // without any model call or vault write — the safe wiring-verification path
+      // (the FULL-corpus run is USER-ACTION, P6-T5). GET so EventSource can drive it.
+      if (url.pathname === "/v1/memory/synthesize" && request.method === "GET") {
+        const dryRun = url.searchParams.get("dry") === "1";
+        const { runSynthesis } = await import("./memory/pipeline");
+        const enc = new TextEncoder();
+        const stream = new ReadableStream({
+          async start(controller) {
+            const send = (e: unknown) =>
+              controller.enqueue(enc.encode(`data: ${JSON.stringify(e)}\n\n`));
+            try {
+              const summary = await runSynthesis({ dryRun }, (ev) => send(ev));
+              send({ type: "summary", ...summary });
+              controller.enqueue(enc.encode("data: [DONE]\n\n"));
+            } catch (e) {
+              send({ type: "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",
+          },
+        });
+      }
+
       if (url.pathname === "/fit" && request.method === "GET") {
         // Fit assessment for the status page: this-machine report at the
         // admission ceiling + the Apple SKU matrix at a fixed 32k.
@@ -1157,8 +1402,9 @@ export function createServer(
         return Response.json({
           name: "mlx-bun", version: pkgVersion, model: ctx.modelId,
           endpoints: [
-            "POST /v1/chat/completions", "POST /v1/messages", "POST /v1/responses",
-            "GET /v1/models", "GET/POST/DELETE /v1/adapters",
+            "POST /v1/chat/completions", "POST /v1/completions", "POST /v1/messages",
+            "POST /v1/responses", "POST /v1/embeddings", "GET /v1/models",
+            "GET/POST/DELETE /v1/adapters", "GET /health",
             "GET /stats", "GET /fit", "GET /library", "GET /downloads",
           ],
         });
@@ -1226,16 +1472,95 @@ export function createServer(
         });
       }
 
-      if (url.pathname === "/v1/models" && request.method === "GET") {
+      // mlx_lm.server parity: GET /health → the exact body it writes
+      // ('{"status": "ok"}', note the space) so byte-for-byte health checks pass.
+      if (url.pathname === "/health" && request.method === "GET") {
+        return new Response('{"status": "ok"}', {
+          headers: { "content-type": "application/json" },
+        });
+      }
+
+      // GET /v1/models — the served model FIRST (with our capability extras),
+      // then every other servable model the registry knows (mlx-lm scans the
+      // HF cache here; our registry is that scan, filtered to supported
+      // architectures). GET /v1/models/<id> filters to that id — same list
+      // shape, matching mlx_lm.server's handle_models_request.
+      if (
+        (url.pathname === "/v1/models" || url.pathname.startsWith("/v1/models/")) &&
+        request.method === "GET"
+      ) {
+        const filterId = url.pathname.length > "/v1/models/".length - 1
+          ? decodeURIComponent(url.pathname.slice("/v1/models/".length))
+          : null;
+        const created = Math.floor(startedAt / 1000);
+        const data: Array<Record<string, unknown>> = [{
+          id: ctx.modelId, object: "model", created, owned_by: "mlx-bun",
+          context_window: ctx.model.config.text.maxPositionEmbeddings,
+          // Capability flags for clients (CLI/external pi) that build a
+          // provider from discovery — `reasoning` gates the thinking toggle,
+          // `vision` the image input declaration (tower loaded or lazily
+          // loadable; same signal the web embed uses).
+          reasoning: ctx.template.supportsThinking,
+          vision: !!(ctx.vision || ctx.loadVision),
+        }];
+        try {
+          const { Registry } = await import("./registry");
+          const reg = new Registry();
+          try {
+            if (reg.list().length === 0) await reg.scan();
+            // Canonical rows only — a repo with N snapshots is ONE model id.
+            for (const m of reg.listCanonical()) {
+              if (m.repoId === ctx.modelId) continue;
+              if (!isSupportedModelRecord(m.modelType, m.repoId)) continue;
+              data.push({ id: m.repoId, object: "model", created });
+            }
+          } finally {
+            reg.close();
+          }
+        } catch { /* registry unavailable → served model only */ }
+        return Response.json({
+          object: "list",
+          data: filterId ? data.filter((m) => m.id === filterId) : data,
+        });
+      }
+
+      // OpenAI embeddings API. Works when the SERVED model is an embedding model
+      // (plain Qwen3 / Qwen3-Embedding) — consistent with the single-model server
+      // design: `mlx-bun serve <embedding-model>` to use this. Optional non-standard
+      // `instruction` applies Qwen3-Embedding's query format. Embedding is a pure
+      // forward (no decode loop / gateway), so it runs inline.
+      if (url.pathname === "/v1/embeddings" && request.method === "POST") {
+        if (!isEmbeddingModel(ctx.model))
+          return Response.json({
+            error: {
+              message: `served model "${ctx.modelId}" is not an embedding model; ` +
+                `serve an embedding model (e.g. Qwen3-Embedding) to use /v1/embeddings`,
+              type: "invalid_request_error",
+            },
+          }, { status: 400 });
+        let body: { input?: string | string[]; instruction?: string };
+        try {
+          body = (await request.json()) as typeof body;
+        } catch {
+          return Response.json({ error: { message: "invalid JSON body" } }, { status: 400 });
+        }
+        const inputs = Array.isArray(body.input) ? body.input : body.input != null ? [body.input] : [];
+        if (inputs.length === 0 || !inputs.every((s) => typeof s === "string"))
+          return Response.json({
+            error: { message: "`input` must be a string or array of strings", type: "invalid_request_error" },
+          }, { status: 400 });
+        const instruction = typeof body.instruction === "string" ? body.instruction : undefined;
+        const results = embedMany(ctx.model, ctx.tokenizer, inputs, instruction);
+        let totalTokens = 0;
+        const data = results.map((r, index) => {
+          totalTokens += r.tokens;
+          return { object: "embedding", index, embedding: Array.from(r.vector) };
+        });
         return Response.json({
           object: "list",
-          data: [{
-            id: ctx.modelId, object: "model", created: 0, owned_by: "mlx-bun",
-            context_window: ctx.model.config.text.maxPositionEmbeddings,
-            // Capability flags for clients (CLI/external pi) that build a
-            // provider from discovery — `reasoning` gates the thinking toggle.
-            reasoning: ctx.template.supportsThinking,
-          }],
+          data,
+          model: ctx.modelId,
+          usage: { prompt_tokens: totalTokens, total_tokens: totalTokens },
         });
       }
 
@@ -1284,7 +1609,10 @@ export function createServer(
         if (!body.id || !body.path)
           return Response.json({ error: { message: "id and path required" } }, { status: 400 });
         try {
-          const info = await enqueue(() => ctx.adapters.mount(body.id!, body.path!));
+          // Under the gateway lock: mount/unmount mutate the shared adapter
+          // registry (unmount disposes arrays a running generation could still
+          // hold) — one mutual-exclusion domain with generation (D3).
+          const info = await gateway.runExclusive(() => ctx.adapters.mount(body.id!, body.path!));
           return Response.json({
             id: info.id, mounted_layers: info.mountedLayers,
             rank: info.rank, scale: info.scale,
@@ -1295,7 +1623,7 @@ export function createServer(
       }
       if (url.pathname.startsWith("/v1/adapters/") && request.method === "DELETE") {
         const id = decodeURIComponent(url.pathname.slice("/v1/adapters/".length));
-        const removed = await enqueue(async () => ctx.adapters.unmount(id));
+        const removed = await gateway.runExclusive(async () => ctx.adapters.unmount(id));
         return removed > 0
           ? Response.json({ id, removed_layers: removed })
           : Response.json({ error: { message: `adapter ${id} not mounted` } }, { status: 404 });
@@ -1314,7 +1642,10 @@ export function createServer(
         if (sids[0] === sids[1] && sids[0] === ctx.tokenizer.bosTokenId) sids = sids.slice(1);
         try {
           const NB = 80;
-          const result = await enqueue(async () => {
+          // gateway.runExclusive: this raw forward must not run concurrently
+          // with batched decode steps or a serial generation (GPU + shared
+          // model state are single-owner — D3, one lock).
+          const result = await gateway.runExclusive(async () => {
             const cache = ctx.model.makeCache();
             try {
               const logits = ctx.model.forward(sids, cache); // [1, L, V]
@@ -1369,7 +1700,9 @@ export function createServer(
             const genOpts: GenerateOptions = useCurve
               ? { curve, seed: baseSeed + i, maxTokens, ...kvScheme }
               : { temperature: recipe.temperature, topP: recipe.topP, topK: recipe.topK, seed: baseSeed + i, maxTokens, ...kvScheme };
-            await enqueue(() => runGeneration(ids, genOpts, (t) => { toks.push(t); }));
+            // runExclusive: runGeneration touches the GPU + prompt cache, so
+            // it needs the gateway lock (this endpoint bypasses gateway.run).
+            await gateway.runExclusive(() => runGeneration(ids, genOpts, (t) => { toks.push(t); }));
             const text = ctx.tokenizer.decode(toks, true).trim();
             samples.push({ text, junk: curveJunk(text) });
           }
@@ -1387,6 +1720,10 @@ export function createServer(
       const handleChat = async (body: ChatRequest): Promise<Response> => {
         if (!Array.isArray(body.messages) || body.messages.length === 0)
           return Response.json({ error: { message: "messages required" } }, { status: 400 });
+        // mlx-lm validates logprobs params up front (ValueError → 400)
+        const lpParamError = validateLogprobsParams(body);
+        if (lpParamError)
+          return Response.json({ error: { message: lpParamError } }, { status: 400 });
 
         const id = `chatcmpl-${crypto.randomUUID()}`;
         const created = Math.floor(Date.now() / 1000);
@@ -1402,8 +1739,33 @@ export function createServer(
         // Vision is Gemma4-only (no switchable thinking channel), so it's false.
         let startInThinking = false;
         let vision: Parameters<typeof runGeneration>[3];
+        let diffusionPixels: import("./mlx/array").MlxArray | null = null;
         try {
-          if (hasImages) {
+          if (hasImages && ctx.model instanceof DiffusionGemmaModel) {
+            // DiffusionGemma image-text-to-text: its OWN dedicated SigLIP tower +
+            // encoder vision merge feed the denoising engine (NOT the AR
+            // forwardEmbeddings path). v1 supports a single image.
+            const dm = ctx.model;
+            if (!dm.visionTower)
+              return Response.json(
+                { error: { message: "this checkpoint has no vision tower" } }, { status: 400 },
+              );
+            const { messages, images } = await extractImages(normalizeMessages(body.messages));
+            if (images.length !== 1)
+              return Response.json(
+                { error: { message: "DiffusionGemma image input supports exactly one image" } },
+                { status: 400 },
+              );
+            const rendered = ctx.template.render(messages, { tools, addGenerationPrompt: true });
+            const rawIds = ctx.tokenizer.encode(rendered, /* addSpecialTokens */ false);
+            const { pixels, softTokens } = await dm.visionTower.preprocess(images[0]!);
+            promptIds = spliceImageTokens(rawIds, [softTokens], {
+              image: ctx.visionTokenIds.imageTokenId,
+              boi: ctx.visionTokenIds.boiTokenId,
+              eoi: ctx.visionTokenIds.eoiTokenId,
+            });
+            diffusionPixels = pixels;
+          } else if (hasImages) {
             // Loads (and caches) the tower on first image request — text-only
             // sessions never pay for it.
             const tower = getVisionTower(ctx);
@@ -1431,7 +1793,17 @@ export function createServer(
             { status: 400 },
           );
         }
-        const options = toOptions(body);
+        let options: ReturnType<typeof toOptions>;
+        try {
+          options = toOptions(body);
+        } catch (e) {
+          // bad logit_bias (non-numeric keys/values) — mlx-lm's coercion error
+          vision?.embeddings.dispose();
+          vision?.imageMask.dispose();
+          diffusionPixels?.dispose();
+          return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
+        }
+        if (diffusionPixels) options.visionPixels = diffusionPixels;
         // Admission: reject what cannot finish within the memory budget
         // (the GPU OOM it would otherwise hit is uncatchable and kills
         // the process — Phase 6 finding).
@@ -1439,6 +1811,7 @@ export function createServer(
         if (requiredCtx > admission.maxSafeContext) {
           vision?.embeddings.dispose();
           vision?.imageMask.dispose();
+          diffusionPixels?.dispose();
           return Response.json(
             {
               error: {
@@ -1455,18 +1828,41 @@ export function createServer(
           );
         }
         try {
-          const adapterIds = ctx.adapters.resolveSpec(body.adapter);
+          // A request's explicit `adapter` (incl. "none") wins over the
+          // startup default from `serve --adapter <dir>`.
+          const adapterIds = ctx.adapters.resolveSpec(body.adapter ?? serverOptions.defaultAdapter);
           if (adapterIds.length) options.adapters = adapterIds;
         } catch (e) {
           return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
         }
 
-        // What lane this request takes (vision / adapters / repetition penalty /
+        // logprobs capture: non-stream only — stream chunks never carry
+        // logprobs (mirroring mlx_lm.server, whose streaming generate_response
+        // calls pass no token_logprobs/top_tokens), so streaming requests skip
+        // the capture cost entirely and stay batchable.
+        const wantLogprobs = body.logprobs === true;
+        const topLogprobs =
+          typeof body.top_logprobs === "number" && body.top_logprobs > 0
+            ? body.top_logprobs : 0;
+        const captureLogprobs = !body.stream && (wantLogprobs || topLogprobs > 0);
+        if (captureLogprobs) {
+          options.logprobs = wantLogprobs;
+          options.topLogprobs = topLogprobs;
+        }
+
+        // What lane this request takes (vision / adapters / logits processors /
         // a user-fixed seed → serial; everything else batches when --batch N).
         const shape = {
           hasVision: !!vision,
           hasAdapters: !!options.adapters?.length,
           hasRepetitionPenalty: !!options.repetitionPenalty,
+          // min_p / XTC / logit_bias / presence+frequency penalties: serial-only
+          // in v1 (see GenerationGateway.RequestShape.hasLogitsExtras).
+          hasLogitsExtras: !!(
+            options.minP || options.xtcProbability || options.logitBias ||
+            options.presencePenalty || options.frequencyPenalty
+          ),
+          wantsLogprobs: captureLogprobs,
           userSeed: body.seed !== undefined,
           kvQuant: !!(options.kvConfig?.length || options.kvBits),
         };
@@ -1581,9 +1977,15 @@ export function createServer(
             const router = toolRouter(tools);
             const stopper = new StopMatcher(options.stopSequences);
             const thinking = new ThinkingTagSplitter(ctx.template.thinkingFormat === "think-tag", startInThinking);
+            // mlx-lm collects logprobs across EVERY generated token (reasoning
+            // and tool tokens included), not just visible content — same here.
+            const lpc = captureLogprobs
+              ? new LogprobsCollector(wantLogprobs, topLogprobs, (tid) => ctx.tokenizer.idToToken(tid))
+              : null;
             let content = "";
             let reasoning = "";
-            const s = await gateway.run(promptIds, options, (token) => {
+            const s = await gateway.run(promptIds, options, (token, lpInfo) => {
+              lpc?.push(token, lpInfo);
               const rawContent = router.push(token);
               reasoning += router.takeReasoning(); // gemma-channel thinking
               const parts = thinking.push(stopper.push(rawContent));
@@ -1610,6 +2012,7 @@ export function createServer(
               ? "tool_calls"
               : stopper.stopped ? "stop"
               : s.generatedTokens >= (options.maxTokens ?? 1024) ? "length" : "stop";
+            const logprobsBlock = lpc?.payload() ?? null;
             return Response.json({
               id, object: "chat.completion", created, model: ctx.modelId,
               choices: [{
@@ -1620,6 +2023,7 @@ export function createServer(
                   ...(reasoning ? { reasoning } : {}),
                   ...(toolCalls.length ? { tool_calls: toolCalls } : {}),
                 },
+                ...(logprobsBlock ? { logprobs: logprobsBlock } : {}),
                 finish_reason: finish,
               }],
               usage: {
@@ -1645,6 +2049,195 @@ export function createServer(
         return handleChat(body);
       }
 
+      // Raw text completion (mlx_lm.server's /v1/completions, request_type
+      // "text"): NO chat template — the prompt string is tokenized directly
+      // (tokenizer.encode with the tokenizer's own special-token handling,
+      // exactly mlx-lm's `tokenizer.encode(request.prompt)`). Rides the same
+      // GenerationGateway + admission + adapter path as chat; no tool router
+      // or thinking splitter (raw text in, raw text out).
+      if (url.pathname === "/v1/completions" && request.method === "POST") {
+        let body: TextCompletionRequest;
+        try {
+          body = (await request.json()) as TextCompletionRequest;
+        } catch {
+          return Response.json({ error: { message: "invalid JSON body" } }, { status: 400 });
+        }
+        if (typeof body.prompt !== "string" || body.prompt.length === 0)
+          return Response.json(
+            {
+              error: {
+                message: "prompt (a non-empty string) is required " +
+                  "(token-array prompts are not accepted, matching mlx_lm.server)",
+              },
+            },
+            { status: 400 },
+          );
+        // mlx-lm validates logprobs params up front (ValueError → 400)
+        const lpParamError = validateLogprobsParams(body);
+        if (lpParamError)
+          return Response.json({ error: { message: lpParamError } }, { status: 400 });
+        const id = `cmpl-${crypto.randomUUID()}`;
+        const created = Math.floor(Date.now() / 1000);
+        const promptIds = ctx.tokenizer.encode(body.prompt);
+        let options: ReturnType<typeof toOptions>;
+        try {
+          options = toOptions(body as unknown as ChatRequest);
+        } catch (e) {
+          // bad logit_bias (non-numeric keys/values) — mlx-lm's coercion error
+          return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
+        }
+        // mlx_lm.server's default max_tokens is 512 (its --max-tokens CLI
+        // default). The chat lane's very generous default is wrong for raw
+        // completion: with no template an EOS may never come.
+        options.maxTokens = body.max_completion_tokens ?? body.max_tokens ?? serverOptions.defaultMaxTokens ?? 512;
+        const requiredCtx = promptIds.length + options.maxTokens;
+        if (requiredCtx > admission.maxSafeContext)
+          return Response.json(
+            {
+              error: {
+                message:
+                  `request needs ${requiredCtx} tokens of context ` +
+                  `(prompt ${promptIds.length} + max_tokens ${options.maxTokens}) but the ` +
+                  `memory budget caps safe context at ${admission.maxSafeContext} — ` +
+                  `shorten the prompt or lower max_tokens`,
+                type: "memory_admission",
+                code: "context_over_budget",
+              },
+            },
+            { status: 400 },
+          );
+        try {
+          const adapterIds = ctx.adapters.resolveSpec(body.adapter ?? serverOptions.defaultAdapter);
+          if (adapterIds.length) options.adapters = adapterIds;
+        } catch (e) {
+          return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
+        }
+        // logprobs capture — same mlx-lm block as chat (generate_response is
+        // shared in the reference too); non-stream only, stream chunks never
+        // carry logprobs.
+        const wantLogprobs = body.logprobs === true;
+        const topLogprobs =
+          typeof body.top_logprobs === "number" && body.top_logprobs > 0
+            ? body.top_logprobs : 0;
+        const captureLogprobs = !body.stream && (wantLogprobs || topLogprobs > 0);
+        if (captureLogprobs) {
+          options.logprobs = wantLogprobs;
+          options.topLogprobs = topLogprobs;
+        }
+        const shape = {
+          hasVision: false,
+          hasAdapters: !!options.adapters?.length,
+          hasRepetitionPenalty: !!options.repetitionPenalty,
+          hasLogitsExtras: !!(
+            options.minP || options.xtcProbability || options.logitBias ||
+            options.presencePenalty || options.frequencyPenalty
+          ),
+          wantsLogprobs: captureLogprobs,
+          userSeed: body.seed !== undefined,
+          kvQuant: !!(options.kvConfig?.length || options.kvBits),
+        };
+        const finishReason = (stopped: boolean, generated: number): "stop" | "length" =>
+          stopped ? "stop" : generated >= (options.maxTokens ?? 512) ? "length" : "stop";
+
+        if (body.stream) {
+          const stream = new ReadableStream<Uint8Array>({
+            async start(controller) {
+              const enc = new TextEncoder();
+              const send = (obj: unknown) =>
+                controller.enqueue(enc.encode(`data: ${JSON.stringify(obj)}\n\n`));
+              const chunk = (text: string, finish: string | null) => ({
+                id, object: "text_completion", created, model: ctx.modelId,
+                choices: [{ index: 0, text, finish_reason: finish }],
+              });
+              try {
+                const decoder = new StreamDecoder(ctx.tokenizer);
+                const stopper = new StopMatcher(options.stopSequences);
+                // ≥25 ms macrotask hop, same reason as the chat lane: keep the
+                // serial decode loop from starving the socket (SSE bursts).
+                let lastFlush = performance.now();
+                const s = await gateway.run(promptIds, options, (token) => {
+                  const text = stopper.push(decoder.push(token));
+                  if (text) send(chunk(text, null));
+                  if (stopper.stopped) return false; // halt generation
+                  if (text) {
+                    const now = performance.now();
+                    if (now - lastFlush >= 25) {
+                      lastFlush = now;
+                      return new Promise<void>((r) => setImmediate(r));
+                    }
+                  }
+                }, undefined, shape);
+                if (!stopper.stopped) {
+                  let tail = stopper.push(decoder.flush());
+                  if (!stopper.stopped) tail += stopper.flush();
+                  if (tail) send(chunk(tail, null));
+                }
+                // final chunk: finish_reason + usage (mlx-lm gates usage behind
+                // stream_options.include_usage; we always attach it, matching
+                // our chat lane — an additive superset OpenAI clients ignore)
+                send({
+                  ...chunk("", finishReason(stopper.stopped, s.generatedTokens)),
+                  usage: {
+                    prompt_tokens: s.promptTokens,
+                    completion_tokens: s.generatedTokens,
+                    total_tokens: s.promptTokens + s.generatedTokens,
+                    prompt_tokens_details: { cached_tokens: s.cachedTokens },
+                  },
+                });
+                controller.enqueue(enc.encode("data: [DONE]\n\n"));
+              } 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",
+            },
+          });
+        }
+
+        try {
+          const decoder = new StreamDecoder(ctx.tokenizer);
+          const stopper = new StopMatcher(options.stopSequences);
+          const lpc = captureLogprobs
+            ? new LogprobsCollector(wantLogprobs, topLogprobs, (tid) => ctx.tokenizer.idToToken(tid))
+            : null;
+          let text = "";
+          const s = await gateway.run(promptIds, options, (token, lpInfo) => {
+            lpc?.push(token, lpInfo);
+            text += stopper.push(decoder.push(token));
+            if (stopper.stopped) return false; // halt generation
+          }, undefined, shape);
+          if (!stopper.stopped) {
+            let tail = stopper.push(decoder.flush());
+            if (!stopper.stopped) tail += stopper.flush();
+            text += tail;
+          }
+          const logprobsBlock = lpc?.payload() ?? null;
+          return Response.json({
+            id, object: "text_completion", created, model: ctx.modelId,
+            choices: [{
+              index: 0, text,
+              ...(logprobsBlock ? { logprobs: logprobsBlock } : {}),
+              finish_reason: finishReason(stopper.stopped, s.generatedTokens),
+            }],
+            usage: {
+              prompt_tokens: s.promptTokens,
+              completion_tokens: s.generatedTokens,
+              total_tokens: s.promptTokens + s.generatedTokens,
+              prompt_tokens_details: { cached_tokens: s.cachedTokens },
+            },
+          });
+        } catch (e) {
+          return Response.json({ error: { message: (e as Error).message } }, { status: 500 });
+        }
+      }
+
       // Anthropic Messages API (Phase 11) — on by default, mirroring
       // optiq serve (--anthropic defaults True; the drop-in claim
       // depends on it). Oracle: optiq/anthropic_shim.py, ported in
@@ -1902,7 +2495,7 @@ export function createServer(
         const snapMatch = body.model_id.match(/(models--[^/]+)\/snapshots\//);
         let org = "local", name: string;
         if (snapMatch) {
-          const parts = snapMatch[1].split("--"); // ["models", org, ...name]
+          const parts = snapMatch[1]!.split("--"); // ["models", org, ...name]
           org = parts[1] ?? "local";
           name = parts.slice(2).join("--");
         } else if (body.model_id.includes("/") && !body.model_id.startsWith("/") && !body.model_id.startsWith("~")) {