mlx-bun 0.4.0 → 1.0.0

package/src/engine/model-host.ts

@@ -0,0 +1,513 @@
+// Own one loaded model, its draft and lazy media towers. Application entry
+// points supply the artifact and options; no HTTP or scheduling lives here.
+import type { KvQuantSpec, ModelConfig } from "@mlx-bun/inference/artifacts/config";
+import type { Weights } from "@mlx-bun/inference/artifacts";
+import type { RuntimeModel, RuntimeOpenOptions, ResolvedModelProfile, ModelImplementationProvider, GenerationDefaults,
+  CompositionRequest } from "@mlx-bun/inference/models";
+import type { ChatTemplate, LoadedTokenizer, SentinelTokens } from "@mlx-bun/inference/input";
+import type { AdapterManager } from "@mlx-bun/inference/adapters";
+import type { AudioTokenIds, VisionTokenIds, VisionEncoder } from "@mlx-bun/inference/input/vision";
+import type { AudioEncoder } from "@mlx-bun/inference/contracts/mlx";
+import type { MediaPreparation } from "./media-preparation";
+
+/** An audio encoder the context owns; released with it. */
+export type LoadedAudioEncoder = AudioEncoder & { dispose?(): void };
+import { sidecarShipsAudioTower } from "@mlx-bun/hub/registry";
+import { fit } from "@mlx-bun/inference/execution/fit";
+import { cleanupFailure, disposeResources } from "@mlx-bun/inference/runtime/resources";
+import type { Composition, DisposableResource, MemoryPlan } from "@mlx-bun/inference/contracts/portable";
+import type { DraftKind } from "@mlx-bun/inference/generation/speculative/draft-kind";
+import type { DraftProviderRegistry, LoadedDraft } from "@mlx-bun/inference/generation/speculative/draft-registry";
+
+export interface ServedModelInfo { readonly config: ModelConfig; readonly weightsBytes: number; }
+
+export interface AdapterService extends Pick<AdapterManager, keyof AdapterManager> {}
+
+export interface ModelContext<Model = RuntimeModel> {
+  /** Releases owned artifacts, draft and loaded towers after execution stops. */
+  dispose(): void;
+  /** State serialization supplied by this model/backend implementation. */
+  stateCodecs?: import("@mlx-bun/inference/state").CacheCodecProvider;
+  model: Model;
+  /** Declared external artifact/family profile that selected model
+   * construction. Request-level methods are resolved separately. */
+  profile: ResolvedModelProfile;
+  tokenizer: LoadedTokenizer;
+  template: ChatTemplate | null;
+  /** The model's declared marker tokens for tool calls and the reasoning
+   * channel, resolved against `tokenizer`; absent when generated text carries
+   * them as ordinary text. */
+  sentinels?: SentinelTokens | null;
+  /** The model's declared chat-generation defaults; absent means the generic ones. */
+  generationDefaults?: GenerationDefaults;
+  modelId: string;
+  /** Lazily-loaded vision tower cache — null until the first image request
+   *  (see `getVisionTower`). The tower (SigLIP ~hundreds of MB, encoder-free
+   *  smaller) is not loaded for text-only sessions. */
+  vision: VisionEncoder | null;
+  /** Loads + selects the vision tower on demand; null when the model has no
+   *  (supported) vision sidecar. Invoked at most once, then cached in
+   *  `vision`. */
+  loadVision: (() => VisionEncoder) | null;
+  /** Null when neither the config nor the model's declaration supplies the image soft tokens. */
+  visionTokenIds: VisionTokenIds | null;
+  /** Lazily-loaded Conformer audio tower — null until the first audio
+   *  request (see `getAudioTower`). Same sidecar file as vision
+   *  (optiq_vision.safetensors), separate tower; text-only sessions never
+   *  pay for it. */
+  audio: LoadedAudioEncoder | null;
+  /** Loads the audio tower on demand; null when the model can't do audio
+   *  (no `audio_config` in config.json or no sidecar). No flags — audio
+   *  auto-enables exactly like vision. */
+  loadAudio: (() => LoadedAudioEncoder) | null;
+  /** null when the model has no `audio_config` (audio-incapable). */
+  audioTokenIds: AudioTokenIds | null;
+  /** Family media preparation, bound once at load; borrows this context. */
+  media: MediaPreparation;
+  adapters: AdapterService;
+  /** Per-layer KV quantization from the repo's kv_config.json (null if
+   *  absent): what `--kv-quant config` composes into the loaded model's KV
+   *  scheme (`composition.kv`). */
+  kvConfig: KvQuantSpec[] | null;
+  /** Model-author recommended sampling from generation_config.json —
+   *  optiq serve injects these as server defaults (gen_config.py);
+   *  explicit request fields always win. */
+  genDefaults: GenSamplingDefaults;
+  /** Selected draft provider; null means no draft configured. */
+  draft?: {
+    provider: import("@mlx-bun/inference/generation/speculative").DraftProvider;
+    numDraftTokens: number;
+    /** Each round's draft count comes from the graph's verify costs (`--num-draft-tokens adaptive`). */
+    adaptiveDepth?: boolean;
+    /** The adaptive gate's settings the graph measured for this drafter's kind (`adaptiveDraftGates`). */
+    gate?: import("@mlx-bun/inference/contracts/mlx").AdaptiveDraftGateSettings;
+    /** The provider is the checkpoint's own draft head, not a configured drafter. */
+    native?: true;
+    /** The registered kind that loaded it (a graph may declare an adaptive width for it). */
+    kind?: string;
+  } | null;
+  /** The composition this load resolved before building the graph: device,
+   * KV scheme, draft depth, prefill chunk and
+   * rows per forward. Absent for a context an engine-owned implementation or a
+   * caller built. */
+  readonly composition?: Composition;
+  /** Present only for a runtime that plans its memory before opening weights:
+   * the exact header-derived process equation it runs under. */
+  memoryPlan?: MemoryPlan | null;
+  /** Live telemetry of such a runtime, surfaced by `/stats`. */
+  runtimeDiagnostics?(): Record<string, unknown>;
+}
+
+export type LoadedModelContext = ModelContext<ServedModelInfo>;
+
+export interface ModelHostSource {
+  readonly modelDir: string;
+  readonly modelId: string | undefined;
+  readonly options: Omit<LoadContextOptions, "implementations">;
+}
+
+export interface LoadContextOptions<Model extends ServedModelInfo = RuntimeModel> {
+  /** Raw generation and embeddings do not require a chat template. Default true. */
+  requireChatTemplate?: boolean;
+  /** Engine-owned complete implementations may own their loader and methods.
+   * Selection happens before opening the default resident or streamed weights.
+   * They receive these options in their source and resolve no composition here. */
+  implementations?: ModelImplementationProvider<ModelHostSource, Promise<ModelContext<Model>>, void>;
+  profiles?: import("@mlx-bun/inference/models").ResolveModelProfileOptions;
+  memoryBudgetBytes?: number;
+  /** The KV flags the composition's scheme resolves from (`--kv-quant`, its
+   * TurboQuant spelling, `--quantized-kv-start`); absent means bf16. */
+  kv?: CompositionRequest["kv"];
+  /** Paged KV (`--paged-kv`, `--paged-kv-block-size`); absent means unpaged. */
+  pagedKv?: CompositionRequest["pagedKv"];
+  /** Most rows one forward carries (`--batch`). Default 8. */
+  maxRows?: number;
+  /** Resource overrides for runtimes that plan memory up front; models
+   * without one ignore them. */
+  runtime?: RuntimeOpenOptions;
+  /** Snapshot dir of a draft model for speculative decoding
+   * (`--draft-model`). Loaded alongside the target; the pair must share
+   * a tokenizer family. */
+  draftModelDir?: string;
+  /** Drafts per round (`--num-draft-tokens`, mlx_lm.server default 3). */
+  numDraftTokens?: number;
+  /** `--num-draft-tokens adaptive`: choose each round's count, up to numDraftTokens,
+   * from the graph's measured verify costs. Refused without a drafter or costs. */
+  adaptiveDraftDepth?: boolean;
+  /** Draft-provider kind override (`--draft-kind`); any kind the registry knows. */
+  draftKind?: DraftKind | (string & {});
+  /** The draft providers this load may select and detect among; defaults to
+   * the library's. Extend a registry to add a provider kind. */
+  draftProviders?: DraftProviderRegistry;
+  /** A ready draft provider and its round width, instead of loading one from
+   * `draftModelDir`/`draftKind`. The context owns it from the call on. */
+  draftProvider?: LoadedDraft;
+  ngramMax?: number;
+  ngramMin?: number;
+}
+
+export interface GenSamplingDefaults {
+  temperature?: number;
+  topP?: number;
+  topK?: number;
+  repetitionPenalty?: number;
+}
+
+export { detectDraftKind, type DraftKind } from "@mlx-bun/inference/generation/speculative/draft-kind";
+
+/** Serving requires a template; one-shot consumers retain main's raw fallback. */
+export async function loadContextTemplate(modelDir: string, required: boolean,
+  load: (directory: string) => Promise<ChatTemplate>): Promise<ChatTemplate | null> {
+  try { return await load(modelDir); }
+  catch (error) { if (required) throw error; return null; }
+}
+
+export function requireChatTemplate<T extends { template: ChatTemplate | null; modelId: string }>(context: T): asserts context is T & { template: ChatTemplate } {
+  if (!context.template) throw new Error(`model ${context.modelId} has no chat template`);
+}
+
+async function loadGenSamplingDefaults(modelDir: string): Promise<GenSamplingDefaults> {
+  const file = Bun.file(`${modelDir}/generation_config.json`);
+  if (!(await file.exists())) return {};
+  try {
+    const raw = (await file.json()) as Record<string, unknown>;
+    const num = (v: unknown) => (typeof v === "number" ? v : undefined);
+    return {
+      temperature: num(raw.temperature),
+      topP: num(raw.top_p),
+      topK: num(raw.top_k),
+      repetitionPenalty: num(raw.repetition_penalty),
+    };
+  } catch {
+    return {};
+  }
+}
+
+export function loadContext<Model extends ServedModelInfo>(
+  modelDir: string, modelId: string | undefined,
+  opts: LoadContextOptions<Model> & {
+    implementations: ModelImplementationProvider<ModelHostSource, Promise<ModelContext<Model>>, void>;
+  },
+): Promise<ModelContext<Model>>;
+export function loadContext(modelDir: string, modelId?: string, opts?: LoadContextOptions): Promise<ModelContext>;
+export function loadContext<Model extends ServedModelInfo>(
+  modelDir: string, modelId: string | undefined, opts: LoadContextOptions<Model>,
+): Promise<LoadedModelContext>;
+export async function loadContext(
+  modelDir: string, modelId?: string,
+  opts: LoadContextOptions<ServedModelInfo> = {},
+): Promise<LoadedModelContext> {
+  const { loadModelConfig } = await import("@mlx-bun/inference/artifacts/config");
+  const { resolveModelProfile } = await import("@mlx-bun/inference/models/profile");
+  const config = await loadModelConfig(modelDir);
+  const profile = resolveModelProfile(config, opts.profiles);
+  if (opts.implementations) {
+    const implementation = opts.implementations.select(config, profile);
+    const { implementations: _provider, ...options } = opts;
+    return implementation.create({ modelDir, modelId, options }, config, profile);
+  }
+  const [{ Weights }, { createModel, declaredGraph, openPlannedRuntime, plansMemory, loadModelChatTemplate,
+      sentinelDeclarationFor, mediaTokenDeclarationFor, generationDefaultsFor, resolveComposition, stateCodecsOf },
+    { loadTokenizer, resolveSentinelTokens, resolveVisionTokenIds, resolveAudioTokenIds },
+    { AdapterManager }, { bindLegacyDraftTarget }, { defaultDraftProviders }] = await Promise.all([
+    import("@mlx-bun/inference/artifacts"), import("@mlx-bun/inference/models"),
+    import("@mlx-bun/inference/input"),
+    import("@mlx-bun/inference/adapters"), import("@mlx-bun/inference/generation/speculative/binding"),
+    import("@mlx-bun/inference/generation/speculative/draft-kind"),
+  ]);
+  const planned = plansMemory(profile);
+  const draftRegistry = opts.draftProviders ?? defaultDraftProviders();
+  if (opts.draftProvider && (opts.draftModelDir !== undefined || opts.draftKind !== undefined))
+    throw new Error("a ready draft provider cannot combine with --draft-model/--draft-kind");
+  // A kind that bundles its companion in the target artifact (`--draft-kind
+  // mtp`: the artifact's own mtp/ subfolder — single-repo packaging, the
+  // companion is a complete model dir the provider already loads) resolves
+  // there when no --draft-model is given. Explicit --draft-model still wins; a
+  // missing bundle is a clear refusal.
+  const bundledCompanion = opts.draftKind === undefined ? undefined : draftRegistry.get(opts.draftKind)?.bundledCompanion;
+  if (bundledCompanion && !opts.draftModelDir) {
+    const bundled = `${modelDir}/${bundledCompanion}`;
+    if (await Bun.file(`${bundled}/config.json`).exists()) {
+      opts = { ...opts, draftModelDir: bundled };
+    } else {
+      throw new Error(
+        `--draft-kind ${opts.draftKind} needs a companion: pass --draft-model <dir> or use ` +
+        `an artifact that bundles one at <model>/${bundledCompanion}/ (none at ${bundled})`,
+      );
+    }
+  }
+  const externalDraft = opts.draftModelDir !== undefined || opts.draftKind !== undefined || opts.draftProvider !== undefined;
+  const resolvedDraftKind = opts.draftModelDir
+    ? opts.draftKind ?? await draftRegistry.detect(opts.draftModelDir)
+    : opts.draftKind;
+  if (planned && externalDraft && opts.runtime?.nativeDraft === true)
+    throw new Error("the checkpoint-native draft head and --draft-model/--draft-kind are mutually exclusive");
+  const owned = new Set<DisposableResource>();
+  const release = (resource: DisposableResource) => { owned.delete(resource); resource.dispose(); };
+  try {
+    let weights: Weights | null = null;
+    let model!: RuntimeModel;
+    // A runtime that plans memory opens its own graph and owns its native draft head.
+    let runtime: Awaited<ReturnType<typeof openPlannedRuntime>> | null = null;
+    if (planned) {
+      const nativeDraft = !externalDraft && opts.runtime?.nativeDraft !== false;
+      const plannedCaches = resolveComposition(config, { kv: opts.kv, pagedKv: opts.pagedKv });
+      runtime = await openPlannedRuntime(modelDir, profile, {
+        ...opts.runtime,
+        memoryBudgetBytes: opts.memoryBudgetBytes ?? opts.runtime?.memoryBudgetBytes,
+        // Preserve the existing loader's draft residency equation.
+        batchSize: nativeDraft || externalDraft ? 1 : opts.runtime?.batchSize,
+        // An explicit alternate drafter replaces the native head.
+        nativeDraft,
+        // Its caches compose the KV scheme and paging the composition resolves below from the same flags.
+        kv: plannedCaches.kv, pagedKv: plannedCaches.pagedKv,
+      });
+      model = runtime.model;
+      owned.add(runtime.model);
+    } else {
+      weights = await Weights.open(modelDir);
+      owned.add(weights);
+    }
+    // memoryBudget enforcement at load: Weights.open only mmaps
+    // (no GPU allocation yet), so a model whose weights can never serve
+    // within the budget is refused HERE — before any unified-memory
+    // commitment — with an actionable error instead of a Metal OOM later.
+    if (!planned && opts.memoryBudgetBytes) {
+      const weightsBytes = [...weights!.shards.files.values()]
+        .reduce((a, f) => a + f.mmap.size, 0);
+      const report = fit(config, weightsBytes, 1, undefined, undefined, 0, opts.memoryBudgetBytes);
+      if (report.maxSafeContext < 1)
+        throw new Error(
+          `model does not fit the memory budget: weights ${(weightsBytes / 1e9).toFixed(2)} GB ` +
+          `+ prefill transient leave no room for any context within ` +
+          `${(opts.memoryBudgetBytes / 1e9).toFixed(2)} GB`,
+        );
+    }
+    const tokenizer = await loadTokenizer(modelDir);
+
+    // Speculative decoding: load the draft (mlx_lm.server --draft-model). The
+    // draft artifact's KIND selects the provider, each of which recognizes and
+    // loads its own artifact; all share ONE serve loop
+    // (packages/inference/src/generation/speculative/run.ts). `--draft-kind`
+    // overrides the detect. The loaded drafter settles the composition's draft
+    // depth, so it loads before the target graph is built; the pairing probe
+    // below needs that graph.
+    let draft: ModelContext["draft"] = null;
+    let validateDraft = true;
+    if (opts.draftProvider) {
+      draft = opts.draftProvider;
+      owned.add(draft.provider);
+    } else if (resolvedDraftKind !== undefined) {
+      const kind = draftRegistry.get(resolvedDraftKind);
+      if (kind && !kind.artifact && opts.draftModelDir)
+        throw new Error(`--draft-kind ${resolvedDraftKind} is model-free — drop --draft-model (it would be ignored)`);
+      // Every artifact kind names an artifact to load — refuse instead of
+      // silently serving without speculation.
+      if (kind?.artifact && !opts.draftModelDir)
+        throw new Error(`--draft-kind ${resolvedDraftKind} requires --draft-model`);
+      draft = await draftRegistry.load(resolvedDraftKind, {
+        dir: opts.draftModelDir, target: { vocabSize: config.text.vocabSize, tokenizer },
+        numDraftTokens: opts.numDraftTokens, ngram: { max: opts.ngramMax, min: opts.ngramMin },
+      });
+      owned.add(draft.provider);
+      // A model-free provider has no pairing to probe and no weights to budget.
+      validateDraft = kind!.artifact;
+    }
+    // A checkpoint-native draft head (a planned runtime's is the production
+    // default) comes from the graph's own declaration: it uses the
+    // already-planned bounded auxiliary tier and the same tokenizer, so there is
+    // no second artifact or compatibility probe to load. A planned runtime has
+    // opened its graph above, so its head settles the depth here; a resident
+    // graph declares one only once it is built, below.
+    const declaredDraftHead = async (): Promise<ModelContext["draft"]> => {
+      const declared = await draftRegistry.native(model);
+      return declared ? { ...declared, native: true } : null;
+    };
+    if (!draft && planned) draft = await declaredDraftHead();
+    // Resolved once, before the resident graph is built, which receives it. A
+    // planned runtime built its graph from its own memory plan above.
+    const composition = resolveComposition(config, { kv: opts.kv, pagedKv: opts.pagedKv, maxRows: opts.maxRows,
+      draft: draft && { numDraftTokens: draft.numDraftTokens, adaptive: opts.adaptiveDraftDepth } });
+    if (!planned) {
+      model = createModel(weights!, config, composition, profile);
+      owned.add(model);
+    }
+    // Generation must stop on the tokenizer's eos_token — the chat turn
+    // terminator (e.g. Qwen <|im_end|> = 248046). Some configs (Qwen3.5-4B)
+    // declare a different eos_token_id in config.json than the chat format
+    // emits, so without this a turn never ends and generation runs away,
+    // hallucinating both sides of the dialogue until max_tokens. mlx-lm stops on
+    // the tokenizer eos; union it in. No-op when already present (Gemma, 27B).
+    if (tokenizer.eosTokenId != null && !config.eosTokenIds.includes(tokenizer.eosTokenId))
+      config.eosTokenIds = [...config.eosTokenIds, tokenizer.eosTokenId];
+
+    if (draft && validateDraft) {
+      const { provider } = draft;
+      // Fail-fast pairing validation (2026-07-07 review): a provider's grouped
+      // rows validate the (target, drafter) pairing when they open — a target
+      // that lacks the port they read, a DeepSpec or DFlash 2 target-layer-count
+      // mismatch — so a mismatch refuses at load instead of failing every drafted
+      // request. Opening prefill rows with no members allocates no row tensors.
+      {
+        const probeCaches = model.makeCache();
+        try {
+          provider.grouped.openPrefill({ target: bindLegacyDraftTarget(model, probeCaches), checkpoints: [] }).dispose();
+        } catch (err) {
+          release(provider);
+          throw new Error(
+            `--draft-model is incompatible with this target: ${err instanceof Error ? err.message : String(err)}`,
+          );
+        } finally {
+          for (const c of probeCaches) c.dispose();
+        }
+      }
+      if (opts.memoryBudgetBytes) {
+        const targetBytes = weights
+          ? [...weights.shards.files.values()].reduce((a, f) => a + f.mmap.size, 0)
+          : model.weightsBytes;
+        // Draft weights shrink the target's envelope. Draft KV is not modeled
+        // (small relative to its weights at serve contexts); admission stays
+        // approximately conservative via the combined-weights term.
+        const report = fit(config, targetBytes + provider.weightsBytes, 1, undefined, undefined, 0, opts.memoryBudgetBytes);
+        if (report.maxSafeContext < 1) {
+          release(provider);
+          throw new Error(
+            `target + draft do not fit the memory budget (draft adds ` +
+              `${(provider.weightsBytes / 1e9).toFixed(2)} GB)`,
+          );
+        }
+      }
+    }
+
+    if (!draft && !planned) draft = await declaredDraftHead();
+
+    if (draft) owned.add(draft.provider);
+    if (opts.adaptiveDraftDepth) {
+      if (!draft) throw new Error("--num-draft-tokens adaptive needs a drafter (--draft-model, --draft-kind or the checkpoint's MTP head)");
+      const graph = declaredGraph(model);
+      if (!graph.verifyRoundCosts)
+        throw new Error("--num-draft-tokens adaptive needs a graph measured on this machine (one that declares its verify costs); this model's graph does not");
+      // A drafter kind whose adaptive width the graph measured drafts up to it, unless the request named a width, and
+      // gates with the settings the graph measured for it.
+      const width = opts.numDraftTokens === undefined && draft.kind !== undefined ? graph.adaptiveDraftWidths?.get(draft.kind) : undefined;
+      const gate = draft.kind !== undefined ? graph.adaptiveDraftGates?.get(draft.kind) : undefined;
+      draft = { ...draft, adaptiveDepth: true, ...(width === undefined ? {} : { numDraftTokens: width }), ...(gate ? { gate } : {}) };
+    }
+    const adapters = new AdapterManager(model);
+    owned.add({ dispose() { disposeResources(adapters.list().map(({ id }) => ({ dispose() { adapters.unmount(id); } }))); } });
+    // A compiled decode step borrows graph constants and weights. Retire it first,
+    // after the owning engine has drained, including if later loading fails.
+    owned.add({ dispose() { declaredGraph(model).releaseCompiledDecode?.(); } });
+    const template = await loadContextTemplate(modelDir, opts.requireChatTemplate ?? true, dir => loadModelChatTemplate(dir, profile));
+    const id = modelId ?? modelDir.split("/").filter(Boolean).at(-1)!;
+    // Soft-token ids: the checkpoint's config, else the token text the profile declares.
+    const mediaTokens = mediaTokenDeclarationFor(profile);
+    const visionTokenIds = resolveVisionTokenIds(config.raw, tokenizer, mediaTokens?.vision ?? null);
+    const audioTokenIds = resolveAudioTokenIds(config.raw, tokenizer, mediaTokens?.audio ?? null);
+    const sentinelTexts = sentinelDeclarationFor(profile);
+    const { bindMediaPreparation } = await import("./media-preparation");
+    // Encoders load lazily (getVisionTower/getAudioTower): text-only sessions
+    // never pay for a tower. The graph knows which its checkpoint ships.
+    const encoders = await declaredGraph(model).mediaEncoders?.(modelDir, { shipsAudioTower: sidecarShipsAudioTower })
+      ?? { vision: null, audio: null };
+    const stateCodecs = stateCodecsOf(model);
+    const context: Omit<ModelContext, "dispose"> = {
+      ...(stateCodecs ? { stateCodecs } : {}),
+      draft,
+      model,
+      profile,
+      sentinels: sentinelTexts ? resolveSentinelTokens(tokenizer, sentinelTexts) : null,
+      generationDefaults: generationDefaultsFor(profile),
+      composition,
+      memoryPlan: runtime?.memoryPlan ?? null,
+      ...(runtime ? { runtimeDiagnostics: runtime.diagnostics } : {}),
+      adapters,
+      kvConfig: config.kvQuant,
+      genDefaults: await loadGenSamplingDefaults(modelDir),
+      tokenizer,
+      template,
+      modelId: id,
+      vision: null,
+      loadVision: encoders.vision,
+      visionTokenIds,
+      // Audio mirrors vision: loaded on the first audio request only
+      // (`02d723a:docs/design/generic-model-support.md` §6.6).
+      audio: null,
+      loadAudio: encoders.audio,
+      audioTokenIds,
+      // The family route is chosen once; it borrows the towers through this
+      // context's lazy slots, so the context stays their only owner.
+      media: await bindMediaPreparation({
+        modelId: id, tokenizer, template, visionTokenIds, audioTokenIds,
+        visionTower: () => getVisionTower(context), audioTower: () => getAudioTower(context),
+      }, model),
+    };
+    return ownModelContext(context, [...owned].reverse());
+  } catch (error) { return cleanupFailure(error, () => disposeResources([...owned].reverse())); }
+}
+
+/** Lazily load + cache the vision tower on first use. A sidecar that fails
+ *  to load is a capability gap, not a fatal error: returns null and the
+ *  request is answered with a 400 (the loader is cleared so we don't retry
+ *  a known-bad load every request). */
+export function getVisionTower(ctx: Pick<ModelContext<unknown>, "vision" | "loadVision">): VisionEncoder | null {
+  if (ctx.vision) return ctx.vision;
+  if (!ctx.loadVision) return null;
+  try {
+    ctx.vision = ctx.loadVision();
+    return ctx.vision;
+  } catch (e) {
+    console.warn(`vision sidecar not loadable (${(e as Error).message}) — serving text-only`);
+    ctx.loadVision = null;
+    return null;
+  }
+}
+
+/** Lazily load + cache the audio tower on first use. Unlike vision's
+ *  warn-and-continue (text-only degrade is correct for requests WITHOUT
+ *  images), this is only ever consulted for requests WITH audio — the
+ *  caller turns null into an explicit 400, never a silent text-only
+ *  degrade. A failed load is not retried every request. (The stub-sidecar
+ *  case — the local 12B state — never gets here: the graph checks the
+ *  sidecar header and returns a null loader.) */
+export function getAudioTower(ctx: Pick<ModelContext<unknown>, "audio" | "loadAudio">): LoadedAudioEncoder | null {
+  if (ctx.audio) return ctx.audio;
+  if (!ctx.loadAudio) return null;
+  try {
+    ctx.audio = ctx.loadAudio();
+    return ctx.audio;
+  } catch (e) {
+    console.warn(`audio sidecar not loadable (${(e as Error).message})`);
+    ctx.loadAudio = null;
+    return null;
+  }
+}
+
+/** Who releases a served context. "owned": its server disposes it once, after
+ * execution drains, including when startup fails. "borrowed": its server never
+ * disposes it; the caller keeps it usable and disposes it after the server closed. */
+export type ContextOwnership = "owned" | "borrowed";
+
+/** The one release rule every serving composition applies to its context. */
+export function releaseContext(context: { dispose(): void }, ownership: ContextOwnership): void {
+  if (ownership === "owned") context.dispose();
+}
+
+/** Attach an idempotent owner to a loaded context. Supplied resources are owned,
+ * never borrowed; all are attempted even when an individual cleanup fails. */
+export function ownModelContext<Model extends ServedModelInfo>(
+  context: Omit<ModelContext<Model>, "dispose">, resources: readonly DisposableResource[],
+): ModelContext<Model> {
+  const owned = [...resources];
+  let disposed = false;
+  return Object.assign(context, {
+    dispose() {
+      if (disposed) return;
+      disposed = true;
+      const media = [context.vision, context.audio].filter((value): value is NonNullable<typeof value> => value != null);
+      context.vision = null; context.audio = null;
+      context.loadVision = null; context.loadAudio = null;
+      disposeResources([...media.map(value => ({ dispose() { value.dispose?.(); } })), ...owned]);
+    },
+  });
+}

