@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/artifacts/config.ts

@@ -0,0 +1,341 @@
+// Model config + quantization metadata reader for the load path.
+// Parses HF config.json (Gemma 4 unified layout: text_config nested),
+// the OptiQ per-layer quantization map, and kv_config.json when present.
+
+import { dialectFor } from "./config-dialects";
+import {
+  parseGlm52Config,
+  type Glm52Config,
+} from "./glm52-config";
+
+export interface RopeParams {
+  ropeTheta: number;
+  ropeType: string;
+  partialRotaryFactor: number;
+  /** ProportionalRoPE frequency multiplier (rope_utils: factor, default 1). */
+  factor: number;
+}
+
+export interface TextConfig {
+  hiddenSize: number;
+  numHiddenLayers: number;
+  numAttentionHeads: number;
+  numKeyValueHeads: number;
+  headDim: number;
+  /** Global (full-attention) layers can use different KV geometry. */
+  numGlobalKeyValueHeads: number;
+  globalHeadDim: number;
+  attentionKEqV: boolean;
+  intermediateSize: number;
+  hiddenActivation: string;
+  rmsNormEps: number;
+  vocabSize: number;
+  maxPositionEmbeddings: number;
+  slidingWindow: number;
+  /** "sliding_attention" | "full_attention" per layer. */
+  layerTypes: string[];
+  /** e2b/e4b: per-layer input embedding width (0 = disabled). */
+  hiddenSizePerLayerInput: number;
+  vocabSizePerLayerInput: number;
+  /** Trailing layers that share KV with earlier same-type layers. */
+  numKvSharedLayers: number;
+  /** 26B-A4B MoE block: dense MLP + routed experts per layer. */
+  enableMoeBlock: boolean;
+  numExperts: number;
+  topKExperts: number;
+  moeIntermediateSize: number;
+  /** Qwen3-MoE sparse-block gating (qwen3_moe): a layer is MoE unless it is
+   *  in `mlpOnlyLayers` and `(idx+1) % decoderSparseStep == 0` selects it.
+   *  `normTopkProb` renormalizes the top-k gate weights. Defaults are inert
+   *  for non-qwen3_moe families. */
+  decoderSparseStep: number;
+  mlpOnlyLayers: number[];
+  normTopkProb: boolean;
+  /** Qwen3.5 hybrid (gated-DeltaNet) geometry; 0/false for other families. */
+  linearNumValueHeads: number;
+  linearNumKeyHeads: number;
+  linearKeyHeadDim: number;
+  linearValueHeadDim: number;
+  linearConvKernelDim: number;
+  /** Every `fullAttentionInterval`-th layer is full-attention (Qwen3.5). */
+  fullAttentionInterval: number;
+  /** Qwen3.5 full-attention output gate: q_proj emits 2× and gates the output. */
+  attnOutputGate: boolean;
+  /** Rotary fraction of head_dim (Qwen3.5 full attention: 0.25 → 64 of 256). */
+  partialRotaryFactor: number;
+  ropeParameters: Record<string, RopeParams>;
+  finalLogitSoftcapping: number | null;
+  tieWordEmbeddings: boolean;
+  bosTokenId: number;
+  eosTokenId: number | number[];
+  /** DiffusionGemma: fixed denoising-canvas length (256). Undefined elsewhere. */
+  canvasLength?: number;
+}
+
+export interface QuantSpec {
+  bits: number;
+  groupSize: number;
+  mode: string;
+  /** `mode: "trellis"` (packed trellis-coded weights, packages/inference/src/layers/trellis-linear.ts):
+   *  bits = k per weight, groupSize = block T, plus the codec geometry. */
+  trellis?: TrellisSpec;
+}
+
+export interface TrellisSpec {
+  L: number;
+  code: string;
+  /** 1 = coded along the input dim (stored [out, in·k/32]); 0 = along the
+   *  output dim (stored [in, out·k/32], i.e. Wᵀ coded along its last axis). */
+  axis: 0 | 1;
+}
+
+export interface QuantizationConfig {
+  default: QuantSpec;
+  /** Per-module overrides keyed by module path (e.g. "language_model.model.embed_tokens"). */
+  perLayer: Map<string, QuantSpec>;
+}
+
+export interface KvQuantSpec {
+  layerIdx: number;
+  bits: number;
+  groupSize: number;
+}
+
+/** TurboQuant scheme (`02d723a:docs/design/turboquant.md`): rotation-based KV
+ *  quantization, a distinct axis from the uniform/per-layer affine kvQuant
+ *  above (composes with it in principle; v1 ships as a standalone CLI-only
+ *  runtime lever, same class as uniform kvBits — see
+ *  packages/inference/src/generation/autoregressive.ts maybeQuantizeKv). Valid
+ *  kBits: {2,4,5,8}; vBits: {2,3,4,5,8} (`02d723a:docs/design/turboquant.md`'s
+ *  supported Lloyd-Max tables). */
+export interface TurboQuantScheme {
+  kBits: number;
+  vBits: number;
+}
+
+export const TURBOQUANT_VALID_KBITS = [2, 4, 5, 8] as const;
+export const TURBOQUANT_VALID_VBITS = [2, 3, 4, 5, 8] as const;
+
+/** Parse `turbo` / `turbo:k<bits>v<bits>` (default k8v3). Returns null for
+ *  non-turbo strings (caller falls through to the uniform/config parser);
+ *  throws Error with a user-facing message for a malformed turbo spec. */
+export function parseTurboQuantScheme(raw: string): TurboQuantScheme | null {
+  if (raw !== "turbo" && !raw.startsWith("turbo:")) return null;
+  if (raw === "turbo") return { kBits: 8, vBits: 3 };
+  const m = /^turbo:k(\d+)v(\d+)$/.exec(raw);
+  if (!m) throw new Error(`--kv-quant turbo spec must look like "turbo:k<bits>v<bits>" (got "${raw}")`);
+  const kBits = Number(m[1]);
+  const vBits = Number(m[2]);
+  if (!(TURBOQUANT_VALID_KBITS as readonly number[]).includes(kBits))
+    throw new Error(`--kv-quant turbo: kBits must be one of ${TURBOQUANT_VALID_KBITS.join(",")} (got ${kBits})`);
+  if (!(TURBOQUANT_VALID_VBITS as readonly number[]).includes(vBits))
+    throw new Error(`--kv-quant turbo: vBits must be one of ${TURBOQUANT_VALID_VBITS.join(",")} (got ${vBits})`);
+  return { kBits, vBits };
+}
+
+export interface ModelConfig {
+  modelDir: string;
+  modelType: string;
+  architectures: string[];
+  dtype: string;
+  text: TextConfig;
+  quantization: QuantizationConfig | null;
+  kvQuant: KvQuantSpec[] | null;
+  hasVisionSidecar: boolean;
+  eosTokenIds: number[];
+  raw: Record<string, unknown>;
+}
+
+function parseRope(raw: Record<string, any>): Record<string, RopeParams> {
+  const out: Record<string, RopeParams> = {};
+  for (const [kind, p] of Object.entries(raw ?? {})) {
+    // transformers ≥5.10 configs mix flat scalars into the per-attention-type
+    // map (`rope_parameters: { full_attention: {...}, rope_theta: null,
+    // rope_type: "default", ... }` — the DeepSpec drafter ships this shape);
+    // only object entries are attention-type params.
+    if (!p || typeof p !== "object") continue;
+    out[kind] = {
+      ropeTheta: p.rope_theta,
+      ropeType: p.rope_type ?? "default",
+      partialRotaryFactor: p.partial_rotary_factor ?? 1.0,
+      factor: p.factor ?? 1.0,
+    };
+  }
+  return out;
+}
+
+export function parseQuantization(raw: Record<string, any> | undefined): QuantizationConfig | null {
+  if (!raw) return null;
+  const def: QuantSpec = {
+    bits: raw.bits,
+    groupSize: raw.group_size,
+    mode: raw.mode ?? "affine",
+  };
+  const perLayer = new Map<string, QuantSpec>();
+  for (const [key, value] of Object.entries(raw)) {
+    if (key === "bits" || key === "group_size" || key === "mode") continue;
+    if (typeof value === "object" && value !== null) {
+      const spec: QuantSpec = {
+        bits: (value as any).bits ?? def.bits,
+        groupSize: (value as any).group_size ?? def.groupSize,
+        mode: (value as any).mode ?? def.mode,
+      };
+      const tr = (value as any).trellis;
+      if (spec.mode === "trellis") {
+        if (!tr || typeof tr !== "object") throw new Error(`quantization.${key}: mode "trellis" needs a trellis {L, code, axis} block`);
+        spec.trellis = { L: Number(tr.L), code: String(tr.code ?? "1mad"), axis: tr.axis === 0 ? 0 : 1 };
+      }
+      perLayer.set(key, spec);
+    } else if (value === false) {
+      // mlx convention: `"layer": false` means not quantized
+      perLayer.set(key, { bits: 0, groupSize: 0, mode: "none" });
+    }
+  }
+  return { default: def, perLayer };
+}
+
+/** Quant spec for a module path, or null if the module is unquantized. */
+export function quantFor(q: QuantizationConfig | null, modulePath: string): QuantSpec | null {
+  if (!q) return null;
+  const spec = q.perLayer.get(modulePath) ?? q.default;
+  return spec.mode === "none" ? null : spec;
+}
+
+export async function loadModelConfig(modelDir: string): Promise<ModelConfig> {
+  const raw = (await Bun.file(`${modelDir}/config.json`).json()) as Record<string, any>;
+  const modelType = raw.model_type;
+  const dialect = dialectFor(modelType);
+  // The converted Colibri artifact keeps stop metadata split across the two
+  // HF files. Preserve the old config-only behavior for every other family;
+  // GLM-5.2 explicitly unions both lists through its dedicated parser.
+  let glm52: Glm52Config | null = null;
+  if (dialect.container === "colibri") {
+    const generationFile = Bun.file(`${modelDir}/generation_config.json`);
+    const generation = await generationFile.exists()
+      ? await generationFile.json() as Record<string, any>
+      : null;
+    glm52 = parseGlm52Config(modelDir, raw, generation);
+  }
+  // Gemma 4 unified nests the LM config; plain text models keep it at top level.
+  // Some architectures keep it flat but omit all dims — backfill from their defaults.
+  const baseT = (raw.text_config ?? raw) as Record<string, any>;
+  const t = dialect.textDefaults ? { ...dialect.textDefaults(), ...baseT } : baseT;
+  const qwenRope = (t.rope_parameters ?? {}) as Record<string, any>;
+  dialect.validate?.(t);
+
+  const text: TextConfig = {
+    hiddenSize: glm52?.hiddenSize ?? t.hidden_size,
+    numHiddenLayers: glm52?.numHiddenLayers ?? t.num_hidden_layers,
+    numAttentionHeads: glm52?.numAttentionHeads ?? t.num_attention_heads,
+    numKeyValueHeads: glm52?.numKeyValueHeads ?? t.num_key_value_heads,
+    headDim: glm52?.qkHeadDim ?? t.head_dim ?? Math.floor(t.hidden_size / t.num_attention_heads),
+    numGlobalKeyValueHeads:
+      glm52?.numKeyValueHeads ?? t.num_global_key_value_heads ?? t.num_key_value_heads,
+    globalHeadDim:
+      glm52?.qkHeadDim ??
+      t.global_head_dim ??
+      t.head_dim ??
+      Math.floor(t.hidden_size / t.num_attention_heads),
+    attentionKEqV: t.attention_k_eq_v ?? false,
+    intermediateSize: glm52?.intermediateSize ?? t.intermediate_size,
+    hiddenActivation: t.hidden_activation ?? t.hidden_act ?? "gelu_pytorch_tanh",
+    rmsNormEps: glm52?.rmsNormEps ?? t.rms_norm_eps,
+    vocabSize: glm52?.vocabSize ?? t.vocab_size,
+    maxPositionEmbeddings:
+      glm52?.maxPositionEmbeddings ?? t.max_position_embeddings,
+    slidingWindow: t.sliding_window ?? 0,
+    layerTypes: glm52
+      ? Array(glm52.numHiddenLayers).fill("full_attention")
+      : t.layer_types ?? (dialect.fullAttentionByDefault ? Array(t.num_hidden_layers).fill("full_attention") : []),
+    hiddenSizePerLayerInput: t.hidden_size_per_layer_input ?? 0,
+    vocabSizePerLayerInput:
+      t.vocab_size_per_layer_input ?? glm52?.vocabSize ?? t.vocab_size,
+    numKvSharedLayers: t.num_kv_shared_layers ?? 0,
+    enableMoeBlock: glm52 ? true : t.enable_moe_block ?? false,
+    numExperts: glm52?.numRoutedExperts ?? t.num_experts ?? 0,
+    // gemma4 uses `top_k_experts`; qwen3_moe uses `num_experts_per_tok`.
+    topKExperts:
+      glm52?.numExpertsPerToken ?? t.top_k_experts ?? t.num_experts_per_tok ?? 0,
+    moeIntermediateSize: glm52?.moeIntermediateSize ?? t.moe_intermediate_size ?? 0,
+    decoderSparseStep: glm52 ? 1 : t.decoder_sparse_step ?? 1,
+    mlpOnlyLayers: glm52
+      ? Array.from({ length: glm52.firstKDenseReplace }, (_, layer) => layer)
+      : t.mlp_only_layers ?? [],
+    normTopkProb: glm52?.normTopkProb ?? t.norm_topk_prob ?? false,
+    linearNumValueHeads: t.linear_num_value_heads ?? 0,
+    linearNumKeyHeads: t.linear_num_key_heads ?? 0,
+    linearKeyHeadDim: t.linear_key_head_dim ?? 0,
+    linearValueHeadDim: t.linear_value_head_dim ?? 0,
+    linearConvKernelDim: t.linear_conv_kernel_dim ?? 0,
+    fullAttentionInterval: t.full_attention_interval ?? 0,
+    attnOutputGate: t.attn_output_gate ?? false,
+    partialRotaryFactor: glm52
+      ? glm52.qkRopeHeadDim / glm52.qkHeadDim
+      : qwenRope.partial_rotary_factor ?? t.partial_rotary_factor ?? 1.0,
+    ropeParameters: glm52
+      ? {
+          full_attention: {
+            ropeTheta: glm52.ropeTheta,
+            ropeType: "default",
+            partialRotaryFactor: glm52.qkRopeHeadDim / glm52.qkHeadDim,
+            factor: 1.0,
+          },
+        }
+      : dialect.rope === "flat"
+      ? {
+          full_attention: {
+            ropeTheta: qwenRope.rope_theta ?? t.rope_theta ?? 10000,
+            ropeType: "default",
+            partialRotaryFactor: qwenRope.partial_rotary_factor ?? 1.0,
+            factor: 1.0,
+          },
+        }
+      : t.rope_parameters
+        ? parseRope(t.rope_parameters)
+        : dialect.rope === "theta"
+          ? {
+              full_attention: {
+                ropeTheta: t.rope_theta ?? 10000,
+                ropeType: "default",
+                partialRotaryFactor: 1.0,
+                factor: 1.0,
+              },
+            }
+          : {},
+    finalLogitSoftcapping: t.final_logit_softcapping ?? null,
+    tieWordEmbeddings: t.tie_word_embeddings ?? raw.tie_word_embeddings ?? false,
+    bosTokenId: t.bos_token_id ?? raw.bos_token_id,
+    eosTokenId: glm52?.eosTokenIds ?? t.eos_token_id ?? raw.eos_token_id,
+    canvasLength: raw.canvas_length ?? t.canvas_length,
+  };
+
+  let kvQuant: KvQuantSpec[] | null = null;
+  if (await Bun.file(`${modelDir}/kv_config.json`).exists()) {
+    const kv = (await Bun.file(`${modelDir}/kv_config.json`).json()) as any[];
+    kvQuant = kv.map((e) => ({
+      layerIdx: e.layer_idx,
+      bits: e.bits,
+      groupSize: e.group_size,
+    }));
+  }
+
+  const eos = glm52?.eosTokenIds ?? raw.eos_token_id ?? text.eosTokenId;
+  return {
+    modelDir,
+    modelType,
+    architectures: raw.architectures ?? [],
+    dtype: raw.dtype ?? "bfloat16",
+    text,
+    // Colibri's config retains the SOURCE FP8 quantization_config, while the
+    // serving shards are converted U8 + `.qs` tensors whose int4/int8/group
+    // geometry is inferred and validated by Glm52Container. Advertising the
+    // source FP8 block as mlx affine quantization would be actively wrong.
+    quantization: glm52
+      ? null
+      : parseQuantization(raw.quantization ?? raw.quantization_config),
+    kvQuant,
+    hasVisionSidecar: await Bun.file(`${modelDir}/optiq_vision.safetensors`).exists(),
+    eosTokenIds: Array.isArray(eos) ? eos : [eos],
+    raw,
+  };
+}

