mlx-bun 0.2.0 → 0.3.0

package/src/backends/mlx/speculative.ts

@@ -0,0 +1,114 @@
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import type { Cache } from "../../model/gemma4";
+import type { RuntimeModel } from "../../model/factory";
+import type { Sampler } from "../../sampler";
+import type { DraftProvider, DraftSource } from "../../spec/source";
+import { runtimeConfig, type RuntimeConfig } from "../../runtime-config";
+import type { MlxModelMemory } from "./autoregressive";
+import type { SpeculativeTransaction } from "../../inference/rollback";
+import { bindCacheRollback } from "./rollback";
+import type { GraphDescriptor } from "../../inference/graph";
+import { bindLegacyDraftTarget } from "./draft-target";
+
+/** Bound target operations and draft construction for one speculative run.
+ * A replacement graph supplies this entire port, including any hidden taps.
+ * The verifier never needs a RuntimeModel or an artifact-family check. */
+export interface MlxSpeculativeBinding {
+  readonly runtime?: RuntimeConfig;
+  readonly memory?: MlxModelMemory;
+  readonly descriptor: GraphDescriptor;
+  readonly eosTokenIds: readonly number[];
+  readonly prefillTailSplit: boolean;
+  makeCache(): Cache[];
+  openDraft(sampler: Sampler, caches: Cache[]): DraftSource;
+  bindRollback(caches: Cache[]): SpeculativeTransaction;
+  forward(ids: MlxArray, caches: Cache[], tapLayers?: number[]):
+    Promise<{ hidden: MlxArray; ctxML: MlxArray | null }>;
+  projectLogits(hidden: MlxArray): MlxArray;
+  /** Establish the implementation's verify kernel context; restore on close. */
+  pinVerify?(): { close(): void };
+}
+
+/** Legacy mutable tap/kernel fields require the gateway's exclusive lease. */
+export function bindLegacySpeculativeModel(model: RuntimeModel, provider: DraftProvider): MlxSpeculativeBinding {
+  const runtime = runtimeConfig();
+  return {
+    runtime,
+    memory: model,
+    descriptor: Object.freeze({ id: `legacy-spec:${model.config.modelType}`, backend: "mlx",
+      graphAbi: "mlx-hidden-bsh-v1", stateAbi: "legacy-cache-array-v1", artifact: "legacy-resident-model" }),
+    eosTokenIds: model.config.eosTokenIds,
+    prefillTailSplit: runtime.flag("MLX_BUN_PREFILL_TAIL_SPLIT", true),
+    makeCache: model.makeCache.bind(model),
+    openDraft: (sampler, caches) => provider.open({ sampler, target: bindLegacyDraftTarget(model, caches) }),
+    bindRollback: bindCacheRollback,
+    forward: (ids, caches, tapLayers) => legacyForwardWithTaps(model, ids, caches, tapLayers),
+    projectLogits: model.logitsFromHidden.bind(model),
+    ...("setSpecKernelPinned" in model ? {
+      pinVerify() {
+        model.setSpecKernelPinned(true);
+        return { close() { model.setSpecKernelPinned(false); } };
+      },
+    } : {}),
+  };
+}
+
+export function assertMlxSpeculativeBinding(binding: MlxSpeculativeBinding): void {
+  const descriptor = binding.descriptor;
+  if (descriptor.backend !== "mlx" || descriptor.graphAbi !== "mlx-hidden-bsh-v1" ||
+      descriptor.stateAbi !== "legacy-cache-array-v1")
+    throw new Error(`speculative binding ${descriptor.id} has an incompatible backend, graph, or state ABI`);
+}
+
+/** Run one target forward, and when `tapLayers` is set (⟹ a Gemma4 target — a
+ *  DSpark-style source), also return the captured multi-layer context
+ *  [1,L,m*H] (tapLayers concatenated on the feature axis; index nLayers is the
+ *  post-finalNorm sentinel). Non-tapping sources get ctxML=null and never
+ *  touch model.hiddenTap. Mirrors generate-dflash.ts forwardTapped. */
+async function legacyForwardWithTaps(
+  model: RuntimeModel,
+  ids: MlxArray,
+  caches: Cache[],
+  tapLayers: number[] | undefined,
+): Promise<{ hidden: MlxArray; ctxML: MlxArray | null }> {
+  if (!tapLayers) {
+    const asyncModel = model as RuntimeModel & {
+      forwardHiddenAsync?: (
+        ids: MlxArray,
+        caches: Cache[],
+      ) => Promise<MlxArray>;
+    };
+    const hidden = typeof asyncModel.forwardHiddenAsync === "function"
+      ? await asyncModel.forwardHiddenAsync(ids, caches)
+      : model.forwardHidden(ids, caches);
+    return { hidden, ctxML: null };
+  }
+  if (!("hiddenTap" in model)) throw new Error("target graph does not support hidden taps");
+  const m = model;
+  const previousTap = m.hiddenTap;
+  const cap = new Map<number, MlxArray>();
+  m.hiddenTap = { layers: new Set(tapLayers), captured: cap };
+  let hidden: MlxArray | null = null;
+  try {
+    hidden = model.forwardHidden(ids, caches);
+    const perLayer = tapLayers.map((li) => {
+      const a = cap.get(li);
+      if (!a) throw new Error(`spec tap: layer ${li} not captured`);
+      return a;
+    });
+    const ctxML = ops.concatAxis(perLayer, 2); // [1,L,m*H]
+    for (const [, a] of cap) a.dispose();
+    cap.clear(); // consumed — the finally must not double-dispose
+    const out = { hidden, ctxML };
+    hidden = null; // ownership returned to the caller
+    return out;
+  } finally {
+    // On any throw (forward mid-capture, missing layer, concat), free the
+    // partially-captured tap tensors and the orphaned hidden. On success both
+    // are already gone (cap cleared, hidden nulled).
+    hidden?.dispose();
+    for (const [, a] of cap) a.dispose();
+    m.hiddenTap = previousTap;
+  }
+}

