mlx-bun 0.0.13 → 0.3.0

package/src/backends/mlx/model-prompt.ts

@@ -0,0 +1,164 @@
+import { ensureWav } from "../../audio/transcode";
+import { DiffusionGemmaModel } from "../../model/diffusion-gemma";
+import { Gemma4Model } from "../../model/gemma4";
+import { Qwen35Model } from "../../model/qwen3_5";
+import { spliceImageTokens } from "../../vision/diffusion-vision";
+import { buildMultimodalPrompt, buildVisionPrompt, extractAudio, extractImages, extractVideos,
+  type VisionEncoder, type VisionTokenIds } from "../../vision/prompt";
+import { buildQwen3VLVisionPrompt } from "../../vision/qwen3vl-prompt";
+import type { Qwen3VLVisionTower } from "../../vision/qwen3vl-tower";
+import { normalizeMessages, type ChatRequestParams } from "../../serve/chat-request";
+import { getAudioTower, getVisionTower, type ServerContext } from "../../serve/model-host";
+import { RequestError } from "../../serve/pipeline";
+import type { RequestOwnership } from "../../serve/request-plan";
+import type { RequestPrep } from "../../serve/request-prep";
+import type { BuiltPrompt } from "../../serve/model-binding";
+
+export async function buildModelPrompt(
+    ctx: ServerContext,
+    prep: RequestPrep,
+    body: ChatRequestParams,
+    tools: ChatRequestParams["tools"] | null,
+    ownership: RequestOwnership,
+  ): Promise<BuiltPrompt> {
+    const toolList = tools ?? null;
+    const partOf = (types: string[]) => body.messages.some(
+      (m) => Array.isArray(m.content) &&
+        m.content.some((p: any) => types.includes(p.type)),
+    );
+    const hasImages = partOf(["image_url", "image"]);
+    // Same shapes extractAudio accepts: OpenAI-canonical input_audio plus
+    // optiq's audio / audio_url aliases (docs/design/generic-model-support.md §3.2).
+    const hasAudio = partOf(["input_audio", "audio", "audio_url"]);
+    // Video content parts (video_url / video with base64 data) —
+    // Qwen3.5-family only (decoded to sampled frames via the
+    // AVFoundation sidecar, vision/video-frames.ts).
+    const hasVideos = partOf(["video_url", "video"]);
+    const noMedia: Pick<BuiltPrompt, "startInThinking" | "probeStableLen" | "diffusionPixels"> =
+      { startInThinking: false, probeStableLen: false, diffusionPixels: null };
+
+    // Video is Qwen3.5-family only and never composes with audio —
+    // one early guard so no downstream branch can silently drop a
+    // video part (the gemma/diffusion builders don't know the type).
+    if (hasVideos && (hasAudio || !(ctx.model instanceof Qwen35Model))) {
+      throw new RequestError(400, hasAudio
+        ? "video and audio content parts cannot be combined"
+        : `model ${ctx.modelId} does not accept video input — video ` +
+          `content parts need a Qwen3.5-family model (e.g. Qwen3.8-27B)`);
+    }
+    if (hasAudio) {
+      // Audio (and MIXED image+audio) input — A4 of
+      // docs/design/generic-model-support.md. One buildMultimodalPrompt call
+      // splices both media kinds in document order. A request WITH
+      // audio on a model whose tower is unavailable is an explicit 400
+      // — never a silent text-only degrade (that leniency is only for
+      // requests without the media, getVisionTower's contract).
+      const audioTower = getAudioTower(ctx);
+      if (!audioTower || !ctx.audioTokenIds) {
+        throw new RequestError(400,
+          `model ${ctx.modelId} has no audio tower — audio input needs ` +
+          `a model whose config.json carries audio_config and whose ` +
+          `sidecar ships the audio tensors (e.g. gemma-4 e4b OptiQ)`);
+      }
+      let visionSide: { tower: VisionEncoder; tokenIds: VisionTokenIds } | undefined;
+      if (hasImages) {
+        const tower = getVisionTower(ctx);
+        if (!tower) throw new RequestError(400, "model has no vision sidecar");
+        visionSide = { tower, tokenIds: ctx.visionTokenIds };
+      }
+      const { messages: withAudioParts, images } =
+        await extractImages(normalizeMessages(body.messages));
+      const { messages, audio } = await extractAudio(withAudioParts);
+      // Non-WAV containers (mp3/m4a/flac/ogg/aiff/…) transcode to
+      // 16 kHz WAV via CoreAudio; RIFF bytes pass through untouched.
+      // Failures throw into the prompt-build 400.
+      const wavs = await Promise.all(audio.map(ensureWav));
+      // The towers are only ever non-null for Gemma4 (loader gates).
+      const mp = await buildMultimodalPrompt(
+        ctx.model as Gemma4Model,
+        {
+          ...(visionSide ? { vision: visionSide } : {}),
+          audio: { tower: audioTower, tokenIds: ctx.audioTokenIds },
+        },
+        ctx.tokenizer, ctx.template, messages, images, wavs, toolList,
+      );
+      // bidirMask is null whenever audio is present (§3.3 Q1: mixed
+      // prompts run fully causal); the union mask does the per-layer
+      // id zeroing either way.
+      return {
+        ...noMedia,
+        promptIds: mp.ids,
+        vision: {
+          embeddings: mp.embeddings,
+          ...(mp.bidirMask ? { imageMask: mp.bidirMask } : {}),
+          multimodalMask: mp.multimodalMask,
+        },
+      };
+    }
+    if (hasImages && ctx.model instanceof DiffusionGemmaModel) {
+      // DiffusionGemma image-text-to-text: its OWN dedicated SigLIP tower +
+      // encoder vision merge feed the denoising engine (NOT the AR
+      // forwardEmbeddings path). v1 supports a single image.
+      const dm = ctx.model;
+      if (!dm.visionTower) throw new RequestError(400, "this checkpoint has no vision tower");
+      const { messages, images } = await extractImages(normalizeMessages(body.messages));
+      if (images.length !== 1)
+        throw new RequestError(400, "DiffusionGemma image input supports exactly one image");
+      const rendered = ctx.template.render(messages, { tools: toolList, addGenerationPrompt: true });
+      const rawIds = ctx.tokenizer.encode(rendered, /* addSpecialTokens */ false);
+      const { pixels, softTokens } = await dm.visionTower.preprocess(images[0]!);
+      const diffusionPixels = ownership.own(pixels);
+      const promptIds = spliceImageTokens(rawIds, [softTokens], {
+        image: ctx.visionTokenIds.imageTokenId,
+        boi: ctx.visionTokenIds.boiTokenId,
+        eoi: ctx.visionTokenIds.eoiTokenId,
+      });
+      return { ...noMedia, promptIds, vision: undefined, diffusionPixels };
+    }
+    if ((hasImages || hasVideos) && ctx.model instanceof Qwen35Model) {
+      // Qwen3.8 vision + video: the tower is a Qwen3VLVisionTower
+      // riding the shared lazy slot (makeVisionLoader's qwen branch);
+      // image/video spans splice into input embeddings and the request
+      // carries mRoPE positions + delta (PLAN 14v/14w). Videos decode
+      // to sampled frames via the AVFoundation sidecar.
+      const tower = getVisionTower(ctx) as unknown as Qwen3VLVisionTower | null;
+      if (!tower) throw new RequestError(400, "model has no vision sidecar");
+      const { messages: withVideos, images } =
+        await extractImages(normalizeMessages(body.messages));
+      const { messages, videos } = await extractVideos(withVideos);
+      const vp = await buildQwen3VLVisionPrompt(
+        ctx.model, tower, ctx.tokenizer, ctx.template, messages, images,
+        {
+          imageTokenId: (ctx.model.config.raw.image_token_id as number) ?? 248056,
+          videoTokenId: (ctx.model.config.raw.video_token_id as number) ?? 248057,
+          visionStartId: (ctx.model.config.raw.vision_start_token_id as number) ?? 248053,
+          visionEndId: (ctx.model.config.raw.vision_end_token_id as number) ?? 248054,
+        },
+        prep.templateOptionsFor(body, toolList),
+        videos,
+      );
+      return { ...noMedia, promptIds: vp.ids, vision: { embeddings: vp.embeddings, mrope: vp.mrope } };
+    }
+    if (hasImages) {
+      // Loads (and caches) the tower on first image request — text-only
+      // sessions never pay for it.
+      const tower = getVisionTower(ctx);
+      if (!tower) throw new RequestError(400, "model has no vision sidecar");
+      const { messages, images } = await extractImages(normalizeMessages(body.messages));
+      // The tower is only ever non-null for Gemma4 (sidecar gate in
+      // makeVisionLoader), so the model narrow is safe here.
+      const vp = await buildVisionPrompt(
+        ctx.model as Gemma4Model, tower, ctx.tokenizer, ctx.template,
+        messages, images, ctx.visionTokenIds, toolList,
+      );
+      return { ...noMedia, promptIds: vp.ids, vision: { embeddings: vp.embeddings, imageMask: vp.imageMask } };
+    }
+    const text = prep.promptIdsFor(body, toolList);
+    return {
+      promptIds: text.ids,
+      vision: undefined,
+      startInThinking: text.startInThinking,
+      probeStableLen: true,
+      diffusionPixels: null,
+    };
+  }

