mlx-bun 0.0.4

package/src/model/generated/index.ts

@@ -0,0 +1,18 @@
+// Registry of generated per-architecture specializations (hand-kept;
+// the files themselves come from scripts/gen-model.ts). Dispatch is by
+// config fingerprint; anything unregistered runs the monolith.
+
+import type { Weights } from "../../weights";
+import type { ModelConfig } from "../../config";
+import { Gemma4Model } from "../gemma4";
+import * as g12b from "./gemma4-12b";
+import * as ge4b from "./gemma4-e4b";
+import * as g26b from "./gemma4-26b";
+
+type ModelCtor = new (weights: Weights, config: ModelConfig) => Gemma4Model;
+
+export const GENERATED = new Map<string, ModelCtor>([
+  [g12b.FINGERPRINT, g12b.GeneratedGemma4],
+  [ge4b.FINGERPRINT, ge4b.GeneratedGemma4],
+  [g26b.FINGERPRINT, g26b.GeneratedGemma4],
+]);

package/src/model/minicpm5.ts

@@ -0,0 +1,278 @@
+// Concrete model graph for MiniCPM5-1B-OptiQ-4bit.
+// Port target: mlx_lm.models.llama from the oracle venv.
+
+import type { ModelConfig } from "../config";
+import type { Weights } from "../weights";
+import { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import {
+  argmaxLastPosition,
+  disposeTriple,
+  disposing,
+  KVCache,
+  LoraState,
+  QuantizedKVCache,
+  QuantizedEmbedding,
+  QuantizedLinear,
+  quantizedSdpa,
+  RMSNorm,
+  type Cache,
+  type Mask,
+} from "./gemma4-base";
+import { flashAttention, getTrainingAttn, flashSupported } from "./flash-attention";
+
+class LlamaAttention {
+  readonly qProj: QuantizedLinear;
+  readonly kProj: QuantizedLinear;
+  readonly vProj: QuantizedLinear;
+  readonly oProj: QuantizedLinear;
+  readonly nHeads: number;
+  readonly nKvHeads: number;
+  readonly headDim: number;
+  readonly scale: number;
+  readonly ropeBase: number;
+
+  constructor(weights: Weights, config: ModelConfig, prefix: string) {
+    const t = config.text;
+    this.nHeads = t.numAttentionHeads;
+    this.nKvHeads = t.numKeyValueHeads;
+    this.headDim = t.headDim;
+    this.scale = Math.pow(this.headDim, -0.5);
+    this.ropeBase = t.ropeParameters.full_attention?.ropeTheta ?? 10000;
+    this.qProj = QuantizedLinear.load(weights, `${prefix}.q_proj`, config);
+    this.kProj = QuantizedLinear.load(weights, `${prefix}.k_proj`, config);
+    this.vProj = QuantizedLinear.load(weights, `${prefix}.v_proj`, config);
+    this.oProj = QuantizedLinear.load(weights, `${prefix}.o_proj`, config);
+  }
+
+  forward(x: MlxArray, mask: Mask, cache: Cache): MlxArray {
+    const [B, L] = x.shape as [number, number, number];
+    let q = this.qProj.forward(x);
+    let k = this.kProj.forward(x);
+    let v = this.vProj.forward(x);
+
+    q = disposing(q, ops.reshape(q, [B, L, this.nHeads, this.headDim]));
+    k = disposing(k, ops.reshape(k, [B, L, this.nKvHeads, this.headDim]));
+    v = disposing(v, ops.reshape(v, [B, L, this.nKvHeads, this.headDim]));
+
+    q = disposing(q, ops.transposeAxes(q, [0, 2, 1, 3]));
+    k = disposing(k, ops.transposeAxes(k, [0, 2, 1, 3]));
+    v = disposing(v, ops.transposeAxes(v, [0, 2, 1, 3]));
+
+    // RoPE: a number offset takes the static fast::rope; an array offset
+    // (per-row positions for left-padded batched decode) takes the dynamic
+    // variant — same kernel, offset read from the array. Captured once: rope
+    // runs before updateAndFetch, so K and Q share the pre-write offset.
+    const offsetArr = cache.ropeOffsetArr;
+    const ropeStep = (x: MlxArray): MlxArray =>
+      offsetArr
+        ? ops.ropeDynamic(x, this.headDim, this.ropeBase, offsetArr, null)
+        : ops.rope(x, this.headDim, this.ropeBase, cache.offset, null);
+    q = disposing(q, ropeStep(q));
+    k = disposing(k, ropeStep(k));
+    let attn: MlxArray;
+    if (cache instanceof QuantizedKVCache) {
+      const [keys, values] = cache.updateAndFetchQuantized(k, v);
+      k.dispose();
+      v.dispose();
+      attn = quantizedSdpa(q, keys, values, this.scale, mask, cache.groupSize, cache.bits);
+      disposeTriple(keys);
+      disposeTriple(values);
+    } else {
+      const [keys, values] = cache.updateAndFetch(k, v);
+      k.dispose();
+      v.dispose();
+      // Training: ops.sdpa's dK vjp is wrong, so route the backward through
+      // the validated flash kernel. MiniCPM5 is full-attention → window 0.
+      const ta = getTrainingAttn();
+      if (ta === "flash" && flashSupported(q) && (mask.mode === "causal" || mask.mode === "array")) {
+        attn = flashAttention(q, keys, values, this.scale, true, 0);
+      } else {
+        attn = ops.sdpa(q, keys, values, this.scale, mask.mode, mask.arr);
+      }
+      keys.dispose();
+      values.dispose();
+    }
+    q.dispose();
+    const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
+    attn.dispose();
+    const merged = ops.reshape(attnT, [B, L, -1]);
+    attnT.dispose();
+    const out = this.oProj.forward(merged);
+    merged.dispose();
+    return out;
+  }
+}
+
+class LlamaMLP {
+  readonly gate: QuantizedLinear;
+  readonly up: QuantizedLinear;
+  readonly down: QuantizedLinear;
+
+  constructor(weights: Weights, config: ModelConfig, prefix: string) {
+    this.gate = QuantizedLinear.load(weights, `${prefix}.gate_proj`, config);
+    this.up = QuantizedLinear.load(weights, `${prefix}.up_proj`, config);
+    this.down = QuantizedLinear.load(weights, `${prefix}.down_proj`, config);
+  }
+
+  forward(x: MlxArray): MlxArray {
+    const gate = this.gate.forward(x);
+    const up = this.up.forward(x);
+    const sig = ops.sigmoid(gate);
+    const silu = ops.mul(gate, sig);
+    gate.dispose();
+    sig.dispose();
+    const hidden = ops.mul(silu, up);
+    silu.dispose();
+    up.dispose();
+    const out = this.down.forward(hidden);
+    hidden.dispose();
+    return out;
+  }
+}
+
+class LlamaLayer {
+  readonly attn: LlamaAttention;
+  readonly mlp: LlamaMLP;
+  readonly inputNorm: RMSNorm;
+  readonly postAttnNorm: RMSNorm;
+
+  constructor(weights: Weights, config: ModelConfig, prefix: string) {
+    this.attn = new LlamaAttention(weights, config, `${prefix}.self_attn`);
+    this.mlp = new LlamaMLP(weights, config, `${prefix}.mlp`);
+    this.inputNorm = new RMSNorm(weights.tensor(`${prefix}.input_layernorm.weight`), config.text.rmsNormEps);
+    this.postAttnNorm = new RMSNorm(weights.tensor(`${prefix}.post_attention_layernorm.weight`), config.text.rmsNormEps);
+  }
+
+  forward(x: MlxArray, mask: Mask, cache: Cache): MlxArray {
+    const xn = this.inputNorm.forward(x);
+    const a = this.attn.forward(xn, mask, cache);
+    xn.dispose();
+    const h = ops.add(x, a);
+    a.dispose();
+    const hn = this.postAttnNorm.forward(h);
+    const m = this.mlp.forward(hn);
+    hn.dispose();
+    const out = ops.add(h, m);
+    h.dispose();
+    m.dispose();
+    return out;
+  }
+}
+
+export class MiniCPM5Model {
+  readonly config: ModelConfig;
+  readonly weightsBytes: number;
+  readonly prefixBase = "model";
+  readonly loraState = new LoraState();
+  readonly embed: QuantizedEmbedding;
+  readonly layers: LlamaLayer[];
+  readonly finalNorm: RMSNorm;
+  readonly lmHead: QuantizedLinear;
+
+  constructor(weights: Weights, config: ModelConfig) {
+    if (config.text.tieWordEmbeddings)
+      throw new Error("llama tied embeddings are not supported yet");
+    this.config = config;
+    this.weightsBytes = [...weights.shards.files.values()]
+      .reduce((a, f) => a + f.mmap.size, 0);
+    this.embed = QuantizedEmbedding.load(weights, "model.embed_tokens", config);
+    this.layers = Array.from(
+      { length: config.text.numHiddenLayers },
+      (_, i) => new LlamaLayer(weights, config, `model.layers.${i}`),
+    );
+    this.finalNorm = new RMSNorm(weights.tensor("model.norm.weight"), config.text.rmsNormEps);
+    this.lmHead = QuantizedLinear.load(weights, "lm_head", config);
+  }
+
+  loraTargets(): Map<string, QuantizedLinear> {
+    const out = new Map<string, QuantizedLinear>();
+    for (let i = 0; i < this.layers.length; i++) {
+      const l = this.layers[i]!;
+      const p = `model.layers.${i}`;
+      out.set(`${p}.self_attn.q_proj`, l.attn.qProj);
+      out.set(`${p}.self_attn.k_proj`, l.attn.kProj);
+      out.set(`${p}.self_attn.v_proj`, l.attn.vProj);
+      out.set(`${p}.self_attn.o_proj`, l.attn.oProj);
+      out.set(`${p}.mlp.gate_proj`, l.mlp.gate);
+      out.set(`${p}.mlp.up_proj`, l.mlp.up);
+      out.set(`${p}.mlp.down_proj`, l.mlp.down);
+    }
+    return out;
+  }
+
+  makeCache(): Cache[] {
+    return this.layers.map(() => new KVCache());
+  }
+
+  forwardHidden(ids: MlxArray, cache: Cache[]): MlxArray {
+    const h = this.embed.encode(ids);
+    return this.forwardLayers(h, cache);
+  }
+
+  forwardEmbeddings(_embeds: MlxArray, _cache: Cache[], _bidir: MlxArray | null): MlxArray {
+    throw new Error("llama vision/input-embedding path is not supported");
+  }
+
+  protected forwardLayers(h0: MlxArray, cache: Cache[]): MlxArray {
+    const h = this.runLayerRange(h0, 0, this.layers.length, cache);
+    h0.dispose(); // runLayerRange keeps its input alive; we own h0 here
+    return disposing(h, this.finalNorm.forward(h));
+  }
+
+  /** Run decoder layers `[aIdx, bIdx)` on hidden `h`, returning the residual
+   *  stream after the last layer in the range (NO finalNorm). The input `h` is
+   *  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).
+   *  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 {
+    if (aIdx >= bIdx) return h;
+    const L = h.shape[1]!;
+    const mask = cache[aIdx]!.makeMask(L, null);
+    let cur = h;
+    for (let i = aIdx; i < bIdx; i++) {
+      const next = this.layers[i]!.forward(cur, mask, cache[i]!);
+      if (i > aIdx) cur.dispose(); // keep the input boundary leaf; drop interiors
+      cur = next;
+    }
+    mask.arr?.dispose();
+    return cur;
+  }
+
+  logitsFromHidden(h: MlxArray): MlxArray {
+    return this.lmHead.forward(h);
+  }
+
+  forward(tokens: number[] | MlxArray, cache: Cache[]): MlxArray {
+    const ids = Array.isArray(tokens)
+      ? ops.fromInt32(tokens, [1, tokens.length])
+      : tokens;
+    const h = this.forwardHidden(ids, cache);
+    if (Array.isArray(tokens)) ids.dispose();
+    const logits = this.logitsFromHidden(h);
+    h.dispose();
+    return logits;
+  }
+
+  generate(promptTokens: number[], maxTokens: number, eosIds: number[] = []): number[] {
+    const cache = this.makeCache();
+    const out: number[] = [];
+    try {
+      let tokens = promptTokens;
+      for (let step = 0; step < maxTokens; step++) {
+        const logits = this.forward(tokens, cache);
+        const next = argmaxLastPosition(logits);
+        logits.dispose();
+        if (eosIds.includes(next)) break;
+        out.push(next);
+        tokens = [next];
+      }
+    } finally {
+      for (const c of cache) c.dispose();
+    }
+    return out;
+  }
+}

package/src/model/qwen3-delta.ts

@@ -0,0 +1,224 @@
+// Gated DeltaNet recurrence for Qwen3.5 linear-attention layers.
+//
+// Port target: mlx_lm.models.gated_delta (compute_g + the `gated_delta_step`
+// Metal kernel) and the recurrent-state cache. mlx-lm runs the GPU kernel by
+// default (use_kernel = not training), and its float accumulation / simd_sum
+// reduction order differ from the pure-ops fallback — so BIT-EXACT parity with
+// mlx-lm requires the SAME kernel, dispatched with the SAME grid/threadgroup.
+// We port the non-vectorized, non-masked variant (g.ndim == 3, mask is None),
+// which is the B=1 single-stream path (ssm_mask is None at batch 1).
+//
+// Numerics (must match the reference dtypes exactly — mlx infers the kernel's
+// pointer element types from the input arrays):
+//   q, k        bf16  [B, T, Hk, Dk]   (after inv_scale * rms_norm(., None))
+//   v           bf16  [B, T, Hv, Dv]
+//   g           f32   [B, T, Hv]        (= exp(-exp(A_log_f32) * softplus(a+dt_bias)))
+//   beta        bf16  [B, T, Hv]        (= sigmoid(b))
+//   state_in    f32   [B, Hv, Dv, Dk]
+//   y (out)     bf16  [B, T, Hv, Dv]    (InT)
+//   state_out   f32   [B, Hv, Dv, Dk]   (StT)
+// GQA is handled inside the kernel (hk_idx = hv_idx / (Hv/Hk)); q/k stay at Hk.
+
+import { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import { MetalKernel } from "../mlx/metal-kernel";
+import * as ops from "../mlx/ops";
+import type { Cache, Mask } from "./gemma4-base";
+
+// Verbatim body of mlx-lm's gated_delta_step (has_mask=False, vectorized=False),
+// with the only change being that the time count `T` arrives as a 1-element
+// int32 input (`Tin`) and is read once — numerically identical to the
+// reference (T only drives loop bounds and integer offsets), but avoids a
+// kernel recompile per distinct sequence length.
+const SOURCE = String.raw`
+    const int T = Tin[0];
+    auto n = thread_position_in_grid.z;
+    auto b_idx = n / Hv;
+    auto hv_idx = n % Hv;
+    auto hk_idx = hv_idx / (Hv / Hk);
+    constexpr int n_per_t = Dk / 32;
+
+    // q, k: [B, T, Hk, Dk]
+    auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk;
+    auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
+
+    // v, y: [B, T, Hv, Dv]
+    auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
+    y += b_idx * T * Hv * Dv + hv_idx * Dv;
+
+    auto dk_idx = thread_position_in_threadgroup.x;
+    auto dv_idx = thread_position_in_grid.y;
+
+    // state_in, state_out: [B, Hv, Dv, Dk]
+    auto i_state = state_in + (n * Dv + dv_idx) * Dk;
+    auto o_state = state_out + (n * Dv + dv_idx) * Dk;
+
+    float state[n_per_t];
+    for (int i = 0; i < n_per_t; ++i) {
+      auto s_idx = n_per_t * dk_idx + i;
+      state[i] = static_cast<float>(i_state[s_idx]);
+    }
+
+    // g: [B, T, Hv]
+    auto g_ = g + b_idx * T * Hv;
+    auto beta_ = beta + b_idx * T * Hv;
+
+    for (int t = 0; t < T; ++t) {
+      if (true) {
+        float kv_mem = 0.0f;
+        for (int i = 0; i < n_per_t; ++i) {
+          auto s_idx = n_per_t * dk_idx + i;
+          state[i] = state[i] * g_[hv_idx];
+          kv_mem += state[i] * k_[s_idx];
+        }
+        kv_mem = simd_sum(kv_mem);
+
+        auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];
+
+        float out = 0.0f;
+        for (int i = 0; i < n_per_t; ++i) {
+          auto s_idx = n_per_t * dk_idx + i;
+          state[i] = state[i] + k_[s_idx] * delta;
+          out += state[i] * q_[s_idx];
+        }
+        out = simd_sum(out);
+        if (thread_index_in_simdgroup == 0) {
+          y[dv_idx] = static_cast<InT>(out);
+        }
+      } else {
+        y[dv_idx] = static_cast<InT>(0);
+      }
+      // Increment data pointers to next time step
+      q_ += Hk * Dk;
+      k_ += Hk * Dk;
+      v_ += Hv * Dv;
+      y += Hv * Dv;
+      g_ += Hv;
+      beta_ += Hv;
+    }
+    for (int i = 0; i < n_per_t; ++i) {
+      auto s_idx = n_per_t * dk_idx + i;
+      o_state[s_idx] = static_cast<StT>(state[i]);
+    }
+`;
+
+let kernel: MetalKernel | null = null;
+function getKernel(): MetalKernel {
+  if (!kernel)
+    kernel = new MetalKernel({
+      name: "gated_delta_step",
+      inputNames: ["q", "k", "v", "g", "beta", "state_in", "Tin"],
+      outputNames: ["y", "state_out"],
+      source: SOURCE,
+      ensureRowContiguous: true,
+    });
+  return kernel;
+}
+
+/** compute_g: exp(-exp(A_log_f32) * softplus(a + dt_bias)). Output is f32
+ *  (the f32 exp(A_log) promotes the bf16 softplus term) — the kernel reads g
+ *  as float, so this dtype is load-bearing. */
+export function computeG(aLog: MlxArray, a: MlxArray, dtBias: MlxArray): MlxArray {
+  const aLogF = aLog.astype(Dtype.float32);
+  const expA = ops.exp(aLogF); // f32 [Hv]
+  aLogF.dispose();
+  const negExpA = ops.neg(expA); // f32 [Hv]
+  expA.dispose();
+  const adt = ops.add(a, dtBias); // bf16 [B,S,Hv]
+  const sp = ops.softplus(adt); // bf16 [B,S,Hv]
+  adt.dispose();
+  const prod = ops.mul(negExpA, sp); // f32 [B,S,Hv]
+  negExpA.dispose();
+  sp.dispose();
+  const g = ops.exp(prod); // f32 [B,S,Hv]
+  prod.dispose();
+  return g;
+}
+
+/** gated_delta_update (use_kernel path): returns [y, newState].
+ *   q, k: [B, S, Hk, Dk] bf16   v: [B, S, Hv, Dv] bf16
+ *   a, b: [B, S, Hv] bf16        aLog, dtBias: [Hv]
+ *   state: [B, Hv, Dv, Dk] f32 (or null → zeros) */
+export function gatedDeltaUpdate(
+  q: MlxArray, k: MlxArray, v: MlxArray, a: MlxArray, b: MlxArray,
+  aLog: MlxArray, dtBias: MlxArray, state: MlxArray | null,
+): [MlxArray, MlxArray] {
+  const [B, , Hk, Dk] = q.shape as [number, number, number, number];
+  const [, , Hv, Dv] = v.shape as [number, number, number, number];
+
+  const beta = ops.sigmoid(b); // bf16
+  const g = computeG(aLog, a, dtBias); // f32
+
+  let stateIn = state;
+  let ownState = false;
+  if (!stateIn) {
+    stateIn = ops.zeros([B, Hv, Dv, Dk], Dtype.float32);
+    ownState = true;
+  }
+
+  const T = q.shape[1]!;
+  const tArr = MlxArray.fromInt32(new Int32Array([T]), [1]);
+  const [y, stateOut] = getKernel().apply([q, k, v, g, beta, stateIn, tArr], {
+    outputs: [
+      { shape: [B, T, Hv, Dv], dtype: q.dtype },
+      { shape: [B, Hv, Dv, Dk], dtype: Dtype.float32 },
+    ],
+    grid: [32, Dv, B * Hv],
+    threadGroup: [32, 4, 1],
+    templateInts: { Dk, Dv, Hk, Hv },
+    templateDtypes: { InT: q.dtype, StT: Dtype.float32 },
+  });
+  beta.dispose();
+  g.dispose();
+  tArr.dispose();
+  if (ownState) stateIn.dispose();
+  return [y!, stateOut!];
+}
+
+/** Recurrent cache for a gated-DeltaNet layer — port of mlx-lm
+ *  cache.ArraysCache(size=2): slot 0 = causal-conv state [B, K-1, conv_dim],
+ *  slot 1 = recurrent state [B, Hv, Dv, Dk] f32. The linear-attn layer reads
+ *  and writes these slots directly; `advance(N)` only tracks token count for
+ *  B=1 single-stream (lengths / left_padding are batched-decode concerns).
+ *  Not a KVCache, so maybeQuantizeKv skips it. */
+export class SSMCache implements Cache {
+  conv: MlxArray | null = null;
+  recurrent: MlxArray | null = null;
+  offset = 0;
+
+  /** Linear-attn layers don't go through the KV update path. */
+  updateAndFetch(): [MlxArray, MlxArray] {
+    throw new Error("SSMCache has no KV updateAndFetch (gated-DeltaNet layer)");
+  }
+
+  makeMask(_N: number, _windowSize: number | null): Mask {
+    // B=1 single-stream: ssm_mask is None (no left-pad / variable lengths).
+    return { mode: "", arr: null };
+  }
+
+  advance(n: number): void {
+    this.offset += n;
+  }
+
+  state(): MlxArray[] {
+    const out: MlxArray[] = [];
+    if (this.conv) out.push(this.conv);
+    if (this.recurrent) out.push(this.recurrent);
+    return out;
+  }
+
+  isTrimmable(): boolean {
+    return false; // recurrent state can't be rolled back per token
+  }
+
+  trim(_n: number): void {
+    throw new Error("SSMCache is not trimmable");
+  }
+
+  dispose(): void {
+    this.conv?.dispose();
+    this.recurrent?.dispose();
+    this.conv = null;
+    this.recurrent = null;
+  }
+}