mlx-bun 0.4.0 → 0.5.0

package/src/backends/mlx/autoregressive.ts

@@ -5,6 +5,9 @@ import type { RuntimeModel } from "../../model/factory";
 import { CompiledDecode } from "../../model/compiled-decode";
 import { runtimeConfig, type RuntimeConfig } from "../../runtime-config";
 import { bindMlxGraph } from "./graph";
+import type { PrefillPolicy } from "../../inference/prefill";
+import { resolveMlxPrefillPolicy } from "./prefill-policy";
+import type { KvSchemeOptions } from "../../kv-scheme";
 
 /** One declaration shared by planning and both native execution lanes. */
 export function legacyCompiledDecodeAvailable(model: RuntimeModel): boolean {
@@ -33,16 +36,33 @@ export interface MlxDecodeStep {
 /** Model-owned execution for committed tokens. Recheck the chunk limit after
  * each forward: native arithmetic can change at a cache-length boundary. */
 export interface MlxTokenAppend {
-  maxChunkSize(state: readonly Cache[]): number;
+  /** Affine formats whose committed append retains one-token arithmetic. */
+  readonly affineKvBits?: readonly number[];
+  readonly turboQuantFormats?: readonly { readonly kBits: number; readonly vBits: number }[];
+  /** Maximum positions per row at this cohort size; omitted rows means B=1. */
+  maxChunkSize(state: readonly Cache[], rows?: number): number;
   forwardHidden(ids: MlxArray, state: Cache[]): MlxArray | Promise<MlxArray>;
 }
 
+/** The model declares which codecs retain its committed-span arithmetic. */
+export function supportsCommittedAppendCache(
+  append: Pick<MlxTokenAppend, "affineKvBits" | "turboQuantFormats"> | null | undefined,
+  options: Pick<KvSchemeOptions, "kvBits" | "kvConfig" | "turboQuant">,
+): boolean {
+  if (options.turboQuant) return !!append?.turboQuantFormats?.some(format =>
+    format.kBits === options.turboQuant!.kBits && format.vBits === options.turboQuant!.vBits);
+  return options.kvConfig?.length
+    ? options.kvConfig.every(layer => append?.affineKvBits?.includes(layer.bits))
+    : !options.kvBits || !!append?.affineKvBits?.includes(options.kvBits);
+}
+
 /** One replaceable binding owns all graph-specific operations used by the AR
  * loop. The current method requires MLX tensors and the legacy Cache[] ABI.
  * Weights are borrowed for the binding's lifetime; generated caches are owned
  * by the run. Caller-provided caches and media remain borrowed. */
 export interface MlxAutoregressiveBinding {
   readonly runtime?: RuntimeConfig;
+  readonly prefillPolicy?: PrefillPolicy;
   readonly graph: AutoregressiveGraph<MlxArray, Cache[], MlxArray>;
   readonly eosTokenIds: readonly number[];
   readonly memory: MlxModelMemory;
@@ -63,6 +83,7 @@ export function bindLegacyAutoregressiveModel(model: RuntimeModel): MlxAutoregre
   const runtime = runtimeConfig();
   return {
     runtime,
+    prefillPolicy: resolveMlxPrefillPolicy(model.config, runtime),
     graph: bindMlxGraph<Cache[]>(model, {
       id: `legacy:${model.config.modelType}`, artifact: "legacy-resident-model",
       stateAbi: "legacy-cache-array-v1", // not a persistence identity

package/src/backends/mlx/batch-group.ts

@@ -1,3 +1,5 @@
+import { runMixedTokenIteration, type MlxForwardWork, type MlxPreparationWork } from "./mixed-iteration";
+import type { MixedTokenModel } from "../../model/token-groups";
 import type { OrdinaryContinuation } from "./continuation";
 import type { MlxRequestStatePolicy } from "./request-state-policy";
 import { AdmissionRejected } from "../../engine/admission";
@@ -63,6 +65,8 @@ import { MlxPrefillCohort, type PrefillState } from "./prefill-cohort";
 // optimization is a later refinement (batching-v2-plan item a).
 
 import { MlxArray } from "../../mlx/array";
+import type { PrefillPolicy } from "../../inference/prefill";
+import { resolveMlxPrefillPolicy } from "./prefill-policy";
 import * as ops from "../../mlx/ops";
 import { activeMemory, cacheMemory, clearCache, Dtype, peakMemory } from "../../mlx/ffi";
 import { runtimeConfig, withRuntimeConfig, type RuntimeConfig } from "../../runtime-config";
@@ -96,6 +100,8 @@ import { SSMCache } from "../../model/qwen3-delta";
 import type { RuntimeModel } from "../../model/factory";
 import type { GrammarController } from "../../grammar";
 import { independentGreedySampling } from "../../sampler";
+import { ExecutionTasks } from "../../engine/execution-tasks";
+import { CancellationSource, GenerationCancelled } from "../../engine/cancellation";
 import { acquireModelWiredLimit } from "../../generate";
 import { batchRowKvBytes } from "../../serve/kv-budget";
 import type { PromptResponseTrace } from "../../serve/prompt-response-trace";
@@ -138,8 +144,9 @@ export interface MlxGroupPreparation {
   readonly rows: readonly Row[];
   readonly tokenWeight?: number;
   readonly canAdmit?: boolean;
+  readonly supportsMixedWork?: boolean;
   admit?(row: Row): void;
-  advance(): Promise<boolean>;
+  advance(work?: MlxPreparationWork): Promise<boolean>;
   dispose(): void;
 }
 export interface MlxGroupMethodHost {
@@ -153,8 +160,10 @@ export interface MlxGroupMethodHost {
   finish(row: Row, reason: "stop" | "length"): void;
 }
 export interface MlxGroupedMethod {
+  /** Maximum target tokens in the next iteration, including candidates. */
+  readonly runningTokens?: number;
   prepare(row: Row): MlxGroupPreparation;
-  advance(): Promise<void>;
+  advance(work?: MlxForwardWork): Promise<void>;
   filterRows(keep: readonly number[], discard: boolean): void;
   dispose(): void;
 }
@@ -171,6 +180,9 @@ export type BatchRequest = BatchRequestFields & (
 );
 
 interface BatchRequestFields {
+  /** Borrowed prepared input; its owner retains it through preparation. */
+  promptInput?: import("./prompt-input").MlxPromptInput;
+  cacheSessionId?: string;
   /** Optional ordinary continuation policy; owns persistence and sampler recovery. */
   continuation?: OrdinaryContinuation;
   statePolicy?: MlxRequestStatePolicy;
@@ -225,6 +237,7 @@ interface BatchRequestFields {
 
 export interface BatchStats {
   spec?: import("../../generate").GenerateStats["spec"];
+  fill?: import("../../generate").GenerateStats["fill"];
   promptTokens: number;
   generatedTokens: number;
   /** Prompt tokens served from the prompt cache (Phase 3.2): a joiner's solo
@@ -239,6 +252,8 @@ export interface BatchStats {
   decodeMs: number;
 }
 
+let nextMixedWorkId = 0;
+
 /** The slice of PromptCache the scheduler drives (structural — the server's
  *  PromptCache satisfies it). take() on admission (any joiner: a restored
  *  prefix + suffix prefill is byte-safe whether the row later merges or
@@ -251,6 +266,7 @@ export type RowPromptCache = import("./checkpoint-state").MlxPrefixCache;
 
 export interface Row {
   spec?: import("../../generate").GenerateStats["spec"];
+  fill?: import("../../generate").GenerateStats["fill"];
   req: BatchRequest;
   resolve: (s: BatchStats) => void;
   reject: (e: unknown) => void;
@@ -342,6 +358,7 @@ type LayerInner =
 type Row1 = { keys: MlxArray; values: MlxArray };
 
 export class MlxBatchExecutionGroup {
+  readonly #tasks = new ExecutionTasks();
   readonly #runtime: RuntimeConfig;
   readonly #noPipeline: boolean;
   readonly #stepTrace: boolean;
@@ -363,7 +380,9 @@ export class MlxBatchExecutionGroup {
   #pendingReal: boolean[] | null = null;
   #method: MlxGroupedMethod | undefined;
   #methodKey: string | undefined;
+  #decodeStateKey: string | undefined;
   #steps = 0; // decode-step counter (clearCache cadence)
+  #cacheMaintenanceSteps = 0;
   #looping = false;
   #closed = false;
   #driver: Promise<void> | null = null;
@@ -372,6 +391,7 @@ export class MlxBatchExecutionGroup {
   readonly #lock: ExclusiveLock | undefined;
   readonly #admissionHeld: (() => boolean) | undefined;
   readonly #prefillChunkSize: number;
+  readonly #prefillPolicy: PrefillPolicy;
   readonly #prefillBatchTokenLimit: number;
   readonly #prefillTailSplit: boolean;
   readonly #kvBudgetBytes: number | undefined;
@@ -409,8 +429,8 @@ export class MlxBatchExecutionGroup {
     this.#maxBatch = Math.max(1, Math.floor(opts.maxBatch));
     this.#lock = opts.lock;
     this.#admissionHeld = opts.admissionHeld;
-    this.#prefillChunkSize = Math.max(1, Math.floor(opts.prefillChunkSize ??
-      this.#runtime.number("MLX_BUN_RD_PREFILL_CHUNK", 2048)));
+    this.#prefillPolicy = resolveMlxPrefillPolicy(model.config, this.#runtime, opts.prefillChunkSize);
+    this.#prefillChunkSize = this.#prefillPolicy.chunkSize(0);
     this.#prefillBatchTokenLimit = opts.prefillBatchTokenLimit ?? 2048;
     this.#prefillTailSplit = this.#runtime.flag("MLX_BUN_PREFILL_TAIL_SPLIT", true);
     this.#kvBudgetBytes = opts.kvBudgetBytes;
@@ -535,6 +555,7 @@ export class MlxBatchExecutionGroup {
     if (this.#closed) return Promise.reject(new Error("scheduler closed"));
     if (req.signal?.aborted) return Promise.reject(req.signal.reason);
     if (this.#pending.length >= this.#maxQueued) return Promise.reject(new AdmissionRejected());
+    req = { ...req, prefillChunkSize: req.prefillChunkSize ?? this.#prefillPolicy.chunkSize(req.promptIds.length) };
     return new Promise<BatchStats>((resolve, reject) => {
       let abortListener: (() => void) | null = null;
       const cleanup = () => {
@@ -562,6 +583,21 @@ export class MlxBatchExecutionGroup {
     this.#ensureLoop();
   }
 
+  /** Read-only native preparation borrows a boundary of the live execution
+   * group. Model mutation still requires the gateway's exclusive drain. */
+  runPreparation<T>(work: () => Promise<T>, signal?: AbortSignal): Promise<T> {
+    if (this.#closed) return Promise.reject(new Error("scheduler closed"));
+    if (signal?.aborted) return Promise.reject(signal.reason);
+    const cancellation = signal ? new CancellationSource() : undefined;
+    const abort = () => cancellation?.cancel("requested");
+    signal?.addEventListener("abort", abort, { once: true });
+    const result = this.#tasks.enqueue(work, cancellation).catch(error => {
+      throw error instanceof GenerationCancelled && signal?.aborted ? signal.reason : error;
+    }).finally(() => signal?.removeEventListener("abort", abort));
+    this.#ensureLoop();
+    return result;
+  }
+
   #ensureLoop(): void {
     if (this.#wake) { this.#wake(); return; }
     if (this.#looping || this.#closed) return;
@@ -580,6 +616,8 @@ export class MlxBatchExecutionGroup {
   async #drive(): Promise<void> {
     const scheduler = this;
     const group: ExecutionGroup = {
+      get pendingTasks() { return scheduler.#tasks.pending; },
+      advanceTask: () => this.#tasks.advance(),
       get active() { return scheduler.#running.length; },
       get queued() { return scheduler.#pending.length; },
       get preparing() { return scheduler.#prefill !== null; },
@@ -606,8 +644,46 @@ export class MlxBatchExecutionGroup {
       canBurst: () => this.#contextCompatible(this.#pending[0]!) &&
         (this.#kvBudgetBytes === undefined ||
           this.projectedKvBytes + this.#rowKvBytes(this.#pending[0]!) <= this.#kvBudgetBytes),
-      advancePreparation: () => this.#advancePreparation(),
-      advance: () => this.#method ? this.#method.advance() : this.#step(),
+      get mixedPreparation() {
+        return scheduler.#decodeStateKey === undefined && scheduler.#prefill?.supportsMixedWork !== false && scheduler.#runtime.flag("MLX_BUN_MIXED_PREFILL", false) &&
+          (!scheduler.#method || scheduler.#method.runningTokens !== undefined) &&
+          typeof (scheduler.model as RuntimeModel & Partial<MixedTokenModel>).forwardHiddenMixed === "function"
+          ? { runningTokens: scheduler.#method?.runningTokens ?? scheduler.#running.length,
+            minimumPreparationTokens: scheduler.#prefill?.rows.length ?? 0 }
+          : undefined;
+      },
+      get maxIterationTokens() { return scheduler.#runtime.number("MLX_BUN_MIXED_TOKEN_BUDGET", 256); },
+      advanceMixed: async tokenBudget => {
+        this.#promptCache?.reclaim?.();
+        const advanced = await runMixedTokenIteration({
+          prepare: forward => this.#advancePreparation({ forward,
+            maxTokens: tokenBudget - (this.#method?.runningTokens ?? this.#running.length) }),
+          decode: forward => this.#method ? this.#method.advance(forward) : this.#step(forward),
+          forward: async (ids, cache, options) => options?.captureLayer
+            ? (this.model as RuntimeModel & MixedTokenModel).forwardHiddenMixed([{ ids, cache, ...options }])[0]!
+            : this.#forwardHidden(ids, cache),
+          mixed: groups => {
+            const workId = `mixed:${++nextMixedWorkId}`;
+            const tokens = groups.map(group => group.ids.shape[0]! * group.ids.shape[1]!);
+            const closes = [...this.#running, ...(this.#prefill?.rows ?? [])].map(row =>
+              row.req.trace?.begin("engine.mixed_forward", { workId,
+                decodeTokens: tokens[0]!, prefillTokens: tokens[1]!, packedTokens: tokens[0]! + tokens[1]! }));
+            try { return (this.model as RuntimeModel & MixedTokenModel).forwardHiddenMixed(groups); }
+            finally { for (const close of closes) close?.(); }
+          },
+        });
+        if (!advanced && this.#running.length) {
+          if (this.#method) await this.#method.advance(); else await this.#step();
+        }
+      },
+      advancePreparation: () => {
+        this.#promptCache?.reclaim?.();
+        return this.#advancePreparation();
+      },
+      advance: () => {
+        if (this.#cacheMaintenanceSteps++ % 256 === 0) this.#promptCache?.reclaim?.();
+        return this.#method ? this.#method.advance() : this.#step();
+      },
       failActive: (error) => {
         for (const row of this.#running) row.reject(error);
         this.#applyFilter([], true); // failed state never enters the prefix store
@@ -629,12 +705,12 @@ export class MlxBatchExecutionGroup {
       await driveExecutionGroup(group, {
         now: () => performance.now(),
         yield: () => new Promise<void>((resolve) => setImmediate(resolve)),
-      }, this.#runtime.value("MLX_BUN_EARLY_FIRST_TOKEN") === "1");
+      }, this.#runtime.flag("MLX_BUN_EARLY_FIRST_TOKEN", true));
     } finally { this.#looping = false; }
   }
 
   #contextCompatible(row: Row): boolean {
-    return (!this.#running.length && !this.#prefill) || ((row.req.context?.key ?? "") === this.#contextKey && row.req.method?.key === this.#methodKey && row.req.statePolicy?.key === this.#statePolicy?.key);
+    return (!this.#running.length && !this.#prefill) || ((row.req.context?.key ?? "") === this.#contextKey && row.req.method?.key === this.#methodKey && row.req.statePolicy?.key === this.#statePolicy?.key && row.req.promptInput?.decodeState?.key === this.#decodeStateKey);
   }
 
   #leaveContext(): void {
@@ -649,6 +725,7 @@ export class MlxBatchExecutionGroup {
     const row = this.#pending.shift()!;
     try {
       this.#statePolicy = row.req.statePolicy;
+      this.#decodeStateKey = row.req.promptInput?.decodeState?.key;
       row.cacheNamespace = typeof row.req.cacheNamespace === "function"
         ? row.req.cacheNamespace() : row.req.cacheNamespace;
       const key = row.req.context?.key ?? "";
@@ -705,11 +782,11 @@ export class MlxBatchExecutionGroup {
     for (const row of rows) row.reject(error);
   }
 
-  async #advancePreparation(): Promise<void> {
+  async #advancePreparation(work?: MlxPreparationWork): Promise<void> {
     const p = this.#prefill!, rows = p.rows;
     this.#preparationPublishedOutput = false;
     try {
-      if (await p.advance()) this.#prefill = null;
+      if (await p.advance(work)) this.#prefill = null;
     } catch (error) {
       this.#prefill = null;
       let failure = error;
@@ -725,6 +802,7 @@ export class MlxBatchExecutionGroup {
   }
 
   #failAll(error: unknown): void {
+    this.#tasks.rejectAll(error);
     const p = this.#prefill;
     const rows = new Set([...this.#running, ...this.#pending, ...(p?.rows ?? [])]);
     const retain = this.#adoptedRetain;
@@ -809,7 +887,7 @@ export class MlxBatchExecutionGroup {
       row.cachedTokens = Math.min(saved.cacheTokens.length, row.promptTokens);
       const stop = await this.#emit(row, saved.pendingToken);
       if (stop !== "continue") {
-        this.#putOrDispose(p.solo, saved.cacheTokens, p.retain, row.cacheNamespace);
+        this.#putOrDispose(p.solo, saved.cacheTokens, p.retain, row.cacheNamespace, row.req.cacheSessionId);
         p.solo = [];
         this.#finish(row, stop);
         return;
@@ -831,7 +909,7 @@ export class MlxBatchExecutionGroup {
     // later pipeline step. Grammar and explicit eager-output modes retain
     // their read-before-advance order.
     if (!this.#inners && !p.row.req.grammar && !this.#noPipeline &&
-        this.#runtime.value("MLX_BUN_EARLY_FIRST_TOKEN") !== "1" && !forceAttribution) {
+        !this.#runtime.flag("MLX_BUN_EARLY_FIRST_TOKEN", true) && !forceAttribution) {
       let owned: MlxArray | null = sampled;
       try {
         ops.asyncEvalAll([sampled]);
@@ -861,7 +939,7 @@ export class MlxBatchExecutionGroup {
         const stop = await this.#emit(p.row, tok);
         // Token 0 was sampled but never fed — the caches cover exactly the
         // prompt, a clean prompt-only entry (put-or-dispose).
-        this.#putOrDispose(p.solo, p.row.req.promptIds, p.retain, p.row.cacheNamespace);
+        this.#putOrDispose(p.solo, p.row.req.promptIds, p.retain, p.row.cacheNamespace, p.row.req.cacheSessionId);
         this.#finish(p.row, stop === "continue" ? "stop" : stop);
         return;
       }
@@ -869,7 +947,7 @@ export class MlxBatchExecutionGroup {
 
     const stop = await this.#emit(p.row, tok);
     if (stop !== "continue") {
-      this.#putOrDispose(p.solo, p.row.req.promptIds, p.retain, p.row.cacheNamespace);
+      this.#putOrDispose(p.solo, p.row.req.promptIds, p.retain, p.row.cacheNamespace, p.row.req.cacheSessionId);
       this.#finish(p.row, stop);
       return;
     }
@@ -1147,7 +1225,7 @@ export class MlxBatchExecutionGroup {
    *  Rows that finish get one extra harmless KV write from the already-built
    *  step; filter drops the row (mlx-lm behaves identically). Length-finished
    *  rows are known in advance and are NOT sampled (placeholder slot). */
-  async #step(): Promise<void> {
+  async #step(forward?: MlxForwardWork): Promise<void> {
     if (this.#stepTrace) {
       const now = performance.now();
       if (STEP_T.lastEnd) STEP_T.gap += now - STEP_T.lastEnd;
@@ -1156,6 +1234,11 @@ export class MlxBatchExecutionGroup {
     const rows = this.#running;
     const B = rows.length;
     const inners = this.#inners!;
+    const decodeState = rows[0]?.req.promptInput?.decodeState;
+    if (decodeState) {
+      const states = rows.map(row => row.req.promptInput!.decodeState!);
+      forward = async (ids, caches) => decodeState.forward(ids, caches, states);
+    }
 
     // A row is live if it still needs tokens sampled; a row whose pending
     // unread token is its last (sampled == maxTokens) only awaits emission.
@@ -1175,7 +1258,7 @@ export class MlxBatchExecutionGroup {
     const hasLiveGrammar = anyLive && rows.some(
       (r) => r.req.grammar && !r.req.grammar.isTerminated,
     );
-    if (hasLiveGrammar) return this.#stepGrammar();
+    if (hasLiveGrammar) return this.#stepGrammar(forward);
 
     let nextToks: MlxArray | null = null;
     let nextReal: boolean[] | null = null;
@@ -1209,7 +1292,7 @@ export class MlxBatchExecutionGroup {
       let lg: MlxArray | null = null;
       let evalWith: MlxArray[] = [];
       if (
-        this.#compiled && B === 1 && unpadded && this.#running[0]!.req.compiledDecode !== false &&
+        !forward && this.#compiled && B === 1 && unpadded && this.#running[0]!.req.compiledDecode !== false &&
         (!this.#pendingToks || this.#pendingToks.dtype === Dtype.uint32) &&
         // A filtered-to-one BATCHED rot-quant cache subclasses the serial
         // class (so supports() passes) but carries batched ring state —
@@ -1251,7 +1334,7 @@ export class MlxBatchExecutionGroup {
           const ids = this.#pendingToks
             ? ops.reshape(this.#pendingToks, [B, 1]) // feed the unread tokens
             : ops.fromInt32(rows.map((r) => r.current), [B, 1]); // pipeline cold
-          const h = await this.#forwardHidden(ids, fwd);
+          const h = await (forward ? forward(ids, fwd) : this.#forwardHidden(ids, fwd));
           ids.dispose();
           lg = this.model.logitsFromHidden(h); // [B,1,V]
           h.dispose();
@@ -1407,7 +1490,7 @@ export class MlxBatchExecutionGroup {
    *  finish("stop") in #emitRows. On a cold start (first step after prefill)
    *  there is no pending array to read; the prefill already accepted token 0
    *  and fired the fill, so we just await ready() + sample. */
-  async #stepGrammar(): Promise<void> {
+  async #stepGrammar(forward?: MlxForwardWork): Promise<void> {
     const rows = this.#running;
     const B = rows.length;
     const inners = this.#inners!;
@@ -1454,7 +1537,7 @@ export class MlxBatchExecutionGroup {
         const ids = prev
           ? ops.reshape(prev, [B, 1]) // feed the unread tokens (device array)
           : ops.fromInt32(rows.map((r) => r.current), [B, 1]); // pipeline cold
-        const h = await this.#forwardHidden(ids, fwd);
+        const h = await (forward ? forward(ids, fwd) : this.#forwardHidden(ids, fwd));
         ids.dispose();
         const lg = this.model.logitsFromHidden(h); // [B,1,V]
         h.dispose();
@@ -1588,6 +1671,7 @@ export class MlxBatchExecutionGroup {
     const first = row.firstTokenAt || now;
     row.resolve({
       ...(row.spec ? { spec: row.spec } : {}),
+      ...(row.fill ? { fill: row.fill } : {}),
       promptTokens: row.promptTokens,
       generatedTokens: row.generated,
       cachedTokens: row.cachedTokens,
@@ -1602,7 +1686,7 @@ export class MlxBatchExecutionGroup {
    *  offset lines up (defensive — a mismatch means the accounting is wrong
    *  and the entry would corrupt future hits), else dispose. `retain` rides
    *  along per the PromptCacheEntry contract (runs after dispose). */
-  #putOrDispose(caches: Cache[], tokens: number[], retain?: () => void, namespace = ""): void {
+  #putOrDispose(caches: Cache[], tokens: number[], retain?: () => void, namespace = "", sessionId?: string): void {
     if (this.#promptCache) {
       // A never-merged row can still own a row layout. Persistence receives
       // its existing serial representation, just like merged-row retirement.
@@ -1623,7 +1707,7 @@ export class MlxBatchExecutionGroup {
         (c) => typeof (c as { offset?: unknown }).offset === "number",
       ) as { offset: number } | undefined;
       if (withOff && withOff.offset === tokens.length) {
-        this.#promptCache.put(tokens, caches, namespace, retain);
+        this.#promptCache.put(tokens, caches, namespace, retain, undefined, sessionId);
         return;
       }
     }
@@ -1634,13 +1718,13 @@ export class MlxBatchExecutionGroup {
   /** Extract a finishing MERGED row's KV into fresh serial caches and put()
    *  them keyed by [promptIds + fed] — the batch-lane mirror of mlx-lm
    *  server.py:872 (extract_cache → prompt_cache.insert_cache). Gates:
-   *  promptTokens >= 256 (the boundary-snapshot substantiality gate,
-   *  server.ts) and an exact coverage key (!fedTainted). Refusal
+   *  promptTokens >= 256 or an explicit cache session, plus an exact
+   *  coverage key (!fedTainted). Refusal
    *  (#extractRowCaches null) just disposes-by-omission — the row's KV dies
    *  with the filter, exactly the pre-extraction behavior. */
   #extractAndPut(b: number, row: Row): void {
     if (!this.#promptCache || !row.merged || row.fedTainted) return;
-    if (row.promptTokens < 256) return;
+    if (row.promptTokens < 256 && !row.req.cacheSessionId) return;
     const tokens = [...row.req.promptIds, ...row.fed];
     const caches = this.#extractRowCaches(b, tokens.length);
     if (!caches) return;
@@ -1649,7 +1733,7 @@ export class MlxBatchExecutionGroup {
     // buffers free once the copies land, instead of being pinned by a lazy
     // graph inside an idle cache entry.
     withResource(leaseCacheStates(caches), state => ops.asyncEvalAll([...state]));
-    this.#putOrDispose(caches, tokens, undefined, row.cacheNamespace);
+    this.#putOrDispose(caches, tokens, undefined, row.cacheNamespace, row.req.cacheSessionId);
   }
 
   /** Row `b` of every layer as OWNED serial-class caches, or null when a
@@ -1717,7 +1801,7 @@ export class MlxBatchExecutionGroup {
           inners as Cache[],
           [...solo.req.promptIds, ...solo.fed],
           this.#adoptedRetain ?? undefined,
-          solo.cacheNamespace,
+          solo.cacheNamespace, solo.req.cacheSessionId,
         );
       } else {
         for (const c of inners) c.dispose();

package/src/backends/mlx/cache-layout.ts

@@ -8,15 +8,15 @@ import { PagedKvRows } from "../../lab/paged-kv/paged-kv-rows";
 
 /** Bind encoded state to its row layout. Execution owners only move rows. */
 export function ownedCacheLayout(cache: Cache): BatchableCache | undefined {
-  if (cache instanceof PagedKVCache) return new PagedKvRows(cache.capacityTokens, cache.blockSize);
+  if (cache instanceof PagedKVCache) return new PagedKvRows(cache.capacityTokens, cache.blockSize, cache.direct, cache.quantization);
   if (isBatchableCache(cache)) return cache.makeEmptyBatch();
-  if (cache instanceof TurboQuantKVCache) return new BatchedTurboQuantKVCache(cache.kBits, cache.vBits);
+  if (cache instanceof TurboQuantKVCache) return new BatchedTurboQuantKVCache(cache.kBits, cache.vBits, cache.fusedDecode);
   return undefined;
 }
 
 /** Prefill reuses decode's layouts; shape selection belongs to storage. */
 export function prefillCacheLayout(cache: Cache): BatchableCache {
-  if (cache instanceof PagedKVCache) return new PagedKvRows(cache.capacityTokens, cache.blockSize);
+  if (cache instanceof PagedKVCache) return new PagedKvRows(cache.capacityTokens, cache.blockSize, cache.direct, cache.quantization);
   if (isBatchableCache(cache)) return cache.makeEmptyBatch();
   if (cache instanceof RotatingKVCache) return new DelayedRotatingQuantizedKVCache(cache.maxSize, 64, 4, Infinity, () => {});
   if (cache instanceof RotatingQuantizedKVCache) return new RotatingAffineLayout(cache.maxSize, cache.groupSize, cache.bits);

package/src/backends/mlx/continuation-persistence.ts

@@ -19,13 +19,14 @@ export class ContinuationPersistence {
   #nextBarrier = 0;
   constructor(readonly store: ContinuationStore, options: {
     maxBytes: number;
-    /** All blocking tensor serialization runs through the owner's idle gate. */
-    runStep: <T>(step: () => T) => Promise<T>;
+    /** Optional owner-thread preparation hook; disk writes use the CPU worker. */
+    runStep?: <T>(step: () => T) => Promise<T>;
   }) {
+    const runStep = options.runStep ?? (async <T>(step: () => T): Promise<T> => step());
     this.#queue = new SpillQueue(options.maxBytes, cacheBytes, async item => {
       const job = this.#jobs.get(item.caches)!;
       if (job.cleanup) {
-        await options.runStep(() => store.removeGenerationCheckpoints(job.attempt.key));
+        await runStep(() => store.removeGenerationCheckpoints(job.attempt.key));
         return true;
       }
       const current = () => this.#attempts.get(job.attempt.key) === job.attempt && !job.attempt.completed;
@@ -33,7 +34,7 @@ export class ContinuationPersistence {
       const stored = await store.storeGenerationCheckpoint(item.tokens, item.caches, job.metadata!, options.runStep);
       // No newer queued write can run until this callback returns. An old
       // in-flight rename can therefore be removed without deleting new work.
-      if (!current()) await options.runStep(() => store.removeGenerationCheckpoints(job.attempt.key));
+      if (!current()) await runStep(() => store.removeGenerationCheckpoints(job.attempt.key));
       return stored;
     }, disposeResources, item => !this.#jobs.get(item.caches)?.cleanup);
   }

package/src/backends/mlx/encoder-cache.ts

@@ -0,0 +1,25 @@
+import type { ObjectCache } from "../../contracts/object-cache";
+import { cloneAttachments, disposeAttachments, type CheckpointAttachment } from "./checkpoint-state";
+import type { MlxArray } from "../../mlx/array";
+
+/** Encoder-owned tensor schema over the cache's exact-object interface. */
+export class EncoderCache {
+  constructor(private readonly objects: ObjectCache<CheckpointAttachment[]>,
+    readonly identity: string) {}
+
+  async take(key: string): Promise<MlxArray | null> {
+    const lease = await this.objects.take(`encoder-v1:${this.identity}:${key}`);
+    if (!lease) return null;
+    try {
+      const tensor = lease.value[0]!.tensors[0]!;
+      return tensor.slice(tensor.shape.map(() => 0), [...tensor.shape]);
+    } finally { lease.dispose(); }
+  }
+
+  /** Borrow the result; storage retains its own immutable native view. */
+  put(key: string, features: MlxArray): void {
+    const owned = cloneAttachments([{ schema: "encoder-features-v1", metadata: {}, tensors: [features] }]);
+    try { this.objects.put(`encoder-v1:${this.identity}:${key}`, owned); }
+    catch (error) { disposeAttachments(owned); throw error; }
+  }
+}

package/src/backends/mlx/fill-append.ts

@@ -6,8 +6,8 @@ import { leaseCacheStates } from "./state-views";
 
 /** Append an already accepted span without intermediate vocabulary heads.
  * Chunking changes the graph shape, so callers must qualify its numerics for
- * their model. A verify round must use chunkSize=0: recurrent rollback records
- * one forward per round. The complete append precedes emission/cancellation
+ * their model. Verification uses its separate round operation because
+ * recurrent rollback records one forward per round. Append precedes emission/cancellation
  * boundaries, preserving the cache-covered prefix on a consumer break.
  */
 export async function appendFillHidden(
@@ -17,15 +17,29 @@ export async function appendFillHidden(
   chunkSize: (cache: readonly Cache[]) => number,
 ): Promise<MlxArray> {
   const input = ops.fromInt32(ids, [1, ids.length]);
+  try { return await appendHiddenRows(forward, cache, input, chunkSize); }
+  finally { input.dispose(); }
+}
+
+/** Append equal-length committed token rows [B,L]. The caller owns the input;
+ * the returned hidden rows belong to the caller. Model policy receives B so
+ * a larger cohort cannot silently cross its qualified projection geometry. */
+export async function appendHiddenRows(
+  forward: (ids: MlxArray, cache: Cache[]) => MlxArray | Promise<MlxArray>,
+  cache: Cache[],
+  input: MlxArray,
+  chunkSize: (cache: readonly Cache[], rows: number) => number,
+): Promise<MlxArray> {
+  const [rows, length] = input.shape as [number, number];
   const hidden: MlxArray[] = [];
   try {
-    for (let start = 0; start < ids.length;) {
-      const limit = chunkSize(cache);
+    for (let start = 0; start < length;) {
+      const limit = chunkSize(cache, rows);
       if (!Number.isSafeInteger(limit) || limit < 1)
         throw new Error("Model append chunk limit must be a positive integer");
-      const end = Math.min(ids.length, start + limit);
-      if (start === 0 && end === ids.length) return await forward(input, cache);
-      const part = input.slice([0, start], [1, end]);
+      const end = Math.min(length, start + limit);
+      if (start === 0 && end === length) return await forward(input, cache);
+      const part = input.slice([0, start], [rows, end]);
       try {
         const h = await forward(part, cache);
         hidden.push(h);
@@ -40,6 +54,5 @@ export async function appendFillHidden(
     return ops.concatAxis(hidden, 1);
   } finally {
     for (const h of hidden) h.dispose();
-    input.dispose();
   }
 }