mlx-bun 0.3.0 → 0.5.0

package/src/backends/mlx/gateway-binding.ts

@@ -1,3 +1,8 @@
+import { speculativePrefixNamespace, captureSpeculativeOptions } from "../../spec/cache-identity";
+import { createOrdinaryContinuationRequest } from "./continuation-request";
+import type { MlxSerialServices } from "./serial-executor";
+import { bindPagedRequestState, pagedPrefixNamespace, type MlxRequestStatePolicy } from "./request-state-policy";
+import type { ExecutionContext } from "../../contracts/scheduling";
 import type { ModelConfig } from "../../config";
 import type { KvScheme } from "../../kv-scheme";
 import type { RuntimeModel } from "../../model/factory";
@@ -5,67 +10,165 @@ import { DiffusionGemmaModel } from "../../model/diffusion-gemma";
 import { UniversalDenseModel } from "../../model/universal/dense";
 import { KVCache, RotatingKVCache, isBatchableCache, isPlainKvCache, isRotatingPlainCache } from "../../model/gemma4-base";
 import { SSMCache } from "../../model/qwen3-delta";
+import { Gemma4Model } from "../../model/gemma4";
 import { Qwen35Model } from "../../model/qwen3_5";
 import { runtimeConfig, type RuntimeConfig } from "../../runtime-config";
 import { disposeResources } from "../../engine/resources";
 import { legacyCompiledDecodeAvailable } from "./autoregressive";
-import { MlxBatchExecutionGroup, type MlxBatchExecutionGroupOptions } from "./batch-group";
+import { MlxBatchExecutionGroup, type MlxBatchExecutionGroupOptions, type MlxGroupMethodRequest } from "./batch-group";
+import type { DraftProvider } from "../../spec/source";
+import { constraintDraftProvider } from "../../spec/ngram-source";
+import { targetRowLayoutFactory } from "./target-layout-capability";
+import { bindSpeculativeGroupRequests } from "./speculative-group";
+import { bindFillGroupRequests } from "./fill-group";
 import type { GenerateOptions } from "../../generate";
 import type { ExecutionRequirements, ResolvedExecution } from "../../contracts/execution";
 import { resolveExecution } from "../../engine/execution-plan";
+import { bindEmbeddingsInput, type MlxPromptInput } from "./prompt-input";
+import { bindQwenMediaInput } from "./qwen-prompt-input";
+import type { Vision } from "../../serve/generation-gateway";
 
 export interface MlxBatchGroup extends Pick<MlxBatchExecutionGroup,
-  "activeRows" | "pendingRows" | "projectedKvBytes" | "kvBudgetBytes" | "submit" | "kick" | "close"> {}
+  "activeRows" | "pendingRows" | "projectedKvBytes" | "kvBudgetBytes" | "submit" | "kick" | "close"> {
+  runPreparation?<T>(work: () => Promise<T>, signal?: AbortSignal): Promise<T>;
+}
 
 /** A model implementation owns capability checks and the execution group.
  * Scheduling never inspects concrete model/cache classes. */
 export interface MlxGatewayBinding {
+  mediaInput?(input: Vision): MlxPromptInput;
+  configureContinuation?(services: MlxSerialServices): void;
+  continuationRequest?(execution: ResolvedExecution | undefined, options: GenerateOptions, prompt: number[],
+    onToken: Parameters<typeof createOrdinaryContinuationRequest>[0]["onToken"]): ReturnType<typeof createOrdinaryContinuationRequest> | undefined;
   readonly config: ModelConfig;
   readonly runtime: RuntimeConfig;
   plan(request: ExecutionRequirements, options: GenerateOptions,
     scheduling: { continuous: boolean; quantizedBatch: boolean; checkpoints: boolean }): ResolvedExecution;
+  bindAdapterContext?(adapters: string[], key: string): ExecutionContext;
   cachesBatchable(): boolean;
   kvBatchable(scheme: KvScheme): boolean;
+  statePolicy?(execution: ResolvedExecution | undefined, options: GenerateOptions, capacityTokens: number): MlxRequestStatePolicy | undefined;
+  prefixNamespace?(execution: ResolvedExecution | undefined, options: GenerateOptions, adapters: string): string | null;
+  methodRequest?(execution: ResolvedExecution | undefined, options: GenerateOptions): MlxGroupMethodRequest | undefined;
   createBatchGroup(options: MlxBatchExecutionGroupOptions): MlxBatchGroup;
 }
 
