mlx-bun 0.2.0 → 0.4.0

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

package/src/backends/mlx/model-prompt.ts

@@ -0,0 +1,164 @@
+import { ensureWav } from "../../audio/transcode";
+import { DiffusionGemmaModel } from "../../model/diffusion-gemma";
+import { Gemma4Model } from "../../model/gemma4";
+import { Qwen35Model } from "../../model/qwen3_5";
+import { spliceImageTokens } from "../../vision/diffusion-vision";
+import { buildMultimodalPrompt, buildVisionPrompt, extractAudio, extractImages, extractVideos,
+  type VisionEncoder, type VisionTokenIds } from "../../vision/prompt";
+import { buildQwen3VLVisionPrompt } from "../../vision/qwen3vl-prompt";
+import type { Qwen3VLVisionTower } from "../../vision/qwen3vl-tower";
+import { normalizeMessages, type ChatRequestParams } from "../../serve/chat-request";
+import { getAudioTower, getVisionTower, type ServerContext } from "../../serve/model-host";
+import { RequestError } from "../../serve/pipeline";
+import type { RequestOwnership } from "../../serve/request-plan";
+import type { RequestPrep } from "../../serve/request-prep";
+import type { BuiltPrompt } from "../../serve/model-binding";
+
+export async function buildModelPrompt(
+    ctx: ServerContext,
+    prep: RequestPrep,
+    body: ChatRequestParams,
+    tools: ChatRequestParams["tools"] | null,
+    ownership: RequestOwnership,
+  ): Promise<BuiltPrompt> {
+    const toolList = tools ?? null;
+    const partOf = (types: string[]) => body.messages.some(
+      (m) => Array.isArray(m.content) &&
+        m.content.some((p: any) => types.includes(p.type)),
+    );
+    const hasImages = partOf(["image_url", "image"]);
+    // Same shapes extractAudio accepts: OpenAI-canonical input_audio plus
+    // optiq's audio / audio_url aliases (docs/design/generic-model-support.md §3.2).
+    const hasAudio = partOf(["input_audio", "audio", "audio_url"]);
+    // Video content parts (video_url / video with base64 data) —
+    // Qwen3.5-family only (decoded to sampled frames via the
+    // AVFoundation sidecar, vision/video-frames.ts).
+    const hasVideos = partOf(["video_url", "video"]);
+    const noMedia: Pick<BuiltPrompt, "startInThinking" | "probeStableLen" | "diffusionPixels"> =
+      { startInThinking: false, probeStableLen: false, diffusionPixels: null };
+
+    // Video is Qwen3.5-family only and never composes with audio —
+    // one early guard so no downstream branch can silently drop a
+    // video part (the gemma/diffusion builders don't know the type).
+    if (hasVideos && (hasAudio || !(ctx.model instanceof Qwen35Model))) {
+      throw new RequestError(400, hasAudio
+        ? "video and audio content parts cannot be combined"
+        : `model ${ctx.modelId} does not accept video input — video ` +
+          `content parts need a Qwen3.5-family model (e.g. Qwen3.8-27B)`);
+    }
+    if (hasAudio) {
+      // Audio (and MIXED image+audio) input — A4 of
+      // docs/design/generic-model-support.md. One buildMultimodalPrompt call
+      // splices both media kinds in document order. A request WITH
+      // audio on a model whose tower is unavailable is an explicit 400
+      // — never a silent text-only degrade (that leniency is only for
+      // requests without the media, getVisionTower's contract).
+      const audioTower = getAudioTower(ctx);
+      if (!audioTower || !ctx.audioTokenIds) {
+        throw new RequestError(400,
+          `model ${ctx.modelId} has no audio tower — audio input needs ` +
+          `a model whose config.json carries audio_config and whose ` +
+          `sidecar ships the audio tensors (e.g. gemma-4 e4b OptiQ)`);
+      }
+      let visionSide: { tower: VisionEncoder; tokenIds: VisionTokenIds } | undefined;
+      if (hasImages) {
+        const tower = getVisionTower(ctx);
+        if (!tower) throw new RequestError(400, "model has no vision sidecar");
+        visionSide = { tower, tokenIds: ctx.visionTokenIds };
+      }
+      const { messages: withAudioParts, images } =
+        await extractImages(normalizeMessages(body.messages));
+      const { messages, audio } = await extractAudio(withAudioParts);
+      // Non-WAV containers (mp3/m4a/flac/ogg/aiff/…) transcode to
+      // 16 kHz WAV via CoreAudio; RIFF bytes pass through untouched.
+      // Failures throw into the prompt-build 400.
+      const wavs = await Promise.all(audio.map(ensureWav));
+      // The towers are only ever non-null for Gemma4 (loader gates).
+      const mp = await buildMultimodalPrompt(
+        ctx.model as Gemma4Model,
+        {
+          ...(visionSide ? { vision: visionSide } : {}),
+          audio: { tower: audioTower, tokenIds: ctx.audioTokenIds },
+        },
+        ctx.tokenizer, ctx.template, messages, images, wavs, toolList,
+      );
+      // bidirMask is null whenever audio is present (§3.3 Q1: mixed
+      // prompts run fully causal); the union mask does the per-layer
+      // id zeroing either way.
+      return {
+        ...noMedia,
+        promptIds: mp.ids,
+        vision: {
+          embeddings: mp.embeddings,
+          ...(mp.bidirMask ? { imageMask: mp.bidirMask } : {}),
+          multimodalMask: mp.multimodalMask,
+        },
+      };
+    }
+    if (hasImages && ctx.model instanceof DiffusionGemmaModel) {
+      // DiffusionGemma image-text-to-text: its OWN dedicated SigLIP tower +
+      // encoder vision merge feed the denoising engine (NOT the AR
+      // forwardEmbeddings path). v1 supports a single image.
+      const dm = ctx.model;
+      if (!dm.visionTower) throw new RequestError(400, "this checkpoint has no vision tower");
+      const { messages, images } = await extractImages(normalizeMessages(body.messages));
+      if (images.length !== 1)
+        throw new RequestError(400, "DiffusionGemma image input supports exactly one image");
+      const rendered = ctx.template.render(messages, { tools: toolList, addGenerationPrompt: true });
+      const rawIds = ctx.tokenizer.encode(rendered, /* addSpecialTokens */ false);
+      const { pixels, softTokens } = await dm.visionTower.preprocess(images[0]!);
+      const diffusionPixels = ownership.own(pixels);
+      const promptIds = spliceImageTokens(rawIds, [softTokens], {
+        image: ctx.visionTokenIds.imageTokenId,
+        boi: ctx.visionTokenIds.boiTokenId,
+        eoi: ctx.visionTokenIds.eoiTokenId,
+      });
+      return { ...noMedia, promptIds, vision: undefined, diffusionPixels };
+    }
+    if ((hasImages || hasVideos) && ctx.model instanceof Qwen35Model) {
+      // Qwen3.8 vision + video: the tower is a Qwen3VLVisionTower
+      // riding the shared lazy slot (makeVisionLoader's qwen branch);
+      // image/video spans splice into input embeddings and the request
+      // carries mRoPE positions + delta (PLAN 14v/14w). Videos decode
+      // to sampled frames via the AVFoundation sidecar.
+      const tower = getVisionTower(ctx) as unknown as Qwen3VLVisionTower | null;
+      if (!tower) throw new RequestError(400, "model has no vision sidecar");
+      const { messages: withVideos, images } =
+        await extractImages(normalizeMessages(body.messages));
+      const { messages, videos } = await extractVideos(withVideos);
+      const vp = await buildQwen3VLVisionPrompt(
+        ctx.model, tower, ctx.tokenizer, ctx.template, messages, images,
+        {
+          imageTokenId: (ctx.model.config.raw.image_token_id as number) ?? 248056,
+          videoTokenId: (ctx.model.config.raw.video_token_id as number) ?? 248057,
+          visionStartId: (ctx.model.config.raw.vision_start_token_id as number) ?? 248053,
+          visionEndId: (ctx.model.config.raw.vision_end_token_id as number) ?? 248054,
+        },
+        prep.templateOptionsFor(body, toolList),
+        videos,
+      );
+      return { ...noMedia, promptIds: vp.ids, vision: { embeddings: vp.embeddings, mrope: vp.mrope } };
+    }
+    if (hasImages) {
+      // Loads (and caches) the tower on first image request — text-only
+      // sessions never pay for it.
+      const tower = getVisionTower(ctx);
+      if (!tower) throw new RequestError(400, "model has no vision sidecar");
+      const { messages, images } = await extractImages(normalizeMessages(body.messages));
+      // The tower is only ever non-null for Gemma4 (sidecar gate in
+      // makeVisionLoader), so the model narrow is safe here.
+      const vp = await buildVisionPrompt(
+        ctx.model as Gemma4Model, tower, ctx.tokenizer, ctx.template,
+        messages, images, ctx.visionTokenIds, toolList,
+      );
+      return { ...noMedia, promptIds: vp.ids, vision: { embeddings: vp.embeddings, imageMask: vp.imageMask } };
+    }
+    const text = prep.promptIdsFor(body, toolList);
+    return {
+      promptIds: text.ids,
+      vision: undefined,
+      startInThinking: text.startInThinking,
+      probeStableLen: true,
+      diffusionPixels: null,
+    };
+  }

