mlx-bun 0.2.0 → 0.3.0

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,47 @@
+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";
+
+/** 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;
+    }
+    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({
+      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/gateway-binding.ts

@@ -0,0 +1,71 @@
+import type { ModelConfig } from "../../config";
+import type { KvScheme } from "../../kv-scheme";
+import type { RuntimeModel } from "../../model/factory";
+import { DiffusionGemmaModel } from "../../model/diffusion-gemma";
+import { UniversalDenseModel } from "../../model/universal/dense";
+import { KVCache, RotatingKVCache, isBatchableCache, isPlainKvCache, isRotatingPlainCache } from "../../model/gemma4-base";
+import { SSMCache } from "../../model/qwen3-delta";
+import { Qwen35Model } from "../../model/qwen3_5";
+import { runtimeConfig, type RuntimeConfig } from "../../runtime-config";
+import { disposeResources } from "../../engine/resources";
+import { legacyCompiledDecodeAvailable } from "./autoregressive";
+import { MlxBatchExecutionGroup, type MlxBatchExecutionGroupOptions } from "./batch-group";
+import type { GenerateOptions } from "../../generate";
+import type { ExecutionRequirements, ResolvedExecution } from "../../contracts/execution";
+import { resolveExecution } from "../../engine/execution-plan";
+
+export interface MlxBatchGroup extends Pick<MlxBatchExecutionGroup,
+  "activeRows" | "pendingRows" | "projectedKvBytes" | "kvBudgetBytes" | "submit" | "kick" | "close"> {}
+
+/** A model implementation owns capability checks and the execution group.
+ * Scheduling never inspects concrete model/cache classes. */
+export interface MlxGatewayBinding {
+  readonly config: ModelConfig;
+  readonly runtime: RuntimeConfig;
+  plan(request: ExecutionRequirements, options: GenerateOptions,
+    scheduling: { continuous: boolean; quantizedBatch: boolean; checkpoints: boolean }): ResolvedExecution;
+  cachesBatchable(): boolean;
+  kvBatchable(scheme: KvScheme): boolean;
+  createBatchGroup(options: MlxBatchExecutionGroupOptions): MlxBatchGroup;
+}
+
+export function bindMlxGateway(model: RuntimeModel): MlxGatewayBinding {
+  const runtime = runtimeConfig();
+  return {
+    config: model.config, runtime,
+    plan(request, options, scheduling) {
+      return resolveExecution(request, {
+        ...scheduling,
+        method: model instanceof DiffusionGemmaModel ? "denoising" : "autoregressive",
+        compiledDecode: legacyCompiledDecodeAvailable(model),
+        grammarBatch: runtime.value("MLX_BUN_GRAMMAR_BATCH") !== "0",
+        speculativeKvQuant: model instanceof Qwen35Model && options.kvBits === 4 &&
+          !options.kvConfig?.length && options.quantizedKvStart === 0 &&
+          runtime.flag("MLX_BUN_QWEN_SPEC_KV4", false),
+      }, {
+        pagedKv: !!options.pagedKv, fill: !!options.fill,
+        compiledDecode: runtime.flag("MLX_BUN_COMPILED_DECODE", true),
+        grammarJump: runtime.flag("MLX_BUN_GRAMMAR_JUMP", false),
+      });
+    },
+    cachesBatchable() {
+      if (model instanceof UniversalDenseModel)
+        return !model.args.maskArray && !model.args.layerTypes?.includes("sliding_attention");
+      const caches = model.makeCache();
+      try {
+        const ssm = runtime.value("MLX_BUN_BATCH_SSM") !== "0";
+        return caches.every((cache) => cache instanceof KVCache || cache instanceof RotatingKVCache ||
+          isBatchableCache(cache) || (ssm && cache instanceof SSMCache));
+      } finally { disposeResources(caches); }
+    },
+    kvBatchable(scheme) {
+      if (scheme.kind !== "affine-config") return false;
+      const caches = model.makeCache();
+      try {
+        return scheme.batchable(model.config,
+          (layer) => isPlainKvCache(caches[layer]) || isRotatingPlainCache(caches[layer]));
+      } finally { disposeResources(caches); }
+    },
+    createBatchGroup: (options) => new MlxBatchExecutionGroup(model, options),
+  };
+}

package/src/backends/mlx/graph.ts

