mlx-bun 0.5.0 → 1.0.0

package/src/model/factory.ts

@@ -1,215 +0,0 @@
-// Model construction consumes one declared profile. Exact artifact profiles
-// outrank family profiles; dedicated/generated graphs outrank the universal
-// fallback. Request-level methods are resolved elsewhere and are never changed
-// by this module.
-
-import { loadModelConfig, type ModelConfig } from "../config";
-import { totalmem } from "node:os";
-import { Weights } from "../weights";
-import { Gemma4Model } from "./gemma4";
-import { configFingerprint } from "./fingerprint";
-import { GENERATED } from "./generated";
-import { MiniCPM5Model } from "./minicpm5";
-import { Qwen35Model } from "./qwen3_5";
-import { Qwen3Model } from "./qwen3";
-import { Qwen3MoeModel } from "./qwen3-moe";
-import { DiffusionGemmaModel } from "./diffusion-gemma";
-import { Glm52Model } from "./glm52";
-import {
-  GLM52_G5_DEFAULT_CONTEXT_TOKENS,
-  GLM52_G5_DEFAULT_MAX_GENERATION_TOKENS,
-  GLM52_G5_DEFAULT_PROCESS_LIMIT_BYTES,
-  planGlm52MemoryForArtifact,
-  type Glm52MemoryPlan,
-} from "./glm52-memory";
-import { UniversalDenseModel } from "./universal/dense";
-import { WhisperModel } from "./whisper";
-import { genericArgsFor } from "./universal/archs";
-import {
-  ModelImplementationRegistry,
-  type ModelImplementation,
-  type ModelImplementationProvider,
-} from "./implementation";
-import {
-  assertResolvedModelProfile,
-  resolveModelProfile,
-  type ResolvedModelProfile,
-  type ResolveModelProfileOptions,
-  type ModelGraph,
-} from "./profile";
-
-export type RuntimeModel =
-  | Gemma4Model | MiniCPM5Model | Qwen35Model | Qwen3Model | Qwen3MoeModel
-  | DiffusionGemmaModel | Glm52Model | UniversalDenseModel;
-
-function residentImplementation(
-  id: string,
-  graph: ModelGraph,
-  create: ModelImplementation<Weights, RuntimeModel>["create"],
-): ModelImplementation<Weights, RuntimeModel> {
-  return { id, graph, loader: "safetensors",
-    loop: graph === "diffusion-gemma" ? "diffusion" : "autoregressive", create };
-}
-
-/** Engine-owned registrations. Exact quant profiles can name additional
- * implementations in a composed registry, without changing sessions or files. */
-export const MLX_MODEL_IMPLEMENTATIONS = new ModelImplementationRegistry<Weights, RuntimeModel>([
-  residentImplementation("diffusion-gemma", "diffusion-gemma", (weights, config) => new DiffusionGemmaModel(weights, config)),
-  residentImplementation("minicpm5", "minicpm5", (weights, config) => new MiniCPM5Model(weights, config)),
-  residentImplementation("qwen3.5", "qwen3.5", (weights, config) => new Qwen35Model(weights, config)),
-  residentImplementation("qwen3-moe", "qwen3-moe", (weights, config) => new Qwen3MoeModel(weights, config)),
-  residentImplementation("qwen3", "qwen3", (weights, config) => new Qwen3Model(weights, config)),
-  residentImplementation("gemma4", "gemma4", (weights, config) => new Gemma4Model(weights, config)),
-  residentImplementation("gemma4-generated", "gemma4", (weights, config) => {
-    const Graph = GENERATED.get(configFingerprint(config));
-    if (!Graph) throw new Error("no generated Gemma graph for this config; refusing to fall back");
-    return new Graph(weights, config);
-  }),
-  residentImplementation("universal-dense", "universal-dense", (weights, config) =>
-    new UniversalDenseModel(weights, config, genericArgsFor(config)!)),
-]);
-
-export interface ModelOpenOptions<Model = RuntimeModel> extends Glm52RuntimeOpenOptions {
-  readonly profiles?: ResolveModelProfileOptions;
-  readonly implementations?: ModelImplementationProvider<Weights, Model>;
-}
-
-export interface Glm52RuntimeOpenOptions {
-  /** Whole-process ceiling. Defaults to the smaller of the validated 25 GiB
-   * preset and physical RAM. */
-  memoryBudgetBytes?: number;
-  /** Context reserved by the exact GLM resource equation. */
-  contextTokens?: number;
-  /** Generated-token allowance within contextTokens. */
-  maxGenerationTokens?: number;
-  /** Maximum ordinary continuous-batch rows. */
-  batchSize?: number;
-  /** Native checkpoint MTP row; on by default. */
-  enableMtp?: boolean;
-  mtpDraftTokens?: number;
-  /** Test/diagnostic override; normal callers use physical RAM. */
-  machineBytes?: number;
-  /** Native expert-I/O dylib override. Packaged/dev defaults are automatic. */
-  libraryPath?: string;
-}
-
-/** Open the direct Colibri artifact through the validated bounded expert
- * runtime. The header-only plan runs before any resident tensor or expert slab
- * is opened, so an impossible configuration fails without committing memory. */
-export async function openGlm52RuntimeModel(
-  modelDir: string,
-  options: Glm52RuntimeOpenOptions = {},
-): Promise<{ model: Glm52Model; plan: Glm52MemoryPlan }> {
-  const machineBytes = options.machineBytes ?? totalmem();
-  const processLimitBytes = options.memoryBudgetBytes ??
-    Math.min(GLM52_G5_DEFAULT_PROCESS_LIMIT_BYTES, machineBytes);
-  const plan = await planGlm52MemoryForArtifact(modelDir, {
-    machineBytes,
-    processLimitBytes,
-    contextTokens: options.contextTokens ?? GLM52_G5_DEFAULT_CONTEXT_TOKENS,
-    maxGenerationTokens:
-      options.maxGenerationTokens ?? GLM52_G5_DEFAULT_MAX_GENERATION_TOKENS,
-    batchSize: options.batchSize ?? 1,
-    enableMtp: options.enableMtp !== false,
-    mtpDraftTokens: options.mtpDraftTokens,
-  });
-  const model = await Glm52Model.openStreamed(modelDir, {
-    budgetBytes: plan.processLimitBytes,
-    reserveBytes: plan.runtimeReserveBytes,
-    workingSlots: plan.mainWorkingSlots,
-    maxSlotsPerLayer: 1,
-    workers: 2,
-    libraryPath: options.libraryPath,
-    decodeKernel: "metal",
-    enableMtp: plan.enableMtp,
-    mtpDraftTokens: plan.mtpDraftTokens,
-  });
-  if (model.expertRuntime?.plan.plannedBytes !== plan.plannedProcessBytes) {
-    const actual = model.expertRuntime?.plan.plannedBytes ?? 0;
-    model.dispose();
-    throw new Error(
-      `GLM runtime plan ${actual} != preflight resource equation ` +
-      `${plan.plannedProcessBytes}`,
-    );
-  }
-  return { model, plan };
-}
-
-/** Open a Whisper checkpoint (encoder-decoder speech model) through the same
- * profile resolution as text models. Whisper never enters the chat loop, so
- * it is not a RuntimeModel: the transcription engine (src/audio/whisper-*)
- * owns it. */
-export async function openWhisperModel(
-  modelDir: string, options: { readonly profiles?: ResolveModelProfileOptions } = {},
-): Promise<{ model: WhisperModel; profile: ResolvedModelProfile; config: ModelConfig }> {
-  const config = await loadModelConfig(modelDir);
-  const profile = resolveModelProfile(config, options.profiles);
-  if (profile.profile.execution.graph !== "whisper")
-    throw new Error(`${modelDir} is not a Whisper checkpoint (profile ${profile.profile.id})`);
-  const weights = await Weights.open(modelDir);
-  try {
-    return { model: new WhisperModel(weights, config), profile, config };
-  } catch (error) {
-    weights.dispose();
-    throw error;
-  }
-}
-
-/**
- * Artifact-aware construction. GLM-5.2's Colibri snapshot has no ordinary
- * model.safetensors.index.json, so it must bypass Weights.open and use its
- * dedicated header catalog/tensor source.
- */
-export function openModel<Model>(
-  modelDir: string,
-  options: ModelOpenOptions<Model> & { readonly implementations: ModelImplementationProvider<Weights, Model> },
-): Promise<Model>;
-export function openModel(modelDir: string, options?: ModelOpenOptions): Promise<RuntimeModel>;
-export async function openModel<Model>(
-  modelDir: string,
-  options: ModelOpenOptions<Model> = {},
-): Promise<Model | RuntimeModel> {
-  const config = await loadModelConfig(modelDir);
-  const profile = resolveModelProfile(config, options.profiles);
-  if (profile.profile.execution.loader === "colibri") {
-    if (options.implementations || profile.profile.execution.implementation !== undefined)
-      throw new Error("custom Colibri implementations require a streamed loader binding; refusing to fall back");
-    return (await openGlm52RuntimeModel(modelDir, options)).model;
-  }
-  if (profile.profile.execution.loop === "encoder-decoder")
-    throw new Error(
-      `${modelDir} is a speech model (${profile.profile.id}); open it with openWhisperModel ` +
-      "or serve it as a transcription model — it has no chat/completion loop",
-    );
-  // Resolve before opening weights; missing or incompatible code allocates nothing.
-  const implementation = (options.implementations ?? MLX_MODEL_IMPLEMENTATIONS).select(config, profile);
-  const weights = await Weights.open(modelDir);
-  try { return implementation.create(weights, config, profile); }
-  catch (error) { weights.dispose(); throw error; }
-}
-
-/** Construct the binding named by a validated profile. A supplied registry can
- * return any backend-owned interface; the compatibility default returns the
- * legacy model union. The caller retains ownership of the supplied weights. */
-export function createModel<Model>(
-  weights: Weights, config: ModelConfig, resolved: ResolvedModelProfile | undefined,
-  implementations: ModelImplementationProvider<Weights, Model>,
-): Model;
-export function createModel(
-  weights: Weights, config: ModelConfig, resolved?: ResolvedModelProfile,
-  implementations?: ModelImplementationProvider<Weights, RuntimeModel>,
-): RuntimeModel;
-export function createModel<Model>(
-  weights: Weights,
-  config: ModelConfig,
-  resolved: ResolvedModelProfile = resolveModelProfile(config),
-  implementations: ModelImplementationProvider<Weights, Model> | ModelImplementationProvider<Weights, RuntimeModel> = MLX_MODEL_IMPLEMENTATIONS,
-): Model | RuntimeModel {
-  assertResolvedModelProfile(config, resolved);
-  if (resolved.profile.execution.loader === "colibri")
-    throw new Error(
-      "glm_moe_dsa uses the direct Colibri container; construct it with " +
-      "openModel(modelDir) or Glm52Model.open(modelDir)",
-    );
-  return implementations.select(config, resolved).create(weights, config, resolved);
-}

