mlx-bun 0.4.0 → 1.0.0

package/src/model/qwen3.ts

@@ -1,277 +0,0 @@
-// Concrete model graph for plain Qwen3 (Qwen3ForCausalLM).
-// Port target: mlx_lm.models.qwen3. A standard dense decoder — GQA softmax
-// attention with per-head q/k RMSNorm (applied over head_dim BEFORE RoPE),
-// full-head RoPE (theta 1e6), SwiGLU MLP, RMSNorm, tied embeddings.
-//
-// Primary use here is as a TEXT-EMBEDDING backbone (mlx-community Qwen3-Embedding-*):
-// `embedPooled` returns the last-token hidden, L2-normalized — the pooling
-// Qwen3-Embedding was trained with (an <|endoftext|> token is appended by the
-// caller; its hidden is the sentence vector). The full causal-LM surface
-// (forward/generate via the tied head) is wired too so it conforms to RuntimeModel.
-
-import type { ModelConfig } from "../config";
-import type { Weights } from "../weights";
-import { MlxArray } from "../mlx/array";
-import * as ops from "../mlx/ops";
-import {
-  argmaxLastPosition,
-  disposing,
-  KVCache,
-  LoraState,
-  QuantizedEmbedding,
-  QuantizedLinear,
-  RMSNorm,
-  type Cache,
-  type Mask,
-} from "./gemma4-base";
-import { compiledSwiglu } from "./minicpm5";
-
-export class Qwen3Attention {
-  readonly qProj: QuantizedLinear;
-  readonly kProj: QuantizedLinear;
-  readonly vProj: QuantizedLinear;
-  readonly oProj: QuantizedLinear;
-  readonly qNorm: RMSNorm;
-  readonly kNorm: RMSNorm;
-  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 ?? 1000000;
-    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);
-    this.qNorm = new RMSNorm(weights.tensor(`${prefix}.q_norm.weight`), t.rmsNormEps);
-    this.kNorm = new RMSNorm(weights.tensor(`${prefix}.k_norm.weight`), t.rmsNormEps);
-  }
-
-  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);
-
-    // Reshape to heads, then q/k RMSNorm over head_dim BEFORE transpose+RoPE
-    // (mlx-lm qwen3 reference order).
-    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, this.qNorm.forward(q));
-    k = disposing(k, this.kNorm.forward(k));
-
-    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]));
-
-    q = disposing(q, ops.rope(q, this.headDim, this.ropeBase, cache.offset, null));
-    k = disposing(k, ops.rope(k, this.headDim, this.ropeBase, cache.offset, null));
-
-    const [keys, values] = cache.updateAndFetch(k, v);
-    k.dispose();
-    v.dispose();
-    const 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;
-  }
-}
-
-export class Qwen3MLP {
-  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);
-    // Oracle (mlx_lm/models/qwen3.py MLP): down_proj(swiglu(gate_proj(x), up_proj(x)))
-    // where swiglu is the `@mx.compile` closure (silu(gate)*up). Compiled = the
-    // same libmlx kernel mlx-lm dispatches (L1-faithful) AND fuses the three
-    // standalone ops into one dispatch. Bit-identical to the old unfused path.
-    const hidden = compiledSwiglu(gate, up);
-    gate.dispose();
-    up.dispose();
-    const out = this.down.forward(hidden);
-    hidden.dispose();
-    return out;
-  }
-}
-
-export class Qwen3Layer {
-  readonly attn: Qwen3Attention;
-  readonly mlp: Qwen3MLP;
-  readonly inputNorm: RMSNorm;
-  readonly postAttnNorm: RMSNorm;
-
-  constructor(weights: Weights, config: ModelConfig, prefix: string) {
-    this.attn = new Qwen3Attention(weights, config, `${prefix}.self_attn`);
-    this.mlp = new Qwen3MLP(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 Qwen3Model {
-  readonly config: ModelConfig;
-  readonly weightsBytes: number;
-  readonly prefixBase = "model";
-  readonly loraState = new LoraState();
-  readonly embed: QuantizedEmbedding;
-  readonly layers: Qwen3Layer[];
-  readonly finalNorm: RMSNorm;
-  /** Separate lm_head when NOT tied; otherwise the tied embedding's asLinear. */
-  readonly lmHead: QuantizedLinear | null;
-
-  constructor(weights: Weights, config: ModelConfig) {
-    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 Qwen3Layer(weights, config, `model.layers.${i}`),
-    );
-    this.finalNorm = new RMSNorm(weights.tensor("model.norm.weight"), config.text.rmsNormEps);
-    this.lmHead = config.text.tieWordEmbeddings
-      ? null
-      : 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("qwen3 input-embedding path is not supported");
-  }
-
-  protected forwardLayers(h0: MlxArray, cache: Cache[]): MlxArray {
-    const L = h0.shape[1]!;
-    const mask = cache[0]!.makeMask(L, null);
-    let cur = h0;
-    for (let i = 0; i < this.layers.length; i++) {
-      const next = this.layers[i]!.forward(cur, mask, cache[i]!);
-      cur.dispose();
-      cur = next;
-    }
-    mask.arr?.dispose();
-    return disposing(cur, this.finalNorm.forward(cur));
-  }
-
-  /** Last-token pooled, L2-normalized sentence embedding for one sequence.
-   *  ids [1, L] → vector [1, hidden]. The caller appends the EOS/eod token
-   *  (Qwen3-Embedding's pooling token) to `ids`. */
-  embedPooled(ids: MlxArray): MlxArray {
-    const cache = this.makeCache();
-    try {
-      const h = this.forwardHidden(ids, cache); // [1, L, hidden]
-      const [B, L, H] = h.shape as [number, number, number];
-      const last = h.slice([0, L - 1, 0], [B, L, H]); // [1, 1, hidden]
-      h.dispose();
-      const vec = disposing(last, ops.reshape(last, [B, H])); // [1, hidden]
-      // L2 normalize: x / max(||x||_2, eps)
-      const sq = ops.square(vec);
-      const ss = disposing(sq, ops.sumAxis(sq, -1, true));
-      const norm = disposing(ss, ops.sqrt(ss));
-      const out = ops.div(vec, norm);
-      vec.dispose();
-      norm.dispose();
-      return out;
-    } finally {
-      for (const c of cache) c.dispose();
-    }
-  }
-
-  logitsFromHidden(h: MlxArray): MlxArray {
-    return this.lmHead ? this.lmHead.forward(h) : this.embed.asLinear(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_5-checkpoint.ts

@@ -1,164 +0,0 @@
-// 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;
-    },
-  };
-}