mlx-bun 0.2.0 → 0.4.0

package/src/backends/mlx/transitioning-kv-rows.ts

@@ -0,0 +1,165 @@
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import type { BatchableCache, Cache, Mask } from "../../model/gemma4-base";
+import { KvTensorRows } from "./kv-tensor-rows";
+import { leaseCacheStates, minimumReusableOffset } from "./state-views";
+import { withResource, disposeResources } from "../../engine/resources";
+import { cloneKvCaches } from "../../kv-store";
+
+/** Encoded layouts retain physical row alignment across precision changes. */
+export interface TransitionedKvLayout extends BatchableCache {
+  alignRows?(leftPad: readonly number[]): void;
+}
+export interface SpeculativeTransitionedKvLayout extends TransitionedKvLayout {
+  specRoundBegin(): void;
+  specRoundCommit(): void;
+  specRoundRollback(keep: number | readonly number[]): void;
+}
+export interface KvRowTransition<Layout extends TransitionedKvLayout> {
+  readonly signature: string;
+  maintain(rows: Cache[]): void;
+  converted(row: Cache): boolean;
+  makeLayout(): Layout;
+  packRows?(layout: Layout, rows: readonly Cache[]): void;
+  prepareRows?(rows: Cache[]): Cache[];
+  extractRow?(row: Cache): Cache;
+  rollbackRow?(row: Cache, before: number, keep: number, preserveColumns: boolean): void;
+}
+
+
+export interface TransitioningKvPositions {
+  readonly rowOffsets: readonly number[];
+  readonly leftPad: readonly number[];
+  readonly offset: number;
+  readonly batchSize: number | null;
+  readonly ropeOffsetArr: MlxArray | undefined;
+  sync(rows: readonly Cache[], preserve: boolean): void;
+  filterRows(keep: readonly number[]): void;
+  makeMask(tokens: number, window: number | null): Mask;
+  dispose(): void;
+}
+export class FullKvPositions extends KvTensorRows implements TransitioningKvPositions {
+  sync(rows: readonly Cache[], preserve: boolean): void {
+    this.mergeRows(preserve
+      ? [{ planes: [], rowOffsets: rows.map(row => row.offset), leftPad: [...this.leftPad] }]
+      : rows.map(row => ({ planes: [], rowOffsets: [row.offset], leftPad: [0] })));
+  }
+}
+
+/** Owns membership, committed positions and precision transitions. Codec and
+ * attention adapters own tensor access; scheduling and persistence use the
+ * existing batch and checkpoint contracts. Row conversion never examines an
+ * uncommitted speculative suffix. */
+export abstract class TransitioningKvRows<Layout extends TransitionedKvLayout> implements BatchableCache {
+  protected readonly positions: TransitioningKvPositions;
+  protected rows: Cache[] = [];
+  protected packed?: Layout;
+  #prefilling = false;
+
+  readonly stateNeedsDispose = true;
+  constructor(readonly transition: KvRowTransition<Layout>, row?: Cache, positions: TransitioningKvPositions = new FullKvPositions()) {
+    this.positions = positions;
+    if (row) { this.rows = transition.prepareRows?.([row]) ?? [row]; this.syncPositions(); }
+  }
+  signature(): string { return this.transition.signature; }
+  get prefillMaintenance() { return this; }
+  beginPrefill(): void { this.#prefilling = true; }
+  commitPrefill(indices: readonly number[]): void {
+    if (this.packed) return;
+    for (const index of indices) {
+      const row = [this.rows[index]!];
+      try { this.transition.maintain(row); }
+      finally { this.rows[index] = row[0]!; }
+    }
+  }
+  endPrefill(): void { this.#prefilling = false; }
+  get minimumReusableOffset(): number { return this.packed?.minimumReusableOffset ?? minimumReusableOffset(this.rows); }
+  get rowOffsets(): readonly number[] { return this.packed?.rowOffsets ?? this.positions.rowOffsets; }
+  get leftPad(): readonly number[] { return this.packed?.leftPad ?? this.positions.leftPad; }
+  get offset(): number { return this.packed?.offset ?? this.positions.offset; }
+  get batchSize(): number | null { return this.packed?.batchSize ?? this.positions.batchSize; }
+  get ropeOffsetArr(): MlxArray | undefined { return this.packed ? this.packed.ropeOffsetArr : this.positions.ropeOffsetArr; }
+  bytesPerToken(): number { return this.packed?.bytesPerToken?.() ?? Math.max(0, ...this.rows.map(row => row.bytesPerToken?.() ?? 0)); }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * tokens; }
+  makeMask(tokens: number, window: number | null): Mask { return (this.packed ?? this.positions).makeMask(tokens, window); }
+  isTrimmable(): boolean { return this.packed?.isTrimmable() ?? this.rows.every(row => row.isTrimmable()); }
+  trim(count: number): void {
+    if (this.packed) this.packed.trim(count);
+    else { for (const row of this.rows) row.trim(count); this.syncPositions(true); }
+  }
+  protected syncPositions(preserve = false): void { this.positions.sync(this.rows, preserve); }
+  protected advance(): void {
+    if (this.packed) return;
+    if (!this.#prefilling) this.transition.maintain(this.rows);
+    if (this.rows.length && this.rows.every(row => this.transition.converted(row))) {
+      const packed = this.transition.makeLayout();
+      try {
+        if (this.transition.packRows) this.transition.packRows(packed, this.rows);
+        else packed.mergeRows(this.rows);
+        packed.alignRows?.(this.positions.leftPad);
+      }
+      catch (error) { packed.dispose(); throw error; }
+      const previous = this.rows; this.rows = []; this.packed = packed;
+      this.positions.dispose(); disposeResources(previous);
+    }
+  }
+  abstract updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray];
+  abstract makeEmptyBatch(): TransitioningKvRows<Layout>;
+  mergeRows(rows: readonly Cache[]): void {
+    const owned: Cache[] = [];
+    try {
+      for (const row of rows) {
+        if (row instanceof TransitioningKvRows) {
+          for (let index = 0; index < row.batchSize!; index++) owned.push(row.extractRow(index));
+        } else owned.push(...cloneKvCaches([row]));
+      }
+    } catch (error) { disposeResources(owned); throw error; }
+    let prepared: Cache[];
+    try { prepared = this.transition.prepareRows?.(owned) ?? owned; }
+    catch (error) { disposeResources(owned); throw error; }
+    this.dispose(); this.rows = prepared; this.syncPositions();
+  }
+  filterRows(keep: readonly number[]): void {
+    if (this.packed) { this.packed.filterRows(keep); return; }
+    const previous = this.rows; this.rows = keep.map(index => previous[index]!);
+    disposeResources(previous.filter((_, index) => !keep.includes(index))); this.positions.filterRows(keep);
+  }
+  extractRow(row: number): Cache { return this.packed?.extractRow(row) ?? (this.transition.extractRow?.(this.rows[row]!) ?? cloneKvCaches([this.rows[row]!])[0]!); }
+  state(): MlxArray[] {
+    return withResource(leaseCacheStates(this.packed ? [this.packed] : this.rows), arrays => {
+      const views: MlxArray[] = [];
+      try { for (const array of arrays) views.push(ops.contiguous(array)); return views; }
+      catch (error) { disposeResources(views); throw error; }
+    });
+  }
+  dispose(): void {
+    this.#prefilling = false;
+    const previous = this.rows; this.rows = [];
+    this.packed?.dispose(); this.packed = undefined; this.positions.dispose(); disposeResources(previous);
+  }
+}
+
+/** Rollback is a method capability, separate from ordinary cache membership. */
+export abstract class SpeculativeTransitioningKvRows<Layout extends SpeculativeTransitionedKvLayout> extends TransitioningKvRows<Layout> {
+  #beforeRound?: number[];
+  specRoundBegin(): void {
+    this.advance(); this.#beforeRound = [...this.rowOffsets]; this.packed?.specRoundBegin();
+  }
+  specRoundCommit(): void {
+    this.packed?.specRoundCommit(); this.#beforeRound = undefined; this.advance();
+  }
+  specRoundRollback(keep: number | readonly number[]): void {
+    if (this.packed) this.packed.specRoundRollback(keep);
+    else {
+      for (const [index, row] of this.rows.entries()) {
+        const count = typeof keep === "number" ? keep : keep[index]!;
+        const before = this.#beforeRound![index]!;
+        if (this.transition.rollbackRow) this.transition.rollbackRow(row, before, count, this.rows.length > 1);
+        else row.trim(row.offset - before - count);
+      }
+      this.syncPositions(true);
+    }
+    this.#beforeRound = undefined; this.advance();
+  }
+  override dispose(): void { this.#beforeRound = undefined; super.dispose(); }
+}

package/src/cli.ts

@@ -35,6 +35,14 @@ import {
 } 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.
@@ -103,13 +111,20 @@ 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>
+                            Queue an in-flight generation snapshot to the SSD
+                            cache every N emitted tokens. Repeating the same
+                            request after restart replays and resumes it.
+                            Requires --ssd-cache; supports ordinary requests
+                            with shared cache layouts.
   --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
@@ -155,8 +170,8 @@ Model & quality:
                             auto-detected — a full same-tokenizer model
                             (mlx_lm.server parity), a Gemma "-assistant"
                             KV-borrowing drafter, or a DSpark checkpoint.
-                            Serial lane only — with --batch N a mounted draft
-                            routes every request serial, like mlx_lm.server.
+                            Supported draft providers use shared batched
+                            execution; ordinary requests can batch too.
   --draft-kind <kind>       Override draft-artifact detection:
                             two-model | assistant | dspark | deepspec | mtp | ngram
                             (deepspec = DeepSeek's released Gemma4DSparkModel
@@ -181,7 +196,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.
@@ -191,12 +206,11 @@ Model & quality:
                             only; sliding-window layers stay bf16)]
   --paged-kv                OPTIONAL vLLM-style paged KV storage (block pool +
                             gather before the stock SDPA; env MLX_BUN_PAGED_KV=1).
-                            v1: serial only (pins --batch 1 unless --batch is
-                            given), Gemma4-family, bf16 — refuses --batch N>1,
-                            --kv-quant, --draft-model; bypasses the prompt
+                            Gemma4-family, bf16; supports --batch N.
+                            Refuses --kv-quant and --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]
@@ -212,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]
@@ -391,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
@@ -475,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
 