package/src/backends/mlx/model-serving.ts

@@ -0,0 +1,79 @@
+import type { ServerContext, ServingContext } from "../../serve/model-host";
+import type { ModelServingBinding } from "../../serve/model-binding";
+import { Glm52Model } from "../../model/glm52";
+import { embedMany, isEmbeddingModel } from "../../embed";
+import { bindMlxGateway } from "./gateway-binding";
+import { bindLegacySerialModel, createMlxSerialExecutor } from "./serial-executor";
+import { buildModelPrompt } from "./model-prompt";
+import { MLX_VERSION, deviceArchitecture } from "../../mlx/ffi";
+
+export function modelPromptBuilder(context: ServingContext): ModelServingBinding["buildPrompt"] {
+  return context.serving?.buildPrompt ?? ((body, tools, ownership, prep) =>
+    buildModelPrompt(context as ServerContext, prep, body, tools, ownership));
+}
+
+/** Existing model classes enter the interface once here. New implementations
+ * provide ctx.serving directly and need no forward/makeCache methods. */
+export function modelServingBinding(context: ServingContext): ModelServingBinding {
+  if (context.serving) return context.serving;
+  const ctx = context as ServerContext;
+  const model = ctx.model;
+  if (typeof model.makeCache !== "function" || typeof model.forward !== "function")
+    throw new Error("model implementation must supply its serving binding");
+  const serial = bindLegacySerialModel(model, ctx.draft ?? undefined);
+  const glm = model instanceof Glm52Model;
+  const embedding = isEmbeddingModel(model);
+  return {
+    stateCompatibility: `mlx-${MLX_VERSION}-${deviceArchitecture()}`,
+    gateway: bindMlxGateway(model),
+    createSerial: (services) => createMlxSerialExecutor(serial, services),
+    buildPrompt: (body, tools, ownership, prep) => buildModelPrompt(ctx, prep, body, tools, ownership),
+    restore: (store, entry) => store.restore(entry, model),
+    async signal(ids, count, minimum) {
+      const caches = model.makeCache();
+      try {
+        const logits = model.forward(ids, caches);
+        let f: Float32Array;
+        let vocab: number;
+        try {
+          const [, length, width] = logits.shape as [number, number, number];
+          vocab = width;
+          const last = logits.slice([0, length - 1, 0], [1, length, width]);
+          try { f = last.toFloat32(); } finally { last.dispose(); }
+        } finally { logits.dispose(); }
+        let mx = -Infinity;
+        for (const value of f) if (value > mx) mx = value;
+        let sum = 0;
+        for (const value of f) sum += Math.exp(value - mx);
+        const lse = mx + Math.log(sum);
+        const bins = new Array<number>(count).fill(0);
+        for (const value of f) {
+          const t = Math.max(0, Math.min(1, (value - lse - minimum) / -minimum));
+          const bin = Math.min(count - 1, Math.floor(t * count));
+          bins[bin] = (bins[bin] ?? 0) + 1;
+        }
+        return { bins, vocab };
+      } finally { for (const cache of caches) cache.dispose(); }
+    },
+    discovery: { adapters: !glm, training: !glm, dsa: glm && model.capabilities.dsa, embeddings: embedding },
+    ...(embedding ? { embed: (inputs: string[], instruction?: string) => embedMany(model, ctx.tokenizer, inputs, instruction) } : {}),
+    diagnostics() {
+      const plan = ctx.glmMemoryPlan;
+      if (!plan) return {};
+      const runtime = glm ? model.expertRuntime : undefined;
+      return { glm52: {
+        preset: plan.preset, planned_process_bytes: plan.plannedProcessBytes,
+        process_limit_bytes: plan.processLimitBytes, context_tokens: plan.contextTokens,
+        max_generation_tokens: plan.maxGenerationTokens, batch_size: plan.batchSize,
+        dsa: glm && model.capabilities.dsa, mtp: ctx.draft?.provider.id === "glm52-native-mtp",
+        mtp_draft_tokens: plan.mtpDraftTokens, resident_weight_bytes: plan.lineItems.residentWeightsBytes,
+        main_expert_slab_bytes: plan.lineItems.mainExpertSlabBytes,
+        mtp_expert_slab_bytes: plan.lineItems.mtpExpertSlabBytes,
+        expert_runtime: runtime ? {
+          main_residency: runtime.manager.snapshot(), mtp_residency: runtime.mtp?.manager.snapshot() ?? null,
+          last_turn: runtime.lastTelemetry, last_repin: runtime.lastRepin,
+        } : null,
+      } };
+    },
+  };
+}

