mlx-bun 0.3.0 → 0.5.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, cache.direct, cache.quantization);
+  if (isBatchableCache(cache)) return cache.makeEmptyBatch();
+  if (cache instanceof TurboQuantKVCache) return new BatchedTurboQuantKVCache(cache.kBits, cache.vBits, cache.fusedDecode);
+  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, cache.direct, cache.quantization);
+  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,107 @@
+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;
+    /** Optional owner-thread preparation hook; disk writes use the CPU worker. */
+    runStep?: <T>(step: () => T) => Promise<T>;
+  }) {
+    const runStep = options.runStep ?? (async <T>(step: () => T): Promise<T> => step());
+    this.#queue = new SpillQueue(options.maxBytes, cacheBytes, async item => {
+      const job = this.#jobs.get(item.caches)!;
+      if (job.cleanup) {
+        await 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 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/encoder-cache.ts

@@ -0,0 +1,25 @@
+import type { ObjectCache } from "../../contracts/object-cache";
+import { cloneAttachments, disposeAttachments, type CheckpointAttachment } from "./checkpoint-state";
+import type { MlxArray } from "../../mlx/array";
+
+/** Encoder-owned tensor schema over the cache's exact-object interface. */
+export class EncoderCache {
+  constructor(private readonly objects: ObjectCache<CheckpointAttachment[]>,
+    readonly identity: string) {}
+
+  async take(key: string): Promise<MlxArray | null> {
+    const lease = await this.objects.take(`encoder-v1:${this.identity}:${key}`);
+    if (!lease) return null;
+    try {
+      const tensor = lease.value[0]!.tensors[0]!;
+      return tensor.slice(tensor.shape.map(() => 0), [...tensor.shape]);
+    } finally { lease.dispose(); }
+  }
+
+  /** Borrow the result; storage retains its own immutable native view. */
+  put(key: string, features: MlxArray): void {
+    const owned = cloneAttachments([{ schema: "encoder-features-v1", metadata: {}, tensors: [features] }]);
+    try { this.objects.put(`encoder-v1:${this.identity}:${key}`, owned); }
+    catch (error) { disposeAttachments(owned); throw error; }
+  }
+}

package/src/backends/mlx/fill-append.ts

@@ -6,8 +6,8 @@ import { leaseCacheStates } from "./state-views";
 
 /** Append an already accepted span without intermediate vocabulary heads.
  * Chunking changes the graph shape, so callers must qualify its numerics for
- * their model. A verify round must use chunkSize=0: recurrent rollback records
- * one forward per round. The complete append precedes emission/cancellation
+ * their model. Verification uses its separate round operation because
+ * recurrent rollback records one forward per round. Append precedes emission/cancellation
  * boundaries, preserving the cache-covered prefix on a consumer break.
  */
 export async function appendFillHidden(
@@ -17,15 +17,29 @@ export async function appendFillHidden(
   chunkSize: (cache: readonly Cache[]) => number,
 ): Promise<MlxArray> {
   const input = ops.fromInt32(ids, [1, ids.length]);
+  try { return await appendHiddenRows(forward, cache, input, chunkSize); }
+  finally { input.dispose(); }
+}
+
+/** Append equal-length committed token rows [B,L]. The caller owns the input;
+ * the returned hidden rows belong to the caller. Model policy receives B so
+ * a larger cohort cannot silently cross its qualified projection geometry. */
+export async function appendHiddenRows(
+  forward: (ids: MlxArray, cache: Cache[]) => MlxArray | Promise<MlxArray>,
+  cache: Cache[],
+  input: MlxArray,
+  chunkSize: (cache: readonly Cache[], rows: number) => number,
+): Promise<MlxArray> {
+  const [rows, length] = input.shape as [number, number];
   const hidden: MlxArray[] = [];
   try {
-    for (let start = 0; start < ids.length;) {
-      const limit = chunkSize(cache);
+    for (let start = 0; start < length;) {
+      const limit = chunkSize(cache, rows);
       if (!Number.isSafeInteger(limit) || limit < 1)
         throw new Error("Model append chunk limit must be a positive integer");
-      const end = Math.min(ids.length, start + limit);
-      if (start === 0 && end === ids.length) return await forward(input, cache);
-      const part = input.slice([0, start], [1, end]);
+      const end = Math.min(length, start + limit);
+      if (start === 0 && end === length) return await forward(input, cache);
+      const part = input.slice([0, start], [rows, end]);
       try {
         const h = await forward(part, cache);
         hidden.push(h);
@@ -40,6 +54,5 @@ export async function appendFillHidden(
     return ops.concatAxis(hidden, 1);
   } finally {
     for (const h of hidden) h.dispose();
-    input.dispose();
   }
 }