mlx-bun 0.2.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/diffusion.ts

@@ -0,0 +1,39 @@
+import type { DenoisingGraph } from "../../inference/denoising";
+import type { MlxArray } from "../../mlx/array";
+import type { Cache } from "../../model/gemma4-base";
+import type { DiffusionGemmaModel } from "../../model/diffusion-gemma";
+import type { MlxModelMemory } from "./autoregressive";
+import { disposeResources } from "../../engine/resources";
+import { runtimeConfig, type RuntimeConfig } from "../../runtime-config";
+
+export interface MlxDenoisingBinding<State = Cache[]> {
+  readonly runtime?: RuntimeConfig;
+  readonly graph: DenoisingGraph<MlxArray, State>;
+  readonly memory: MlxModelMemory;
+  readonly adapters?: { active: string[] };
+}
+
+export function bindLegacyDenoisingModel(model: DiffusionGemmaModel): MlxDenoisingBinding {
+  return {
+    runtime: runtimeConfig(),
+    memory: model, adapters: model.loraState,
+    graph: {
+      descriptor: Object.freeze({ id: "legacy-diffusion-gemma", backend: "mlx",
+        graphAbi: "mlx-denoising-v1", stateAbi: "legacy-cache-array-v1", artifact: "legacy-resident-model" }),
+      vocabSize: model.config.text.vocabSize, canvasLength: model.canvasLength, embedScale: model.embedScale,
+      prefill: (ids, vision) => vision ? model.prefillVision(ids, vision) : model.prefill(ids),
+      extendPrefill: model.extendPrefill.bind(model),
+      decoderLogits: model.decoderLogits.bind(model),
+      dequantEmbedWeight: model.dequantEmbedWeight.bind(model),
+      softEmbeddings: model.softEmbeddings.bind(model),
+      closeState: disposeResources,
+    },
+  };
+}
+
+export function assertMlxDenoisingGraph<State>(graph: DenoisingGraph<MlxArray, State>): void {
+  const descriptor = graph.descriptor;
+  if (descriptor.backend !== "mlx" || descriptor.graphAbi !== "mlx-denoising-v1" ||
+      !descriptor.stateAbi)
+    throw new Error(`denoising graph ${descriptor.id} has an incompatible backend, graph, or state ABI`);
+}

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

@@ -0,0 +1,32 @@
+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 } 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. */
+export function bindLegacyDraftTarget(model: RuntimeModel, caches: Cache[]): TargetView {
+  if (model instanceof Gemma4Model) {
+    let sliding = -1, full = -1;
+    for (let i = 0; i < model.numDonors; i++) {
+      if (model.layers[i]!.layerType === "sliding_attention") sliding = i;
+      else full = i;
+    }
+    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) }),
+    }) });
+  }
+  if (model instanceof Qwen35Model) return Object.freeze({ identity: model, qwenMtp: Object.freeze({
+    hiddenSize: model.config.text.hiddenSize, layerCount: model.config.text.numHiddenLayers,
+    embed: model.embed.encode.bind(model.embed), logitsFromHidden: model.logitsFromHidden.bind(model),
+  }) });
+  return Object.freeze({ identity: model });
+}

package/src/backends/mlx/evaluation.ts