package/src/backends/mlx/prefill.ts

@@ -0,0 +1,33 @@
+import type { PrefillStep } from "../../inference/prefill";
+import type { MlxArray } from "../../mlx/array";
+import type { Cache } from "../../model/gemma4";
+import * as ops from "../../mlx/ops";
+import { clearCache } from "../../mlx/ffi";
+import { leaseCacheStates } from "./state-views";
+import { withResource } from "../../engine/resources";
+
+/** State views with explicit ownership are released even on deferred errors.
+ * Owned views precede borrowed state, preserving the serial eval ordering. */
+export function evalCacheState(cache: Cache[]): void {
+  withResource(leaseCacheStates(cache), (state) => ops.evalAll([...state]));
+}
+
+/** Forward once. Drains evaluate state, apply KV maintenance, then clear the
+ * allocator cache. Finals return owned hidden state without forcing eval. */
+export async function executeMlxPrefillStep(
+  forward: (ids: MlxArray, cache: Cache[]) => MlxArray | Promise<MlxArray>,
+  cache: Cache[], prompt: readonly number[], step: PrefillStep,
+  maintain: () => void,
+): Promise<MlxArray | null> {
+  const chunk = prompt.slice(step.start, step.end);
+  const ids = ops.fromInt32(chunk, [1, chunk.length]);
+  let hidden: MlxArray;
+  try { hidden = await forward(ids, cache); }
+  finally { ids.dispose(); }
+  if (step.kind === "final") return hidden;
+  hidden.dispose();
+  evalCacheState(cache);
+  maintain();
+  clearCache();
+  return null;
+}

