mlx-bun 0.0.9 → 0.0.10

package/src/prompt-cache.ts

@@ -17,6 +17,11 @@ export interface PromptCacheEntry {
   /** Namespace key — adapter spec for LoRA requests ("" = base model).
    *  KV computed under one adapter must never seed another's prefill. */
   ns: string;
+  /** Release hook for entries restored from the SSD cold tier: unmaps the
+   *  backing COW file mapping. MUST run only after `caches` are disposed
+   *  (they may alias file pages); every dispose path here honors that, and
+   *  callers who take() an entry carry the thunk with it. */
+  retain?: () => void;
 }
 
 function cacheBytes(caches: Cache[]): number {
@@ -26,7 +31,7 @@ function cacheBytes(caches: Cache[]): number {
 }
 
 /** Length of the longest common prefix of two token sequences. */
-function commonPrefixLength(a: number[], b: number[]): number {
+export function commonPrefixLength(a: number[], b: number[]): number {
   const n = Math.min(a.length, b.length);
   let i = 0;
   while (i < n && a[i] === b[i]) i++;
@@ -39,9 +44,22 @@ export class PromptCache {
   #clock = 0;
   hits = 0;
   misses = 0;
+  /** Cold-tier spill hook: called with an entry ABOUT to be evicted (still
+   *  alive — the hook may serialize its caches); the cache disposes it after.
+   *  Failures are the hook's problem; eviction proceeds regardless. */
+  readonly #spill: ((entry: PromptCacheEntry) => void) | null;
 
-  constructor(maxBytes: number) {
+  constructor(maxBytes: number, spill: ((entry: PromptCacheEntry) => void) | null = null) {
     this.maxBytes = maxBytes;
+    this.#spill = spill;
+  }
+
+  #disposeEntry(entry: PromptCacheEntry, spillFirst: boolean): void {
+    if (spillFirst && this.#spill) {
+      try { this.#spill(entry); } catch {}
+    }
+    for (const c of entry.caches) c.dispose();
+    entry.retain?.();
   }
 
   get totalBytes(): number {
@@ -88,26 +106,54 @@ export class PromptCache {
     return entry;
   }
 
-  /** Insert (or reinsert) an entry; evicts LRU entries over the byte cap.
-   *  If the entry itself exceeds the cap it is disposed, not stored. */
-  put(tokens: number[], caches: Cache[], ns = ""): void {
+  /** Longest usable prefix length available for `prompt` WITHOUT taking
+   *  (the SSD tier compares this against its own best before restoring). */
+  peekPrefixLen(prompt: number[], ns = ""): number {
+    let best = 0;
+    for (const { entry: e } of this.#entries) {
+      if (e.ns !== ns) continue;
+      const p = Math.min(commonPrefixLength(e.tokens, prompt), prompt.length - 1);
+      if (p <= best) continue;
+      const trimNeeded = e.tokens.length - p; // same usability rule as take()
+      if (trimNeeded > 0 && !e.caches.every((c) => c.isTrimmable())) continue;
+      best = p;
+    }
+    return best;
+  }
+
+  /** Read-only exact-token lookup (the write-behind snapshot path: the
+   *  entry stays owned by the cache; the caller only reads array state,
+   *  under the gateway lock so no generation is mutating it). */
+  findExact(tokens: number[], ns = ""): PromptCacheEntry | null {
+    for (const { entry: e } of this.#entries) {
+      if (e.ns !== ns || e.tokens.length !== tokens.length) continue;
+      if (commonPrefixLength(e.tokens, tokens) === tokens.length) return e;
+    }
+    return null;
+  }
+
+  /** Insert (or reinsert) an entry; evicts LRU entries over the byte cap
+   *  (spilling each to the cold tier first, when one is attached). If the
+   *  entry itself exceeds the cap it is spilled + disposed, not stored. */
+  put(tokens: number[], caches: Cache[], ns = "", retain?: () => void): void {
     const bytes = cacheBytes(caches);
+    const entry: PromptCacheEntry = { tokens, caches, ns, retain };
     if (bytes > this.maxBytes) {
-      for (const c of caches) c.dispose();
+      this.#disposeEntry(entry, true);
       return;
     }
-    this.#entries.push({ entry: { tokens, caches, ns }, bytes, lastUsed: ++this.#clock });
+    this.#entries.push({ entry, bytes, lastUsed: ++this.#clock });
     while (this.totalBytes > this.maxBytes && this.#entries.length > 1) {
       let lruIdx = 0;
       for (let i = 1; i < this.#entries.length; i++)
         if (this.#entries[i]!.lastUsed < this.#entries[lruIdx]!.lastUsed) lruIdx = i;
       const [evicted] = this.#entries.splice(lruIdx, 1);
-      for (const c of evicted!.entry.caches) c.dispose();
+      this.#disposeEntry(evicted!.entry, true);
     }
   }
 
   clear(): void {
-    for (const e of this.#entries) for (const c of e.entry.caches) c.dispose();
+    for (const e of this.#entries) this.#disposeEntry(e.entry, false);
     this.#entries = [];
   }
 }

package/src/registry.ts

@@ -465,7 +465,10 @@ function visionConfigTypeOf(config: Record<string, any>): string | null {
   return typeof vt === "string" && vt.endsWith("_vision") ? vt : null;
 }
 
-async function scanSnapshot(dir: string, repoId: string): Promise<ModelRecord | null> {
+/** Build a ModelRecord straight from a snapshot directory (also the scan()
+ *  workhorse). Exported for explicit-path loads (`serve --model <dir>`),
+ *  which must not depend on the registry DB having scanned that dir. */
+export async function scanSnapshot(dir: string, repoId: string): Promise<ModelRecord | null> {
   const configPath = join(dir, "config.json");
   if (!existsSync(configPath)) return null;
 

package/src/serve/batch-scheduler.ts

@@ -21,9 +21,12 @@
 // mask). Full layers share one leftPad/offset (all rows advance together); the
 // rotating caches self-track per-row leftPad/offset as the ring wraps. The
 // per-row absolute position stays consistent across both (full: offset-leftPad;
-// rot: offsetArr) — see docs/design/parallel-slots.md. Models whose caches
-// don't support the dynamic-B ops (e.g. Qwen3.5's SSMCache) never reach this
-// scheduler — GenerationGateway.willBatch gates on cache capability.
+// rot: offsetArr) — see docs/design/parallel-slots.md. Hybrid gated-DeltaNet
+// models (Qwen3.5) add "ssm" layers: SSMCache state is plain [B,...] with no
+// temporal axis and no padding (rows solo-prefill unpadded, decode feeds one
+// real token per row), so merge/filter are B-axis concat/take and the cache
+// passes through the step unwrapped (no mask, no per-row RoPE — full layers
+// carry positions). Gate: GenerationGateway.willBatch on cache capability.
 //
 // Engine mechanics (the serial decode loop's hygiene, transplanted —
 // batching-v2-plan step 3):
@@ -57,12 +60,23 @@ import { clearCache } from "../mlx/ffi";
 import { KVCache, RotatingKVCache, type Cache } from "../model/gemma4-base";
 import { BatchedDecodeMaskCache, mergeKVRows, filterKVRows } from "../model/batched-mask";
 import { BatchedRotatingCache } from "../model/batched-rotating";
+import { SSMCache } from "../model/qwen3-delta";
 import type { RuntimeModel } from "../model/factory";
 
 /** Decode-pipeline kill switch (read once at load, like the serial loop's
  *  MLX_BUN_COMPILED_DECODE): 1 ⇒ read each step's tokens synchronously. */
 const NO_PIPELINE = process.env.MLX_BUN_BATCH_NO_PIPELINE === "1";
 
+// NOTE — oMLX-style "adaptive burst decode" (several steps per event-loop
+// yield, omlx/engine_core.py _step_burst) was ported here 2026-07-02 and
+// REFUTED by measurement: on this runtime it REGRESSED cpm5 B=4 aggregate
+// 345→289 tok/s, batch-lane B=1 149→121, and TTFT ~+100 ms (first-token SSE
+// flush waits out the burst budget). Their win exists because each Python
+// step hand-off ping-pongs the GIL with asyncio/uvicorn (~1 ms/token); Bun's
+// setImmediate hop costs microseconds, so bursting here only delays socket
+// flushes. Don't re-add without new evidence — the per-yield step below is
+// the measured optimum (docs/design/batching-perf-path.md P4 notes).
+
 /** A token sampler for one row: (logits [1,V], step) → token array [1] on
  *  device. Greedy is `(l) => ops.argmaxAxis(l, -1)`; richer closures fold in
  *  temperature / top-p / logits processors + history (built per request from
@@ -127,7 +141,7 @@ export interface BatchSchedulerOptions {
   prefillChunkSize?: number;
 }
 
-type LayerInner = KVCache | BatchedRotatingCache;
+type LayerInner = KVCache | BatchedRotatingCache | SSMCache;
 type Row1 = { keys: MlxArray; values: MlxArray };
 
 export class BatchScheduler {
@@ -146,7 +160,7 @@ export class BatchScheduler {
   readonly #lock: ExclusiveLock | undefined;
   readonly #admissionHeld: (() => boolean) | undefined;
   readonly #prefillChunkSize: number;
-  readonly #kinds: ("full" | "rot")[]; // per-layer attention type
+  readonly #kinds: ("full" | "rot" | "ssm")[]; // per-layer attention type
   readonly #rotMaxSize: number[]; // per-layer sliding window (rot layers only)
 
   constructor(private readonly model: RuntimeModel, opts: BatchSchedulerOptions) {
@@ -155,7 +169,9 @@ export class BatchScheduler {
     this.#admissionHeld = opts.admissionHeld;
     this.#prefillChunkSize = Math.max(1, Math.floor(opts.prefillChunkSize ?? 2048));
     const proto = model.makeCache(); // fresh caches hold no buffers
-    this.#kinds = proto.map((c) => (c instanceof RotatingKVCache ? "rot" : "full"));
+    this.#kinds = proto.map((c) =>
+      c instanceof RotatingKVCache ? "rot" : c instanceof SSMCache ? "ssm" : "full",
+    );
     this.#rotMaxSize = proto.map((c) => (c instanceof RotatingKVCache ? c.maxSize : 0));
     for (const c of proto) c.dispose();
   }
@@ -314,6 +330,16 @@ export class BatchScheduler {
     const newInners: LayerInner[] = [];
     let newFullPad = this.#fullLeftPad;
     for (let layer = 0; layer < this.#kinds.length; layer++) {
+      if (this.#kinds[layer] === "ssm") {
+        // No temporal axis, no left-pad: B-axis concat of the state slots.
+        // mergeRows steals the solo arrays when the batch starts cold, so the
+        // unconditional p.solo dispose below stays safe either way.
+        newInners.push(SSMCache.mergeRows(
+          (prev?.[layer] as SSMCache | undefined) ?? null,
+          p.solo[layer] as SSMCache,
+        ));
+        continue;
+      }
       const soloC = p.solo[layer] as KVCache | RotatingKVCache;
       const [sk, sv] = soloC.temporalView();
       const newRow: Row1 = { keys: sk, values: sv };
@@ -391,10 +417,13 @@ export class BatchScheduler {
     let nextToks: MlxArray | null = null;
     if (anyLive) {
       // Per-layer forward cache: rot layers use the persistent
-      // BatchedRotatingCache directly; full layers get a fresh
-      // BatchedDecodeMaskCache wrapper.
+      // BatchedRotatingCache directly; ssm layers are already [B,...] state
+      // with no padding (no mask, no per-row RoPE — pass through); full
+      // layers get a fresh BatchedDecodeMaskCache wrapper.
       const fwd: Cache[] = inners.map((c) =>
-        c instanceof BatchedRotatingCache ? c : new BatchedDecodeMaskCache(c, B, this.#fullLeftPad, null),
+        c instanceof BatchedRotatingCache || c instanceof SSMCache
+          ? c
+          : new BatchedDecodeMaskCache(c, B, this.#fullLeftPad, null),
       );
       try {
         const ids = this.#pendingToks
@@ -524,6 +553,9 @@ export class BatchScheduler {
       if (inner instanceof BatchedRotatingCache) {
         inner.filter(keep); // in-place (mutates + drops rows + reduces padding)
         out.push(inner);
+      } else if (inner instanceof SSMCache) {
+        inner.filter(keep); // in-place B-axis take on both state slots
+        out.push(inner);
       } else {
         const [k0, v0] = inner.temporalView();
         const f = filterKVRows(k0, v0, keep);

package/src/serve/generation-gateway.ts

@@ -35,11 +35,13 @@
 // in v1 (cachedTokens = 0).
 
 import { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
 import type { RuntimeModel } from "../model/factory";
 import { DiffusionGemmaModel } from "../model/diffusion-gemma";
 import { KVCache, RotatingKVCache } from "../model/gemma4-base";
+import { SSMCache } from "../model/qwen3-delta";
 import type { GenerateOptions, GenerateStats, TokenLogprobs } from "../generate";
-import { makeSampler, toLogprobs } from "../sampler";
+import { makeSampler, makeLogitsProcessors, toLogprobs } from "../sampler";
 import { BatchScheduler } from "./batch-scheduler";
 
 /** Async mutex: acquire() resolves to a release fn; releases run FIFO. */
@@ -126,13 +128,18 @@ export class GenerationGateway {
 
   /** Cache-capability gate (mirrors mlx-lm server.py's all-caches-have-merge
    *  check): the scheduler's dynamic-B ops (temporalView/merge/filter) exist
-   *  only on KVCache and RotatingKVCache. A hybrid model (Qwen3.5's SSMCache)
-   *  routes serial until the batched-SSM port (batching-v2-plan item h). */
+   *  on KVCache, RotatingKVCache, and SSMCache (hybrid gated-DeltaNet models
+   *  — Qwen3.5 — batch via SSMCache.mergeRows/filter; MLX_BUN_BATCH_SSM=0 is
+   *  the kill switch back to serial routing). */
   #modelCachesBatchable(): boolean {
     if (this.#cacheBatchable === null) {
+      const ssmOk = process.env.MLX_BUN_BATCH_SSM !== "0";
       const proto = this.model.makeCache(); // fresh caches hold no buffers
       this.#cacheBatchable = proto.every(
-        (c) => c instanceof KVCache || c instanceof RotatingKVCache,
+        (c) =>
+          c instanceof KVCache ||
+          c instanceof RotatingKVCache ||
+          (ssmOk && c instanceof SSMCache),
       );
       for (const c of proto) c.dispose();
     }
@@ -146,11 +153,15 @@ export class GenerationGateway {
     if (this.model instanceof DiffusionGemmaModel) return false;
     if (!this.batchingEnabled) return false;
     if (!this.#modelCachesBatchable()) return false;
+    // repetitionPenalty / logits extras (min_p, XTC, logit_bias,
+    // presence/frequency penalties) batch: the per-row sampler folds in the
+    // same makeSampler/makeLogitsProcessors closures the serial lane runs,
+    // over a per-row device-side token history. (Load-bearing beyond opt-in
+    // knobs: some models — Qwen3.5 — ship a default repetition penalty in
+    // generation_config.json, which used to route EVERY request serial.)
     return (
       !shape.hasVision &&
       !shape.hasAdapters &&
-      !shape.hasRepetitionPenalty &&
-      !shape.hasLogitsExtras &&
       !shape.wantsLogprobs &&
       !shape.userSeed &&
       !shape.kvQuant
@@ -187,27 +198,75 @@ export class GenerationGateway {
     vision: Vision | undefined,
     shape: RequestShape,
   ): Promise<GenerateStats> {
-    if (!this.willBatch(shape))
-      return this.runExclusive(() => this.serialRun(promptIds, options, onToken, vision));
+    if (!this.willBatch(shape)) {
+      // Serial decode is an unbroken microtask chain (FFI + generator
+      // resumes): without a periodic macrotask hop the event loop starves for
+      // the WHOLE generation — /stats, /health, and new-connection accepts
+      // stall until the request finishes (measured 2.5 s on a 512-token cpm5
+      // run). Rate-limited to ≥25 ms — the SSE flush hop's budget, ~0.1% of
+      // decode — and applied HERE so every serial caller (streaming or not,
+      // all four API surfaces) is covered. The batch lane needs nothing: its
+      // drive loop yields per step (measured /stats ~8 ms mid-generation).
+      let lastHop = performance.now();
+      const hop = (): Promise<void> => new Promise<void>((r) => setImmediate(r));
+      const hoppingOnToken: OnToken = (token, lp) => {
+        const r = onToken(token, lp);
+        if (r === false) return false;
+        if (performance.now() - lastHop < 25) return r;
+        lastHop = performance.now();
+        if (r instanceof Promise) return r.then((v) => (v === false ? false : hop().then(() => v)));
+        return hop().then(() => r);
+      };
+      return this.runExclusive(() => this.serialRun(promptIds, options, hoppingOnToken, vision));
+    }
 
-    // Per-row sampler, mirroring generate()'s sampleStep (no logits processors:
-    // repetition penalty was excluded by willBatch). Greedy (temperature 0) is
-    // argmax; temp>0 uses this request's own seed → independent per-row RNG.
+    // Per-row sampler, mirroring generate()'s sampleStep: logits processors
+    // (logit_bias / repetition / presence / frequency penalties — the same
+    // device-side closures the serial lane runs) fold in over a per-row token
+    // history seeded with the prompt, then the sampler. Greedy (temperature 0)
+    // is argmax; temp>0 uses this request's own seed → independent per-row RNG.
     const sampler = makeSampler(options);
+    const processors = makeLogitsProcessors(options);
+    // Device-side history (prompt + generated), maintained only when a
+    // processor needs it — generate()'s pushHistory, one per row. The sample
+    // closure is called once per sampled token in order, so extending it
+    // inside the closure keeps history exact.
+    let history: MlxArray | null = null;
+    if (processors.length > 0) history = ops.fromInt32(promptIds, [promptIds.length]);
     const sample = (logits1V: MlxArray, step: number): MlxArray => {
-      const lp = toLogprobs(logits1V);
+      let cur = logits1V; // caller-owned; only intermediates are disposed here
+      for (const p of processors) {
+        const next = p(history, cur);
+        if (cur !== logits1V) cur.dispose();
+        cur = next;
+      }
+      const lp = toLogprobs(cur);
+      if (cur !== logits1V) cur.dispose();
       const tok = sampler(lp, step);
       lp.dispose();
+      if (history) {
+        const t1 = ops.reshape(tok, [1]);
+        const prev = history;
+        history = ops.concatAxis([prev, t1], 0);
+        prev.dispose();
+        t1.dispose();
+      }
       return tok;
     };
 
-    const st = await this.#ensureScheduler().submit({
-      promptIds,
-      maxTokens: options.maxTokens ?? 512,
-      eosTokenIds: options.eosTokenIds ?? this.model.config.eosTokenIds,
-      sample,
-      onToken,
-    });
+    let st;
+    try {
+      st = await this.#ensureScheduler().submit({
+        promptIds,
+        maxTokens: options.maxTokens ?? 512,
+        eosTokenIds: options.eosTokenIds ?? this.model.config.eosTokenIds,
+        sample,
+        onToken,
+      });
+    } finally {
+      history?.dispose();
+      history = null;
+    }
 
     return {
       promptTokens: st.promptTokens,

package/src/server.ts

@@ -41,6 +41,8 @@ import {
   TOOL_CALL_END, TOOL_CALL_START,
 } from "./tool-call";
 import { PromptCache } from "./prompt-cache";
+import { SsdCacheStore } from "./ssd-cache";
+import { configFingerprint } from "./model/fingerprint";
 import {
   anthropicToChatBody, chatJsonToAnthropic, translateOpenAiSse,
   type AnthropicRequest,
@@ -123,6 +125,16 @@ export interface ServerOptions {
    *  request's explicit `adapter` (including "none") always wins. Hot-swap
    *  via /v1/adapters is unchanged. */
   defaultAdapter?: string;
+  /** SSD cold tier for the prompt/KV cache (`--ssd-cache <dir>`): prefix
+   *  KV survives RAM eviction AND server restarts — the coding-agent
+   *  long-context TTFT win (docs/design/ssd-kv-cold-tier.md). Off unless
+   *  a directory is given. Serial lane only (where the prompt cache lives). */
+  ssdCacheDir?: string;
+  /** Byte cap for the SSD tier (default 32 GiB). */
+  ssdCacheMaxBytes?: number;
+  /** Verify every tensor hash on restore (`--ssd-cache-verify`) — reads all
+   *  bytes eagerly, defeating lazy fault-in; integrity paranoia only. */
+  ssdCacheVerify?: boolean;
 }
 
 export interface ServerContext {
@@ -976,7 +988,63 @@ export function createServer(
   // lane (see GenerationGateway.willBatch).
   const batch = Math.max(1, Math.floor(serverOptions.batch ?? 1));
 
-  const promptCache = new PromptCache(serverOptions.promptCacheBytes ?? 2e9);
+  // KV-quant scheme, resolved once: kv_config.json by default (optiq
+  // serve's headline behavior), overridable to uniform bits or off.
+  // KV quant scheme. The serial default is the model's mixed-precision config
+  // (optiq parity). But `--batch N` is a bf16 continuous-batching MODE (= mlx-lm
+  // B=N parity), so when KV quant is left UNSET under `--batch N` it defaults to
+  // bf16 — the batch path engages out of the box. An EXPLICIT --kv-quant
+  // (config / bits) is still honored, but those requests then route to the
+  // serial path (batched quantized KV is the L2 follow-up); warned below.
+  const configScheme = ctx.kvConfig?.length ? { kvConfig: ctx.kvConfig } : {};
+  const kvScheme: Pick<GenerateOptions, "kvBits" | "kvConfig" | "quantizedKvStart"> =
+    serverOptions.kvQuant === "off" ? {}
+    : serverOptions.kvQuant === "config" ? configScheme
+    : typeof serverOptions.kvQuant === "number"
+      ? { kvBits: serverOptions.kvQuant, quantizedKvStart: 0 }
+    : batch > 1 ? {} // unset + batching → bf16 mode (Option B)
+    : configScheme; // unset + serial → optiq-parity mixed-precision default
+  if (batch > 1 && (kvScheme.kvConfig?.length || kvScheme.kvBits))
+    console.warn(
+      `[batch] --batch ${batch} with explicit --kv-quant: batched serving (v1) is ` +
+        `bf16-only, so kv-quant requests route to the serial path (no batching for them). ` +
+        `Omit --kv-quant to batch in bf16. (docs/design/parallel-slots.md)`,
+    );
+
+  // SSD cold tier (docs/design/ssd-kv-cold-tier.md): prefix KV survives RAM
+  // eviction and restarts. Compatibility key = configFingerprint (graph
+  // shape) + the EFFECTIVE kv scheme (flags pick bf16 vs config vs uniform
+  // on the same model — restored caches must match what serving produces) +
+  // tokenizer hash (ids must keep meaning the same text).
+  let ssdStore: SsdCacheStore | null = null;
+  if (serverOptions.ssdCacheDir) {
+    if ((serverOptions.promptCacheBytes ?? 2e9) <= 0)
+      throw new Error("--ssd-cache requires the RAM prompt cache (--prompt-cache 0 disables it)");
+    const schemeKey = kvScheme.kvBits
+      ? `kv${kvScheme.kvBits}`
+      : kvScheme.kvConfig?.length ? "config" : "bf16";
+    const tokJson = readFileSync(`${ctx.model.config.modelDir}/tokenizer.json`);
+    ssdStore = new SsdCacheStore({
+      dir: serverOptions.ssdCacheDir,
+      maxBytes: serverOptions.ssdCacheMaxBytes ?? 32 * 2 ** 30,
+      configFingerprint: `${configFingerprint(ctx.model.config)}-${schemeKey}`,
+      tokenizerHash: Bun.hash(tokJson).toString(16),
+      modelId: ctx.modelId,
+      verify: serverOptions.ssdCacheVerify,
+    });
+    const recovered = ssdStore.scan();
+    console.log(
+      `[ssd-cache] ${serverOptions.ssdCacheDir} — ${recovered} entr${recovered === 1 ? "y" : "ies"} recovered, ` +
+      `${(ssdStore.totalBytes / 2 ** 30).toFixed(2)} GiB of ${(ssdStore.maxBytes / 2 ** 30).toFixed(0)} GiB cap`,
+    );
+  }
+
+  const promptCache = new PromptCache(
+    serverOptions.promptCacheBytes ?? 2e9,
+    // RAM eviction spills to the cold tier (whole-prefix entries, no work
+    // on the token loop — the spill runs at eviction time only).
+    ssdStore ? (entry) => { ssdStore!.store(entry.tokens, entry.caches, entry.ns); } : null,
+  );
   // Responses-API store for previous_response_id resumption (Phase 11):
   // TTL + byte-capped LRU, port of optiq/response_store.py. Pairs with
   // the prompt cache: a resumed conversation re-renders the same prefix,
@@ -1001,7 +1069,31 @@ export function createServer(
     // Cache entries are adapter-specific: KV computed under one adapter
     // must never seed another's (or the base's) prefill.
     const cacheNs = options.adapters?.join("+") ?? "";
-    const entry = vision ? null : promptCache.take(promptIds, cacheNs);
+    let entry = vision ? null : promptCache.take(promptIds, cacheNs);
+    // Cold tier: restore from SSD when it has a STRICTLY longer usable
+    // prefix than RAM offered (zero-copy COW mmap; pages fault in lazily
+    // during the continuation prefill). The restored entry carries an
+    // unmap thunk that PromptCache runs after its caches are disposed.
+    if (!vision && ssdStore) {
+      const ssdHit = ssdStore.find(promptIds, cacheNs);
+      if (ssdHit && ssdHit.prefixLen > (entry?.tokens.length ?? 0)) {
+        const loaded = ssdStore.restore(ssdHit.entry, ctx.model);
+        if (loaded) {
+          const trimNeeded = loaded.tokens.length - ssdHit.prefixLen;
+          if (trimNeeded > 0 && !loaded.caches.every((c) => c.isTrimmable())) {
+            // Divergent tail on an untrimmable kind (ring post-wrap, SSM):
+            // this file can't seed the prompt — fall back to the RAM hit.
+            for (const c of loaded.caches) c.dispose();
+            loaded.mmap.unmap();
+          } else {
+            if (trimNeeded > 0) for (const c of loaded.caches) c.trim(trimNeeded);
+            if (entry) promptCache.put(entry.tokens, entry.caches, cacheNs, entry.retain); // return the RAM hit untouched
+            entry = { tokens: loaded.tokens.slice(0, ssdHit.prefixLen), caches: loaded.caches,
+              ns: cacheNs, retain: () => loaded.mmap.unmap() };
+          }
+        }
+      }
+    }
     const caches = entry?.caches ?? ctx.model.makeCache();
     try {
       const gen = generate(ctx.model, promptIds, {
@@ -1016,11 +1108,13 @@ export function createServer(
       if (vision) {
         for (const c of caches) c.dispose();
       } else {
-        promptCache.put(s.cacheTokens, caches, cacheNs);
+        promptCache.put(s.cacheTokens, caches, cacheNs, entry?.retain);
+        scheduleSsdSnapshot(s.cacheTokens, cacheNs);
       }
       return s;
     } catch (e) {
       for (const c of caches) c.dispose();
+      entry?.retain?.();
       throw e;
     } finally {
       vision?.embeddings.dispose();
@@ -1029,6 +1123,29 @@ export function createServer(
     }
   };
 
+  // Write-behind persistence (restart survival — the oMLX boundary-snapshot
+  // idea at whole-entry granularity, spill-on-evict alone can't survive a
+  // clean exit): after a request completes, snapshot its still-RAM-resident
+  // entry to SSD off the request path. Debounced + coalesced per namespace —
+  // an agent hammering one conversation persists the settled state once
+  // things go quiet, not every turn. Runs under the gateway lock (byte
+  // extraction must not race a generation mutating the entry); if the entry
+  // was evicted (already spilled) or extended (a newer schedule pending)
+  // meanwhile, findExact misses and nothing is written.
+  const ssdPending = new Map<string, ReturnType<typeof setTimeout>>();
+  const scheduleSsdSnapshot = (tokens: number[], ns: string): void => {
+    if (!ssdStore || tokens.length === 0) return;
+    const prev = ssdPending.get(ns);
+    if (prev) clearTimeout(prev);
+    ssdPending.set(ns, setTimeout(() => {
+      ssdPending.delete(ns);
+      void gateway.runExclusive(async () => {
+        const e = promptCache.findExact(tokens, ns);
+        if (e) ssdStore!.store(e.tokens, e.caches, ns);
+      }).catch(() => { /* cold tier is best-effort */ });
+    }, 1000));
+  };
+
   // The lane picker: routes each request to the serial path (runGeneration,
   // above) or the continuous-batching scheduler, keeping the two off the GPU
   // (and shared loraState) at the same time. See src/serve/generation-gateway.ts.
@@ -1075,28 +1192,8 @@ export function createServer(
     return jobStore;
   };
 
-  // KV-quant scheme, resolved once: kv_config.json by default (optiq
-  // serve's headline behavior), overridable to uniform bits or off.
-  // KV quant scheme. The serial default is the model's mixed-precision config
-  // (optiq parity). But `--batch N` is a bf16 continuous-batching MODE (= mlx-lm
-  // B=N parity), so when KV quant is left UNSET under `--batch N` it defaults to
-  // bf16 — the batch path engages out of the box. An EXPLICIT --kv-quant
-  // (config / bits) is still honored, but those requests then route to the
-  // serial path (batched quantized KV is the L2 follow-up); warned below.
-  const configScheme = ctx.kvConfig?.length ? { kvConfig: ctx.kvConfig } : {};
-  const kvScheme: Pick<GenerateOptions, "kvBits" | "kvConfig" | "quantizedKvStart"> =
-    serverOptions.kvQuant === "off" ? {}
-    : serverOptions.kvQuant === "config" ? configScheme
-    : typeof serverOptions.kvQuant === "number"
-      ? { kvBits: serverOptions.kvQuant, quantizedKvStart: 0 }
-    : batch > 1 ? {} // unset + batching → bf16 mode (Option B)
-    : configScheme; // unset + serial → optiq-parity mixed-precision default
-  if (batch > 1 && (kvScheme.kvConfig?.length || kvScheme.kvBits))
-    console.warn(
-      `[batch] --batch ${batch} with explicit --kv-quant: batched serving (v1) is ` +
-        `bf16-only, so kv-quant requests route to the serial path (no batching for them). ` +
-        `Omit --kv-quant to batch in bf16. (docs/design/parallel-slots.md)`,
-    );
+  // (kvScheme is resolved earlier, above the prompt-cache construction —
+  // the SSD tier's compatibility fingerprint folds the effective scheme in.)
 
   // XTC never removes EOS or the newline token: mlx_lm.server passes
   // [tokenizer.eos_token_id, tokenizer.encode("\n")] as xtc_special_tokens.
@@ -1442,6 +1539,17 @@ export function createServer(
             hits: promptCache.hits,
             misses: promptCache.misses,
           },
+          ...(ssdStore ? {
+            ssd_cache: {
+              dir: serverOptions.ssdCacheDir,
+              entries: ssdStore.entries,
+              bytes: ssdStore.totalBytes,
+              max_bytes: ssdStore.maxBytes,
+              restores: ssdStore.stats.restores,
+              spills: ssdStore.stats.spills,
+              restore_ms_last: Math.round(ssdStore.stats.restoreMsLast),
+            },
+          } : {}),
           response_store: {
             entries: responseStore.size,
             bytes: responseStore.totalBytes,
@@ -1867,6 +1975,8 @@ export function createServer(
           kvQuant: !!(options.kvConfig?.length || options.kvBits),
         };
         const batched = gateway.willBatch(shape);
+        if (process.env.MLX_BUN_LANE_DEBUG === "1")
+          console.error(`[lane] batched=${batched} shape=${JSON.stringify(shape)} t=${Date.now() % 100000}`);
 
         if (body.stream) {
           const stream = new ReadableStream<Uint8Array>({