@@ -825,14 +840,14 @@ function openChatUi(url: string, hostPort: string): void {
 /** 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. Installs decode runtime options and returns the kv-quant mode. See
- *  docs/design/parity-tier-dag.md + unified-engine-frontier-plan.md.
+ *  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. */
@@ -897,7 +912,7 @@ function applyDecodeRoute(): { kvQuant?: "off" | "config" | number; turboQuant?:
   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;
@@ -940,23 +955,41 @@ 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");
@@ -979,16 +1012,11 @@ 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:
-  // with no explicit --batch, pin --batch 1 (the default is 8); an explicit
-  // --batch N>1 is refused loudly by createServer rather than downgraded.
+  // storage, default off (docs/design/kv-cache.md). The state policy owns
+  // block pools and cache reuse; paging preserves the selected batch size.
   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) {
-      serverOptions.batch = 1;
-      console.log("[paged-kv] serial-only in v1 — pinning --batch 1");
-    }
   }
   // Bind loopback unless asked otherwise (mlx_lm.server parity); --host
   // 0.0.0.0 is the explicit opt-in for LAN exposure. The chat-UI open and
@@ -1026,7 +1054,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,
@@ -1538,7 +1566,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
@@ -1580,7 +1608,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: {
@@ -1596,7 +1624,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;
@@ -1613,25 +1641,6 @@ 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");
@@ -1651,9 +1660,8 @@ switch (cmd) {
     }
     // Speculative decoding (mlx_lm.server parity): --draft-model <path|query>
     // resolves like the main model; --num-draft-tokens defaults to 3 (the
-    // server's default upstream — mlx_lm.generate's is 2). Serial-lane-only:
-    // with --batch N, a mounted draft routes every request serial (upstream
-    // is_batchable = draft is None).
+    // server's default upstream — mlx_lm.generate's is 2). Qualified providers
+    // use shared execution; ordinary requests retain their own placement.
     const draftQuery = opt("draft-model");
     let draftModelDir: string | undefined;
     if (draftQuery !== null) {
@@ -1886,7 +1894,7 @@ switch (cmd) {
     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");
@@ -1931,7 +1939,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();
@@ -2006,7 +2014,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) => {
@@ -2078,7 +2086,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",

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(); })();
+    },
+  };
+}