package/src/backends/mlx/state-views.ts

@@ -0,0 +1,26 @@
+import type { StateView } from "../../contracts/resources";
+import { disposeResources, ownResource } from "../../engine/resources";
+import type { MlxArray } from "../../mlx/array";
+import type { Cache } from "../../model/gemma4-base";
+
+/** Compatibility boundary for legacy state(). Numerical code consumes a lease,
+ * never the legacy ownership marker. No copy, evaluation or fence is added. */
+export function leaseCacheState(cache: Cache): StateView<MlxArray> {
+  return ownResource<readonly MlxArray[]>(cache.state(), cache.stateNeedsDispose ? disposeResources : () => {});
+}
+
+/** Preserve the existing prefill evaluation order: temporary views first,
+ * borrowed state second. Partial acquisition releases every acquired view. */
+export function leaseCacheStates(caches: readonly Cache[]): StateView<MlxArray> {
+  const leases: StateView<MlxArray>[] = [];
+  const close = () => disposeResources(leases.map((lease) => ({ dispose: () => lease.close() })));
+  try {
+    for (const owned of [true, false]) {
+      for (const cache of caches) if (!!cache.stateNeedsDispose === owned) leases.push(leaseCacheState(cache));
+    }
+    return ownResource(leases.flatMap((lease) => [...lease.borrow()]), close);
+  } catch (error) {
+    try { close(); } catch (cleanupError) { throw new AggregateError([error, cleanupError], "state view acquisition failed"); }
+    throw error;
+  }
+}

package/src/backends/mlx/text-engine.ts

