mlx-bun 0.3.0 → 0.5.0

package/src/serve/generated-token-history.ts

@@ -0,0 +1,57 @@
+/** Token provenance for rendered multi-turn prompts. A generated BPE token
+ * sequence need not equal encode(decode(ids)). Preserve the sequence that
+ * produced the saved inference state when its rendered text is unchanged.
+ * This stores only text and IDs; KV ownership and persistence live elsewhere. */
+export interface PromptTokenHistory {
+  remember(tokens: readonly number[]): void;
+  resolve(rendered: string, canonical: number[]): number[];
+}
+
+interface TokenCodec {
+  encode(text: string, addSpecialTokens?: boolean): number[];
+  decode(ids: number[], skipSpecialTokens?: boolean): string;
+}
+
+export class GeneratedTokenHistory implements PromptTokenHistory {
+  readonly #entries = new Map<string, number[]>();
+  #chars = 0;
+  constructor(readonly codec: TokenCodec, readonly maxChars = 2 ** 21, readonly maxEntries = 16) {}
+
+  remember(tokens: readonly number[]): void {
+    if (!tokens.length) return;
+    // Repeated requests can publish the same generated checkpoint. Its text
+    // and IDs are already owned here; refresh recency without decoding the
+    // entire conversation again on the request's completion path.
+    for (const [text, ids] of this.#entries) {
+      if (ids.length !== tokens.length || !ids.every((id, i) => id === tokens[i])) continue;
+      this.#entries.delete(text);
+      this.#entries.set(text, ids);
+      return;
+    }
+    const ids = [...tokens], text = this.codec.decode(ids, false);
+    // An unfinished UTF-8 token is not a text boundary. A later checkpoint
+    // can supply a complete prefix instead.
+    if (!text.length || text.endsWith("\uFFFD") || text.length > this.maxChars) return;
+    if (this.#entries.delete(text)) this.#chars -= text.length;
+    this.#entries.set(text, ids); this.#chars += text.length;
+    while (this.#entries.size > this.maxEntries || this.#chars > this.maxChars) {
+      const oldest = this.#entries.keys().next().value!;
+      this.#entries.delete(oldest); this.#chars -= oldest.length;
+    }
+  }
+
+  resolve(rendered: string, canonical: number[]): number[] {
+    let best: { text: string; ids: number[] } | undefined;
+    for (const [text, ids] of this.#entries) {
+      if (text.length >= rendered.length || (best && text.length <= best.text.length)) continue;
+      if (rendered.startsWith(text)) best = { text, ids };
+    }
+    if (!best) return canonical;
+    // Normal canonical output keeps the existing tokenizer fast path.
+    if (best.ids.every((id, i) => canonical[i] === id)) return canonical;
+    const ids = [...best.ids, ...this.codec.encode(rendered.slice(best.text.length), false)];
+    // Decoder cleanup/normalization (e.g. SentencePiece's leading-space
+    // rules) can make a standalone suffix unsafe to append.
+    return this.codec.decode(ids, false) === rendered ? ids : canonical;
+  }
+}

package/src/serve/generation-gateway.ts

@@ -29,7 +29,7 @@ import { cleanupFailure, disposeResources } from "../engine/resources";
 // server.py's all-caches-have-merge check). GLM's capability owns compressed
 // MLA/DSA merge/extract/filter and byte projection; Qwen3.5's SSM cache owns
 // its row state. The remaining request-shape exclusions include vision, LoRA,
-// serial-only KV modes, and user-fixed seed
+// serial-only KV modes
 // (reproducibility ⇒ solo, matching mlx-lm's _is_batchable). Temperature /
 // top-p / top-k DO batch (each row samples with its own seed). Full-attention
 // AND sliding-window (Gemma) models both batch — the scheduler assembles each
@@ -42,9 +42,9 @@ import { MlxArray } from "../mlx/array";
 import type { RuntimeModel } from "../model/factory";
 import type { GenerateOptions, GenerateStats, TokenLogprobs } from "../generate";
 import type { KvScheme } from "../kv-scheme";
+import { createRowSampling } from "../backends/mlx/row-sampling";
 import { makeStepSampler } from "../sampler";
 import type { RowPromptCache } from "../backends/mlx/batch-group";
-import type { MlxMemoryBudget } from "../backends/mlx/memory-guard";
 import { bindMlxGateway, type MlxGatewayBinding, type MlxBatchGroup } from "../backends/mlx/gateway-binding";
 import type { RuntimeConfig } from "../runtime-config";
 import type { ExecutionRequirements, ResolvedExecution } from "../contracts/execution";
@@ -94,11 +94,12 @@ class AsyncMutex {
   }
 }
 
