mlx-bun 0.5.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/mixed-iteration.ts

@@ -1,62 +0,0 @@
-import type { MlxArray } from "../../mlx/array";
-import type { Cache } from "../../model/gemma4-base";
-import type { TokenGroup, TokenWorkOptions } from "../../model/token-groups";
-
-export type MlxForwardWork = (ids: MlxArray, cache: Cache[], options?: TokenWorkOptions) => Promise<MlxArray>;
-export interface MlxPreparationWork {
-  /** Real prompt tokens across all rows, not a padded sequence width. */
-  readonly maxTokens?: number;
-  readonly forward?: MlxForwardWork;
-}
-
-/** Join two method-owned work producers at their forward boundary. Decode
- * finishes sampling/publication/retirement before preparation can join rows.
- * The model owns packing; this coordinator never merges caches or samples.
- * A restored/cancelled preparation need not produce a forward at all. */
-export async function runMixedTokenIteration(options: {
-  prepare(forward: MlxForwardWork): Promise<void>;
-  decode(forward: MlxForwardWork): Promise<void>;
-  forward: MlxForwardWork;
-  mixed(groups: readonly TokenGroup[]): MlxArray[];
-}): Promise<boolean> {
-  let notify!: () => void;
-  const offered = new Promise<void>(resolve => { notify = resolve; });
-  let pending: { group: TokenGroup; resolve(hidden: MlxArray): void; reject(error: unknown): void } | undefined;
-  let captured = false;
-  let preparationError: unknown;
-  let preparationFailed = false;
-  const preparing = options.prepare(async (ids, cache, policy) => {
-    if (captured) return options.forward(ids, cache, policy);
-    captured = true;
-    return new Promise<MlxArray>((resolve, reject) => {
-      pending = { group: { ids, cache, ...policy }, resolve, reject }; notify();
-    });
-  }).catch(error => { preparationFailed = true; preparationError = error; }).finally(notify);
-  await offered;
-  if (!pending) {
-    await preparing;
-    if (preparationFailed) throw preparationError;
-    return false;
-  }
-  const held = pending;
-  let preparedHidden: MlxArray | undefined;
-  let mixed = false;
-  try {
-    await options.decode(async (ids, cache, policy) => {
-      if (mixed) return options.forward(ids, cache, policy);
-      const outputs = options.mixed([{ ids, cache, ...policy }, held.group]);
-      mixed = true; preparedHidden = outputs[1]!;
-      return outputs[0]!;
-    });
-    // A final unread token can retire decode without another model call.
-    preparedHidden ??= await options.forward(held.group.ids, held.group.cache, held.group);
-    held.resolve(preparedHidden); preparedHidden = undefined;
-    await preparing;
-    if (preparationFailed) throw preparationError;
-    return true;
-  } catch (error) {
-    preparedHidden?.dispose(); held.reject(error);
-    await preparing;
-    throw error;
-  }
-}

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