-export function bindMlxGateway(model: RuntimeModel): MlxGatewayBinding {
+export function bindMlxGateway(model: RuntimeModel, draft?: { provider: DraftProvider; numDraftTokens: number }): MlxGatewayBinding {
   const runtime = runtimeConfig();
+  let continuationServices: MlxSerialServices | undefined;
+  const kvBatchCapabilities = { delayedAffine: model instanceof Qwen35Model || model instanceof Gemma4Model };
+  const cachesBatchable = () => {
+    if (model instanceof UniversalDenseModel)
+      return !model.args.maskArray && !model.args.layerTypes?.includes("sliding_attention");
+    const caches = model.makeCache();
+    try {
+      const ssm = runtime.value("MLX_BUN_BATCH_SSM") !== "0";
+      return caches.every((cache) => cache instanceof KVCache || cache instanceof RotatingKVCache ||
+        isBatchableCache(cache) || (ssm && cache instanceof SSMCache));
+    } finally { disposeResources(caches); }
+  };
+  const supportsTargetRows = () => {
+    const caches = model.makeCache();
+    try { return caches.every(cache => targetRowLayoutFactory(cache) !== undefined); }
+    finally { disposeResources(caches); }
+  };
+  const speculative = draft?.provider.grouped && cachesBatchable() && supportsTargetRows()
+    ? bindSpeculativeGroupRequests(model, draft.provider, draft.numDraftTokens) : undefined;
+  const grammarProvider = runtime.flag("MLX_BUN_GRAMMAR_JUMP", false) && cachesBatchable() && supportsTargetRows()
+    ? constraintDraftProvider() : undefined;
+  const grammarProposals = grammarProvider ? bindSpeculativeGroupRequests(model, grammarProvider,
+    Math.max(1, Math.trunc(runtime.number("MLX_BUN_GRAMMAR_DRAFT_TOKENS", 3)))) : undefined;
+  const adapterState = "loraState" in model ? model.loraState : undefined;
+  const fillRequests = supportsTargetRows() ? bindFillGroupRequests(model) : undefined;
+  const mediaInput = model instanceof Gemma4Model ? (input: Vision) =>
+    bindEmbeddingsInput((ids, caches, start) => start > 0 ? model.forwardHidden(ids, caches)
+      : model.forwardEmbeddings(input.embeddings,
+      caches, input.imageMask ?? null, ids, input.multimodalMask ?? null))
+    : model instanceof Qwen35Model ? (input: Vision) => bindQwenMediaInput(model, input.embeddings, input.mrope!)
+    : undefined;
   return {
+    mediaInput,
     config: model.config, runtime,
+    configureContinuation: services => { continuationServices = services; },
+    continuationRequest(execution, options, prompt, onToken) {
+      if (!execution?.checkpoint || execution.mechanism !== "continuous") return undefined;
+      const services = continuationServices;
+      if (!services?.checkpoints || !services.checkpointPersistence)
+        throw new Error("qualified continuation requires bound persistence services");
+      return createOrdinaryContinuationRequest({ options, prompt, onToken, execution,
+        store: services.checkpoints, persistence: services.checkpointPersistence,
+        restore: entry => services.checkpoints!.restore(entry, model),
+        interval: services.checkpointEveryTokens!, identity: services.identity });
+    },
+    statePolicy: (execution, options, capacity) => {
+      // Media token IDs alone do not identify the prepared embeddings.
+      // Preserve the existing uncached media policy through the state port.
+      if (execution?.method === "autoregressive" && !execution.promptCache)
+        return { key: "uncached-prepared-input", create: () => model.makeCache() };
+      return execution?.pagedKv ? bindPagedRequestState(model, options, capacity, continuationServices?.promptCache, runtime) : undefined;
+    },
+    prefixNamespace: (execution, options, adapters) => {
+      if (execution?.pagedKv) return pagedPrefixNamespace(options, adapters, runtime.flag("MLX_BUN_PAGED_ATTN", false));
+      if (execution?.method !== "speculative") return adapters;
+      const namespace = (execution.grammarJump ? grammarProvider : draft?.provider)?.grouped?.checkpointNamespace?.();
+      return namespace === undefined ? null : speculativePrefixNamespace(namespace, adapters, captureSpeculativeOptions(options));
+    },
+    methodRequest: (execution, options) => execution?.method === "speculative"
+      ? (execution.grammarJump ? grammarProposals : speculative)?.(
+        options.fill && !execution.fill ? { ...options, fill: undefined } : options)
+      : execution?.fill ? fillRequests?.(options) : undefined,
+    ...(adapterState ? { bindAdapterContext(adapters: string[], key: string): ExecutionContext {
+      const selected = [...adapters];
+      return { key, enter() {
+        const previous = adapterState.active;
+        adapterState.active = selected;
+        return () => { adapterState.active = previous; };
+      } };
+    } } : {}),
     plan(request, options, scheduling) {
+      const sharedMethod = request.hasDraft ? speculative : grammarProposals;
+      const provider = request.hasDraft ? draft?.provider : grammarProvider;
       return resolveExecution(request, {
         ...scheduling,
+        sharedCheckpoints: !!continuationServices?.checkpointPersistence &&
+          !request.hasDraft && !request.hasVision && !request.hasGrammar &&
+          !request.wantsLogprobs && !options.fill && !options.pagedKv,
+        adapterBatch: !!adapterState, pagedBatch: model instanceof Gemma4Model,
+        mediaBatch: !!mediaInput,
+        mediaPrefixCache: runtime.flag("MLX_BUN_MEDIA_PREFIX_CACHE", true),
+        groupedMethods: sharedMethod ? ["autoregressive", "speculative"] : ["autoregressive"],
+        sharedGrammarProposals: !!grammarProposals,
+        sharedFill: !!fillRequests && !!options.fill,
+        sharedSpeculativeEcho: !!options.fill?.plan.echo && provider?.grouped?.supportsExternalTokens === true,
+        speculativeLogprobs: scheduling.continuous && !!sharedMethod,
+        sharedSpeculativeAdapters: scheduling.continuous && !!sharedMethod && !!adapterState &&
+          provider?.grouped?.supportsTargetAdapters === true,
+        turboQuantBatch: scheduling.quantizedBatch,
+        speculativeTurboQuant: scheduling.continuous && !!sharedMethod && !!options.turboQuant,
         method: model instanceof DiffusionGemmaModel ? "denoising" : "autoregressive",
         compiledDecode: legacyCompiledDecodeAvailable(model),
         grammarBatch: runtime.value("MLX_BUN_GRAMMAR_BATCH") !== "0",
-        speculativeKvQuant: model instanceof Qwen35Model && options.kvBits === 4 &&
-          !options.kvConfig?.length && options.quantizedKvStart === 0 &&
-          runtime.flag("MLX_BUN_QWEN_SPEC_KV4", false),
+        speculativeKvQuant: (!(model instanceof Qwen35Model) || runtime.flag("MLX_BUN_QWEN_SPEC_KV4", true)) && (
+          (scheduling.continuous && !!sharedMethod && (options.kvBits === 4 || options.kvBits === 8 || !!options.kvConfig?.length)) ||
+          (!options.kvConfig?.length && model instanceof Qwen35Model &&
+            (options.kvBits === 4 || (scheduling.continuous && !!sharedMethod && options.kvBits === 8)) &&
+            (options.quantizedKvStart === 0 || (scheduling.continuous && !!sharedMethod)))
+        ),
       }, {
         pagedKv: !!options.pagedKv, fill: !!options.fill,
         compiledDecode: runtime.flag("MLX_BUN_COMPILED_DECODE", true),
         grammarJump: runtime.flag("MLX_BUN_GRAMMAR_JUMP", false),
       });
     },
-    cachesBatchable() {
-      if (model instanceof UniversalDenseModel)
-        return !model.args.maskArray && !model.args.layerTypes?.includes("sliding_attention");
-      const caches = model.makeCache();
-      try {
-        const ssm = runtime.value("MLX_BUN_BATCH_SSM") !== "0";
-        return caches.every((cache) => cache instanceof KVCache || cache instanceof RotatingKVCache ||
-          isBatchableCache(cache) || (ssm && cache instanceof SSMCache));
-      } finally { disposeResources(caches); }
-    },
+    cachesBatchable,
     kvBatchable(scheme) {
-      if (scheme.kind !== "affine-config") return false;
       const caches = model.makeCache();
       try {
         return scheme.batchable(model.config,
-          (layer) => isPlainKvCache(caches[layer]) || isRotatingPlainCache(caches[layer]));
+          (layer) => isPlainKvCache(caches[layer]) || isRotatingPlainCache(caches[layer]), caches.length,
+          kvBatchCapabilities);
       } finally { disposeResources(caches); }
     },
-    createBatchGroup: (options) => new MlxBatchExecutionGroup(model, options),
+    createBatchGroup: (options) => new MlxBatchExecutionGroup(model, { ...options, kvBatchCapabilities }),
   };
 }

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

