mlx-bun 0.0.13 → 0.3.0

package/src/model/batched-rotating.ts

@@ -20,11 +20,17 @@
 // temporal coordinates then ROLLED to the ring's physical layout — see
 // buildBatchedRotatingMask (a column-by-column port of mlx-lm make_mask,
 // including the `roll(shift=idx+1)`). Gated model-free against mlx-lm
-// (tests/batched-rotating.test.ts) + end-to-end vs a long-context Gemma oracle.
+// (tests/unit/batched-rotating.test.ts) + end-to-end vs a long-context Gemma oracle.
 
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
 import { RotatingKVCache, type Cache, type Mask } from "./gemma4-base";
+import { BatchedRotatingState } from "./batched-rotating-state";
+import {
+  mergeStorageRows,
+  plainRowStorage,
+  temporalStorageView,
+} from "./batched-row-storage";
 
 const STEP = 256;
 
@@ -78,39 +84,46 @@ export function buildBatchedRotatingMask(
 export class BatchedRotatingCache implements Cache {
   keys: MlxArray | null = null;
   values: MlxArray | null = null;
-  /** Per-row absolute position (drives RoPE) and buffer padding. */
-  offsetArr: number[];
-  leftPad: number[];
-  #idx = 0;
-  #offset = 0; // scalar total tokens processed (mlx-lm _offset)
-  #rotated = false;
+  readonly #rows: BatchedRotatingState;
   #ropeArr: MlxArray | null = null;
   #ropeForOffset = -1;
   readonly maxSize: number;
 
   constructor(maxSize: number, leftPad: number[]) {
     this.maxSize = maxSize;
-    this.leftPad = [...leftPad];
-    this.offsetArr = leftPad.map((l) => -l);
+    this.#rows = new BatchedRotatingState(maxSize, leftPad);
   }
 
+  /** Same signature as the serial RotatingKVCache — the scheduler's merge
+   *  guards recognize a batched ring as "rotating-plain + RowBatchCache"
+   *  (batch-scheduler #mergeJoiner prevRot branch); without this the next
+   *  join sees "unknown", skips the prevRot rows, and replaces the batch's
+   *  sliding-layer KV with a one-row cache (B≥3 collapse). Guard order
+   *  everywhere is isRowBatchCache FIRST, so the shared string never
+   *  misroutes a batched cache into a serial-only path. */
+  signature(): string { return "kv:rotating-plain"; }
+
+  get offsetArr(): number[] { return this.#rows.offsets; }
+  get leftPad(): number[] { return this.#rows.leftPad; }
+  get batchSize(): number { return this.#rows.batchSize; }
+
   /** Current batch size — tracks filter() (which shrinks the per-row arrays). */
   get #B(): number {
-    return this.leftPad.length;
+    return this.#rows.batchSize;
   }
 
   /** mlx-lm uses `cache.offset` (the per-row array) for the scalar interface
    *  too; we expose the scalar total as `offset` and the per-row positions via
    *  `ropeOffsetArr` (the model's per-row RoPE path). */
   get offset(): number {
-    return this.#offset;
+    return this.#rows.totalOffset;
   }
 
   get ropeOffsetArr(): MlxArray {
-    if (this.#ropeArr && this.#ropeForOffset === this.#offset) return this.#ropeArr;
+    if (this.#ropeArr && this.#ropeForOffset === this.offset) return this.#ropeArr;
     this.#ropeArr?.dispose();
     this.#ropeArr = MlxArray.fromInt32(Int32Array.from(this.offsetArr), [this.#B]);
-    this.#ropeForOffset = this.#offset;
+    this.#ropeForOffset = this.offset;
     return this.#ropeArr;
   }
 
@@ -120,7 +133,7 @@ export class BatchedRotatingCache implements Cache {
       mode: "array",
       arr: buildBatchedRotatingMask(
         this.#B, N, this.leftPad, this.maxSize, window,
-        this.#idx, this.#offset, this.#rotated,
+        this.#rows.ringIndex, this.offset, this.#rows.rotated,
       ),
     };
   }
@@ -131,7 +144,7 @@ export class BatchedRotatingCache implements Cache {
     const vD = v.shape[3]!;
     if (S !== 1)
       throw new Error("BatchedRotatingCache supports N=1 decode updates only (solo-prefill then merge)");
-    const prev = this.#offset;
+    const prev = this.offset;
 
     // Grow the buffer (in STEP chunks) until it reaches maxSize.
     if (!this.keys || (prev >= this.keys.shape[2]! && this.keys.shape[2]! < this.maxSize)) {
@@ -148,7 +161,7 @@ export class BatchedRotatingCache implements Cache {
         this.keys = newK;
         this.values = newV;
       }
-      this.#idx = prev;
+      this.#rows.markGrown(prev);
     }
 
     // Trim any overshoot past maxSize (decrements left padding persistently).
@@ -160,34 +173,27 @@ export class BatchedRotatingCache implements Cache {
       this.values!.dispose();
       this.keys = tk;
       this.values = tv;
-      this.#idx = this.maxSize;
-      this.leftPad = this.leftPad.map((x) => x - trimSize);
+      this.#rows.trimOvershoot(trimSize);
     }
 
     // Rotate when the write head reaches the end of the ring.
-    if (this.#idx === this.maxSize) {
-      this.#rotated = true;
-      this.#idx = 0;
-    }
-    if (this.#rotated) this.leftPad = this.leftPad.map((x) => x - S);
+    const writeIndex = this.#rows.beginWrite(S);
 
     // Write the new K/V at the ring head.
     const [, , SK, DK] = this.keys!.shape as [number, number, number, number];
-    const k2 = ops.sliceUpdate(this.keys!, k, [0, 0, this.#idx, 0], [B, H, this.#idx + S, DK]);
-    const v2 = ops.sliceUpdate(this.values!, v, [0, 0, this.#idx, 0], [B, H, this.#idx + S, vD]);
+    const k2 = ops.sliceUpdate(this.keys!, k, [0, 0, writeIndex, 0], [B, H, writeIndex + S, DK]);
+    const v2 = ops.sliceUpdate(this.values!, v, [0, 0, writeIndex, 0], [B, H, writeIndex + S, vD]);
     this.keys!.dispose();
     this.values!.dispose();
     this.keys = k2;
     this.values = v2;
-    this.#offset += S;
-    this.offsetArr = this.offsetArr.map((x) => x + S);
-    this.#idx += S;
+    this.#rows.commitWrite(S);
 
     // Return the populated prefix (ring not yet full) or the whole buffer.
-    if (this.#offset < this.maxSize) {
+    if (this.offset < this.maxSize) {
       return [
-        this.keys.slice([0, 0, 0, 0], [B, H, this.#offset, DK]),
-        this.values.slice([0, 0, 0, 0], [B, H, this.#offset, vD]),
+        this.keys.slice([0, 0, 0, 0], [B, H, this.offset, DK]),
+        this.values.slice([0, 0, 0, 0], [B, H, this.offset, vD]),
       ];
     }
     return [
@@ -206,27 +212,10 @@ export class BatchedRotatingCache implements Cache {
   /** Ring contents in temporal order, cut to the valid length (extract). */
   temporalView(): [MlxArray, MlxArray] {
     if (!this.keys || !this.values) throw new Error("cache is empty");
-    const valid = Math.min(this.#offset, this.maxSize);
-    const order = (a: MlxArray): MlxArray => {
-      const [B, H, Sbuf, D] = a.shape as [number, number, number, number];
-      let t: MlxArray;
-      if (this.#idx === Sbuf) {
-        t = a.slice([0, 0, 0, 0], [B, H, Sbuf, D]);
-      } else if (this.#idx < this.#offset) {
-        const tail = a.slice([0, 0, this.#idx, 0], [B, H, Sbuf, D]);
-        const head = a.slice([0, 0, 0, 0], [B, H, this.#idx, D]);
-        t = ops.concatAxis([tail, head], 2);
-        tail.dispose();
-        head.dispose();
-      } else {
-        t = a.slice([0, 0, 0, 0], [B, H, this.#idx, D]);
-      }
-      const [, , , Dd] = t.shape as [number, number, number, number];
-      const cut = t.slice([0, 0, 0, 0], [B, H, valid, Dd]);
-      t.dispose();
-      return cut;
-    };
-    return [order(this.keys), order(this.values)];
+    return [
+      temporalStorageView(plainRowStorage, this.keys, this.#rows),
+      temporalStorageView(plainRowStorage, this.values, this.#rows),
+    ];
   }
 
   /** mlx-lm `BatchRotatingKVCache.extract` (models/cache.py:1417): row `i`
@@ -242,34 +231,13 @@ export class BatchedRotatingCache implements Cache {
   extractRow(i: number): RotatingKVCache | null {
     if (!this.keys || !this.values) return null;
     const pad = Math.max(0, this.leftPad[i]!);
-    const valid = Math.min(this.#offset, this.maxSize);
-    const cut = (a: MlxArray): MlxArray => {
-      const [, H, Sbuf, D] = a.shape as [number, number, number, number];
-      const row = a.slice([i, 0, 0, 0], [i + 1, H, Sbuf, D]);
-      // temporal order — same three cases as temporalView, single row.
-      let t: MlxArray;
-      if (this.#idx === Sbuf) {
-        t = row;
-      } else if (this.#idx < this.#offset) {
-        const tail = row.slice([0, 0, this.#idx, 0], [1, H, Sbuf, D]);
-        const head = row.slice([0, 0, 0, 0], [1, H, this.#idx, D]);
-        t = ops.concatAxis([tail, head], 2);
-        tail.dispose();
-        head.dispose();
-        row.dispose();
-      } else {
-        t = row.slice([0, 0, 0, 0], [1, H, this.#idx, D]);
-        row.dispose();
-      }
-      const cutV = t.slice([0, 0, pad, 0], [1, H, valid, D]);
-      t.dispose();
-      const own = ops.contiguous(cutV);
-      cutV.dispose();
-      return own;
-    };
     const c = new RotatingKVCache(this.maxSize);
-    const k = cut(this.keys);
-    const v = cut(this.values);
+    const k = temporalStorageView(plainRowStorage, this.keys, this.#rows, {
+      row: i, from: pad, copy: true,
+    });
+    const v = temporalStorageView(plainRowStorage, this.values, this.#rows, {
+      row: i, from: pad, copy: true,
+    });
     c.restoreState(k, v, this.offsetArr[i]!, k.shape[2]!);
     return c;
   }
@@ -277,31 +245,27 @@ export class BatchedRotatingCache implements Cache {
   /** Keep only `keep` rows along the batch axis (eviction). */
   filter(keep: number[]): void {
     if (this.keys && this.values) {
-      const idxArr = MlxArray.fromInt32(Int32Array.from(keep), [keep.length]);
-      const k = ops.takeAxis(this.keys, idxArr, 0);
-      const v = ops.takeAxis(this.values, idxArr, 0);
-      idxArr.dispose();
+      const k = plainRowStorage.takeRows(this.keys, keep);
+      const v = plainRowStorage.takeRows(this.values, keep);
       this.keys.dispose();
       this.values.dispose();
       this.keys = k;
       this.values = v;
     }
-    this.offsetArr = keep.map((i) => this.offsetArr[i]!);
-    this.leftPad = keep.map((i) => this.leftPad[i]!);
+    this.#rows.filter(keep);
     this.releaseRopeArr();
   }
 
+  filterRows(keep: readonly number[]): void { this.filter([...keep]); }
+
   state(): MlxArray[] {
     return this.keys && this.values ? [this.keys, this.values] : [];
   }
   isTrimmable(): boolean {
-    return this.#offset < this.maxSize;
+    return this.#rows.trimmable;
   }
   trim(n: number): void {
-    const k = Math.min(this.#offset, n);
-    this.#offset -= k;
-    this.#idx -= k;
-    this.offsetArr = this.offsetArr.map((x) => x - k);
+    this.#rows.trim(n);
   }
   dispose(): void {
     this.keys?.dispose();
@@ -322,24 +286,11 @@ export class BatchedRotatingCache implements Cache {
     const width = Math.max(...lens, 0);
     const leftPad = lens.map((l) => width - l);
     const cache = new BatchedRotatingCache(maxSize, leftPad);
-    cache.offsetArr = [...offsets];
+    cache.#rows.restoreMerged(width, offsets);
     if (width === 0) return cache;
 
-    const padTo = (a: MlxArray, pad: number): MlxArray => {
-      if (pad === 0) return a.slice([0, 0, 0, 0], a.shape as number[]);
-      const [B, H, , D] = a.shape as [number, number, number, number];
-      const z = ops.zeros([B, H, pad, D], a.dtype);
-      const out = ops.concatAxis([z, a], 2);
-      z.dispose();
-      return out;
-    };
-    const ks = rows.map((r, i) => padTo(r.keys, leftPad[i]!));
-    const vs = rows.map((r, i) => padTo(r.values, leftPad[i]!));
-    cache.keys = ops.concatAxis(ks, 0);
-    cache.values = ops.concatAxis(vs, 0);
-    for (const a of [...ks, ...vs]) a.dispose();
-    cache.#idx = width;
-    cache.#offset = width;
+    cache.keys = mergeStorageRows(plainRowStorage, rows.map((row) => row.keys), leftPad);
+    cache.values = mergeStorageRows(plainRowStorage, rows.map((row) => row.values), leftPad);
     return cache;
   }
 }

package/src/model/batched-row-storage.ts

@@ -0,0 +1,127 @@
+import { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { BatchedRotatingState } from "./batched-rotating-state";
+
+export interface RowStorage<T> {
+  shape(value: T): readonly [number, number, number, number];
+  slice(value: T, batchFrom: number, batchTo: number, tokenFrom: number, tokenTo: number): T;
+  concatTokens(values: readonly T[]): T;
+  concatRows(values: readonly T[]): T;
+  padLeft(value: T, tokens: number): T;
+  takeRows(value: T, keep: readonly number[]): T;
+  copy(value: T): T;
+  dispose(value: T): void;
+}
+
+export const plainRowStorage: RowStorage<MlxArray> = {
+  shape: (value) => value.shape as [number, number, number, number],
+  slice(value, batchFrom, batchTo, tokenFrom, tokenTo) {
+    const [, H, , D] = value.shape as [number, number, number, number];
+    return value.slice([batchFrom, 0, tokenFrom, 0], [batchTo, H, tokenTo, D]);
+  },
+  concatTokens: (values) => ops.concatAxis([...values], 2),
+  concatRows: (values) => ops.concatAxis([...values], 0),
+  padLeft(value, tokens) {
+    if (tokens === 0) return value.slice([0, 0, 0, 0], value.shape as number[]);
+    const [B, H, , D] = value.shape as [number, number, number, number];
+    const zeros = ops.zeros([B, H, tokens, D], value.dtype);
+    const out = ops.concatAxis([zeros, value], 2);
+    zeros.dispose();
+    return out;
+  },
+  takeRows(value, keep) {
+    const indices = MlxArray.fromInt32(Int32Array.from(keep), [keep.length]);
+    const out = ops.takeAxis(value, indices, 0);
+    indices.dispose();
+    return out;
+  },
+  copy: (value) => ops.copyOf(value),
+  dispose: (value) => value.dispose(),
+};
+
+export const quantizedRowStorage: RowStorage<ops.QuantizedTensor> = {
+  shape: (value) => value.packed.shape as [number, number, number, number],
+  slice(value, batchFrom, batchTo, tokenFrom, tokenTo) {
+    const cut = (array: MlxArray): MlxArray => {
+      const [, H, , D] = array.shape as [number, number, number, number];
+      return array.slice([batchFrom, 0, tokenFrom, 0], [batchTo, H, tokenTo, D]);
+    };
+    return { packed: cut(value.packed), scales: cut(value.scales), biases: cut(value.biases) };
+  },
+  concatTokens(values) {
+    return {
+      packed: ops.concatAxis(values.map((value) => value.packed), 2),
+      scales: ops.concatAxis(values.map((value) => value.scales), 2),
+      biases: ops.concatAxis(values.map((value) => value.biases), 2),
+    };
+  },
+  concatRows(values) {
+    return {
+      packed: ops.concatAxis(values.map((value) => value.packed), 0),
+      scales: ops.concatAxis(values.map((value) => value.scales), 0),
+      biases: ops.concatAxis(values.map((value) => value.biases), 0),
+    };
+  },
+  padLeft(value, tokens) {
+    const pad = (array: MlxArray): MlxArray => plainRowStorage.padLeft(array, tokens);
+    return { packed: pad(value.packed), scales: pad(value.scales), biases: pad(value.biases) };
+  },
+  takeRows(value, keep) {
+    const take = (array: MlxArray): MlxArray => plainRowStorage.takeRows(array, keep);
+    return { packed: take(value.packed), scales: take(value.scales), biases: take(value.biases) };
+  },
+  copy(value) {
+    return {
+      packed: ops.copyOf(value.packed),
+      scales: ops.copyOf(value.scales),
+      biases: ops.copyOf(value.biases),
+    };
+  },
+  dispose(value) {
+    value.packed.dispose();
+    value.scales.dispose();
+    value.biases.dispose();
+  },
+};
+
+export function mergeStorageRows<T>(
+  storage: RowStorage<T>,
+  rows: readonly T[],
+  leftPad: readonly number[],
+): T {
+  const padded = rows.map((row, index) => storage.padLeft(row, leftPad[index]!));
+  try {
+    return storage.concatRows(padded);
+  } finally {
+    for (const row of padded) storage.dispose(row);
+  }
+}
+
+/** De-roll storage into temporal order, then cut one row or the whole batch. */
+export function temporalStorageView<T>(
+  storage: RowStorage<T>,
+  value: T,
+  state: BatchedRotatingState,
+  options: { row?: number; from?: number; to?: number; copy?: boolean } = {},
+): T {
+  const [batch, , tokens] = storage.shape(value);
+  const batchFrom = options.row ?? 0;
+  const batchTo = options.row === undefined ? batch : options.row + 1;
+  const parts = state.temporalRanges(tokens).map(([from, to]) =>
+    storage.slice(value, batchFrom, batchTo, from, to)
+  );
+  const ordered = parts.length === 1 ? parts[0]! : storage.concatTokens(parts);
+  if (parts.length > 1) for (const part of parts) storage.dispose(part);
+  const cut = storage.slice(
+    ordered,
+    0,
+    batchTo - batchFrom,
+    options.from ?? 0,
+    options.to ?? state.validLength,
+  );
+  storage.dispose(ordered);
+  if (!options.copy) return cut;
+  const owned = storage.copy(cut);
+  storage.dispose(cut);
+  return owned;
+}

package/src/model/compiled-decode.ts

@@ -1,4 +1,4 @@
-// Compiled decode step (Phase A of docs/design/optimization_plan.md).
+// Compiled decode step (Phase A of docs/archive/investigations/optimization_plan.md).
 //
 // The single-token decode graph is identical token-to-token except for
 // integer state (RoPE offset, cache write position, active length). This
@@ -33,6 +33,7 @@
 import { MlxArray } from "../mlx/array";
 import { CompiledFunction } from "../mlx/compile";
 import * as ops from "../mlx/ops";
+import { runtimeValue } from "../runtime-config";
 import {
   Gemma4Model,
   KVCache,
@@ -207,7 +208,7 @@ function slotDesc(c: AnyCache, plan: DecodeStepPlan): SlotDesc {
  *  inside quantizedSdpa's dispatch). Offsets/capacities don't — they're
  *  array values or shapeless dims. */
 function closureKey(descs: SlotDesc[]): string {
-  const flags = `nf=${process.env.MLX_BUN_NO_FUSED_SDPA === "1" ? 1 : 0}`;
+  const flags = `nf=${runtimeValue("MLX_BUN_NO_FUSED_SDPA") === "1" ? 1 : 0}`;
   return descs.map((d) => `${d.kind}:${d.groupSize}:${d.bits}`).join(",") + "|" + flags;
 }
 
@@ -555,7 +556,7 @@ export class CompiledDecode {
       anyCaches
         .map((c, i) => (phases[i] === "ring" ? `${slotDesc(c, plans[i]!).kind}:${slotDesc(c, plans[i]!).groupSize}:${slotDesc(c, plans[i]!).bits}` : "js"))
         .join(",") +
-      `|nf=${process.env.MLX_BUN_NO_FUSED_SDPA === "1" ? 1 : 0}`;
+      `|nf=${runtimeValue("MLX_BUN_NO_FUSED_SDPA") === "1" ? 1 : 0}`;
     if (this.#broken.has(key)) throw new Error(`compiled decode: known-broken closure ${key}`);
 
     let closures = this.#segClosures.get(key);

package/src/model/diffusion-gemma.ts

@@ -1,3 +1,4 @@
+import { cleanupFailure, disposeResources } from "../engine/resources";
 // DiffusionGemma-26B-A4B-it (model_type "diffusion_gemma") — the first
 // NON-autoregressive model in the codebase. Instead of a left-to-right AR loop
 // it runs an ENCODER prefill over the prompt (builds a KV cache) and then a
@@ -10,7 +11,7 @@
 // Ported verbatim from optiq/vlm/_mlxvlm/models/diffusion_gemma/language.py.
 // Reuses mlx-bun's quantized primitives (QuantizedLinear/Embedding/SwitchLinear,
 // RMSNorm, KVCache/RotatingKVCache) from gemma4-base. Architecture deltas vs
-// gemma4 (verified against the checkpoint, see docs/design/diffusion-gemma-port.md):
+// gemma4 (verified against the checkpoint, see docs/design/generic-model-support.md):
 //   - attention scale=1.0, NO attention softcap (only final-logit softcap 30 fp32)
 //   - QK/V-norm POST-proj / PRE-RoPE; v_norm is RMSNormNoScale; V gets no RoPE
 //   - full-attention layers (5,11,17,23,29) reuse k as v (no v_proj), hd=512 kv=2,
@@ -626,8 +627,8 @@ export class DiffusionGemmaModel {
   /** Prefill from a pre-spliced image+text prompt → populated cache. */
   prefillVision(splicedIds: number[], pixels: MlxArray): Cache[] {
     const cache = this.makeCache();
-    this.encodeVisionPrefill(splicedIds, pixels, cache);
-    return cache;
+    try { this.encodeVisionPrefill(splicedIds, pixels, cache); return cache; }
+    catch (error) { cleanupFailure(error, () => disposeResources(cache)); }
   }
 
   /** One KVCache (full) / RotatingKVCache (sliding) per layer — the encoder
@@ -899,16 +900,15 @@ export class DiffusionGemmaModel {
    *  loop reuses it across all steps of a canvas block). */
   prefill(promptIds: number[]): Cache[] {
     const cache = this.makeCache();
-    this.encodePrompt(promptIds, cache);
-    return cache;
+    try { this.encodePrompt(promptIds, cache); return cache; }
+    catch (error) { cleanupFailure(error, () => disposeResources(cache)); }
   }
 
   /** One denoising-step decoder pass -> softcapped logits [1, canvas, vocab]. */
   decoderLogits(canvasArr: MlxArray, cache: Cache[], scEmbeddings: MlxArray | null): MlxArray {
     const h = this.decodeCanvasArr(canvasArr, cache, null, scEmbeddings);
-    const logits = this.logitsFromHidden(h);
-    h.dispose();
-    return logits;
+    try { return this.logitsFromHidden(h); }
+    finally { h.dispose(); }
   }
 
   /** Dequantized embedding table (once per call), for self-conditioning soft

package/src/model/factory.ts

@@ -1,8 +1,7 @@
-// Model construction with generated-specialization dispatch
-// (docs/design/optimization_plan.md Phase C): pick the generated class whose config
-// fingerprint matches, else the monolith. Generated classes subclass
-// Gemma4Model and only override forwardLayers (with their own
-// cache-signature guard), so the choice is always safe.
+// Model construction consumes one declared profile. Exact artifact profiles
+// outrank family profiles; dedicated/generated graphs outrank the universal
+// fallback. Request-level methods are resolved elsewhere and are never changed
+// by this module.
 
 import { loadModelConfig, type ModelConfig } from "../config";
 import { totalmem } from "node:os";
@@ -24,13 +23,56 @@ import {
   type Glm52MemoryPlan,
 } from "./glm52-memory";
 import { UniversalDenseModel } from "./universal/dense";
-import { genericArgsFor, GENERIC_MODEL_TYPES, remapModelType } from "./universal/archs";
-import { isDiffusionGemmaConfig, isGlm52Config, isMiniCPM5Config, isQwen35Config, isQwen3Config, isQwen3MoeConfig } from "./support";
+import { genericArgsFor } from "./universal/archs";
+import {
+  ModelImplementationRegistry,
+  type ModelImplementation,
+  type ModelImplementationProvider,
+} from "./implementation";
+import {
+  assertResolvedModelProfile,
+  resolveModelProfile,
+  type ResolvedModelProfile,
+  type ResolveModelProfileOptions,
+  type ModelGraph,
+} from "./profile";
 
 export type RuntimeModel =
   | Gemma4Model | MiniCPM5Model | Qwen35Model | Qwen3Model | Qwen3MoeModel
   | DiffusionGemmaModel | Glm52Model | UniversalDenseModel;
 
+function residentImplementation(
+  id: string,
+  graph: ModelGraph,
+  create: ModelImplementation<Weights, RuntimeModel>["create"],
+): ModelImplementation<Weights, RuntimeModel> {
+  return { id, graph, loader: "safetensors",
+    loop: graph === "diffusion-gemma" ? "diffusion" : "autoregressive", create };
+}
+
+/** Engine-owned registrations. 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>([
+  residentImplementation("diffusion-gemma", "diffusion-gemma", (weights, config) => new DiffusionGemmaModel(weights, config)),
+  residentImplementation("minicpm5", "minicpm5", (weights, config) => new MiniCPM5Model(weights, config)),
+  residentImplementation("qwen3.5", "qwen3.5", (weights, config) => new Qwen35Model(weights, config)),
+  residentImplementation("qwen3-moe", "qwen3-moe", (weights, config) => new Qwen3MoeModel(weights, config)),
+  residentImplementation("qwen3", "qwen3", (weights, config) => new Qwen3Model(weights, config)),
+  residentImplementation("gemma4", "gemma4", (weights, config) => new Gemma4Model(weights, config)),
+  residentImplementation("gemma4-generated", "gemma4", (weights, config) => {
+    const Graph = GENERATED.get(configFingerprint(config));
+    if (!Graph) throw new Error("no generated Gemma graph for this config; refusing to fall back");
+    return new Graph(weights, config);
+  }),
+  residentImplementation("universal-dense", "universal-dense", (weights, config) =>
+    new UniversalDenseModel(weights, config, genericArgsFor(config)!)),
+]);
+
+export interface ModelOpenOptions<Model = RuntimeModel> extends Glm52RuntimeOpenOptions {
+  readonly profiles?: ResolveModelProfileOptions;
+  readonly implementations?: ModelImplementationProvider<Weights, Model>;
+}
+
 export interface Glm52RuntimeOpenOptions {
   /** Whole-process ceiling. Defaults to the smaller of the validated 25 GiB
    * preset and physical RAM. */
@@ -97,55 +139,51 @@ export async function openGlm52RuntimeModel(
  * model.safetensors.index.json, so it must bypass Weights.open and use its
  * dedicated header catalog/tensor source.
  */
-export async function openModel(
+export function openModel<Model>(
   modelDir: string,
-  options: Glm52RuntimeOpenOptions = {},
-): Promise<RuntimeModel> {
+  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);
-  if (isGlm52Config(config))
+  const profile = resolveModelProfile(config, options.profiles);
+  if (profile.profile.execution.loader === "colibri") {
+    if (options.implementations || profile.profile.execution.implementation !== undefined)
+      throw new Error("custom Colibri implementations require a streamed loader binding; refusing to fall back");
     return (await openGlm52RuntimeModel(modelDir, options)).model;
-  return createModel(await Weights.open(modelDir), config);
+  }
+  // Resolve before opening weights; missing or incompatible code allocates nothing.
+  const implementation = (options.implementations ?? MLX_MODEL_IMPLEMENTATIONS).select(config, profile);
+  const weights = await Weights.open(modelDir);
+  try { return implementation.create(weights, config, profile); }
+  catch (error) { weights.dispose(); throw error; }
 }
 
-/** Dispatch ladder (docs/design/generic-model-support.md §3.3):
- *  dedicated → generated/monolith (gemma4*) → generic (Tier-0 universal)
- *  → reject with a helpful error. Generic never shadows a dedicated port. */
-export function createModel(weights: Weights, config: ModelConfig): RuntimeModel {
-  // 1. Dedicated classes.
-  if (isGlm52Config(config)) {
+/** Construct the binding named by a validated profile. 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, resolved: ResolvedModelProfile | undefined,
+  implementations: ModelImplementationProvider<Weights, Model>,
+): Model;
+export function createModel(
+  weights: Weights, config: ModelConfig, resolved?: ResolvedModelProfile,
+  implementations?: ModelImplementationProvider<Weights, RuntimeModel>,
+): RuntimeModel;
+export function createModel<Model>(
+  weights: Weights,
+  config: ModelConfig,
+  resolved: ResolvedModelProfile = resolveModelProfile(config),
+  implementations: ModelImplementationProvider<Weights, Model> | ModelImplementationProvider<Weights, RuntimeModel> = MLX_MODEL_IMPLEMENTATIONS,
+): Model | RuntimeModel {
+  assertResolvedModelProfile(config, resolved);
+  if (resolved.profile.execution.loader === "colibri")
     throw new Error(
       "glm_moe_dsa uses the direct Colibri container; construct it with " +
       "openModel(modelDir) or Glm52Model.open(modelDir)",
     );
-  }
-  // DiffusionGemma is non-autoregressive — generate() detects it and routes to
-  // the denoising engine instead of the AR decode loop. It exposes the shared
-  // RuntimeModel surface (config/weightsBytes/loraState/makeCache) with the
-  // AR-only methods as throwing stubs (never called on this model).
-  if (isDiffusionGemmaConfig(config)) return new DiffusionGemmaModel(weights, config);
-  if (isMiniCPM5Config(config)) return new MiniCPM5Model(weights, config);
-  if (isQwen35Config(config)) return new Qwen35Model(weights, config);
-  if (isQwen3MoeConfig(config)) return new Qwen3MoeModel(weights, config);
-  if (isQwen3Config(config)) return new Qwen3Model(weights, config);
-  // 2. gemma4*: fingerprint-matched generated specialization, else monolith.
-  //    Always prefer the model-specific generated file when its config
-  //    fingerprint matches — the parity regime is selected by the OUTPUT-changing
-  //    flags (kv-quant), which the generated forward already gates, NOT by
-  //    detouring an optimized model through the generic monolith.
-  if (config.modelType.startsWith("gemma4")) {
-    const cls = GENERATED.get(configFingerprint(config)) ?? Gemma4Model;
-    return new cls(weights, config);
-  }
-  // 3. Generic Tier-0 fallback: the universal-dense descriptor table
-  // (replaces the old hard "llama = MiniCPM5 only" / unknown-type throws).
-  const generic = genericArgsFor(config);
-  if (generic) return new UniversalDenseModel(weights, config, generic);
-  // 4. Reject, naming the arch (post-remap) and the declared surface.
-  const arch = remapModelType(config.modelType);
-  throw new Error(
-    `unsupported model_type "${config.modelType}"` +
-    (arch !== config.modelType ? ` (mlx-lm remaps it to "${arch}")` : "") +
-    ` — targeted: gemma4*, diffusion_gemma, glm_moe_dsa, qwen3_5, qwen3, qwen3_moe, MiniCPM5;` +
-    ` generic (Tier-0): ${[...GENERIC_MODEL_TYPES].sort().join(", ")}`,
-  );
+  return implementations.select(config, resolved).create(weights, config, resolved);
 }

package/src/model/fingerprint.ts

@@ -1,4 +1,4 @@
-// Config fingerprint for generated-model dispatch (docs/design/optimization_plan.md
+// Config fingerprint for generated-model dispatch (docs/archive/investigations/optimization_plan.md
 // Phase C): a stable hash over every config field that changes the
 // decode graph structure, plus the kv_config quant layout. A generated
 // specialization is used only when the fingerprint of the loaded config