@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/models/factory.ts

@@ -0,0 +1,351 @@
+// Model construction consumes one declared profile and one composition. Exact
+// artifact profiles outrank family profiles; the profile names the family's
+// implementation, and the family's graph table picks the graph for the
+// composition. Request-level methods are resolved elsewhere and are never
+// changed by this module.
+
+import { loadModelConfig, type ModelConfig } from "../artifacts/config";
+import { totalmem } from "node:os";
+import { Weights } from "../artifacts/weights";
+import { Gemma4Model } from "./gemma4/model";
+import { configFingerprint } from "../artifacts/fingerprint";
+import { qwen35GraphFingerprint } from "./qwen/checkpoint";
+import {
+  DECODE_ROWS as QWEN38_M4PRO_DECODE_ROWS, QWEN38_TRELLIS_M4PRO_GRAPH, Qwen38TrellisM4Pro,
+  anePrograms as qwen38TrellisM4ProAnePrograms,
+} from "./qwen/qwen38-27b-trellis-m4pro";
+import {
+  QWEN38_TRELLIS_M1MAX_GRAPH, Qwen38TrellisM1Max,
+  anePrograms as qwen38TrellisM1MaxAnePrograms,
+} from "./qwen/qwen38-27b-trellis-m1max";
+import { GEMMA_GRAPHS, type Gemma4Graph } from "./gemma4/graphs";
+import { MiniCPM5Model } from "./minicpm5/model";
+import { Qwen35Model } from "./qwen/qwen3_5";
+import { Qwen3Model } from "./qwen/qwen3";
+import { Qwen3MoeModel } from "./qwen/qwen3-moe";
+import { DiffusionGemmaModel } from "./diffusion-gemma/model";
+import { Glm52Model, type Glm52LoadedModel, type Glm52LatentCaches } from "./glm52/model";
+import {
+  GLM52_G5_DEFAULT_CONTEXT_TOKENS,
+  GLM52_G5_DEFAULT_MAX_GENERATION_TOKENS,
+  GLM52_G5_DEFAULT_PROCESS_LIMIT_BYTES,
+  planGlm52MemoryForArtifact,
+  type Glm52MemoryPlan,
+} from "../artifacts/glm52/memory";
+import { UniversalDenseModel } from "./universal/dense";
+import { WhisperModel } from "./whisper/model";
+import { genericArgsFor } from "./universal/archs";
+import { familyForGraph } from "./families";
+import { resolveComposition } from "./composition";
+import { bindGraphCaches } from "./cache-composition";
+import { phaseRowLimit } from "./graph";
+import type { Composition } from "../contracts/portable/composition";
+import type { Cache } from "../contracts/mlx/cache";
+import {
+  GraphTable,
+  ModelImplementationRegistry,
+  type GraphEntry,
+  type ModelImplementation,
+  type ModelImplementationProvider,
+} from "./implementation";
+import {
+  assertResolvedModelProfile,
+  resolveModelProfile,
+  type ResolvedModelProfile,
+  type ResolveModelProfileOptions,
+  type ModelGraph,
+  type ModelLoader,
+  type GenerationLoop,
+} from "./profile";
+
+type WrittenGraph = (
+  | Gemma4Model | Gemma4Graph | MiniCPM5Model | Qwen35Model | Qwen38TrellisM4Pro | Qwen38TrellisM1Max | Qwen3Model | Qwen3MoeModel
+  | DiffusionGemmaModel | Glm52Model | UniversalDenseModel);
+
+export type RuntimeModel = WrittenGraph & { makeCache(): Cache[]; readonly composition?: Composition };
+
+/** The family's single fallback graph. */
+function generic(id: string, create: GraphEntry<Weights, WrittenGraph>["create"]): GraphEntry<Weights, RuntimeModel> {
+  return loadedEntry({ id, serves: {}, create });
+}
+
+function loadedEntry(entry: GraphEntry<Weights, WrittenGraph>): GraphEntry<Weights, RuntimeModel> {
+  return { ...entry, create: (weights, config, composition) => loadedGraph(entry.create(weights, config, composition), composition) };
+}
+
+/** Attach the loader's cache and scheduler configuration to a shipped graph.
+ * Its written implementation only provides the layer layout and operations. */
+function loadedGraph(graph: WrittenGraph, composition: Composition): RuntimeModel {
+  return Object.assign(bindGraphCaches(graph, composition), { composition });
+}
+
+/** Resident graph files and their default recommendation keys. Explicit
+ * choices are unrestricted; see {@link GraphTable}. */
+export const MLX_GRAPH_TABLES: Readonly<Record<ModelGraph, GraphTable<Weights, RuntimeModel>>> = Object.freeze({
+  "diffusion-gemma": new GraphTable([generic("diffusion-gemma", (weights, config, composition) => new DiffusionGemmaModel(weights, config))]),
+  "minicpm5": new GraphTable([generic("minicpm5", (weights, config, composition) => new MiniCPM5Model(weights, config, phaseRowLimit(composition)))]),
+  // Written graph recommendations for the published Trellis artifact.
+  // Explicit caller choices bypass these artifact and device keys.
+  "qwen3.5": new GraphTable([
+    loadedEntry({ id: QWEN38_TRELLIS_M4PRO_GRAPH, serves: {
+        fingerprint:"cfa205c8f5af046a",device:"applegpu_g16s",maxRows:QWEN38_M4PRO_DECODE_ROWS,
+        kv:[{kind:"bf16"},{kind:"affine",bits:4,groupSize:64}],
+      },
+      create: (weights,config) => new Qwen38TrellisM4Pro(weights,config),
+      anePrograms: () => qwen38TrellisM4ProAnePrograms() }),
+    loadedEntry({ id: QWEN38_TRELLIS_M1MAX_GRAPH, serves: {
+      fingerprint:"cfa205c8f5af046a",device:"applegpu_g13s",
+    }, create: (weights, config) => new Qwen38TrellisM1Max(weights, config),
+      anePrograms: () => qwen38TrellisM1MaxAnePrograms() }),
+    generic("qwen3.5", (weights, config, composition) => new Qwen35Model(weights, config, phaseRowLimit(composition), composition.device)),
+  ], qwen35GraphFingerprint),
+  "qwen3-moe": new GraphTable([generic("qwen3-moe", (weights, config, composition) => new Qwen3MoeModel(weights, config, phaseRowLimit(composition)))]),
+  "qwen3": new GraphTable([generic("qwen3", (weights, config, composition) => new Qwen3Model(weights, config, phaseRowLimit(composition)))]),
+  "gemma4": new GraphTable([...GEMMA_GRAPHS.map(loadedEntry), generic("gemma4", (weights, config, composition) => new Gemma4Model(weights, config, phaseRowLimit(composition)))],
+    configFingerprint),
+  "universal-dense": new GraphTable([
+    generic("universal-dense", (weights, config, composition) =>
+      new UniversalDenseModel(weights, config, genericArgsFor(config)!, phaseRowLimit(composition), composition.device))]),
+});
+
+/** A resident family's native construction: its graph table selects the graph for
+ * the composition it receives. Its loader and loop are the family's own declaration. */
+function residentImplementation(graph: ModelGraph, table: GraphTable<Weights, RuntimeModel>): ModelImplementation<Weights, RuntimeModel> {
+  const family = familyForGraph(graph);
+  if (!family) throw new Error(`graph table ${graph} names a graph no model family declares`);
+  return { id: graph, graph, loader: family.loader, loop: family.loop,
+    create: (weights, config, _profile, composition) =>
+      table.select(config, composition).create(weights, config, composition) };
+}
+
+/** Engine-owned registrations, one per resident graph of `MODEL_FAMILIES`. Exact
+ * quant profiles can name additional implementations in a composed registry,
+ * without changing sessions or files. */
+export const MLX_MODEL_IMPLEMENTATIONS = new ModelImplementationRegistry<Weights, RuntimeModel>(
+  Object.entries(MLX_GRAPH_TABLES).map(([graph, table]) => residentImplementation(graph, table)));
+
+export interface ModelOpenOptions<Model = RuntimeModel> extends Glm52RuntimeOpenOptions {
+  readonly profiles?: ResolveModelProfileOptions;
+  readonly implementations?: ModelImplementationProvider<Weights, Model>;
+  /** The composition the graph is built for; absent resolves the default record
+   * (`resolveComposition(config)`). The direct Colibri runtime plans its own
+   * resources from the options above and takes none. */
+  readonly composition?: Composition;
+  /** A named shipped graph or the caller's own load function. Explicit choices
+   * bypass device and configuration recommendations and need no registration. */
+  readonly graph?: string | ((weights: Weights, config: ModelConfig) => Model);
+}
+
+/** A loader entry cannot provide the requested artifact/operation pairing. */
+export class UnsupportedModelLoadError extends Error {
+  readonly code = "unsupported-model-load";
+  constructor(readonly profile: ResolvedModelProfile, message: string) {
+    super(message);
+    this.name = "UnsupportedModelLoadError";
+  }
+}
+
+export interface Glm52RuntimeOpenOptions {
+  /** Whole-process ceiling. Defaults to the smaller of the validated 25 GiB
+   * preset and physical RAM. */
+  memoryBudgetBytes?: number;
+  /** Context reserved by the exact GLM resource equation. */
+  contextTokens?: number;
+  /** Generated-token allowance within contextTokens. */
+  maxGenerationTokens?: number;
+  /** Maximum ordinary continuous-batch rows. */
+  batchSize?: number;
+  /** Native checkpoint MTP row; on by default. */
+  enableMtp?: boolean;
+  mtpDraftTokens?: number;
+  /** Test/diagnostic override; normal callers use physical RAM. */
+  machineBytes?: number;
+  /** Native expert-I/O dylib override. Packaged/dev defaults are automatic. */
+  libraryPath?: string;
+  /** The composition's KV scheme and paging, which the graph's latent caches
+   * compose (default f32, unpaged). */
+  caches?: Glm52LatentCaches;
+}
+
+/** Open the direct Colibri artifact through the validated bounded expert
+ * runtime. The header-only plan runs before any resident tensor or expert slab
+ * is opened, so an impossible configuration fails without committing memory. */
+export async function openGlm52RuntimeModel(
+  modelDir: string,
+  options: Glm52RuntimeOpenOptions = {},
+): Promise<{ model: Glm52LoadedModel; plan: Glm52MemoryPlan }> {
+  const machineBytes = options.machineBytes ?? totalmem();
+  const processLimitBytes = options.memoryBudgetBytes ??
+    Math.min(GLM52_G5_DEFAULT_PROCESS_LIMIT_BYTES, machineBytes);
+  const plan = await planGlm52MemoryForArtifact(modelDir, {
+    machineBytes,
+    processLimitBytes,
+    contextTokens: options.contextTokens ?? GLM52_G5_DEFAULT_CONTEXT_TOKENS,
+    maxGenerationTokens:
+      options.maxGenerationTokens ?? GLM52_G5_DEFAULT_MAX_GENERATION_TOKENS,
+    batchSize: options.batchSize ?? 1,
+    enableMtp: options.enableMtp !== false,
+    mtpDraftTokens: options.mtpDraftTokens,
+  });
+  const model = await Glm52Model.openStreamed(modelDir, {
+    budgetBytes: plan.processLimitBytes,
+    reserveBytes: plan.runtimeReserveBytes,
+    workingSlots: plan.mainWorkingSlots,
+    maxSlotsPerLayer: 1,
+    workers: 2,
+    libraryPath: options.libraryPath,
+    decodeKernel: "metal",
+    enableMtp: plan.enableMtp,
+    mtpDraftTokens: plan.mtpDraftTokens,
+    ...(options.caches ? { caches: options.caches } : {}),
+  });
+  if (model.expertRuntime?.plan.plannedBytes !== plan.plannedProcessBytes) {
+    const actual = model.expertRuntime?.plan.plannedBytes ?? 0;
+    model.dispose();
+    throw new Error(
+      `GLM runtime plan ${actual} != preflight resource equation ` +
+      `${plan.plannedProcessBytes}`,
+    );
+  }
+  return { model, plan };
+}
+
+/** Open a Whisper checkpoint (encoder-decoder speech model) through the same
+ * profile resolution as text models. Whisper never enters the chat loop, so
+ * it is not a RuntimeModel: the transcription engine
+ * (packages/inference/src/transcription/whisper/)
+ * owns it. */
+export async function openWhisperModel(
+  modelDir: string, options: { readonly profiles?: ResolveModelProfileOptions } = {},
+): Promise<{ model: WhisperModel; profile: ResolvedModelProfile; config: ModelConfig }> {
+  const config = await loadModelConfig(modelDir);
+  const profile = resolveModelProfile(config, options.profiles);
+  if (profile.profile.execution.graph !== "whisper")
+    throw new Error(`${modelDir} is not a Whisper checkpoint (profile ${profile.profile.id})`);
+  const weights = await Weights.open(modelDir);
+  try {
+    return { model: new WhisperModel(weights, config), profile, config };
+  } catch (error) {
+    weights.dispose();
+    throw error;
+  }
+}
+
+/**
+ * Artifact-aware construction. GLM-5.2's Colibri snapshot has no ordinary
+ * model.safetensors.index.json, so it must bypass Weights.open and use its
+ * dedicated header catalog/tensor source.
+ */
+export function openModel<Model>(
+  modelDir: string,
+  options: ModelOpenOptions<Model> & { readonly graph: (weights: Weights, config: ModelConfig) => Model },
+): Promise<Model>;
+export function openModel<Model>(
+  modelDir: string,
+  options: ModelOpenOptions<Model> & { readonly implementations: ModelImplementationProvider<Weights, Model> },
+): Promise<Model>;
+export function openModel(modelDir: string, options?: ModelOpenOptions): Promise<RuntimeModel>;
+export async function openModel<Model>(
+  modelDir: string,
+  options: ModelOpenOptions<Model> = {},
+): Promise<Model | RuntimeModel> {
+  const config = await loadModelConfig(modelDir);
+  const composition = options.composition ?? resolveComposition(config);
+  // A caller's graph is the complete arrangement. Loading it does not require
+  // a family profile or participation in the default recommendation table.
+  if (typeof options.graph === "function") {
+    const weights = await Weights.open(modelDir);
+    try { return options.graph(weights, config); }
+    catch (error) { weights.dispose(); throw error; }
+  }
+  const profile = resolveModelProfile(config, options.profiles);
+  if (options.graph !== undefined) {
+    const entry = MLX_GRAPH_TABLES[profile.profile.execution.graph]?.explicit(options.graph);
+    if (!entry) throw new Error(`no shipped graph table for ${profile.profile.execution.graph}`);
+    const weights = await Weights.open(modelDir);
+    try { return entry.create(weights, config, composition); }
+    catch (error) { weights.dispose(); throw error; }
+  }
+  return ARTIFACT_LOADERS[profile.profile.execution.loader][profile.profile.execution.loop]
+    .open({ modelDir, config, profile, composition, options });
+}
+
+interface ArtifactLoad<Model> {
+  readonly modelDir: string;
+  readonly config: ModelConfig;
+  readonly profile: ResolvedModelProfile;
+  readonly composition: Composition;
+  readonly options: ModelOpenOptions<Model>;
+}
+
+interface ArtifactLoader {
+  open<Model>(load: ArtifactLoad<Model>): Promise<Model | RuntimeModel>;
+  create<Model>(weights: Weights, config: ModelConfig, profile: ResolvedModelProfile,
+    composition: Composition, implementations: ModelImplementationProvider<Weights, Model>): Model;
+}
+
+const residentLoader: ArtifactLoader = {
+  async open({ modelDir, config, profile, composition, options }) {
+    const implementation = (options.implementations ?? MLX_MODEL_IMPLEMENTATIONS).select(config, profile);
+    const weights = await Weights.open(modelDir);
+    try { return implementation.create(weights, config, profile, composition); }
+    catch (error) { weights.dispose(); throw error; }
+  },
+  create: (weights, config, profile, composition, implementations) =>
+    implementations.select(config, profile).create(weights, config, profile, composition),
+};
+
+const streamedLoader: ArtifactLoader = {
+  async open({ modelDir, profile, composition, options }) {
+    if (options.implementations || profile.profile.execution.implementation !== undefined)
+      throw new UnsupportedModelLoadError(profile, "custom Colibri implementations require a streamed loader binding; refusing to fall back");
+    return (await openGlm52RuntimeModel(modelDir, options.composition
+      ? { ...options, caches: { kv: composition.kv, pagedKv: composition.pagedKv } } : options)).model;
+  },
+  create(_weights, _config, profile) {
+    throw new UnsupportedModelLoadError(profile, `${profile.profile.id} uses the direct Colibri container; construct it with ` +
+      "openModel(modelDir) or Glm52Model.open(modelDir)");
+  },
+};
+
+const speechLoader: ArtifactLoader = {
+  async open({ modelDir, profile }) {
+    throw new UnsupportedModelLoadError(profile, `${modelDir} is a speech model (${profile.profile.id}); open it with openWhisperModel ` +
+      "or serve it as a transcription model — it has no chat/completion loop");
+  },
+  create(_weights, _config, profile) {
+    throw new UnsupportedModelLoadError(profile, `${profile.profile.id} has an encoder-decoder graph; open it with openWhisperModel`);
+  },
+};
+
+/** Artifact loading follows declared table facts, including streamed weights
+ * and the separate speech operation. It never selects by a model predicate. */
+const ARTIFACT_LOADERS: Readonly<Record<ModelLoader, Readonly<Record<GenerationLoop, ArtifactLoader>>>> = {
+  safetensors: { autoregressive: residentLoader, diffusion: residentLoader, "encoder-decoder": speechLoader },
+  colibri: { autoregressive: streamedLoader, diffusion: streamedLoader, "encoder-decoder": speechLoader },
+};
+
+/** Construct the binding named by a validated profile for `composition`, which
+ * the implementation receives (absent: the default record,
+ * `resolveComposition(config)`). A supplied registry can return any
+ * backend-owned interface; the compatibility default returns the legacy model
+ * union. The caller retains ownership of the supplied weights. */
+export function createModel<Model>(
+  weights: Weights, config: ModelConfig, composition: Composition | undefined,
+  resolved: ResolvedModelProfile | undefined, implementations: ModelImplementationProvider<Weights, Model>,
+): Model;
+export function createModel(
+  weights: Weights, config: ModelConfig, composition?: Composition, resolved?: ResolvedModelProfile,
+  implementations?: ModelImplementationProvider<Weights, RuntimeModel>,
+): RuntimeModel;
+export function createModel<Model>(
+  weights: Weights,
+  config: ModelConfig,
+  composition: Composition = resolveComposition(config),
+  resolved: ResolvedModelProfile = resolveModelProfile(config),
+  implementations: ModelImplementationProvider<Weights, Model> | ModelImplementationProvider<Weights, RuntimeModel> = MLX_MODEL_IMPLEMENTATIONS,
+): Model | RuntimeModel {
+  assertResolvedModelProfile(config, resolved);
+  return ARTIFACT_LOADERS[resolved.profile.execution.loader][resolved.profile.execution.loop]
+    .create<Model | RuntimeModel>(weights, config, resolved, composition, implementations);
+}