@@ -0,0 +1,73 @@
+import { createCompletionMethod } from "../../serve/session-completion-engine";
+import type { ExecutionPlanner, InferenceEngine, Timer } from "../../contracts/generation";
+import { createInferenceEngine, type EngineOptions } from "../../engine/engine";
+import { throwIfCancelled } from "../../engine/cancellation";
+import type { GenerateOptions, GenerateStats } from "../../generate";
+import type { GenerationGateway, RequestShape } from "../../serve/generation-gateway";
+
+/** Data-only API for text requests whose policy/shape was already
+ * resolved by the caller. Native request resources still belong to the prepared
+ * completion API; they are deliberately absent from this token API. */
+export type TextGenerationOptions = Omit<GenerateOptions,
+  "signal" | "cache" | "grammar" | "fill" | "adapters" | "promptEmbeddings" |
+  "imageMask" | "multimodalMask" | "visionPixels" | "onPrefillDone" |
+  "onDecodeCheckpoint" | "initialPendingToken" | "initialGeneratedTokens" |
+  "originalPromptTokens" | "snapshotAt" | "checkpointEveryTokens"
+>;
+
+export interface TextGenerationRequest {
+  readonly promptIds: readonly number[];
+  readonly options: TextGenerationOptions;
+  readonly shape: RequestShape;
+}
+
+export const hostTimer: Timer = {
+  after(milliseconds, callback) {
+    const timer = setTimeout(callback, milliseconds);
+    return () => { clearTimeout(timer); };
+  },
+};
+
+/** Reuses the gateway's placement, exclusive lease, prompt cache, batching,
+ * and quant-specific execution. The portable session never imports them.
+ * Closing this engine cancels its sessions; it does not dispose a shared gateway. */
+export function createTextInferenceEngine(
+  gateway: Pick<GenerationGateway, "place" | "run">,
+  options: Omit<EngineOptions, "timer"> & { timer?: Timer } = {},
+): InferenceEngine<TextGenerationRequest, GenerateStats> {
+  const planner: ExecutionPlanner<TextGenerationRequest, GenerateStats> = {
+    async plan(request, cancellation) {
+      throwIfCancelled(cancellation);
+      if (request.shape.hasVision || request.shape.hasAdapters || request.shape.hasGrammar)
+        throw new Error("native request resources require the prepared completion API");
+      const placement = gateway.place(request.shape, request.options);
+      const maxTokens = request.options.maxTokens ?? 512;
+      return {
+        id: `text:${placement.mechanism}`, outputTokenLimit: maxTokens,
+        method: createCompletionMethod(gateway, {
+          prompt: [...request.promptIds], options: request.options,
+          shape: request.shape, placement,
+        }),
+      };
+    },
+  };
+  const engine = createInferenceEngine(planner, { ...options, timer: options.timer ?? hostTimer });
+  return {
+    open(request, control) {
+      // Snapshot before demand-start: callers may reuse/mutate request arrays.
+      // TextGenerationOptions contains data only; no native handles or callbacks.
+      const snapshot = structuredClone(request);
+      if (!snapshot.promptIds.length || snapshot.promptIds.some((id) => !Number.isInteger(id) || id < 0 || id > 0xffffffff))
+        return Promise.reject(new Error("prompt must contain valid token ids"));
+      const maxTokens = snapshot.options.maxTokens ?? 512;
+      if (!Number.isSafeInteger(maxTokens) || maxTokens <= 0)
+        return Promise.reject(new Error("text generation requires a positive token limit"));
+      if ((snapshot.options.topLogprobs ?? 0) > (options.maxTopLogprobs ?? 20))
+        return Promise.reject(new Error("top logprobs exceed session delivery capacity"));
+      if (control.output === "collect" && (snapshot.options.logprobs || snapshot.options.topLogprobs))
+        return Promise.reject(new Error("token logprobs require stream consumption"));
+      return engine.open(snapshot, control);
+    },
+    close: () => engine.close(),
+  };
+}

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,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
@@ -181,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.
@@ -196,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]
@@ -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,7 +1012,7 @@ 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") || runtimeValue("MLX_BUN_PAGED_KV") === "1") {
@@ -1026,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,
@@ -1538,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
@@ -1580,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: {
@@ -1596,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;
@@ -1886,7 +1919,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 +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();
@@ -2006,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) => {
@@ -2078,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",

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>;
+}

package/src/contracts/execution.ts

@@ -0,0 +1,49 @@
+/** Request facts only. Native resources remain with the preparation owner. */
+export interface ExecutionRequirements {
+  readonly hasVision: boolean;
+  readonly hasAdapters: boolean;
+  readonly hasRepetitionPenalty: boolean;
+  readonly userSeed: boolean;
+  readonly kvQuant: boolean;
+  readonly turboQuant: boolean;
+  readonly hasLogitsExtras: boolean;
+  readonly hasGrammar: boolean;
+  readonly wantsLogprobs: boolean;
+  readonly hasDraft: boolean;
+}
+
+export interface ExecutionCapabilities {
+  readonly method: "autoregressive" | "denoising";
+  readonly continuous: boolean;
+  readonly quantizedBatch: boolean;
+  readonly grammarBatch: boolean;
+  readonly checkpoints: boolean;
+  /** Model-qualified speculative execution can retain this request's KV codec. */
+  readonly speculativeKvQuant?: boolean;
+  /** A graph-owned compiled step exists; cache geometry can still decline it. */
+  readonly compiledDecode?: boolean;
+}
+
+export interface ExecutionFeatures {
+  readonly pagedKv: boolean;
+  readonly fill: boolean;
+  readonly compiledDecode?: boolean;
+  readonly grammarJump?: boolean;
+}
+
+/** Selected once; consumers execute/report these values without reselecting. */
+export interface ResolvedExecution {
+  /** Implementation-owned method ID. The built-in planner retains its known
+   * methods; another model may register a different set without editing this contract. */
+  readonly method: string;
+  readonly mechanism: "serial" | "continuous";
+  readonly pagedKv: boolean;
+  readonly promptCache: boolean;
+  readonly checkpoint: boolean;
+  readonly fill: boolean;
+  /** Permission to attempt the bound compiled step, never a promise that
+   * every dynamic row/state shape supports replay. */
+  readonly compiledDecode: boolean;
+  readonly grammarJump: boolean;
+  readonly reasons: readonly string[];
+}