-/** Embeddings-prefill payload for media prompts (vision and/or audio — the
- *  name predates audio). Any request carrying one routes to the serial lane
- *  (shape.hasVision) and prefills through generate()'s promptEmbeddings
- *  path, bypassing the prompt cache. */
+/** Prepared embeddings for vision/audio prompts. Qualified models consume
+ * this through shared preparation and decode; other bindings retain serial
+ * execution. Token-only prefix caching remains disabled for media. */
 export type Vision = {
+  /** Producer identity includes all media and preceding rendered tokens. */
+  prefixIdentity?: string;
   embeddings: MlxArray;
   /** bool [L] image-token mask for the bidirectional attention overlay.
    *  Absent when the prompt carries ANY audio — audio(-containing) prompts
@@ -109,12 +110,12 @@ export type Vision = {
    *  where zeroing falls back to imageMask. */
   multimodalMask?: MlxArray;
   /** Qwen3.5/3.8 vision: the request's mRoPE positions + decode delta,
-   *  installed on the model for exactly this serial run (server-side). */
+   *  owned by shared forward input, or scoped to the explicit serial run. */
   mrope?: import("../model/qwen3-mrope").MropeRequestState;
 };
 
 /** Per-token sink: returning `false` halts this generation (stop sequence).
- *  `logprobs` is only populated on the serial lane when the request asked for
+ *  `logprobs` is populated on either lane when the request asked for
  *  logprobs capture (GenerateOptions.logprobs / topLogprobs). */
 export type OnToken = (
   token: number,
@@ -172,7 +173,6 @@ export class GenerationGateway {
     private readonly serialRun: SerialRun,
     private readonly opts: {
       kvBudgetBytes?: number;
-      memoryBudget?: MlxMemoryBudget;
       checkpoints?: boolean;
       stateCodecs?: import("../kv-store").CacheCodecProvider;
       /** The server-wide KV scheme (server.ts kvScheme) — threaded to the
@@ -185,9 +185,7 @@ export class GenerationGateway {
        *  back on finish. Safe to share with the serial lane — both use it
        *  only inside this gateway's mutual-exclusion domain. */
       promptCache?: RowPromptCache;
-      /** Finish cache writes before taking the generation lock. Writers
-       * need that same lock, so this hook must never run inside it. */
-      beforeSerial?: () => Promise<void>;
+      adapterNamespace?: (adapters: string[]) => string;
     } = {},
   ) {
     this.#binding = "createBatchGroup" in model ? model : bindMlxGateway(model);
@@ -235,16 +233,20 @@ export class GenerationGateway {
     await this.onIdle();
   }
 
-  /** Resolves once the engine is idle (poll-based, ~`pollMs` granularity —
-   *  cheap vs the multi-hundred-ms flushes it paces; resolves without a
-   *  timer when already idle). The SSD write-behind gates each per-tensor
-   *  flush step on this so durability work never steals blocking
-   *  GPU-sync + writeSync slices from an active decode (the 2026-07-07
-   *  decode@ctx contamination — see server.ts's write-behind block). */
+  /** Observe idleness. Call runWhenIdle for background work that must also
+   * acquire ownership atomically with the idle check. */
   async onIdle(pollMs = 20): Promise<void> {
     while (this.busy) await new Promise<void>((r) => setTimeout(r, pollMs));
   }
 
+  /** Background work waits without requesting a batch drain. The final busy
+   * check and mutex acquisition have no await between them, so a foreground
+   * request cannot enter between the check and ownership registration. */
+  async runWhenIdle<T>(fn: () => Promise<T>, pollMs = 20): Promise<T> {
+    while (this.busy) await new Promise<void>(resolve => setTimeout(resolve, pollMs));
+    return this.runExclusive(fn);
+  }
+
   /** Cumulative rows routed to the batch lane since server start. The serial
    *  lane never advances this — /stats' race-free lane-routing observable
    *  (active/pending are instantaneous and read 0 once a request finishes). */
@@ -272,7 +274,7 @@ export class GenerationGateway {
     const frozenShape = Object.freeze(shape);
     const execution = this.#binding.plan(frozenShape, options, {
       continuous: this.batchingEnabled,
-      quantizedBatch: shape.kvQuant && this.#kvBatchable(),
+      quantizedBatch: (shape.kvQuant || shape.turboQuant) && this.#kvBatchable(),
       checkpoints: this.opts.checkpoints === true,
     });
     return Object.freeze({ shape: frozenShape, mechanism: execution.mechanism, execution });
@@ -317,6 +319,19 @@ export class GenerationGateway {
     try { return await fn(); } finally { lease.dispose(); }
   }
 
+  /** Preparation reads model weights without replacing active model state.
+   * A compatible backend can perform it between existing decode iterations. */
+  runPreparation<T>(work: () => Promise<T>, signal?: AbortSignal): Promise<T> {
+    if (this.#closed) return Promise.reject(new Error("generation gateway is closed"));
+    if (signal?.aborted) return Promise.reject(signal.reason);
+    const group = this.batchingEnabled ? this.#ensureScheduler() : undefined;
+    return group?.runPreparation
+      ? group.runPreparation(work, signal)
+      : this.runExclusive(work, undefined, signal);
+  }
+
+  get mediaBatchingEnabled(): boolean { return this.batchingEnabled && !!this.#binding.mediaInput; }
+
   /** Run one generation on the appropriate lane. onToken is invoked per emitted
    *  token (its `false` halts); resolves with stats when the generation ends. */
   async run(
@@ -330,10 +345,32 @@ export class GenerationGateway {
     trace?: PromptResponseTrace,
   ): Promise<GenerateStats> {
     let reservation;
+    const closeRequestWait = trace?.begin("engine.request_wait", { capacity: this.#batch });
     try { reservation = await acquireReservation(this.#requests, signal); }
     catch (error) { cleanupFailure(error, () => disposeUnstartedRequest(options, vision)); }
-    try { return await this.#run(promptIds, options, onToken, vision, shape, placement, signal, trace); }
-    finally { reservation.dispose(); }
+    finally { closeRequestWait?.(); }
+    let releasePrefix: (() => void) | undefined;
+    try {
+      try {
+        const adapters = options.adapters?.length
+          ? this.opts.adapterNamespace?.(options.adapters) ?? JSON.stringify(options.adapters) : "";
+        const baseNamespace = this.#binding.prefixNamespace?.(placement.execution, options, adapters) ??
+          (placement.execution?.method === "speculative" ? null : adapters);
+        const namespace = baseNamespace !== null && vision?.prefixIdentity
+          ? JSON.stringify(["prepared-prefix-v1", vision.prefixIdentity, baseNamespace]) : baseNamespace;
+        if ((!vision || placement.execution?.promptCache) && namespace !== null) {
+          const closePrefetch = this.opts.promptCache?.prefetch ? trace?.begin("cache.prefetch") : undefined;
+          try { releasePrefix = await this.opts.promptCache?.prefetch?.(promptIds, namespace, options.cacheSessionId); }
+          finally { closePrefetch?.(); }
+        }
+      } catch (error) { cleanupFailure(error, () => disposeUnstartedRequest(options,
+        placement.mechanism === "continuous" ? undefined : vision)); }
+      return await this.#run(promptIds, options, onToken, vision, shape, placement, signal, trace);
+    } finally {
+      disposeResources([{ dispose: () => releasePrefix?.() }, reservation,
+        ...(placement.mechanism === "continuous" && vision
+          ? [vision.embeddings, vision.imageMask, vision.multimodalMask].filter(value => value != null) : [])]);
+    }
   }
 
   async #run(
@@ -346,7 +383,8 @@ export class GenerationGateway {
     signal?: AbortSignal,
     trace?: PromptResponseTrace,
   ): Promise<GenerateStats> {
-    const disposeUnstarted = () => disposeUnstartedRequest(options, vision);
+    const disposeUnstarted = () => disposeUnstartedRequest(options,
+      placement.mechanism === "continuous" ? undefined : vision);
     if (placement.shape !== shape) {
       disposeUnstarted();
       throw new Error("generation placement does not belong to this request shape");
@@ -390,7 +428,7 @@ export class GenerationGateway {
         });
       };
       let started = false;
-      return Promise.resolve().then(() => this.opts.beforeSerial?.()).then(() => this.runExclusive(async () => {
+      return this.runExclusive(async () => {
         signal?.throwIfAborted();
         started = true;
         const stats = await this.serialRun(
@@ -398,7 +436,7 @@ export class GenerationGateway {
         );
         signal?.throwIfAborted();
         return stats;
-      }, trace, signal))
+      }, trace, signal)
         // Covers a serial waiter aborted before serialRun takes ownership.
         // generate() also disposes defensively; the operation is idempotent.
         .finally(() => {
@@ -407,16 +445,25 @@ export class GenerationGateway {
         });
     }
 
+    const adapters = options.adapters?.length ? [...options.adapters] : undefined;
+    const adapterKey = adapters ? JSON.stringify(adapters) : "";
+    const adapterNamespace = () => adapters ? this.opts.adapterNamespace?.(adapters) ?? adapterKey : "";
+    const cacheNamespace = vision?.prefixIdentity
+      ? () => JSON.stringify(["prepared-prefix-v1", vision.prefixIdentity, adapterNamespace()])
+      : adapters ? adapterNamespace : "";
+    const context = adapters
+      ? this.#binding.bindAdapterContext?.(adapters, `adapters:${adapterKey}`) : undefined;
+
     // The scheduler owns grammar ready/accept sequencing. StepSampler owns the
     // shared processors -> mask -> logprobs -> sample -> history contract.
-    const stepSampler = makeStepSampler(options, {
-      tokenRepresentation: "device",
-      grammarWait: "external",
-      historyUpdate: "after-sample",
-      initialHistory: promptIds,
-    });
-    const sample = (logits1V: MlxArray, step: number): MlxArray =>
-      stepSampler.sample(logits1V, step).token;
+    const method = this.#binding.methodRequest?.(placement.execution, options);
+    const continuation = method ? undefined : this.#binding.continuationRequest?.(placement.execution, options, promptIds, onToken);
+    const sampling = method ? undefined : continuation ?? createRowSampling(makeStepSampler(options, {
+      tokenRepresentation: "device", grammarWait: "external",
+      historyUpdate: "after-sample", initialHistory: promptIds,
+      captureSelectedLogprob: options.logprobs === true,
+      captureTopLogprobs: options.topLogprobs,
+    }), onToken);
 
     let st;
     this.#rowsSubmitted++;
@@ -425,13 +472,15 @@ export class GenerationGateway {
     });
     try {
       st = await this.#ensureScheduler().submit({
-        promptIds,
+        promptIds, context, cacheNamespace, cacheSessionId: options.cacheSessionId, prefillChunkSize: options.prefillChunkSize,
+        ...(vision ? { promptInput: this.#binding.mediaInput!(vision) } : {}),
+        continuation: continuation?.continuation,
+        statePolicy: this.#binding.statePolicy?.(placement.execution, options, promptIds.length + (options.maxTokens ?? 512)),
         compiledDecode: placement.execution?.compiledDecode,
         maxTokens: options.maxTokens ?? 512,
         eosTokenIds: options.eosTokenIds ?? this.#binding.config.eosTokenIds,
-        sample,
-        plainGreedy: stepSampler.isPlainGreedy,
-        onToken,
+        ...(method ? { method, onToken } : { sample: sampling!.sample,
+          plainGreedy: sampling!.plainGreedy, onToken: sampling!.onToken }),
         onAdmitted: closeAdmission,
         trace,
         ...(signal ? { signal } : {}),
@@ -445,7 +494,7 @@ export class GenerationGateway {
       });
     } finally {
       closeAdmission?.();
-      stepSampler.dispose();
+      sampling?.dispose();
       // The scheduler never owns per-row grammar state.
       options.grammar?.dispose();
     }
@@ -460,6 +509,8 @@ export class GenerationGateway {
       prefillTps: st.prefillMs > 0 ? ((st.promptTokens - st.cachedTokens) / st.prefillMs) * 1000 : 0,
       decodeTps: st.decodeMs > 0 && st.generatedTokens > 1 ? ((st.generatedTokens - 1) / st.decodeMs) * 1000 : 0,
       cacheTokens: [],
+      ...(st.spec ? { spec: st.spec } : {}),
+      ...(st.fill ? { fill: st.fill } : {}),
     };
   }
 
@@ -470,7 +521,6 @@ export class GenerationGateway {
         maxBatch: this.#batch,
         stateCodecs: this.opts.stateCodecs,
         kvBudgetBytes: this.opts.kvBudgetBytes,
-        memoryBudget: this.opts.memoryBudget,
         // Phase 3.1: the batchable kvConfig composition is applied by the
         // scheduler (solo rows convert at serial chunk boundaries, then
         // merge as quantized triples). Only threaded when the scheme

package/src/serve/model-binding.ts

@@ -21,9 +21,14 @@ export interface BuiltPrompt {
   diffusionPixels: MlxArray | null;
 }
 
+/** Enter the model's execution domain for native preparation. Fetching and
+ * container transcoding stay outside this callback. */
+export type PromptNativeWork = <T>(work: () => Promise<T>) => Promise<T>;
+
 export interface ModelPromptBuilder {
   (body: ChatRequestParams, tools: ChatRequestParams["tools"] | null,
-    ownership: RequestOwnership, prep: RequestPrep): Promise<BuiltPrompt>;
+    ownership: RequestOwnership, prep: RequestPrep, nativeWork?: PromptNativeWork,
+    objects?: import("../contracts/object-cache").ObjectCache<import("../backends/mlx/checkpoint-state").CheckpointAttachment[]>): Promise<BuiltPrompt>;
 }
 
 /** Model-owned native services. A replacement supplies this one binding;
@@ -35,6 +40,7 @@ export interface ModelServingBinding {
   createSerial(services: MlxSerialServices): SerialRun;
   readonly buildPrompt: ModelPromptBuilder;
   restore(store: SsdCacheStore, entry: SsdIndexEntry): ReturnType<SsdCacheStore["restore"]>;
+  restoreAsync?(store: SsdCacheStore, entry: SsdIndexEntry): ReturnType<SsdCacheStore["restoreAsync"]>;
   signal(promptIds: number[], bins: number, minimum: number): Promise<{ bins: number[]; vocab: number }>;
   diagnostics(): Record<string, unknown>;
   readonly discovery: { readonly adapters: boolean; readonly training: boolean;

package/src/serve/model-host.ts

@@ -73,10 +73,9 @@ export interface ServerContext<Model = RuntimeModel> {
    *  optiq serve injects these as server defaults (gen_config.py);
    *  explicit request fields always win. */
   genDefaults: GenSamplingDefaults;
-  /** Speculative decoding (`serve --draft-model`, mlx_lm.server parity).
-   *  Server-level: when set, EVERY request routes to the serial lane
-   *  (upstream: is_batchable = draft is None) and spec-eligible ones decode
-   *  through src/spec/serve-loop.ts. null = no draft configured. */
+  /** Spec-eligible requests use src/spec/serve-loop.ts; ordinary fallbacks
+   *  use their own feature-compatible scheduling mechanism.
+   *  null = no draft configured. */
   draft?: {
     provider: import("../spec/source").DraftProvider;
     numDraftTokens: number;

package/src/serve/preparation.ts

@@ -17,11 +17,11 @@ export async function acquireReservation(pool: Pick<AdmissionPool, "acquire">, s
   finally { signal?.removeEventListener("abort", abort); }
 }
 
-/** One retained media preparation at a time; grammar rows may form a batch. */
+/** Bound retained preparations independently from the execution queue. */
 export function createPreparationExecutor(
-  run: PreparationExecutor["run"], grammarCapacity: number,
+  run: PreparationExecutor["run"], grammarCapacity: number, mediaCapacity = 1,
 ): PreparationExecutor & { close(): void } {
-  const media = new AdmissionPool(1);
+  const media = new AdmissionPool(Math.max(1, mediaCapacity));
   const constraints = new AdmissionPool(Math.max(1, grammarCapacity));
   return {
     run,

package/src/serve/prompt-response-trace.ts

@@ -10,16 +10,27 @@ export type P2RTracePhase =
   | "request.prompt_prepare"
   | "completion.total"
   | "completion.placement"
+  | "engine.request_wait"
   | "engine.admission_wait"
+  | "engine.mixed_forward"
   | "cache.lookup_restore"
+  | "cache.prefetch"
   | "prefill.total"
   | "prefill.batch_setup"
   | "prefill.chunk"
+  | "prefill.row_sync"
+  | "prefill.forward"
+  | "prefill.evaluate"
+  | "prefill.checkpoint"
+  | "prefill.project"
+  | "prefill.complete"
+  | "prefill.companion"
   | "prefill.kv_maintenance"
   | "token_zero.total"
   | "token_zero.forward"
   | "token_zero.head"
   | "token_zero.sample"
+  | "response.token_route"
   | "response.first_write"
   | "response.final_write";
 
@@ -42,6 +53,8 @@ export interface P2RTraceRecord {
   route: string;
   clock: "monotonic-ms";
   outcome: P2RTraceOutcome;
+  /** Process-local monotonic origin; aligns concurrent requests in one server. */
+  startedAtMs?: number;
   totalMs: number;
   events: P2RTraceEvent[];
   attributes?: P2RTraceAttributes;
@@ -129,6 +142,7 @@ export class PromptResponseTrace {
       route: this.route,
       clock: "monotonic-ms",
       outcome,
+      startedAtMs: this.#startedAt,
       totalMs,
       events: [...this.#events].sort((a, b) => a.startMs - b.startMs),
       ...(attributes ? { attributes } : {}),

package/src/serve/request-plan.ts

@@ -33,12 +33,13 @@ export interface RequestPlanInput {
   promptIds: number[];
   options: GenerateOptions & { stopSequences: string[] };
   requestedMaxTokens: number;
-  maxSafeContext: number;
+  contextLimit: number | null;
   stream: boolean;
   wantLogprobs: boolean;
   topLogprobs: number;
   adapterIds: string[];
   hasVision: boolean;
+  hasPreparedPrefixIdentity?: boolean;
   userSeed: boolean;
   hasGrammar: boolean;
   hasDraft: boolean;
@@ -114,11 +115,12 @@ export type PlanRequestResult = RequestPlan | RequestRejection;
  * options, adapter selection, and lane shape are derived here once for chat
  * and raw completions. */
 export function planRequest(input: RequestPlanInput): PlanRequestResult {
-  const available = input.maxSafeContext - input.promptIds.length;
+  const available = input.contextLimit === null ? Infinity
+    : input.contextLimit - input.promptIds.length;
   if (available < 1) {
     return new RequestRejection(
       `prompt is ${input.promptIds.length} tokens but the memory budget caps ` +
-        `safe context at ${input.maxSafeContext} — no room to generate; ` +
+        `safe context at ${input.contextLimit} — no room to generate; ` +
         "shorten the prompt or raise --memory-budget",
       input.ownership,
     );
@@ -136,6 +138,7 @@ export function planRequest(input: RequestPlanInput): PlanRequestResult {
 
   const shape: RequestShape = {
     hasVision: input.hasVision,
+    hasPreparedPrefixIdentity: input.hasPreparedPrefixIdentity,
     hasAdapters: !!options.adapters?.length,
     hasRepetitionPenalty: !!options.repetitionPenalty,
     hasLogitsExtras: !!(

package/src/serve/request-prep.ts

@@ -5,7 +5,7 @@
 // tool-call stream router. Resolved ONCE per server (createRequestPrep) and
 // captured by the route handlers. Extracted from src/server.ts (repo-taming
 // Phase 4).
-import type { ToolDefinition } from "../chat-template";
+import type { ChatMessage, ToolDefinition } from "../chat-template";
 import {
   fillEchoConfig,
   fillMaxSpan,
@@ -27,6 +27,7 @@ import {
   type ChatRequestParams,
 } from "./chat-request";
 import type { ServingContext } from "./model-host";
+import type { PromptTokenHistory } from "./generated-token-history";
 import { selectToolStreamMode, ToolAwareStream, type ToolStreamMode } from "./token-streams";
 
 /** The ServerOptions fields request preparation reads (server-wide
@@ -49,6 +50,7 @@ export function createRequestPrep(input: {
   /** `--max-tokens` / the GLM memory plan's generation cap; undefined = the
    *  surface's own default (chat 65 536, raw completion 512). */
   defaultGeneratedTokens: number | undefined;
+  tokenHistory?: PromptTokenHistory;
 }) {
   const { ctx, serverOptions, kvScheme, defaultGeneratedTokens } = input;
 
@@ -92,10 +94,10 @@ export function createRequestPrep(input: {
       ? Math.min(genTemp, NO_THINK_TEMPERATURE)
       : genTemp;
     return {
-      // No invented cap: unset means "generate until EOS or the admitted
-      // context is exhausted" — admission (planRequest) is the real limit
-      // and clamps to maxSafeContext − promptLen. Only an explicit request
-      // max_tokens or the operator's --max-tokens narrows it further.
+      cacheSessionId: typeof req.session_id === "string" ? req.session_id :
+        typeof req.prompt_cache_key === "string" ? req.prompt_cache_key : undefined,
+      // Unset means generate until EOS/stop. Only a client cap or an explicit
+      // operator limit constrains completion length; fit estimates do not.
       maxTokens: req.max_completion_tokens ?? req.max_tokens ??
         defaultGeneratedTokens,
       temperature: req.temperature ?? serverOptions.defaultTemperature ?? defaultTemp,
@@ -193,7 +195,8 @@ export function createRequestPrep(input: {
     const ids = ctx.tokenizer.encode(rendered);
     // template includes <bos>; tokenizer post-processor also prepends one
     const trimmed = ids[0] === ids[1] && ids[0] === ctx.tokenizer.bosTokenId ? ids.slice(1) : ids;
-    return { ids: trimmed, startInThinking: promptEndsInOpenThink(rendered) };
+    return { ids: input.tokenHistory?.resolve(rendered, trimmed) ?? trimmed,
+      startInThinking: promptEndsInOpenThink(rendered) };
   };
 
   /** STABLE cache boundary: the prompt's tail is a generation primer (e.g.
@@ -221,18 +224,32 @@ export function createRequestPrep(input: {
     if (primer !== undefined) return Math.max(0, trimmed.length - primer);
     let stableLen = trimmed.length;
     try {
-      const probe = ctx.tokenizer.encode(
-        ctx.template.render(
-          [...normalizeMessages(req.messages), { role: "assistant", content: "x" }],
-          opts,
-        ),
-      );
-      const probeTrimmed =
-        probe[0] === probe[1] && probe[0] === ctx.tokenizer.bosTokenId ? probe.slice(1) : probe;
-      let i = 0;
-      const n = Math.min(trimmed.length, probeTrimmed.length);
-      while (i < n && trimmed[i] === probeTrimmed[i]) i++;
-      stableLen = i;
+      const canonical = ctx.tokenizer.encode(ctx.template.render(normalizeMessages(req.messages), opts));
+      const current = canonical[0] === canonical[1] && canonical[0] === ctx.tokenizer.bosTokenId
+        ? canonical.slice(1) : canonical;
+      const replies: ChatMessage[] = [{ role: "assistant", content: "x" }];
+      // Some templates retain an empty think primer for text replies but
+      // omit it before a tool call. Both reply forms must preserve the
+      // boundary retained by an untrimmable recurrent cache.
+      if (tools?.length) replies.push({ role: "assistant", content: null, tool_calls: [{
+        id: "cache-boundary-probe", type: "function",
+        function: { name: tools[0]!.function.name, arguments: {} },
+      }] });
+      let i = current.length;
+      for (const reply of replies) {
+        const probe = ctx.tokenizer.encode(ctx.template.render([
+          ...normalizeMessages(req.messages), reply, { role: "user", content: "x" },
+        ], opts));
+        const probeTrimmed = probe[0] === probe[1] && probe[0] === ctx.tokenizer.bosTokenId ? probe.slice(1) : probe;
+        let common = 0;
+        const n = Math.min(i, probeTrimmed.length);
+        while (common < n && current[common] === probeTrimmed[common]) common++;
+        i = common;
+      }
+      // Generated-history provenance may use a different BPE segmentation
+      // inside old turns. Measure the primer canonically, then subtract it
+      // from the actual request IDs that inference will receive.
+      stableLen = Math.max(0, trimmed.length - (current.length - i));
       // Memoize the mode's primer length (sanity-capped: a "primer" longer
       // than 64 tokens means the probe diverged for content reasons —
       // don't generalize that).
@@ -273,10 +290,10 @@ export function createRequestPrep(input: {
    *  when the feature is off or the request's shape refuses it.
    *
    *  Refusals owned here (body-level); the ones only ChatStage can see —
-   *  a compiled grammar, media prompts, a mounted draft model, a quantized KV
-   *  scheme — are applied there. Continuous (batch) placement needs no
-   *  refusal: generate() is the only fill site, so a batched request simply
-   *  does not fill.
+   *  a compiled grammar, media prompts and a mounted draft model — are applied
+   *  there. The execution binding owns cache-format and method support. Strict
+   *  shared execution keeps the sampler's absolute token positions, including
+   *  known tokens; an explicit seed therefore remains replayable.
    *
    *  MLX_BUN_FILL=echo additionally arms the echo index (K3c, Lab tier): the
    *  strict rows keep policy "assert", and copied spans ride the same
@@ -289,9 +306,7 @@ export function createRequestPrep(input: {
     const mode = resolveFillMode();
     if (mode === "off") return null;
     if (!tools?.length) return null;
-    // A user-fixed seed means reproducibility: the sampler's step index would
-    // skip the injected positions, so an identical request could not replay.
-    if (req.seed !== undefined) return null;
+    // Seed support is selected by the execution binding after placement.
     // Injected tokens are never sampled — they have no logprob row.
     if (req.logprobs === true) return null;
     if (typeof req.top_logprobs === "number" && req.top_logprobs > 0) return null;

package/src/serve/text-completion-stage.ts

@@ -16,7 +16,7 @@ export class TextCompletionStage {
   constructor(
     private readonly ctx: ServingContext,
     private readonly prep: RequestPrep,
-    private readonly maxSafeContext: number,
+    private readonly contextLimit: number | null,
     /** `--max-tokens` / the GLM plan's reservation. No invented fallback:
      *  unset runs to EOS or the admitted context (with no template an EOS
      *  may never come — that run is bounded by admission, the real limit). */
@@ -81,7 +81,7 @@ export class TextCompletionStage {
         promptIds,
         options,
         requestedMaxTokens,
-        maxSafeContext: this.maxSafeContext,
+        contextLimit: this.contextLimit,
         stream: request.stream,
         wantLogprobs,
         topLogprobs,