@@ -0,0 +1,105 @@
+import type { TaskModel, GenOpts } from "../../eval/runner";
+import type { CompletionClient } from "../../contracts/completion";
+import { generate } from "../../generate";
+import { clearCache } from "../../mlx/ffi";
+import * as ops from "../../mlx/ops";
+import { runtimeValue } from "../../runtime-config";
+import { DiffusionGemmaModel } from "../../model/diffusion-gemma";
+
+/** Bit-exact greedy decode: raw `model.forward` + argmax loop, which matches mlx-lm
+ *  token-for-token (gen8k: 8000/8000). The product `generate()` decode wrapper
+ *  (forwardHidden/logitsFromHidden + pipelined sampler) diverges from the raw forward
+ *  on near-ties past ~32 tokens, so the EVAL decodes bit-exactly itself rather than
+ *  alter the serving path. Greedy + full-precision KV only. */
+export function greedyDecodeBitExact(tm: TaskModel, ids: number[], maxTokens: number): string {
+  if (tm.activeAdapters) tm.model.loraState.active = tm.activeAdapters; // apply the mounted adapter
+  const cache = tm.model.makeCache();
+  try {
+    // Chunked, LM-head-free prefill. The old single-shot `forward(prompt)`
+    // materialized [1, L, vocab] logits just to discard them — at a 248k
+    // vocab that's ~1 GB per 1k prompt tokens, and with no allocator-cache
+    // clears the Metal buffer cache retained every giant freed tensor
+    // (the 27B capability-suite 73 GB swap-thrash). Prefill only needs the
+    // cache state, which forwardHidden advances identically (the LM head
+    // never touches the cache) — bit-exact by construction. Chunk size and
+    // clear cadence mirror mlx-lm generate_step / our generate().
+    const PREFILL_CHUNK = 2048;
+    const prompt = ids.slice(0, -1);
+    for (let s = 0; s < prompt.length; s += PREFILL_CHUNK) {
+      const chunk = prompt.slice(s, s + PREFILL_CHUNK);
+      const cIds = ops.fromInt32(chunk, [1, chunk.length]);
+      tm.model.forwardHidden(cIds, cache).dispose();
+      cIds.dispose();
+      clearCache();
+    }
+    let last = ids[ids.length - 1]!;
+    const eos = new Set(tm.config.eosTokenIds);
+    const out: number[] = [];
+    for (let i = 0; i < maxTokens; i++) {
+      const lg = tm.model.forward([last], cache);
+      const am = ops.argmaxAxis(lg, -1); // GPU argmax (mlx-lm greedy: lowest-index tie-break)
+      lg.dispose();
+      const best = ops.itemUint32(am);
+      am.dispose();
+      if (i % 256 === 0) clearCache(); // mlx-lm decode cadence
+      if (eos.has(best)) break; // mlx-lm halts on EOS and excludes it from the output
+      out.push(best);
+      last = best;
+    }
+    return tm.tokenizer.decode(out, true);
+  } finally {
+    for (const c of cache) c.dispose();
+  }
+}
+
+/** Complete `body` via the model's real quantized-KV path. Greedy by default;
+ *  pass `opts.sampler` to drive temperature/HLG/etc. */
+async function generateTextNative(tm: TaskModel, body: string, opts: GenOpts = {}): Promise<string> {
+  const maxTokens = opts.maxTokens ?? 256;
+  const enableThinking = opts.enableThinking ?? runtimeValue("MLX_BUN_EVAL_THINK") === "1";
+  const templated = opts.useChat !== false && tm.template !== null;
+  const text = templated
+    ? tm.template!.render([{ role: "user", content: body }], { addGenerationPrompt: true, enableThinking })
+    : body;
+  // The chat template ALREADY emits the BOS; encoding it with add_special_tokens on
+  // would prepend a SECOND BOS (<s><s>…), corrupting generation. Raw bodies (no
+  // template) get the BOS added as normal.
+  const ids = tm.tokenizer.encode(text, /* addSpecialTokens */ !templated);
+
+  // optiq's PUBLISHED eval generates with a FULL-PRECISION KV cache — its eval calls
+  // plain mlx_lm.generate (kv-quant lives in the serving runtime, not the eval). So
+  // for like-for-like parity we default the eval to UNQUANTIZED KV; MMLU's argmax
+  // path was already unquantized, which is why it matched while generation didn't.
+  // MLX_BUN_EVAL_KV_QUANT=1 generates through the model's quantized KV (serving
+  // config); an explicit opts.kvScheme (cli generate's resolved tier route) wins.
+  const envKv: GenOpts["kvScheme"] | undefined =
+    (runtimeValue("MLX_BUN_EVAL_KV_QUANT") === "1" && tm.config.kvQuant?.length)
+      ? { kvConfig: tm.config.kvQuant, quantizedKvStart: 0 } : undefined;
+  const kvScheme = opts.kvScheme ?? envKv;
+  const kvActive = !!(kvScheme && (kvScheme.kvBits || kvScheme.kvConfig?.length || kvScheme.turboQuant));
+
+  // Parity default — greedy + full-precision KV + no sampler arm: decode bit-exactly
+  // via the raw forward (matches mlx-lm token-for-token). Sampler arms / kv-quant fall
+  // through to the product generate(). DiffusionGemma is non-autoregressive (no AR
+  // forward) — always go through generate(), which routes it to the denoising engine.
+  const isDiffusion = tm.model instanceof DiffusionGemmaModel;
+  if (!isDiffusion && !kvActive && !opts.sampler && !tm.samplerOverride) {
+    return greedyDecodeBitExact(tm, ids, maxTokens);
+  }
+  const gen = generate(tm.model, ids, {
+    maxTokens,
+    temperature: 0, // greedy default — deterministic head-to-head arms
+    adapters: tm.activeAdapters ?? [], // per-request activation of the mounted eval-adapter
+    ...(kvActive ? kvScheme : {}),
+    ...(opts.sampler ?? {}), // overrides temperature when provided
+    ...(tm.samplerOverride ?? {}), // arm override wins over the task's own sampler
+  });
+
+  const out: number[] = [];
+  for await (const { token } of gen) out.push(token);
+  return tm.tokenizer.decode(out, true);
+}
+
+export function createEvalCompletionClient(tm: TaskModel): CompletionClient<{ body: string; options?: GenOpts }, string> {
+  return { complete: ({ body, options }) => generateTextNative(tm, body, options) };
+}

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

@@ -0,0 +1,45 @@
+import type { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import type { Cache } from "../../model/gemma4";
+import { withResource } from "../../engine/resources";
+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
+ * boundaries, preserving the cache-covered prefix on a consumer break.
+ */
+export async function appendFillHidden(
+  forward: (ids: MlxArray, cache: Cache[]) => MlxArray | Promise<MlxArray>,
+  cache: Cache[],
+  ids: number[],
+  chunkSize: (cache: readonly Cache[]) => number,
+): Promise<MlxArray> {
+  const input = ops.fromInt32(ids, [1, ids.length]);
+  const hidden: MlxArray[] = [];
+  try {
+    for (let start = 0; start < ids.length;) {
+      const limit = chunkSize(cache);
+      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]);
+      try {
+        const h = await forward(part, cache);
+        hidden.push(h);
+        // Bound transient graphs while retaining all positions for an optional
+        // diagnostic trace. Normal generation projects only the last position.
+        withResource(leaseCacheStates(cache), (state) => ops.evalAll([h, ...state]));
+      } finally {
+        part.dispose();
+      }
+      start = end;
+    }
+    return ops.concatAxis(hidden, 1);
+  } finally {
+    for (const h of hidden) h.dispose();
+    input.dispose();
+  }
+}

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