@@ -1,178 +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 MultimodalTowers } 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, PromptNativeWork } from "../../serve/model-binding";
-import { EncoderCache } from "./encoder-cache";
-
-export async function buildModelPrompt(
-    ctx: ServerContext,
-    prep: RequestPrep,
-    body: ChatRequestParams,
-    tools: ChatRequestParams["tools"] | null,
-    ownership: RequestOwnership,
-    nativeWork: PromptNativeWork = (work) => work(),
-    objects?: import("../../contracts/object-cache").ObjectCache<import("./checkpoint-state").CheckpointAttachment[]>,
-  ): 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 { 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 towers = await nativeWork(async () => {
-        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: MultimodalTowers<{ softTokens: number }>["vision"];
-        if (hasImages) {
-          const tower = getVisionTower(ctx);
-          if (!tower) throw new RequestError(400, "model has no vision sidecar");
-          visionSide = { tower, tokenIds: ctx.visionTokenIds,
-            cache: objects && tower.cacheIdentity ? new EncoderCache(objects, tower.cacheIdentity) : undefined };
-        }
-        return {
-          ...(visionSide ? { vision: visionSide } : {}),
-          audio: { tower: audioTower, tokenIds: ctx.audioTokenIds,
-            cache: objects ? new EncoderCache(objects, audioTower.cacheIdentity) : undefined },
-        };
-      });
-      const mp = await buildMultimodalPrompt(
-        ctx.model as Gemma4Model, towers, ctx.tokenizer, ctx.template,
-        messages, images, wavs, toolList, nativeWork,
-      );
-      // 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, prefixIdentity: mp.prefixIdentity,
-          ...(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 nativeWork(() => 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 { messages: withVideos, images } =
-        await extractImages(normalizeMessages(body.messages));
-      const { messages, videos } = await extractVideos(withVideos);
-      const { tower, encoderCache } = await nativeWork(async () => {
-        const tower = getVisionTower(ctx) as unknown as Qwen3VLVisionTower | null;
-        if (!tower) throw new RequestError(400, "model has no vision sidecar");
-        return { tower, encoderCache: objects ? new EncoderCache(objects, tower.cacheIdentity) : undefined };
-      });
-      const vp = await buildQwen3VLVisionPrompt(
-        ctx.model as Qwen35Model, 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, nativeWork, encoderCache,
-      );
-      return { ...noMedia, promptIds: vp.ids, vision: { embeddings: vp.embeddings, mrope: vp.mrope, prefixIdentity: vp.prefixIdentity } };
-    }
-    if (hasImages) {
-      // Loads (and caches) the tower on first image request — text-only
-      // sessions never pay for it.
-      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 { tower, encoderCache } = await nativeWork(async () => {
-        const tower = getVisionTower(ctx);
-        if (!tower) throw new RequestError(400, "model has no vision sidecar");
-        return { tower, encoderCache: objects && tower.cacheIdentity
-          ? new EncoderCache(objects, tower.cacheIdentity) : undefined };
-      });
-      const vp = await buildVisionPrompt(
-        ctx.model as Gemma4Model, tower, ctx.tokenizer, ctx.template,
-        messages, images, ctx.visionTokenIds, toolList, nativeWork, encoderCache,
-      );
-      return { ...noMedia, promptIds: vp.ids, vision: { embeddings: vp.embeddings, imageMask: vp.imageMask, prefixIdentity: vp.prefixIdentity } };
-    }
-    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,80 +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, nativeWork, objects) =>
-    buildModelPrompt(context as ServerContext, prep, body, tools, ownership, nativeWork, objects));
-}
-
-/** 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, nativeWork, objects) => buildModelPrompt(ctx, prep, body, tools, ownership, nativeWork, objects),
-    restore: (store, entry) => store.restore(entry, model),
-    restoreAsync: (store, entry) => store.restoreAsync(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/paged-cache-codec.ts

@@ -1,27 +0,0 @@
-import type { CacheCodec } from "../../kv-store";
-import { BlockPool, PagedKVCache } from "../../lab/paged-kv/paged-kv";
-
-/** Page layout owns its persistence codec, independently of SSD transport. */
-export const pagedCacheCodec: CacheCodec = {
-    matches: cache => cache instanceof PagedKVCache,
-    snapshot(cache, context) {
-      const c = cache as PagedKVCache, p = c.pool!;
-      for (const array of p.arrays()) context.push(array, false);
-      return { kind: "paged", offset: c.offset, tensors: context.slots,
-        ...(p.quantization ? p.quantization : {}),
-        paged: { capacityTokens: c.capacityTokens, blockSize: c.blockSize, blockTable: [...c.blockTable],
-          numBlocks: p.numBlocks, headDim: p.headDim, vHeadDim: p.vHeadDim, dtype: p.dtype, direct: c.direct } };
-    },
-    clone: cache => (cache as PagedKVCache).clone(),
-    load(entry, context) {
-      const layout = entry.paged!;
-      const quantization = entry.bits ? { bits: entry.bits as 4 | 8, groupSize: entry.groupSize! } : undefined;
-      const cache = new PagedKVCache(layout.capacityTokens, layout.blockSize, layout.direct, quantization);
-      const pool = new BlockPool({ ...layout, numKvHeads: entry.tensors[0]!.shape[1]!, quantization });
-      try { pool.restoreArrays(entry.tensors.map(context.arr), layout.blockTable); }
-      catch (error) { pool.dispose(); throw error; }
-      cache.pool = pool; cache.offset = entry.offset; cache.blockTable = [...layout.blockTable];
-      return cache;
-    },
-    headerTrimmable: () => true,
-  };