mlx-bun 0.4.0 → 1.0.0

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

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

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

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

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

@@ -1,89 +0,0 @@
-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?.(); },
-    });
-  }
-}