package/src/models/families.ts

@@ -0,0 +1,70 @@
+// The model-family registry: every family the engine serves, in one list. A new
+// model is one record here plus its graph and its native implementation
+// (`factory.ts`). Everything else that used to restate a family (support tiers,
+// profiles, engine capabilities, declarations) is derived from this list.
+
+import type { ModelConfig } from "../artifacts/config";
+import { DIFFUSION_GEMMA_FAMILY } from "./diffusion-gemma/family";
+import type { EngineCapability, ModelFamily, ModelGraph } from "./family";
+import { GEMMA4_FAMILY } from "./gemma4/family";
+import { GLM52_FAMILY } from "./glm52/family";
+import { MINICPM5_FAMILY } from "./minicpm5/family";
+import { QWEN35_FAMILY, QWEN3_FAMILY, QWEN3_MOE_FAMILY } from "./qwen/families";
+import { GENERIC_MODEL_TYPES, remapModelType } from "./universal/archs";
+import { UNIVERSAL_FAMILY } from "./universal/family";
+import { WHISPER_FAMILY } from "./whisper/family";
+
+/** First match wins. Targeted families accept disjoint configs; the universal
+ * family comes last so it never shadows one. */
+export const MODEL_FAMILIES: readonly ModelFamily[] = Object.freeze([
+  GEMMA4_FAMILY,
+  DIFFUSION_GEMMA_FAMILY,
+  GLM52_FAMILY,
+  QWEN35_FAMILY,
+  QWEN3_FAMILY,
+  QWEN3_MOE_FAMILY,
+  WHISPER_FAMILY,
+  MINICPM5_FAMILY,
+  UNIVERSAL_FAMILY,
+]);
+
+{
+  const graphs = MODEL_FAMILIES.map(family => family.graph);
+  const duplicate = graphs.find((graph, index) => graphs.indexOf(graph) !== index);
+  if (duplicate) throw new Error(`model families: duplicate graph ${duplicate}`);
+}
+
+/** The family whose graph runs this config, or null. Throws where the family's
+ * own parser rejects a malformed config. */
+export function familyOf(config: ModelConfig): ModelFamily | null {
+  return MODEL_FAMILIES.find(family => family.accepts(config)) ?? null;
+}
+
+/** The family a registry record's `model_type` names before any config is read. */
+export function familyOfModelType(modelType: string): ModelFamily | null {
+  if (typeof modelType !== "string") return null; // a record with no model_type names no family
+  return MODEL_FAMILIES.find(family => family.hasModelType(modelType)) ?? null;
+}
+
+export function familyForGraph(graph: ModelGraph): ModelFamily | null {
+  return MODEL_FAMILIES.find(family => family.graph === graph) ?? null;
+}
+
+/** What the engine can do independent of any model. */
+const BASE_CAPABILITIES = [
+  "autoregressive", "colibri-container", "diffusion", "encoder-decoder",
+  "mixed-precision-kv", "native-mtp", "safetensors", "vision-sidecar",
+] as const satisfies readonly EngineCapability[];
+
+/** The base capabilities plus what every registered family's graph brings. */
+export const ENGINE_CAPABILITIES: readonly EngineCapability[] = Object.freeze(
+  [...new Set([...BASE_CAPABILITIES, ...MODEL_FAMILIES.flatMap(family => family.capabilities)])].sort());
+
+/** Why no family serves this config, naming what is served. */
+export function unsupportedModelMessage(config: ModelConfig): string {
+  const arch = remapModelType(config.modelType);
+  return `unsupported model_type "${config.modelType}"` +
+    (arch !== config.modelType ? ` (mlx-lm remaps it to "${arch}")` : "") +
+    ` — targeted: ${MODEL_FAMILIES.flatMap(family => family.label ?? []).join(", ")};` +
+    ` generic (Tier-0): ${[...GENERIC_MODEL_TYPES].sort().join(", ")}`;
+}