@@ -0,0 +1,38 @@
+import type { AutoregressiveGraph, GraphDescriptor, LogitSelection } from "../../inference/graph";
+import type { MlxArray } from "../../mlx/array";
+
+/** Structural input: generated and hand-written graphs need no union membership. */
+export interface MlxGraphOperations<State> {
+  forwardHidden(ids: MlxArray, state: State): MlxArray;
+  forwardHiddenAsync?: (ids: MlxArray, state: State) => Promise<MlxArray>;
+  logitsFromHidden(hidden: MlxArray): MlxArray;
+}
+
+/** Bind once. No array wrapper, await, readback, or synchronization is added
+ * to synchronous forward calls. Existing quant-specific kernels stay inside
+ * the supplied operations. This adapter does not take ownership of weights. */
+export function bindMlxGraph<State>(
+  operations: MlxGraphOperations<State>,
+  descriptor: Omit<GraphDescriptor, "backend" | "graphAbi">,
+): AutoregressiveGraph<MlxArray, State, MlxArray> {
+  const forward = operations.forwardHiddenAsync
+    ? operations.forwardHiddenAsync.bind(operations)
+    : operations.forwardHidden.bind(operations);
+  const project = operations.logitsFromHidden.bind(operations);
+  return {
+    descriptor: Object.freeze({ ...descriptor, backend: "mlx", graphAbi: "mlx-hidden-bsh-v1" }),
+    forwardHidden: forward,
+    projectLogits(hidden: MlxArray, selection: LogitSelection) {
+      if (selection.type === "all") return project(hidden);
+      const shape = hidden.shape;
+      if (shape.length !== 3) throw new Error("MLX hidden state must have shape [batch, positions, width]");
+      const [batch, positions, width] = shape as [number, number, number];
+      const start = selection.type === "last" ? positions - 1 : selection.start;
+      const end = selection.type === "last" ? positions : selection.end;
+      if (!Number.isInteger(start) || !Number.isInteger(end) || start < 0 || end <= start || end > positions)
+        throw new Error("invalid logits position selection");
+      const selected = hidden.slice([0, start, 0], [batch, end, width]);
+      try { return project(selected); } finally { selected.dispose(); }
+    },
+  };
+}

package/src/backends/mlx/memory-client.ts