package/src/engine/model-serving.ts

@@ -0,0 +1,89 @@
+import type { MlxHeadForward, MlxHeadRange, MlxSequenceRange, MlxSequenceForward } from "@mlx-bun/inference/contracts/mlx";
+import type { Composition } from "@mlx-bun/inference/contracts/portable";
+import type { ModelContext, LoadedModelContext } from "./model-host";
+import type { ModelBinding } from "./model-binding";
+
+/** The loaded composition as `/stats` reports it: the record's fields, with
+ * snake_case keys like the rest of that view. */
+function compositionWire(composition: Composition): unknown {
+  const wire = (value: unknown): unknown => Array.isArray(value) ? value.map(wire)
+    : value !== null && typeof value === "object"
+      ? Object.fromEntries(Object.entries(value).map(([key, field]) => [key.replace(/[A-Z]/g, c => `_${c.toLowerCase()}`), wire(field)]))
+      : value;
+  return wire(composition);
+}
+/** Existing model classes enter the interface once here. New implementations
+ * supply a binding explicitly and need no forward/makeCache methods. */
+export async function modelServingBinding(context: LoadedModelContext, supplied?: ModelBinding): Promise<ModelBinding> {
+  if (supplied) return supplied;
+  const [{ declaredGraph, embeddingDeclarationFor }, { embedMany, embeddingTerminatorId, isEmbeddingModel }, { bindMlxGateway }, { MLX_VERSION, deviceArchitecture }, ops, { sequenceHiddenOperation, sequenceLogitsOperation, scoreDecisionAnswers }] = await Promise.all([
+    import("@mlx-bun/inference/models"), import("@mlx-bun/inference/embeddings"),
+    import("@mlx-bun/inference/execution"), import("@mlx-bun/mlx/ffi"), import("@mlx-bun/mlx/ops"), import("@mlx-bun/inference/scoring"),
+  ]);
+  const ctx = context as ModelContext;
+  const model = ctx.model;
+  const scoring = model as { readonly sequence?: readonly MlxSequenceRange[]; readonly sequenceHidden?: MlxSequenceForward;
+    readonly sequenceLogits?: readonly MlxHeadRange[]; readonly logitsFromHidden?: MlxHeadForward };
+  if (typeof model.makeCache !== "function")
+    throw new Error("model implementation must supply its serving binding");
+  const graph = declaredGraph(model);
+  const hasSequence = !!scoring.sequence?.some(range => typeof range.operation === "function") || typeof scoring.sequenceHidden === "function";
+  const hasHead = !!scoring.sequenceLogits?.some(range => typeof range.operation === "function") || typeof scoring.logitsFromHidden === "function";
+  const adapters = typeof model.loraTargets === "function" && model.loraTargets().size > 0;
+  const gateway = bindMlxGateway(model, ctx.draft ?? undefined);
+  const capabilities = gateway.capabilities;
+  const embedding = isEmbeddingModel(model);
+  const declaration = embedding ? embeddingDeclarationFor(ctx.profile) : null;
+  if (embedding && !declaration) throw new Error(`profile ${ctx.profile.profile.id} declares no embedding recipe for its embedding graph`);
+  const terminator = declaration ? embeddingTerminatorId(ctx.tokenizer, declaration.terminator) : 0;
+  return {
+    stateCompatibility: `mlx-${MLX_VERSION}-${deviceArchitecture()}`,
+    gateway,
+    async signal(ids, count, minimum) {
+      if (!hasSequence || !hasHead) throw new TypeError("model declares no full-sequence scoring operation for signal");
+      const caches = model.makeCache();
+      try {
+        // The whole prompt as one window (the scoring forward), the head over every position.
+        const input = ops.fromInt32(ids, [1, ids.length]);
+        let hidden;
+        try { hidden = sequenceHiddenOperation(scoring, input, caches, ids.length); } finally { input.dispose(); }
+        let logits;
+        try { logits = sequenceLogitsOperation(model, hidden); } finally { hidden.dispose(); }
+        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, training: adapters && ((hasSequence && hasHead) || typeof graph.trainable?.denoising?.canvasLogits === "function"),
+      dsa: capabilities.sparseAttention, embeddings: embedding, decisions: hasSequence && hasHead },
+    ...(embedding ? { embed: (inputs: string[], instruction?: string) => embedMany(model, ctx.tokenizer, terminator, inputs, instruction) } : {}),
+    ...(hasSequence && hasHead ? { decide: async ({ prompt, candidates, vision, signal, execute }: Parameters<NonNullable<ModelBinding["decide"]>>[0]) => {
+      if (vision && !graph.bindMediaInput) throw new TypeError("graph declares no media input for decisions");
+      return (await scoreDecisionAnswers({ makeCache: () => model.makeCache(),
+        forward: (ids, cache, rows) => sequenceHiddenOperation(scoring, ids, cache, rows),
+        logits: hidden => sequenceLogitsOperation(model, hidden) }, prompt, candidates,
+        { signal, execute, ...(vision ? { input: graph.bindMediaInput!(vision) } : {}) })).probabilities;
+    } } : {}),
+    diagnostics() {
+      const runtime = ctx.runtimeDiagnostics?.();
+      return { ...(ctx.composition ? { composition: compositionWire(ctx.composition) } : {}), ...(runtime ? { runtime } : {}) };
+    },
+  };
+}

package/src/{serve → engine}/preparation.ts

@@ -1,12 +1,12 @@
 /** Host-side admission before grammar/media preparation can allocate native
  * resources. The implementation shares the generation execution domain. */
 export interface PreparationExecutor {
-  reserve?(kind: "media" | "constraint", signal?: AbortSignal): Promise<import("../contracts/resources").DisposableResource>;
+  reserve?(kind: "media" | "constraint" | "decision", signal?: AbortSignal): Promise<import("@mlx-bun/inference/contracts/portable").DisposableResource>;
   run<T>(prepare: () => Promise<T>, signal?: AbortSignal): Promise<T>;
 }
 
-import { AdmissionPool } from "../engine/admission";
-import { CancellationSource } from "../engine/cancellation";
+import { AdmissionPool } from "@mlx-bun/inference/execution/admission";
+import { CancellationSource } from "@mlx-bun/inference/execution/cancellation";
 
 export async function acquireReservation(pool: Pick<AdmissionPool, "acquire">, signal?: AbortSignal) {
   signal?.throwIfAborted();
@@ -14,20 +14,22 @@ export async function acquireReservation(pool: Pick<AdmissionPool, "acquire">, s
   const abort = () => cancellation.cancel("requested");
   signal?.addEventListener("abort", abort, { once: true });
   try { return await pool.acquire(cancellation); }
+  catch (error) { throw signal?.aborted ? signal.reason : error; }
   finally { signal?.removeEventListener("abort", abort); }
 }
 
-/** One retained media preparation at a time; grammar rows may form a batch. */
+/** Bound retained preparations independently from the execution queue. */
 export function createPreparationExecutor(
-  run: PreparationExecutor["run"], grammarCapacity: number,
+  run: PreparationExecutor["run"], grammarCapacity: number, mediaCapacity = 1,
 ): PreparationExecutor & { close(): void } {
-  const media = new AdmissionPool(1);
+  const media = new AdmissionPool(Math.max(1, mediaCapacity));
   const constraints = new AdmissionPool(Math.max(1, grammarCapacity));
+  const decisions = new AdmissionPool(Math.max(1, grammarCapacity));
   return {
     run,
     async reserve(kind, signal) {
-      return acquireReservation(kind === "media" ? media : constraints, signal);
+      return acquireReservation(kind === "media" ? media : kind === "decision" ? decisions : constraints, signal);
     },
-    close() { media.close(); constraints.close(); },
+    close() { media.close(); constraints.close(); decisions.close(); },
   };
 }

package/src/{serve → engine}/session-completion-engine.ts

@@ -1,9 +1,9 @@
-import { cleanupFailure } from "../engine/resources";
-import type { ExecutionPlan, InferenceMethod, Timer } from "../contracts/generation";
-import { CancellationSource, throwIfCancelled } from "../engine/cancellation";
-import { createInferenceEngine } from "../engine/engine";
-import type { GenerateStats } from "../generate";
-import type { CompletionEngine } from "./completion-executor";
+import { cleanupFailure } from "@mlx-bun/inference/runtime/resources";
+import type { ExecutionPlan, InferenceMethod, Timer } from "@mlx-bun/inference/contracts/portable";
+import { CancellationSource, throwIfCancelled } from "@mlx-bun/inference/execution/cancellation";
+import { createInferenceEngine } from "@mlx-bun/inference/execution/engine";
+import type { GenerateStats } from "@mlx-bun/inference/generation";
+import type { CompletionEngine } from "./completion";
 
 const timer: Timer = {
   after(ms, callback) {