package/src/models/family.ts

@@ -0,0 +1,125 @@
+// A model family is one record: what accepts a parsed config, which graph runs
+// it, how it is composed, and every fact the layers above read about it. The
+// records are listed once in `models/families.ts`; support tiers, profiles,
+// engine capabilities, implementation composition and the declarations
+// (training, generation, embedding, sentinels, media tokens, chat fallback)
+// are all derived from that list. Records are native-free: they never import a
+// graph class, so listing a model's support never loads MLX.
+
+import type { ModelConfig } from "../artifacts/config";
+import type { TemplateFallback } from "../input/chat-template";
+import type { MediaTokenTexts, SentinelTokenTexts } from "../input/special-tokens";
+
+/** An open string: a graph is a name a family declares and an implementation registers under. */
+export type ModelGraph = string;
+/** A named thing the engine can do. Families declare the ones their graph needs. */
+export type EngineCapability = string;
+export type FidelityTier = "l1" | "l2" | "l3";
+
+export type FidelityTarget =
+  | Readonly<{ tier: "l1"; oracle: "mlx-lm"; claim: "bit-exact" }>
+  | Readonly<{ tier: "l1"; oracle: "mlx-whisper"; claim: "bit-exact" }>
+  | Readonly<{ tier: "l2"; oracle: "mlx-optiq"; claim: "bit-exact" }>
+  | Readonly<{ tier: "l3"; oracle: null; claim: "measured" }>;
+
+export type ModelLoader = "safetensors" | "colibri";
+export type GenerationLoop = "autoregressive" | "diffusion" | "encoder-decoder";
+/** What a profile was declared for: an exact artifact, a family's own graphs, or
+ * the universal graph. Which of a family's graphs a load builds (its generic
+ * graph, or a written one built for the composition) is the
+ * family's graph table's choice (`models/implementation.ts`), not a profile's. */
+export type ModelSpecialization = "artifact" | "dedicated" | "generic";
+
+export const L1: FidelityTarget = Object.freeze({ tier: "l1", oracle: "mlx-lm", claim: "bit-exact" });
+/** L1 for speech: the capability's own reference implementation (mlx-whisper)
+ * plays the mlx-lm role; no mlx-lm arm exists for Whisper. */
+export const L1_WHISPER: FidelityTarget = Object.freeze({ tier: "l1", oracle: "mlx-whisper", claim: "bit-exact" });
+export const L2: FidelityTarget = Object.freeze({ tier: "l2", oracle: "mlx-optiq", claim: "bit-exact" });
+export const L3: FidelityTarget = Object.freeze({ tier: "l3", oracle: null, claim: "measured" });
+
+/** Fine-tuning defaults a graph may declare. The CLI reads these from the
+ * resolved profile; nothing outside the model decides them by name. */
+export interface TrainingDefaults {
+  /** Default maximum training sequence length in tokens. */
+  readonly maxSeqLength: number;
+}
+
+/** Applies to every graph that declares nothing of its own. */
+export const GENERIC_TRAINING_DEFAULTS: TrainingDefaults = Object.freeze({ maxSeqLength: 4096 });
+
+/** Request defaults a graph declares for chat generation. The server applies
+ * them under explicit request and server settings; nothing outside the model
+ * decides them by name. */
+export interface GenerationDefaults {
+  /** `enable_thinking` when neither the request nor the server chose; absent leaves the template's own default. */
+  readonly enableThinking?: boolean;
+  /** Upper bound on the configured sampling temperature for replies with thinking off
+   * (model authors publish one think-mode temperature and recommend a cooler one for direct replies). */
+  readonly noThinkTemperatureCap: number;
+}
+
+/** Applies to every graph that declares nothing of its own. */
+export const GENERIC_GENERATION_DEFAULTS: GenerationDefaults = Object.freeze({ noThinkTemperatureCap: 0.7 });
+
+/** How a graph's pooled text embedding is formed from tokenized input. */
+export interface EmbeddingDeclaration {
+  /** Special token appended to every input; its final hidden state is the vector. */
+  readonly terminator: string;
+}
+
+/** The experts a sparse graph routes each token through. */
+export interface MoeDeclaration {
+  readonly experts: number;
+  readonly topK: number;
+}
+
+/** A MoE family whose config spells its experts with the shared text fields. */
+export function configuredMoe(config: ModelConfig): MoeDeclaration | null {
+  const text = config.text;
+  return text.enableMoeBlock && text.numExperts > 0 ? { experts: text.numExperts, topK: text.topKExperts } : null;
+}
+
+export interface ModelFamily {
+  /** The graph this family runs: what profiles select, what an implementation
+   * registers under, and the key of every declaration below. */
+  readonly graph: ModelGraph;
+  /** The one structural rule: does this family's graph run this parsed config?
+   * May throw for a config too malformed to read (the generic parser does). */
+  accepts(config: ModelConfig): boolean;
+  /** Whether a bare `model_type` (a registry record, no config read) is this
+   * family's. False where the rule needs more than the type: MiniCPM5 shares
+   * `llama` and is told apart only by its dimensions. */
+  hasModelType(modelType: string): boolean;
+  /** `targeted`: a dedicated forward with its own oracle path. `generic`: the
+   * universal module (L1 only). */
+  readonly tier: "targeted" | "generic";
+  /** Names a targeted family in the unsupported-model message. */
+  readonly label?: string;
+  readonly loader: ModelLoader;
+  readonly loop: GenerationLoop;
+  readonly fidelity: FidelityTarget;
+  /** `dedicated` for a graph written for the family, `generic` for the universal one. */
+  readonly specialization: "dedicated" | "generic";
+  /** The family profile's id. */
+  readonly profileId: string;
+  /** Capabilities the graph needs of the engine, beyond its loader and loop's own. */
+  readonly capabilities: readonly EngineCapability[];
+  /** A speech-to-text family, served by the transcription engine. */
+  readonly transcribes?: true;
+  readonly moe?: (config: ModelConfig) => MoeDeclaration | null;
+  readonly trainingDefaults?: TrainingDefaults;
+  /** Present when the graph declares `GraphCapabilities.embeddings`. */
+  readonly embedding?: EmbeddingDeclaration;
+  /** Present when generated text delimits tool calls and reasoning with special tokens. */
+  readonly sentinels?: SentinelTokenTexts;
+  /** Soft-token markers of the graph's tower prompts. */
+  readonly mediaTokens?: MediaTokenTexts;
+  readonly generationDefaults?: GenerationDefaults;
+  /** Renders chats for artifacts that ship no template of their own. */
+  readonly chatTemplateFallback?: TemplateFallback;
+}
+
+/** Declare a family: typed against the record, frozen. */
+export function defineFamily(family: ModelFamily): ModelFamily {
+  return Object.freeze(family);
+}