package/src/artifacts/draft-vocabulary.ts

@@ -0,0 +1,30 @@
+// A frequency-ranked draft vocabulary: `draft_vocab.json`, `{"ids": [...]}`, the
+// target token ids a drafter may propose (FR-Spec, Zhao et al. 2025). The Qwen3.8
+// Trellis artifact ships one in its MTP companion directory (`mtp/`); the
+// companion's drafts and, where a graph's composition declares it, a tapped
+// drafter's select over those rows of the target head.
+
+import { readFileSync } from "node:fs";
+
+export const DRAFT_VOCAB_FILE = "draft_vocab.json";
+
+const lists = new Map<string, readonly number[]>();
+
+/** The list at `path`: distinct ids in ascending order, read once per path.
+ *  Throws when the file holds no ids, an id is not a nonnegative integer, or one
+ *  falls outside the target's `vocabSize`. */
+export function readDraftVocabulary(path: string, vocabSize: number): readonly number[] {
+  let ids = lists.get(path);
+  if (!ids) {
+    const parsed = JSON.parse(readFileSync(path, "utf8")) as { ids?: unknown };
+    if (!Array.isArray(parsed.ids) || parsed.ids.length === 0)
+      throw new Error(`draft vocabulary ${path}: expected {"ids":[...]}`);
+    const sorted = [...new Set(parsed.ids as number[])].sort((a, b) => a - b);
+    if (!sorted.every(id => Number.isInteger(id) && id >= 0))
+      throw new Error(`draft vocabulary ${path}: ids must be nonnegative integers`);
+    lists.set(path, ids = Object.freeze(sorted));
+  }
+  if (ids.at(-1)! >= vocabSize)
+    throw new Error(`draft vocabulary ${path}: id ${ids.at(-1)} outside the target vocabulary ${vocabSize}`);
+  return ids;
+}

