@mlx-bun/inference 0.0.0-stage → 1.0.0

package/examples/fit-model.ts

@@ -0,0 +1,25 @@
+import { readdir, stat } from "node:fs/promises";
+import { join } from "node:path";
+import { loadModelConfig } from "@mlx-bun/inference/artifacts/config";
+import { fit, thisMachine, type MachineSpec } from "@mlx-bun/inference/execution/fit";
+
+/** Estimate whether one checkpoint directory fits a machine at a context length.
+ *  Weight bytes here are every `.safetensors` file in the directory: a sidecar
+ *  such as a vision tower counts as resident weight, and MoE expert bytes are
+ *  not separated, so the decode estimate assumes every weight is read per token.
+ *  A registry that classifies files (such as `@mlx-bun/hub`) supplies exact
+ *  language-weight and expert bytes. Estimates are advisory. */
+export async function fitModel(modelDir: string, contextTokens: number, machine: MachineSpec = thisMachine()) {
+  const config = await loadModelConfig(modelDir);
+  let weightsBytes = 0;
+  for (const name of await readdir(modelDir))
+    if (name.endsWith(".safetensors")) weightsBytes += (await stat(join(modelDir, name))).size;
+  return fit(config, weightsBytes, contextTokens, machine);
+}
+
+if (import.meta.main) {
+  const [modelDir, context = "8192"] = process.argv.slice(2);
+  if (!modelDir) throw new Error("usage: bun packages/inference/examples/fit-model.ts <model-directory> [context-tokens]");
+  const { fits, maxSafeContext, totalBytes, predictedDecodeTps } = await fitModel(modelDir, Number(context));
+  console.log({ fits, maxSafeContext, totalBytes, predictedDecodeTps });
+}

package/examples/qwen3-forward.ts

@@ -0,0 +1,42 @@
+import { bindGraphCaches } from "@mlx-bun/inference/models";
+import { phaseRowLimit, rowOperationFor } from "@mlx-bun/inference/models/graph";
+import { MlxPhaseDispatch } from "@mlx-bun/inference/generation";
+import { loadModelConfig, Weights } from "@mlx-bun/inference/artifacts";
+import { resolveComposition } from "@mlx-bun/inference";
+
+import { Qwen3Model } from "@mlx-bun/inference/models/qwen3";
+import { argmaxLastPosition, lastPositionLogits } from "@mlx-bun/inference/scoring";
+import * as ops from "@mlx-bun/mlx/ops";
+
+/** Explicit forward passes through the graph's named phases: the caller
+ * supplies tokens, names the phase of each call and owns the live cache. A
+ * prompt shorter than the prefill chunk is one prompt tail. */
+export async function forwardTokens(directory: string, promptTokens: number[]) {
+  const config = await loadModelConfig(directory);
+  const weights = await Weights.open(directory);
+  try {
+    const graph = bindGraphCaches(new Qwen3Model(weights, config, phaseRowLimit(resolveComposition(config))), resolveComposition(config));
+    const state = graph.makeCache();
+    try {
+      using prompt = ops.fromInt32(promptTokens, [1, promptTokens.length]);
+      using promptHidden = await new MlxPhaseDispatch(graph).prompt("tail", 1, (promptTokens.length))(prompt, state, null);
+      using prefill = rowOperationFor(graph.logits,promptTokens.length)(promptHidden);
+      const token = argmaxLastPosition(prefill);
+      using next = ops.fromInt32([token], [1, 1]);
+      using nextHidden = await rowOperationFor(graph.decode,1)(next, state, null);
+      using continuation = rowOperationFor(graph.logits,1)(nextHidden);
+      return { token, nextToken: argmaxLastPosition(continuation),
+        logits: lastPositionLogits(continuation), offset: state[0]!.offset };
+    } finally {
+      for (const cache of state) cache.dispose();
+    }
+  } finally {
+    weights.dispose();
+  }
+}
+
+if (import.meta.main) {
+  const directory = process.argv[2], ids = process.argv[3];
+  if (!directory || !ids) throw new Error('Usage: bun qwen3-forward.ts <checkpoint-directory> <token-ids-as-JSON>');
+  console.log(await forwardTokens(directory, JSON.parse(ids)));
+}