package/src/models/gemma4/assistant-target.ts

@@ -0,0 +1,46 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { Cache, KvDonorAttention } from "../../contracts/mlx/cache";
+import { captureKvDonorAttention } from "../../state/kv-attention-view";
+import type { AssistantAttention } from "../../contracts/mlx/attention";
+import { disposeResources } from "../../runtime/resources";
+
+/** Bind captured cache validity to Q-only attention. Encoded layouts can
+ * supply this same numerical port without changing the draft graph. */
+export function readAssistantDonors(sliding: Cache, full: Cache) {
+  const views: KvDonorAttention[] = [], masks: MlxArray[] = [];
+  const dispose = () => disposeResources([...masks, ...views]);
+  try {
+    for (const cache of [sliding, full]) views.push(captureKvDonorAttention(cache));
+    const bind = (view: KvDonorAttention): AssistantAttention => {
+      // All draft steps borrow the same target snapshot. Build each validity
+      // mask once, then reuse it across matching layers and the entire chain.
+      const byWindow = new Map<number | null, MlxArray | null>();
+      return { attend(query, scale, window) {
+        if (!byWindow.has(window)) {
+          const width = view.width, B = view.offsets.length;
+          const starts = view.starts.map((start, row) => Math.max(start,
+            window === null ? start : view.ends[row]! - window));
+          let mask: MlxArray | null = null;
+          if (starts.some(start => start !== 0) || view.ends.some(end => end !== width)) {
+            using columns = ops.arange(0, width, 1, Dtype.int32);
+            using lower = ops.fromInt32(starts, [B, 1, 1, 1]);
+            using upper = ops.fromInt32([...view.ends], [B, 1, 1, 1]);
+            using before = ops.less(columns, lower);
+            using after = ops.greaterEqual(columns, upper);
+            using invalid = ops.logicalOr(before, after);
+            using floats = invalid.astype(view.dtype);
+            mask = ops.mulScalar(floats, -1e9);
+            masks.push(mask);
+          }
+          byWindow.set(window, mask);
+        }
+        const mask = byWindow.get(window)!;
+        return view.attend(query, scale, { mode: mask ? "array" : "", arr: mask });
+      } };
+    };
+    return { positions: views[0]!.offsets.map(offset => offset - 1),
+      sliding: bind(views[0]!), full: bind(views[1]!), dispose };
+  } catch (error) { dispose(); throw error; }
+}