mlx-bun 0.0.13 → 0.3.0

package/src/cli.ts

@@ -13,7 +13,8 @@ import { existsSync } from "node:fs";
 import { basename, join, resolve } from "node:path";
 import { Registry } from "./registry";
 import { loadModelConfig, parseTurboQuantScheme, type TurboQuantScheme } from "./config";
-import { fit, skuMatrix, thisMachine } from "./fit";
+import { fit, skuMatrix, thisMachine, type FitKvScheme } from "./fit";
+import { resolveKvScheme } from "./kv-scheme";
 import { EvalDB } from "./evaldb";
 import pkg from "../package.json" with { type: "json" };
 import { renderHelp } from "./tui";
@@ -27,8 +28,21 @@ import {
   planGlm52MemoryForArtifact,
 } from "./model/glm52-memory";
 import type { JobEvent } from "./jobs/types";
+import {
+  configureRuntime,
+  runtimeValue,
+  type RuntimeKey,
+} from "./runtime-config";
 
 const argv = process.argv.slice(2);
+// Private self-exec protocol for managed jobs in the single-file binary.
+// Literal imports in job-entry keep the numerical runners inside the bundle.
+if (argv[0] === "__job") {
+  const { runJobEntry } = await import("./jobs/job-entry");
+  await runJobEntry(argv[1]);
+  process.exit(0);
+}
+
 // The appliance path: naked `mlx-bun` (or only options, e.g.
 // `mlx-bun --port 9000`) runs `serve` — first run downloads a model and
 // opens the chat UI. `--help`/`--version` and explicit subcommands still win.
@@ -97,13 +111,19 @@ const SERVER_FLAGS = `Server options:
                             restarts (long-context agent TTFT drops from a
                             full re-prefill to a zero-copy mmap restore).
                             Off unless set.
-  --ssd-cache-max <GB>      SSD tier byte cap  [default: 32 GB]
+  --ssd-cache-max <GB>      Optional SSD tier byte cap; 0 = unlimited
+                            [default: unlimited]
   --ssd-demote-idle <sec>   Demote prompt-cache entries idle this long to the
                             SSD tier, freeing their GPU memory (prefixes stay
                             reachable via zero-copy restore). 0 disables.
                             [default: 300 when --ssd-cache is on]
   --ssd-cache-verify        Verify tensor hashes on every restore (reads all
                             bytes eagerly — integrity paranoia only)
+  --generation-checkpoint <tokens>
+                            Snapshot an in-flight serial generation to the SSD
+                            cache every N emitted tokens. Repeating the same
+                            request after restart replays and resumes it.
+                            Requires --ssd-cache and --batch 1.
   --isolate                 Run the inference engine as a CHILD process and
                             keep this process a pure UI/API proxy — the UI
                             stays instant under GPU load and survives engine
@@ -118,10 +138,10 @@ const SERVER_FLAGS = `Server options:
   --unix <path>             (internal) listen on a unix socket instead of a
                             TCP port — the engine half of --isolate.
   --batch <n>               Max concurrent requests decoded together
-                            [default: 8]. A LONE request runs the exact
-                            serial engine (same bits, same speed); the cap
-                            engages only when concurrent requests arrive
-                            (agent fan-out). 1 pins strict serial.
+                            [default: 8]. A lone scheduler request uses the
+                            parity- and latency-gated B=1 fast path; concurrent
+                            requests can form B=N (agent fan-out). 1 pins the
+                            preserved strict serial executor.
                             --decode-concurrency is accepted as the
                             mlx_lm.server alias (semantics differ; see
                             docs/reference/server-config.md)
@@ -152,9 +172,13 @@ Model & quality:
                             Serial lane only — with --batch N a mounted draft
                             routes every request serial, like mlx_lm.server.
   --draft-kind <kind>       Override draft-artifact detection:
-                            two-model | assistant | dspark | deepspec | ngram
+                            two-model | assistant | dspark | deepspec | mtp | ngram
                             (deepspec = DeepSeek's released Gemma4DSparkModel
-                            drafters; ngram = MODEL-FREE prompt lookup — no
+                            drafters; mtp = a NATIVE multi-token-prediction
+                            head split from the target's own release, e.g.
+                            mlx-community/Qwen3.8-27B-MTP-bf16 — auto-detected
+                            from model_type *_mtp;
+                            ngram = MODEL-FREE prompt lookup — no
                             --draft-model, drafts copied from the request's
                             own context, lossless)  [default: auto]
   --num-draft-tokens <n>    Drafts per verify round  [default: 3; ngram: 10;
@@ -171,7 +195,7 @@ Model & quality:
                             (bf16), 4 / 8 (uniform bits), or turbo[:k<bits>v
                             <bits>] (rotation-based TurboQuant — default
                             k8v3; kBits in {2,4,5,8}, vBits in {2,3,4,5,8};
-                            docs/design/turboquant-kv.md)
+                            docs/design/turboquant.md)
                             [default: off — quantized KV measured 5-20% slower
                             decode than bf16 at ≤16k (on mlx-lm too); it buys
                             memory headroom, so opt in when context is tight.
@@ -186,7 +210,7 @@ Model & quality:
                             --kv-quant, --draft-model; bypasses the prompt
                             cache; runs uncompiled decode. Bit-exact with the
                             plain path; expect a small decode cost at batch=1
-                            (the gather copy). docs/design/paged-kv-cache.md
+                            (the gather copy). docs/design/kv-cache.md
                             [default: off]
   --paged-kv-block-size <n> Tokens per KV block  [default: 256 = the plain
                             cache's growth step]
@@ -202,13 +226,14 @@ Model & quality:
                             pass --temp 0 for its unset-request behavior)
                             [default: the model's generation_config.json]
   --max-tokens <n>          Completion cap when a request omits max_tokens
-                            (mlx_lm.server flag; its default is 512 there)
-                            [default: 128 GLM-5.2; otherwise 65536 chat /
-                            512 raw completion]
+                            (mlx_lm.server flag; --max-tokens 512 reproduces
+                            its default) [default: 128 GLM-5.2 (memory-plan
+                            reservation); otherwise none — a defaulted
+                            request runs to EOS or the admitted context]
   --hlg-sampling on|off     Piecewise tone-curve sampling (HLG): rolls off the
                             top, boosts the mids, gentles the tail. Gain folds
                             from --temperature. [default: off]
-                            See docs/design/hlg-sampling.md
+                            See docs/archive/hlg-sampling.md
   --hlg-width <nats>        HLG mid-region half-width  [default: 4]
   --hlg-shoulder <nats>     HLG highlight rolloff scale  [default: 4]
   --hlg-toe <nats>          HLG shadow rolloff scale  [default: 6]
@@ -343,6 +368,8 @@ Usage: mlx-bun fit <query> [options]
 
 Options:
   --ctx <tokens>       Context size to assess  [default: 8192]
+  --kv-quant <mode>    Bill the KV cache quantized: 4 | 8 | config | off
+                       (mirrors serve --kv-quant; default bf16)
   --skus               Also print the Apple silicon SKU matrix
 
 Solves weights + KV growth + prefill transient against wired memory and
@@ -379,7 +406,7 @@ Usage: mlx-bun benchmark [query] [options]
 
 Runs mlx-bun by itself (no other systems required) and records the
 result to the eval DB. For quotable, cross-stack numbers use
-./benchmark.sh from the repo — it preflight-gates on an idle machine
+bun scripts/bench-serve.ts all from the repo — it preflight-gates on an idle machine
 and runs the mlx-lm/optiq comparison legs.
 
   [query]              Model to benchmark (--query <q> is accepted for the
@@ -463,7 +490,7 @@ Options:
   --clear              Archive the eval DB (timestamped backup) and start fresh
 
 Table view: when, model, bench kind, KV mode, decode tok/s, TTFT,
-peak memory, commit. Runs are written by ./benchmark.sh.`,
+peak memory, commit. Runs are written by bun scripts/bench-serve.ts all.`,
 
   train: `mlx-bun train — fine-tune a LoRA adapter on your data
 
@@ -576,6 +603,9 @@ knapsack per-layer bit allocation; implies quantization, no -q needed):
   --candidate-bits <l> Comma list the knapsack may pick from [default: 4,8]
   --calibration-mix <m> "optiq" or a JSONL path  [default: optiq]
   --n-calibration <n>  Calibration samples  [default: 2]
+  --rotate-weights     Fold the model's offline TurboQuant rotation before
+                       quantization (auto-detects Llama/Qwen3.5/Qwen MTP)
+  --rotation-seed <n>  Deterministic rotation seed  [default: 42]
 
 Not supported (mlx_lm.convert flags we don't implement — the command
 exits with an error rather than guessing): --dtype, -d/--dequantize,
@@ -805,19 +835,19 @@ function openChatUi(url: string, hostPort: string): void {
 }
 
 /** Server/runtime flags shared by every mode that loads a model
- *  (serve, pi). Env levers are set here so they're in place before the
+ *  (serve, pi). Runtime overrides are installed here before the
  *  generate/compiled-decode modules read them. */
 /** Resolve the decode ROUTE: a tier alias (--l1/--l2) sets the whole route,
  *  and an explicit per-fork flag (--kv-quant/--fused-sdpa/…) overrides the
- *  alias. Sets the decode env levers and returns the kv-quant mode. See
- *  docs/design/parity-tier-dag.md + unified-engine-frontier-plan.md.
+ *  alias. Installs decode runtime options and returns the kv-quant mode. See
+ *  docs/design/unified-engine-frontier-plan.md + unified-engine-frontier-plan.md.
  *  Each tier is a GUARANTEE about which reference you reproduce bit-for-bit:
  *    --l1  bit-for-bit IDENTICAL to mlx-lm    — drop-in replacement for mlx-lm
  *          (bf16 KV; THE naked default since 2026-07-05)
  *    --l2  bit-for-bit IDENTICAL to mlx-optiq — drop-in replacement for
  *          mlx-optiq: quantized KV per kv_config.json + fused N-tiled SDPA for
  *          L>1 + STOCK unfused L=1 decode (the optiq-golden composition,
- *          scripts/regen-kvq-goldens.ts + regen-mixed-kv-goldens.ts).
+ *          scripts/regen.ts kvq + regen-mixed-kv-goldens.ts).
  *  There is no --l3: output-changing experiments live in the Lab (env flags
  *  with a bench + expiry) until one beats the L1 baseline in a paired A/B —
  *  docs/design/unified-engine-frontier-plan.md §6-7. */
@@ -863,8 +893,9 @@ function applyDecodeRoute(): { kvQuant?: "off" | "config" | number; turboQuant?:
   const pick = (name: string, base: boolean | undefined): boolean | null => {
     const ex = onOff(name); return ex !== null ? ex : (base ?? null); // explicit flag wins, else tier preset
   };
-  const set = (env: string, val: boolean | null, invert = false) => {
-    if (val !== null) process.env[env] = (invert ? !val : val) ? "1" : "0";
+  const set = (name: RuntimeKey, val: boolean | null, invert = false) => {
+    if (val !== null)
+      configureRuntime({ [name]: (invert ? !val : val) ? "1" : "0" });
   };
   // Explicit --kv-quant picks the composition its ORACLE uses: config (the
   // optiq mixed scheme) runs optiq's fused SDPA (the L2 golden composition);
@@ -874,14 +905,14 @@ function applyDecodeRoute(): { kvQuant?: "off" | "config" | number; turboQuant?:
   set("MLX_BUN_COMPILED_DECODE", pick("compiled-decode", p.compiled));
   set("MLX_BUN_NO_FUSED_SDPA", pick("fused-sdpa", fusedSdpaBase), true); // inverted env
   // Compiled activations (the faithful geglu/swiglu — mlx-lm's @mx.compile kernel).
-  // One fork drives both env vars; qwen3/qwen3.5/universal compile unconditionally,
+  // One fork drives both runtime options; qwen3/qwen3.5/universal compile unconditionally,
   // so this toggles gemma geglu + minicpm5 swiglu (the only unfused-capable sites).
   set("MLX_BUN_COMPILED_GEGLU", pick("compiled-activations", p.compiledAct));
   set("MLX_BUN_COMPILED_SWIGLU", pick("compiled-activations", p.compiledAct));
   if (kv === "off") return { kvQuant: "off" };
   if (kv === "config") return { kvQuant: "config" };
   if (kv) {
-    // TurboQuant (docs/design/turboquant-kv.md): a distinct axis from the
+    // TurboQuant (docs/design/turboquant.md): a distinct axis from the
     // affine kvQuant above — dequantize-on-fetch, stock unfused ops.sdpa,
     // no kvQuant value at all (mutually exclusive with kvBits/kvConfig).
     let turboQuant: TurboQuantScheme | null;
@@ -909,11 +940,12 @@ function serverRuntimeFlags(): { port: number; serverOptions: import("./server")
     process.exit(1);
   };
   const route = applyDecodeRoute(); // --l1/--l2 tier alias, with per-fork flags overriding
-  if (flag("force-wire")) process.env.MLX_BUN_FORCE_WIRE = "1";
+  if (flag("force-wire")) configureRuntime({ MLX_BUN_FORCE_WIRE: "1" });
   // Remote image_url/audio_url media fetches block private/loopback/link-
   // local destinations by default (SSRF — src/media-fetch.ts); this is the
   // LAN-hosts escape hatch.
-  if (flag("allow-private-media")) process.env.MLX_BUN_ALLOW_PRIVATE_MEDIA = "1";
+  if (flag("allow-private-media"))
+    configureRuntime({ MLX_BUN_ALLOW_PRIVATE_MEDIA: "1" });
 
   const serverOptions: import("./server").ServerOptions = {};
   const budgetGB = Number(opt("memory-budget", "0"));
@@ -923,29 +955,47 @@ function serverRuntimeFlags(): { port: number; serverOptions: import("./server")
   if (kvBudgetGB > 0) serverOptions.kvBudgetBytes = kvBudgetGB * 1e9;
   const pcRaw = opt("prompt-cache"); // null if absent; an explicit value (incl. 0) wins → `--prompt-cache 0` DISABLES the cache
   if (pcRaw !== null) serverOptions.promptCacheBytes = Math.max(0, Number(pcRaw)) * 2 ** 30;
-  // Engine-child mode (runtime-isolation.md): listen on a unix socket
+  // Engine-child mode (docs/reference/server-config.md): listen on a unix socket
   // instead of TCP — the isolated parent proxies here. Internal flag,
   // spawned by --isolate; usable directly for socket-level integration.
   const unixRaw = opt("unix");
   if (unixRaw !== null) serverOptions.unixSocket = unixRaw;
-  // SSD cold tier (docs/design/ssd-kv-cold-tier.md) — off unless a dir is given.
+  // SSD cold tier (docs/design/kv-cache.md) — off unless a dir is given.
   const ssdDir = opt("ssd-cache");
   if (ssdDir !== null) {
     serverOptions.ssdCacheDir = ssdDir;
     const maxRaw = opt("ssd-cache-max");
-    if (maxRaw !== null) serverOptions.ssdCacheMaxBytes = Math.max(1, Number(maxRaw)) * 2 ** 30;
+    if (maxRaw !== null) {
+      const maxGb = Number(maxRaw);
+      if (!Number.isFinite(maxGb) || maxGb < 0) {
+        console.error(`--ssd-cache-max expects a number >= 0 (got "${maxRaw}")`);
+        process.exit(1);
+      }
+      serverOptions.ssdCacheMaxBytes = maxGb === 0
+        ? Number.POSITIVE_INFINITY
+        : maxGb * 2 ** 30;
+    }
     if (flag("ssd-cache-verify")) serverOptions.ssdCacheVerify = true;
     const demoteRaw = opt("ssd-demote-idle");
     if (demoteRaw !== null) serverOptions.ssdDemoteIdleSec = Math.max(0, Number(demoteRaw));
+    const checkpointRaw = opt("generation-checkpoint");
+    if (checkpointRaw !== null) {
+      const n = Number(checkpointRaw);
+      if (!Number.isInteger(n) || n < 1) {
+        console.error(`--generation-checkpoint expects an integer >= 1 (got "${checkpointRaw}")`);
+        process.exit(1);
+      }
+      serverOptions.generationCheckpointTokens = n;
+    }
   } else {
     // Refuse-loudly house rule: sub-flags of an absent tier must not vanish.
-    for (const f of ["ssd-cache-max", "ssd-demote-idle"] as const)
+    for (const f of ["ssd-cache-max", "ssd-demote-idle", "generation-checkpoint"] as const)
       if (opt(f) !== null) console.warn(`--${f} has no effect without --ssd-cache — ignored`);
     if (flag("ssd-cache-verify"))
       console.warn("--ssd-cache-verify has no effect without --ssd-cache — ignored");
   }
   // --batch N: max concurrent requests batched through the mlx-lm-parity
-  // engine (N=1 = today's serial path). --decode-concurrency is accepted for
+  // engine (N=1 pins the strict serial executor). --decode-concurrency is accepted for
   // drop-in compatibility with mlx_lm.server, but the semantics differ: there
   // it caps per-BatchGenerator decode parallelism (default 32); in mlx-bun it
   // caps continuous batching (default 8; `--batch 1` pins the serial path —
@@ -962,10 +1012,10 @@ function serverRuntimeFlags(): { port: number; serverOptions: import("./server")
   if (route.kvQuant !== undefined) serverOptions.kvQuant = route.kvQuant;
   if (route.turboQuant !== undefined) serverOptions.turboQuant = route.turboQuant;
   // --paged-kv (env: MLX_BUN_PAGED_KV=1): OPTIONAL vLLM-style block-pool KV
-  // storage, default off (docs/design/paged-kv-cache.md). Serial-only in v1:
+  // storage, default off (docs/design/kv-cache.md). Serial-only in v1:
   // with no explicit --batch, pin --batch 1 (the default is 8); an explicit
   // --batch N>1 is refused loudly by createServer rather than downgraded.
-  if (flag("paged-kv") || process.env.MLX_BUN_PAGED_KV === "1") {
+  if (flag("paged-kv") || runtimeValue("MLX_BUN_PAGED_KV") === "1") {
     const bsRaw = opt("paged-kv-block-size");
     serverOptions.pagedKv = bsRaw !== null ? { blockSize: Number(bsRaw) } : {};
     if (batchRaw === null) {
@@ -1009,7 +1059,7 @@ function serverRuntimeFlags(): { port: number; serverOptions: import("./server")
   const maxTok = numFlag("max-tokens", 1, 10_000_000);
   if (maxTok !== null) serverOptions.defaultMaxTokens = Math.floor(maxTok);
   // HLG sampling — piecewise tone curve on the logprobs (default off). Knobs in
-  // nats; the mid gain folds from --temperature. docs/design/hlg-sampling.md.
+  // nats; the mid gain folds from --temperature. docs/archive/hlg-sampling.md.
   if (onOff("hlg-sampling") === true) {
     serverOptions.hlg = {
       enabled: true,
@@ -1086,7 +1136,7 @@ async function mountStartupAdapter(
 function runtimeSummary(o: import("./server").ServerOptions): string {
   const kv = o.turboQuant ? `turbo-k${o.turboQuant.kBits}v${o.turboQuant.vBits}`
     : o.kvQuant === "off" ? "off" : typeof o.kvQuant === "number" ? `kv${o.kvQuant}` : "config";
-  const lever = (env: string, dflt: string) => process.env[env] ?? dflt;
+  const lever = (name: RuntimeKey, dflt: string) => runtimeValue(name) ?? dflt;
   return `kv-quant ${kv} · compiled-decode ${lever("MLX_BUN_COMPILED_DECODE", "1") === "1" ? "on" : "off"}` +
     (lever("MLX_BUN_COMPILED_GEGLU", "1") === "0" ? " · compiled-activations off" : "") +
     (o.batch && o.batch > 1 ? ` · batch ${o.batch}` : "") +
@@ -1376,7 +1426,7 @@ switch (cmd) {
 
   case "fit": {
     const query = positional(0);
-    if (!query) { console.error("usage: mlx-bun fit <query> [--ctx N] [--skus]"); process.exit(1); }
+    if (!query) { console.error("usage: mlx-bun fit <query> [--ctx N] [--kv-quant 4|8|config] [--skus]"); process.exit(1); }
     const reg = new Registry();
     if (reg.list().length === 0) await reg.scan();
     const m = reg.resolve(query);
@@ -1443,9 +1493,32 @@ switch (cmd) {
       console.log();
       break;
     }
-    const r = fit(config, m.sizeBytes, ctx, thisMachine(), undefined, m.expertsBytes);
+    // --kv-quant mirrors serve: bill the quantized cache's true bytes so
+    // the reported window matches what serving with the same flag admits.
+    const kvQuantOpt = opt("kv-quant");
+    let fitKvScheme: FitKvScheme | undefined;
+    if (kvQuantOpt === "4" || kvQuantOpt === "8") {
+      fitKvScheme = resolveKvScheme({ override: Number(kvQuantOpt) }).fitOptions;
+    } else if (kvQuantOpt === "config") {
+      if (!config.kvQuant?.length) {
+        console.error(`${m.repoId} ships no per-layer kv-quant config — use --kv-quant 4|8`);
+        process.exit(1);
+      }
+      fitKvScheme = resolveKvScheme({
+        override: "config",
+        config: config.kvQuant,
+        missingConfig: "error",
+      }).fitOptions;
+    } else if (kvQuantOpt != null && kvQuantOpt !== "off") {
+      console.error(`--kv-quant ${kvQuantOpt}: expected 4, 8, config, or off`);
+      process.exit(1);
+    }
+    const r = fit(config, m.sizeBytes, ctx, thisMachine(), undefined, m.expertsBytes, undefined, fitKvScheme);
     h1("will it fit?");
-    console.log(`  ${style.bold(m.repoId)} ${style.dim(`@ ${ctx.toLocaleString()} context · this machine`)}`);
+    const kvNote = fitKvScheme
+      ? style.dim(` · kv-quant ${fitKvScheme.kvBits ? `${fitKvScheme.kvBits}-bit` : "config"}`)
+      : "";
+    console.log(`  ${style.bold(m.repoId)} ${style.dim(`@ ${ctx.toLocaleString()} context · this machine`)}${kvNote}`);
     console.log();
     const expertsNote = config.text.enableMoeBlock && m.expertsBytes > 0
       ? style.dim(`  (experts ${gb(m.expertsBytes)}; top ${config.text.topKExperts}/${config.text.numExperts} read per token)`)
@@ -1470,7 +1543,7 @@ switch (cmd) {
           { header: "fits", paint: (c) => (c.includes("fits") ? style.green(c) : style.dim(c)) },
           { header: "max context", align: "right" }, { header: "decode", align: "right" },
         ],
-        skuMatrix(config, m.sizeBytes, ctx, m.expertsBytes).map((row) => [
+        skuMatrix(config, m.sizeBytes, ctx, m.expertsBytes, fitKvScheme).map((row) => [
           row.sku, `${row.ramGB} GB`, row.fits ? "fits" : "—",
           row.fits ? row.maxContext.toLocaleString() : "—",
           row.fits ? `~${row.decodeTps.toFixed(0)} tok/s` : "—",
@@ -1498,7 +1571,7 @@ switch (cmd) {
         process.exit(1);
       }
     }
-    // Runtime isolation (--isolate, runtime-isolation.md): THIS process
+    // Runtime isolation (--isolate, docs/reference/server-config.md): THIS process
     // becomes a pure reactor — it binds the TCP port and reverse-proxies to
     // an engine CHILD (the full server on a unix socket, spawned from the
     // same argv minus parent-only flags). Zero MLX calls in the parent: the
@@ -1540,7 +1613,7 @@ switch (cmd) {
       };
       const poolMaxRaw = opt("model-pool");
       const poolMax = Math.max(1, Number(poolMaxRaw ?? 1) || 1);
-      const { engine, pool } = startProxyServer({
+      const { engine, close } = startProxyServer({
         port: rt.port,
         ...(rt.serverOptions.hostname ? { hostname: rt.serverOptions.hostname } : {}),
         engine: {
@@ -1556,7 +1629,7 @@ switch (cmd) {
       const shownHost0 = rt.serverOptions.hostname ?? "localhost";
       console.log(style.dim(`  isolated mode: engine child starting (socket ${sock})`));
       console.log(`  ${style.green("●")} ${style.bold("proxy up")} http://${shownHost0}:${rt.port} ${style.dim(`· ${m0.repoId} loading behind it · model pool ${poolMax}`)}`);
-      const onExit = () => { pool?.stopAll(); engine.stop(); process.exit(0); };
+      const onExit = () => { void close().finally(() => process.exit(0)); };
       process.on("SIGINT", onExit);
       process.on("SIGTERM", onExit);
       await engine.ready;
@@ -1573,10 +1646,29 @@ switch (cmd) {
     const config = await loadModelConfig(m.path);
     const report = fit(config, m.sizeBytes, 8192, undefined, undefined, m.expertsBytes);
     sFit.done(`${style.bold(m.repoId)} ${style.dim(`· ${gb(m.sizeBytes)} · ~${report.predictedDecodeTps.toFixed(0)} tok/s predicted`)}${picked ? style.dim(" · auto-picked (override: mlx-bun serve <query>)") : ""}`);
+    // Near-ceiling advisory (admission doctrine: advise, never refuse). A
+    // model near the DEFAULT Metal wired ceiling (~75% of RAM) dies with an
+    // UNCATCHABLE async-GPU-OOM on long prefills (M4 24 GB × 17 GB model,
+    // 2026-08-20) — surface the standard remedy up front instead.
+    {
+      const { maxRecommendedWorkingSetSize } = await import("./mlx/ffi");
+      const ceiling = maxRecommendedWorkingSetSize();
+      const sysctlRaw = Bun.spawnSync(["sysctl", "-n", "iogpu.wired_limit_mb"]).stdout.toString().trim();
+      const limitIsDefault = sysctlRaw === "0" || sysctlRaw === "";
+      if (limitIsDefault && ceiling > 0 && m.sizeBytes > 0.8 * ceiling) {
+        const ramGb = Number(Bun.spawnSync(["sysctl", "-n", "hw.memsize"]).stdout.toString().trim()) / 2 ** 30;
+        const suggestMb = Math.floor((ramGb - 2.5) * 1024);
+        console.error(style.dim(
+          `  ! ${gb(m.sizeBytes)} model vs ~${gb(ceiling)} default GPU ceiling — long prompts can\n` +
+          `    hit an uncatchable Metal OOM. If that happens, raise the limit first:\n` +
+          `      sudo sysctl iogpu.wired_limit_mb=${suggestMb}   (resets on reboot)`,
+        ));
+      }
+    }
     const sNative = step("native runtime");
     await ensureNative(sNative);
     sNative.done("native runtime ready");
-    const { createServer, loadContext } = await import("./server");
+    const { createServer, loadContext, shutdownServer } = await import("./server");
     if (flag("expert-offload")) {
       if (m.expertsBytes === 0) {
         console.error(style.dim("--expert-offload ignored: this model has no experts (dense)"));
@@ -1609,9 +1701,9 @@ switch (cmd) {
     }
     const draftKindRaw = opt("draft-kind");
     const draftKind = (draftKindRaw ?? undefined) as
-      | "dspark" | "deepspec" | "assistant" | "two-model" | "ngram" | undefined;
-    if (draftKind !== undefined && !["dspark", "deepspec", "assistant", "two-model", "ngram"].includes(draftKind)) {
-      console.error(`--draft-kind expects two-model|assistant|dspark|deepspec|ngram (got "${draftKindRaw}")`);
+      | "dspark" | "deepspec" | "assistant" | "two-model" | "ngram" | "mtp" | undefined;
+    if (draftKind !== undefined && !["dspark", "deepspec", "assistant", "two-model", "ngram", "mtp"].includes(draftKind)) {
+      console.error(`--draft-kind expects two-model|assistant|dspark|deepspec|mtp|ngram (got "${draftKindRaw}")`);
       process.exit(1);
     }
     // Prompt-lookup window bounds (ngram kind only; validated here so a typo
@@ -1639,7 +1731,10 @@ switch (cmd) {
     const ctx = await loadContext(m.path, m.repoId, {
       memoryBudgetBytes: rt.serverOptions.memoryBudgetBytes,
       glm: glm52RuntimeFlags(rt.serverOptions),
-      ...(draftModelDir || draftKind === "ngram"
+      // ngram is model-free; mtp resolves the bundled <model>/mtp/ companion
+      // in loadContext when no --draft-model is given (cli.md contract —
+      // this gate silently dropped that case until 2026-08-20).
+      ...(draftModelDir || draftKind === "ngram" || draftKind === "mtp"
         ? {
             ...(draftModelDir ? { draftModelDir } : {}),
             numDraftTokens,
@@ -1652,6 +1747,30 @@ switch (cmd) {
     sLoad.done(`weights loaded ${style.dim(`in ${(performance.now() - t0).toFixed(0)} ms`)}${draftModelDir ? style.dim(` · draft: ${draftModelDir.split("/").filter(Boolean).at(-1)}`) : draftKind === "ngram" ? style.dim(" · draft: ngram (prompt lookup)") : ""}`);
     await mountStartupAdapter(ctx, rt.serverOptions);
     const server = createServer(ctx, rt.port, { ...rt.serverOptions, owner: "serve" });
+    let shutdownStarted = false;
+    const onServeSignal = (signal: "SIGINT" | "SIGTERM"): void => {
+      if (shutdownStarted) return;
+      shutdownStarted = true;
+      const rawTimeout = Number(runtimeValue("MLX_BUN_SHUTDOWN_TIMEOUT_MS"));
+      const timeoutMs = Number.isFinite(rawTimeout) && rawTimeout > 0 ? rawTimeout : 120_000;
+      console.log(`\n[server] ${signal}: flushing prompt-cache durability`);
+      void shutdownServer(server, timeoutMs).then((result) => {
+        if (!result.durability.durable) {
+          const d = result.durability;
+          console.warn(
+            `[server] cache flush incomplete after ${Math.round(d.elapsedMs)} ms: ` +
+              `${d.pendingSnapshots} snapshots, ${d.pendingSpills} spills, ` +
+              `${d.droppedSpills} dropped, ${d.failedSpills} failed`,
+          );
+        }
+        process.exit(0);
+      }).catch((error) => {
+        console.warn(`[server] shutdown flush failed: ${(error as Error).message}`);
+        process.exit(0);
+      });
+    };
+    process.once("SIGINT", () => onServeSignal("SIGINT"));
+    process.once("SIGTERM", () => onServeSignal("SIGTERM"));
     const shownHost = rt.serverOptions.hostname ?? "localhost";
     console.log();
     box([
@@ -1682,7 +1801,7 @@ switch (cmd) {
     // three entry points (raw / OpenAI API / chat UI) differ only in how params
     // are POPULATED; here every param is explicit. Decode-path levers mirror
     // serve so you can pin the route (mlx-lm compat = --l1).
-    // --l1/--l2 + per-fork overrides set the decode env levers; the
+    // --l1/--l2 + per-fork overrides set the decode runtime options; the
     // returned KV scheme is applied below so `generate --l2` runs the same
     // quantized-KV route (and produces the same tokens) as `serve --l2`.
     const route = applyDecodeRoute();
@@ -1706,13 +1825,12 @@ switch (cmd) {
     // bf16, same silent fallback as serve), N = uniform bits from decode
     // start. Unset (no tier, no --kv-quant) keeps generate's historical
     // default: bf16 bit-exact greedy (generateText's parity path).
-    const kvScheme =
-      route.turboQuant ? { turboQuant: route.turboQuant, quantizedKvStart: 0 }
-      : route.kvQuant === "off" ? {}
-      : route.kvQuant === "config"
-        ? (tm.config.kvQuant?.length ? { kvConfig: tm.config.kvQuant } : {})
-      : typeof route.kvQuant === "number"
-        ? { kvBits: route.kvQuant, quantizedKvStart: 0 }
+    const kvScheme = route.turboQuant || route.kvQuant !== undefined
+      ? resolveKvScheme({
+          override: route.kvQuant,
+          turboQuant: route.turboQuant,
+          config: tm.config.kvQuant,
+        }).generationOptions
       : undefined;
     const num = (n: string): number | undefined => { const v = opt(n); return v == null ? undefined : Number(v); };
     const text = await generateText(tm, prompt, {
@@ -1794,14 +1912,14 @@ switch (cmd) {
   case "benchmark": {
     const { banner, step, box, style } = await import("./tui");
     banner(pkg.version);
-    serverRuntimeFlags(); // env levers (--compiled-decode, --fused-sdpa, ...) apply to the run
+    serverRuntimeFlags(); // decode options (--compiled-decode, --fused-sdpa, ...) apply to the run
     const { m, picked } = await resolveModelAuto(opt("model") ?? positional(0) ?? opt("query"));
     const tokens = Number(opt("tokens", "256"));
     const runs = Number(opt("runs", "3"));
     const promptTokens = Number(opt("prompt-tokens", "0"));
     const kvMode = opt("kv-quant", "off")!;
     console.log(`  ${style.dim("measures THIS machine, OUR stack only — no other systems needed.")}`);
-    console.log(`  ${style.dim("numbers on a loaded machine are not quotable; ./benchmark.sh is the preflight-gated harness.")}`);
+    console.log(`  ${style.dim("numbers on a loaded machine are not quotable; bun scripts/bench-serve.ts all is the preflight-gated harness.")}`);
     const sNative = step("native runtime");
     await ensureNative(sNative);
     sNative.done("native runtime ready");
@@ -1846,7 +1964,7 @@ switch (cmd) {
       : kvMode !== "off" ? { kvBits: Number(kvMode), quantizedKvStart: 0 } : {};
 
     // Warmup materializes weights so prefill timing measures prefill,
-    // not lazy page-in (same protocol as scripts/bench.ts).
+    // not lazy page-in (same protocol as bench.ts (deleted 2026-08-23; git history)).
     const sWarm = step("warmup");
     {
       const wCache = model.makeCache();
@@ -1921,7 +2039,7 @@ switch (cmd) {
     }
     const { style } = await import("./tui");
     if (rows.length === 0) {
-      console.log("no eval runs recorded yet — ./benchmark.sh writes them");
+      console.log("no eval runs recorded yet — bun scripts/bench-serve.ts all writes them");
       break;
     }
     const ago = (ts: number) => {
@@ -1993,7 +2111,7 @@ switch (cmd) {
     // already run their own pi connect it with `mlx-bun harness pi`.
     const OURS_VAL = new Set([
       "--query", "-q", "--model", "--port", "--host", "--memory-budget", "--kv-budget", "--prompt-cache",
-      "--ssd-cache", "--ssd-cache-max", "--ssd-demote-idle", "--kv-quant", "--unix",
+      "--ssd-cache", "--ssd-cache-max", "--ssd-demote-idle", "--generation-checkpoint", "--kv-quant", "--unix",
       "--batch", "--decode-concurrency", "--adapter", "--adapter-path",
       "--draft-model", "--num-draft-tokens", "--draft-kind", "--ngram-max", "--ngram-min",
       "--compiled-decode", "--compiled-activations", "--fused-sdpa", "--thinking",
@@ -2446,6 +2564,12 @@ switch (cmd) {
       console.error(`--candidate-bits expects a comma list of integers in [2, 8] (got "${opt("candidate-bits")}")`);
       process.exit(1);
     }
+    const rotateWeights = flag("rotate-weights");
+    const rotationSeed = Number(opt("rotation-seed", "42"));
+    if (!Number.isInteger(rotationSeed)) {
+      console.error(`--rotation-seed expects an integer (got "${opt("rotation-seed")}")`);
+      process.exit(1);
+    }
     const mlxPath = opt("mlx-path", "mlx_model")!;
     const { existsSync } = await import("node:fs");
     if (existsSync(mlxPath)) {
@@ -2487,7 +2611,7 @@ switch (cmd) {
         ? `quantizing (mixed, target ${targetBpw} bpw — sensitivity sweep, ~minutes)`
         : `quantizing (${qBits}-bit, group ${qGroup})`,
     );
-    const { quantizeModelDir } = await import("./quantize");
+    const { quantizeModelDir, automaticRotationWeightTransform } = await import("./quantize");
     try {
       const r = await quantizeModelDir(
         srcDir,
@@ -2500,6 +2624,9 @@ switch (cmd) {
           ...(candidateBits ? { candidateBits } : {}),
           ...(opt("calibration-mix") ? { calibrationMix: opt("calibration-mix")! } : {}),
           ...(opt("n-calibration") ? { nCalibration: Number(opt("n-calibration")) } : {}),
+          ...(rotateWeights
+            ? { weightTransform: automaticRotationWeightTransform({ seed: rotationSeed }) }
+            : {}),
         },
         (e) => sQ.update(e.message),
       );
@@ -2511,6 +2638,7 @@ switch (cmd) {
         `source    ${style.dim(srcDir)}`,
         `model     ${style.bold(r.outDir)} ${style.dim(`· ${gb(r.write.totalSize)}`)}`,
         `quant     ${style.dim(targetBpw !== undefined ? `mixed ${r.achievedBpw.toFixed(2)} bpw (target ${targetBpw})` : `${qBits}-bit g${qGroup} affine`)}`,
+        ...(rotateWeights ? [`transform ${style.dim(`TurboQuant rotation seed ${rotationSeed}`)}`] : []),
         "",
         `serve it   ${style.accent(`mlx-bun serve ${r.outDir}`)}`,
       ]);

package/src/client.ts

@@ -0,0 +1,66 @@
+import type { CompletionClient } from "./contracts/completion";
+import type { EngineHost } from "./contracts/host";
+
+export interface CompletionCall {
+  readonly body: Readonly<Record<string, unknown>>;
+  readonly route?: "chat/completions" | "completions";
+  readonly signal?: AbortSignal;
+}
+
+export interface CompletionResponse {
+  choices: Array<{ text?: string; message?: { content?: string | null; [key: string]: unknown }; [key: string]: unknown }>;
+  usage?: Record<string, unknown>;
+  [key: string]: unknown;
+}
+
+/** A borrowed transport. Direct services, HTTP and isolated workers receive
+ * the same request and return the same result; POSTs are never retried here. */
+export function createCompletionClient(options: {
+  baseUrl: string;
+  headers?: Readonly<Record<string, string>>;
+  host?: Pick<EngineHost<Request, Response>, "forward">;
+}): CompletionClient<CompletionCall, CompletionResponse> {
+  const baseUrl = options.baseUrl.replace(/\/$/, "");
+  const headers = { ...options.headers, "content-type": "application/json" };
+  const send = options.host ? options.host.forward.bind(options.host) : fetch;
+  return {
+    async complete({ body, route = "chat/completions", signal }) {
+      signal?.throwIfAborted();
+      const response = await send(new Request(`${baseUrl}/${route}`, {
+        method: "POST", headers, body: JSON.stringify({ ...body, stream: false }), signal,
+      }));
+      if (!response.ok) {
+        const detail = await response.text();
+        throw new Error(`completion failed (${response.status}): ${detail}`);
+      }
+      const result = await response.json() as CompletionResponse;
+      if (!Array.isArray(result.choices)) throw new Error("completion response has no choices");
+      return result;
+    },
+  };
+}
+
+/** Adapts an in-process request handler. Closing waits for active handlers;
+ * a handler owns any response stream it returns. */
+export function createDirectHost(
+  handler: (request: Request) => Promise<Response>, shutdown: () => Promise<void> = async () => {},
+): EngineHost<Request, Response> {
+  let closed = false;
+  let closing: Promise<void> | undefined;
+  const active = new Set<Promise<Response>>();
+  return {
+    ready: Promise.resolve(),
+    forward(request) {
+      if (closed) return Promise.reject(new Error("engine host is closed"));
+      if (request.signal.aborted) return Promise.reject(request.signal.reason);
+      const work = Promise.resolve().then(() => handler(request));
+      active.add(work);
+      void work.then(() => active.delete(work), () => active.delete(work));
+      return work;
+    },
+    close() {
+      closed = true;
+      return closing ??= (async () => { await Promise.allSettled(active); await shutdown(); })();
+    },
+  };
+}