package/examples/qwen3-generate.ts

@@ -0,0 +1,28 @@
+import { bindGraphCaches } from "@mlx-bun/inference/models";
+import { phaseRowLimit } from "@mlx-bun/inference/models/graph";
+import { loadModelConfig, Weights, loadTokenizer, generate, resolveComposition } from "@mlx-bun/inference";
+import { Qwen3Model } from "@mlx-bun/inference/models/qwen3";
+
+/** The caller chooses a Qwen3 checkpoint and owns its weights. */
+export async function generateText(directory: string, prompt: string, maxTokens = 32) {
+  const config = await loadModelConfig(directory);
+  const tokenizer = await loadTokenizer(directory);
+  const weights = await Weights.open(directory);
+  try {
+    const graph = bindGraphCaches(new Qwen3Model(weights, config, phaseRowLimit(resolveComposition(config))), resolveComposition(config));
+    const generation = generate(graph, tokenizer.encode(prompt), {
+      maxTokens, temperature: 0, eosTokenIds: config.eosTokenIds,
+    });
+    const tokens: number[] = [];
+    for await (const { token } of generation) tokens.push(token);
+    return { text: tokenizer.decode(tokens), tokens, stats: generation.stats };
+  } finally {
+    weights.dispose();
+  }
+}
+
+if (import.meta.main) {
+  const directory = process.argv[2];
+  if (!directory) throw new Error('Usage: bun qwen3-generate.ts <checkpoint-directory> [prompt]');
+  console.log(await generateText(directory, process.argv[3] ?? "Hello"));
+}

package/examples/trellis-expand.ts

@@ -0,0 +1,23 @@
+import { Dtype, MlxArray, ops } from "@mlx-bun/mlx";
+import { vectorTrellisExpand, vectorTrellisExpandEligible,
+  type TrellisGeometry } from "@mlx-bun/inference/kernels/trellis";
+
+/** Borrow packed weights; return an owned lazy array for the caller's graph. */
+export function expandWeights(codes: MlxArray, scales: MlxArray, geometry: TrellisGeometry) {
+  if (!vectorTrellisExpandEligible(geometry, Dtype.bfloat16))
+    throw new Error("This geometry is not supported by the vector expansion kernel");
+  return vectorTrellisExpand(codes, scales, geometry);
+}
+
+if (import.meta.main) {
+  // Small generated inputs demonstrate the layout; no checkpoint is needed.
+  const geometry: TrellisGeometry = {
+    k: 3, L: 12, T: 256, axis: 1,
+    rows: 64, cols: 512, inFeatures: 512, outFeatures: 64,
+  };
+  using codes = ops.zeros([geometry.rows, geometry.cols * geometry.k / 32], Dtype.uint32);
+  using scales = MlxArray.fromFloat32(new Float32Array(geometry.rows).fill(1), [geometry.rows]);
+  using weights = expandWeights(codes, scales, geometry);
+  weights.eval();
+  console.log(weights.shape);
+}

package/package.json