package/src/backends/mlx/request-policy.ts

@@ -0,0 +1,20 @@
+import type { GenerateOptions } from "../../generate";
+
+/** Copy policy data without cloning native caches, controllers or callbacks. */
+export function snapshotGenerationPolicy<T extends GenerateOptions & { stopSequences?: readonly string[] }>(options: T): T {
+  const snapshot = { ...options };
+  const freeze = (value: unknown): void => {
+    if (!value || typeof value !== "object") return;
+    for (const item of Object.values(value)) freeze(item);
+    Object.freeze(value);
+  };
+  for (const key of ["stopSequences", "eosTokenIds", "adapters", "xtcSpecialTokens", "logitBias",
+    "kvConfig", "turboQuant", "pagedKv", "hlg", "curve", "decodePolicy"] as const) {
+    if (snapshot[key] === undefined) continue;
+    const value = structuredClone(snapshot[key]);
+    freeze(value);
+    Object.assign(snapshot, { [key]: value });
+  }
+  return Object.freeze(snapshot);
+}
+

package/src/backends/mlx/rollback.ts

@@ -0,0 +1,51 @@
+import type { SpeculativeTransaction } from "../../inference/rollback";
+import type { Cache } from "../../model/gemma4-base";
+
+/** Bind the full legacy recurrent transaction or the trim implementation once.
+ * Ring-capacity checks still happen at each safe boundary before any writes. */
+export function bindCacheRollback(caches: readonly Cache[]): SpeculativeTransaction {
+  const layers = caches.map((cache) => {
+    const { specRoundBegin: begin, specRoundCommit: commit, specRoundRollback: rollback } = cache;
+    const hasRound = [begin, commit, rollback].some((method) => method !== undefined);
+    if (hasRound && !(typeof begin === "function" && typeof commit === "function" && typeof rollback === "function"))
+      throw new Error("cache exposes an incomplete speculative transaction");
+    return {
+      fits(drafts: number) {
+        if ("maxSize" in cache && typeof cache.maxSize === "number")
+          return cache.offset + drafts + 1 < cache.maxSize;
+        return hasRound || cache.isTrimmable();
+      },
+      begin: () => begin?.call(cache),
+      resolve(accepted: number, drafts: number) {
+        if (accepted < drafts) {
+          if (rollback) rollback.call(cache, accepted + 1);
+          else cache.trim(drafts - accepted);
+        } else commit?.call(cache);
+      },
+    };
+  });
+  let active: number | undefined;
+  let invalid = false;
+  const validate = (drafts: number) => {
+    if (invalid) throw new Error("speculative state is invalid and must be discarded");
+    if (!Number.isSafeInteger(drafts) || drafts < 0) throw new Error("invalid speculative draft count");
+  };
+  return {
+    canBegin(drafts) { validate(drafts); return active === undefined && layers.every((layer) => layer.fits(drafts)); },
+    begin(drafts) {
+      validate(drafts);
+      if (active !== undefined || !layers.every((layer) => layer.fits(drafts)))
+        throw new Error("speculative state cannot begin this round");
+      active = drafts;
+      try { for (const layer of layers) layer.begin(); }
+      catch (error) { invalid = true; throw error; }
+    },
+    resolve(accepted) {
+      validate(accepted);
+      if (active === undefined || accepted > active) throw new Error("invalid speculative acceptance count");
+      try { for (const layer of layers) layer.resolve(accepted, active); }
+      catch (error) { invalid = true; throw error; }
+      active = undefined;
+    },
+  };
+}

