mlx-bun 0.3.0 → 0.4.0

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

@@ -0,0 +1,67 @@
+import type { Cache } from "../../model/gemma4-base";
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import { legacyCacheCodecs, type CacheCodecProvider, type CacheHeaderEntry,
+  type SnapshotContext } from "../../kv-store";
+import { disposeResources } from "../../engine/resources";
+import type { CheckpointAttachment } from "./checkpoint-state";
+type TensorSlot = SnapshotContext["slots"][number];
+
+/** Reuse the backend cache codecs for method-owned companion state. Tensor
+ * slots address attachment indices here; the persistence tier supplies the
+ * outer file offsets and checksums. No host serialization of tensors occurs. */
+export function captureCacheAttachment(schema: string, caches: readonly Cache[],
+  metadata: CheckpointAttachment["metadata"], codecs: CacheCodecProvider = legacyCacheCodecs): CheckpointAttachment {
+  const tensors: MlxArray[] = [], headers: Array<CacheHeaderEntry | null> = [];
+  try {
+    for (const cache of caches) {
+      if (!cache.offset) { headers.push(null); continue; }
+      const slots: TensorSlot[] = [];
+      const liveSlice = (array: MlxArray, length: number, axis = 2) => {
+        const end = [...array.shape]; end[axis] = length;
+        return array.slice(array.shape.map(() => 0), end);
+      };
+      const push = (array: MlxArray, owned: boolean) => {
+        const value = owned ? array : ops.copyOf(array);
+        const off = tensors.length; tensors.push(value);
+        slots.push({ off, bytes: value.nbytes, shape: [...value.shape], dtype: value.dtype, hash: "" });
+      };
+      const context: SnapshotContext = { slots, push, liveSlice,
+        liveMlaSlice: (array, length) => liveSlice(array, length, 1),
+        pushTriple(triple, length) {
+          for (const array of [triple.packed, triple.scales, triple.biases])
+            push(length === null ? array : liveSlice(array, length), length !== null);
+        },
+      };
+      const entry = codecs.forCache(cache).snapshot(cache, context);
+      if (cache.minimumReusableOffset !== undefined) entry.minimumReusableOffset = cache.minimumReusableOffset;
+      headers.push(entry);
+    }
+    return { schema, metadata: { ...metadata, cacheCodecs: codecs.id, cacheHeaders: JSON.stringify(headers) }, tensors };
+  } catch (error) { disposeResources(tensors); throw error; }
+}
+
+/** Returned caches own immutable aliases. The caller keeps an external SSD
+ * lease until its normal state materialization boundary has completed. */
+export function restoreCacheAttachment(attachment: CheckpointAttachment, makeCache: () => Cache[],
+  codecs: CacheCodecProvider = legacyCacheCodecs): Cache[] {
+  if (attachment.metadata.cacheCodecs !== codecs.id) throw new Error("companion cache codec mismatch");
+  const headers = JSON.parse(attachment.metadata.cacheHeaders as string) as Array<CacheHeaderEntry | null>;
+  const caches = makeCache(), pending: MlxArray[] = [];
+  const arr = (slot: TensorSlot) => {
+    const value = ops.copyOf(attachment.tensors[slot.off]!); pending.push(value); return value;
+  };
+  try {
+    for (const [index, entry] of headers.entries()) {
+      if (!entry) continue;
+      const cache = codecs.forHeader(entry).load(entry, { path: attachment.schema, arr, grownArr: arr,
+        triple: (slots, at) => ({ packed: arr(slots[at]!), scales: arr(slots[at + 1]!), biases: arr(slots[at + 2]!) }),
+      });
+      pending.length = 0;
+      cache.minimumReusableOffset = entry.minimumReusableOffset ?? 0;
+      const previous = caches[index]; caches[index] = cache;
+      previous?.dispose();
+    }
+    return caches;
+  } catch (error) { disposeResources([...pending, ...caches]); throw error; }
+}

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

