mlx-bun 0.3.0 → 0.5.0

package/src/server.ts

@@ -1,4 +1,7 @@
+import { ContinuationPersistence } from "./backends/mlx/continuation-persistence";
 import { modelServingBinding } from "./backends/mlx/model-serving";
+import { handleAudioRoute } from "./serve/audio-routes";
+import { TranscriptionService, defaultWhisperRecord } from "./serve/transcription-service";
 import type { ServingContext } from "./serve/model-host";
 export { generationCheckpointKey } from "./serve/checkpoint-identity";
 // OpenAI-compatible HTTP server: /v1/chat/completions (+ SSE streaming)
@@ -25,7 +28,7 @@ import {
   type GenerateOptions,
 } from "./generate";
 import { createPreparationExecutor } from "./serve/preparation";
-import { cloneKvCaches, legacyCacheCodecs, SpillQueue } from "./kv-store";
+import { cloneKvCaches, legacyCacheCodecs } from "./kv-store";
 import type { Cache } from "./model/gemma4";
 import { resolveKvScheme } from "./kv-scheme";
 import { runtimeValue } from "./runtime-config";
@@ -48,10 +51,11 @@ import { handleLabRoute } from "./serve/lab-routes";
 import { handleModelAdminRoute } from "./serve/model-admin-routes";
 import { handleStaticRoute } from "./serve/static-routes";
 
-import { PromptCache, cacheBytes } from "./prompt-cache";
+import { PromptCache } from "./prompt-cache";
+import { CostSizeRetention } from "./storage/retention-policy";
+import { TieredPromptCache } from "./tiered-prompt-cache";
 import { SsdCacheStore } from "./ssd-cache";
 import {
-  SsdDurabilityCoordinator,
   type DurabilityFlushResult,
   type DurabilitySnapshotStats,
 } from "./ssd-durability";
@@ -66,7 +70,7 @@ import {
   type ResponsesRequest,
 } from "./responses";
 import { fit } from "./fit";
-import { setMemoryLimit } from "./mlx/ffi";
+import { activeMemory, maxRecommendedWorkingSetSize, setMemoryLimit } from "./mlx/ffi";
 import { makePiWsHandler, type PiWsData } from "./pi-web";
 import { ChatStage } from "./serve/chat-stage";
 import { TextCompletionStage } from "./serve/text-completion-stage";
@@ -76,6 +80,7 @@ import {
   chatCompletionJson, chatCompletionStream, textCompletionJson, textCompletionStream,
 } from "./serve/openai-wire";
 import { createRequestPrep } from "./serve/request-prep";
