mlx-bun 0.0.11 → 0.0.12

package/src/mlx/custom-vjp.ts

@@ -15,6 +15,24 @@ import { JSCallback, ptr, read } from "bun:ffi";
 import { C, takeMlxError } from "./ffi";
 import { MlxArray } from "./array";
 
+type CallbackHandle = Pick<JSCallback, "ptr" | "close">;
+type NativeCall<F extends (...args: any[]) => any> =
+  (...args: Parameters<F>) => ReturnType<F>;
+
+/** Constructor-only dependency seam for deterministic native failure tests. */
+export interface CustomVjpConstructorHooks {
+  callback?: (
+    callback: ConstructorParameters<typeof JSCallback>[0],
+    definition: ConstructorParameters<typeof JSCallback>[1],
+  ) => CallbackHandle;
+  forwardClosureNew?: NativeCall<typeof C.mlx_closure_new_func_payload>;
+  vjpClosureNew?: NativeCall<typeof C.mlx_closure_custom_new_func_payload>;
+  transformedNew?: NativeCall<typeof C.mlx_closure_new>;
+  transform?: NativeCall<typeof C.mlx_custom_vjp>;
+  closureFree?: NativeCall<typeof C.mlx_closure_free>;
+  customClosureFree?: NativeCall<typeof C.mlx_closure_custom_free>;
+}
+
 function readVec(vec: bigint): MlxArray[] {
   const n = Number(C.mlx_vector_array_size(vec));
   const out: MlxArray[] = [];
@@ -41,19 +59,30 @@ function setVec(outPtr: number, arrays: MlxArray[]): void {
  *  pure functions of their arguments — the vjp re-reads everything from the
  *  graph during the enclosing backward. */
 export class CustomVjp {
-  readonly #fwdCb: JSCallback;
-  readonly #vjpCb: JSCallback;
+  readonly #fwdCb: CallbackHandle;
+  readonly #vjpCb: CallbackHandle;
   readonly #fwdClosure: bigint;
   readonly #vjpClosure: bigint;
   readonly #combined: bigint;
+  readonly #freeClosure: NativeCall<typeof C.mlx_closure_free>;
+  readonly #freeCustomClosure: NativeCall<typeof C.mlx_closure_custom_free>;
   #error: string | null = null;
   #disposed = false;
 
   constructor(
     fwd: (inputs: MlxArray[]) => MlxArray[],
     vjp: (primals: MlxArray[], cotangents: MlxArray[], outputs: MlxArray[]) => MlxArray[],
+    hooks: CustomVjpConstructorHooks = {},
   ) {
-    this.#fwdCb = new JSCallback(
+    const makeCallback = hooks.callback ?? ((callback, definition) => new JSCallback(callback, definition));
+    const forwardClosureNew = hooks.forwardClosureNew ?? C.mlx_closure_new_func_payload;
+    const vjpClosureNew = hooks.vjpClosureNew ?? C.mlx_closure_custom_new_func_payload;
+    const transformedNew = hooks.transformedNew ?? C.mlx_closure_new;
+    const transform = hooks.transform ?? C.mlx_custom_vjp;
+    this.#freeClosure = hooks.closureFree ?? C.mlx_closure_free;
+    this.#freeCustomClosure = hooks.customClosureFree ?? C.mlx_closure_custom_free;
+
+    const fwdCb = makeCallback(
       (outPtr: number, inVec: bigint, _payload: number): number => {
         try {
           const inputs = readVec(inVec);
@@ -70,37 +99,56 @@ export class CustomVjp {
       { args: ["ptr", "u64", "ptr"], returns: "i32" },
     );
 
-    this.#vjpCb = new JSCallback(
-      (resPtr: number, primalsVec: bigint, cotsVec: bigint, outsVec: bigint, _payload: number): number => {
-        try {
-          const primals = readVec(primalsVec);
-          const cots = readVec(cotsVec);
-          const outs = readVec(outsVec);
-          const grads = vjp(primals, cots, outs);
-          setVec(resPtr, grads);
-          for (const a of [...primals, ...cots, ...outs]) a.dispose();
-          for (const g of grads) g.dispose();
-          return 0;
-        } catch (e) {
-          this.#error = e instanceof Error ? e.message : String(e);
-          return 1;
-        }
-      },
-      { args: ["ptr", "u64", "u64", "u64", "ptr"], returns: "i32" },
-    );
-
-    this.#fwdClosure = C.mlx_closure_new_func_payload(this.#fwdCb.ptr as never, null, null);
-    this.#vjpClosure = C.mlx_closure_custom_new_func_payload(this.#vjpCb.ptr as never, null, null);
+    let vjpCb: CallbackHandle;
+    try {
+      vjpCb = makeCallback(
+        (resPtr: number, primalsVec: bigint, cotsVec: bigint, outsVec: bigint, _payload: number): number => {
+          try {
+            const primals = readVec(primalsVec);
+            const cots = readVec(cotsVec);
+            const outs = readVec(outsVec);
+            const grads = vjp(primals, cots, outs);
+            setVec(resPtr, grads);
+            for (const a of [...primals, ...cots, ...outs]) a.dispose();
+            for (const g of grads) g.dispose();
+            return 0;
+          } catch (e) {
+            this.#error = e instanceof Error ? e.message : String(e);
+            return 1;
+          }
+        },
+        { args: ["ptr", "u64", "u64", "u64", "ptr"], returns: "i32" },
+      );
+    } catch (e) {
+      // Callback construction is itself fallible. The native-closure guard
+      // below starts only after both callbacks exist, so this partial state
+      // must release the already-created forward callback here.
+      fwdCb.close();
+      throw e;
+    }
 
-    const slot = new BigUint64Array([C.mlx_closure_new()]);
-    if (C.mlx_custom_vjp(ptr(slot), this.#fwdClosure, this.#vjpClosure) !== 0) {
-      C.mlx_closure_free(this.#fwdClosure);
-      C.mlx_closure_custom_free(this.#vjpClosure);
-      this.#fwdCb.close();
-      this.#vjpCb.close();
-      throw new Error(`mlx_custom_vjp failed: ${takeMlxError() ?? ""}`);
+    let fwdClosure: bigint | null = null;
+    let vjpClosure: bigint | null = null;
+    let slot: BigUint64Array | null = null;
+    try {
+      fwdClosure = forwardClosureNew(fwdCb.ptr as never, null, null);
+      vjpClosure = vjpClosureNew(vjpCb.ptr as never, null, null);
+      slot = new BigUint64Array([transformedNew()]);
+      if (transform(ptr(slot), fwdClosure, vjpClosure) !== 0)
+        throw new Error(`mlx_custom_vjp failed: ${takeMlxError() ?? ""}`);
+      this.#fwdCb = fwdCb;
+      this.#vjpCb = vjpCb;
+      this.#fwdClosure = fwdClosure;
+      this.#vjpClosure = vjpClosure;
+      this.#combined = read.u64(ptr(slot), 0);
+    } catch (e) {
+      if (slot) this.#freeClosure(read.u64(ptr(slot), 0));
+      if (fwdClosure !== null) this.#freeClosure(fwdClosure);
+      if (vjpClosure !== null) this.#freeCustomClosure(vjpClosure);
+      fwdCb.close();
+      vjpCb.close();
+      throw e;
     }
-    this.#combined = read.u64(ptr(slot), 0);
   }
 
   apply(inputs: MlxArray[]): MlxArray[] {
@@ -127,9 +175,9 @@ export class CustomVjp {
   dispose(): void {
     if (this.#disposed) return;
     this.#disposed = true;
-    C.mlx_closure_free(this.#combined);
-    C.mlx_closure_free(this.#fwdClosure);
-    C.mlx_closure_custom_free(this.#vjpClosure);
+    this.#freeClosure(this.#combined);
+    this.#freeClosure(this.#fwdClosure);
+    this.#freeCustomClosure(this.#vjpClosure);
     this.#fwdCb.close();
     this.#vjpCb.close();
   }

package/src/model/factory.ts

@@ -4,8 +4,9 @@
 // Gemma4Model and only override forwardLayers (with their own
 // cache-signature guard), so the choice is always safe.
 
-import type { ModelConfig } from "../config";
-import type { Weights } from "../weights";
+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";
@@ -14,19 +15,109 @@ 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 { genericArgsFor, GENERIC_MODEL_TYPES, remapModelType } from "./universal/archs";
-import { isDiffusionGemmaConfig, isMiniCPM5Config, isQwen35Config, isQwen3Config, isQwen3MoeConfig } from "./support";
+import { isDiffusionGemmaConfig, isGlm52Config, isMiniCPM5Config, isQwen35Config, isQwen3Config, isQwen3MoeConfig } from "./support";
 
 export type RuntimeModel =
   | Gemma4Model | MiniCPM5Model | Qwen35Model | Qwen3Model | Qwen3MoeModel
-  | DiffusionGemmaModel | UniversalDenseModel;
+  | DiffusionGemmaModel | Glm52Model | UniversalDenseModel;
+
+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 };
+}
+
+/**
+ * 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 async function openModel(
+  modelDir: string,
+  options: Glm52RuntimeOpenOptions = {},
+): Promise<RuntimeModel> {
+  const config = await loadModelConfig(modelDir);
+  if (isGlm52Config(config))
+    return (await openGlm52RuntimeModel(modelDir, options)).model;
+  return createModel(await Weights.open(modelDir), config);
+}
 
 /** Dispatch ladder (docs/design/generic-model-support.md §3.3):
  *  dedicated → generated/monolith (gemma4*) → generic (Tier-0 universal)
  *  → reject with a helpful error. Generic never shadows a dedicated port. */
 export function createModel(weights: Weights, config: ModelConfig): RuntimeModel {
   // 1. Dedicated classes.
+  if (isGlm52Config(config)) {
+    throw new Error(
+      "glm_moe_dsa uses the direct Colibri container; construct it with " +
+      "openModel(modelDir) or Glm52Model.open(modelDir)",
+    );
+  }
   // DiffusionGemma is non-autoregressive — generate() detects it and routes to
   // the denoising engine instead of the AR decode loop. It exposes the shared
   // RuntimeModel surface (config/weightsBytes/loraState/makeCache) with the
@@ -54,7 +145,7 @@ export function createModel(weights: Weights, config: ModelConfig): RuntimeModel
   throw new Error(
     `unsupported model_type "${config.modelType}"` +
     (arch !== config.modelType ? ` (mlx-lm remaps it to "${arch}")` : "") +
-    ` — targeted: gemma4*, diffusion_gemma, qwen3_5, qwen3, qwen3_moe, MiniCPM5;` +
+    ` — targeted: gemma4*, diffusion_gemma, glm_moe_dsa, qwen3_5, qwen3, qwen3_moe, MiniCPM5;` +
     ` generic (Tier-0): ${[...GENERIC_MODEL_TYPES].sort().join(", ")}`,
   );
 }

package/src/model/gemma4-base.ts

@@ -234,6 +234,38 @@ export interface Cache {
   dispose(): void;
 }
 
+/**
+ * Capability contract for cache layouts that own their dynamic-row batching.
+ * The scheduler uses this for non-K/V state families (for example GLM's
+ * compressed MLA/DSA state) instead of teaching the scheduler their tensor
+ * layout. `mergeRows` writes into an empty cache and does not consume inputs;
+ * extracted rows are independent owned caches.
+ */
+export interface BatchableCache extends Cache {
+  readonly batchSize: number | null;
+  readonly rowOffsets: readonly number[];
+  readonly leftPad: readonly number[];
+  readonly maxTokens?: number;
+  makeEmptyBatch(): BatchableCache;
+  mergeRows(rows: readonly Cache[]): void;
+  extractRow(row: number): Cache;
+  filterRows(keep: readonly number[]): void;
+  projectedBytes(tokens: number): number;
+}
+
+export function isBatchableCache(cache: Cache): cache is BatchableCache {
+  const candidate = cache as Partial<BatchableCache>;
+  return (
+    typeof candidate.makeEmptyBatch === "function" &&
+    typeof candidate.mergeRows === "function" &&
+    typeof candidate.extractRow === "function" &&
+    typeof candidate.filterRows === "function" &&
+    typeof candidate.projectedBytes === "function" &&
+    Array.isArray(candidate.rowOffsets) &&
+    Array.isArray(candidate.leftPad)
+  );
+}
+
 /** How one decode step interacts with a cache under compiled decode
  *  (returned by prepareDecodeStep, consumed by compiled-decode.ts):
  *  - "concat": the compiled graph fetches concat(active prefix, new kv)