mlx-bun 0.4.0 → 1.0.0

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

@@ -1,67 +0,0 @@
-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

@@ -1,31 +0,0 @@
-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

@@ -1,53 +0,0 @@
-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

@@ -1,106 +0,0 @@
-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

@@ -1,43 +0,0 @@
-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

@@ -1,77 +0,0 @@
-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

@@ -1,39 +0,0 @@
-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

@@ -1,32 +0,0 @@
-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

@@ -1,105 +0,0 @@
-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

@@ -1,45 +0,0 @@
-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

@@ -1,32 +0,0 @@
-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]); }
-}