mlx-bun 0.2.0 → 0.4.0

package/src/model/batched-turboquant-kv.ts

@@ -0,0 +1,93 @@
+import { decodedKvDonorAttention } from "./decoded-kv-donor";
+import type { MlxArray } from "../mlx/array";
+import { KvTensorRows, type KvTensorRowView } from "../backends/mlx/kv-tensor-rows";
+import { disposeResources } from "../engine/resources";
+import { TurboQuantKVCache, type BatchableCache, type Cache, type Mask, type RotatedValueAttentionState, type PaddedPrefillCache, type PrefillPadding } from "./gemma4-base";
+import { TurboQuantCodec, disposeTurboQuant, type TurboQuantTensor } from "./turboquant-codec";
+
+const fields = ["kIdx", "kScales", "kZeros", "vPacked", "vScales"] as const;
+function encoded(planes: readonly MlxArray[]): TurboQuantTensor {
+  return { kIdx: planes[0]!, kScales: planes[1]!, kZeros: planes[2]!, vPacked: planes[3]!, vScales: planes[4]! };
+}
+
+/** The existing TurboQuant codec over shared tensor-row storage. Row changes
+ * copy encoded bytes without decoding, rotating or requantizing their values. */
+export class BatchedTurboQuantKVCache implements BatchableCache, PaddedPrefillCache {
+  readonly #storage = new KvTensorRows();
+  readonly #codec: TurboQuantCodec;
+  #reuseOffsets: number[] = [];
+  #headDim: number | null = null;
+  constructor(readonly kBits: number, readonly vBits: number) {
+    this.#codec = new TurboQuantCodec(kBits, vBits, process.env.MLX_BUN_TURBOQUANT_FUSED_DECODE === "1");
+  }
+  restorePrefillEnds(ends: readonly number[] | undefined): void { this.#storage.restorePrefillEnds(ends); }
+  preparePrefill(padding: PrefillPadding): void {
+    this.#storage.preparePrefill(padding);
+    if (!this.#reuseOffsets.length) this.#reuseOffsets = padding.lengths.map(() => 0);
+  }
+  finalizePrefill(): void { this.#storage.finalizePrefill(); }
+  captureDonorRows(): import("./gemma4-base").KvDonorRows {
+    if (!this.#storage.planes.length || this.#headDim === null) throw new Error("cache is empty");
+    const [keys, values] = this.#codec.decode(encoded(this.#storage.planes), this.offset, this.#headDim, false);
+    return { keys, values, ...this.#storage.captureDonorValidity() };
+  }
+  captureDonorAttention(): import("./gemma4-base").KvDonorAttention {
+    return decodedKvDonorAttention(this.captureDonorRows());
+  }
+  get rotatedValueAttention(): RotatedValueAttentionState { return this; }
+  signature(): string { return `kv:batched-turboquant:${this.kBits}:${this.vBits}`; }
+  get minimumReusableOffset(): number { return Math.max(0, ...this.#reuseOffsets); }
+  get headDim(): number | null { return this.#headDim; }
+  get rowOffsets(): readonly number[] { return this.#storage.rowOffsets; }
+  get leftPad(): readonly number[] { return this.#storage.leftPad; }
+  get batchSize(): number | null { return this.#storage.batchSize; }
+  get offset(): number { return this.#storage.offset; }
+  get ropeOffsetArr(): MlxArray | undefined { return this.#storage.ropeOffsetArr; }
+  makeEmptyBatch(): BatchedTurboQuantKVCache { return new BatchedTurboQuantKVCache(this.kBits, this.vBits); }
+  bytesPerToken(): number { return this.#storage.bytesPerToken(); }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * tokens; }
+  makeMask(tokens: number, window: number | null): Mask { return this.#storage.makeMask(tokens, window); }
+  isTrimmable(): boolean { return true; }
+  trim(count: number): void { this.#storage.trim(count); }
+  specRoundBegin(): void { this.#storage.specRoundBegin(); }
+  specRoundCommit(): void { this.#storage.specRoundCommit(); }
+  specRoundRollback(keep: number | readonly number[]): void { this.#storage.specRoundRollback(keep); }
+  #append(k: MlxArray, v: MlxArray, defer: boolean): [MlxArray, MlxArray] {
+    const packed = this.#codec.encode(k, v);
+    try {
+      this.#storage.append(fields.map(field => packed[field]));
+      this.#headDim = k.shape[3]!;
+      return this.#codec.decode(encoded(this.#storage.planes), this.offset, this.#headDim, defer);
+    } finally { disposeTurboQuant(packed); }
+  }
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.#append(k, v, false); }
+  updateAndFetchDeferredV(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.#append(k, v, true); }
+  mergeRows(rows: readonly Cache[]): void {
+    const views: KvTensorRowView[] = [], held: MlxArray[] = [];
+    const reuseOffsets: number[] = [];
+    let dim = this.#headDim;
+    try {
+      for (const row of rows) {
+        if (row instanceof BatchedTurboQuantKVCache) {
+          views.push(row.#storage); dim ??= row.headDim; reuseOffsets.push(...row.#reuseOffsets);
+        } else if (row instanceof TurboQuantKVCache) {
+          const planes = row.state(); held.push(...planes); reuseOffsets.push(row.minimumReusableOffset);
+          views.push({ planes, rowOffsets: [row.offset], leftPad: [0] }); dim ??= row.headDim;
+        } else throw new Error(`TurboQuant row layout cannot merge ${row.signature()}`);
+      }
+      this.#storage.mergeRows(views); this.#headDim = dim; this.#reuseOffsets = reuseOffsets;
+    } finally { disposeResources(held); }
+  }
+  alignRows(leftPad: readonly number[]): void { this.#storage.alignRows(leftPad); }
+  filterRows(keep: readonly number[]): void {
+    this.#storage.filterRows(keep); this.#reuseOffsets = keep.map(row => this.#reuseOffsets[row]!);
+  }
+  state(): MlxArray[] { return this.#storage.planes; }
+  extractRow(row: number): TurboQuantKVCache {
+    const result = new TurboQuantKVCache(this.kBits, this.vBits), count = this.rowOffsets[row]!;
+    if (count) result.restoreState(encoded(this.#storage.extractRow(row)), count, this.#headDim!);
+    result.minimumReusableOffset = this.#reuseOffsets[row] ?? 0;
+    return result;
+  }
+  dispose(): void { this.#storage.dispose(); this.#headDim = null; this.#reuseOffsets = []; }
+}

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
@@ -277,7 +277,10 @@ function decodePhase(c: AnyCache): "concat" | "ring" {
 }
 
 function disposeShared(s: SharedKv): void {
-  if (s.kind === "plain") {
+  if (s.kind === "view") {
+    s.attention.dispose();
+    s.offsetArr?.dispose();
+  } else if (s.kind === "plain") {
     s.keys.dispose();
     s.values.dispose();
   } else {
@@ -381,6 +384,9 @@ export class CompiledDecode {
   readonly #segmented: boolean;
 
   private constructor(readonly model: Gemma4Model) {
+    // A fresh resume has no prefill to evaluate lazy FP32 constants. Capture
+    // their values, not their construction graph, in every compiled closure.
+    model.materializeGraphConstants();
     const t = model.config.text;
     this.#segmented =
       t.numKvSharedLayers === 0 && !t.enableMoeBlock && model.perLayerWidth === 0;

package/src/model/decoded-kv-donor.ts

@@ -0,0 +1,12 @@
+import * as ops from "../mlx/ops";
+import type { KvDonorAttention, KvDonorRows } from "./gemma4-base";
+
+/** Takes ownership of decoded donor planes. Assistant attention uses decoded
+ * values, preserving its SDPA arithmetic rather than moving inverse rotation. */
+export function decodedKvDonorAttention(view: KvDonorRows): KvDonorAttention {
+  const { keys, values, offsets, starts, ends } = view;
+  return { width: keys.shape[2]!, dtype: keys.dtype, offsets, starts, ends,
+    attend: (q, scale, mask) => ops.sdpa(q, keys, values, scale, mask.mode, mask.arr),
+    dispose() { keys.dispose(); values.dispose(); },
+  };
+}

package/src/model/delayed-quantized-kv.ts

@@ -0,0 +1,72 @@
+import { captureFullKvDonorAttention } from "../backends/mlx/full-kv-row-donor";
+import { appendFullKvRows } from "../backends/mlx/full-kv-row-append";
+import { fullRowPadding, fullRowPhysicalLength } from "./full-prefill-row";
+import { captureKvAttention } from "./kv-attention-view";
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { FullTransitioningKvRows } from "../backends/mlx/full-transitioning-kv-rows";
+import { BatchedQuantizedKVCache } from "./batched-quantized-kv";
+import { QuantizedKVCache, type Cache, type Mask, type KvAttentionState, type KvAttentionView } from "./gemma4-base";
+
+/** Affine storage keeps its native quantized-attention arithmetic while rows
+ * cross the conversion boundary independently. Once all rows convert, their
+ * packed planes use the existing batched attention implementation. */
+export class DelayedQuantizedKVCache extends FullTransitioningKvRows<BatchedQuantizedKVCache> implements KvAttentionState {
+  constructor(readonly groupSize: number, readonly bits: number, readonly start: number,
+    readonly maintain: (rows: Cache[]) => void, row?: Cache) {
+    super({ signature: `kv:delayed-quant:${bits}:${groupSize}:${start}`, maintain,
+      converted: row => row instanceof QuantizedKVCache,
+      makeLayout: () => new BatchedQuantizedKVCache(groupSize, bits) }, row);
+  }
+  get attentionState(): KvAttentionState { return this; }
+  captureDonorAttention() {
+    return this.packed?.captureDonorAttention() ?? captureFullKvDonorAttention(this.rows, this.leftPad, this.offset);
+  }
+  makeEmptyBatch(): DelayedQuantizedKVCache { return new DelayedQuantizedKVCache(this.groupSize, this.bits, this.start, this.maintain); }
+  updateAndFetch(): [MlxArray, MlxArray] { throw new Error("mixed precision rows use their attention state"); }
+  appendAndFetch(k: MlxArray, v: MlxArray): KvAttentionView {
+    this.advance();
+    if (this.packed) return captureKvAttention(this.packed, k, v);
+    if (this.rows.every(row => !row.quantizedAttention)) {
+      const [keys, values] = appendFullKvRows(this.rows, this.leftPad, this.offset + k.shape[2]!, k, v);
+      try { this.syncPositions(true); } catch (error) { keys.dispose(); values.dispose(); throw error; }
+      return { attend: (q, scale, mask) => ops.sdpa(q, keys, values, scale, mask.mode, mask.arr),
+        dispose() { keys.dispose(); values.dispose(); } };
+    }
+    const views: KvAttentionView[] = [], pads = this.rows.map((row, index) => this.leftPad[index]! - fullRowPadding(row));
+    const lengths = this.rows.map(cache => fullRowPhysicalLength(cache) + k.shape[2]!);
+    try {
+      for (const [row, cache] of this.rows.entries()) {
+        const slice = (a: MlxArray) => a.slice([row, 0, 0, 0], [row + 1, a.shape[1]!, a.shape[2]!, a.shape[3]!]);
+        using kr = slice(k), vr = slice(v);
+        views.push(captureKvAttention(cache, kr, vr));
+      }
+      this.syncPositions(true);
+    } catch (error) { for (const view of views) view.dispose(); throw error; }
+    return {
+      attend(q, scale, mask) {
+        const outputs: MlxArray[] = [];
+        try {
+          for (const [row, view] of views.entries()) {
+            using qr = q.slice([row, 0, 0, 0], [row + 1, q.shape[1]!, q.shape[2]!, q.shape[3]!]);
+            // Capture alignment with the KV tensors, never consult live rows.
+            const start = mask.arr?.shape.map(() => 0), end = mask.arr ? [...mask.arr.shape] : undefined;
+            if (start && end) {
+              if (end.length === 4 && end[0]! > 1) { start[0] = row; end[0] = row + 1; }
+              start[start.length - 1] = pads[row]!;
+              end[end.length - 1] = pads[row]! + lengths[row]!;
+            }
+            using arr = mask.arr ? mask.arr.slice(start!, end!) : null;
+            outputs.push(view.attend(qr, scale, { mode: mask.mode, arr }));
+          }
+          return outputs.length === 1 ? ops.contiguous(outputs[0]!) : ops.concatAxis(outputs, 0);
+        } finally { for (const output of outputs) output.dispose(); }
+      },
+      dispose() { for (const view of views) view.dispose(); },
+    };
+  }
+  updateAndAttend(q: MlxArray, k: MlxArray, v: MlxArray, scale: number, mask: Mask): MlxArray {
+    const view = this.appendAndFetch(k, v);
+    try { return view.attend(q, scale, mask); } finally { view.dispose(); }
+  }
+}

package/src/model/delayed-rotating-quantized-kv.ts

@@ -0,0 +1,93 @@
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { RotatingKVCache, type Cache, type KvAttentionState, type KvAttentionView, type PaddedPrefillCache, type PrefillPadding } from "./gemma4-base";
+import { SpeculativeTransitioningKvRows } from "../backends/mlx/transitioning-kv-rows";
+import { AlignedRotatingCache, alignRotatingRows, RotatingAffineLayout, SpeculativeRotatingAffineLayout, RotatingKvPositions } from "./rotating-kv-layout";
+import { BatchedRotatingQuantCache } from "./batched-rotating-quant";
+import { captureKvAttention, combineKvDonorAttention } from "./kv-attention-view";
+
+/** Precision changes preserve each row's physical columns. The model owns
+ * queries and scale; the shared lifecycle owns membership and conversion. */
+export class DelayedRotatingQuantizedKVCache extends SpeculativeTransitioningKvRows<RotatingAffineLayout> implements KvAttentionState, PaddedPrefillCache {
+  constructor(readonly maxSize: number, readonly groupSize: number, readonly bits: number, readonly start: number,
+    readonly maintain: (rows: Cache[]) => void, row?: Cache, readonly speculative = false) {
+    super({ signature: `kv:delayed-rotating-quant:${maxSize}:${bits}:${groupSize}:${start}`, maintain,
+      converted: row => row instanceof AlignedRotatingCache && row.inner instanceof BatchedRotatingQuantCache,
+      makeLayout: () => speculative ? new SpeculativeRotatingAffineLayout(maxSize, groupSize, bits) : new RotatingAffineLayout(maxSize, groupSize, bits), prepareRows: alignRotatingRows,
+      packRows: (layout, rows) => layout.adoptAlignedRows(rows as readonly AlignedRotatingCache[]),
+      extractRow: row => (row as AlignedRotatingCache).extract(speculative ? maxSize : undefined),
+      rollbackRow: (row, before, keep, preserve) => (row as AlignedRotatingCache).rollback(before, keep, preserve) }, row, new RotatingKvPositions(maxSize));
+  }
+  get attentionState(): KvAttentionState { return this; }
+  captureDonorAttention(): import("./gemma4-base").KvDonorAttention {
+    if (this.packed) return this.packed.captureDonorAttention();
+    const views: import("./gemma4-base").KvDonorAttention[] = [];
+    try {
+      for (const row of this.rows) views.push((row as AlignedRotatingCache).captureDonorAttention());
+      return combineKvDonorAttention(views);
+    } catch (error) { for (const view of views) view.dispose(); throw error; }
+  }
+  override makeMask(tokens: number, window: number | null) {
+    if (this.speculative && this.leftPad.every(pad => pad === 0) &&
+        this.offset + tokens <= Math.min(this.maxSize, window ?? this.maxSize))
+      return { mode: tokens === 1 ? "" as const : "causal" as const, arr: null };
+    return super.makeMask(tokens, window);
+  }
+  preparePrefill(padding: PrefillPadding): void {
+    if (!this.batchSize) {
+      const rows = padding.lengths.map(() => new RotatingKVCache(this.maxSize));
+      try { this.mergeRows(rows); } finally { for (const row of rows) row.dispose(); }
+    }
+    if (this.packed) { this.packed.preparePrefill(padding); return; }
+    // Every physical row follows the same block geometry, even when only a
+    // sibling has right padding. The token lengths still belong to each row.
+    const rightPaddedBatch = padding.rightPadding?.some(pad => pad > 0);
+    for (const [row, cache] of this.rows.entries()) (cache as AlignedRotatingCache).inner.preparePrefill({
+      lengths: [padding.lengths[row]!],
+      ...(padding.leftPadding ? { leftPadding: [padding.leftPadding[row]!] } : {}),
+      ...(padding.rightPadding ? { rightPadding: [padding.rightPadding[row]!] } : {}),
+    }, rightPaddedBatch);
+    this.syncPositions(true);
+  }
+  finalizePrefill(): void {
+    if (this.packed) { this.packed.finalizePrefill(); return; }
+    for (const cache of this.rows) (cache as AlignedRotatingCache).inner.finalizePrefill();
+    this.syncPositions(true);
+  }
+  get maxTokens(): number { return this.maxSize; }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * Math.min(tokens, this.maxSize); }
+  makeEmptyBatch(): DelayedRotatingQuantizedKVCache {
+    return new DelayedRotatingQuantizedKVCache(this.maxSize, this.groupSize, this.bits, this.start, this.maintain, undefined, this.speculative);
+  }
+  updateAndFetch(): [MlxArray, MlxArray] { throw new Error("mixed precision rows use their attention state"); }
+  releaseRopeArr(): void { this.packed?.releaseRopeArr(); }
+  appendAndFetch(k: MlxArray, v: MlxArray): KvAttentionView {
+    this.advance();
+    if (this.packed) return captureKvAttention(this.packed, k, v);
+    const views: KvAttentionView[] = [];
+    try {
+      for (const [row, cache] of this.rows.entries()) {
+        const slice = (a: MlxArray) => a.slice([row, 0, 0, 0], [row + 1, a.shape[1]!, a.shape[2]!, a.shape[3]!]);
+        using kr = slice(k), vr = slice(v);
+        views.push(captureKvAttention(cache, kr, vr));
+      }
+      this.syncPositions(true);
+    } catch (error) { for (const view of views) view.dispose(); throw error; }
+    return {
+      attend(q, scale, mask) {
+        const outputs: MlxArray[] = [];
+        try {
+          for (const [row, view] of views.entries()) {
+            using qr = q.slice([row, 0, 0, 0], [row + 1, q.shape[1]!, q.shape[2]!, q.shape[3]!]);
+            const start = mask.arr?.shape.map(() => 0), end = mask.arr ? [...mask.arr.shape] : undefined;
+            if (start && end && end.length === 4 && end[0]! > 1) { start[0] = row; end[0] = row + 1; }
+            using arr = mask.arr ? mask.arr.slice(start!, end!) : null;
+            outputs.push(view.attend(qr, scale, { mode: mask.mode, arr }));
+          }
+          return outputs.length === 1 ? ops.contiguous(outputs[0]!) : ops.concatAxis(outputs, 0);
+        } finally { for (const output of outputs) output.dispose(); }
+      },
+      dispose() { for (const view of views) view.dispose(); },
+    };
+  }
+}

package/src/model/delayed-turboquant-kv.ts

@@ -0,0 +1,55 @@
+import { captureFullKvDonorRows } from "../backends/mlx/full-kv-row-donor";
+import { decodedKvDonorAttention } from "./decoded-kv-donor";
+import { appendFullKvRows } from "../backends/mlx/full-kv-row-append";
+import type { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import * as ops from "../mlx/ops";
+import { unrotateValues } from "../mlx/turboquant-ops";
+import { type Cache, type RotatedValueAttentionState, TurboQuantKVCache } from "./gemma4-base";
+import { BatchedTurboQuantKVCache } from "./batched-turboquant-kv";
+import { FullTransitioningKvRows } from "../backends/mlx/full-transitioning-kv-rows";
+
+/** TurboQuant tensor access and value-domain handling over the shared
+ * row-transition lifecycle. Model projections and attention remain batched. */
+export class DelayedTurboQuantKVCache extends FullTransitioningKvRows<BatchedTurboQuantKVCache> implements RotatedValueAttentionState {
+  #rotated: boolean[] = [];
+  constructor(readonly kBits: number, readonly vBits: number, readonly start: number,
+    readonly maintain: (rows: Cache[]) => void, row?: Cache) {
+    super({ signature: `kv:delayed-turboquant:${kBits}:${vBits}:${start}`, maintain,
+      converted: row => row instanceof TurboQuantKVCache,
+      makeLayout: () => new BatchedTurboQuantKVCache(kBits, vBits) }, row);
+  }
+  captureDonorRows(): import("./gemma4-base").KvDonorRows {
+    return this.packed?.captureDonorRows() ?? captureFullKvDonorRows(this.rows, this.leftPad, this.offset);
+  }
+  captureDonorAttention(): import("./gemma4-base").KvDonorAttention {
+    return decodedKvDonorAttention(this.captureDonorRows());
+  }
+  get rotatedValueAttention(): RotatedValueAttentionState { return this; }
+  #append(k: MlxArray, v: MlxArray, deferred: boolean): [MlxArray, MlxArray] {
+    this.advance();
+    if (this.packed) {
+      this.#rotated = Array.from({ length: this.batchSize! }, () => true);
+      return deferred ? this.packed.updateAndFetchDeferredV(k, v) : this.packed.updateAndFetch(k, v);
+    }
+    this.#rotated = this.rows.map(row => !!row.rotatedValueAttention);
+    const result = appendFullKvRows(this.rows, this.leftPad, this.offset + k.shape[2]!, k, v, deferred);
+    try { this.syncPositions(true); } catch (error) { result[0].dispose(); result[1].dispose(); throw error; }
+    return result;
+  }
+
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.#append(k, v, false); }
+  updateAndFetchDeferredV(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.#append(k, v, true); }
+  captureValueTransform(): (output: MlxArray) => MlxArray {
+    const rotated = [...this.#rotated];
+    if (rotated.every(value => value)) return unrotateValues;
+    if (rotated.every(value => !value)) return output => ops.contiguous(output);
+    return output => {
+      using restored = unrotateValues(output);
+      using flags = ops.fromInt32(rotated.map(value => Number(value)), [rotated.length, 1, 1, 1]);
+      using mask = flags.astype(Dtype.bool);
+      return ops.where(mask, restored, output);
+    };
+  }
+  makeEmptyBatch(): DelayedTurboQuantKVCache { return new DelayedTurboQuantKVCache(this.kBits, this.vBits, this.start, this.maintain); }
+}

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

@@ -24,16 +24,55 @@ import {
 } from "./glm52-memory";
 import { UniversalDenseModel } from "./universal/dense";
 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. */
@@ -100,47 +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);
-  const profile = resolveModelProfile(config);
-  if (profile.profile.execution.loader === "colibri")
+  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, profile);
+  }
+  // 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; }
 }
 
-/** Construct the graph named by a previously validated profile. */
+/** 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),
-): RuntimeModel {
+  implementations: ModelImplementationProvider<Weights, Model> | ModelImplementationProvider<Weights, RuntimeModel> = MLX_MODEL_IMPLEMENTATIONS,
+): Model | RuntimeModel {
   assertResolvedModelProfile(config, resolved);
-  const execution = resolved.profile.execution;
-  switch (execution.graph) {
-    case "glm5.2":
-      throw new Error(
-        "glm_moe_dsa uses the direct Colibri container; construct it with " +
-        "openModel(modelDir) or Glm52Model.open(modelDir)",
-      );
-    // DiffusionGemma is non-autoregressive. The profile pairs this graph with
-    // the denoising loop; generate() recognizes the constructed model surface.
-    case "diffusion-gemma": return new DiffusionGemmaModel(weights, config);
-    case "minicpm5": return new MiniCPM5Model(weights, config);
-    case "qwen3.5": return new Qwen35Model(weights, config);
-    case "qwen3-moe": return new Qwen3MoeModel(weights, config);
-    case "qwen3": return new Qwen3Model(weights, config);
-    case "gemma4": {
-      // Generated specializations remain exact config matches. Output-changing
-      // request methods stay independently resolved and cannot be fallbacks.
-      const cls = execution.specialization === "generated"
-        ? GENERATED.get(configFingerprint(config))!
-        : Gemma4Model;
-      return new cls(weights, config);
-    }
-    case "universal-dense":
-      return new UniversalDenseModel(weights, config, genericArgsFor(config)!);
-  }
+  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)",
+    );
+  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

package/src/model/flash-attention.ts

@@ -545,7 +545,7 @@ export function flashBackward(
     // (dK_j[d] → [kv_row][d]); no post-transpose. (An earlier "fix" transposed
     // it on the false belief the buffer was swapped — that transpose is only the
     // identity at Tkv==D and corrupted dK for Tkv≠D; removed after
-    // scripts/experiments/flash-dkv-debug.ts showed the raw output is correct.)
+    // flash-dkv-debug.ts (deleted 2026-08-23; git history) showed the raw output is correct.)
     return [dq[0]!, dkv[0]!, dkv[1]!];
   } finally {
     Dvec.dispose();
@@ -621,8 +621,8 @@ export function flashAttention(
 //  - This flash kernel is the L2 path: a port of mlx-optiq's
 //    flash_attention_metal (O(L) backward memory). Its parity oracle is OPTIQ,
 //    not finite differences. Confirmed: optiq's flash dK == ops.sdpa to f16
-//    (scripts/flash-optiq-check.py, rel 0.0%), and mlx-bun's flash == ops.sdpa
-//    (tests/flash-attention.test.ts) ⟹ mlx-bun flash == optiq flash (L2 parity).
+//    (scripts/oracle/flash-optiq-check.py, rel 0.0%), and mlx-bun's flash == ops.sdpa
+//    (tests/unit/flash-attention.test.ts) ⟹ mlx-bun flash == optiq flash (L2 parity).
 //
 // Two real port bugs were fixed to reach that parity (the original diverged
 // ~100% from optiq on dK): a spurious dK transpose in flashBackward (corrupted

package/src/model/full-prefill-padding.ts

@@ -0,0 +1,56 @@
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import type { PrefillPadding } from "./gemma4-base";
+import { rowRollIndices } from "./batched-row-storage";
+import { disposeResources } from "../engine/resources";
+
+interface Positions { rowOffsets: number[]; leftPad: number[] }
+
+/** mlx-lm BatchKVCache.prepare/finalize (MIT), over arbitrary stored planes.
+ * Padding changes token positions; the representation's codec is uninvolved. */
+export class FullPrefillPadding {
+  constructor(readonly tokenAxis = 2) {}
+  #ends?: number[];
+  get endOffsets(): readonly number[] | undefined { return this.#ends; }
+  restoreEnds(ends: readonly number[] | undefined): void { this.#ends = ends ? [...ends] : undefined; }
+  validOffset(offset: number, row: number): number { return Math.min(offset, this.#ends?.[row] ?? offset); }
+
+  prepare(positions: Positions, padding: PrefillPadding): void {
+    if (!positions.rowOffsets.length) {
+      positions.rowOffsets = padding.lengths.map(() => 0);
+      positions.leftPad = [...positions.rowOffsets];
+    }
+    if (padding.leftPadding) {
+      positions.leftPad = positions.leftPad.map((pad, row) => pad + padding.leftPadding![row]!);
+      positions.rowOffsets = positions.rowOffsets.map((offset, row) => offset - padding.leftPadding![row]!);
+    }
+    // Like rotating storage, retain absolute valid ends. A cohort can finish
+    // before its planned width when a longer request leaves during prefill.
+    this.#ends = padding.rightPadding?.some(pad => pad > 0)
+      ? padding.lengths.map((length, row) => positions.rowOffsets[row]! + length) : undefined;
+  }
+
+  /** Borrow planes; return owned replacements only when padding needs moving.
+   * Stage every plane before committing positions, including on FFI failure. */
+  finalize(planes: readonly MlxArray[], positions: Positions): MlxArray[] | undefined {
+    if (!this.#ends) return;
+    const padding = this.#ends.map((end, row) => Math.max(0, positions.rowOffsets[row]! - end));
+    if (!padding.some(pad => pad > 0)) { this.#ends = undefined; return; }
+    const next: MlxArray[] = [];
+    try {
+      if (planes.length) {
+        const width = planes[0]!.shape[this.tokenAxis]!;
+        using indices = rowRollIndices(width, padding);
+        using compressed = this.tokenAxis === 1 ? ops.reshape(indices, [padding.length, width, 1]) : null;
+        for (const plane of planes) next.push(ops.takeAlongAxis(plane, compressed ?? indices, this.tokenAxis));
+      }
+    } catch (error) { disposeResources(next); throw error; }
+    positions.rowOffsets = positions.rowOffsets.map((offset, row) => offset - padding[row]!);
+    positions.leftPad = positions.leftPad.map((pad, row) => pad + padding[row]!);
+    this.#ends = undefined;
+    return next;
+  }
+
+  filter(keep: readonly number[]): void { if (this.#ends) this.#ends = keep.map(row => this.#ends![row]!); }
+  clear(): void { this.#ends = undefined; }
+}