package/src/backends/mlx/model-serving.ts

@@ -0,0 +1,79 @@
+import type { ServerContext, ServingContext } from "../../serve/model-host";
+import type { ModelServingBinding } from "../../serve/model-binding";
+import { Glm52Model } from "../../model/glm52";
+import { embedMany, isEmbeddingModel } from "../../embed";
+import { bindMlxGateway } from "./gateway-binding";
+import { bindLegacySerialModel, createMlxSerialExecutor } from "./serial-executor";
+import { buildModelPrompt } from "./model-prompt";
+import { MLX_VERSION, deviceArchitecture } from "../../mlx/ffi";
+
+export function modelPromptBuilder(context: ServingContext): ModelServingBinding["buildPrompt"] {
+  return context.serving?.buildPrompt ?? ((body, tools, ownership, prep) =>
+    buildModelPrompt(context as ServerContext, prep, body, tools, ownership));
+}
+
+/** Existing model classes enter the interface once here. New implementations
+ * provide ctx.serving directly and need no forward/makeCache methods. */
+export function modelServingBinding(context: ServingContext): ModelServingBinding {
+  if (context.serving) return context.serving;
+  const ctx = context as ServerContext;
+  const model = ctx.model;
+  if (typeof model.makeCache !== "function" || typeof model.forward !== "function")
+    throw new Error("model implementation must supply its serving binding");
+  const serial = bindLegacySerialModel(model, ctx.draft ?? undefined);
+  const glm = model instanceof Glm52Model;
+  const embedding = isEmbeddingModel(model);
+  return {
+    stateCompatibility: `mlx-${MLX_VERSION}-${deviceArchitecture()}`,
+    gateway: bindMlxGateway(model, ctx.draft ?? undefined),
+    createSerial: (services) => createMlxSerialExecutor(serial, services),
+    buildPrompt: (body, tools, ownership, prep) => buildModelPrompt(ctx, prep, body, tools, ownership),
+    restore: (store, entry) => store.restore(entry, model),
+    async signal(ids, count, minimum) {
+      const caches = model.makeCache();
+      try {
+        const logits = model.forward(ids, caches);
+        let f: Float32Array;
+        let vocab: number;
+        try {
+          const [, length, width] = logits.shape as [number, number, number];
+          vocab = width;
+          const last = logits.slice([0, length - 1, 0], [1, length, width]);
+          try { f = last.toFloat32(); } finally { last.dispose(); }
+        } finally { logits.dispose(); }
+        let mx = -Infinity;
+        for (const value of f) if (value > mx) mx = value;
+        let sum = 0;
+        for (const value of f) sum += Math.exp(value - mx);
+        const lse = mx + Math.log(sum);
+        const bins = new Array<number>(count).fill(0);
+        for (const value of f) {
+          const t = Math.max(0, Math.min(1, (value - lse - minimum) / -minimum));
+          const bin = Math.min(count - 1, Math.floor(t * count));
+          bins[bin] = (bins[bin] ?? 0) + 1;
+        }
+        return { bins, vocab };
+      } finally { for (const cache of caches) cache.dispose(); }
+    },
+    discovery: { adapters: !glm, training: !glm, dsa: glm && model.capabilities.dsa, embeddings: embedding },
+    ...(embedding ? { embed: (inputs: string[], instruction?: string) => embedMany(model, ctx.tokenizer, inputs, instruction) } : {}),
+    diagnostics() {
+      const plan = ctx.glmMemoryPlan;
+      if (!plan) return {};
+      const runtime = glm ? model.expertRuntime : undefined;
+      return { glm52: {
+        preset: plan.preset, planned_process_bytes: plan.plannedProcessBytes,
+        process_limit_bytes: plan.processLimitBytes, context_tokens: plan.contextTokens,
+        max_generation_tokens: plan.maxGenerationTokens, batch_size: plan.batchSize,
+        dsa: glm && model.capabilities.dsa, mtp: ctx.draft?.provider.id === "glm52-native-mtp",
+        mtp_draft_tokens: plan.mtpDraftTokens, resident_weight_bytes: plan.lineItems.residentWeightsBytes,
+        main_expert_slab_bytes: plan.lineItems.mainExpertSlabBytes,
+        mtp_expert_slab_bytes: plan.lineItems.mtpExpertSlabBytes,
+        expert_runtime: runtime ? {
+          main_residency: runtime.manager.snapshot(), mtp_residency: runtime.mtp?.manager.snapshot() ?? null,
+          last_turn: runtime.lastTelemetry, last_repin: runtime.lastRepin,
+        } : null,
+      } };
+    },
+  };
+}