@@ -1,6 +1,131 @@
 {
   "name": "@mlx-bun/inference",
-  "version": "0.0.0-stage",
-  "stub": true,
-  "description": "Temporary package placeholder for staged publishing"
-}
+  "version": "1.0.0",
+  "description": "Composable inference kernels and graphs for MLX in Bun",
+  "type": "module",
+  "license": "MIT",
+  "engines": {
+    "bun": ">=1.4.2"
+  },
+  "exports": {
+    ".": "./src/index.ts",
+    "./kernels/trellis": "./src/kernels/trellis/index.ts",
+    "./artifacts": "./src/artifacts/index.ts",
+    "./input": "./src/input/index.ts",
+    "./layers": "./src/layers/index.ts",
+    "./runtime/config": "./src/runtime/config.ts",
+    "./layers/lora": "./src/layers/lora.ts",
+    "./state": "./src/state/index.ts",
+    "./contracts": "./src/contracts/index.ts",
+    "./kernels/turboquant": "./src/kernels/turboquant/index.ts",
+    "./kernels/delta": "./src/kernels/delta/index.ts",
+    "./kernels/attention/paged": "./src/kernels/attention/paged.ts",
+    "./state/paged": "./src/state/paged/cache.ts",
+    "./models/qwen3": "./src/models/qwen/qwen3.ts",
+    "./models/qwen3-moe": "./src/models/qwen/qwen3-moe.ts",
+    "./models/qwen3_5": "./src/models/qwen/qwen3_5.ts",
+    "./models/qwen/blocks": "./src/models/qwen/blocks.ts",
+    "./models/qwen/written-layers": "./src/models/qwen/written-layers.ts",
+    "./models/qwen38-27b-trellis-m4pro": "./src/models/qwen/qwen38-27b-trellis-m4pro.ts",
+    "./models/qwen38-27b-trellis-m1max": "./src/models/qwen/qwen38-27b-trellis-m1max.ts",
+    "./models/gemma4": "./src/models/gemma4/model.ts",
+    "./models/gemma4/graphs": "./src/models/gemma4/graphs.ts",
+    "./models/gemma4/blocks": "./src/models/gemma4/written-blocks.ts",
+    "./models/minicpm5": "./src/models/minicpm5/model.ts",
+    "./models/universal": "./src/models/universal/dense.ts",
+    "./models/universal/blocks": "./src/models/universal/blocks.ts",
+    "./models/graph": "./src/models/graph.ts",
+    "./kernels/attention/flash": "./src/kernels/attention/flash.ts",
+    "./scoring": "./src/scoring/index.ts",
+    "./models/gemma4/compiled-decode": "./src/models/gemma4/compiled-decode.ts",
+    "./artifacts/experts": "./src/artifacts/experts/index.ts",
+    "./models": "./src/models/index.ts",
+    "./models/memory-plan": "./src/models/memory-plan.ts",
+    "./models/glm52": "./src/models/glm52/model.ts",
+    "./models/diffusion-gemma": "./src/models/diffusion-gemma/model.ts",
+    "./models/whisper": "./src/models/whisper/model.ts",
+    "./artifacts/glm52": "./src/artifacts/glm52/index.ts",
+    "./input/vision": "./src/input/vision/index.ts",
+    "./kernels/glm52": "./src/kernels/glm52/streamed-experts.ts",
+    "./sampling": "./src/sampling/index.ts",
+    "./sampling/grammar": "./src/sampling/grammar.ts",
+    "./kernels/sampling": "./src/kernels/sampling/index.ts",
+    "./adapters": "./src/adapters/index.ts",
+    "./adapters/available": "./src/adapters/available.ts",
+    "./embeddings": "./src/embeddings/index.ts",
+    "./generation": "./src/generation/index.ts",
+    "./generation/diffusion": "./src/generation/diffusion.ts",
+    "./generation/fill": "./src/generation/fill/index.ts",
+    "./generation/bindings/autoregressive": "./src/generation/bindings/autoregressive.ts",
+    "./generation/bindings/denoising": "./src/generation/bindings/denoising.ts",
+    "./input/audio": "./src/input/audio/index.ts",
+    "./input/vision/qwen3vl": "./src/input/vision/qwen3vl-preprocess.ts",
+    "./input/vision/qwen3vl-prompt": "./src/input/vision/qwen3vl-prompt.ts",
+    "./input/media-fetch": "./src/input/media-fetch.ts",
+    "./models/audio/conformer": "./src/models/audio/conformer.ts",
+    "./transcription/silero-vad": "./src/transcription/silero-vad.ts",
+    "./models/vision/qwen3vl": "./src/models/vision/qwen3vl.ts",
+    "./models/vision/siglip": "./src/models/vision/siglip.ts",
+    "./models/vision/unified": "./src/models/vision/unified.ts",
+    "./transcription": "./src/transcription/index.ts",
+    "./transcription/format": "./src/transcription/whisper/format.ts",
+    "./generation/speculative": "./src/generation/speculative/index.ts",
+    "./generation/speculative/draft-kind": "./src/generation/speculative/draft-kind.ts",
+    "./generation/speculative/draft-registry": "./src/generation/speculative/draft-registry.ts",
+    "./generation/speculative/loader": "./src/models/speculative/loader.ts",
+    "./generation/speculative/binding": "./src/generation/speculative/bindings/binding.ts",
+    "./models/drafters": "./src/models/drafters.ts",
+    "./models/drafter-loaders": "./src/models/drafter-loaders.ts",
+    "./models/gemma4/assistant": "./src/models/gemma4/assistant.ts",
+    "./models/qwen/mtp": "./src/models/qwen/mtp.ts",
+    "./models/glm52/mtp": "./src/models/glm52/mtp.ts",
+    "./models/speculative/dflash": "./src/models/speculative/dflash.ts",
+    "./models/speculative/deepspec": "./src/models/speculative/deepspec.ts",
+    "./execution": "./src/execution/index.ts",
+    "./kernels/quantization": "./src/kernels/quantization/affine-rows.ts",
+    "./contracts/portable": "./src/contracts/portable/index.ts",
+    "./contracts/mlx": "./src/contracts/mlx/index.ts",
+    "./artifacts/config": "./src/artifacts/config.ts",
+    "./artifacts/auxiliary-files": "./src/artifacts/auxiliary-files.ts",
+    "./models/support": "./src/models/support.ts",
+    "./state/kv-scheme": "./src/state/kv-scheme.ts",
+    "./artifacts/glm52/memory": "./src/artifacts/glm52/memory.ts",
+    "./runtime/trace": "./src/runtime/trace.ts",
+    "./execution/engine": "./src/execution/engine.ts",
+    "./execution/admission": "./src/execution/admission.ts",
+    "./execution/cancellation": "./src/execution/cancellation.ts",
+    "./execution/fit": "./src/execution/fit.ts",
+    "./runtime/resources": "./src/runtime/resources.ts",
+    "./execution/request-policy": "./src/execution/request-policy.ts",
+    "./models/diffusion-gemma/vision": "./src/models/diffusion-gemma/vision.ts",
+    "./input/tool-call": "./src/input/tool-call.ts",
+    "./models/profile": "./src/models/profile.ts"
+  },
+  "files": [
+    "src",
+    "examples",
+    "dist/native",
+    "LICENSE",
+    "THIRD_PARTY_NOTICES.md"
+  ],
+  "scripts": {
+    "typecheck": "bun ../../scripts/typecheck.ts --noEmit && bun ../../scripts/typecheck.ts --project src/contracts/portable/tsconfig.json",
+    "test": "bun test",
+    "build:native": "sh scripts/build-native.sh",
+    "stage:native": "bun scripts/stage-native.ts",
+    "prepack": "bun scripts/stage-native.ts"
+  },
+  "dependencies": {
+    "@mlx-bun/mlx": "1.0.0",
+    "@huggingface/tokenizers": "^0.1.3",
+    "@huggingface/jinja": "^0.5.9",
+    "fast-png": "^8.0.0",
+    "@mlc-ai/web-xgrammar": "^0.1.27"
+  },
+  "os": [
+    "darwin"
+  ],
+  "cpu": [
+    "arm64"
+  ]
+}

