mlx-bun 0.0.13 → 0.3.0

package/src/model/generated/gemma4-26b.ts

@@ -4,14 +4,15 @@
 // fingerprint: c9dd67ed5a525231
 //
 // Branch-resolved, unrolled forward pass for this architecture
-// (docs/design/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
-// generator's gate: tests/generated-parity.test.ts. The cache-signature
+// (docs/archive/investigations/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
+// generator's gate: tests/parity/generated-parity.test.ts. The cache-signature
 // guard falls back to the monolith for anything this file wasn't
 // generated for (bf16 compat runs, vision bidir prefill).
 
 import { MlxArray } from "../../mlx/array";
 import * as ops from "../../mlx/ops";
 import {
+  cacheSignature,
   disposing,
   fusedSdpaRuntimeOk,
   KVCache,
@@ -293,101 +294,40 @@ function donorFullQuantG64B4KeqVMoeScal(layer: DecoderLayer, x: MlxArray, mask:
 }
 
 export class GeneratedGemma4 extends Gemma4Model {
-  /** Per-layer cache classes must match what this file was generated
-   *  for; anything else runs the monolith path. */
+  /** Per-layer cache signatures must match what this file was generated for. */
   #matches(cache: Cache[]): boolean {
     return (
       cache.length === 30 &&
-      cache[0] instanceof RotatingQuantizedKVCache &&
-      cache[0].bits === 8 &&
-      cache[0].groupSize === 64 &&
-      cache[1] instanceof RotatingQuantizedKVCache &&
-      cache[1].bits === 4 &&
-      cache[1].groupSize === 64 &&
-      cache[2] instanceof RotatingQuantizedKVCache &&
-      cache[2].bits === 4 &&
-      cache[2].groupSize === 64 &&
-      cache[3] instanceof RotatingQuantizedKVCache &&
-      cache[3].bits === 8 &&
-      cache[3].groupSize === 64 &&
-      cache[4] instanceof RotatingQuantizedKVCache &&
-      cache[4].bits === 8 &&
-      cache[4].groupSize === 64 &&
-      cache[5] instanceof QuantizedKVCache &&
-      cache[5].bits === 8 &&
-      cache[5].groupSize === 64 &&
-      cache[6] instanceof RotatingQuantizedKVCache &&
-      cache[6].bits === 4 &&
-      cache[6].groupSize === 64 &&
-      cache[7] instanceof RotatingQuantizedKVCache &&
-      cache[7].bits === 4 &&
-      cache[7].groupSize === 64 &&
-      cache[8] instanceof RotatingQuantizedKVCache &&
-      cache[8].bits === 8 &&
-      cache[8].groupSize === 64 &&
-      cache[9] instanceof RotatingQuantizedKVCache &&
-      cache[9].bits === 8 &&
-      cache[9].groupSize === 64 &&
-      cache[10] instanceof RotatingQuantizedKVCache &&
-      cache[10].bits === 4 &&
-      cache[10].groupSize === 64 &&
-      cache[11] instanceof QuantizedKVCache &&
-      cache[11].bits === 8 &&
-      cache[11].groupSize === 64 &&
-      cache[12] instanceof RotatingQuantizedKVCache &&
-      cache[12].bits === 8 &&
-      cache[12].groupSize === 64 &&
-      cache[13] instanceof RotatingQuantizedKVCache &&
-      cache[13].bits === 4 &&
-      cache[13].groupSize === 64 &&
-      cache[14] instanceof RotatingQuantizedKVCache &&
-      cache[14].bits === 4 &&
-      cache[14].groupSize === 64 &&
-      cache[15] instanceof RotatingQuantizedKVCache &&
-      cache[15].bits === 4 &&
-      cache[15].groupSize === 64 &&
-      cache[16] instanceof RotatingQuantizedKVCache &&
-      cache[16].bits === 4 &&
-      cache[16].groupSize === 64 &&
-      cache[17] instanceof QuantizedKVCache &&
-      cache[17].bits === 4 &&
-      cache[17].groupSize === 64 &&
-      cache[18] instanceof RotatingQuantizedKVCache &&
-      cache[18].bits === 4 &&
-      cache[18].groupSize === 64 &&
-      cache[19] instanceof RotatingQuantizedKVCache &&
-      cache[19].bits === 4 &&
-      cache[19].groupSize === 64 &&
-      cache[20] instanceof RotatingQuantizedKVCache &&
-      cache[20].bits === 4 &&
-      cache[20].groupSize === 64 &&
-      cache[21] instanceof RotatingQuantizedKVCache &&
-      cache[21].bits === 4 &&
-      cache[21].groupSize === 64 &&
-      cache[22] instanceof RotatingQuantizedKVCache &&
-      cache[22].bits === 4 &&
-      cache[22].groupSize === 64 &&
-      cache[23] instanceof QuantizedKVCache &&
-      cache[23].bits === 4 &&
-      cache[23].groupSize === 64 &&
-      cache[24] instanceof RotatingQuantizedKVCache &&
-      cache[24].bits === 4 &&
-      cache[24].groupSize === 64 &&
-      cache[25] instanceof RotatingQuantizedKVCache &&
-      cache[25].bits === 4 &&
-      cache[25].groupSize === 64 &&
-      cache[26] instanceof RotatingQuantizedKVCache &&
-      cache[26].bits === 4 &&
-      cache[26].groupSize === 64 &&
-      cache[27] instanceof RotatingQuantizedKVCache &&
-      cache[27].bits === 4 &&
-      cache[27].groupSize === 64 &&
-      cache[28] instanceof RotatingQuantizedKVCache &&
-      cache[28].bits === 4 &&
-      cache[28].groupSize === 64 &&
-      cache[29] instanceof QuantizedKVCache &&
-      cache[29].bits === 4 &&
-      cache[29].groupSize === 64
+      cacheSignature(cache[0]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[1]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[2]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[3]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[4]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[5]) === "kv:quant:8:64" &&
+      cacheSignature(cache[6]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[7]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[8]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[9]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[10]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[11]) === "kv:quant:8:64" &&
+      cacheSignature(cache[12]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[13]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[14]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[15]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[16]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[17]) === "kv:quant:4:64" &&
+      cacheSignature(cache[18]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[19]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[20]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[21]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[22]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[23]) === "kv:quant:4:64" &&
+      cacheSignature(cache[24]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[25]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[26]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[27]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[28]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[29]) === "kv:quant:4:64"
     );
   }
 

package/src/model/generated/gemma4-e4b.ts

@@ -4,14 +4,15 @@
 // fingerprint: 418e9adc386ea67c
 //
 // Branch-resolved, unrolled forward pass for this architecture
-// (docs/design/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
-// generator's gate: tests/generated-parity.test.ts. The cache-signature
+// (docs/archive/investigations/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
+// generator's gate: tests/parity/generated-parity.test.ts. The cache-signature
 // guard falls back to the monolith for anything this file wasn't
 // generated for (bf16 compat runs, vision bidir prefill).
 
 import { MlxArray } from "../../mlx/array";
 import * as ops from "../../mlx/ops";
 import {
+  cacheSignature,
   disposing,
   fusedSdpaRuntimeOk,
   KVCache,
@@ -529,83 +530,34 @@ function sharerFullQuantG64B4VPliScal(layer: DecoderLayer, x: MlxArray, mask: Ma
 }
 
 export class GeneratedGemma4 extends Gemma4Model {
-  /** Per-layer cache classes must match what this file was generated
-   *  for; anything else runs the monolith path. */
+  /** Per-layer cache signatures must match what this file was generated for. */
   #matches(cache: Cache[]): boolean {
     return (
       cache.length === 24 &&
-      cache[0] instanceof RotatingQuantizedKVCache &&
-      cache[0].bits === 4 &&
-      cache[0].groupSize === 64 &&
-      cache[1] instanceof RotatingQuantizedKVCache &&
-      cache[1].bits === 4 &&
-      cache[1].groupSize === 64 &&
-      cache[2] instanceof RotatingQuantizedKVCache &&
-      cache[2].bits === 4 &&
-      cache[2].groupSize === 64 &&
-      cache[3] instanceof RotatingQuantizedKVCache &&
-      cache[3].bits === 8 &&
-      cache[3].groupSize === 64 &&
-      cache[4] instanceof RotatingQuantizedKVCache &&
-      cache[4].bits === 8 &&
-      cache[4].groupSize === 64 &&
-      cache[5] instanceof QuantizedKVCache &&
-      cache[5].bits === 4 &&
-      cache[5].groupSize === 64 &&
-      cache[6] instanceof RotatingQuantizedKVCache &&
-      cache[6].bits === 4 &&
-      cache[6].groupSize === 64 &&
-      cache[7] instanceof RotatingQuantizedKVCache &&
-      cache[7].bits === 4 &&
-      cache[7].groupSize === 64 &&
-      cache[8] instanceof RotatingQuantizedKVCache &&
-      cache[8].bits === 4 &&
-      cache[8].groupSize === 64 &&
-      cache[9] instanceof RotatingQuantizedKVCache &&
-      cache[9].bits === 8 &&
-      cache[9].groupSize === 64 &&
-      cache[10] instanceof RotatingQuantizedKVCache &&
-      cache[10].bits === 8 &&
-      cache[10].groupSize === 64 &&
-      cache[11] instanceof QuantizedKVCache &&
-      cache[11].bits === 8 &&
-      cache[11].groupSize === 64 &&
-      cache[12] instanceof RotatingQuantizedKVCache &&
-      cache[12].bits === 4 &&
-      cache[12].groupSize === 64 &&
-      cache[13] instanceof RotatingQuantizedKVCache &&
-      cache[13].bits === 4 &&
-      cache[13].groupSize === 64 &&
-      cache[14] instanceof RotatingQuantizedKVCache &&
-      cache[14].bits === 8 &&
-      cache[14].groupSize === 64 &&
-      cache[15] instanceof RotatingQuantizedKVCache &&
-      cache[15].bits === 4 &&
-      cache[15].groupSize === 64 &&
-      cache[16] instanceof RotatingQuantizedKVCache &&
-      cache[16].bits === 4 &&
-      cache[16].groupSize === 64 &&
-      cache[17] instanceof QuantizedKVCache &&
-      cache[17].bits === 4 &&
-      cache[17].groupSize === 64 &&
-      cache[18] instanceof RotatingQuantizedKVCache &&
-      cache[18].bits === 4 &&
-      cache[18].groupSize === 64 &&
-      cache[19] instanceof RotatingQuantizedKVCache &&
-      cache[19].bits === 4 &&
-      cache[19].groupSize === 64 &&
-      cache[20] instanceof RotatingQuantizedKVCache &&
-      cache[20].bits === 4 &&
-      cache[20].groupSize === 64 &&
-      cache[21] instanceof RotatingQuantizedKVCache &&
-      cache[21].bits === 4 &&
-      cache[21].groupSize === 64 &&
-      cache[22] instanceof RotatingQuantizedKVCache &&
-      cache[22].bits === 4 &&
-      cache[22].groupSize === 64 &&
-      cache[23] instanceof QuantizedKVCache &&
-      cache[23].bits === 4 &&
-      cache[23].groupSize === 64
+      cacheSignature(cache[0]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[1]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[2]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[3]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[4]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[5]) === "kv:quant:4:64" &&
+      cacheSignature(cache[6]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[7]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[8]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[9]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[10]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[11]) === "kv:quant:8:64" &&
+      cacheSignature(cache[12]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[13]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[14]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[15]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[16]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[17]) === "kv:quant:4:64" &&
+      cacheSignature(cache[18]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[19]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[20]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[21]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[22]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[23]) === "kv:quant:4:64"
     );
   }
 

package/src/model/glm52-cache.ts

@@ -88,7 +88,7 @@ function trimToTokens(array: MlxArray, tokens: number): MlxArray {
   const stop = [...array.shape];
   stop[1] = tokens;
   const view = array.slice(start, stop);
-  const exact = ops.contiguous(view);
+  const exact = ops.copyOf(view); // TRUE copy: see 2026-08-20 contiguous-view pin class
   view.dispose();
   return exact;
 }
@@ -225,6 +225,15 @@ export class Glm52Cache implements BatchableCache {
       throw new Error("native MTP cache cannot contain target DSA state");
   }
 
+  signature(): string {
+    if (this.role === "mtp") return "kv:mla:mtp";
+    return this.dsa ? "kv:mla:target:dsa" : "kv:mla:target";
+  }
+
+  bytesPerToken(): number {
+    return (this.kvLoraRank + this.ropeHeadDim + (this.dsa?.headDim ?? 0)) * 4;
+  }
+
   /** Append one equal-width row block without expanding it to per-head K/V. */
   append(latent: MlxArray, rope: MlxArray, dsa: MlxArray | null = null): void {
     const latentShape = latent.shape;
@@ -620,15 +629,16 @@ export class Glm52Cache implements BatchableCache {
     const nextWidth = this.offset - removablePad;
     const indices = ops.fromInt32([...keep], [keep.length]);
     const filterFamily = (array: MlxArray): MlxArray => {
-      const selected = ops.takeAxis(array, indices, 0);
-      if (removablePad === 0) return selected;
-      const view = selected.slice(
-        [0, removablePad, 0],
-        [keep.length, this.offset, array.shape[2]!],
-      );
-      const exact = ops.contiguous(view);
+      const trimmed = removablePad === 0
+        ? array
+        : array.slice(
+            [0, removablePad, 0],
+            [this.batchSize!, this.offset, array.shape[2]!],
+          );
+      const selected = ops.takeAxis(trimmed, indices, 0);
+      const exact = ops.copyOf(selected); // TRUE copy: see 2026-08-20 contiguous-view pin class
       selected.dispose();
-      view.dispose();
+      if (trimmed !== array) trimmed.dispose();
       return exact;
     };
 

package/src/model/glm52-dsa.ts

@@ -291,7 +291,7 @@ function getDsaRankKeyKernel(): MetalKernel {
           ? ~bits
           : (bits ^ 0x80000000u);
         const uint descending = ~ascending;
-        keys[i] = (ulong(descending) << 32) | ulong(i);
+        keys[i] = descending;
       `,
     });
   }
@@ -308,7 +308,8 @@ function getDsaContractOrderKeyKernel(): MetalKernel {
         const uint i = thread_position_in_grid.x;
         const uint position = selected_positions[i];
         const uint threshold_class = selected_scores[i] > threshold ? 0u : 1u;
-        keys[i] = (ulong(threshold_class) << 32) | ulong(position);
+        // MLX axis lengths are signed int32; positions leave the high bit free.
+        keys[i] = (threshold_class << 31) | position;
       `,
     });
   }
@@ -318,8 +319,8 @@ function getDsaContractOrderKeyKernel(): MetalKernel {
 /**
  * Deterministic on-device equivalent of `selectDsaThresholdTiesF32`.
  *
- * The first uint64 key sorts by score descending and then position ascending,
- * making the top-k set deterministic even at the threshold. The second key
+ * Stable sorting of the first uint32 key orders scores descending and keeps
+ * equal-score positions ascending, including threshold ties. The second key
  * restores Colibri's observable two-scan order: all scores strictly above the
  * threshold in position order, followed by threshold ties in position order.
  * The caller owns the returned arrays and must call `dispose()` once the last
@@ -340,18 +341,21 @@ export function selectGlm52DsaDevice(
     ? scores
     : scores.astype(Dtype.float32);
   const [rankKeys] = getDsaRankKeyKernel().apply([scoresF32], {
-    outputs: [{ shape: [contextLength], dtype: Dtype.uint64 }],
+    outputs: [{ shape: [contextLength], dtype: Dtype.uint32 }],
     grid: [contextLength, 1, 1],
     threadGroup: [Math.min(256, contextLength), 1, 1],
   });
-  const partition = ops.argpartitionAxis(rankKeys!, topK - 1, 0);
-  const selectedUnordered = partition.slice([0], [topK]);
+  // Stable argsort retains ascending positions for equal score bits, so the
+  // ranking key needs no position field. MLX 0.31.2's GPU argpartition also
+  // sorts fully, but does not promise this tie order in its public contract.
+  const rankedPositions = ops.argsortAxis(rankKeys!, 0);
+  const selectedUnordered = rankedPositions.slice([0], [topK]);
   const selectedScores = ops.takeAlongAxis(scoresF32, selectedUnordered, 0);
   const threshold = ops.minAxis(selectedScores, 0, false);
   const [orderKeys] = getDsaContractOrderKeyKernel().apply(
     [selectedScores, selectedUnordered, threshold],
     {
-      outputs: [{ shape: [topK], dtype: Dtype.uint64 }],
+      outputs: [{ shape: [topK], dtype: Dtype.uint32 }],
       grid: [topK, 1, 1],
       threadGroup: [Math.min(256, topK), 1, 1],
     },
@@ -360,7 +364,7 @@ export function selectGlm52DsaDevice(
   const positions = ops.takeAlongAxis(selectedUnordered, order, 0);
 
   rankKeys!.dispose();
-  partition.dispose();
+  rankedPositions.dispose();
   selectedUnordered.dispose();
   selectedScores.dispose();
   orderKeys!.dispose();

package/src/model/glm52-memory.ts

@@ -4,6 +4,7 @@ import {
   type Glm52Config,
 } from "./glm52-config";
 import { ColibriGlm52Container } from "./glm52-container";
+import type { MemoryPlan } from "../admission-plan";
 
 export const GLM52_G5_GIB = 1024 ** 3;
 export const GLM52_G5_DEFAULT_MACHINE_BYTES = 32 * GLM52_G5_GIB;
@@ -73,8 +74,9 @@ export interface Glm52MemoryLineItems {
   readonly safetyMarginBytes: number;
 }
 
-export interface Glm52MemoryPlan {
+export interface Glm52MemoryPlan extends MemoryPlan {
   readonly schemaVersion: 1;
+  readonly strategy: "glm52-colibri";
   readonly preset: "g5-32gb-quality";
   readonly machineBytes: number;
   readonly processLimitBytes: number;
@@ -374,11 +376,14 @@ export function planGlm52Memory(
 
   return {
     schemaVersion: 1,
+    strategy: "glm52-colibri",
+    fits: true,
     preset: "g5-32gb-quality",
     machineBytes,
     processLimitBytes,
     osReserveBytes,
     contextTokens,
+    maxSafeContext: contextTokens,
     maxGenerationTokens,
     batchSize,
     enableMtp,
@@ -392,6 +397,28 @@ export function planGlm52Memory(
     mtpResidentSlots,
     mtpTotalSlots,
     lineItems,
+    weightsBytes: sum("GLM planned weights", [
+      lineItems.residentWeightsBytes,
+      lineItems.mainExpertSlabBytes,
+      lineItems.mtpExpertSlabBytes,
+    ]),
+    kvBytes: sum("GLM planned KV", [
+      lineItems.targetKvBytes,
+      lineItems.mtpKvBytes,
+    ]),
+    transientBytes: sum("GLM planned transient", [
+      lineItems.reconstructedKvTransientBytes,
+      lineItems.verifyBatchTransientBytes,
+    ]),
+    reserveBytes: sum("GLM planned reserves", [
+      lineItems.allocatorReserveBytes,
+      lineItems.bunNativeReserveBytes,
+      lineItems.safetyMarginBytes,
+    ]),
+    totalBytes: plannedProcessBytes,
+    usableBytes: processLimitBytes,
+    predictedDecodeTps: null,
+    allocatorLimitBytes: lineItems.allocatorReserveBytes,
     runtimeReserveBytes,
     plannedProcessBytes,
     processHeadroomBytes: processLimitBytes - plannedProcessBytes,

package/src/model/implementation.ts

@@ -0,0 +1,69 @@
+import type { ModelConfig } from "../config";
+import {
+  assertResolvedModelProfile,
+  type GenerationLoop,
+  type ModelGraph,
+  type ModelLoader,
+  type ResolvedModelProfile,
+} from "./profile";
+
+/** Engine code implements this port. Source and Model remain backend-specific;
+ * selection introduces no tensor conversion or dispatch in the token loop.
+ * create borrows source; a failed construction must release its own resources. */
+export interface ModelImplementation<Source, Model> {
+  readonly id: string;
+  readonly graph: ModelGraph;
+  readonly loader: ModelLoader;
+  readonly loop: GenerationLoop;
+  create(source: Source, config: ModelConfig, profile: ResolvedModelProfile): Model;
+}
+
+export interface ModelImplementationProvider<Source, Model> {
+  /** Validate identity and composition, then return one compatible binding.
+   * A declared implementation must never silently fall back to another. */
+  select(config: ModelConfig, profile: ResolvedModelProfile): ModelImplementation<Source, Model>;
+}
+
+/** Immutable composition, with identity matching delegated to model profiles.
+ * Exact declarations are requirements: absent/incompatible code is an error,
+ * never an invitation to retry with a different numerical implementation. */
+export class ModelImplementationRegistry<Source, Model> implements ModelImplementationProvider<Source, Model> {
+  private readonly entries: ReadonlyMap<string, ModelImplementation<Source, Model>>;
+
+  constructor(implementations: readonly ModelImplementation<Source, Model>[]) {
+    const entries = new Map<string, ModelImplementation<Source, Model>>();
+    for (const implementation of implementations) {
+      if (!implementation.id.trim()) throw new Error("model implementation id must not be empty");
+      if (entries.has(implementation.id))
+        throw new Error(`duplicate model implementation ${implementation.id}`);
+      entries.set(implementation.id, Object.freeze({
+        id: implementation.id,
+        graph: implementation.graph,
+        loader: implementation.loader,
+        loop: implementation.loop,
+        create: implementation.create.bind(implementation),
+      }));
+    }
+    this.entries = entries;
+  }
+
+  /** Return a new composition; active models retain their existing binding. */
+  with(...implementations: readonly ModelImplementation<Source, Model>[]): ModelImplementationRegistry<Source, Model> {
+    return new ModelImplementationRegistry([...this.entries.values(), ...implementations]);
+  }
+
+  select(config: ModelConfig, profile: ResolvedModelProfile): ModelImplementation<Source, Model> {
+    assertResolvedModelProfile(config, profile);
+    const execution = profile.profile.execution;
+    // These default IDs preserve the existing dedicated/generated/family choice.
+    const id = execution.implementation ?? (execution.specialization === "generated"
+      ? `${execution.graph}-generated` : execution.graph);
+    const implementation = this.entries.get(id);
+    if (!implementation)
+      throw new Error(`model profile ${profile.profile.id} requires unavailable implementation ${id}; refusing to fall back`);
+    if (implementation.graph !== execution.graph || implementation.loader !== execution.loader ||
+        implementation.loop !== execution.loop)
+      throw new Error(`model implementation ${id} is incompatible with profile ${profile.profile.id}; refusing to fall back`);
+    return implementation;
+  }
+}

package/src/model/minicpm5.ts

@@ -24,12 +24,13 @@ import {
   type Mask,
 } from "./gemma4-base";
 import { CompiledFunction } from "../mlx/compile";
+import { runtimeValue } from "../runtime-config";
 
 /** Verbatim port of mlx_lm/models/activations.py:
  *    `@partial(mx.compile, shapeless=True) def swiglu(gate, x): return nn.silu(gate) * x`
  *  mx.compile fuses sigmoid + mul + mul into ONE kernel (mlx-lm's decode
  *  `CV2ISigmoid…Multiply`) instead of our three separate dispatches. Traced
- *  once, replayed thereafter. Used on the decode (M=1) path only. */
+ *  once, replayed thereafter for decode and prefill. */
 let _swigluClosure: CompiledFunction | null = null;
 export function compiledSwiglu(gate: MlxArray, up: MlxArray): MlxArray {
   if (!_swigluClosure) {
@@ -199,15 +200,13 @@ export class LlamaMLP {
 
   /** MLP: returns mlp(x); the caller adds the residual when given. */
   forward(x: MlxArray, residual?: MlxArray): MlxArray {
-    const H = x.shape[x.shape.length - 1]!;
-    const M = x.shape.reduce((a, b) => a * b, 1) / H;
     const gate = this.gate.forward(x);
     const up = this.up.forward(x);
-    // Decode (M=1): fuse silu·up into one kernel (mlx_lm activations.py swiglu).
-    // Plain sigmoid+mul+mul otherwise (M>1 prefill/training, or inside a compile
-    // trace) — identical math either way. MLX_BUN_COMPILED_SWIGLU=0 disables.
+    // Match mlx_lm activations.py at every sequence length. Separate kernels
+    // differ numerically on MLX 0.32.2 at larger prefill shapes. An enclosing
+    // compiled trace owns fusion itself; the explicit opt-out stays available.
     let hidden: MlxArray;
-    if (M === 1 && !isCompiledTrace() && process.env.MLX_BUN_COMPILED_SWIGLU !== "0") {
+    if (!isCompiledTrace() && runtimeValue("MLX_BUN_COMPILED_SWIGLU") !== "0") {
       hidden = compiledSwiglu(gate, up);
       gate.dispose();
       up.dispose();
@@ -320,7 +319,7 @@ export class MiniCPM5Model {
    *  NEVER disposed — the caller owns it (for segmented backward `h` is a leaf
    *  whose lifetime is managed by the autograd closure); intra-range
    *  intermediates ARE disposed. Building block for `forwardLayers` and the
-   *  segmented-backward training path (docs/design/segmented-backward-training.md).
+   *  segmented-backward training path (docs/design/orpo-training.md).
    *  MiniCPM5 is plain full-attention with no KV-sharing, so each layer reads
    *  its own stateless `cache[i]` and a single causal mask. */
   runLayerRange(h: MlxArray, aIdx: number, bIdx: number, cache: Cache[]): MlxArray {