package/src/backends/mlx/prefill-cohort.ts

@@ -0,0 +1,89 @@
+import type { Cache } from "../../model/gemma4-base";
+import type { RuntimeModel } from "../../model/factory";
+import type { MlxArray } from "../../mlx/array";
+import { nextPrefillStep } from "../../inference/prefill";
+import type { CacheCodecProvider } from "../../kv-store";
+import { disposeResources, cleanupFailure } from "../../engine/resources";
+import type { Row, RowPromptCache } from "./batch-group";
+import type { KvMaintenance } from "./kv-maintenance";
+import { MlxPrefillRows, type MlxPrefillState } from "./prefill-rows";
+
+export interface PrefillState extends MlxPrefillState {
+  snapAt: number | null;
+  continuation?: import("./continuation").OrdinaryContinuationState;
+  closePrefill?: () => void;
+}
+
+interface PrefillHost {
+  model: RuntimeModel;
+  chunkSize: number;
+  tailSplit: boolean;
+  promptCache?: RowPromptCache;
+  stateCodecs?: CacheCodecProvider;
+  maintain?: KvMaintenance;
+  forward(ids: MlxArray, caches: Cache[]): Promise<MlxArray>;
+  /** Target projection preserves the forward batch geometry before sampling. */
+  project(hidden: MlxArray, caches: Cache[], completed: readonly PrefillState[]): MlxArray;
+  /** Borrows [1,1,V] logits; takes ownership of the completed request caches. */
+  complete(state: PrefillState, logits: MlxArray): Promise<void>;
+  resume?(state: PrefillState): Promise<void>;
+  reject(row: Row, error: unknown): void;
+}
+
+/** Ordinary sampling/cache policy over the shared target preparation driver. */
+export class MlxPrefillCohort extends MlxPrefillRows<PrefillState> {
+  constructor(readonly host: PrefillHost) {
+    super({
+      stateCodecs: host.stateCodecs, maintain: host.maintain,
+      open(row) {
+        let owned: { caches: Cache[]; retain?: () => void } | undefined;
+        let closeCache: (() => void) | undefined;
+        let closePrefill: (() => void) | undefined;
+        try {
+          closeCache = row.req.trace?.begin("cache.lookup_restore", { mechanism: "continuous" });
+          const continuation = row.req.continuation?.restore(row.cacheNamespace ?? "");
+          if (continuation) {
+            owned = { caches: continuation.caches };
+            return { row, solo: continuation.caches, pos: row.req.promptIds.length,
+              snapAt: null, continuation };
+          }
+          const hit = host.promptCache?.take(row.req.promptIds, row.cacheNamespace) ?? null;
+          if (hit) owned = hit;
+          closeCache?.(); closeCache = undefined;
+          if (hit) row.cachedTokens = hit.tokens.length;
+          const length = row.req.promptIds.length;
+          const boundary = Math.min(row.req.snapshotAt ?? length, length - 1);
+          const snapAt = host.promptCache && boundary >= 256 && boundary > (hit?.tokens.length ?? 0) ? boundary : null;
+          closePrefill = row.req.trace?.begin("prefill.total", {
+            mechanism: "continuous", promptTokens: row.promptTokens, cachedTokens: row.cachedTokens,
+          });
+          const closeSetup = row.req.trace?.begin("prefill.batch_setup", { mechanism: "continuous" });
+          owned ??= { caches: row.req.statePolicy?.create() ?? host.model.makeCache() }; closeSetup?.();
+          return { row, solo: owned.caches, pos: hit?.tokens.length ?? 0, retain: owned.retain, snapAt, closePrefill };
+        } catch (error) {
+          return cleanupFailure(error, () => disposeResources([...(owned?.caches ?? []),
+            { dispose: () => owned?.retain?.() }, { dispose: () => closePrefill?.() }]));
+        } finally { closeCache?.(); }
+      },
+      ready: state => !!state.continuation,
+      plan(state) { return nextPrefillStep({ length: state.row.promptTokens, position: state.pos,
+        chunkSize: state.row.req.prefillChunkSize ?? host.chunkSize,
+        tailSplit: host.tailSplit, snapshotAt: state.snapAt }); },
+      forward: host.forward.bind(host), project: host.project.bind(host),
+      complete: (state, logits) => state.continuation ? host.resume!(state) : host.complete(state, logits!), reject: host.reject.bind(host),
+      checkpoint(state, capture) {
+        let snapshot: Cache[] | undefined;
+        try {
+          snapshot = capture();
+          host.promptCache!.put(state.row.req.promptIds.slice(0, state.pos), snapshot, state.row.cacheNamespace);
+          snapshot = undefined;
+        } catch (error) {
+          if (snapshot) disposeResources(snapshot);
+          console.warn(`batch-lane boundary snapshot skipped: ${(error as Error).message}`);
+        }
+        state.snapAt = null;
+      },
+      close(state) { const close = state.closePrefill; state.closePrefill = undefined; close?.(); },
+    });
+  }
+}