package/src/artifacts/expert-offload-build.ts

@@ -0,0 +1,75 @@
+// Expert offload — BUILD. Produces the page-aligned offload file
+// (experts.bin + manifest.json) that `activateExpertOffload` serves: each
+// expert tensor's bytes are copied to a 16 KB-aligned offset so the GPU can
+// gather it directly from a read-only file mmap (safetensors packs tensors at
+// arbitrary offsets). Kept apart from the runtime module so the hot model path
+// never imports the converter code.
+
+import { closeSync, existsSync, mkdirSync, openSync, readFileSync, statSync, writeFileSync, writeSync } from "node:fs";
+import { ShardedSafetensors } from "./safetensors";
+import type { OffloadEntry, OffloadManifest } from "./expert-offload";
+
+const PAGE = 16384;
+const alignUp = (n: number, a: number) => Math.ceil(n / a) * a;
+const isExpertTensor = (n: string) => n.includes("switch_glu") || n.includes(".experts.");
+const layerOf = (n: string) => Number(n.match(/layers\.(\d+)/)?.[1] ?? -1);
+
+/** Reuse `<modelPath>/.mlx-bun-offload` when present and valid for this model,
+ * else build it. Returns the offload dir to pass to `activateExpertOffload`. */
+export async function ensureOffloadFile(
+  modelPath: string,
+  onProgress?: (message: string) => void,
+  outDir = `${modelPath}/.mlx-bun-offload`,
+): Promise<string> {
+  const bin = `${outDir}/experts.bin`;
+  const man = `${outDir}/manifest.json`;
+  if (existsSync(bin) && existsSync(man)) {
+    try {
+      const manifest = JSON.parse(readFileSync(man, "utf8")) as OffloadManifest;
+      if (manifest.model === modelPath && statSync(bin).size === manifest.totalBytes) {
+        onProgress?.(`reusing ${bin}`);
+        return outDir;
+      }
+    } catch { /* fall through to rebuild */ }
+  }
+  await buildOffloadFile(modelPath, outDir, onProgress);
+  return outDir;
+}
+
+/** Pack every expert tensor (layers below `layerCap`) into a page-aligned file. */
+export async function buildOffloadFile(
+  modelPath: string,
+  outDir: string,
+  onProgress?: (message: string) => void,
+  layerCap = Infinity,
+): Promise<OffloadManifest> {
+  const shards = await ShardedSafetensors.open(modelPath);
+  const names = shards.tensorNames
+    .filter(isExpertTensor)
+    .filter(name => layerOf(name) < layerCap)
+    .sort((a, b) => layerOf(a) - layerOf(b) || a.localeCompare(b));
+  if (!names.length) throw new Error(`${modelPath}: no expert tensors (not an MoE model?)`);
+
+  mkdirSync(outDir, { recursive: true });
+  const fd = openSync(`${outDir}/experts.bin`, "w");
+  const pad = Buffer.alloc(PAGE);
+  let offset = 0;
+  let count = 0;
+  const tensors: OffloadEntry[] = [];
+  try {
+    for (const name of names) {
+      const aligned = alignUp(offset, PAGE);
+      if (aligned > offset) { writeSync(fd, pad, 0, aligned - offset); offset = aligned; }
+      const info = shards.info(name);
+      const bytes = shards.view(name); // zero-copy view into the source mmap
+      writeSync(fd, bytes);
+      tensors.push({ name, offset, length: bytes.length, dtype: info.dtype, shape: info.shape });
+      offset += bytes.length;
+      if (++count % 30 === 0) onProgress?.(`packing experts ${(offset / 1e9).toFixed(1)} GB (${count}/${names.length})`);
+    }
+  } finally { closeSync(fd); }
+  const manifest: OffloadManifest = { page: PAGE, model: modelPath, totalBytes: offset, tensors };
+  writeFileSync(`${outDir}/manifest.json`, JSON.stringify(manifest));
+  onProgress?.(`wrote ${(offset / 1e9).toFixed(2)} GB, ${tensors.length} tensors, ${new Set(names.map(layerOf)).size} layers`);
+  return manifest;
+}