@@ -0,0 +1,31 @@
+import { TurboQuantKVCache, RotatingKVCache, RotatingQuantizedKVCache, isBatchableCache, type BatchableCache, type Cache } from "../../model/gemma4-base";
+import { DelayedRotatingQuantizedKVCache } from "../../model/delayed-rotating-quantized-kv";
+import { RotatingAffineLayout } from "../../model/rotating-kv-layout";
+import { BatchedTurboQuantKVCache } from "../../model/batched-turboquant-kv";
+import { targetRowLayoutFactory, type TargetRowLayout } from "./target-layout-capability";
+import { PagedKVCache } from "../../lab/paged-kv/paged-kv";
+import { PagedKvRows } from "../../lab/paged-kv/paged-kv-rows";
+
+/** Bind encoded state to its row layout. Execution owners only move rows. */
+export function ownedCacheLayout(cache: Cache): BatchableCache | undefined {
+  if (cache instanceof PagedKVCache) return new PagedKvRows(cache.capacityTokens, cache.blockSize);
+  if (isBatchableCache(cache)) return cache.makeEmptyBatch();
+  if (cache instanceof TurboQuantKVCache) return new BatchedTurboQuantKVCache(cache.kBits, cache.vBits);
+  return undefined;
+}
+
+/** Prefill reuses decode's layouts; shape selection belongs to storage. */
+export function prefillCacheLayout(cache: Cache): BatchableCache {
+  if (cache instanceof PagedKVCache) return new PagedKvRows(cache.capacityTokens, cache.blockSize);
+  if (isBatchableCache(cache)) return cache.makeEmptyBatch();
+  if (cache instanceof RotatingKVCache) return new DelayedRotatingQuantizedKVCache(cache.maxSize, 64, 4, Infinity, () => {});
+  if (cache instanceof RotatingQuantizedKVCache) return new RotatingAffineLayout(cache.maxSize, cache.groupSize, cache.bits);
+  return targetCacheLayout(cache);
+}
+
+/** Target transaction layouts retain each source cache's numerical codec. */
+export function targetCacheLayout(cache: Cache): TargetRowLayout {
+  const make = targetRowLayoutFactory(cache);
+  if (!make) throw new Error(`No target row layout for ${cache.signature()}`);
+  return make();
+}

package/src/backends/mlx/checkpoint-state.ts

@@ -0,0 +1,53 @@
+import type { Cache } from "../../model/gemma4-base";
+import * as ops from "../../mlx/ops";
+import { leaseCacheStates } from "./state-views";
+import { withResource } from "../../engine/resources";
+import type { MlxArray } from "../../mlx/array";
+import { disposeResources, cleanupFailure } from "../../engine/resources";
+
+/** Generation-method state accompanying the target cache. The method owns
+ * its schema and alignment; storage owns these immutable tensor snapshots.
+ * Attachments require an exact prefix boundary and are never trimmed. */
+export interface CheckpointAttachment {
+  schema: string;
+  metadata: Record<string, string | number | boolean>;
+  tensors: MlxArray[];
+}
+
+export function attachmentBytes(attachments: readonly CheckpointAttachment[] = []): number {
+  return attachments.reduce((sum, entry) =>
+    sum + entry.tensors.reduce((bytes, tensor) => bytes + tensor.nbytes, 0), 0);
+}
+
+export function disposeAttachments(attachments: readonly CheckpointAttachment[] = []): void {
+  disposeResources(attachments.flatMap((entry) => entry.tensors));
+}
+
+export function cloneAttachments(attachments: readonly CheckpointAttachment[] = []): CheckpointAttachment[] {
+  const held: MlxArray[] = [];
+  try {
+    return attachments.map((entry) => ({
+      schema: entry.schema, metadata: { ...entry.metadata },
+      tensors: entry.tensors.map((tensor) => {
+        const copy = tensor.slice(tensor.shape.map(() => 0), [...tensor.shape]);
+        held.push(copy);
+        return copy;
+      }),
+    }));
+  } catch (error) { return cleanupFailure(error, () => disposeResources(held)); }
+}
+
+/** Shared cache port for MLX execution methods and their companion state. */
+export type MlxPrefixCache = import("../../contracts/prefix-cache").PrefixCache<
+  import("../../model/gemma4").Cache[], CheckpointAttachment[]
+>;
+
+/** Resolve snapshot copies before storage retains them. A deferred compact
+ * row still owns its full parent batch; persistence may wait until idle.
+ * This only settles GPU state and performs no serialization or disk I/O. */
+export function materializeCheckpoint(caches: Cache[], attachments: readonly CheckpointAttachment[] = []): void {
+  withResource(leaseCacheStates(caches), state => {
+    const tensors = [...state, ...attachments.flatMap(attachment => attachment.tensors)];
+    if (tensors.length) ops.evalAll(tensors);
+  });
+}

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