package/src/config.ts

@@ -75,6 +75,17 @@ export interface QuantSpec {
   bits: number;
   groupSize: number;
   mode: string;
+  /** `mode: "trellis"` (packed trellis-coded weights, src/model/trellis-linear.ts):
+   *  bits = k per weight, groupSize = block T, plus the codec geometry. */
+  trellis?: TrellisSpec;
+}
+
+export interface TrellisSpec {
+  L: number;
+  code: string;
+  /** 1 = coded along the input dim (stored [out, in·k/32]); 0 = along the
+   *  output dim (stored [in, out·k/32], i.e. Wᵀ coded along its last axis). */
+  axis: 0 | 1;
 }
 
 export interface QuantizationConfig {
@@ -89,12 +100,12 @@ export interface KvQuantSpec {
   groupSize: number;
 }
 
-/** TurboQuant scheme (docs/design/turboquant-kv.md): rotation-based KV
+/** TurboQuant scheme (docs/design/turboquant.md): rotation-based KV
  *  quantization, a distinct axis from the uniform/per-layer affine kvQuant
  *  above (composes with it in principle; v1 ships as a standalone CLI-only
  *  runtime lever, same class as uniform kvBits — see
  *  src/generate.ts maybeQuantizeKv). Valid kBits: {2,4,5,8}; vBits:
- *  {2,3,4,5,8} (docs/design/turboquant-kv.md's supported Lloyd-Max tables). */
+ *  {2,3,4,5,8} (docs/design/turboquant.md's supported Lloyd-Max tables). */
 export interface TurboQuantScheme {
   kBits: number;
   vBits: number;
@@ -162,11 +173,17 @@ export function parseQuantization(raw: Record<string, any> | undefined): Quantiz
   for (const [key, value] of Object.entries(raw)) {
     if (key === "bits" || key === "group_size" || key === "mode") continue;
     if (typeof value === "object" && value !== null) {
-      perLayer.set(key, {
+      const spec: QuantSpec = {
         bits: (value as any).bits ?? def.bits,
         groupSize: (value as any).group_size ?? def.groupSize,
         mode: (value as any).mode ?? def.mode,
-      });
+      };
+      const tr = (value as any).trellis;
+      if (spec.mode === "trellis") {
+        if (!tr || typeof tr !== "object") throw new Error(`quantization.${key}: mode "trellis" needs a trellis {L, code, axis} block`);
+        spec.trellis = { L: Number(tr.L), code: String(tr.code ?? "1mad"), axis: tr.axis === 0 ? 0 : 1 };
+      }
+      perLayer.set(key, spec);
     } else if (value === false) {
       // mlx convention: `"layer": false` means not quantized
       perLayer.set(key, { bits: 0, groupSize: 0, mode: "none" });

package/src/contracts/completion.ts

@@ -0,0 +1,16 @@
+import type { Cancellation } from "./generation";
+
+/** Application-facing calls. The implementation owns tokenization, numerical
+ * method and transport; callers provide request data and consume results. */
+export interface CompletionClient<Request, Result> {
+  complete(request: Request): Promise<Result>;
+}
+
+export interface BatchCompletionClient<Request, Result> extends CompletionClient<Request, Result> {
+  /** Results preserve input order; the implementation chooses execution groups. */
+  completeBatch(requests: readonly Request[]): Promise<Result[]>;
+}
+
+export interface TaskClient<Request, Progress, Result> {
+  run(request: Request, report: (progress: Progress) => void, cancellation?: Cancellation): Promise<Result>;
+}