mlx-bun 0.2.0 → 0.3.0

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

@@ -30,7 +30,7 @@ import { runtimeValue } from "../runtime-config";
  *    `@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) {
@@ -200,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() && runtimeValue("MLX_BUN_COMPILED_SWIGLU") !== "0") {
+    if (!isCompiledTrace() && runtimeValue("MLX_BUN_COMPILED_SWIGLU") !== "0") {
       hidden = compiledSwiglu(gate, up);
       gate.dispose();
       up.dispose();
@@ -321,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 {

package/src/model/profile.ts

@@ -58,6 +58,9 @@ export interface ModelExecutionComposition {
   readonly graph: ModelGraph;
   readonly loop: GenerationLoop;
   readonly specialization: ModelSpecialization;
+  /** Engine-owned code registration. Model files never supply executable code.
+   * Omit to retain the existing graph implementation. */
+  readonly implementation?: string;
 }
 
 /** A profile declares construction only. Request methods such as MTP, KV
@@ -323,6 +326,10 @@ function executionCapabilities(profile: ModelProfile): EngineCapability[] {
 
 function validateProfile(profile: ModelProfile): void {
   if (!profile.id) throw new Error("model profile id must not be empty");
+  if (profile.execution.implementation !== undefined &&
+      (!profile.execution.implementation.trim() || !profile.artifactFingerprint ||
+       profile.execution.specialization !== "artifact"))
+    throw new Error(`model profile ${profile.id} must bind a named implementation to an exact artifact`);
   if ((profile.artifactFingerprint === undefined) !== (profile.configFingerprint === undefined))
     throw new Error(
       `model profile ${profile.id} must declare artifactFingerprint and configFingerprint together`,
@@ -413,6 +420,7 @@ export function assertResolvedModelProfile(
   config: ModelConfig,
   resolved: ResolvedModelProfile,
 ): void {
+  validateProfile(resolved.profile);
   const fingerprint = configFingerprint(config);
   if (resolved.artifact.configFingerprint !== fingerprint)
     throw new Error(
@@ -424,8 +432,10 @@ export function assertResolvedModelProfile(
       `model profile ${resolved.profile.id} selects ${resolved.profile.execution.graph}, ` +
       `which is incompatible with ${config.modelType}`,
     );
-  if (resolved.exactArtifact &&
-      resolved.profile.artifactFingerprint !== resolved.artifact.fingerprint)
+  if (resolved.exactArtifact !== (resolved.profile.artifactFingerprint !== undefined) ||
+      (resolved.exactArtifact &&
+       (resolved.profile.artifactFingerprint !== resolved.artifact.fingerprint ||
+        resolved.profile.configFingerprint !== fingerprint)))
     throw new Error(
       `model profile ${resolved.profile.id} does not match artifact ` +
       `${resolved.artifact.fingerprint ?? "<unidentified>"}`,

package/src/model/qwen-append.ts

@@ -0,0 +1,12 @@
+// MLX 0.32.2, mlx/backend/metal/scaled_dot_product_attention.cpp (MIT):
+// https://github.com/ml-explore/mlx/blob/v0.32.2/mlx/backend/metal/scaled_dot_product_attention.cpp
+// For the qualified Qwen GQA-6/D256 graph on applegpu_g16s, one-pass/vector
+// selection and two-pass block counts change after these inclusive KV lengths.
+// A committed span must stay within one arithmetic regime to match M=1.
+const ATTENTION_BOUNDARIES = [1023, 1024, 8192, 32768, 65536];
+
+export function qwenAppendChunkSize(cacheTokens: number): number {
+  for (const end of ATTENTION_BOUNDARIES)
+    if (cacheTokens < end) return Math.min(4, end - cacheTokens);
+  return 4;
+}

package/src/model/qwen-conv.ts

@@ -0,0 +1,52 @@
+// Experimental Qwen GDN convolution. Its two outputs are evaluated together:
+// consuming the activation also materializes the small, independently owned
+// state tail, without keeping a concatenated prefill buffer alive for it.
+// Arithmetic follows MLX 0.31.2 depthwise_conv_1d and unary sigmoid (MIT,
+// Copyright Apple Inc.), preserving the bf16 intermediate rounding.
+import type { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import { MetalKernel } from "../mlx/metal-kernel";
+
+/** Borrow all inputs; return owned activation and an independent state tail. */
+export type QwenConvolution = (x: MlxArray, state: MlxArray, weight: MlxArray) => [MlxArray, MlxArray];
+
+let kernel: MetalKernel | null = null;
+export const fusedQwenConvolution: QwenConvolution = (x, state, weight) => {
+  const [B, M, D] = x.shape as [number, number, number];
+  const K = weight.shape[1]!;
+  if (x.shape.length !== 3 || weight.shape.length !== 3 || weight.shape[0] !== D ||
+      weight.shape[2] !== 1 || K !== 4 || state.shape.join() !== [B, K - 1, D].join() ||
+      ![Dtype.bfloat16, Dtype.float32].includes(x.dtype) || state.dtype !== x.dtype || weight.dtype !== x.dtype)
+    throw new Error("fused Qwen convolution requires matching bf16/f32 inputs, [B,M,D], [B,3,D], [D,4,1]");
+  kernel ??= new MetalKernel({ name: "qwen_causal_conv_silu_tail",
+    inputNames: ["state", "x", "w"], outputNames: ["out", "tail"],
+    source: String.raw`
+      const uint c = thread_position_in_grid.x;
+      const uint t = thread_position_in_grid.y;
+      const uint b = thread_position_in_grid.z;
+      if (c >= D) return;
+      if (t < M) {
+        float acc = 0.0f;
+        for (uint i = 0; i < K; ++i) {
+          const uint row = t + i;
+          const T value = row < K-1 ? state[(b*(K-1)+row)*D+c] : x[(b*M+row-(K-1))*D+c];
+          acc += float(value) * float(w[c*K+i]);
+        }
+        const T conv = T(acc);
+        // Preserve the pinned MLX sigmoid's typed operations and bf16
+        // intermediate rounding. Promoting this expression to f32 differs.
+        const auto y = 1 / (1 + metal::exp(metal::abs(conv)));
+        const T sig = conv < 0 ? y : 1-y;
+        out[(b*M+t)*D+c] = conv * sig;
+      }
+      if (t < K-1) {
+        const uint row = M + t;
+        tail[(b*(K-1)+t)*D+c] = row < K-1 ? state[(b*(K-1)+row)*D+c] : x[(b*M+row-(K-1))*D+c];
+      }
+    ` });
+  return kernel.apply([state, x, weight], {
+    outputs: [{ shape: [B, M, D], dtype: x.dtype }, { shape: [B, K - 1, D], dtype: x.dtype }],
+    grid: [D, Math.max(M, K - 1), B], threadGroup: [256, 1, 1],
+    templateDtypes: { T: x.dtype }, templateInts: { D, M, K },
+  }) as [MlxArray, MlxArray];
+};

package/src/model/qwen3_5-checkpoint.ts

@@ -0,0 +1,164 @@
+// Qwen3.5 checkpoint-generation normalization — the load-seam port of
+// mlx_lm.models.qwen3_5's `Model.sanitize` + `TextModel.sanitize`.
+//
+// Two checkpoint generations ship the same graph under different conventions:
+//
+//   pre-5.8 (mlx-lm-converted)    5.8-family (transformers 5.8.0 export)
+//   ──────────────────────────    ───────────────────────────────────────
+//   language_model.model.*        model.language_model.*
+//   language_model.lm_head.*      lm_head.*                (top level)
+//   vision_tower.*                model.visual.*
+//   (no mtp tensors in-repo)      mtp.*                    (in-repo)
+//   conv1d.weight [C, K, 1]       conv1d.weight [C, 1, K]  (HF layout)
+//   RMSNorm gains stored as γ     RMSNorm gains stored as γ−1
+//
+// Our graph speaks the FIRST column (`PREFIX = "language_model"` in
+// qwen3_5.ts) — which is exactly mlx-lm's post-sanitize name space — so this
+// module maps the artifact onto it instead of teaching every consumer two
+// namings. The rules below are a line-for-line copy of the oracle; see
+// mlx_lm/models/qwen3_5.py `Model.sanitize` (naming/drop) and
+// `TextModel.sanitize` (γ shift, conv layout).
+//
+// The γ shift is the dangerous one: a spurious +1.0 silently corrupts an old
+// artifact. The oracle's discriminator is NOT the naming — it is
+//   should_shift = any("mtp." in k) or any conv1d.weight with shape[-1] != 1
+// evaluated over the whole (post-rename, pre-drop) weight map. Reproduced
+// verbatim. Measured on the two artifacts in play (2026-09-01, M1 Max):
+//   • ~/.cache/…/models--mjriii--Qwen3.8-27B/snapshots/staged — no mtp.*,
+//     conv1d [10240,4,1] → no shift; layer-0 input_layernorm ≡ 1.0.
+//   • /Volumes/MLX-Models/models/mjriii/Qwen3.8-27B-{tqalloc-norot,rtn4-g64}
+//     — 31 mtp.* tensors, conv1d [10240,1,4] → shift; layer-0
+//     input_layernorm mean −0.0334 (= γ−1), byte-identical to the raw
+//     Qwen/Qwen3.8-27B 5.8 checkpoint (the quantizer copies norms verbatim).
+
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import type { Weights, WeightsView } from "../weights";
+
+/** Header-level facts about one stored tensor (no bytes are read). */
+export interface CheckpointTensor {
+  readonly name: string;
+  readonly shape: readonly number[];
+}
+
+export interface Qwen35Sanitized {
+  /** canonical (graph) name → the name the artifact actually stores. */
+  readonly names: Map<string, string>;
+  /** Canonical names stored as γ−1: a +1.0 is due at load. */
+  readonly normShift: Set<string>;
+  /** Canonical names in HF conv layout: a moveaxis(2, 1) is due at load. */
+  readonly convMoveaxis: Set<string>;
+  /** True when the artifact already IS our canonical layout — no view needed
+   *  and the pre-existing load path runs untouched. */
+  readonly identity: boolean;
+}
+
+/** mlx-lm `TextModel.sanitize` norm_keys, verbatim and in order. */
+const NORM_KEYS = [
+  ".input_layernorm.weight",
+  ".post_attention_layernorm.weight",
+  "model.norm.weight",
+  ".q_norm.weight",
+  ".k_norm.weight",
+] as const;
+
+/** mlx-lm `Model.sanitize`: drop vision, fold every naming onto
+ *  `language_model.…`. Returns null for a dropped tensor. */
+function canonicalName(key: string): string | null {
+  if (key.startsWith("vision_tower") || key.startsWith("model.visual")) return null;
+  if (key.startsWith("model.language_model"))
+    return "language_model.model" + key.slice("model.language_model".length);
+  if (key.startsWith("language_model.")) return key;
+  return "language_model." + key;
+}
+
+/**
+ * Port of mlx_lm.models.qwen3_5 `Model.sanitize` ∘ `TextModel.sanitize`,
+ * name/shape half. Pure: takes header facts, returns the load plan.
+ */
+export function sanitizeQwen35Checkpoint(
+  tensors: Iterable<CheckpointTensor>,
+): Qwen35Sanitized {
+  // Pass 1 — rename (vision dropped), exactly as Model.sanitize.
+  const renamed: { canonical: string; source: string; shape: readonly number[] }[] = [];
+  let sawForeignName = false;
+  for (const t of tensors) {
+    const canonical = canonicalName(t.name);
+    if (canonical === null) {
+      sawForeignName ||= t.name.startsWith("model.visual");
+      continue;
+    }
+    if (canonical !== t.name) sawForeignName = true;
+    renamed.push({ canonical, source: t.name, shape: t.shape });
+  }
+
+  // Pass 2 — TextModel.sanitize's two flags, over the renamed map BEFORE the
+  // mtp drop (the oracle computes them first, then filters).
+  const hasMtpWeights = renamed.some((t) => t.canonical.includes("mtp."));
+  const hasUnsanitizedConv1d = renamed.some(
+    (t) => t.canonical.includes("conv1d.weight") && t.shape[t.shape.length - 1] !== 1,
+  );
+  const shouldShiftNormWeights = hasMtpWeights || hasUnsanitizedConv1d;
+
+  const names = new Map<string, string>();
+  const normShift = new Set<string>();
+  const convMoveaxis = new Set<string>();
+  for (const t of renamed) {
+    if (t.canonical.includes("mtp.")) continue; // oracle drops these outright
+    names.set(t.canonical, t.source);
+    if (t.canonical.includes("conv1d.weight") && t.shape[t.shape.length - 1] !== 1)
+      convMoveaxis.add(t.canonical);
+    if (
+      shouldShiftNormWeights &&
+      NORM_KEYS.some((sfx) => t.canonical.endsWith(sfx)) &&
+      t.shape.length === 1
+    )
+      normShift.add(t.canonical);
+  }
+
+  return {
+    names,
+    normShift,
+    convMoveaxis,
+    identity: !sawForeignName && !hasMtpWeights && !shouldShiftNormWeights &&
+      convMoveaxis.size === 0,
+  };
+}
+
+/**
+ * The {@link WeightsView} for a qwen3.5 artifact, or null when the artifact is
+ * already canonical (then nothing is installed and the load path is bit-for-bit
+ * what it was before this seam existed).
+ */
+export function qwen35WeightsView(weights: Weights): WeightsView | null {
+  const plan = sanitizeQwen35Checkpoint(
+    weights.tensorNames.map((name) => ({ name, shape: weights.info(name).shape })),
+  );
+  if (plan.identity) return null;
+  return {
+    names: plan.names,
+    fixup(canonical: string, arr: MlxArray): MlxArray | null {
+      // Oracle order: moveaxis first, then the γ shift (disjoint sets — a
+      // conv1d weight is 3-D and never a norm gain).
+      if (plan.convMoveaxis.has(canonical)) {
+        // `v.moveaxis(2, 1)` on [C, 1, K] → [C, K, 1]. Materialized: the
+        // conv1d kernel consumes it every forward, and the transposed view
+        // would otherwise pin the source shard buffer.
+        const moved = ops.transposeAxes(arr, [0, 2, 1]);
+        const out = ops.contiguous(moved);
+        moved.dispose();
+        return out;
+      }
+      if (plan.normShift.has(canonical)) {
+        // `v + 1.0` — a weak python scalar, so the add happens at the gain's
+        // own dtype (bf16 here); scalarLike builds the constant host-side at
+        // that dtype for the same single-kernel add the reference dispatches.
+        const one = ops.scalarLike(1, arr);
+        const out = ops.add(arr, one);
+        one.dispose();
+        return out;
+      }
+      return null;
+    },
+  };
+}