@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/artifacts/safetensors-writer.ts

@@ -0,0 +1,249 @@
+// Sharded safetensors writer: takes a list of named mlx arrays and writes
+// them to disk as one or more `model*.safetensors` files via mlx's native
+// saver (mlx_save_safetensors), matching the on-disk layout the loader
+// (packages/inference/src/artifacts/weights.ts → ShardedSafetensors) reads back.
+//
+// Layout parity with mlx-lm / the MiniCPM5-OptiQ target:
+//   - single shard  → `model.safetensors`, NO index file.
+//   - multiple shards → `model-00001-of-0000N.safetensors`, plus a
+//     `model.safetensors.index.json` { metadata:{total_size,
+//     total_parameters}, weight_map: <name → shard> }.
+//   - metadata map carries {"format":"mlx"} so the file round-trips through
+//     mlx_load_safetensors with the same provenance HF/mlx writes.
+//
+// Memory discipline: each array is eval'd immediately before insertion so its
+// bytes are materialized one tensor at a time (the caller hands us already-
+// disposed-after-write arrays), bounding peak memory to ~one tensor over the
+// resident weight set rather than the whole model.
+
+import { ptr } from "bun:ffi";
+import { renameSync, writeFileSync } from "node:fs";
+import { join } from "node:path";
+import { C, takeMlxError } from "@mlx-bun/mlx/ffi";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+
+const cstr = (s: string) => Buffer.from(s + "\0", "utf8");
+
+/** Default greedy shard ceiling: 5 GiB (mlx-lm's MAX shard size). */
+export const DEFAULT_SHARD_BYTES = 5 << 30;
+
+export interface NamedTensor {
+  name: string;
+  array: MlxArray;
+}
+
+export interface ShardInfo {
+  /** Filename written into outDir (e.g. "model.safetensors"). */
+  file: string;
+  /** Tensor names placed in this shard, in insertion order. */
+  names: string[];
+  /** Byte size of the tensors in this shard (sum of nbytes). */
+  bytes: number;
+}
+
+export interface SafetensorsIndex {
+  metadata: { total_size: number; total_parameters: number };
+  weight_map: Record<string, string>;
+}
+
+export interface WriteResult {
+  shards: ShardInfo[];
+  totalSize: number;
+  totalParams: number;
+  /** The index object (also written to disk when shards.length > 1). */
+  index: SafetensorsIndex;
+}
+
+export interface WriteOpts {
+  /** Greedy shard ceiling in bytes (default DEFAULT_SHARD_BYTES). */
+  shardBytes?: number;
+}
+
+/** Insert one array into a native string→array map (key cstring, handle u64).
+ *  Evals the array first so its bytes are materialized exactly here. */
+function insertArray(map: bigint, name: string, array: MlxArray): void {
+  array.eval();
+  if (C.mlx_map_string_to_array_insert(map, ptr(cstr(name)), array.handle) !== 0)
+    throw new Error(`map insert ${name} failed: ${takeMlxError() ?? ""}`);
+}
+
+/** Build the {"format":"mlx"} metadata map. Caller frees it. */
+function newMetaMap(): bigint {
+  const map = C.mlx_map_string_to_string_new();
+  if (C.mlx_map_string_to_string_insert(map, ptr(cstr("format")), ptr(cstr("mlx"))) !== 0)
+    throw new Error(`metadata insert failed: ${takeMlxError() ?? ""}`);
+  return map;
+}
+
+/** Pack the shard's tensors into a fresh native map and save to `path`. */
+function saveShard(path: string, tensors: NamedTensor[]): void {
+  const arrMap = C.mlx_map_string_to_array_new();
+  const metaMap = newMetaMap();
+  try {
+    for (const { name, array } of tensors) insertArray(arrMap, name, array);
+    if (C.mlx_save_safetensors(ptr(cstr(path)), arrMap, metaMap) !== 0)
+      throw new Error(`mlx_save_safetensors(${path}) failed: ${takeMlxError() ?? ""}`);
+  } finally {
+    C.mlx_map_string_to_array_free(arrMap);
+    C.mlx_map_string_to_string_free(metaMap);
+  }
+}
+
+/** Shard filename for shard i of n (1-based, 5-digit, matches mlx-lm). */
+function shardName(i: number, n: number): string {
+  if (n === 1) return "model.safetensors";
+  const pad = (x: number) => String(x).padStart(5, "0");
+  return `model-${pad(i + 1)}-of-${pad(n)}.safetensors`;
+}
+
+/** Greedily bin tensors into shards no larger than `shardBytes`. A single
+ *  tensor larger than the ceiling still gets its own shard (never split). */
+function planShards(tensors: NamedTensor[], shardBytes: number): NamedTensor[][] {
+  const shards: NamedTensor[][] = [];
+  let cur: NamedTensor[] = [];
+  let curBytes = 0;
+  for (const t of tensors) {
+    const nb = t.array.nbytes;
+    if (cur.length > 0 && curBytes + nb > shardBytes) {
+      shards.push(cur);
+      cur = [];
+      curBytes = 0;
+    }
+    cur.push(t);
+    curBytes += nb;
+  }
+  if (cur.length > 0) shards.push(cur);
+  return shards.length > 0 ? shards : [[]];
+}
+
+/**
+ * Write `tensors` to `outDir` as sharded safetensors. Returns shard layout,
+ * total byte size, total parameter count, and the (single source of truth)
+ * index object. An index file is written only when there is more than one
+ * shard — matching the loader's expectation that a lone `model.safetensors`
+ * carries no sidecar index.
+ */
+export function writeShardedSafetensors(
+  outDir: string,
+  tensors: NamedTensor[],
+  opts: WriteOpts = {},
+): WriteResult {
+  const shardBytes = opts.shardBytes ?? DEFAULT_SHARD_BYTES;
+  const plan = planShards(tensors, shardBytes);
+  const n = plan.length;
+
+  const shards: ShardInfo[] = [];
+  const weight_map: Record<string, string> = {};
+  let totalSize = 0;
+  let totalParams = 0;
+
+  for (let i = 0; i < n; i++) {
+    const group = plan[i]!;
+    const file = shardName(i, n);
+    let bytes = 0;
+    for (const t of group) {
+      bytes += t.array.nbytes;
+      totalParams += t.array.size;
+      weight_map[t.name] = file;
+    }
+    totalSize += bytes;
+    saveShard(join(outDir, file), group);
+    shards.push({ file, names: group.map((t) => t.name), bytes });
+  }
+
+  // weight_map sorted by name for a stable, diff-friendly index.
+  const sortedMap: Record<string, string> = {};
+  for (const name of Object.keys(weight_map).sort()) sortedMap[name] = weight_map[name]!;
+
+  const index: SafetensorsIndex = {
+    metadata: { total_size: totalSize, total_parameters: totalParams },
+    weight_map: sortedMap,
+  };
+
+  if (n > 1) {
+    writeFileSync(
+      join(outDir, "model.safetensors.index.json"),
+      JSON.stringify(index, null, 2),
+    );
+  }
+
+  return { shards, totalSize, totalParams, index };
+}
+
+/**
+ * Incremental sharded writer for model-scale streaming transforms (27B fold/
+ * quantize): tensors are added one at a time; when the current shard reaches
+ * the byte ceiling it is SAVED AND ITS ARRAYS DISPOSED immediately, so peak
+ * memory is one shard plus the tensor in flight — never the whole model
+ * (the whole-list writeShardedSafetensors path holds every array until the
+ * caller disposes them, which OOMs a 51 GB model on a 32 GB box).
+ *
+ * Shards are written under temp names (the final `-of-NNNNN` count is unknown
+ * until finish()) and renamed at finish, which also writes the index.
+ * add() takes ownership of the array (evaluates it now, disposes after save).
+ */
+export class ShardedWriter {
+  readonly #outDir: string;
+  readonly #shardBytes: number;
+  #cur: NamedTensor[] = [];
+  #curBytes = 0;
+  #done: { tmp: string; names: string[]; bytes: number }[] = [];
+  #totalParams = 0;
+
+  constructor(outDir: string, opts: WriteOpts = {}) {
+    this.#outDir = outDir;
+    this.#shardBytes = opts.shardBytes ?? DEFAULT_SHARD_BYTES;
+  }
+
+  add(name: string, array: MlxArray): void {
+    array.eval();
+    const nb = array.nbytes;
+    if (this.#cur.length > 0 && this.#curBytes + nb > this.#shardBytes) this.#flush();
+    this.#cur.push({ name, array });
+    this.#curBytes += nb;
+    this.#totalParams += array.size;
+  }
+
+  #flush(): void {
+    if (this.#cur.length === 0) return;
+    const tmp = `model-shard-${String(this.#done.length + 1).padStart(5, "0")}.tmp.safetensors`;
+    saveShard(join(this.#outDir, tmp), this.#cur);
+    this.#done.push({
+      tmp,
+      names: this.#cur.map((t) => t.name),
+      bytes: this.#curBytes,
+    });
+    for (const t of this.#cur) t.array.dispose();
+    this.#cur = [];
+    this.#curBytes = 0;
+  }
+
+  finish(): WriteResult {
+    this.#flush();
+    const n = this.#done.length;
+    const shards: ShardInfo[] = [];
+    const weight_map: Record<string, string> = {};
+    let totalSize = 0;
+    for (let i = 0; i < n; i++) {
+      const d = this.#done[i]!;
+      const file = shardName(i, n);
+      renameSync(join(this.#outDir, d.tmp), join(this.#outDir, file));
+      for (const name of d.names) weight_map[name] = file;
+      totalSize += d.bytes;
+      shards.push({ file, names: d.names, bytes: d.bytes });
+    }
+    const sortedMap: Record<string, string> = {};
+    for (const name of Object.keys(weight_map).sort()) sortedMap[name] = weight_map[name]!;
+    const index: SafetensorsIndex = {
+      metadata: { total_size: totalSize, total_parameters: this.#totalParams },
+      weight_map: sortedMap,
+    };
+    if (n > 1) {
+      writeFileSync(
+        join(this.#outDir, "model.safetensors.index.json"),
+        JSON.stringify(index, null, 2),
+      );
+    }
+    return { shards, totalSize, totalParams: this.#totalParams, index };
+  }
+}

package/src/artifacts/safetensors.ts

@@ -0,0 +1,164 @@
+// Safetensors reader: mmap the file (via libc — see mmap.ts for why not
+// Bun.mmap), parse the JSON header, expose tensors as raw pointers into
+// the mapping. No tensor bytes are ever copied; pages materialize only
+// when something (the GPU) touches them.
+//
+// Format: [u64 LE header_len][header_len bytes of JSON][data...]
+// Header: { "<name>": { dtype, shape, data_offsets: [begin, end] }, ... }
+// with optional "__metadata__"; offsets are relative to the data section.
+
+import { MmapFile } from "./mmap";
+
+import type { SafetensorsDtype } from "@mlx-bun/mlx/array";
+export type { SafetensorsDtype } from "@mlx-bun/mlx/array";
+
+export const DTYPE_SIZE: Record<SafetensorsDtype, number> = {
+  BOOL: 1, U8: 1, I8: 1, U16: 2, I16: 2, U32: 4, I32: 4, U64: 8, I64: 8,
+  F16: 2, BF16: 2, F32: 4, F64: 8,
+};
+
+export interface TensorInfo {
+  name: string;
+  dtype: SafetensorsDtype;
+  shape: number[];
+  /** Byte range relative to the start of the data section. */
+  begin: number;
+  end: number;
+}
+
+export class SafetensorsFile {
+  readonly path: string;
+  readonly metadata: Record<string, string>;
+  readonly tensors: Map<string, TensorInfo>;
+  readonly mmap: MmapFile;
+  readonly #dataStart: number;
+
+  private constructor(
+    path: string,
+    mmap: MmapFile,
+    tensors: Map<string, TensorInfo>,
+    metadata: Record<string, string>,
+    dataStart: number,
+  ) {
+    this.path = path;
+    this.mmap = mmap;
+    this.tensors = tensors;
+    this.metadata = metadata;
+    this.#dataStart = dataStart;
+  }
+
+  static open(path: string): SafetensorsFile {
+    const mmap = MmapFile.open(path);
+    if (mmap.size < 8) throw new Error(`${path}: too small for safetensors`);
+
+    const lenBytes = mmap.view(0, 8);
+    const headerLen = Number(
+      new DataView(lenBytes.buffer, lenBytes.byteOffset, 8).getBigUint64(0, true),
+    );
+    const dataStart = 8 + headerLen;
+    if (dataStart > mmap.size)
+      throw new Error(`${path}: header length ${headerLen} exceeds file size`);
+
+    const headerJson = new TextDecoder().decode(mmap.view(8, headerLen));
+    const header = JSON.parse(headerJson) as Record<
+      string,
+      { dtype: SafetensorsDtype; shape: number[]; data_offsets: [number, number] }
+    >;
+
+    const tensors = new Map<string, TensorInfo>();
+    let metadata: Record<string, string> = {};
+    for (const [name, entry] of Object.entries(header)) {
+      if (name === "__metadata__") {
+        metadata = entry as unknown as Record<string, string>;
+        continue;
+      }
+      const { dtype, shape, data_offsets: [begin, end] } = entry;
+      if (!(dtype in DTYPE_SIZE))
+        throw new Error(`${path}: tensor ${name} has unsupported dtype ${dtype}`);
+      const expected = shape.reduce((a, b) => a * b, 1) * DTYPE_SIZE[dtype];
+      if (end - begin !== expected)
+        throw new Error(
+          `${path}: tensor ${name} byte range ${end - begin} != shape ${JSON.stringify(shape)} × ${dtype}`,
+        );
+      if (dataStart + end > mmap.size)
+        throw new Error(`${path}: tensor ${name} extends past end of file`);
+      tensors.set(name, { name, dtype, shape, begin, end });
+    }
+    return new SafetensorsFile(path, mmap, tensors, metadata, dataStart);
+  }
+
+  #info(name: string): TensorInfo {
+    const t = this.tensors.get(name);
+    if (!t) throw new Error(`${this.path}: no tensor named ${name}`);
+    return t;
+  }
+
+  /** Raw pointer to the tensor's first byte (for zero-copy FFI handoff). */
+  pointer(name: string): number {
+    return this.mmap.pointer(this.#dataStart + this.#info(name).begin);
+  }
+
+  /** JS view of a tensor's bytes (small tensors / tests only). */
+  view(name: string): Uint8Array {
+    const t = this.#info(name);
+    return this.mmap.view(this.#dataStart + t.begin, t.end - t.begin);
+  }
+}
+
+/** A sharded model directory: index json + N safetensors files. */
+export class ShardedSafetensors {
+  readonly files = new Map<string, SafetensorsFile>();
+  readonly tensorToFile = new Map<string, SafetensorsFile>();
+
+  static async open(modelDir: string): Promise<ShardedSafetensors> {
+    const self = new ShardedSafetensors();
+    const indexPath = `${modelDir}/model.safetensors.index.json`;
+    if (await Bun.file(indexPath).exists()) {
+      const index = (await Bun.file(indexPath).json()) as {
+        weight_map: Record<string, string>;
+      };
+      for (const [tensor, file] of Object.entries(index.weight_map)) {
+        let sf = self.files.get(file);
+        if (!sf) {
+          sf = SafetensorsFile.open(`${modelDir}/${file}`);
+          self.files.set(file, sf);
+        }
+        self.tensorToFile.set(tensor, sf);
+      }
+    } else {
+      // Single-file artifacts: HF's model.safetensors, or mlx-whisper's
+      // weights.safetensors (mlx-community/whisper-*).
+      const file = (await Bun.file(`${modelDir}/model.safetensors`).exists())
+        ? "model.safetensors"
+        : (await Bun.file(`${modelDir}/weights.safetensors`).exists())
+          ? "weights.safetensors"
+          : "model.safetensors";
+      const sf = SafetensorsFile.open(`${modelDir}/${file}`);
+      self.files.set(file, sf);
+      for (const name of sf.tensors.keys()) self.tensorToFile.set(name, sf);
+    }
+    return self;
+  }
+
+  get tensorNames(): string[] {
+    return [...this.tensorToFile.keys()].sort();
+  }
+
+  #file(name: string): SafetensorsFile {
+    const sf = this.tensorToFile.get(name);
+    if (!sf) throw new Error(`no tensor named ${name}`);
+    return sf;
+  }
+
+  info(name: string): TensorInfo {
+    return this.#file(name).tensors.get(name)!;
+  }
+
+  pointer(name: string): number {
+    return this.#file(name).pointer(name);
+  }
+
+  view(name: string): Uint8Array {
+    return this.#file(name).view(name);
+  }
+}

package/src/artifacts/turboquant-fold.ts

@@ -0,0 +1,42 @@
+// The TurboQuant fold record: `turboquant_fold.json`, the auxiliary file the
+// Trellis quantizer (packages/quantize/src/trellis-quantizer.ts) writes beside a
+// rotated checkpoint. It records the R1 rotation the weights were folded with
+// and the final-norm gain the fold moved into `lm_head`. Records written before
+// the quantizer recorded the gain carry no `final_gain`.
+
+import { existsSync, readFileSync } from "node:fs";
+import { join } from "node:path";
+
+export const TURBOQUANT_FOLD_FILE = "turboquant_fold.json";
+
+/** What a rotated checkpoint records about its fold. */
+export interface TurboQuantFold {
+  /** Seed of R1's sign vector. */
+  readonly seed: number;
+  /** The residual stream carries the R1 rotation. */
+  readonly r1: boolean;
+  readonly hiddenSize: number;
+  /** `final_gain`: the effective final-norm gain γ (stored + 1 of the source's
+   *  final norm) the fold moved into `lm_head`, one per hidden channel. Absent
+   *  when the record does not carry it. */
+  readonly finalGain?: Float32Array;
+}
+
+/** The fold record in `modelDir`, or null when the artifact carries none.
+ *  Throws when the record lacks a field or its gain does not span `hiddenSize`. */
+export function readTurboQuantFold(modelDir: string): TurboQuantFold | null {
+  const path = join(modelDir, TURBOQUANT_FOLD_FILE);
+  if (!existsSync(path)) return null;
+  const raw = JSON.parse(readFileSync(path, "utf8")) as Record<string, unknown>;
+  const { seed, r1, hiddenSize } = raw;
+  if (!Number.isInteger(seed)) throw new Error(`${path}: seed must be an integer (got ${String(seed)})`);
+  if (typeof r1 !== "boolean") throw new Error(`${path}: r1 must be a boolean (got ${String(r1)})`);
+  if (!Number.isInteger(hiddenSize) || (hiddenSize as number) <= 0)
+    throw new Error(`${path}: hiddenSize must be a positive integer (got ${String(hiddenSize)})`);
+  const fold = { seed: seed as number, r1, hiddenSize: hiddenSize as number };
+  if (raw.final_gain === undefined) return Object.freeze(fold);
+  const gain = raw.final_gain;
+  if (!Array.isArray(gain) || gain.length !== hiddenSize || !gain.every(value => typeof value === "number" && Number.isFinite(value)))
+    throw new Error(`${path}: final_gain must be ${String(hiddenSize)} finite numbers`);
+  return Object.freeze({ ...fold, finalGain: Float32Array.from(gain as number[]) });
+}

package/src/artifacts/weight-audit.ts

@@ -0,0 +1,43 @@
+// Universal-tier primitives: dense (unquantized) linear/embedding modules
+// (mlx nn.Linear / nn.Embedding semantics), quantized-or-dense
+// loader helpers, norm loaders, and the load-time weight audit.
+//
+// The quantized paths delegate to the proven QuantizedLinear/QuantizedEmbedding
+// primitives (packages/inference/src/layers/); the
+// dense paths are verbatim ports of mlx.nn.Linear (`mx.addmm(bias, x, W.T)`
+// / `x @ W.T`) and mlx.nn.Embedding (`weight[ids]`, `x @ weight.T`).
+
+import { MlxArray } from "@mlx-bun/mlx/array";
+import type { Weights } from "./weights";
+
+/** Records every tensor a universal load consumes; `finish` diffs against
+ *  the shard index so a descriptor mistake is a LOAD error (unconsumed /
+ *  missing tensors named), never a silently-wrong model
+ *  (`02d723a:docs/design/generic-model-support.md` §4.1). */
+export class WeightAudit {
+  readonly consumed = new Set<string>();
+
+  use(name: string): void {
+    this.consumed.add(name);
+  }
+
+  /** `drop` = the arch's sanitize rules (tensors mlx-lm discards on load,
+   *  e.g. rotary_emb.inv_freq, or lm_head.weight under tied embeddings). */
+  finish(weights: Weights, drop: RegExp[]): void {
+    const unconsumed = weights.tensorNames.filter(
+      (n) => !this.consumed.has(n) && !drop.some((re) => re.test(n)),
+    );
+    if (unconsumed.length > 0)
+      throw new Error(
+        `weight audit: ${unconsumed.length} tensor(s) in the checkpoint were not consumed ` +
+        `by the universal module (descriptor mismatch?): ${unconsumed.slice(0, 12).join(", ")}` +
+        (unconsumed.length > 12 ? ", …" : ""),
+      );
+  }
+}
+
+/** Load a tensor and record it in the audit. */
+export function tensorUsed(weights: Weights, audit: WeightAudit, name: string): MlxArray {
+  audit.use(name);
+  return weights.tensor(name);
+}