mlx-bun 0.3.0 → 0.5.0

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

@@ -4,7 +4,7 @@ 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";
+  type MultimodalTowers } from "../../vision/prompt";
 import { buildQwen3VLVisionPrompt } from "../../vision/qwen3vl-prompt";
 import type { Qwen3VLVisionTower } from "../../vision/qwen3vl-tower";
 import { normalizeMessages, type ChatRequestParams } from "../../serve/chat-request";
@@ -12,7 +12,8 @@ import { getAudioTower, getVisionTower, type ServerContext } from "../../serve/m
 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";
+import type { BuiltPrompt, PromptNativeWork } from "../../serve/model-binding";
+import { EncoderCache } from "./encoder-cache";
 
 export async function buildModelPrompt(
     ctx: ServerContext,
@@ -20,6 +21,8 @@ export async function buildModelPrompt(
     body: ChatRequestParams,
     tools: ChatRequestParams["tools"] | null,
     ownership: RequestOwnership,
+    nativeWork: PromptNativeWork = (work) => work(),
+    objects?: import("../../contracts/object-cache").ObjectCache<import("./checkpoint-state").CheckpointAttachment[]>,
   ): Promise<BuiltPrompt> {
     const toolList = tools ?? null;
     const partOf = (types: string[]) => body.messages.some(
@@ -53,19 +56,6 @@ export async function buildModelPrompt(
       // 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);
@@ -74,13 +64,30 @@ export async function buildModelPrompt(
       // 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,
-        {
+      const towers = await nativeWork(async () => {
+        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: MultimodalTowers<{ softTokens: number }>["vision"];
+        if (hasImages) {
+          const tower = getVisionTower(ctx);
+          if (!tower) throw new RequestError(400, "model has no vision sidecar");
+          visionSide = { tower, tokenIds: ctx.visionTokenIds,
+            cache: objects && tower.cacheIdentity ? new EncoderCache(objects, tower.cacheIdentity) : undefined };
+        }
+        return {
           ...(visionSide ? { vision: visionSide } : {}),
-          audio: { tower: audioTower, tokenIds: ctx.audioTokenIds },
-        },
-        ctx.tokenizer, ctx.template, messages, images, wavs, toolList,
+          audio: { tower: audioTower, tokenIds: ctx.audioTokenIds,
+            cache: objects ? new EncoderCache(objects, audioTower.cacheIdentity) : undefined },
+        };
+      });
+      const mp = await buildMultimodalPrompt(
+        ctx.model as Gemma4Model, towers, ctx.tokenizer, ctx.template,
+        messages, images, wavs, toolList, nativeWork,
       );
       // bidirMask is null whenever audio is present (§3.3 Q1: mixed
       // prompts run fully causal); the union mask does the per-layer
@@ -89,7 +96,7 @@ export async function buildModelPrompt(
         ...noMedia,
         promptIds: mp.ids,
         vision: {
-          embeddings: mp.embeddings,
+          embeddings: mp.embeddings, prefixIdentity: mp.prefixIdentity,
           ...(mp.bidirMask ? { imageMask: mp.bidirMask } : {}),
           multimodalMask: mp.multimodalMask,
         },
@@ -106,7 +113,7 @@ export async function buildModelPrompt(
         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 { pixels, softTokens } = await nativeWork(() => dm.visionTower!.preprocess(images[0]!));
       const diffusionPixels = ownership.own(pixels);
       const promptIds = spliceImageTokens(rawIds, [softTokens], {
         image: ctx.visionTokenIds.imageTokenId,
@@ -121,13 +128,16 @@ export async function buildModelPrompt(
       // 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 { tower, encoderCache } = await nativeWork(async () => {
+        const tower = getVisionTower(ctx) as unknown as Qwen3VLVisionTower | null;
+        if (!tower) throw new RequestError(400, "model has no vision sidecar");
+        return { tower, encoderCache: objects ? new EncoderCache(objects, tower.cacheIdentity) : undefined };
+      });
       const vp = await buildQwen3VLVisionPrompt(
-        ctx.model, tower, ctx.tokenizer, ctx.template, messages, images,
+        ctx.model as Qwen35Model, 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,
@@ -135,23 +145,27 @@ export async function buildModelPrompt(
           visionEndId: (ctx.model.config.raw.vision_end_token_id as number) ?? 248054,
         },
         prep.templateOptionsFor(body, toolList),
-        videos,
+        videos, nativeWork, encoderCache,
       );
-      return { ...noMedia, promptIds: vp.ids, vision: { embeddings: vp.embeddings, mrope: vp.mrope } };
+      return { ...noMedia, promptIds: vp.ids, vision: { embeddings: vp.embeddings, mrope: vp.mrope, prefixIdentity: vp.prefixIdentity } };
     }
     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 { tower, encoderCache } = await nativeWork(async () => {
+        const tower = getVisionTower(ctx);
+        if (!tower) throw new RequestError(400, "model has no vision sidecar");
+        return { tower, encoderCache: objects && tower.cacheIdentity
+          ? new EncoderCache(objects, tower.cacheIdentity) : undefined };
+      });
       const vp = await buildVisionPrompt(
         ctx.model as Gemma4Model, tower, ctx.tokenizer, ctx.template,
-        messages, images, ctx.visionTokenIds, toolList,
+        messages, images, ctx.visionTokenIds, toolList, nativeWork, encoderCache,
       );
-      return { ...noMedia, promptIds: vp.ids, vision: { embeddings: vp.embeddings, imageMask: vp.imageMask } };
+      return { ...noMedia, promptIds: vp.ids, vision: { embeddings: vp.embeddings, imageMask: vp.imageMask, prefixIdentity: vp.prefixIdentity } };
     }
     const text = prep.promptIdsFor(body, toolList);
     return {

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

@@ -8,8 +8,8 @@ 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));
+  return context.serving?.buildPrompt ?? ((body, tools, ownership, prep, nativeWork, objects) =>
+    buildModelPrompt(context as ServerContext, prep, body, tools, ownership, nativeWork, objects));
 }
 
 /** Existing model classes enter the interface once here. New implementations
@@ -25,10 +25,11 @@ export function modelServingBinding(context: ServingContext): ModelServingBindin
   const embedding = isEmbeddingModel(model);
   return {
     stateCompatibility: `mlx-${MLX_VERSION}-${deviceArchitecture()}`,
-    gateway: bindMlxGateway(model),
+    gateway: bindMlxGateway(model, ctx.draft ?? undefined),
     createSerial: (services) => createMlxSerialExecutor(serial, services),
-    buildPrompt: (body, tools, ownership, prep) => buildModelPrompt(ctx, prep, body, tools, ownership),
+    buildPrompt: (body, tools, ownership, prep, nativeWork, objects) => buildModelPrompt(ctx, prep, body, tools, ownership, nativeWork, objects),
     restore: (store, entry) => store.restore(entry, model),
+    restoreAsync: (store, entry) => store.restoreAsync(entry, model),
     async signal(ids, count, minimum) {
       const caches = model.makeCache();
       try {

package/src/backends/mlx/paged-cache-codec.ts

@@ -0,0 +1,27 @@
+import type { CacheCodec } from "../../kv-store";
+import { BlockPool, PagedKVCache } from "../../lab/paged-kv/paged-kv";
+
+/** Page layout owns its persistence codec, independently of SSD transport. */
+export const pagedCacheCodec: CacheCodec = {
+    matches: cache => cache instanceof PagedKVCache,
+    snapshot(cache, context) {
+      const c = cache as PagedKVCache, p = c.pool!;
+      for (const array of p.arrays()) context.push(array, false);
+      return { kind: "paged", offset: c.offset, tensors: context.slots,
+        ...(p.quantization ? p.quantization : {}),
+        paged: { capacityTokens: c.capacityTokens, blockSize: c.blockSize, blockTable: [...c.blockTable],
+          numBlocks: p.numBlocks, headDim: p.headDim, vHeadDim: p.vHeadDim, dtype: p.dtype, direct: c.direct } };
+    },
+    clone: cache => (cache as PagedKVCache).clone(),
+    load(entry, context) {
+      const layout = entry.paged!;
+      const quantization = entry.bits ? { bits: entry.bits as 4 | 8, groupSize: entry.groupSize! } : undefined;
+      const cache = new PagedKVCache(layout.capacityTokens, layout.blockSize, layout.direct, quantization);
+      const pool = new BlockPool({ ...layout, numKvHeads: entry.tensors[0]!.shape[1]!, quantization });
+      try { pool.restoreArrays(entry.tensors.map(context.arr), layout.blockTable); }
+      catch (error) { pool.dispose(); throw error; }
+      cache.pool = pool; cache.offset = entry.offset; cache.blockTable = [...layout.blockTable];
+      return cache;
+    },
+    headerTrimmable: () => true,
+  };

package/src/backends/mlx/prefill-cohort.ts

@@ -0,0 +1,100 @@
+import type { Cache } from "../../model/gemma4-base";
+import type { RuntimeModel } from "../../model/factory";
+import type { MlxArray } from "../../mlx/array";
+import { nextPrefillStep } from "../../inference/prefill";
+import type { CacheCodecProvider } from "../../kv-store";
+import { disposeResources, cleanupFailure } from "../../engine/resources";
+import type { Row, RowPromptCache } from "./batch-group";
+import type { KvMaintenance } from "./kv-maintenance";
+import { MlxPrefillRows, type MlxPrefillState } from "./prefill-rows";
+
+export interface PrefillState extends MlxPrefillState {
+  snapAt: number | null;
+  continuation?: import("./continuation").OrdinaryContinuationState;
+  closePrefill?: () => void;
+}
+
+interface PrefillHost {
+  model: RuntimeModel;
+  chunkSize: number;
+  tailSplit: boolean;
+  promptCache?: RowPromptCache;
+  stateCodecs?: CacheCodecProvider;
+  maintain?: KvMaintenance;
+  forward(ids: MlxArray, caches: Cache[]): Promise<MlxArray>;
+  /** Target projection preserves the forward batch geometry before sampling. */
+  project(hidden: MlxArray, caches: Cache[], completed: readonly PrefillState[]): MlxArray;
+  /** Borrows [1,1,V] logits; takes ownership of the completed request caches. */
+  complete(state: PrefillState, logits: MlxArray): Promise<void>;
+  resume?(state: PrefillState): Promise<void>;
+  reject(row: Row, error: unknown): void;
+}
+
+/** Ordinary sampling/cache policy over the shared target preparation driver. */
+export class MlxPrefillCohort extends MlxPrefillRows<PrefillState> {
+  override get canAdmit(): boolean {
+    return super.canAdmit && this.supportsMixedWork;
+  }
+  get supportsMixedWork(): boolean { return !this.rows.some(row => row.req.promptInput); }
+  constructor(readonly host: PrefillHost) {
+    super({
+      stateCodecs: host.stateCodecs, maintain: host.maintain,
+      open(row) {
+        let owned: { caches: Cache[]; retain?: () => void } | undefined;
+        let closeCache: (() => void) | undefined;
+        let closePrefill: (() => void) | undefined;
+        try {
+          closeCache = row.req.trace?.begin("cache.lookup_restore", { mechanism: "continuous" });
+          const continuation = row.req.continuation?.restore(row.cacheNamespace ?? "");
+          if (continuation) {
+            owned = { caches: continuation.caches };
+            return { row, solo: continuation.caches, pos: row.req.promptIds.length,
+              snapAt: null, continuation };
+          }
+          const hit = host.promptCache?.take(row.req.promptIds, row.cacheNamespace, row.req.cacheSessionId) ?? null;
+          if (hit) owned = hit;
+          closeCache?.(); closeCache = undefined;
+          if (hit) row.cachedTokens = hit.tokens.length;
+          const length = row.req.promptIds.length;
+          const boundary = Math.min(row.req.snapshotAt ?? length, length - 1);
+          const snapAt = host.promptCache && boundary >= 256 && boundary > (hit?.tokens.length ?? 0) ? boundary : null;
+          closePrefill = row.req.trace?.begin("prefill.total", {
+            mechanism: "continuous", promptTokens: row.promptTokens, cachedTokens: row.cachedTokens,
+          });
+          const closeSetup = row.req.trace?.begin("prefill.batch_setup", { mechanism: "continuous" });
+          owned ??= { caches: row.req.statePolicy?.create() ?? host.model.makeCache() }; closeSetup?.();
+          if (row.req.promptInput) host.maintain?.(owned.caches);
+          return { row, solo: owned.caches, pos: hit?.tokens.length ?? 0, retain: owned.retain, snapAt, closePrefill };
+        } catch (error) {
+          return cleanupFailure(error, () => disposeResources([...(owned?.caches ?? []),
+            { dispose: () => owned?.retain?.() }, { dispose: () => closePrefill?.() }]));
+        } finally { closeCache?.(); }
+      },
+      ready: state => !!state.continuation,
+      plan(state) {
+        if (state.row.req.promptInput) return { kind: "final", start: state.pos,
+          end: state.row.promptTokens, snapshot: false, batchYield: false, atomic: true };
+        return nextPrefillStep({ length: state.row.promptTokens, position: state.pos,
+        chunkSize: state.row.req.prefillChunkSize ?? host.chunkSize,
+        tailSplit: host.tailSplit, snapshotAt: state.snapAt }); },
+      forward: (ids, caches, states, work) => states[0]?.row.req.promptInput
+        ? Promise.resolve(states[0].row.req.promptInput.forward(ids, caches, states[0].pos))
+        : work ? work(ids, caches) : host.forward(ids, caches),
+      project: host.project.bind(host),
+      complete: (state, logits) => state.continuation ? host.resume!(state) : host.complete(state, logits!), reject: host.reject.bind(host),
+      checkpoint(state, capture) {
+        let snapshot: Cache[] | undefined;
+        try {
+          snapshot = capture();
+          host.promptCache!.put(state.row.req.promptIds.slice(0, state.pos), snapshot, state.row.cacheNamespace, undefined, undefined, state.row.req.cacheSessionId);
+          snapshot = undefined;
+        } catch (error) {
+          if (snapshot) disposeResources(snapshot);
+          console.warn(`batch-lane boundary snapshot skipped: ${(error as Error).message}`);
+        }
+        state.snapAt = null;
+      },
+      close(state) { const close = state.closePrefill; state.closePrefill = undefined; close?.(); },
+    });
+  }
+}

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

@@ -0,0 +1,26 @@
+import type { ModelConfig } from "../../config";
+import type { PrefillPolicy } from "../../inference/prefill";
+import { sdpaFallbackBytes } from "../../kv-scheme";
+import type { RuntimeConfig } from "../../runtime-config";
+
+/** Retain large chunks for short prompts, but bound the materialized attention
+ * score workspace of recurrent-attention models at long context. This changes
+ * work size, never request admission. Explicit chunk settings remain exact. */
+export function resolveMlxPrefillPolicy(config: ModelConfig, runtime: RuntimeConfig,
+  override?: number): PrefillPolicy {
+  const explicit = override ?? (runtime.value("MLX_BUN_RD_PREFILL_CHUNK") === undefined
+    ? undefined : runtime.number("MLX_BUN_RD_PREFILL_CHUNK", 2048));
+  if (explicit !== undefined) {
+    const chunk = Math.max(1, Math.floor(explicit));
+    return { chunkSize: () => chunk };
+  }
+  if (!config.text?.layerTypes?.includes("linear_attention")) return { chunkSize: () => 2048 };
+  // Capture geometry now, independently of later configuration mutation.
+  const geometry = { ...config, text: { ...config.text, layerTypes: [...config.text.layerTypes] } };
+  const workspaceBytes = 1024 ** 3;
+  return { chunkSize(promptTokens) {
+    let chunk = 2048;
+    while (chunk > 8 && sdpaFallbackBytes(geometry, chunk, promptTokens) > workspaceBytes) chunk /= 2;
+    return chunk;
+  } };
+}

package/src/backends/mlx/prefill-rows.ts

@@ -0,0 +1,244 @@
+import type { BatchableCache, Cache } from "../../model/gemma4-base";
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import { clearCache } from "../../mlx/ffi";
+import type { PrefillStep } from "../../inference/prefill";
+import { cloneKvCaches, type CacheCodecProvider } from "../../kv-store";
+import { disposeResources, cleanupFailure } from "../../engine/resources";
+import type { MlxGroupPreparation, Row } from "./batch-group";
+import { evalCacheState } from "./prefill";
+import { prefillCacheLayout } from "./cache-layout";
+import { MlxStateRows } from "./state-rows";
+import type { P2RTracePhase, P2RTraceAttributes } from "../../serve/prompt-response-trace";
+import type { MlxForwardWork, MlxPreparationWork } from "./mixed-iteration";
+import type { KvMaintenance } from "./kv-maintenance";
+
+// Shared work is recorded on every participating request. workId identifies
+// duplicates when constructing a process timeline; row spans are not additive.
+let nextWorkId = 0;
+function traceRows(states: readonly MlxPrefillState[], phase: P2RTracePhase, attributes: P2RTraceAttributes) {
+  if (!states.some(state => state.row.req.trace)) return null;
+  const closes = states.map(state => state.row.req.trace?.begin(phase, attributes));
+  return { [Symbol.dispose]() { for (const close of closes) close?.(); } };
+}
+
+export interface MlxPrefillStep extends PrefillStep {
+  /** A non-causal input cannot be split at a scheduler token budget.
+   * Its preparation must decline joins until this step completes. */
+  readonly atomic?: boolean;
+  /** Some methods maintain their final target chunk before sampling. */
+  readonly maintain?: boolean;
+}
+
+export interface MlxPrefillState {
+  row: Row;
+  solo: Cache[];
+  pos: number;
+  retain?: () => void;
+  planned?: MlxPrefillStep;
+}
+
+/** The method owns initialization, planning, companion state, snapshots and
+ * completion. This driver owns common target batching and work boundaries.
+ * Checkpoint factories return owned cache state; close/dispose are idempotent. */
+export interface MlxPrefillRowsHost<State extends MlxPrefillState> {
+  stateCodecs?: CacheCodecProvider;
+  maintain?: KvMaintenance;
+  open(row: Row): State;
+  plan(state: State): MlxPrefillStep;
+  /** Already prepared state can complete without any target forward. */
+  ready?(state: State): boolean;
+  forward(ids: MlxArray, caches: Cache[], states: readonly State[], work?: MlxForwardWork): Promise<MlxArray>;
+  project(hidden: MlxArray, caches: Cache[], completed: readonly State[]): MlxArray | null;
+  /** Consume method context after target maintenance, before snapshots. */
+  afterForward?(ids: MlxArray, caches: Cache[], states: readonly State[], hidden: MlxArray): void | Promise<void>;
+  checkpoint(state: State, capture: () => Cache[], row: number): void;
+  complete(state: State, logits: MlxArray | null, row: number): Promise<void>;
+  reject(row: Row, error: unknown): void;
+  close(state: State): void;
+  filterRows?(keep: readonly number[]): void;
+  dispose?(): void;
+}
+
+/** One target preparation lifecycle at B1/B>1. Late joins retain each row's
+ * planned chunk endpoint, precision transition and checkpoint boundary. */
+export class MlxPrefillRows<State extends MlxPrefillState> implements MlxGroupPreparation {
+  #states: State[] = [];
+  #rows: Row[] = [];
+  #stateRows?: MlxStateRows;
+  constructor(readonly operations: MlxPrefillRowsHost<State>) {}
+  get rows(): readonly Row[] { return this.#rows; }
+  get canAdmit(): boolean { return this.#rows.length > 0; }
+  /** Keep a request's prefill weight until retirement so a nearly completed
+   * request does not repeatedly attract another long prefill into its tail. */
+  get tokenWeight(): number { return this.#states.reduce((sum, state) => sum + state.row.promptTokens - state.row.cachedTokens, 0); }
+
+  admit(row: Row): void {
+    this.#states.push(this.operations.open(row));
+    this.#rows.push(row);
+  }
+
+  #release(state: State): void {
+    const caches = state.solo; state.solo = [];
+    const retain = state.retain; state.retain = undefined;
+    disposeResources([...caches, { dispose: () => retain?.() }, { dispose: () => this.operations.close(state) }]);
+  }
+  #filter(keep: number[]): void {
+    if (keep.length) this.#stateRows?.filterRows(keep);
+    else { const rows = this.#stateRows; this.#stateRows = undefined; rows?.dispose(); }
+    this.operations.filterRows?.(keep);
+    this.#states = keep.map(row => this.#states[row]!);
+    this.#rows = keep.map(row => this.#rows[row]!);
+  }
+  /** Admission changes membership only between forwards. Existing planned
+   * endpoints survive a join, including partially consumed drain chunks. */
+  #synchronizeRows(): void {
+    if (!this.#stateRows && this.#states.length === 1) return;
+    if (!this.#stateRows) {
+      for (const state of this.#states) this.operations.maintain?.preparePrefill?.(state.solo);
+      const layouts: BatchableCache[] = [];
+      try { for (const cache of this.#states[0]!.solo) layouts.push(prefillCacheLayout(cache)); }
+      catch (error) { return cleanupFailure(error, () => disposeResources(layouts)); }
+      this.#stateRows = new MlxStateRows(layouts);
+      this.#stateRows.mergeRows(this.#states.map(state => state.solo));
+    } else {
+      for (const state of this.#states.slice(this.#stateRows.rowCount)) {
+        this.operations.maintain?.preparePrefill?.(state.solo);
+        this.#stateRows.append(state.solo);
+      }
+    }
+    // Resolve input-backed copies before releasing SSD restore leases.
+    if (this.#states.some(state => state.solo.length)) {
+      evalCacheState(this.#stateRows.caches);
+      for (const state of this.#states) {
+        const caches = state.solo; state.solo = [];
+        const retain = state.retain; state.retain = undefined;
+        disposeResources([...caches, { dispose: () => retain?.() }]);
+      }
+    }
+    for (const cache of this.#stateRows.caches) cache.prefillMaintenance?.beginPrefill();
+  }
+
+  async advance(workLimit?: MlxPreparationWork): Promise<boolean> {
+    let remaining = workLimit?.maxTokens ?? Infinity;
+    // Complete restored inputs before merging: their cache coverage must not
+    // introduce padding or a zero-width forward into unrelated cold rows.
+    if (this.operations.ready && this.#states.some(state => this.operations.ready!(state))) {
+      const remaining: State[] = [];
+      for (const state of this.#states) {
+        if (!this.operations.ready(state)) { remaining.push(state); continue; }
+        try {
+          state.row.req.signal?.throwIfAborted();
+          await this.operations.complete(state, null, -1);
+          state.solo = []; state.retain = undefined;
+        } catch (error) { this.operations.reject(state.row, error); this.#release(state); }
+        finally { this.operations.close(state); }
+      }
+      this.#states = remaining; this.#rows = remaining.map(state => state.row);
+      if (!remaining.length) { this.dispose(); return true; }
+    }
+    {
+      using span = traceRows(this.#states, "prefill.row_sync", { workId: ++nextWorkId, batchSize: this.#states.length });
+      this.#synchronizeRows();
+    }
+    const keep: number[] = [];
+    for (const [index, state] of this.#states.entries()) {
+      if (!state.row.req.signal?.aborted) { keep.push(index); continue; }
+      this.#release(state); this.operations.reject(state.row, state.row.req.signal.reason);
+    }
+    if (keep.length !== this.#states.length) this.#filter(keep);
+    if (!this.#states.length) { this.dispose(); return true; }
+    while (this.#states.length) {
+      for (const state of this.#states) state.planned ??= this.operations.plan(state);
+      const count = Math.min(this.#states.some(state => state.planned!.atomic) ? Infinity
+        : Math.max(1, Math.floor(remaining / this.#states.length)),
+        ...this.#states.map(state => state.planned!.end - state.pos));
+      const work = { workId: ++nextWorkId, batchSize: this.#states.length, tokensPerRow: count };
+      const closes = this.#states.map(state => state.row.req.trace?.begin("prefill.chunk", {
+        mechanism: "continuous", startToken: state.pos, ...work,
+      }));
+      const caches = this.#stateRows?.caches ?? this.#states[0]!.solo;
+      try {
+        using ids = count ? ops.fromInt32(this.#states.flatMap(state => state.row.req.promptIds.slice(state.pos, state.pos + count)),
+          [this.#states.length, count]) : null;
+        let forwarded: MlxArray | null;
+        {
+          using span = traceRows(this.#states, "prefill.forward", work);
+          forwarded = ids ? await this.operations.forward(ids, caches, this.#states, workLimit?.forward) : null;
+        }
+        using hidden = forwarded;
+        const drains: number[] = [], finals: number[] = [];
+        let batchYield = false;
+        for (const [row, state] of this.#states.entries()) {
+          state.pos += count;
+          if (state.pos !== state.planned!.end) continue;
+          if (state.planned!.kind === "final") finals.push(row);
+          else drains.push(row);
+          batchYield ||= state.planned!.batchYield;
+        }
+        const maintained = [...drains, ...finals.filter(row => this.#states[row]!.planned!.maintain)];
+        if (count && (maintained.length || !finals.length)) {
+          {
+            using span = traceRows(this.#states, "prefill.evaluate", work);
+            evalCacheState(caches);
+          }
+          using span = traceRows(this.#states, "prefill.kv_maintenance", work);
+          if (this.#stateRows) {
+            for (const cache of caches) cache.prefillMaintenance?.commitPrefill(maintained);
+          } else if (maintained.length) this.operations.maintain?.(this.#states[0]!.solo);
+          // A scheduler budget can split one planned chunk into many pieces.
+          // Resolve state at each yield, but retain allocator reuse until the
+          // method's maintenance boundary instead of purging it every piece.
+          if (maintained.length) clearCache();
+        }
+        if (ids && this.operations.afterForward) {
+          using span = traceRows(this.#states, "prefill.companion", work);
+          const companion = this.operations.afterForward?.(ids, caches, this.#states, hidden!);
+          if (companion) await companion;
+        }
+        for (const row of [...drains, ...finals]) {
+          const state = this.#states[row]!;
+          if (state.planned!.snapshot) {
+            using span = traceRows([state], "prefill.checkpoint", { ...work, row });
+            this.operations.checkpoint(state, () =>
+              this.#stateRows ? this.#stateRows.extractRow(row) : cloneKvCaches(state.solo, this.operations.stateCodecs), row);
+          }
+          if (!finals.includes(row)) state.planned = undefined;
+        }
+        let projected: MlxArray | null = null;
+        if (hidden && finals.length) {
+          using span = traceRows(this.#states, "prefill.project", work);
+          projected = this.operations.project(hidden, caches, finals.map(row => this.#states[row]!));
+        }
+        using logits = projected;
+        for (const row of finals) {
+          const state = this.#states[row]!;
+          using span = traceRows([state], "prefill.complete", { ...work, row });
+          if (this.#stateRows) state.solo = this.#stateRows.extractRow(row);
+          using part = logits ? logits.slice([row, logits.shape[1]! - 1, 0], [row + 1, logits.shape[1]!, logits.shape[2]!]) : null;
+          try {
+            state.row.req.signal?.throwIfAborted();
+            await this.operations.complete(state, part, row);
+            state.solo = []; state.retain = undefined;
+          } catch (error) {
+            this.operations.reject(state.row, error);
+            this.#release(state);
+          } finally { this.operations.close(state); }
+        }
+        for (const cache of caches) (cache as { releaseRopeArr?: () => void }).releaseRopeArr?.();
+        if (finals.length) this.#filter(this.#states.flatMap((_, row) => finals.includes(row) ? [] : [row]));
+        if (!this.#states.length) { this.dispose(); return true; }
+        remaining -= count * work.batchSize;
+        if (batchYield || finals.length || remaining < this.#states.length) return false;
+      } finally { for (const close of closes) close?.(); }
+    }
+    return true;
+  }
+
+  dispose(): void {
+    const states = this.#states; this.#states = []; this.#rows = [];
+    const rows = this.#stateRows; this.#stateRows = undefined;
+    disposeResources([...(rows ? [rows] : []), ...states.map(state => ({ dispose: () => this.#release(state) })),
+      { dispose: () => this.operations.dispose?.() }]);
+  }
+}

package/src/backends/mlx/prefill.ts

@@ -9,7 +9,7 @@ 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]));
+  withResource(leaseCacheStates(cache), (state) => { if (state.length) ops.evalAll([...state]); });
 }
 
 /** Forward once. Drains evaluate state, apply KV maintenance, then clear the