package/src/model/fingerprint.ts

@@ -1,45 +0,0 @@
-// Config fingerprint for generated-model dispatch (docs/archive/investigations/optimization_plan.md
-// Phase C): a stable hash over every config field that changes the
-// decode graph structure, plus the kv_config quant layout. A generated
-// specialization is used only when the fingerprint of the loaded config
-// matches the one it was generated from; everything else runs the
-// monolith (slow, never broken).
-
-import type { ModelConfig } from "../config";
-
-export function configFingerprint(config: ModelConfig): string {
-  const t = config.text;
-  const payload = JSON.stringify({
-    layerTypes: t.layerTypes,
-    hiddenSize: t.hiddenSize,
-    headDim: t.headDim,
-    globalHeadDim: t.globalHeadDim,
-    numAttentionHeads: t.numAttentionHeads,
-    numKeyValueHeads: t.numKeyValueHeads,
-    numGlobalKeyValueHeads: t.numGlobalKeyValueHeads,
-    attentionKEqV: t.attentionKEqV,
-    numKvSharedLayers: t.numKvSharedLayers,
-    enableMoeBlock: t.enableMoeBlock,
-    numExperts: t.enableMoeBlock ? t.numExperts : 0,
-    topKExperts: t.enableMoeBlock ? t.topKExperts : 0,
-    hiddenSizePerLayerInput: t.hiddenSizePerLayerInput,
-    slidingWindow: t.slidingWindow,
-    finalLogitSoftcapping: t.finalLogitSoftcapping,
-    kvQuant: config.kvQuant ?? null,
-  });
-  const h = new Bun.CryptoHasher("sha256");
-  h.update(payload);
-  return h.digest("hex").slice(0, 16);
-}
-
-/** Exact immutable model-component identity, calculated once while its owner
- * holds the native execution lease. Names/shapes/dtypes disambiguate layouts. */
-export function tensorFingerprint(entries: Iterable<readonly [string,
-  Pick<import("../mlx/array").MlxArray, "dtype" | "shape" | "rawBytesView">]>): string {
-  const hash = new Bun.CryptoHasher("sha256");
-  for (const [name, tensor] of [...entries].sort(([a], [b]) => a.localeCompare(b))) {
-    hash.update(JSON.stringify([name, tensor.dtype, tensor.shape]));
-    hash.update(tensor.rawBytesView());
-  }
-  return hash.digest("hex");
-}