+import { GeneratedTokenHistory } from "./serve/generated-token-history";
 import {
   detectDraftKind,
   loadContext,
@@ -109,13 +114,15 @@ export interface ServerOptions {
    *  (`--kv-budget`, batching-perf-path P3). Joiners whose projected KV
    *  (prompt + max_tokens, window-capped) would exceed it QUEUE until rows
    *  evict; a request over the budget alone is rejected. Unset = no
-   *  aggregate cap (per-request admission via memoryBudget still applies). */
+   *  aggregate cap (an explicit memoryBudget also limits each request). */
   kvBudgetBytes?: number;
   /** KV quantization override. Unset/"off" is bf16. "config" applies the
    *  model's declared mixed-precision kv_config; supported per-layer schemes
    *  compose with continuous scheduling. A number forces uniform bits
    *  (group size 64, start 0) and uses the preserved serial executor. */
   kvQuant?: "off" | "config" | number;
+  /** Optional absolute row offset for KV conversion; omitted preserves start-zero server policy. */
+  quantizedKvStart?: number;
   /** TurboQuant scheme (docs/design/turboquant.md): a separate axis from
    *  kvQuant above, mutually exclusive with it (`--kv-quant turbo[:k<bits>v
    *  <bits>]` sets this instead of kvQuant). Solo-only in v1 — see
@@ -124,19 +131,13 @@ export interface ServerOptions {
   /** OPTIONAL paged KV cache (`--paged-kv`, docs/design/kv-cache.md):
    *  vLLM-style block-pool storage for full-attention layers,
    *  gather-to-contiguous before the stock SDPA. Default off (unset = the
-   *  plain KVCache path, byte-identical). v1 scope: serial batch=1,
-   *  Gemma4-family, bf16 — startup refuses `--batch N>1`, `--kv-quant`,
+   *  plain KVCache path, byte-identical). Gemma4-family, bf16, B1/B>1;
+   *  startup refuses `--kv-quant`,
    *  turbo, and `--draft-model` combinations; paged requests bypass the
    *  prompt cache and run uncompiled. Gated bit-exact vs the plain path. */
   pagedKv?: { blockSize?: number };
-  /** Memory budget for the serving process (admission control — Phase 5).
-   *  Requests whose prompt + max_tokens exceed the budget's max safe
-   *  context are rejected with 400 instead of crashing the GPU: the OOM
-   *  crash class is UNCATCHABLE (Phase 6 — mlx throws from a Metal
-   *  completion handler ⇒ std::terminate; optiq serve died exactly this
-   *  way loading the 26B). Also caps the mlx allocator
-   *  (mlx_set_memory_limit) as defense in depth. Default: machine RAM ×
-   *  WIRED_FRACTION, admission check only, allocator untouched. */
+  /** Explicit opt-in process budget. Enforces the fit estimate's context
+   *  limit and sets the MLX allocator limit. Unset keeps estimates advisory. */
   memoryBudgetBytes?: number;
   /** Who owns this server's lifetime: "serve" (persistent, mlx-bun
    *  serve) or "pi-session" (dies with the pi session that started
@@ -166,6 +167,16 @@ export interface ServerOptions {
   defaultTemperature?: number;
   defaultTopP?: number;
   defaultTopK?: number;
+  /** Whisper checkpoint directory for /v1/audio/* (`--whisper-model`).
+   *  Unset ⇒ the first downloaded `whisper` checkpoint is used on demand. */
+  whisperModelDir?: string;
+  /** Display id for the Whisper checkpoint. */
+  whisperModelId?: string;
+  /** Seconds idle before the Whisper weights are disposed
+   *  (`--whisper-idle-unload`, default 300; 0 keeps them resident). */
+  whisperIdleUnloadSec?: number;
+  /** Keep the Whisper weights resident (`--whisper-resident`). */
+  whisperResident?: boolean;
   /** Completion cap when the request omits max_tokens (`--max-tokens`).
    *  mlx_lm.server's flag; its default there is 512 — ours stays 65,536 so
    *  thinking traces never truncate. `--max-tokens 512` = mlx-lm behavior. */
@@ -287,16 +298,8 @@ export interface ContextAdmissionDecision {
   clamped: boolean;
 }
 
-/** Resolve a request's completion upper bound against the admitted context.
- *
- * `max_tokens` is a ceiling, not a promise that every token will be emitted:
- * when a prompt fits but a broad client-wide `max_tokens` would push the
- * reservation past the safe context, the bound is CLAMPED to the remaining
- * room instead of rejecting an otherwise valid request (the pre-v0.0.13
- * behavior 400'd a prompt with thousands of tokens of generation room over a
- * 17-token overshoot). This never weakens the OOM guard — generation stops at
- * the clamped ceiling, inside the admitted context. Only a prompt that leaves
- * no generation slot at all is rejected. */
+/** Resolve a completion cap against an explicitly enforced context limit.
+ * Serving keeps fit estimates advisory unless the operator chooses a budget. */
 export function admitRequestContext(
   promptTokens: number,
   requestedMaxTokens: number,
@@ -333,8 +336,7 @@ export function createServer(
   // The scheduler chooses its B=1 fast path or B=N step from active rows.
   // Both full-attention (CPM) and
   // sliding-window (Gemma) models batch — the scheduler assembles each layer's
-  // cache by attention type. Non-batchable requests (vision / adapters /
-  // user seed / unsupported explicit kv-quant) drain to the serial executor
+  // cache by attention type. Unsupported model/method combinations drain to the serial executor
   // (see GenerationGateway.place). No inference setting is rewritten.
   // DEFAULT 8 (flipped 2026-07-05, Josh's call, after GATE-B1-SPEED): a
   // lone request through the batch lane IS the serial engine (adopted
@@ -352,8 +354,7 @@ export function createServer(
       throw new Error("--generation-checkpoint expects a positive integer token interval");
     if (!serverOptions.ssdCacheDir)
       throw new Error("--generation-checkpoint requires --ssd-cache <dir>");
-    if (batch !== 1)
-      throw new Error("--generation-checkpoint requires --batch 1 (in-flight resume is serial-only)");
+
   }
   const defaultGeneratedTokens =
     serverOptions.defaultMaxTokens ?? ctx.glmMemoryPlan?.maxGenerationTokens;
@@ -376,32 +377,11 @@ export function createServer(
   const resolvedKvScheme = resolveKvScheme({
     override: serverOptions.kvQuant,
     turboQuant: serverOptions.turboQuant,
+    quantizedKvStart: serverOptions.quantizedKvStart,
     config: ctx.kvConfig,
   });
   const kvScheme = resolvedKvScheme.generationOptions;
-  // Phase 3.1: kvConfig whose layers are all full-attention BATCHES (the
-  // scheduler applies the mixed scheme per row); uniform kvBits and configs
-  // touching rotating layers still route those requests serial. The warning
-  // only fires for the still-serial compositions.
-  if (batch > 1 && kvScheme.kvBits)
-    console.warn(
-      `[batch] --batch ${batch} with uniform --kv-quant ${kvScheme.kvBits}: uniform ` +
-        `quantized KV is serial-only (it touches rotating layers) — those requests ` +
-        `won't batch. Use --kv-quant config for batched mixed-precision KV, or omit ` +
-        `--kv-quant to batch in bf16. (docs/design/unified-engine-frontier-plan.md)`,
-    );
-  // TurboQuant is solo-only in v1 (novel cache class, not batchable by
-  // construction — see GenerationGateway placement's cache-capability gate).
-  if (batch > 1 && kvScheme.turboQuant)
-    console.warn(
-      `[batch] --batch ${batch} with --kv-quant turbo: TurboQuant is serial-only in v1 ` +
-        `— those requests won't batch. Omit --kv-quant to batch in bf16. ` +
-        `(docs/design/turboquant.md)`,
-    );
-  // Quantized KV (any axis) excludes the spec lane: drafted requests fall to
-  // the normal serial path WITH the KV scheme applied rather than losing it
-  // silently (mirrors the affine kvBits/kvConfig exclusion; the spec loop is
-  // bf16-KV-only in v1).
+  // The model execution plan qualifies the method while retaining KV policy.
   if (ctx.draft && (kvScheme.turboQuant || kvScheme.kvBits || kvScheme.kvConfig?.length))
     console.warn(
       `[spec] --draft-model with quantized KV (--kv-quant ${serverOptions.turboQuant ? "turbo" : String(serverOptions.kvQuant)}): ` +
@@ -427,19 +407,12 @@ export function createServer(
   }
   // Paged KV v1 (docs/design/kv-cache.md): explicit refusals, not
   // silent downgrades — the incompatible combos would otherwise degrade
-  // quietly (batch: paged caches can't merge; kv-quant: the swap would
-  // drop the scheme, the exact composition bug the kvQuant gate above
-  // exists to prevent; draft: the spec lane assumes serial-class caches).
+  // quietly (kv-quant: the swap would drop the scheme; speculative
+  // provider composition with paged storage is not implemented).
   if (serverOptions.pagedKv) {
-    if (batch > 1)
+    if (kvScheme.kvConfig?.length || kvScheme.turboQuant)
       throw new Error(
-        `--paged-kv is serial-only in v1 — add --batch 1 (got --batch ${batch}). ` +
-          `Batched paging is the follow-up PR (docs/design/kv-cache.md).`,
-      );
-    if (kvScheme.kvBits || kvScheme.kvConfig?.length || kvScheme.turboQuant)
-      throw new Error(
-        `--paged-kv is bf16-only in v1 — omit --kv-quant (quantized paged ` +
-          `blocks are a documented follow-up, docs/design/kv-cache.md).`,
+        `--paged-kv supports bf16 and uniform affine KV4/KV8; per-layer and TurboQuant pages are not implemented.`,
       );
     if (ctx.draft)
       throw new Error(
@@ -454,11 +427,7 @@ export function createServer(
     const bs = serverOptions.pagedKv.blockSize;
     if (bs !== undefined && (!Number.isInteger(bs) || bs <= 0))
       throw new Error(`--paged-kv-block-size must be a positive integer (got ${bs})`);
-    if (serverOptions.ssdCacheDir)
-      console.warn(
-        "[paged-kv] --ssd-cache has no effect: paged requests bypass the prompt " +
-          "cache (v1 non-goal), so nothing reaches the SSD tier.",
-      );
+
   }
 
   // SSD cold tier (docs/design/kv-cache.md): prefix KV survives RAM
@@ -493,6 +462,8 @@ export function createServer(
       tokenizerHash: Bun.hash(tokJson).toString(16),
       modelId: ctx.modelId,
       verify: serverOptions.ssdCacheVerify,
+      storage: { layout: runtimeValue("MLX_BUN_SSD_LAYOUT") === "blocks" ? "blocks" : "whole",
+        segmented: runtimeValue("MLX_BUN_SSD_SEGMENTED") !== "0" },
     });
     const recovered = ssdStore.scan();
     const capacity = Number.isFinite(ssdStore.maxBytes)
@@ -514,6 +485,15 @@ export function createServer(
           const h = ssdStore!.find(prompt, ns);
           return h ? { prefixLen: h.prefixLen, handle: h.entry } : null;
         },
+        findExact: (tokens: number[], ns: string) => {
+          const hit = ssdStore!.findExact(tokens, ns);
+          return hit ? { prefixLen: hit.prefixLen, handle: hit.entry } : null;
+        },
+        restoreObjectAsync: async (handle: unknown) => {
+          const loaded = await ssdStore!.restoreAsync(handle as import("./ssd-cache").SsdIndexEntry,
+            { makeCache: () => [] });
+          return loaded ? { ...loaded, retain: () => {} } : null;
+        },
         restore: (handle: unknown) => {
           const loaded = serving.restore(ssdStore!, handle as import("./ssd-cache").SsdIndexEntry);
           if (!loaded) return null;
@@ -521,219 +501,138 @@ export function createServer(
           // own their bytes and no mapping outlives loadKvCache — nothing
           // to pin, nothing to unmap. retain stays in the entry contract
           // as a no-op so callers' dispose ordering is unchanged.
-          return { tokens: loaded.tokens, caches: loaded.caches, retain: () => {} };
+          return { tokens: loaded.tokens, caches: loaded.caches, attachments: loaded.attachments, retain: () => {} };
         },
-        store: (tokens: number[], caches: import("./model/gemma4").Cache[], ns: string) => {
+        ...(serving.restoreAsync && runtimeValue("MLX_BUN_SSD_PREFETCH") !== "0" ? {
+          restoreAsync: async (handle: unknown) => {
+            const loaded = await serving.restoreAsync!(ssdStore!, handle as import("./ssd-cache").SsdIndexEntry);
+            return loaded ? { ...loaded, retain: () => {} } : null;
+          },
+        } : {}),
+        store: (tokens: number[], caches: import("./model/gemma4").Cache[], ns: string,
+          attachments?: import("./backends/mlx/checkpoint-state").CheckpointAttachment[]) => {
           if (ssdStore!.hasDurablePrefix(tokens, ns)) return true;
-          return ssdStore!.store(tokens, caches, ns);
+          return ssdStore!.store(tokens, caches, ns, attachments);
         },
       }
     : null;
-  const promptCache = new PromptCache(
-    promptCacheCap,
-    // RAM eviction AND idle demotion spill to the cold tier NON-BLOCKING
-    // (2026-07-06, same contract as the write-behind snapshot below): the
-    // cache hands us OWNED zero-copy clones + copied tokens (made under
-    // the generation lock — microseconds; entries are immutable so the
-    // clones stay consistent), and the flush goes through the BOUNDED
-    // SpillQueue -> storeAsync (generation-locked per tensor, see below) ->
-    // clone disposal on every settle/drop path (that dispose is what
-    // actually frees the demoted GPU memory; the queue's byte cap keeps
-    // starved-gate retention bounded — 2026-07-07 review fix). spillQueue
-    // is declared with the snapshot scheduler below; safe to close over
-    // here because spills only fire at put()/demoteIdle time, long after
-    // init.
-    ssdStore
-      ? {
-          spillOwned: (entry) =>
-            spillQueue!.enqueue({ tokens: entry.tokens, caches: entry.caches, ns: entry.ns }),
-        }
-      : null,
-    coldTier,
-    cloneState,
-  );
+  const promptCache = ssdStore && coldTier
+    ? new TieredPromptCache(promptCacheCap, ssdStore, coldTier, cloneState,
+        runtimeValue("MLX_BUN_SSD_WRITEBEHIND") !== "0")
+    : new PromptCache(promptCacheCap, null, null, cloneState);
+  if (runtimeValue("MLX_BUN_CACHE_RETENTION") === "cost-size") promptCache.retention = new CostSizeRetention();
+  const spillQueue = promptCache instanceof TieredPromptCache ? promptCache.spillQueue : null;
+  const durability = promptCache instanceof TieredPromptCache ? promptCache.durability : 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,
   // so its KV prefill is already cached.
   const responseStore = new ResponseStore();
 
-  // Admission ceiling, resolved once (Phase 5 memoryBudget enforcement).
-  // fit() solves max safe context from weights + KV growth + prefill
-  // transient. The active kv-quant scheme (uniform kvBits / per-layer
-  // kvConfig) is billed at its true bytes/element so a quantized cache
-  // advertises and admits the larger window it actually enables; only
-  // TurboQuant still bills bf16 (conservative — no projector for its
-  // layout yet, and it is solo-only in v1).
+  // Memory estimates remain available in diagnostics. Only an explicit
+  // operator budget or context cap constrains request planning.
   const memoryAdmission = ctx.glmMemoryPlan ?? fit(
     ctx.model.config, ctx.model.weightsBytes, 1,
     undefined, undefined, 0, serverOptions.memoryBudgetBytes,
     resolvedKvScheme.fitOptions,
   );
-  // Isolated benchmark profile cap. It may narrow physical admission, never
-  // enlarge it. The normal request planner still clamps prompt + output.
   const profileContext = runtimeValue("MLX_BUN_RD_CONTEXT_LIMIT");
-  const contextLimit = profileContext === undefined ? null : Number(profileContext);
-  if (contextLimit !== null && (!Number.isSafeInteger(contextLimit) || contextLimit < 1))
+  const profileLimit = profileContext === undefined ? null : Number(profileContext);
+  if (profileLimit !== null && (!Number.isSafeInteger(profileLimit) || profileLimit < 1))
     throw new Error("MLX_BUN_RD_CONTEXT_LIMIT must be a positive integer");
-  const admission = contextLimit === null ? memoryAdmission : {
-    ...memoryAdmission, maxSafeContext: Math.min(contextLimit, memoryAdmission.maxSafeContext),
-  };
-  // A zero ceiling is a warning, never a startup refusal: killing the
-  // server here can only parrot the per-request admission message (which
-  // still fires, with this same ceiling) or be a false positive from the
-  // fit model itself — it can never save anything the request gate
-  // doesn't. Serve until physically incapable.
-  if (admission.maxSafeContext < 1)
-    console.warn(
-      `[admission] memory budget ${(admission.usableBytes / 1e9).toFixed(2)} GB leaves no ` +
-      `safe context for ${ctx.modelId} (weights ${(ctx.model.weightsBytes / 1e9).toFixed(2)} GB) ` +
-      `— serving anyway; generation requests will be refused until the budget is raised`,
-    );
+  const budgetLimit = serverOptions.memoryBudgetBytes !== undefined
+    ? memoryAdmission.maxSafeContext : ctx.glmMemoryPlan?.contextTokens ?? null;
+  const contextLimit = profileLimit === null ? budgetLimit
+    : Math.min(profileLimit, budgetLimit ?? Infinity);
+  const admission = memoryAdmission;
   const allocatorLimit =
     admission.allocatorLimitBytes ?? serverOptions.memoryBudgetBytes;
   if (allocatorLimit) setMemoryLimit(allocatorLimit);
 
-  const memoryBudget = {
-    usableBytes: admission.usableBytes,
-    kvOptions: resolvedKvScheme.fitOptions,
-    promptCache,
-  };
-  const runGeneration = serving.createSerial({
+  let checkpointPersistence: ContinuationPersistence | undefined;
+  const executionServices: import("./backends/mlx/serial-executor").MlxSerialServices = {
+    get checkpointPersistence() { return checkpointPersistence; },
     promptCache, checkpoints: ssdStore,
-    memoryBudget,
     checkpointEveryTokens: serverOptions.generationCheckpointTokens,
     identity: { artifact: ctx.profile.artifact, implementation: ctx.profile.profile.execution,
       stateAbi: "legacy-cache-array-v1", codecs: stateCodecs.id },
     adapterNamespace: (adapters) => ctx.adapters.cacheNamespace(adapters),
     cloneState,
-  });
+  };
+  serving.gateway.configureContinuation?.(executionServices);
+  const runGeneration = serving.createSerial(executionServices);
 
-  // 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.
-  // Keyed by namespace and exact tokens: a sub-second final [prompt+gen] put
-  // must not cancel the pending boundary snapshot, and unrelated prompts of
-  // equal length must not cancel each other. Exact reschedules still
-  // coalesce; stale keys self-clean when findExact misses.
-  //
-  // NON-BLOCKING (2026-07-06, the write-behind persistence contract): the
-  // gateway lock is held only for a zero-copy SNAPSHOT (findExact +
-  // cloneKvCaches — microseconds; entries are immutable so the clones are
-  // consistent forever). The flush itself runs OFF the lock via
-  // storeAsync, and writes chain serially so two multi-hundred-MB flushes
-  // never overlap.
-  //
-  // IDLE-GATED (2026-07-07, the decode@ctx fix): "off the lock" was not
-  // enough — every per-tensor flush step is a blocking GPU sync on the
-  // SAME stream decode uses (ops.contiguous enqueues a kernel, rawBytesView
-  // evals) plus a synchronous multi-MB writeSync, and the setImmediate
-  // pacing interleaved those slices exactly between decode tokens. A ~16k
-  // entry's flush overlapping the bench's ctx repeats depressed decode@ctx
-  // ~9% on e4b (mlx-lm runs no equivalent background work). Now every step
-  // — including the first — runs through gateway.runExclusive(), so the
-  // flush only progresses while NOTHING is generating and pauses when a
-  // request arrives mid-flush. Tradeoffs, accepted: durability waits for a
-  // quiet moment (single-user serving quiesces constantly; sustained
-  // hammering defers the flush AND the spill clones' GPU-memory release —
-  // bounded by the chain). MLX_BUN_SSD_WRITEBEHIND=0
-  // disables write-behind snapshots entirely — the paired-A/B lever + kill
-  // switch (restart survival then degrades to spill-on-evict only).
-  const writeBehindOn = runtimeValue("MLX_BUN_SSD_WRITEBEHIND") !== "0";
-  const ssdFlushStep = <T>(step: () => T): Promise<T> =>
-    gateway.runExclusive(async () => step());
-  // Bounded write-behind queue (2026-07-07 review fix — see SpillQueue in
-  // kv-store.ts): pending clones pin their entries' GPU buffers while the
-  // idle gate starves under sustained traffic, so QUEUED bytes are capped —
-  // over cap the oldest queued spill drops (clone disposed immediately, a
-  // future cache miss, never a wrong result) instead of accumulating past
-  // the prompt-cache cap. Default 2 GB (a quarter of the 8 GB RAM-cache
-  // default); MLX_BUN_SSD_SPILL_QUEUE_GB overrides. The durability
-  // coordinator retains a dirty record until the queue confirms the atomic
-  // store. Explicit flush and graceful shutdown can therefore retry a drop.
-  // `|| 2` would coerce an explicit "0" back to 2 GB — parse so 0 works
-  // (cap 0 = keep only the newest + in-flight clone pinned; the soft cap
-  // never drops the item just enqueued).
-  // The gateway exists before the spill queue because the SSD writer uses
-  // its idle boundary to avoid competing with generation.
+  // Cache policy owns RAM residency and SSD persistence independently of
+  // request scheduling. Published state is immutable.
   const gateway = new GenerationGateway(serving.gateway, batch, runGeneration, {
     kvBudgetBytes: serverOptions.kvBudgetBytes,
-    memoryBudget,
     checkpoints: !!(serverOptions.generationCheckpointTokens && ssdStore),
     stateCodecs,
     kvScheme: resolvedKvScheme,
     promptCache,
-    beforeSerial: async () => { await flushDurability(); },
+    adapterNamespace: (adapters) => ctx.adapters.cacheNamespace(adapters),
   });
   const spillQueueGbRaw = Number(runtimeValue("MLX_BUN_SSD_SPILL_QUEUE_GB"));
   const spillQueueCapBytes =
     (Number.isFinite(spillQueueGbRaw) && spillQueueGbRaw >= 0 ? spillQueueGbRaw : 2) * 1024 ** 3;
-  const spillQueue = ssdStore
-    ? new SpillQueue(
-        spillQueueCapBytes,
-        cacheBytes,
-        (item) => ssdStore!.storeAsync(item.tokens, item.caches, item.ns, ssdFlushStep),
-        (caches) => { for (const c of caches) c.dispose(); },
-      )
-    : null;
-  const durability = ssdStore && spillQueue && writeBehindOn
-    ? new SsdDurabilityCoordinator(
-        gateway,
-        promptCache,
-        spillQueue,
-        cloneState,
-        (tokens, ns) => ssdStore.hasDurablePrefix(tokens, ns),
-      )
-    : null;
-  const durabilityStats = (): DurabilitySnapshotStats =>
-    durability?.stats ?? {
+  if (ssdStore && serverOptions.generationCheckpointTokens)
+    checkpointPersistence = new ContinuationPersistence(ssdStore, {
+      maxBytes: spillQueueCapBytes,
+    });
+  // Keep allocator headroom for the next forward. Only optional cache
+  // snapshots are reclaimed; context, batch size and sampling stay intact.
+  // Do not count MLX's reusable pool as live state or await an SSD write.
+  const cacheWorkingSet = Math.min(maxRecommendedWorkingSetSize(), allocatorLimit ?? Infinity);
+  const cacheOverBudget = () => Math.max(activeMemory(),
+    ctx.model.weightsBytes + Math.max(promptCache.totalBytes, spillQueue?.pendingBytes ?? 0)) > cacheWorkingSet * 0.85;
+  promptCache.pressure = {
+    // File-backed weights may not yet appear in MLX's active allocations
+    // before the first forward. Account for their known footprint as well.
+    overBudget: cacheOverBudget,
+  };
+  const durabilityStats = (): DurabilitySnapshotStats => {
+    const stats = durability?.stats ?? {
       pendingSnapshots: 0,
       pendingSpills: spillQueue?.pendingCount ?? 0,
       pendingSpillBytes: spillQueue?.pendingBytes ?? 0,
       droppedSpills: spillQueue?.droppedCount ?? 0,
       failedSpills: spillQueue?.failedCount ?? 0,
     };
+    const checkpoints = checkpointPersistence?.stats;
+    return { ...stats, pendingSpills: stats.pendingSpills + (checkpoints?.pendingCount ?? 0),
+      pendingSpillBytes: stats.pendingSpillBytes + (checkpoints?.pendingBytes ?? 0),
+      droppedSpills: stats.droppedSpills + (checkpoints?.dropped ?? 0),
+      failedSpills: stats.failedSpills + (checkpoints?.failed ?? 0) };
+  };
   const flushDurability = async (): Promise<DurabilityFlushResult> => {
-    if (durability) return durability.flush();
     const started = performance.now();
-    await spillQueue?.drain();
+    const checkpoints = await checkpointPersistence?.flush();
+    const result = durability ? await durability.flush() : null;
+    if (!result) await spillQueue?.drain();
     const stats = durabilityStats();
     return {
       ...stats,
-      durable: stats.pendingSpills === 0,
-      flushedSnapshots: 0,
-      missingSnapshots: 0,
+      durable: (result?.durable ?? stats.pendingSpills === 0) && (checkpoints?.durable ?? true) && stats.pendingSpills === 0,
+      flushedSnapshots: result?.flushedSnapshots ?? 0,
+      missingSnapshots: result?.missingSnapshots ?? 0,
       elapsedMs: performance.now() - started,
     };
   };
-  // Every put() — serial lane AND batch scheduler — schedules the snapshot
-  // (Layer 0: batch-lane entries survive restarts too, not just evictions).
-  promptCache.onPut = durability ? (tokens, ns) => durability.schedule(tokens, ns) : null;
+  // Remember exact generated IDs independently of cache persistence.
+  const tokenHistory = new GeneratedTokenHistory(ctx.tokenizer);
+  if (ssdStore) for (const tokens of ssdStore.tokenPrefixes()) tokenHistory.remember(tokens);
+  promptCache.onPut = (tokens) => {
+    tokenHistory.remember(tokens);
+  };
 
-  // Idle demotion (Layer 0): entries unused for --ssd-demote-idle seconds
-  // spill to SSD and free their GPU memory — the RAM tier drains between
-  // bursts while every prefix stays reachable (take() restores via
-  // zero-copy mmap). The exclusive section below covers only the sweep's
-  // zero-copy clone + dispose (spillOwned above); the SSD write itself
-  // runs off-lock on ssdWriteChain. Swept only when the engine is TRULY idle: the
-  // runExclusive below registers as a serial waiter, which would DRAIN a
-  // running batch, so the activity check guards it (a momentary exclusive
-  // grab of an idle engine is free). 0 disables.
+  // Cache age, not scheduler activity, decides idle demotion. The cache
+  // queues any missing SSD copy and retains RAM until that copy commits.
   const demoteIdleMs = (serverOptions.ssdDemoteIdleSec ?? (ssdStore ? 300 : 0)) * 1000;
   let demoteTimer: ReturnType<typeof setInterval> | null = null;
   if (ssdStore && demoteIdleMs > 0) {
     demoteTimer = setInterval(() => {
-      if (gateway.busy) return;
-      void gateway.runExclusive(async () => {
-        const n = promptCache.demoteIdle(demoteIdleMs);
-        if (n > 0) console.log(`[ssd-cache] demoted ${n} idle entr${n === 1 ? "y" : "ies"} to disk`);
-      }).catch(() => {});
+      const n = promptCache.demoteIdle(demoteIdleMs);
+      if (n > 0) console.log(`[ssd-cache] demoted ${n} idle entr${n === 1 ? "y" : "ies"} to disk`);
     }, Math.max(30_000, demoteIdleMs / 4));
     demoteTimer.unref?.();
   }
@@ -743,7 +642,33 @@ export function createServer(
 
   // /library response cache (30 s) — registry + config reads only.
   const startedAt = Date.now();
-  const discoveryRoutes = createDiscoveryRoutes(ctx, gateway, startedAt);
+  // Speech-to-text companion (docs/reference/server-api.md "Audio
+  // transcription"): a Whisper checkpoint served beside the chat model,
+  // loaded on the first /v1/audio request and paged out after
+  // --whisper-idle-unload seconds. Explicit --whisper-model wins; otherwise
+  // the first downloaded `whisper` checkpoint is resolved lazily.
+  let transcription: TranscriptionService | null = null;
+  let transcriptionResolved = false;
+  const transcriptionService = async (): Promise<TranscriptionService | null> => {
+    if (transcription || transcriptionResolved) return transcription;
+    transcriptionResolved = true;
+    const rec = serverOptions.whisperModelDir
+      ? { path: serverOptions.whisperModelDir, repoId: serverOptions.whisperModelId ?? serverOptions.whisperModelDir }
+      : await defaultWhisperRecord();
+    if (!rec) return null;
+    transcription = new TranscriptionService({
+      modelDir: rec.path, modelId: rec.repoId,
+      idleUnloadSec: serverOptions.whisperIdleUnloadSec,
+      resident: serverOptions.whisperResident,
+      exclusive: (fn, signal) => gateway.runExclusive(fn, undefined, signal),
+    });
+    return transcription;
+  };
+  const transcriptionInfo = async () => {
+    const svc = await transcriptionService();
+    return svc ? { id: svc.modelId, resident: svc.resident } : null;
+  };
+  const discoveryRoutes = createDiscoveryRoutes(ctx, gateway, startedAt, transcriptionInfo);
 
   // Captured so the WebSocket handler can resolve the bound (possibly
   // ephemeral) port lazily for the loopback pi provider.
@@ -773,16 +698,26 @@ export function createServer(
   //   → inferenceStage.admit → inferenceStage.run → InferenceResult → wire
   // (/v1/completions substitutes textStage; Anthropic and Responses reuse
   // chatStage with their own wire formats.)
-  const prep = createRequestPrep({ ctx, serverOptions, kvScheme, defaultGeneratedTokens });
+  const prep = createRequestPrep({ ctx, serverOptions, kvScheme, defaultGeneratedTokens, tokenHistory });
   const { templateOptionsFor } = prep;
-  const preparation = createPreparationExecutor((work, signal) => gateway.runExclusive(work, undefined, signal), batch);
+  const preparation = createPreparationExecutor((work, signal) => gateway.runPreparation(work, signal),
+    batch, gateway.mediaBatchingEnabled ? batch : 1);
   const chatStage = new ChatStage(
-    ctx, prep, promptCache, admission.maxSafeContext, serverOptions.defaultAdapter,
+    ctx, prep, promptCache, contextLimit, serverOptions.defaultAdapter,
     preparation, serving.buildPrompt);
   const textStage = new TextCompletionStage(
-    ctx, prep, admission.maxSafeContext, defaultGeneratedTokens, serverOptions.defaultAdapter);
+    ctx, prep, contextLimit, defaultGeneratedTokens, serverOptions.defaultAdapter);
   const inferenceStage = new InferenceStage(completionExecutor);
   const openAiMeta = (id: string) => ({ id, created: Math.floor(Date.now() / 1000), model: ctx.modelId });
+  /** Carry application session affinity into the shared cache request fields. */
+  const applyCacheSession = (body: ChatRequestParams, request: Request, original: unknown = body) => {
+    const fields = original as { session_id?: unknown; prompt_cache_key?: unknown };
+    const session = typeof fields.session_id === "string" ? fields.session_id :
+      typeof fields.prompt_cache_key === "string" ? fields.prompt_cache_key :
+      request.headers.get("x-session-affinity") ?? request.headers.get("session_id") ?? undefined;
+    if (runtimeValue("MLX_BUN_SESSION_CACHE") !== "0") body.session_id = session;
+    else { delete body.session_id; delete body.prompt_cache_key; }
+  };
   /** Parse the JSON body under the request's trace; a bad body is a 400. */
   const parseBody = async <T,>(request: Request, trace: PromptResponseTrace | undefined): Promise<T | Response> => {
     const closeBodyParse = trace?.begin("request.body_parse");
@@ -821,6 +756,7 @@ export function createServer(
       // capability flags — true if a tower is loaded or loadable (lazy)
       vision: !!(ctx.vision || ctx.loadVision),
       audio: !!(ctx.audio || ctx.loadAudio),
+      transcription: async () => (await transcriptionService()) !== null,
       thinking: ctx.template.supportsThinking,
       // Model-author sampling defaults (generation_config.json, server-CLI
       // overrides applied) for the sampling popover's per-model recommended
@@ -837,6 +773,12 @@ export function createServer(
     async fetch(request, server) {
       const url = new URL(request.url);
 
+      if (url.pathname === "/admin/cache/session/close" && request.method === "POST") {
+        const body = await parseBody<{ session_id?: string }>(request, undefined);
+        if (body instanceof Response) return body;
+        if (typeof body.session_id === "string") promptCache.closeSession(body.session_id);
+        return Response.json({ closed: typeof body.session_id === "string" });
+      }
       if (url.pathname === "/admin/cache/flush" && request.method === "POST") {
         const result = await flushDurability();
         return Response.json(
@@ -868,6 +810,9 @@ export function createServer(
       const modelAdminResponse = await handleModelAdminRoute(url, request, ctx, gateway);
       if (modelAdminResponse) return modelAdminResponse;
 
+      const audioResponse = await handleAudioRoute(url, request, { service: transcriptionService });
+      if (audioResponse) return audioResponse;
+
       if (url.pathname === "/fit" && request.method === "GET") {
         // Fit assessment for the status page: this-machine report at the
         // admission ceiling + the Apple SKU matrix at a fixed 32k.
@@ -978,30 +923,6 @@ export function createServer(
       }
 
       if (url.pathname === "/stats" && request.method === "GET") {
-        // Active KV scheme across ALL layers. Since Phase 9 rotating
-        // (sliding-window) caches quantize too, so every layer the
-        // scheme names counts — the old display filtered to
-        // full_attention and silently undercounted (e.g. 26B showed
-        // 5/30 quantized when its kv_config.json covers all 30).
-        const layerTypes = ctx.model.config.text.layerTypes;
-        const kvLayers: Record<string, number> = {};
-        let kvMode = "bf16";
-        if (kvScheme.turboQuant) {
-          // v1: full-attention layers only (sliding-window stays bf16 —
-          // docs/design/turboquant.md non-goal).
-          const fullAttn = layerTypes.filter((l) => l !== "sliding_attention").length;
-          kvMode = `turbo k${kvScheme.turboQuant.kBits}v${kvScheme.turboQuant.vBits}`;
-          kvLayers[`turbo-k${kvScheme.turboQuant.kBits}v${kvScheme.turboQuant.vBits}`] = fullAttn;
-        } else if (kvScheme.kvBits) {
-          kvMode = `uniform-kv${kvScheme.kvBits}`;
-          kvLayers[`kv${kvScheme.kvBits}`] = layerTypes.length;
-        } else if (kvScheme.kvConfig) {
-          kvMode = "mixed (kv_config.json)";
-          for (const e of kvScheme.kvConfig)
-            kvLayers[`kv${e.bits}`] = (kvLayers[`kv${e.bits}`] ?? 0) + 1;
-        }
-        const bf16Layers = layerTypes.length - Object.values(kvLayers).reduce((a, b) => a + b, 0);
-        const slidingLayers = layerTypes.filter((l) => l === "sliding_attention").length;
         return Response.json({
           server: {
             owner: serverOptions.owner ?? "embedded",
@@ -1014,6 +935,12 @@ export function createServer(
             max_bytes: promptCache.maxBytes,
             hits: promptCache.hits,
             misses: promptCache.misses,
+            session_hits: promptCache.sessionHits,
+            session_misses: promptCache.sessionMisses,
+            prefix_scans: promptCache.prefixScans,
+            object_hits: promptCache.objectHits,
+            object_misses: promptCache.objectMisses,
+            object_restores: promptCache.objectRestores,
           },
           ...(ssdStore ? {
             ssd_cache: {
@@ -1039,16 +966,10 @@ export function createServer(
             max_bytes: responseStore.maxBytes,
             ttl_ms: responseStore.ttlMs,
           },
-          kv_quant: {
-            mode: kvMode,
-            layers: { ...kvLayers, ...(bf16Layers > 0 ? { bf16: bf16Layers } : {}) },
-            attention: {
-              global: layerTypes.length - slidingLayers,
-              sliding_window: slidingLayers,
-            },
-          },
+          kv_quant: resolvedKvScheme.describe(ctx.model.config),
           admission: {
             max_safe_context: admission.maxSafeContext,
+            enforced_context_tokens: contextLimit,
             memory_budget_bytes:
               ctx.glmMemoryPlan?.processLimitBytes ??
               serverOptions.memoryBudgetBytes ?? null,
@@ -1074,7 +995,6 @@ export function createServer(
       // A parent-managed GPU job holds this response open. The connection owns
       // the native lease, so parent death cannot strand the worker lock.
       if (serverOptions.unixSocket && url.pathname === "/admin/lease" && request.method === "POST") {
-        await flushDurability();
         const lease = await gateway.acquireExecutionLease(request.signal);
         let released = false;
         const release = () => {
@@ -1190,6 +1110,7 @@ export function createServer(
         });
         const body = await parseBody<ChatRequestParams>(request, trace);
         if (body instanceof Response) return body;
+        applyCacheSession(body as ChatRequestParams, request);
         const meta = openAiMeta(id);
         const a = await admit(
           inferenceStage, () => chatStage.run(new ChatRequest(body), id, request.signal), trace, "chat request");
@@ -1210,6 +1131,7 @@ export function createServer(
         });
         const body = await parseBody<TextCompletionParams>(request, trace);
         if (body instanceof Response) return body;
+        applyCacheSession(body as ChatRequestParams, request);
         const meta = openAiMeta(id);
         const a = await admit(
           inferenceStage, () => textStage.run(new TextCompletionRequest(body), id), trace, "text completion");
@@ -1239,6 +1161,7 @@ export function createServer(
         let chatBody: ChatRequestParams;
         try {
           chatBody = anthropicToChatBody(anthropicBody) as unknown as ChatRequestParams;
+          applyCacheSession(chatBody, request, anthropicBody);
         } catch (e) {
           return anthropicError(400, (e as Error).message, {});
         }
@@ -1294,6 +1217,7 @@ export function createServer(
         let chatBody: ChatRequestParams;
         try {
           chatBody = responsesToChatBody(responsesBody) as unknown as ChatRequestParams;
+          applyCacheSession(chatBody, request, responsesBody);
         } catch (e) {
           return responsesError(400, (e as Error).message, {});
         }