mlx-bun 0.3.0 → 0.4.0

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

@@ -1,3 +1,7 @@
+import { createOrdinaryContinuationRequest } from "./continuation-request";
+import type { MlxSerialServices } from "./serial-executor";
+import { bindPagedRequestState, 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,11 +9,15 @@ 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 { targetRowLayoutFactory } from "./target-layout-capability";
+import { bindSpeculativeGroupRequests } from "./speculative-group";
 import type { GenerateOptions } from "../../generate";
 import type { ExecutionRequirements, ResolvedExecution } from "../../contracts/execution";
 import { resolveExecution } from "../../engine/execution-plan";
@@ -20,52 +28,103 @@ export interface MlxBatchGroup extends Pick<MlxBatchExecutionGroup,
 /** A model implementation owns capability checks and the execution group.
  * Scheduling never inspects concrete model/cache classes. */
 export interface MlxGatewayBinding {
+  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;
+  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 adapterState = "loraState" in model ? model.loraState : undefined;
   return {
     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) => execution?.pagedKv ? bindPagedRequestState(model, options, capacity) : undefined,
+    methodRequest: (execution, options) => execution?.method === "speculative" ? speculative?.(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) {
       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,
+        groupedMethods: speculative ? ["autoregressive", "speculative"] : ["autoregressive"],
+        speculativeLogprobs: scheduling.continuous && !!speculative,
+        sharedSpeculativeAdapters: scheduling.continuous && !!speculative && !!adapterState &&
+          draft?.provider.grouped?.supportsTargetAdapters === true,
+        turboQuantBatch: scheduling.quantizedBatch,
+        speculativeTurboQuant: scheduling.continuous && !!speculative && !!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 && !!speculative && (options.kvBits === 4 || options.kvBits === 8 || !!options.kvConfig?.length)) ||
+          (!options.kvConfig?.length && model instanceof Qwen35Model &&
+            (options.kvBits === 4 || (scheduling.continuous && !!speculative && options.kvBits === 8)) &&
+            (options.quantizedKvStart === 0 || (scheduling.continuous && !!speculative)))
+        ),
       }, {
         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,133 @@
+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;
+  /** 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 maintain: KvMaintenance = (cache) => maybeTurboQuantizeKv(cache, scheme, start);
+    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);
+      }
+    };
+    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();
+    }
+  };
+  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/model-serving.ts

@@ -25,7 +25,7 @@ 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),
     restore: (store, entry) => store.restore(entry, model),

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

@@ -0,0 +1,89 @@
+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> {
+  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) ?? 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?.();
+          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) { 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),
+      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);
+          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?.(); },
+    });
+  }
+}