@@ -0,0 +1,9 @@
+import type { MlxArray } from "../../mlx/array";
+
+/** A backend KV representation independent of row storage and scheduling.
+ * Inputs are borrowed. Encoding and decoding return caller-owned arrays.
+ * A codec may retain values in a rotated domain for attention to invert once. */
+export interface KvCodec<Encoded> {
+  encode(keys: MlxArray, values: MlxArray): Encoded;
+  decode(encoded: Encoded, length: number, headDim: number, deferValues?: boolean): [MlxArray, MlxArray];
+}

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

@@ -0,0 +1,161 @@
+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";
+import * as ops from "../../mlx/ops";
+import { clearCache } from "../../mlx/ffi";
+
+import { DelayedRotatingQuantizedKVCache } from "../../model/delayed-rotating-quantized-kv";
+import { DelayedQuantizedKVCache } from "../../model/delayed-quantized-kv";
+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. */
+  preparePrefill?(cache: Cache[]): void;
+}
+
+const unchanged = (_cache: Cache[]): void => {};
+let warnedTurboRotating = false;
+
+/** Port of mlx-lm maybe_quantize_kv_cache + BOTH halves of optiq serve's
+ *  per-layer patched variant (incl. patch_rotating_to_quantized: rotating
+ *  caches convert too — Phase 9):
+ *  - per-layer bits/group_size selection (kvConfig overrides kvBits,
+ *    matching optiq's --kv-config precedence; shipped kv_config.json
+ *    files cover EVERY cache-owning layer, sliding ones included —
+ *    verified 12B 48/48, 26B 30/30, e4b 24/24 distinct caches — so
+ *    rotating quantization engages straight from the config; uniform
+ *    kvBits — like optiq --kv-bits — reaches them too), and
+ *  - STREAMING conversion (optiq streaming_kv_quant / serve.py
+ *    patched_maybe_quantize): eval each layer's quantized triples and
+ *    clear the buffer pool before building the next layer's conversion.
+ *    Lazily batching every layer's toQuantized into one eval pins ALL
+ *    layers' bf16 K/V as graph inputs alongside ALL quantized outputs —
+ *    the exact transient optiq's fix kills (16.35 → 7.60 GB at 32k on a
+ *    24 GB Mac). Numerics untouched: same quantize math, only the eval
+ *    ordering is forced (tests/parity/kv-quant.test.ts, tests/parity/rotating-kvq.test.ts). */
+export function createKvMaintenance(options: Readonly<Omit<KvSchemeOptions, "kvConfig">> & {
+  readonly kvConfig?: readonly Readonly<KvQuantSpec>[];
+}): KvMaintenance {
+  const { kvBits, kvConfig, turboQuant } = options;
+  if (turboQuant) {
+    const start = options.quantizedKvStart ?? 0;
+    const scheme = { ...turboQuant };
+    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, fusedDecode);
+      }
+    };
+    if (start > 0) maintain.prepareBatch = maintain.preparePrefill;
+    return maintain;
+  }
+  if (!kvBits && !kvConfig?.length) return unchanged;
+  const start = options.quantizedKvStart ?? (kvConfig?.length ? 0 : 5000);
+  const groupSize = options.kvGroupSize ?? 64;
+  // Resolve layer policy once when composing execution, never per token.
+  const byLayer = kvConfig?.length
+    ? new Map(kvConfig.map((entry) => [entry.layerIdx, { ...entry }]))
+    : null;
+  const maintain: KvMaintenance = (cache) => {
+    for (let i = 0; i < cache.length; i++) {
+      const c = cache[i]!;
+      const conversion = c.affineConversion ?? (c instanceof KVCache || c instanceof RotatingKVCache ? c : undefined);
+      if (!conversion || conversion.offset < start || conversion.offset === 0) continue;
+      const conversionOffset = conversion.offset;
+      // The populated-cache boundary is part of oracle parity: quantizing an
+      // empty cache would also quantize the first prefill.
+      if (byLayer) {
+        const entry = byLayer.get(i);
+        if (!entry) continue;
+        cache[i] = conversion.toQuantized(entry.groupSize, entry.bits);
+      } else {
+        cache[i] = conversion.toQuantized(groupSize, kvBits!);
+      }
+      if (start > 0) cache[i]!.minimumReusableOffset = conversionOffset;
+      // Materialize one layer before converting the next to bound the live
+      // bf16 source plus quantized destination to one conversion at a time.
+      ops.evalAll(cache[i]!.state());
+      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]!;
+      if (!(row instanceof KVCache || row instanceof QuantizedKVCache || row instanceof RotatingKVCache || row instanceof RotatingQuantizedKVCache)) continue;
+      const spec = byLayer ? byLayer.get(layer) : { bits: kvBits!, groupSize };
+      if (!spec) continue;
+      const rowMaintenance = createKvMaintenance({ kvBits: spec.bits, kvGroupSize: spec.groupSize, quantizedKvStart: start });
+      cache[layer] = row instanceof RotatingKVCache || row instanceof RotatingQuantizedKVCache
+        ? new DelayedRotatingQuantizedKVCache(row.maxSize, spec.groupSize, spec.bits, start, rowMaintenance, row)
+        : new DelayedQuantizedKVCache(spec.groupSize, spec.bits, start, rowMaintenance, row);
+    }
+  };
+  if (start > 0) maintain.prepareBatch = maintain.preparePrefill;
+  return maintain;
+}
+
+function maybeTurboQuantizeKv(cache: Cache[], scheme: TurboQuantScheme, start: number): void {
+  for (let i = 0; i < cache.length; i++) {
+    const c = cache[i]!;
+    if (c instanceof RotatingKVCache) {
+      if (!warnedTurboRotating) {
+        warnedTurboRotating = true;
+        console.warn(
+          "[turbo-quant] sliding-window (RotatingKVCache) layers stay bf16 in v1 " +
+          "(full-attention only) — docs/design/turboquant.md.",
+        );
+      }
+      continue;
+    }
+    const conversion = c.turboConversion ?? (c instanceof KVCache ? c : undefined);
+    if (!conversion || conversion.offset < start || conversion.offset === 0) continue;
+    const conversionOffset = conversion.offset;
+    const tq = c.turboConversion ? c.turboConversion.toTurboQuantized(scheme.kBits, scheme.vBits)
+      : TurboQuantKVCache.fromKVCache(c as KVCache, scheme.kBits, scheme.vBits);
+    if (start > 0) tq.minimumReusableOffset = conversionOffset;
+    cache[i] = tq;
+    // state() allocates fresh trimmed slice views for this cache kind
+    // (see evalCacheState) — dispose after materializing (throw included),
+    // or they leak.
+    const state = tq.state();
+    try {
+      ops.evalAll(state);
+    } finally {
+      for (const a of state) a.dispose();
+    }
+    clearCache();
+  }
+}