@@ -0,0 +1,106 @@
+import { SpillQueue, type SpillItem } from "../../kv-store";
+import { cacheBytes } from "../../prompt-cache";
+import { disposeResources } from "../../engine/resources";
+import type { ContinuationStore, OrdinaryContinuationState } from "./continuation";
+
+type Metadata = Parameters<ContinuationStore["storeGenerationCheckpoint"]>[2];
+interface Attempt { key: string; completed: boolean; pending: number; failed: boolean; version: number; retired: boolean }
+interface Job { attempt: Attempt; metadata?: Metadata; cleanup?: number }
+
+/** One coordinator per checkpoint store. Queue owns snapshots through disk
+ * completion; per-key attempts order cleanup before subsequent writes. */
+export class ContinuationPersistence {
+  readonly #attempts = new Map<string, Attempt>();
+  readonly #jobs = new WeakMap<SpillItem["caches"], Job>();
+  readonly #queue: SpillQueue;
+  readonly #cleanupFailures = new Map<string, Set<number>>();
+  readonly #cleanupPending = new Map<string, Set<number>>();
+  readonly #cleanupSucceeded = new Map<string, number>();
+  #nextBarrier = 0;
+  constructor(readonly store: ContinuationStore, options: {
+    maxBytes: number;
+    /** All blocking tensor serialization runs through the owner's idle gate. */
+    runStep: <T>(step: () => T) => Promise<T>;
+  }) {
+    this.#queue = new SpillQueue(options.maxBytes, cacheBytes, async item => {
+      const job = this.#jobs.get(item.caches)!;
+      if (job.cleanup) {
+        await options.runStep(() => store.removeGenerationCheckpoints(job.attempt.key));
+        return true;
+      }
+      const current = () => this.#attempts.get(job.attempt.key) === job.attempt && !job.attempt.completed;
+      if (!current()) return true;
+      const stored = await store.storeGenerationCheckpoint(item.tokens, item.caches, job.metadata!, options.runStep);
+      // No newer queued write can run until this callback returns. An old
+      // in-flight rename can therefore be removed without deleting new work.
+      if (!current()) await options.runStep(() => store.removeGenerationCheckpoints(job.attempt.key));
+      return stored;
+    }, disposeResources, item => !this.#jobs.get(item.caches)?.cleanup);
+  }
+  /** A completed response is ineligible immediately, before queued unlink. */
+  canRestore(key: string): boolean { return !this.#cleanupPending.has(key) && !this.#cleanupFailures.has(key); }
+  begin(key: string): Attempt {
+    const attempt = { key, completed: false, pending: 0, failed: false, version: 0, retired: false };
+    this.#attempts.set(key, attempt);
+    this.#queue.cancelWhere(item => !this.#jobs.get(item.caches)?.cleanup && this.#jobs.get(item.caches)?.attempt.key === key);
+    return attempt;
+  }
+  /** Takes ownership even when a completed/superseded attempt declines it. */
+  enqueue(attempt: Attempt, state: Omit<OrdinaryContinuationState, "seed">, metadata: Metadata): void {
+    if (attempt.completed || this.#attempts.get(attempt.key) !== attempt) {
+      disposeResources(state.caches); return;
+    }
+    this.#jobs.set(state.caches, { attempt, metadata: { ...metadata } });
+    attempt.pending++;
+    const version = ++attempt.version;
+    let completion: Promise<boolean>;
+    try { completion = this.#queue.enqueue({ caches: state.caches, tokens: [...state.cacheTokens], ns: attempt.key }); }
+    catch (error) { attempt.pending--; disposeResources(state.caches); throw error; }
+    void completion.then(stored => {
+      attempt.pending--;
+      if (version === attempt.version && !attempt.completed) attempt.failed = !stored;
+      this.#releaseSettled(attempt);
+    });
+  }
+  complete(attempt: Attempt): void {
+    attempt.completed = true;
+    this.#queue.cancelWhere(item => !this.#jobs.get(item.caches)?.cleanup && this.#jobs.get(item.caches)?.attempt === attempt);
+    if (this.#attempts.get(attempt.key) === attempt) {
+      this.#attempts.delete(attempt.key);
+      const barrier = ++this.#nextBarrier;
+      const pending = this.#cleanupPending.get(attempt.key) ?? new Set<number>();
+      pending.add(barrier); this.#cleanupPending.set(attempt.key, pending);
+      const caches: SpillItem["caches"] = [];
+      this.#jobs.set(caches, { attempt, cleanup: barrier });
+      void this.#queue.enqueue({ caches, tokens: [], ns: attempt.key }).then(stored => {
+        const key = attempt.key;
+        const outstanding = this.#cleanupPending.get(key)!;
+        outstanding.delete(barrier);
+        if (!outstanding.size) this.#cleanupPending.delete(key);
+        const failures = this.#cleanupFailures.get(key) ?? new Set<number>();
+        if (stored) {
+          this.#cleanupSucceeded.set(key, Math.max(barrier, this.#cleanupSucceeded.get(key) ?? 0));
+          // A successful unlink covers older failures, never a newer barrier.
+          for (const failed of failures) if (failed <= barrier) failures.delete(failed);
+        } else if (barrier > (this.#cleanupSucceeded.get(key) ?? 0)) failures.add(barrier);
+        if (failures.size) this.#cleanupFailures.set(key, failures);
+        else this.#cleanupFailures.delete(key);
+        if (!outstanding.size) this.#cleanupSucceeded.delete(key);
+      });
+    }
+  }
+  /** Request cancellation/disposal preserves queued and durable state. */
+  release(attempt: Attempt): void { attempt.retired = true; this.#releaseSettled(attempt); }
+  #releaseSettled(attempt: Attempt): void {
+    if (attempt.retired && attempt.pending === 0 && !attempt.failed && this.#attempts.get(attempt.key) === attempt)
+      this.#attempts.delete(attempt.key);
+  }
+  get stats() { return { pendingCount: this.#queue.pendingCount, pendingBytes: this.#queue.pendingBytes,
+    dropped: this.#queue.droppedCount, failed: this.#queue.failedCount }; }
+  async flush(): Promise<{ durable: boolean; pendingBytes: number; failed: number }> {
+    await this.#queue.drain();
+    const active = [...this.#attempts.values()].filter(attempt => !attempt.completed);
+    return { durable: this.#cleanupFailures.size === 0 && active.every(attempt => attempt.pending === 0 && !attempt.failed),
+      pendingBytes: this.#queue.pendingBytes, failed: this.#cleanupFailures.size + active.filter(attempt => attempt.failed).length };
+  }
+}

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

@@ -0,0 +1,43 @@
+import { snapshotGenerationPolicy } from "./request-policy";
+import type { GenerateOptions } from "../../generate";
+import type { ResolvedExecution } from "../../contracts/execution";
+import { makeStepSampler } from "../../sampler";
+import { createRowSampling } from "./row-sampling";
+import { bindContinuationPolicy, type OrdinaryContinuation } from "./continuation";
+
+/** Optional policy and sampler recovery for the existing ordinary driver.
+ * KV layout and delayed conversion remain owned by cache maintenance and persistence. */
+export function createOrdinaryContinuationRequest(input: Omit<Parameters<typeof bindContinuationPolicy>[0], "namespace"> & {
+  interval: number;
+  persistence: import("./continuation-persistence").ContinuationPersistence;
+  onToken: Parameters<typeof createRowSampling>[1];
+}) {
+  const options: GenerateOptions = snapshotGenerationPolicy(input.options);
+  if (!Number.isSafeInteger(input.interval) || input.interval < 1) throw new Error("continuation interval must be positive");
+  if (options.grammar || options.logprobs || options.topLogprobs || options.fill || options.pagedKv ||
+      input.execution.method !== "autoregressive")
+    throw new Error("ordinary continuation requires replayable ordinary sampling");
+  const make = (seed = options.seed, history: readonly number[] = input.prompt) => createRowSampling(makeStepSampler(
+    { ...options, seed }, { tokenRepresentation: "device", grammarWait: "external", historyUpdate: "after-sample",
+      initialHistory: [...history] }), input.onToken);
+  let sampling = make();
+  let policy: ReturnType<typeof bindContinuationPolicy> | undefined;
+  const continuation: OrdinaryContinuation = {
+    interval: input.interval,
+    restore(namespace) { policy = bindContinuationPolicy({ ...input, options, namespace }); return policy.restore(); },
+    resumeSampling(state) {
+      const next = make(state.seed, [...state.cacheTokens, state.pendingToken]);
+      const previous = sampling; sampling = next; previous.dispose();
+    },
+    captureOwned: state => policy!.captureOwned(state),
+    complete: () => policy?.complete(),
+  };
+  return {
+    continuation,
+    sample: ((...args) => sampling.sample(...args)) as ReturnType<typeof createRowSampling>["sample"],
+    // History-free greedy sampling needs no recovery work in the fast path.
+    plainGreedy: sampling.plainGreedy,
+    onToken: input.onToken,
+    dispose: () => { policy?.release(); sampling.dispose(); },
+  };
+}

package/src/backends/mlx/continuation.ts

@@ -0,0 +1,77 @@
+import type { ContinuationPersistence } from "./continuation-persistence";
+import type { Cache } from "../../model/gemma4-base";
+import type { GenerateOptions } from "../../generate";
+import type { ResolvedExecution } from "../../contracts/execution";
+import type { SsdCacheStore } from "../../ssd-cache";
+import { generationCheckpointKey } from "../../serve/checkpoint-identity";
+import { disposeResources } from "../../engine/resources";
+
+export interface OrdinaryContinuationState {
+  caches: Cache[];
+  cacheTokens: number[];
+  generatedTokens: number;
+  pendingToken: number;
+  seed: number;
+}
+/** Request-owned policy. Restore and captureOwned transfer cache ownership.
+ * No store or key knowledge enters the driver. */
+export interface OrdinaryContinuation {
+  readonly interval: number;
+  restore(namespace: string): OrdinaryContinuationState | null;
+  resumeSampling(state: Omit<OrdinaryContinuationState, "caches">): void;
+  captureOwned(state: Omit<OrdinaryContinuationState, "seed">): void;
+  complete(): void;
+}
+export type ContinuationStore = Pick<SsdCacheStore, "findGenerationCheckpoint" | "restore" |
+  "storeGenerationCheckpoint" | "removeGenerationCheckpoints">;
+
+/** Shared persistence policy for compatibility and grouped ordinary drivers. */
+export function bindContinuationPolicy(input: {
+  store: ContinuationStore;
+  restore: (entry: NonNullable<ReturnType<ContinuationStore["findGenerationCheckpoint"]>>) => ReturnType<ContinuationStore["restore"]>;
+  prompt: number[]; options: GenerateOptions; execution: ResolvedExecution;
+  identity: unknown; namespace: string;
+  persistence?: ContinuationPersistence;
+}) {
+  const { store, options, prompt, namespace } = input;
+  const key = generationCheckpointKey(prompt, options, namespace, input.execution, input.identity);
+  let seed = options.seed ?? 0;
+  const attempt = input.persistence?.begin(key);
+  return {
+    key,
+    restore(): OrdinaryContinuationState | null {
+      if (input.persistence && !input.persistence.canRestore(key)) return null;
+      const entry = store.findGenerationCheckpoint(prompt, key, namespace);
+      const restored = entry ? input.restore(entry) : null;
+      if (!restored) return null;
+      const metadata = restored.header.generationCheckpoint;
+      if (!metadata || metadata.originalPromptTokens !== prompt.length ||
+          metadata.generatedTokens !== restored.tokens.length - prompt.length) {
+        disposeResources(restored.caches);
+        throw new Error("generation checkpoint has inconsistent continuation metadata");
+      }
+      seed = metadata.seed;
+      return { caches: restored.caches, cacheTokens: restored.tokens,
+        generatedTokens: metadata.generatedTokens, pendingToken: metadata.pendingToken, seed };
+    },
+    async capture(state: Omit<OrdinaryContinuationState, "seed">): Promise<void> {
+      const stored = await store.storeGenerationCheckpoint(state.cacheTokens, state.caches, {
+        key, cacheNs: namespace, originalPromptTokens: prompt.length,
+        generatedTokens: state.generatedTokens, pendingToken: state.pendingToken,
+        seed, seedWasExplicit: options.seedWasExplicit === true,
+      });
+      if (stored) console.log(`[generation-checkpoint] saved ${state.generatedTokens} emitted tokens`);
+    },
+    captureOwned(state: Omit<OrdinaryContinuationState, "seed">): void {
+      if (!input.persistence || !attempt) { disposeResources(state.caches); throw new Error("owned continuation requires persistence coordinator"); }
+      input.persistence.enqueue(attempt, state, { key, cacheNs: namespace, originalPromptTokens: prompt.length,
+        generatedTokens: state.generatedTokens, pendingToken: state.pendingToken,
+        seed, seedWasExplicit: options.seedWasExplicit === true });
+    },
+    release() { if (input.persistence && attempt) input.persistence.release(attempt); },
+    complete() {
+      if (input.persistence && attempt) input.persistence.complete(attempt);
+      else store.removeGenerationCheckpoints(key);
+    },
+  };
+}

package/src/backends/mlx/draft-target.ts

@@ -1,9 +1,10 @@
+import { readAssistantDonors } from "./assistant-target";
 import type { RuntimeModel } from "../../model/factory";
 import { Gemma4Model, type Cache } from "../../model/gemma4";
 import { Qwen35Model } from "../../model/qwen3_5";
 import type { MlxArray } from "../../mlx/array";
 import * as ops from "../../mlx/ops";
-import type { TargetView, AssistantTarget } from "../../spec/source";
+import type { TargetView } from "../../spec/source";
 
 /** Class checks and live cache layout stay at the compatibility boundary.
  * A replacement graph supplies the same ports without becoming these classes. */
@@ -14,27 +15,11 @@ export function bindLegacyDraftTarget(model: RuntimeModel, caches: Cache[]): Tar
       if (model.layers[i]!.layerType === "sliding_attention") sliding = i;
       else full = i;
     }
-    const view = (index: number): [MlxArray, MlxArray] => {
-      const cache = caches[index];
-      if (!cache || !("temporalView" in cache) || typeof cache.temporalView !== "function")
-        throw new Error("assistant target has no plain donor view");
-      return cache.temporalView();
-    };
-    const assistant: AssistantTarget | undefined = sliding >= 0 && full >= 0 ? {
-      position: () => caches[0]!.offset - 1,
-      embedScaled(token) {
-        const ids = ops.fromInt32([token], [1, 1]);
-        let embedded: MlxArray | undefined;
-        try { embedded = model.embed.encode(ids); return ops.mulScalar(embedded, model.embedScale); }
-        finally { ids.dispose(); embedded?.dispose(); }
-      },
-      readDonors() {
-        const slidingView = view(sliding);
-        try { return { sliding: slidingView, full: view(full) }; }
-        catch (error) { for (const array of slidingView) array.dispose(); throw error; }
-      },
-    } : undefined;
-    return Object.freeze({ identity: model, assistant, gemmaTaps: Object.freeze({
+    return Object.freeze({ identity: model, assistantRows: sliding >= 0 && full >= 0 ? {
+      hiddenSize: model.config.text.hiddenSize,
+      embed(ids: MlxArray) { using embedded = model.embed.encode(ids); return ops.mulScalar(embedded, model.embedScale); },
+      readDonors: () => readAssistantDonors(caches[sliding]!, caches[full]!),
+    } : undefined, gemmaTaps: Object.freeze({
       layerCount: model.layers.length,
       projection: Object.freeze({ embed: model.embed, logitsFromHidden: model.logitsFromHidden.bind(model) }),
     }) });

package/src/backends/mlx/full-kv-row-append.ts

@@ -0,0 +1,32 @@
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import type { Cache } from "../../model/gemma4-base";
+import { fullRowPadding } from "../../model/full-prefill-row";
+import { disposeResources } from "../../engine/resources";
+
+/** Assemble full-attention tensors at the model's B while storage retains
+ * per-row formats and positions. Returned arrays own their captured state. */
+export function appendFullKvRows(rows: readonly Cache[], leftPad: readonly number[], width: number,
+  k: MlxArray, v: MlxArray, deferredValues = false): [MlxArray, MlxArray] {
+  const parts: [MlxArray, MlxArray][] = [], padded: MlxArray[] = [];
+  try {
+    for (const [index, row] of rows.entries()) {
+      using kr = k.slice([index, 0, 0, 0], [index + 1, k.shape[1]!, k.shape[2]!, k.shape[3]!]);
+      using vr = v.slice([index, 0, 0, 0], [index + 1, v.shape[1]!, v.shape[2]!, v.shape[3]!]);
+      parts.push(deferredValues && row.rotatedValueAttention
+        ? row.rotatedValueAttention.updateAndFetchDeferredV(kr, vr) : row.updateAndFetch(kr, vr));
+    }
+    const combine = (field: 0 | 1) => {
+      const views = parts.map((part, row) => {
+        const a = part[field], before = leftPad[row]! - fullRowPadding(rows[row]!), after = width - before - a.shape[2]!;
+        if (!before && !after) return a;
+        using prefix = ops.zeros([1, a.shape[1]!, before, a.shape[3]!], a.dtype);
+        using suffix = ops.zeros([1, a.shape[1]!, after, a.shape[3]!], a.dtype);
+        const result = ops.concatAxis([prefix, a, suffix], 2); padded.push(result); return result;
+      });
+      return views.length === 1 ? ops.contiguous(views[0]!) : ops.concatAxis(views, 0);
+    };
+    const keys = combine(0);
+    try { return [keys, combine(1)]; } catch (error) { keys.dispose(); throw error; }
+  } finally { disposeResources([...parts.flat(), ...padded]); }
+}

package/src/backends/mlx/full-kv-row-donor.ts

@@ -0,0 +1,54 @@
+import { captureKvDonorAttention, combineKvDonorAttention } from "../../model/kv-attention-view";
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import type { Cache, KvDonorRows } from "../../model/gemma4-base";
+import { fullRowInner, fullRowPadding } from "../../model/full-prefill-row";
+import { disposeResources } from "../../engine/resources";
+
+/** Read full-attention rows in their existing physical columns. Valid token
+ * coverage excludes prefill padding; no append, conversion or compaction. */
+export function captureFullKvDonorRows(rows: readonly Cache[], leftPad: readonly number[], width: number): KvDonorRows {
+  const parts: KvDonorRows[] = [], padded: MlxArray[] = [];
+  try {
+    for (const row of rows) parts.push(fullRowInner(row).captureDonorRows!());
+    const combine = (field: "keys" | "values") => {
+      const views = parts.map((part, index) => {
+        const a = part[field], before = leftPad[index]! - fullRowPadding(rows[index]!), after = width - before - a.shape[2]!;
+        if (!before && !after) return a;
+        using prefix = ops.zeros([1, a.shape[1]!, before, a.shape[3]!], a.dtype);
+        using suffix = ops.zeros([1, a.shape[1]!, after, a.shape[3]!], a.dtype);
+        const result = ops.concatAxis([prefix, a, suffix], 2); padded.push(result); return result;
+      });
+      return views.length === 1 ? ops.contiguous(views[0]!) : ops.concatAxis(views, 0);
+    };
+    const offsets = rows.map(row => row.offset), starts = [...leftPad];
+    const keys = combine("keys");
+    try { return { keys, values: combine("values"), offsets, starts, ends: offsets.map((offset, row) => starts[row]! + offset) }; }
+    catch (error) { keys.dispose(); throw error; }
+  } finally { disposeResources([...parts.flatMap(part => [part.keys, part.values]), ...padded]); }
+}
+
+
+/** Capture full rows in global columns while each encoded view retains its
+ * own physical width and numerical attention operation. */
+export function captureFullKvDonorAttention(rows: readonly Cache[], leftPad: readonly number[], width: number): import("../../model/gemma4-base").KvDonorAttention {
+  const captured: import("../../model/gemma4-base").KvDonorAttention[] = [];
+  try {
+    for (const [index, row] of rows.entries()) {
+      const view = captureKvDonorAttention(fullRowInner(row));
+      const pad = leftPad[index]! - fullRowPadding(row);
+      captured.push({ width, dtype: view.dtype, offsets: [row.offset],
+        starts: [leftPad[index]!], ends: [leftPad[index]! + row.offset],
+        attend(q, scale, mask) {
+          if (!mask.arr) return view.attend(q, scale, mask);
+          const from = mask.arr.shape.map(() => 0), to = [...mask.arr.shape];
+          from[from.length - 1] = pad; to[to.length - 1] = pad + view.width;
+          using arr = mask.arr.slice(from, to);
+          return view.attend(q, scale, { mode: mask.mode, arr });
+        },
+        dispose: () => view.dispose(),
+      });
+    }
+    return combineKvDonorAttention(captured);
+  } catch (error) { disposeResources(captured); throw error; }
+}

package/src/backends/mlx/full-transitioning-kv-rows.ts

@@ -0,0 +1,65 @@
+import { KVCache, type Cache, type PaddedPrefillCache, type PrefillPadding } from "../../model/gemma4-base";
+import { FullPrefillRow, fullRowInner } from "../../model/full-prefill-row";
+import { cloneKvCaches } from "../../kv-store";
+import { FullKvPositions, SpeculativeTransitioningKvRows, type KvRowTransition, type SpeculativeTransitionedKvLayout } from "./transitioning-kv-rows";
+
+type FullLayout = SpeculativeTransitionedKvLayout & PaddedPrefillCache & {
+  restorePrefillEnds(ends: readonly number[] | undefined): void;
+}
+
+class PrefillPositions extends FullKvPositions {
+  override sync(rows: readonly Cache[], preserve: boolean): void {
+    const offsets = rows.map(row => row instanceof FullPrefillRow ? row.positionOffset : row.offset);
+    this.mergeRows(preserve
+      ? [{ planes: [], rowOffsets: offsets, leftPad: [...this.leftPad] }]
+      : rows.map((row, index) => ({ planes: [], rowOffsets: [offsets[index]!],
+        leftPad: [row instanceof FullPrefillRow ? row.leftPadding : 0] })));
+  }
+}
+
+/** Full-attention padding composes with the existing precision transition
+ * and rollback lifecycle. Codec-specific attention stays in each adapter. */
+export abstract class FullTransitioningKvRows<Layout extends FullLayout>
+  extends SpeculativeTransitioningKvRows<Layout> implements PaddedPrefillCache {
+  constructor(transition: KvRowTransition<Layout>, row?: Cache) {
+    super({ ...transition,
+      converted: row => transition.converted(fullRowInner(row)),
+      extractRow: row => row instanceof FullPrefillRow ? row.extract() : cloneKvCaches([row])[0]!,
+      packRows(layout, rows) {
+        layout.mergeRows(rows.map(fullRowInner));
+        if (rows.some(row => row instanceof FullPrefillRow)) {
+          layout.preparePrefill({ lengths: rows.map(() => 0),
+            leftPadding: rows.map(row => row instanceof FullPrefillRow ? row.leftPadding : 0) });
+          const ends = rows.map(row => row instanceof FullPrefillRow ? row.padding.endOffsets?.[0] : undefined);
+          layout.restorePrefillEnds(ends.some(end => end !== undefined) ? ends.map(end => end ?? Infinity) : undefined);
+        }
+      },
+    }, row, new PrefillPositions());
+  }
+  preparePrefill(padding: PrefillPadding): void {
+    if (this.packed) { this.packed.preparePrefill(padding); return; }
+    if (!this.rows.length) this.rows = padding.lengths.map(() => new KVCache());
+    this.rows = this.rows.map(row => row instanceof FullPrefillRow ? row : new FullPrefillRow(row as KVCache));
+    for (const [index, row] of this.rows.entries()) (row as FullPrefillRow).preparePrefill({
+      lengths: [padding.lengths[index]!],
+      ...(padding.leftPadding ? { leftPadding: [padding.leftPadding[index]!] } : {}),
+      ...(padding.rightPadding ? { rightPadding: [padding.rightPadding[index]!] } : {}),
+    });
+    this.syncPositions();
+  }
+  finalizePrefill(): void {
+    if (this.packed) { this.packed.finalizePrefill(); return; }
+    for (const row of this.rows) if (row instanceof FullPrefillRow) row.finalizePrefill();
+    this.syncPositions();
+  }
+  override filterRows(keep: readonly number[]): void {
+    if (!this.packed && keep.length) {
+      const pads = [...this.leftPad], shared = Math.min(...keep.map(row => pads[row]!));
+      for (const index of keep) {
+        const row = this.rows[index]!;
+        if (row instanceof FullPrefillRow) row.removeLeftPadding(Math.max(0, row.leftPadding - (pads[index]! - shared)));
+      }
+    }
+    super.filterRows(keep);
+  }
+}