package/src/artifacts/expert-offload.ts

@@ -0,0 +1,89 @@
+// E1 expert offload — RUNTIME.
+//
+// Routes expert WEIGHT tensors through a read-only file mmap (clean,
+// page-aligned pages that cost ~0 phys_footprint and that the GPU gathers
+// directly — probes confirmed bit-exact + non-wiring) instead of the
+// anonymous mlx_load_safetensors copy that counts against memory pressure.
+//
+// Activated two ways, both before model construction:
+//   - env MLX_BUN_EXPERT_OFFLOAD=<dir>  (scripts / direct runs)
+//   - activateExpertOffload(dir)        (CLI `--expert-offload`)
+// The file is produced by packages/inference/src/artifacts/expert-offload-build.ts. Tensors not in the
+// manifest (or when never activated) fall back to the resident path, so this
+// is inert by default and safe on a partially-converted file.
+
+import { readFileSync } from "node:fs";
+import { MmapFile } from "./mmap";
+import { MlxArray, SAFETENSORS_TO_MLX } from "@mlx-bun/mlx/array";
+import type { SafetensorsDtype } from "./safetensors";
+import { runtimeValue } from "../runtime/config";
+
+export interface OffloadEntry {
+  name: string;
+  offset: number;
+  length: number;
+  dtype: SafetensorsDtype;
+  shape: number[];
+}
+export interface OffloadManifest {
+  page: number;
+  model: string;
+  totalBytes: number;
+  tensors: OffloadEntry[];
+}
+
+let mm: MmapFile | null = null;
+let manifest: Map<string, OffloadEntry> | null = null;
+
+/** Open the offload file at `dir` (experts.bin + manifest.json) and route
+ *  expert weights through it. Call BEFORE the model is constructed.
+ *
+ *  Routing is process-global. The returned function restores the routing that
+ *  was active before this call, so an owner that activated a file for one
+ *  model can stop later loads from resolving against it. It never unmaps:
+ *  weights already created from the mapping keep borrowing its bytes for the
+ *  process lifetime (see `expertOffloadArray`). */
+export function activateExpertOffload(dir: string): () => void {
+  // Resolve both halves before touching either global: a failed open leaves
+  // whatever routing was active untouched.
+  const parsed = JSON.parse(readFileSync(`${dir}/manifest.json`, "utf8")) as OffloadManifest;
+  const next = new Map(parsed.tensors.map((t) => [t.name, t]));
+  const mapping = MmapFile.open(`${dir}/experts.bin`, "ro");
+  const previous = { mm, manifest };
+  manifest = next;
+  mm = mapping;
+  process.stderr.write(`[expert-offload] mmap ${dir}/experts.bin (${parsed.tensors.length} tensors)\n`);
+  return () => { mm = previous.mm; manifest = previous.manifest; };
+}
+
+export function isExpertOffload(): boolean {
+  return mm !== null;
+}
+
+/** Zero-copy mmap-backed array for `name` if offload is active and the tensor
+ *  is in the manifest; else null. The converter aligns every tensor to 16 KB
+ *  and the mmap base is page-aligned, so the pointer is page-aligned —
+ *  required for the GPU to gather from it.
+ *
+ *  Lifetime: these arrays are created ONCE at model load
+ *  (QuantizedSwitchLinear.load, packages/inference/src/layers/quantized-switch-linear.ts)
+ *  and live as model weights;
+ *  every decode's gather_qmm references them, so GPU command buffers retain
+ *  the mlx buffer past any dispose and mlx may drop the LAST reference on
+ *  the Metal completion thread. That's why this must be fromPointer (native
+ *  free(NULL) dtor — safe from any thread), NOT the old fromView JSCallback
+ *  pin (the 2026-07-06 deadlock class). The memory contract is trivially
+ *  met: `mm` is module-level and never unmapped — the mapping owns the
+ *  bytes for the process (clean file-backed pages, ~0 phys_footprint). */
+export function expertOffloadArray(name: string): MlxArray | null {
+  if (!mm || !manifest) return null;
+  const e = manifest.get(name);
+  if (!e) return null;
+  // Bounds check the whole tensor (pointer() only checks the start).
+  if (e.offset + e.length > mm.size)
+    throw new Error(`expert-offload: ${name} [${e.offset}, +${e.length}) exceeds ${mm.path} (${mm.size} B)`);
+  return MlxArray.fromPointer(mm.pointer(e.offset), e.shape, SAFETENSORS_TO_MLX[e.dtype]);
+}
+
+const expertOffloadPath = runtimeValue("MLX_BUN_EXPERT_OFFLOAD");
+if (expertOffloadPath) activateExpertOffload(expertOffloadPath);

package/src/artifacts/experts/index.ts

@@ -0,0 +1,3 @@
+export * from "./io";
+export * from "./residency";
+export * from "./usage";