@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/adapters/manager.ts

@@ -0,0 +1,316 @@
+// Hot-swap mounted LoRA adapters — port of optiq/adapters/{mount,
+// registry,resolver}.py (serving side; lora/apply.py is the training-
+// side rank logic). Mount N adapters on one quantized base, select per
+// request by id, never reload the base.
+//
+// Deviations from the reference, both deliberate:
+// - No ContextVar/serve-pin: our generation queue is serialized, so the
+//   active adapter is a plain field (LoraState) set by generate().
+// - Residual composition is mlx-lm LoRALinear / optiq apply.py
+//   (`y + (scale·z).astype(x.dtype)`), not mount.py's uncast f32 add —
+//   the cast form is what the adapters were trained behind.
+
+import { MlxArray,cpuStream } from "@mlx-bun/mlx/array";
+import { C } from "@mlx-bun/mlx/ffi";
+import { transposeAxes } from "@mlx-bun/mlx/ops";
+import { ptr,read } from "bun:ffi";
+import { createHash } from "node:crypto";
+import { existsSync,statSync } from "node:fs";
+import { resolve } from "node:path";
+import { SafetensorsFile } from "../artifacts/safetensors";
+import type { LoraState,LoraWeights } from "../layers/lora";
+import { adapterWeightsFile,readAdapterScale } from "./available";
+
+/** The adapter manager needs named mount points, not a concrete model family. */
+export interface LoraTarget {
+  readonly inFeatures: number;
+  readonly outFeatures: number;
+  adapters?: Map<string, LoraWeights> | null;
+  loraState?: LoraState | null;
+}
+export interface AdapterTarget {
+  readonly prefixBase: string;
+  readonly loraState: LoraState;
+  loraTargets(): ReadonlyMap<string, LoraTarget>;
+}
+
+const cstr = (s: string) => Buffer.from(s + "\0", "utf8");
+
+export interface AdapterInfo {
+  id: string;
+  path: string;
+  rank: number | null;
+  scale: number;
+  sizeBytes: number;
+  mountedLayers: number;
+  /** Adapter tensors for modules outside the 7 target suffixes
+   *  (per-layer-input projections etc.) — skipped, like the reference. */
+  skippedTensors: number;
+  /** Actual resident bytes of this adapter's mounted lora_a/lora_b arrays
+   *  (sum of MlxArray.nbytes — the real RAM cost while mounted, not the
+   *  on-disk safetensors size). Web chat's adapter routing table shows this
+   *  per loaded adapter instead of guessing
+   *  from file size. */
+  ramBytes: number;
+}
+
+/** Read every tensor from an adapter safetensors file (small, f32) into
+ *  materialized mlx arrays. Tensor names come from our header parser;
+ *  arrays from mlx's native loader via the proven map-get pattern
+ *  (weights.ts). The map and header mmap are freed before returning. */
+export function loadAdapterTensors(file: string): Map<string, MlxArray> {
+  const sf = SafetensorsFile.open(file);
+  const names = [...sf.tensors.keys()];
+  sf.mmap.unmap();
+
+  const arrMapSlot = new BigUint64Array([C.mlx_map_string_to_array_new()]);
+  const metaMapSlot = new BigUint64Array([C.mlx_map_string_to_string_new()]);
+  const arrMapPtr = ptr(arrMapSlot);
+  const metaMapPtr = ptr(metaMapSlot);
+  const status = C.mlx_load_safetensors(arrMapPtr, metaMapPtr, ptr(cstr(file)), cpuStream);
+  C.mlx_map_string_to_string_free(read.u64(metaMapPtr, 0));
+  const mapHandle = read.u64(arrMapPtr, 0);
+  if (status !== 0) {
+    C.mlx_map_string_to_array_free(mapHandle);
+    throw new Error(`mlx_load_safetensors(${file}) failed`);
+  }
+  const out = new Map<string, MlxArray>();
+  try {
+    for (const name of names) {
+      const slot = new BigUint64Array([C.mlx_array_new()]);
+      const slotPtr = ptr(slot);
+      if (C.mlx_map_string_to_array_get(slotPtr, mapHandle, ptr(cstr(name))) !== 0)
+        throw new Error(`adapter tensor ${name} missing from native map`);
+      const arr = new MlxArray(read.u64(slotPtr, 0));
+      arr.eval(); // materialize — the map (and its Load refs) is freed below
+      out.set(name, arr);
+    }
+  } catch (e) {
+    for (const a of out.values()) a.dispose();
+    C.mlx_map_string_to_array_free(mapHandle);
+    throw e;
+  }
+  C.mlx_map_string_to_array_free(mapHandle);
+  return out;
+}
+
+/** Parse an adapter selection spec: "sft" | "sft+dpo" | "sft,dpo"
+ *  (reference mount.py stacking syntax — residuals sum in order). */
+export function parseAdapterSpec(spec: string): string[] {
+  const sep = spec.includes("+") ? "+" : ",";
+  return spec.split(sep).map((s) => s.trim()).filter(Boolean);
+}
+
+export class AdapterManager {
+  readonly #model: AdapterTarget;
+  readonly #mounted = new Map<string, AdapterInfo>();
+  /** Owned adapter arrays per id (disposed on unmount). */
+  readonly #arrays = new Map<string, MlxArray[]>();
+  readonly #revisions = new Map<string, string>();
+
+  constructor(model: AdapterTarget) {
+    this.#model = model;
+  }
+
+  /** Mount an adapter directory under `id`. Validates layer-name and
+   *  rank/shape compatibility against the base BEFORE any weights are
+   *  attached — a bad adapter fails here, never at request time
+   *  (reference: resolver.py + mount.py add_adapter checks). */
+  async mount(id: string, dir: string): Promise<AdapterInfo> {
+    if (this.#mounted.has(id)) return this.#mounted.get(id)!;
+    if (id.includes("+") || id.includes(","))
+      throw new Error(`adapter id ${JSON.stringify(id)} may not contain '+' or ','`);
+    const path = resolve(dir);
+    if (!existsSync(path)) throw new Error(`adapter dir not found: ${path}`);
+
+    const weightsFile = adapterWeightsFile(path);
+    const { scale, rank: configRank, rsLora } = await readAdapterScale(path);
+    // Hash once at mount, before native allocation. Reusing an adapter name
+    // for different weights or scale must never reuse its prefix/checkpoint.
+    const digest = createHash("sha256").update(JSON.stringify({ scale, rank: configRank, rsLora }));
+    for await (const chunk of Bun.file(weightsFile).stream()) digest.update(chunk);
+    const revision = digest.digest("hex");
+    const tensors = loadAdapterTensors(weightsFile);
+    const targets = this.#model.loraTargets();
+
+    // Group adapter tensors into (modulePath → {a, b}), probing both the
+    // pure-LLM and VLM-wrapped prefixes like the reference _find_weight_pair.
+    const pairs = new Map<string, {
+      a?: MlxArray;
+      b?: MlxArray;
+      aIsPeft?: boolean;
+      bIsPeft?: boolean;
+    }>();
+    let skipped = 0;
+    const ourPrefix = `${this.#model.prefixBase}.layers.`;
+    const altPrefix = ourPrefix.startsWith("language_model.")
+      ? ourPrefix.slice("language_model.".length)
+      : `language_model.${ourPrefix}`;
+    for (const [name, arr] of tensors) {
+      const m = name.match(/^(.*)\.(lora_a|lora_A|lora_b|lora_B)(\.weight)?$/);
+      if (!m) { skipped++; arr.dispose(); continue; }
+      let modulePath = m[1]!;
+      // Standard PEFT state dicts wrap the model path in
+      // `base_model.model.`. The remainder is the underlying HF module path
+      // (`model.layers...` or `language_model.model.layers...`).
+      if (modulePath.startsWith("base_model.model."))
+        modulePath = modulePath.slice("base_model.model.".length);
+      if (modulePath.startsWith(altPrefix))
+        modulePath = ourPrefix + modulePath.slice(altPrefix.length);
+      if (!targets.has(modulePath)) { skipped++; arr.dispose(); continue; }
+      const slot = pairs.get(modulePath) ?? {};
+      const peftName = m[3] === ".weight" || m[2] === "lora_A" || m[2] === "lora_B";
+      if (m[2]!.toLowerCase() === "lora_a") {
+        slot.a = arr;
+        slot.aIsPeft = peftName;
+      } else {
+        slot.b = arr;
+        slot.bIsPeft = peftName;
+      }
+      pairs.set(modulePath, slot);
+    }
+
+    // Validate every pair against the base linear's dims before mounting
+    // anything (all-or-nothing: a bad adapter must not half-mount).
+    const validated: { linear: LoraTarget; lw: LoraWeights }[] = [];
+    const dispose = () => {
+      for (const { a, b } of pairs.values()) { a?.dispose(); b?.dispose(); }
+    };
+    try {
+      for (const [modulePath, pair] of pairs) {
+        const { a, b } = pair;
+        if (!a || !b)
+          throw new Error(`${modulePath}: adapter has only one of lora_a/lora_b`);
+        const linear = targets.get(modulePath)!;
+        if (a.shape.length !== 2 || b.shape.length !== 2)
+          throw new Error(
+            `${modulePath}: LoRA tensors must be rank-2; got ` +
+            `lora_a [${a.shape}] / lora_b [${b.shape}]`,
+          );
+
+        const internalLayout =
+          a.shape[0] === linear.inFeatures &&
+          b.shape[0] === a.shape[1] &&
+          b.shape[1] === linear.outFeatures;
+        const peftLayout =
+          a.shape[1] === linear.inFeatures &&
+          b.shape[1] === a.shape[0] &&
+          b.shape[0] === linear.outFeatures;
+        if (!internalLayout && !peftLayout)
+          throw new Error(
+            `${modulePath}: shape mismatch — lora_a [${a.shape}] / lora_b [${b.shape}] ` +
+            `vs base [in ${linear.inFeatures}, out ${linear.outFeatures}]. Expected ` +
+            `mlx-lm [in, rank] + [rank, out] or PEFT [rank, in] + [out, rank].`,
+          );
+
+        // Full-rank square adapters satisfy both shape predicates. In that
+        // ambiguous case the PEFT spelling is authoritative; otherwise the
+        // unique compatible layout decides.
+        const usePeftLayout =
+          peftLayout && (!internalLayout || (pair.aIsPeft === true && pair.bIsPeft === true));
+        let mountedA = a;
+        let mountedB = b;
+        if (usePeftLayout) {
+          let transposedA: MlxArray | null = null;
+          let transposedB: MlxArray | null = null;
+          try {
+            transposedA = transposeAxes(a, [1, 0]);
+            transposedB = transposeAxes(b, [1, 0]);
+          } catch (e) {
+            transposedA?.dispose();
+            transposedB?.dispose();
+            throw e;
+          }
+          mountedA = transposedA;
+          mountedB = transposedB;
+          pair.a = mountedA;
+          pair.b = mountedB;
+          a.dispose();
+          b.dispose();
+        }
+        const rank = mountedA.shape[1]!;
+        // rsLoRA: effective per-layer scale is α/√rank (matches training).
+        const effScale = rsLora ? scale / Math.sqrt(rank) : scale;
+        validated.push({ linear, lw: { a: mountedA, b: mountedB, scale: effScale, rank } });
+      }
+      if (validated.length === 0)
+        throw new Error(
+          `failed to mount adapter ${JSON.stringify(id)}: no tensors match the ` +
+          `target modules (q/k/v/o/gate/up/down_proj). Check that the adapter ` +
+          `was trained for this base model.`,
+        );
+    } catch (e) {
+      dispose();
+      throw e;
+    }
+
+    for (const { linear, lw } of validated) {
+      (linear.adapters ??= new Map()).set(id, lw);
+      linear.loraState = this.#model.loraState;
+    }
+
+    const ramBytes = validated.reduce((sum, { lw }) => sum + lw.a.nbytes + lw.b.nbytes, 0);
+    const info: AdapterInfo = {
+      id,
+      path,
+      rank: configRank,
+      scale,
+      sizeBytes: statSync(weightsFile).size,
+      mountedLayers: validated.length,
+      skippedTensors: skipped,
+      ramBytes,
+    };
+    this.#mounted.set(id, info);
+    this.#revisions.set(id, revision);
+    this.#arrays.set(id, validated.flatMap(({ lw }) => [lw.a, lw.b]));
+    this.#model.loraState.refreshBindings();
+    return info;
+  }
+
+  /** Remove `id` from every linear (frees its arrays). The mount points
+   *  stay in place for other adapters, like the reference. */
+  unmount(id: string): number {
+    if (!this.#mounted.delete(id)) return 0;
+    this.#revisions.delete(id);
+    let removed = 0;
+    for (const linear of this.#model.loraTargets().values()) {
+      if (linear.adapters?.delete(id)) removed++;
+    }
+    for (const a of this.#arrays.get(id) ?? []) a.dispose();
+    this.#arrays.delete(id);
+    if (this.#model.loraState.active.includes(id))
+      this.#model.loraState.active = this.#model.loraState.active.filter((x) => x !== id);
+    return removed;
+  }
+
+  /** Ordered, content-addressed namespace for the actual mounted composition. */
+  cacheNamespace(ids: readonly string[]): string {
+    return ids.map((id) => {
+      const revision = this.#revisions.get(id);
+      if (!revision) throw new Error(`adapter ${id} is no longer mounted`);
+      return `${id}@${revision}`;
+    }).join("+");
+  }
+
+  list(): AdapterInfo[] {
+    return [...this.#mounted.values()];
+  }
+
+  get(id: string): AdapterInfo | undefined {
+    return this.#mounted.get(id);
+  }
+
+  /** Resolve a request's adapter spec to validated ids ([] for none).
+   *  Unknown ids are an error — selection must fail loudly, not no-op. */
+  resolveSpec(spec: string | null | undefined): string[] {
+    if (!spec || spec.toLowerCase() === "none") return [];
+    const ids = parseAdapterSpec(spec);
+    for (const id of ids)
+      if (!this.#mounted.has(id))
+        throw new Error(
+          `unknown adapter ${JSON.stringify(id)} — mounted: ` +
+          `${[...this.#mounted.keys()].join(", ") || "(none)"}`,
+        );
+    return ids;
+  }
+}

package/src/artifacts/auxiliary-files.ts

@@ -0,0 +1,45 @@
+// Auxiliary checkpoint files: the non-weight files (tokenizer, chat template,
+// generation config, …) a written model directory needs to stay loadable.
+// Checkpoint producers (quantization, adapter fusion) copy them through verbatim.
+
+import { existsSync } from "node:fs";
+import { join } from "node:path";
+
+/** Aux files copied through verbatim from the source dir when present. */
+const AUX_FILES = [
+  "tokenizer.json",
+  "tokenizer_config.json",
+  "tokenizer.model",
+  "spiece.model",
+  "chat_template.jinja",
+  "generation_config.json",
+  "special_tokens_map.json",
+  "added_tokens.json",
+  "vocab.json",
+  "merges.txt",
+  "README.md",
+  "kv_config.json",
+];
+
+/** Copy known aux files (and any `*.model`) from src → out when present. */
+export async function copyAuxFiles(srcDir: string, outDir: string): Promise<void> {
+  const seen = new Set<string>();
+  for (const f of AUX_FILES) {
+    const src = join(srcDir, f);
+    if (existsSync(src)) {
+      await Bun.write(join(outDir, f), Bun.file(src));
+      seen.add(f);
+    }
+  }
+  // Sweep any additional *.model tokenizer sidecars not in the known list.
+  try {
+    const { readdirSync } = await import("node:fs");
+    for (const f of readdirSync(srcDir)) {
+      if (f.endsWith(".model") && !seen.has(f)) {
+        await Bun.write(join(outDir, f), Bun.file(join(srcDir, f)));
+      }
+    }
+  } catch {
+    // src dir unreadable — aux copy is best-effort.
+  }
+}

package/src/artifacts/config-dialects.ts

@@ -0,0 +1,94 @@
+// How a `config.json` differs from the plain Hugging Face layout, by `model_type`.
+// A dialect is a reading rule for the artifact, not a statement about a model
+// family (`models/families.ts` owns those): `loadModelConfig` applies exactly one
+// and branches on its fields, never on a model type.
+
+/** How `rope_parameters` is spelled. `map`: a per-attention-type map, absent means none.
+ *  `theta`: as `map`, and a flat `rope_theta` stands in for a missing map.
+ *  `flat`: one flat dict (`rope_theta`, `partial_rotary_factor`) for full attention. */
+export type RopeLayout = "map" | "theta" | "flat";
+
+export interface ConfigDialect {
+  /** Text-config fields the config omits: the architecture's own defaults. Fields present still win. */
+  readonly textDefaults?: () => Record<string, any>;
+  /** Geometry and stop tokens come from a converted container's own parser. */
+  readonly container?: "colibri";
+  readonly rope: RopeLayout;
+  /** `layer_types` when the config declares none: every layer full attention. */
+  readonly fullAttentionByDefault?: true;
+  /** Refuse a value this loader has not verified rather than mis-read it. */
+  readonly validate?: (text: Record<string, any>) => void;
+}
+
+/** DiffusionGemma's config.json ships only token ids + canvas_length + the
+ *  quant map — the architecture dims live in optiq's `config.py` TextConfig
+ *  defaults (and layer_types / rope_parameters are computed in __post_init__).
+ *  We reproduce those defaults in snake_case so the generic parser below picks
+ *  them up; any field actually present in config.json still overrides. Source:
+ *  optiq/vlm/_mlxvlm/models/diffusion_gemma/config.py. */
+function diffusionGemmaRawDefaults(): Record<string, any> {
+  const numLayers = 30;
+  const pattern = ["sliding_attention", "sliding_attention", "sliding_attention",
+    "sliding_attention", "sliding_attention", "full_attention"];
+  const layer_types = Array.from({ length: numLayers }, (_, i) => pattern[i % pattern.length]);
+  layer_types[numLayers - 1] = "full_attention"; // last forced full
+  return {
+    hidden_size: 2816,
+    num_hidden_layers: numLayers,
+    num_attention_heads: 16,
+    num_key_value_heads: 8,
+    num_global_key_value_heads: 2,
+    head_dim: 256,
+    global_head_dim: 512,
+    intermediate_size: 2112,
+    moe_intermediate_size: 704,
+    hidden_activation: "gelu_pytorch_tanh",
+    rms_norm_eps: 1e-6,
+    vocab_size: 262144,
+    max_position_embeddings: 262144,
+    sliding_window: 1024,
+    layer_types,
+    enable_moe_block: true,
+    num_experts: 128,
+    top_k_experts: 8,
+    final_logit_softcapping: 30.0,
+    tie_word_embeddings: true,
+    bos_token_id: 2,
+    eos_token_id: 1,
+    rope_parameters: {
+      sliding_attention: { rope_type: "default", rope_theta: 10000.0 },
+      full_attention: { rope_type: "proportional", partial_rotary_factor: 0.25, rope_theta: 1000000.0 },
+    },
+  };
+}
+
+/** Qwen3.8 adds output_gate_type ("swish"), which every implementation
+ *  (transformers ground truth, mlx-lm pinned + main) currently ignores —
+ *  the attention output gate is hardcoded o_proj(out·σ(gate)). Only accept
+ *  values verified to mean that; an unknown value must fail at load, not
+ *  silently mis-gate a future checkpoint where the field starts mattering. */
+function verifyOutputGate(text: Record<string, any>): void {
+  if (text.output_gate_type != null && text.output_gate_type !== "swish")
+    throw new Error(
+      `qwen3_5: unverified output_gate_type "${text.output_gate_type}" (only "swish" ` +
+        `is verified to mean the standard sigmoid output gate)`,
+    );
+}
+
+const PLAIN: ConfigDialect = Object.freeze({ rope: "map" });
+const FLAT_THETA: ConfigDialect = Object.freeze({ rope: "theta" });
+
+/** The reading rule for a config's `model_type`. Anything not listed reads as the plain layout. */
+export function dialectFor(modelType: unknown): ConfigDialect {
+  if (modelType === "diffusion_gemma") return { rope: "map", textDefaults: diffusionGemmaRawDefaults };
+  if (modelType === "glm_moe_dsa") return { rope: "map", container: "colibri" };
+  // Qwen3.5 rope_parameters is a flat dict ({type, rope_theta, mrope_section,
+  // partial_rotary_factor}), not the gemma per-attention-type map — and
+  // type "default" means plain partial nn.RoPE (mrope_section ignored for text).
+  if (typeof modelType === "string" && modelType.startsWith("qwen3_5")) return { rope: "flat", validate: verifyOutputGate };
+  // Llama-style flat configs: a scalar rope_theta and no rope_parameters map.
+  if (modelType === "llama") return { rope: "theta", fullAttentionByDefault: true };
+  // Plain Qwen3 (Qwen3ForCausalLM, e.g. Qwen3-Embedding): handled like llama for rope.
+  if (modelType === "qwen3") return FLAT_THETA;
+  return PLAIN;
+}