mlx-bun 0.4.0 → 0.5.0

package/src/server.ts

@@ -1,5 +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)
@@ -26,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";
@@ -49,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";
@@ -67,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";
@@ -77,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,
@@ -163,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. */
@@ -322,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 /
-  // 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
@@ -397,10 +410,9 @@ export function createServer(
   // quietly (kv-quant: the swap would drop the scheme; speculative
   // provider composition with paged storage is not implemented).
   if (serverOptions.pagedKv) {
-    if (kvScheme.kvBits || kvScheme.kvConfig?.length || kvScheme.turboQuant)
+    if (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(
@@ -415,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
@@ -454,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)
@@ -475,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;
@@ -484,6 +503,12 @@ export function createServer(
           // as a no-op so callers' dispose ordering is unchanged.
           return { tokens: loaded.tokens, caches: loaded.caches, attachments: loaded.attachments, retain: () => {} };
         },
+        ...(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;
@@ -491,29 +516,13 @@ export function createServer(
         },
       }
     : 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, attachments: entry.attachments }),
-        }
-      : 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,
@@ -553,59 +562,8 @@ export function createServer(
   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.runWhenIdle(), 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.runWhenIdle(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,
     checkpoints: !!(serverOptions.generationCheckpointTokens && ssdStore),
@@ -617,29 +575,21 @@ export function createServer(
   const spillQueueGbRaw = Number(runtimeValue("MLX_BUN_SSD_SPILL_QUEUE_GB"));
   const spillQueueCapBytes =
     (Number.isFinite(spillQueueGbRaw) && spillQueueGbRaw >= 0 ? spillQueueGbRaw : 2) * 1024 ** 3;
-  // Both queues share the configured total: reserve half for checkpoints
-  // only when that service exists. Each retains the existing soft-cap rule.
-  const checkpointQueueBytes = ssdStore && serverOptions.generationCheckpointTokens
-    ? Math.floor(spillQueueCapBytes / 2) : 0;
   if (ssdStore && serverOptions.generationCheckpointTokens)
-    checkpointPersistence = new ContinuationPersistence(ssdStore, { maxBytes: checkpointQueueBytes, runStep: ssdFlushStep });
-  const spillQueue = ssdStore
-    ? new SpillQueue(
-        spillQueueCapBytes - checkpointQueueBytes,
-        cacheBytes,
-        (item) => ssdStore!.storeAsync(item.tokens, item.caches, item.ns, ssdFlushStep, item.attachments),
-        (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;
+    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,
@@ -668,28 +618,21 @@ export function createServer(
       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?.();
   }
@@ -699,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.
@@ -729,9 +698,10 @@ 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, contextLimit, serverOptions.defaultAdapter,
     preparation, serving.buildPrompt);
@@ -739,6 +709,15 @@ export function createServer(
     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");
@@ -777,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
@@ -793,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(
@@ -824,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.
@@ -934,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",
@@ -970,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: {
@@ -995,14 +966,7 @@ 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,
@@ -1146,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");
@@ -1166,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");
@@ -1195,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, {});
         }
@@ -1250,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, {});
         }

package/src/spec/assistant-rows.ts

@@ -86,6 +86,7 @@ class AssistantRows implements DraftPrefillGroup, DraftRowGroup {
 
 export function assistantGroups(drafter: GemmaAssistantDrafter, namespace: string): GroupedDraftProvider {
   return {
+    checkpointNamespace: () => namespace,
     openPrefill: ({target,checkpoints}) => new AssistantRows(target.assistantRows!,drafter,namespace,checkpoints),
     open: ({target,checkpoints}) => new AssistantRows(target.assistantRows!,drafter,namespace,checkpoints),
   };

package/src/spec/deepspec-rows.ts

@@ -12,7 +12,7 @@ export function bindDeepspecTarget(target: TargetView, drafter: Pick<DeepspecDra
 }
 
 export function deepspecGroups(drafter: DeepspecDrafter, namespace: string): GroupedDraftProvider {
-  return projectedDraftGroups(target => {
+  return projectedDraftGroups(namespace, target => {
     bindDeepspecTarget(target, drafter);
     return {
       namespace, schema: "deepspec-context-v1",

package/src/spec/dflash-source.ts

@@ -48,7 +48,7 @@ export class DflashProvider implements DraftProvider {
     this.id = id;
     this.weightsBytes = weightsBytes;
     this.gamma = drafter.cfg.gamma;
-    this.grouped = projectedDraftGroups(target => {
+    this.grouped = projectedDraftGroups(namespace, target => {
       if (!target.gemmaTaps) throw new Error("DSpark drafter requires a Gemma4 target");
       const projection = target.gemmaTaps.projection;
       const minConf = runtimeValue("MLX_BUN_DSPARK_MINCONF");

package/src/spec/glm52-mtp-source.ts

@@ -30,6 +30,7 @@ function lastToken(feed: readonly number[]): number {
 
 export class Glm52NativeMtpProvider implements DraftProvider {
   readonly grouped: GroupedDraftProvider = {
+    checkpointNamespace: () => "glm52-native-mtp-v1",
     open: options => this.#openRows(options.target, options.sampling, options.checkpoints),
     openPrefill: options => this.#openRows(options.target, null, options.checkpoints),
   };

package/src/spec/ngram-source.ts

@@ -38,11 +38,11 @@
 // speculation is structurally never worse than plain decode by more than the
 // JS scan (~µs against 30k-token histories).
 
-import { flagOn } from "../runtime-config";
+import { runtimeFlag } from "../runtime-config";
 import type { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
 import type { CheckpointAttachment } from "../backends/mlx/checkpoint-state";
-import type { DraftProvider, DraftSource, DraftRowCheckpoint, DraftRowGroup, DraftPrefillGroup, GroupedDraftProvider } from "./source";
+import type { DraftProvider, DraftSource, DraftRowCheckpoint, DraftRowGroup, DraftPrefillGroup, GroupedDraftProvider, DraftRowConstraints } from "./source";
 import { applyStateChanges } from "../engine/resources";
 
 /** The matching policy is independent of the executor's feed convention. */
@@ -78,7 +78,8 @@ class NgramRows implements DraftRowGroup, DraftPrefillGroup {
   readonly namespace: string;
   #histories: number[][] = [];
   #drafts: number[][] = [];
-  constructor(readonly max: number, readonly min: number, checkpoints: readonly (DraftRowCheckpoint | null)[]) {
+  constructor(readonly max: number, readonly min: number, checkpoints: readonly (DraftRowCheckpoint | null)[],
+    readonly constraints?: DraftRowConstraints) {
     this.namespace = namespace(max, min);
     this.append(checkpoints);
   }
@@ -95,7 +96,12 @@ class NgramRows implements DraftRowGroup, DraftPrefillGroup {
     for (let row = 0; row < this.rowCount; row++) this.#histories[row]!.push(...ids.slice(row * length, (row + 1) * length));
   }
   materialize(): void {}
-  draft(pending: readonly number[], depth: number): number[][] {
+  draft(pending: readonly number[], depth: number): number[][] | Promise<number[][]> {
+    if (this.constraints) {
+      for (const [row, history] of this.#histories.entries()) history.push(pending[row]!);
+      return Promise.all(this.#histories.map((_, row) => this.constraints!.propose(row, depth)))
+        .then(proposals => this.#drafts = proposals);
+    }
     this.#drafts = this.#histories.map((history, row) => {
       history.push(pending[row]!);
       return proposeNgram(history, this.max, this.min, depth);
@@ -126,6 +132,7 @@ export class NgramProvider implements DraftProvider {
   readonly max: number;
   readonly min: number;
   readonly grouped: GroupedDraftProvider = {
+    checkpointNamespace: () => namespace(this.max, this.min),
     supportsTargetAdapters: true,
     open: options => new NgramRows(this.max, this.min, options.checkpoints),
     openPrefill: options => new NgramRows(this.max, this.min, options.checkpoints),
@@ -150,7 +157,7 @@ class NgramSource implements DraftSource {
   readonly weightsBytes = 0;
   #hist: number[] = [];
   #lastDrafts: number[] = [];
-  readonly #tailSplit = flagOn("MLX_BUN_PREFILL_TAIL_SPLIT", true);
+  readonly #tailSplit = runtimeFlag("MLX_BUN_PREFILL_TAIL_SPLIT", true);
   readonly checkpoint: NonNullable<DraftSource["checkpoint"]>;
 
   constructor(
@@ -202,3 +209,18 @@ class NgramSource implements DraftSource {
 
   dispose(): void {}
 }
+
+/** Constraint candidates share committed-history ownership with prompt lookup.
+ * They have no serial draft implementation and never perform ngram lookup. */
+export function constraintDraftProvider(): { id: string; grouped: GroupedDraftProvider } {
+  return {
+    id: "grammar",
+    grouped: {
+      checkpointNamespace: () => namespace(3, 1),
+      supportsTargetAdapters: true,
+      openPrefill: options => new NgramRows(3, 1, options.checkpoints),
+      open: options => new NgramRows(3, 1, options.checkpoints,
+        options.constraints ?? { propose: async () => [] }),
+    },
+  };
+}

package/src/spec/projected-draft-rows.ts

@@ -65,9 +65,9 @@ class ProjectedDraftRows implements DraftPrefillGroup, DraftRowGroup {
   dispose(): void { this.context.dispose(); }
 }
 
-export function projectedDraftGroups(bind: (target: TargetView) => ProjectedDraftMethod): GroupedDraftProvider {
+export function projectedDraftGroups(namespace: string, bind: (target: TargetView) => ProjectedDraftMethod): GroupedDraftProvider {
   const open = ({ target, checkpoints }: {
     target: TargetView; checkpoints: readonly (DraftRowCheckpoint | null)[];
   }) => new ProjectedDraftRows(bind(target), checkpoints);
-  return { open, openPrefill: open };
+  return { checkpointNamespace: () => namespace, open, openPrefill: open };
 }

package/src/spec/qwen-mtp-rows.ts

@@ -133,6 +133,17 @@ export class QwenMtpRows {
     if (state.length) ops.evalAll(state);
   }
 
+  /** Align the companion to verified external continuations through the same
+   * true-hidden prefill graph. Different rows retain different prefix lengths. */
+  consume(tokens: MlxArray, context: MlxArray, lengths: readonly number[]): void {
+    this.#cache.specRoundBegin();
+    this.prefill(tokens, context);
+    this.#cache.specRoundRollback([...lengths]);
+    using positions = ops.fromInt32(lengths.map(length => length - 1), [lengths.length, 1, 1]);
+    const hidden = ops.takeAlongAxis(context, positions, 1);
+    this.#hidden!.dispose(); this.#hidden = hidden;
+  }
+
   /** Nonnegative depth, with every request's first pending token already known.
    * Build the complete device dependency chain before reading proposal IDs. */
   draft(pending: readonly number[], depth: number, steps: readonly number[]): number[][] {