mlx-bun 0.4.0 → 0.5.0

package/src/backends/mlx/kv-maintenance.ts

@@ -1,3 +1,5 @@
+import { turboQuantFusedDecode } from "../../model/turboquant-codec";
+import { runtimeConfig, withRuntimeConfig } from "../../runtime-config";
 import type { KvSchemeOptions } from "../../kv-scheme";
 import { KVCache, QuantizedKVCache, RotatingKVCache, RotatingQuantizedKVCache, TurboQuantKVCache, type Cache } from "../../model/gemma4-base";
 import type { KvQuantSpec, TurboQuantScheme } from "../../config";
@@ -10,6 +12,8 @@ import { DelayedTurboQuantKVCache } from "../../model/delayed-turboquant-kv";
 
 export interface KvMaintenance {
   (cache: Cache[]): void;
+  /** Limit committed work at a pending precision transition. */
+  maxAppendTokens?(cache: readonly Cache[]): number;
   /** Bind state requiring row-local maintenance before shared decode. */
   prepareBatch?(cache: Cache[]): void;
   /** Bind all precision policies before a prefill cohort owns row boundaries. */
@@ -43,12 +47,24 @@ export function createKvMaintenance(options: Readonly<Omit<KvSchemeOptions, "kvC
   if (turboQuant) {
     const start = options.quantizedKvStart ?? 0;
     const scheme = { ...turboQuant };
-    const maintain: KvMaintenance = (cache) => maybeTurboQuantizeKv(cache, scheme, start);
+    const runtime = runtimeConfig();
+    const fusedDecode = turboQuantFusedDecode(runtime);
+    const maintain: KvMaintenance = (cache) => withRuntimeConfig(runtime, () => maybeTurboQuantizeKv(cache, scheme, start));
+    if (start > 0) maintain.maxAppendTokens = (cache) => {
+      let remaining = Number.POSITIVE_INFINITY;
+      for (const c of cache) {
+        remaining = Math.min(remaining, c.maxAppendTokens?.() ?? Infinity);
+        const conversion = c.turboConversion ?? (c instanceof KVCache ? c : undefined);
+        if (conversion && conversion.offset < start)
+          remaining = Math.min(remaining, start - conversion.offset);
+      }
+      return remaining;
+    };
     maintain.preparePrefill = (cache) => {
       for (let layer = 0; layer < cache.length; layer++) {
         const row = cache[layer]!;
         if (row instanceof KVCache || row instanceof TurboQuantKVCache)
-          cache[layer] = new DelayedTurboQuantKVCache(scheme.kBits, scheme.vBits, start, maintain, row);
+          cache[layer] = new DelayedTurboQuantKVCache(scheme.kBits, scheme.vBits, start, maintain, row, fusedDecode);
       }
     };
     if (start > 0) maintain.prepareBatch = maintain.preparePrefill;
@@ -83,6 +99,18 @@ export function createKvMaintenance(options: Readonly<Omit<KvSchemeOptions, "kvC
       clearCache();
     }
   };
+  if (start > 0) maintain.maxAppendTokens = (cache) => {
+    let remaining = Number.POSITIVE_INFINITY;
+    for (let layer = 0; layer < cache.length; layer++) {
+      if (byLayer && !byLayer.has(layer)) continue;
+      const c = cache[layer]!;
+      remaining = Math.min(remaining, c.maxAppendTokens?.() ?? Infinity);
+      const conversion = c.affineConversion ?? (c instanceof KVCache || c instanceof RotatingKVCache ? c : undefined);
+      if (conversion && conversion.offset < start)
+        remaining = Math.min(remaining, start - conversion.offset);
+    }
+    return remaining;
+  };
   maintain.preparePrefill = (cache) => {
     for (let layer = 0; layer < cache.length; layer++) {
       const row = cache[layer]!;

package/src/backends/mlx/mixed-iteration.ts

@@ -0,0 +1,62 @@
+import type { MlxArray } from "../../mlx/array";
+import type { Cache } from "../../model/gemma4-base";
+import type { TokenGroup, TokenWorkOptions } from "../../model/token-groups";
+
+export type MlxForwardWork = (ids: MlxArray, cache: Cache[], options?: TokenWorkOptions) => Promise<MlxArray>;
+export interface MlxPreparationWork {
+  /** Real prompt tokens across all rows, not a padded sequence width. */
+  readonly maxTokens?: number;
+  readonly forward?: MlxForwardWork;
+}
+
+/** Join two method-owned work producers at their forward boundary. Decode
+ * finishes sampling/publication/retirement before preparation can join rows.
+ * The model owns packing; this coordinator never merges caches or samples.
+ * A restored/cancelled preparation need not produce a forward at all. */
+export async function runMixedTokenIteration(options: {
+  prepare(forward: MlxForwardWork): Promise<void>;
+  decode(forward: MlxForwardWork): Promise<void>;
+  forward: MlxForwardWork;
+  mixed(groups: readonly TokenGroup[]): MlxArray[];
+}): Promise<boolean> {
+  let notify!: () => void;
+  const offered = new Promise<void>(resolve => { notify = resolve; });
+  let pending: { group: TokenGroup; resolve(hidden: MlxArray): void; reject(error: unknown): void } | undefined;
+  let captured = false;
+  let preparationError: unknown;
+  let preparationFailed = false;
+  const preparing = options.prepare(async (ids, cache, policy) => {
+    if (captured) return options.forward(ids, cache, policy);
+    captured = true;
+    return new Promise<MlxArray>((resolve, reject) => {
+      pending = { group: { ids, cache, ...policy }, resolve, reject }; notify();
+    });
+  }).catch(error => { preparationFailed = true; preparationError = error; }).finally(notify);
+  await offered;
+  if (!pending) {
+    await preparing;
+    if (preparationFailed) throw preparationError;
+    return false;
+  }
+  const held = pending;
+  let preparedHidden: MlxArray | undefined;
+  let mixed = false;
+  try {
+    await options.decode(async (ids, cache, policy) => {
+      if (mixed) return options.forward(ids, cache, policy);
+      const outputs = options.mixed([{ ids, cache, ...policy }, held.group]);
+      mixed = true; preparedHidden = outputs[1]!;
+      return outputs[0]!;
+    });
+    // A final unread token can retire decode without another model call.
+    preparedHidden ??= await options.forward(held.group.ids, held.group.cache, held.group);
+    held.resolve(preparedHidden); preparedHidden = undefined;
+    await preparing;
+    if (preparationFailed) throw preparationError;
+    return true;
+  } catch (error) {
+    preparedHidden?.dispose(); held.reject(error);
+    await preparing;
+    throw error;
+  }
+}

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
@@ -27,8 +27,9 @@ export function modelServingBinding(context: ServingContext): ModelServingBindin
     stateCompatibility: `mlx-${MLX_VERSION}-${deviceArchitecture()}`,
     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

@@ -32,6 +32,10 @@ interface PrefillHost {
 
 /** 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,
@@ -47,7 +51,7 @@ export class MlxPrefillCohort extends MlxPrefillRows<PrefillState> {
             return { row, solo: continuation.caches, pos: row.req.promptIds.length,
               snapAt: null, continuation };
           }
-          const hit = host.promptCache?.take(row.req.promptIds, row.cacheNamespace) ?? null;
+          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;
@@ -59,6 +63,7 @@ export class MlxPrefillCohort extends MlxPrefillRows<PrefillState> {
           });
           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 ?? []),
@@ -66,16 +71,22 @@ export class MlxPrefillCohort extends MlxPrefillRows<PrefillState> {
         } finally { closeCache?.(); }
       },
       ready: state => !!state.continuation,
-      plan(state) { return nextPrefillStep({ length: state.row.promptTokens, position: state.pos,
+      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: host.forward.bind(host), project: host.project.bind(host),
+      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);
+          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);

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

@@ -9,9 +9,23 @@ 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;
 }
@@ -34,7 +48,7 @@ export interface MlxPrefillRowsHost<State extends MlxPrefillState> {
   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[]): Promise<MlxArray>;
+  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>;
@@ -105,7 +119,8 @@ export class MlxPrefillRows<State extends MlxPrefillState> implements MlxGroupPr
     for (const cache of this.#stateRows.caches) cache.prefillMaintenance?.beginPrefill();
   }
 
-  async advance(): Promise<boolean> {
+  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))) {
@@ -122,7 +137,10 @@ export class MlxPrefillRows<State extends MlxPrefillState> implements MlxGroupPr
       this.#states = remaining; this.#rows = remaining.map(state => state.row);
       if (!remaining.length) { this.dispose(); return true; }
     }
-    this.#synchronizeRows();
+    {
+      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; }
@@ -132,15 +150,23 @@ export class MlxPrefillRows<State extends MlxPrefillState> implements MlxGroupPr
     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.map(state => state.planned!.end - state.pos));
+      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, batchSize: this.#states.length,
+        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;
-        using hidden = ids ? await this.operations.forward(ids, caches, this.#states) : 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()) {
@@ -152,25 +178,42 @@ export class MlxPrefillRows<State extends MlxPrefillState> implements MlxGroupPr
         }
         const maintained = [...drains, ...finals.filter(row => this.#states[row]!.planned!.maintain)];
         if (count && (maintained.length || !finals.length)) {
-          evalCacheState(caches);
+          {
+            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);
-          clearCache();
+          // 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) {
+        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) this.operations.checkpoint(state, () =>
-            this.#stateRows ? this.#stateRows.extractRow(row) : cloneKvCaches(state.solo, this.operations.stateCodecs), 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;
         }
-        using logits = hidden && finals.length ? this.operations.project(hidden, caches, finals.map(row => this.#states[row]!)) : null;
+        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 {
@@ -185,7 +228,8 @@ export class MlxPrefillRows<State extends MlxPrefillState> implements MlxGroupPr
         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; }
-        if (batchYield || finals.length) return false;
+        remaining -= count * work.batchSize;
+        if (batchYield || finals.length || remaining < this.#states.length) return false;
       } finally { for (const close of closes) close?.(); }
     }
     return true;

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

@@ -0,0 +1,29 @@
+import type { MlxArray } from "../../mlx/array";
+import type { Cache } from "../../model/gemma4-base";
+
+/** Prepared native input with an indivisible prefill. The preparation owner
+ * retains its tensors; the model binding returns an owned final hidden row.
+ * Decode, sampling and row retirement use the ordinary execution group. */
+export interface MlxPromptInput {
+  forward(ids: MlxArray, caches: Cache[], start?: number): MlxArray;
+  readonly decodeState?: MlxDecodeState;
+}
+
+/** Additional model-owned forward state. Equal keys declare compatible row
+ * state; the backend supplies the current row order after joins/retirement. */
+export interface MlxDecodeState {
+  readonly key: string;
+  forward(ids: MlxArray, caches: Cache[], rows: readonly MlxDecodeState[]): MlxArray;
+}
+
+/** Preserve the ordinary generator's projection geometry for a prepared
+ * embedding sequence: process the whole prompt, then project only its tip. */
+export function bindEmbeddingsInput(
+  forward: (ids: MlxArray, caches: Cache[], start: number) => MlxArray,
+): MlxPromptInput {
+  return { forward(ids, caches, start = 0) {
+    using hidden = forward(ids, caches, start);
+    const [batch, length, width] = hidden.shape;
+    return hidden.slice([0, length! - 1, 0], [batch!, length!, width!]);
+  } };
+}

package/src/backends/mlx/qwen-prompt-input.ts

@@ -0,0 +1,40 @@
+import type { Qwen35Model } from "../../model/qwen3_5";
+import { mropePositionIds, type MropeRequestState } from "../../model/qwen3-mrope";
+import type { MlxArray } from "../../mlx/array";
+import type { Cache } from "../../model/gemma4-base";
+import * as ops from "../../mlx/ops";
+import { Dtype } from "../../mlx/ffi";
+import { bindEmbeddingsInput, type MlxDecodeState, type MlxPromptInput } from "./prompt-input";
+
+class QwenMediaDecodeState implements MlxDecodeState {
+  readonly key = "qwen-interleaved-mrope";
+  constructor(readonly model: Qwen35Model, readonly delta: number) {}
+
+  forward(ids: MlxArray, caches: Cache[], rows: readonly MlxDecodeState[]): MlxArray {
+    const B = rows.length, L = ids.shape[1]!;
+    const cache = caches[this.model.faIdx]!;
+    // The cache supplies logical positions, including left padding, on device.
+    // No offset readback or request-global position mutation is needed.
+    const offsets = (cache as Cache & { ropeOffsetArr?: MlxArray }).ropeOffsetArr;
+    using deltas = ops.fromInt32(rows.map(row => (row as QwenMediaDecodeState).delta), [B]);
+    using scalar = offsets ? null : ops.fromInt32([cache.offset], []);
+    using starts = ops.add(offsets ?? scalar!, deltas);
+    using startRows = ops.reshape(starts, [B, 1]);
+    using indices = ops.arange(0, L, 1, Dtype.int32);
+    using positions = ops.add(startRows, indices);
+    using axis = ops.reshape(positions, [1, B, L]);
+    using grid = ops.concatAxis([axis, axis, axis], 0);
+    return this.model.forwardHiddenAtPositions(ids, caches, grid);
+  }
+}
+
+/** The prompt owns its grid; only its continuation delta survives into decode. */
+export function bindQwenMediaInput(model: Qwen35Model, embeddings: MlxArray,
+  state: MropeRequestState): MlxPromptInput {
+  const input = bindEmbeddingsInput((ids, caches, start) => {
+    using positions = mropePositionIds(state, start, ids.shape[1]!);
+    return start > 0 ? model.forwardHiddenAtPositions(ids, caches, positions)
+      : model.forwardEmbeddingsAtPositions(embeddings, caches, positions);
+  });
+  return { ...input, decodeState: new QwenMediaDecodeState(model, state.delta) };
+}