@@ -0,0 +1,189 @@
+import { existsSync } from "node:fs";
+import { loadTaskModel, greedyDecodeBitExact, type TaskModel } from "../../eval/runner";
+import type { ChatTemplate } from "../../chat-template";
+import { MlxBatchExecutionGroup as BatchScheduler, type RowSampler } from "./batch-group";
+import * as ops from "../../mlx/ops";
+import type { MlxArray } from "../../mlx/array";
+import { adapterDirFor, memoryBatchSize, memoryPromptIds, MODEL_ID,
+  type LocalInput, type MemoryCompletionClient } from "../../memory/model";
+
+class AsyncMutex {
+  #tail: Promise<void> = Promise.resolve();
+  acquire(): Promise<() => void> {
+    let release!: () => void;
+    const gate = new Promise<void>((r) => { release = r; });
+    const wait = this.#tail;
+    this.#tail = this.#tail.then(() => gate);
+    return wait.then(() => release);
+  }
+}
+
+interface MemoryRuntime {
+  tm: TaskModel;
+  template: ChatTemplate;
+  scheduler: BatchScheduler;
+  /** True once the `memory-chunk` adapter is mounted on the shared model. */
+  hasChunkAdapter: boolean;
+}
+
+
+function requireModel(): void {
+  if (!existsSync(`${MODEL_ID}/config.json`))
+    throw new Error(`memory: Gemma-4-e4b is not downloaded (looked under ${MODEL_ID}). Fetch it first: HF_HUB_DISABLE_XET=1 hf download mlx-community/gemma-4-e4b-it-OptiQ-4bit`);
+}
+
+let runtimePromise: Promise<MemoryRuntime> | null = null;
+const opLock = new AsyncMutex(); // serializes whole memory ops (loraState ownership)
+const schedLock = new AsyncMutex(); // the scheduler's exclusive-GPU lock
+
+/** Lazily load the ONE shared model + mount memory-chunk + build the scheduler.
+ *  Never auto-downloads; throws (with the fetch hint) until the snapshot lands. */
+async function getRuntime(): Promise<MemoryRuntime> {
+  if (!runtimePromise) {
+    runtimePromise = (async () => {
+      requireModel();
+      const tm = await loadTaskModel(MODEL_ID); // base e4b, no adapter
+      if (!tm.template) throw new Error(`memory: ${MODEL_ID} has no chat template`);
+      let hasChunkAdapter = false;
+      const chunkDir = adapterDirFor("chunk");
+      if (chunkDir) {
+        // Mount ONCE; activation is per-call via loraState.active (see decodeGreedy
+        // / batchDecode). Mounting alone runs base — it must be activated.
+        const { AdapterManager } = await import("../../lora");
+        await new AdapterManager(tm.model).mount("memory-chunk", chunkDir);
+        hasChunkAdapter = true;
+      }
+      const scheduler = new BatchScheduler(tm.model, {
+        maxBatch: memoryBatchSize(),
+        lock: { acquire: () => schedLock.acquire() },
+      });
+      return { tm, template: tm.template, scheduler, hasChunkAdapter };
+    })();
+  }
+  return runtimePromise;
+}
+
+/** The adapters to activate for a stage. Only `chunk` is adapter-bound (and only
+ *  when `memory-chunk` is mounted); every base stage runs with `active = []`. */
+function activeFor(stage: string, rt: MemoryRuntime): string[] {
+  return stage === "chunk" && rt.hasChunkAdapter ? ["memory-chunk"] : [];
+}
+
+/** One bit-exact greedy decode for a stage (raw-forward argmax, runner parity).
+ *  Sets the stage's adapter for the call and resets `active` to [] after. */
+function decodeGreedy(rt: MemoryRuntime, stage: string, input: LocalInput, maxTokens: number): string {
+  const active = activeFor(stage, rt);
+  rt.tm.model.loraState.active = active;
+  const ids = memoryPromptIds(stage, input, rt.tm.tokenizer, rt.template);
+  try {
+    return greedyDecodeBitExact(
+      { ...rt.tm, activeAdapters: active.length ? active : undefined },
+      ids,
+      maxTokens,
+    );
+  } finally {
+    rt.tm.model.loraState.active = [];
+  }
+}
+
+/** Batched decode for N independent inputs of ONE stage. Sets the stage adapter
+ *  globally for the whole batch (the scheduler reads `loraState.active` per
+ *  forward, so all rows get it uniformly — a batch is necessarily single-adapter),
+ *  submits every row to the shared scheduler with a greedy per-row sampler, and
+ *  resets `active` after all rows finish. Order-preserving: out[i] ↔ inputs[i].
+ *
+ *  Decode-path caveat: this samples via forwardHidden/logitsFromHidden (the
+ *  scheduler), which can diverge from the runner's raw-forward greedy on near-ties
+ *  past ~32 tokens. For the JSON/binary memory stages (validated by
+ *  well-formedness) this is acceptable; the single-call fallback stays bit-exact. */
+async function batchDecode(
+  rt: MemoryRuntime,
+  stage: string,
+  inputs: LocalInput[],
+  maxTokens: number,
+): Promise<string[]> {
+  const active = activeFor(stage, rt);
+  rt.tm.model.loraState.active = active;
+  const eos = rt.tm.config.eosTokenIds;
+  const greedySample: RowSampler = (logits1V: MlxArray) => ops.argmaxAxis(logits1V, -1);
+  try {
+    return await Promise.all(
+      inputs.map((inp) => {
+        const ids = memoryPromptIds(stage, inp, rt.tm.tokenizer, rt.template);
+        const toks: number[] = [];
+        return rt.scheduler
+          .submit({
+            promptIds: ids,
+            maxTokens,
+            eosTokenIds: eos,
+            sample: greedySample,
+            onToken: (t) => { toks.push(t); },
+          })
+          .then(() => rt.tm.tokenizer.decode(toks, true));
+      }),
+    );
+  } finally {
+    rt.tm.model.loraState.active = [];
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Public seams
+// ---------------------------------------------------------------------------
+
+/** Call the local model for a pipeline `stage` with a {system?, user} input.
+ *  Renders through the shared chat-template path, runs the MODEL precheck, and
+ *  decodes ONE completion on the bit-exact greedy path (the safe fallback used by
+ *  every single-call site). NEVER auto-downloads. */
+export async function callLocal(
+  stage: string,
+  input: LocalInput,
+  opts?: { maxTokens?: number },
+): Promise<string> {
+  const rt = await getRuntime();
+  const release = await opLock.acquire();
+  try {
+    return decodeGreedy(rt, stage, input, opts?.maxTokens ?? 256);
+  } finally {
+    release();
+  }
+}
+
+/** N independent calls for ONE stage (homogeneous adapter by construction).
+ *  Routed through the shared in-process BatchScheduler when batch>1 AND N>1, else
+ *  looped on the bit-exact greedy fallback. Order-preserving: out[i] ↔ inputs[i]. */
+export async function callLocalBatch(
+  stage: string,
+  inputs: LocalInput[],
+  opts?: { maxTokens?: number },
+): Promise<string[]> {
+  if (inputs.length === 0) return [];
+  const rt = await getRuntime();
+  const maxTokens = opts?.maxTokens ?? 256;
+  const release = await opLock.acquire();
+  try {
+    if (inputs.length === 1 || memoryBatchSize() <= 1) {
+      const out: string[] = [];
+      for (const inp of inputs) out.push(decodeGreedy(rt, stage, inp, maxTokens));
+      return out;
+    }
+    return await batchDecode(rt, stage, inputs, maxTokens);
+  } finally {
+    release();
+  }
+}
+
+export const memoryCompletionClient: MemoryCompletionClient = {
+  complete: ({ stage, input, maxTokens }) => callLocal(stage, input, { maxTokens }),
+  completeBatch(requests) {
+    if (!requests.length) return Promise.resolve([]);
+    const first = requests[0]!;
+    if (requests.some((request) => request.stage !== first.stage || request.maxTokens !== first.maxTokens))
+      return (async () => {
+        const results: string[] = [];
+        for (const request of requests) results.push(await this.complete(request));
+        return results;
+      })();
+    return callLocalBatch(first.stage, requests.map((request) => request.input), { maxTokens: first.maxTokens });
+  },
+};

package/src/backends/mlx/memory-guard.ts

@@ -0,0 +1,40 @@
+import { DEFAULT_CHUNK, TRANSIENT_PER_TOKEN } from "../../fit";
+import { kvBytesAt, type KvSchemeOptions } from "../../kv-scheme";
+import { activeMemory, clearCache, maxRecommendedWorkingSetSize } from "../../mlx/ffi";
+import type { RuntimeModel } from "../../model/factory";
+import { Qwen35Model } from "../../model/qwen3_5";
+import type { PromptCache } from "../../prompt-cache";
+import { runtimeValue } from "../../runtime-config";
+
+/** Admission's ceiling and the cache that can release memory under the GPU lease. */
+export interface MlxMemoryBudget {
+  readonly usableBytes: number;
+  readonly kvOptions: KvSchemeOptions;
+  readonly promptCache: Pick<PromptCache, "relievePressure" | "totalBytes">;
+}
+
+/** Both execution paths enter after acquiring the request's reusable prefix.
+ * The caller owns the GPU lease and must close the scope on every exit. */
+export function enterMlxMemoryGuard(model: RuntimeModel, budget: MlxMemoryBudget,
+  promptTokens: number, prefillChunkSize?: number): { check(): void; close(): void } {
+  const reserve = (prefillChunkSize ?? DEFAULT_CHUNK) * TRANSIENT_PER_TOKEN +
+    kvBytesAt(model.config, promptTokens, budget.kvOptions);
+  const ceiling = Math.min(budget.usableBytes, maxRecommendedWorkingSetSize());
+  const check = () => {
+    if (activeMemory() + reserve <= ceiling) return;
+    clearCache();
+    if (runtimeValue("MLX_BUN_PREFILL_MEM_LOG") === "1")
+      console.error(`[cache-memory] active=${activeMemory()} prompt=${budget.promptCache.totalBytes} reserve=${reserve} ceiling=${ceiling}`);
+    const count = budget.promptCache.relievePressure(() => {
+      clearCache();
+      return activeMemory() + reserve > ceiling;
+    });
+    if (count) console.log(`[prompt-cache] demoted ${count} entries for GPU headroom`);
+    if (activeMemory() + reserve > ceiling)
+      throw new Error("insufficient GPU headroom after prompt-cache demotion; saved generation checkpoints are retained");
+  };
+  const qwen = model instanceof Qwen35Model ? model : undefined;
+  const previous = qwen?.prefillMemoryGuard;
+  if (qwen) qwen.prefillMemoryGuard = check;
+  return { check, close() { if (qwen) qwen.prefillMemoryGuard = previous ?? null; } };
+}

package/src/backends/mlx/methods.ts

@@ -0,0 +1,90 @@
+import { runtimeConfig, withRuntimeConfig, type RuntimeConfig } from "../../runtime-config";
+import { snapshotGenerationPolicy } from "./request-policy";
+import type { GenerationOutput, InferenceMethod, MethodResult } from "../../contracts/generation";
+import { throwIfCancelled } from "../../engine/cancellation";
+import { generateAutoregressive, generateDenoising, type GenerateOptions, type GenerateStats } from "../../generate";
+import { specRun } from "../../spec/serve-loop";
+import type { MlxAutoregressiveBinding } from "./autoregressive";
+import type { MlxSpeculativeBinding } from "./speculative";
+import type { MlxDenoisingBinding } from "./diffusion";
+
+/** The native algorithm owns cleanup in its finally. This host adapter owns
+ * cancellation subscriptions and waits for the algorithm before closing. */
+function nativeMethod<Metrics>(
+  id: string,
+  execute: (output: GenerationOutput, signal: AbortSignal) => Promise<MethodResult<Metrics>>,
+  runtime: RuntimeConfig,
+): InferenceMethod<Metrics> {
+  return {
+    id,
+    async createRun(cancellation) {
+      throwIfCancelled(cancellation);
+      const abort = new AbortController();
+      const unsubscribe = cancellation.subscribe((reason) => abort.abort(
+        new DOMException(`generation cancelled: ${reason}`, "AbortError"),
+      ));
+      let work: Promise<MethodResult<Metrics>> | undefined;
+      let closing: Promise<void> | undefined;
+      return {
+        execute(output) {
+          if (work || closing) return Promise.reject(new Error("method run can execute only once"));
+          work = withRuntimeConfig(runtime, () => execute(output, abort.signal));
+          return work;
+        },
+        close() {
+          return closing ??= (async () => {
+            abort.abort(new DOMException("method closed", "AbortError"));
+            try { await work; } catch { /* execute owns the algorithm's error */ }
+            finally { unsubscribe(); }
+          })();
+        },
+      };
+    },
+  };
+}
+
+/** Each plan supplies its prepared per-run resources. Caches/media in options
+ * retain the same ownership contract as generateAutoregressive. */
+export function createAutoregressiveMethod(
+  binding: MlxAutoregressiveBinding, promptIds: readonly number[], options: GenerateOptions,
+): InferenceMethod<GenerateStats> {
+  const prompt = [...promptIds];
+  options = snapshotGenerationPolicy(options);
+  return nativeMethod("mlx-autoregressive", async (output, signal) => {
+    const generation = generateAutoregressive(binding, prompt, { ...options, signal });
+    for await (const item of generation)
+      if (await output.commit([item.token], item.logprobs ? [item.logprobs] : undefined) === false) break;
+    const metrics = generation.stats;
+    if (!metrics) throw new Error("AR method did not settle its metrics");
+    return { finishReason: metrics.finishReason ?? (metrics.generatedTokens >= (options.maxTokens ?? 512) ? "length" : "stop"), metrics };
+  }, binding.runtime ?? runtimeConfig());
+}
+
+export function createSpeculativeMethod(
+  binding: MlxSpeculativeBinding, numDraftTokens: number,
+  promptIds: readonly number[], options: GenerateOptions,
+): InferenceMethod<GenerateStats> {
+  const prompt = [...promptIds];
+  options = snapshotGenerationPolicy(options);
+  return nativeMethod("mlx-speculative", async (output, signal) => {
+    const metrics = await specRun(binding, numDraftTokens, prompt, { ...options, signal },
+      (token) => output.commit([token]));
+    return { finishReason: metrics.finishReason ?? (metrics.generatedTokens >= (options.maxTokens ?? 512) ? "length" : "stop"), metrics };
+  }, binding.runtime ?? runtimeConfig());
+}
+
+/** The canvas stays private. This method publishes only the finished result. */
+export function createDenoisingMethod<State>(
+  binding: MlxDenoisingBinding<State>, promptIds: readonly number[], options: GenerateOptions,
+): InferenceMethod<GenerateStats> {
+  const prompt = [...promptIds];
+  options = snapshotGenerationPolicy(options);
+  return nativeMethod("mlx-denoising", async (output, signal) => {
+    const generation = generateDenoising(binding, prompt, { ...options, signal });
+    for await (const item of generation)
+      if (await output.commit([item.token]) === false) break;
+    const metrics = generation.stats;
+    if (!metrics) throw new Error("denoising method did not settle its metrics");
+    return { finishReason: metrics.generatedTokens >= (options.maxTokens ?? 256) ? "length" : "stop", metrics };
+  }, binding.runtime ?? runtimeConfig());
+}