package/src/backends/mlx/serial-executor.ts

@@ -0,0 +1,273 @@
+import { cleanupFailure, disposeResources, ownResource } from "../../engine/resources";
+import { bindGeneration } from "../../generate";
+import type { RuntimeModel } from "../../model/factory";
+import { Qwen35Model } from "../../model/qwen3_5";
+import type { Cache } from "../../model/gemma4";
+import type { PromptCache } from "../../prompt-cache";
+import type { SsdCacheStore } from "../../ssd-cache";
+import type { DraftProvider } from "../../spec/source";
+import type { SerialRun, Vision } from "../../serve/generation-gateway";
+import { generationCheckpointKey } from "../../serve/checkpoint-identity";
+import { runtimeConfig, withRuntimeConfig, type RuntimeConfig } from "../../runtime-config";
+import { bindLegacySpeculativeModel } from "./speculative";
+import { enterMlxMemoryGuard, type MlxMemoryBudget } from "./memory-guard";
+
+/** Native serial execution depends on this bound port, never a model union.
+ * Weights remain borrowed. The gateway supplies the exclusive runtime lease. */
+export interface MlxSerialBinding {
+  readonly runtime: RuntimeConfig;
+  readonly generate: ReturnType<typeof bindGeneration>;
+  readonly speculate?: SerialRun;
+  makeCache(): Cache[];
+  enterMedia?(vision?: Vision): () => void;
+  enterMemoryGuard?(budget: SerialMemoryBudget, promptTokens: number,
+    prefillChunkSize?: number): { check(): void; close(): void };
+}
+
+/** Admission's ceiling and the cache that can release memory under the GPU lease. */
+export type SerialMemoryBudget = MlxMemoryBudget;
+
+/** Family-specific context is bound once at the compatibility boundary. */
+export function bindLegacySerialModel(
+  model: RuntimeModel,
+  draft?: { provider: DraftProvider; numDraftTokens: number },
+): MlxSerialBinding {
+  const speculative = draft ? bindLegacySpeculativeModel(model, draft.provider) : undefined;
+  return {
+    runtime: runtimeConfig(),
+    generate: bindGeneration(model),
+    makeCache: model.makeCache.bind(model),
+    enterMemoryGuard(budget, promptTokens, prefillChunkSize) {
+      return enterMlxMemoryGuard(model, budget, promptTokens, prefillChunkSize);
+    },
+    ...(speculative && draft ? { speculate: (async (prompt, options, onToken) => {
+      const { specRun } = await import("../../spec/serve-loop");
+      return specRun(speculative, draft.numDraftTokens, prompt, options, onToken);
+    }) satisfies SerialRun } : {}),
+    ...(model instanceof Qwen35Model ? { enterMedia(vision?: Vision) {
+      const previous = model.mrope;
+      if (vision?.mrope) model.mrope = vision.mrope;
+      return () => { model.mrope = previous; };
+    } } : {}),
+  };
+}
+
+export interface MlxSerialServices {
+  readonly promptCache: Pick<PromptCache, "take" | "put"> & Partial<Pick<PromptCache, "maxBytes">>;
+  readonly checkpoints: Pick<SsdCacheStore, "findGenerationCheckpoint" | "restore" |
+    "storeGenerationCheckpoint" | "removeGenerationCheckpoints"> | null;
+  readonly checkpointEveryTokens?: number;
+  readonly memoryBudget?: SerialMemoryBudget;
+  /** Artifact, implementation, state ABI and codec identity captured at load. */
+  readonly identity: unknown;
+  adapterNamespace(adapters: string[]): string;
+  cloneState(caches: Cache[]): Cache[];
+}
+
+/** Cache lookup, replay, prompt snapshots and checkpoint persistence belong to
+ * native execution. HTTP only composes this service and frames its output.
+ * The caller transfers prepared native inputs; caches created/taken here are
+ * owned until returned to the prefix store or disposed after execution. */
+export function createMlxSerialExecutor(binding: MlxSerialBinding, services: MlxSerialServices): SerialRun {
+  const { promptCache, checkpoints: ssdStore } = services;
+  return (promptIds, options, onToken, vision, trace, execution) => withRuntimeConfig(binding.runtime, async () => {
+    let caches: Cache[] = [];
+    let retain: (() => void) | undefined;
+    let closeMedia: (() => void) | undefined;
+    let memoryGuard: ReturnType<NonNullable<MlxSerialBinding["enterMemoryGuard"]>> | undefined;
+    const cleanup = ownResource(null, () => disposeResources([
+      { dispose() {
+        disposeResources(caches);
+        retain?.(); // backing release requires successful cache disposal
+      } },
+      { dispose: () => closeMedia?.() },
+      { dispose: () => memoryGuard?.close() },
+      ...[options.grammar, vision?.embeddings, vision?.imageMask,
+        vision?.multimodalMask, options.visionPixels].filter((value) => value != null),
+    ]));
+    try {
+      options.signal?.throwIfAborted();
+      if (!execution) throw new Error("serial execution requires a resolved plan");
+      if (execution.method === "speculative") {
+        if (!binding.speculate) throw new Error("resolved speculation requires a bound verifier");
+        return await binding.speculate(promptIds, {
+          ...options,
+          speculativeCacheBytes: binding.runtime.flag("MLX_BUN_MTP_PROMPT_CACHE", false)
+            ? services.promptCache.maxBytes ?? 0 : 0,
+        }, onToken);
+      }
+      // Cache entries are adapter-specific: KV computed under one adapter
+      // must never seed another's (or the base's) prefill.
+      const cacheNs = options.adapters?.length ? services.adapterNamespace(options.adapters) : "";
+      // Paged-KV request scope (docs/design/kv-cache.md): media
+      // prompts (bidir overlay) and LoRA-adapter requests are v1 non-goals —
+      // they run the PLAIN cache path even under --paged-kv (scope the flag
+      // per request, never 400). Effective value computed ONCE so the
+      // prompt-cache bypass below and the generate() options can't disagree.
+      const pagedKv = execution.pagedKv ? options.pagedKv : undefined;
+      // Paged requests bypass the prompt cache entirely (v1 non-goal:
+      // PagedKVCache has no cloneKvCaches/restore path — the vision
+      // precedent). Fresh caches per request, disposed on completion.
+      const skipPromptCache = !execution.promptCache;
+      const checkpointEvery = services.checkpointEveryTokens;
+      const checkpointEligible = execution.checkpoint;
+      const checkpointKey = checkpointEligible
+        ? generationCheckpointKey(promptIds, options, cacheNs, execution, services.identity)
+        : null;
+      // Both tiers in one call (Layer 0): take() prefers a strictly-longer
+      // SSD prefix, restores it zero-copy, and trims — see PromptCache.take.
+      const closeCacheLookup = trace?.begin("cache.lookup_restore", {
+        mechanism: "serial",
+        bypassed: skipPromptCache,
+      });
+      const checkpointEntry = checkpointKey
+        ? ssdStore!.findGenerationCheckpoint(promptIds, checkpointKey, cacheNs)
+        : null;
+      const restoredCheckpoint = checkpointEntry
+        ? ssdStore!.restore(checkpointEntry, binding)
+        : null;
+      caches = restoredCheckpoint?.caches ?? [];
+      const checkpoint = restoredCheckpoint?.header.generationCheckpoint;
+      if (restoredCheckpoint && !checkpoint)
+        throw new Error("restored generation checkpoint has no continuation metadata");
+      const resuming = Boolean(restoredCheckpoint && checkpoint);
+      const generationPromptIds = resuming ? restoredCheckpoint!.tokens : promptIds;
+      const entry = skipPromptCache || resuming
+        ? null
+        : promptCache.take(promptIds, cacheNs);
+      if (entry) { caches = entry.caches; retain = entry.retain; }
+      if (!restoredCheckpoint && !entry) caches = binding.makeCache();
+      closeCacheLookup?.();
+      // Prompt-boundary snapshot (the multi-turn agent fix, 2026-07-04): the
+      // prompt+gen entry put() below is UNTRIMMABLE at context > sliding
+      // window (wrapped rings) and under quantized KV (mid-group), so any
+      // decode→encode roundtrip drift in the reply the client sends back
+      // turns the next turn into a total miss (measured: 12B turn-2 TTFT
+      // 8.9 s instead of ~0.2 s). A prompt-ONLY entry is always an exact
+      // prefix of the next turn's rendering regardless of reply drift.
+      // Zero-copy (cloneKvCaches = slice views); only for substantial cold
+      // prefills, where the re-prefill it saves is worth an extra entry.
+      // The oracle invariant (mlx-lm insert_segments): a trim-free STRICT
+      // prefix of the prompt exists for EVERY substantial request — cap the
+      // boundary at len-1 so even a stableLen == len prompt (e4b: the
+      // template tail survives the probe render) snapshots prompt[:-1]. An
+      // exact repeat then matches with trimNeeded == 0, bypassing
+      // isTrimmable() entirely — the only reuse path a wrapped ring has.
+      const boundary = Math.min(options.snapshotAt ?? promptIds.length, promptIds.length - 1);
+      // Re-snapshot on EVERY substantial request whose stable boundary extends
+      // past the cached prefix; the clone is zero-copy views, so re-putting
+      // is ~free.
+      const snapshotBoundary =
+        !skipPromptCache && !resuming && boundary >= 256 &&
+        boundary > (entry?.tokens.length ?? 0);
+      if (services.memoryBudget) {
+        memoryGuard = binding.enterMemoryGuard?.(
+          services.memoryBudget, generationPromptIds.length, options.prefillChunkSize,
+        );
+        memoryGuard?.check();
+      }
+      closeMedia = binding.enterMedia?.(vision);
+      if (resuming) {
+        const replay = generationPromptIds.slice(promptIds.length);
+        console.log(
+          `[generation-checkpoint] resuming ${checkpointKey} at ` +
+          `${replay.length} emitted tokens`,
+        );
+        for (const token of replay) {
+          options.signal?.throwIfAborted();
+          if ((await onToken(token)) === false)
+            throw new Error("saved generation prefix triggered a terminal stop while replaying");
+        }
+      }
+      const gen = binding.generate(generationPromptIds, {
+        ...options,
+        decodePolicy: execution,
+        ...(resuming ? { seed: checkpoint!.seed } : {}),
+        fill: execution.fill ? options.fill : undefined,
+        pagedKv, // request-scoped (undefined strips the server-wide flag)
+        cache: caches,
+        ...(resuming
+          ? {
+              initialPendingToken: checkpoint!.pendingToken,
+              initialGeneratedTokens: checkpoint!.generatedTokens,
+              originalPromptTokens: checkpoint!.originalPromptTokens,
+            }
+          : {}),
+        ...(checkpointEligible
+          ? {
+              checkpointEveryTokens: checkpointEvery,
+              onDecodeCheckpoint: async (state: {
+                cacheTokens: number[];
+                caches: Cache[];
+                generatedTokens: number;
+                pendingToken: number;
+              }) => {
+                const stored = await ssdStore!.storeGenerationCheckpoint(
+                  state.cacheTokens,
+                  state.caches,
+                  {
+                    key: checkpointKey!,
+                    cacheNs,
+                    originalPromptTokens: promptIds.length,
+                    generatedTokens: state.generatedTokens,
+                    pendingToken: state.pendingToken,
+                    seed: options.seed ?? 0,
+                    seedWasExplicit: options.seedWasExplicit === true,
+                  },
+                );
+                if (stored)
+                  console.log(
+                    `[generation-checkpoint] saved ${state.generatedTokens} emitted tokens`,
+                  );
+              },
+            }
+          : {}),
+        ...(snapshotBoundary
+          ? {
+              // snapshotAt MUST travel with the hook: generate() splits the
+              // prefill at exactly this many tokens and fires the hook while
+              // the caches hold exactly that prefix — putting boundary
+              // tokens against caches at any other offset is silent KV
+              // corruption.
+              snapshotAt: boundary,
+              onPrefillDone: () => {
+                try {
+                  const snapshot = ownResource(services.cloneState(caches), disposeResources);
+                  try {
+                    promptCache.put(promptIds.slice(0, boundary), snapshot.borrow(), cacheNs);
+                    snapshot.transfer();
+                  } finally { snapshot.close(); }
+                } catch (e) {
+                  console.warn(`prompt-boundary snapshot skipped: ${(e as Error).message}`);
+                }
+              },
+            }
+          : {}),
+        ...(vision
+          ? {
+              promptEmbeddings: vision.embeddings,
+              ...(vision.imageMask ? { imageMask: vision.imageMask } : {}),
+              ...(vision.multimodalMask ? { multimodalMask: vision.multimodalMask } : {}),
+            }
+          : {}),
+      }, { trace, mechanism: "serial" });
+      for await (const t of gen) {
+        if (t.index % 256 === 0) memoryGuard?.check();
+        if ((await onToken(t.token, t.logprobs)) === false) break;
+      }
+      const s = gen.stats!; // set on completion AND on early break
+      if (checkpointKey) ssdStore!.removeGenerationCheckpoints(checkpointKey);
+      if (!skipPromptCache) {
+        // put() fires onPut → the debounced write-behind SSD snapshot
+        // (wired below), covering the batch lane's puts too.
+        promptCache.put(s.cacheTokens, caches, cacheNs, retain);
+        caches = []; retain = undefined; // ownership returned to the prefix store
+      }
+      return s;
+    } catch (error) {
+      cleanupFailure(error, () => cleanup.close());
+    } finally {
+      cleanup.close();
+    }
+  });
+}