package/src/backends/mlx/kv-tensor-rows.ts

@@ -0,0 +1,212 @@
+import type { MlxArray } from "../../mlx/array";
+import { Dtype } from "../../mlx/ffi";
+import * as ops from "../../mlx/ops";
+import { materializeCopy } from "../../mlx/materialize";
+import { FullPrefillPadding } from "../../model/full-prefill-padding";
+import type { Mask, PrefillPadding } from "../../model/gemma4-base";
+import { cleanupFailure, disposeResources } from "../../engine/resources";
+
+/** Borrowed tensor planes with shared token positions. Plane shapes may differ
+ * in head width and dtype. Every plane uses [row, head, token, field]. */
+export interface KvTensorRowView {
+  readonly planes: readonly MlxArray[];
+  readonly rowOffsets: readonly number[];
+  readonly leftPad: readonly number[];
+}
+
+/** Position and membership operations shared by packed KV representations.
+ * A codec supplies tensor planes; this owner never encodes or decodes them. */
+export class KvTensorRows implements KvTensorRowView {
+  planes: MlxArray[] = [];
+  rowOffsets: number[] = [];
+  leftPad: number[] = [];
+  #rope: MlxArray | undefined;
+  #beforeRound: number[] | undefined;
+  readonly #padding = new FullPrefillPadding();
+
+  restorePrefillEnds(ends: readonly number[] | undefined): void { this.#padding.restoreEnds(ends); }
+  preparePrefill(padding: PrefillPadding): void {
+    this.#padding.prepare(this, padding); this.#positionsChanged();
+  }
+  finalizePrefill(): void {
+    const next = this.#padding.finalize(this.planes, this);
+    if (next) this.#replace(next, this.rowOffsets, this.leftPad);
+  }
+
+  /** Snapshot logical validity without changing physical row alignment. */
+  captureDonorValidity(): Pick<import("../../model/gemma4-base").KvDonorRows, "offsets" | "starts" | "ends"> {
+    const offsets = this.rowOffsets.map((offset, row) => Math.max(0, this.#padding.validOffset(offset, row)));
+    return { offsets, starts: [...this.leftPad], ends: offsets.map((offset, row) => offset + this.leftPad[row]!) };
+  }
+  get offset(): number { return Math.max(0, ...this.#ends()); }
+  get batchSize(): number | null { return this.rowOffsets.length || null; }
+  #ends(): number[] { return this.rowOffsets.map((offset, row) => offset + this.leftPad[row]!); }
+  get ropeOffsetArr(): MlxArray | undefined {
+    if (this.leftPad.every(pad => pad === 0) && this.rowOffsets.every(offset => offset === this.rowOffsets[0])) return undefined;
+    return this.#rope ??= ops.fromInt32(this.rowOffsets, [this.rowOffsets.length]);
+  }
+  #positionsChanged(): void { this.#rope?.dispose(); this.#rope = undefined; }
+  bytesPerToken(): number {
+    return this.planes.length && this.rowOffsets.length ? this.planes.reduce((bytes, plane) => bytes + plane.nbytes, 0) /
+      (this.rowOffsets.length * this.planes[0]!.shape[2]!) : 0;
+  }
+  trim(count: number): void {
+    this.rowOffsets = this.rowOffsets.map(offset => offset - count); this.#positionsChanged();
+  }
+  specRoundBegin(): void { this.#beforeRound = [...this.rowOffsets]; }
+  specRoundCommit(): void { this.#beforeRound = undefined; }
+  specRoundRollback(keep: number | readonly number[]): void {
+    this.rowOffsets = this.#beforeRound!.map((offset, row) => offset + (typeof keep === "number" ? keep : keep[row]!));
+    this.#beforeRound = undefined; this.#positionsChanged();
+  }
+  makeMask(N: number, window: number | null): Mask {
+    const B = this.rowOffsets.length, ends = this.#ends(), S = this.offset + N;
+    if (window === null && this.leftPad.every(pad => pad === 0) && ends.every(end => end === ends[0]))
+      return { mode: N === 1 ? "" : "causal", arr: null };
+    using starts = ops.fromInt32(ends, [B, 1, 1, 1]);
+    using relative = ops.arange(0, N, 1, Dtype.int32);
+    using queries = ops.reshape(relative, [1, 1, N, 1]);
+    using position = ops.add(starts, queries);
+    using indices = ops.arange(0, S, 1, Dtype.int32);
+    using keys = ops.reshape(indices, [1, 1, 1, S]);
+    using pads = ops.fromInt32(this.leftPad, [B, 1, 1, 1]);
+    using causal = ops.lessEqual(keys, position);
+    using valid = ops.greaterEqual(keys, pads);
+    const mask = ops.logicalAnd(causal, valid);
+    if (window === null) return { mode: "array", arr: mask };
+    using windowSize = ops.fromInt32([window], []);
+    using limit = ops.add(keys, windowSize);
+    using within = ops.less(position, limit);
+    try { return { mode: "array", arr: ops.logicalAnd(mask, within) }; }
+    finally { mask.dispose(); }
+  }
+  #replace(planes: MlxArray[], offsets: number[], pads: number[]): void {
+    const previous = this.planes;
+    this.planes = planes; this.rowOffsets = offsets; this.leftPad = pads;
+    this.#positionsChanged(); disposeResources(previous);
+  }
+  append(inputs: readonly MlxArray[]): void {
+    const B = this.rowOffsets.length, N = inputs[0]!.shape[2]!, ends = this.#ends();
+    const needed = Math.max(...ends) + N, capacity = Math.ceil(needed / 256) * 256;
+    const common = ends.every(end => end === ends[0]);
+    using positions = common ? null : ops.fromInt32(ends.flatMap(end =>
+      Array.from({ length: N }, (_, position) => end + position)), [B, 1, N, 1]);
+    const next: MlxArray[] = [];
+    try {
+      for (const [index, input] of inputs.entries()) {
+        const old = this.planes[index];
+        let base: MlxArray;
+        if (!old) base = ops.zeros([B, input.shape[1]!, capacity, input.shape[3]!], input.dtype);
+        else if (old.shape[2]! >= needed) base = old;
+        else {
+          using padding = ops.zeros([B, old.shape[1]!, capacity - old.shape[2]!, old.shape[3]!], old.dtype);
+          base = ops.concatAxis([old, padding], 2);
+        }
+        try {
+          const updated = common
+            ? ops.sliceUpdate(base, input, [0, 0, ends[0]!, 0], [B, base.shape[1]!, ends[0]! + N, base.shape[3]!])
+            : ops.putAlongAxis(base, positions!, input, 2);
+          next.push(updated);
+        } finally { if (base !== old) base.dispose(); }
+      }
+    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    this.#replace(next, this.rowOffsets.map(offset => offset + N), this.leftPad);
+  }
+  mergeRows(rows: readonly KvTensorRowView[]): void {
+    const next: MlxArray[] = [], held: MlxArray[] = [];
+    let offsets: number[], pads: number[];
+    try {
+      if (rows.length === 1) {
+        const row = rows[0]!;
+        for (const plane of row.planes) next.push(ops.copyOf(plane));
+        offsets = [...row.rowOffsets]; pads = [...row.leftPad];
+      } else {
+        const widths = rows.map(row => Math.max(0, ...row.rowOffsets.map((offset, i) => offset + row.leftPad[i]!)));
+        const width = Math.max(0, ...widths);
+        if (width) {
+          const prototype = rows.find((_, index) => widths[index]! > 0)!;
+          for (const [field, geometry] of prototype.planes.entries()) {
+            const pieces = rows.map((row, index) => {
+              const active = widths[index]!;
+              if (!active) {
+                const empty = ops.zeros([row.rowOffsets.length, geometry.shape[1]!, width, geometry.shape[3]!], geometry.dtype);
+                held.push(empty); return empty;
+              }
+              const plane = row.planes[field]!;
+              const view = plane.slice([0, 0, 0, 0], [row.rowOffsets.length, plane.shape[1]!, active, plane.shape[3]!]);
+              held.push(view);
+              if (active === width) return view;
+              using padding = ops.zeros([row.rowOffsets.length, plane.shape[1]!, width - active, plane.shape[3]!], plane.dtype);
+              const padded = ops.concatAxis([padding, view], 2); held.push(padded); return padded;
+            });
+            next.push(ops.concatAxis(pieces, 0));
+          }
+        }
+        offsets = rows.flatMap(row => [...row.rowOffsets]);
+        pads = rows.flatMap((row, index) => row.leftPad.map(pad => pad + width - widths[index]!));
+      }
+    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    finally { disposeResources(held); }
+    this.#replace(next, offsets, pads);
+  }
+  /** Preserve caller-owned physical positions when changing representation.
+   * Live tensor bytes are copied; padding never participates in a codec. */
+  alignRows(leftPad: readonly number[]): void {
+    if (leftPad.every((pad, row) => pad === this.leftPad[row])) return;
+    const width = Math.max(0, ...this.rowOffsets.map((offset, row) => offset + leftPad[row]!));
+    const rows: MlxArray[][] = [], planes: MlxArray[] = [];
+    try {
+      for (let row = 0; row < this.rowOffsets.length; row++) rows.push(this.extractRow(row));
+      for (let field = 0; field < this.planes.length; field++) {
+        const parts: MlxArray[] = [];
+        try {
+          for (let row = 0; row < rows.length; row++) {
+            const geometry = this.planes[field]!, live = rows[row]![field];
+            using before = ops.zeros([1, geometry.shape[1]!, leftPad[row]!, geometry.shape[3]!], geometry.dtype);
+            using after = ops.zeros([1, geometry.shape[1]!, width - leftPad[row]! - this.rowOffsets[row]!, geometry.shape[3]!], geometry.dtype);
+            parts.push(ops.concatAxis(live ? [before, live, after] : [before, after], 2));
+          }
+          planes.push(ops.concatAxis(parts, 0));
+        } finally { disposeResources(parts); }
+      }
+    } catch (error) { return cleanupFailure(error, () => disposeResources(planes)); }
+    finally { disposeResources(rows.flat()); }
+    this.#replace(planes, [...this.rowOffsets], [...leftPad]);
+  }
+  filterRows(keep: readonly number[]): void {
+    if (!keep.length) { this.dispose(); return; }
+    const next: MlxArray[] = [];
+    using indices = this.planes.length ? ops.fromInt32([...keep], [keep.length]) : null;
+    const pads = keep.map(row => this.leftPad[row]!);
+    const sharedPadding = Math.min(...pads);
+    try {
+      for (const plane of this.planes) {
+        const selected = ops.takeAxis(plane, indices!, 0);
+        if (!sharedPadding) next.push(selected);
+        else {
+          try { next.push(selected.slice([0, 0, sharedPadding, 0], [...selected.shape])); }
+          finally { selected.dispose(); }
+        }
+      }
+    }
+    catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    this.#replace(next, keep.map(row => this.rowOffsets[row]!), pads.map(pad => pad - sharedPadding));
+    this.#padding.filter(keep);
+  }
+  extractRow(row: number): MlxArray[] {
+    const start = this.leftPad[row]!, count = this.rowOffsets[row]!, result: MlxArray[] = [];
+    if (!count) return result;
+    try {
+      for (const plane of this.planes) {
+        using view = plane.slice([row, 0, start, 0], [row + 1, plane.shape[1]!, start + count, plane.shape[3]!]);
+        result.push(materializeCopy(view));
+      }
+      return result;
+    } catch (error) { return cleanupFailure(error, () => disposeResources(result)); }
+  }
+  dispose(): void {
+    const previous = this.planes;
+    this.planes = []; this.rowOffsets = []; this.leftPad = []; this.#beforeRound = undefined; this.#padding.clear();
+    this.#positionsChanged(); disposeResources(previous);
+  }
+}

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;
+  }
+}