package/src/adapters/available.ts

@@ -0,0 +1,110 @@
+// Finding adapters on disk and reading their configs: file system only, no native library, so a process that owns
+// no model (an isolated host's catalog) can list what a model could mount.
+import { existsSync, readdirSync } from "node:fs";
+import { resolve } from "node:path";
+
+/** Adapter scale from its config: mlx-lm writes lora_parameters.scale;
+ *  PEFT writes lora_alpha + r and optionally use_rslora. For ordinary PEFT
+ *  return alpha/r as the scale. For PEFT rsLoRA return alpha and set rsLora,
+ *  so the caller's one per-layer division produces alpha/√rank exactly once.
+ *  Internal configs already store their pre-√rank scale directly. */
+export async function readAdapterScale(dir: string): Promise<{ scale: number; rank: number | null; rsLora: boolean }> {
+  for (const name of ["optiq_lora_config.json", "adapter_config.json"]) {
+    const f = Bun.file(`${dir}/${name}`);
+    if (await f.exists()) {
+      const cfg = (await f.json()) as Record<string, any>;
+      const lp = cfg.lora_parameters;
+      const rsLora = Boolean(cfg.rs_lora ?? lp?.rs_lora ?? cfg.use_rslora ?? false);
+      if (lp && typeof lp === "object")
+        return { scale: Number(lp.scale ?? 20.0), rank: lp.rank ?? null, rsLora };
+      const alpha = Number(cfg.lora_alpha ?? 16);
+      const r = Number(cfg.r ?? 8);
+      return {
+        scale: rsLora ? alpha : r ? alpha / r : 1.0,
+        rank: r || null,
+        rsLora,
+      };
+    }
+  }
+  throw new Error(`no adapter_config.json in ${dir}`);
+}
+
+export function adapterWeightsFile(dir: string): string {
+  for (const name of ["adapters.safetensors", "adapter_model.safetensors"]) {
+    const p = `${dir}/${name}`;
+    if (existsSync(p)) return p;
+  }
+  throw new Error(`no adapter weights at ${dir}/adapters.safetensors or adapter_model.safetensors`);
+}
+
+/** A mountable adapter found on disk (not yet mounted). */
+export interface AvailableAdapter {
+  id: string; // directory basename — the handle to mount it under
+  path: string; // absolute adapter dir
+  scale: number;
+  rank: number | null;
+  /** Base model repo id (org/name) the adapter was trained on, or null if the
+   *  config doesn't record it. Lets the chat selector hide adapters that don't
+   *  fit the served model (a MiniCPM5 adapter can't mount on Gemma). */
+  baseModel: string | null;
+}
+
+/** The base-model repo id an adapter was trained on, from its config's
+ *  source_model / base_model_name_or_path. These are stored as machine-local
+ *  snapshot PATHS (…/models--ORG--NAME/snapshots/HASH), so recover the stable
+ *  repo id (ORG/NAME) from the path; returns null if unrecorded. */
+async function readAdapterBaseRepoId(dir: string): Promise<string | null> {
+  for (const [name, key] of [
+    ["optiq_lora_config.json", "source_model"],
+    ["adapter_config.json", "base_model_name_or_path"],
+  ] as const) {
+    const f = Bun.file(`${dir}/${name}`);
+    if (await f.exists()) {
+      const cfg = (await f.json()) as Record<string, any>;
+      const src = cfg[key];
+      if (typeof src === "string" && src) {
+        const m = src.match(/models--([^/]+)/);
+        return m ? m[1]!.replace(/--/g, "/") : src;
+      }
+    }
+  }
+  return null;
+}
+
+/** Scan adapter stores for mountable adapters: directories holding an adapter
+ *  weights file. id = dir basename; scale/rank/baseModel read from each adapter's
+ *  config. Dirs without weights (dataset folders) and unreadable stores are skipped.
+ *  Backs GET /v1/adapters/available, which populates the chat adapter selector. */
+export async function listAvailableAdapters(stores: string[]): Promise<AvailableAdapter[]> {
+  const out: AvailableAdapter[] = [];
+  const seen = new Set<string>();
+  for (const store of stores) {
+    let entries;
+    try {
+      entries = readdirSync(store, { withFileTypes: true });
+    } catch {
+      continue; // store missing → skip
+    }
+    for (const e of entries) {
+      if (!e.isDirectory()) continue;
+      const dir = resolve(store, e.name);
+      try {
+        adapterWeightsFile(dir);
+      } catch {
+        continue; // no weights → not an adapter
+      }
+      if (seen.has(dir)) continue;
+      seen.add(dir);
+      let scale = 1;
+      let rank: number | null = null;
+      try {
+        ({ scale, rank } = await readAdapterScale(dir));
+      } catch {
+        // keep defaults if the config is unreadable
+      }
+      const baseModel = await readAdapterBaseRepoId(dir).catch(() => null);
+      out.push({ id: e.name, path: dir, scale, rank, baseModel });
+    }
+  }
+  return out;
+}

package/src/adapters/index.ts

@@ -0,0 +1,3 @@
+export * from "../layers/lora";
+export * from "./available";
+export * from "./manager";