mlx-bun 0.0.11 → 0.0.13

package/src/jobs/runner.ts

@@ -149,6 +149,7 @@ interface QueuedSpawn {
   jobId: string;
   entry: string;
   bin: string;
+  spawn: typeof Bun.spawn;
   onComplete?: (jobId: string, code: number) => void;
 }
 
@@ -171,6 +172,8 @@ export interface SubprocessOpts {
   entry?: string;
   /** Override the runtime binary (tests). Defaults to "bun". */
   bin?: string;
+  /** Override process creation for deterministic failure-path tests. */
+  spawn?: typeof Bun.spawn;
   /** Called on the server (parent) after the child exits — used to invalidate
    *  caches (e.g. the Library) so a finished quantize surfaces immediately. */
   onComplete?: (jobId: string, code: number) => void;
@@ -192,35 +195,46 @@ export function submitSubprocess(
     jobId: row.id,
     entry: opts.entry ?? JOB_ENTRY_PATH,
     bin: opts.bin ?? "bun",
+    spawn: opts.spawn ?? Bun.spawn,
     onComplete: opts.onComplete,
   };
-  if (gpuLeaseHolder === null) {
-    spawnNow(item);
-  } else {
-    spawnQueue.push(item);
-  }
+  spawnQueue.push(item);
+  drainQueue();
   return { jobId: row.id, outputPath };
 }
 
 /** Spawn the child, acquire the lease, stream stdout/stderr into the log,
  *  reconcile terminal status on exit, then release the lease and drain. */
 function spawnNow(item: QueuedSpawn): void {
-  const { store, jobId, entry, bin } = item;
+  const { store, jobId, entry, bin, spawn } = item;
   gpuLeaseHolder = jobId;
 
   // The child opens its OWN JobStore over the SAME DB/logs — a sqlite
   // connection can't cross the process boundary, so we hand the paths via
   // env. A :memory: DB can't be shared with a child; subprocess jobs require
   // a file-backed DB.
-  const proc = Bun.spawn([bin, entry, jobId], {
-    stdout: "pipe",
-    stderr: "pipe",
-    env: {
-      ...process.env,
-      MLX_BUN_JOBS_DB: store.dbPath,
-      MLX_BUN_JOBS_DIR: store.logsDir,
-    },
-  });
+  let proc: Bun.Subprocess<"ignore", "pipe", "pipe">;
+  try {
+    proc = spawn([bin, entry, jobId], {
+      stdout: "pipe",
+      stderr: "pipe",
+      env: {
+        ...process.env,
+        MLX_BUN_JOBS_DB: store.dbPath,
+        MLX_BUN_JOBS_DIR: store.logsDir,
+      },
+    });
+  } catch (e) {
+    try {
+      store.setStatus(jobId, "failed", {
+        error: errString(e),
+        endedAt: nowIso(),
+      });
+    } finally {
+      releaseLease(jobId);
+    }
+    return;
+  }
 
   const row = store.get(jobId);
   const logPath = row?.log_path;
@@ -259,9 +273,14 @@ function releaseLease(jobId: string): void {
 
 /** Run the next queued subprocess job if the lease is free. */
 export function drainQueue(): void {
-  if (gpuLeaseHolder !== null) return;
-  const next = spawnQueue.shift();
-  if (next) spawnNow(next);
+  while (gpuLeaseHolder === null) {
+    const next = spawnQueue.shift();
+    if (!next) return;
+    // A synchronous spawn failure releases the lease. Keep draining
+    // iteratively so a run of bad queue entries cannot recurse or wedge the
+    // first startable job behind them.
+    spawnNow(next);
+  }
 }
 
 /** Read a child stream line-by-line, buffering partial trailing lines, and

package/src/kv-store.ts

@@ -50,6 +50,7 @@ import {
   QuantizedKVCache, RotatingQuantizedKVCache,
   TurboQuantKVCache, type TurboQuantTensor,
 } from "./model/gemma4-base";
+import { Glm52Cache, MLACache } from "./model/glm52-cache";
 import { SSMCache } from "./model/qwen3-delta";
 
 const MAGIC = "MLXBUNKV2\n";
@@ -62,7 +63,16 @@ const PREFIX_LEN = MAGIC.length + 4 + 4 + 8;
 // the header's byte length when the real values replace the placeholders.
 const hash64 = (bytes: Uint8Array): string => Bun.hash(bytes).toString(16).padStart(16, "0");
 
-export type CacheKind = "kv" | "rotating" | "qkv" | "rotating-qkv" | "ssm" | "turboquant";
+export type CacheKind =
+  | "kv"
+  | "rotating"
+  | "qkv"
+  | "rotating-qkv"
+  | "ssm"
+  | "turboquant"
+  | "mla"
+  | "mla-dsa"
+  | "mtp-mla";
 
 interface TensorSlot {
   off: number;
@@ -91,6 +101,11 @@ export interface CacheHeaderEntry {
   /** turboquant: head_dim, needed to unpack kIdx/vPacked on restore
    *  (packed byte width alone doesn't recover the original element count). */
   headDim?: number;
+  /** GLM checkpoint-native compressed-cache geometry. */
+  kvLoraRank?: number;
+  ropeHeadDim?: number;
+  dsaHeadDim?: number;
+  maxTokens?: number;
   /** kv/rotating: [k, v] · qkv/rotating-qkv: [kPacked, kScales, kBiases,
    *  vPacked, vScales, vBiases] · ssm: [conv, recurrent] · turboquant:
    *  [kIdx, kScales, kZeros, vPacked, vScales] */
@@ -142,6 +157,10 @@ function snapshotCache(c: Cache, dataOffset: number): { entry: CacheHeaderEntry;
     const [B, H, , D] = a.shape as [number, number, number, number];
     return a.slice([0, 0, 0, 0], [B, H, upTo, D]);
   };
+  const liveMlaSlice = (a: MlxArray, upTo: number): MlxArray => {
+    const [B, , D] = a.shape as [number, number, number];
+    return a.slice([0, 0, 0], [B, upTo, D]);
+  };
   const pushTriple = (t: ops.QuantizedTensor, upTo: number | null): void => {
     for (const a of [t.packed, t.scales, t.biases]) {
       if (upTo === null) push(a, false);
@@ -150,7 +169,26 @@ function snapshotCache(c: Cache, dataOffset: number): { entry: CacheHeaderEntry;
   };
 
   let entry: CacheHeaderEntry;
-  if (c instanceof TurboQuantKVCache) {
+  if (c instanceof Glm52Cache) {
+    if (!c.latent || !c.rope || c.batchSize === null || c.offset === 0)
+      throw new Error("cannot persist an empty GLM compressed cache");
+    push(liveMlaSlice(c.latent, c.offset), true);
+    push(liveMlaSlice(c.rope, c.offset), true);
+    if (c.dsa) {
+      if (!c.dsa.data || c.dsa.offset !== c.offset)
+        throw new Error("cannot persist misaligned GLM DSA state");
+      push(liveMlaSlice(c.dsa.data, c.offset), true);
+    }
+    entry = {
+      kind: c.role === "mtp" ? "mtp-mla" : c.dsa ? "mla-dsa" : "mla",
+      offset: c.offset,
+      kvLoraRank: c.kvLoraRank,
+      ropeHeadDim: c.ropeHeadDim,
+      ...(c.dsa ? { dsaHeadDim: c.dsa.headDim } : {}),
+      maxTokens: c.maxTokens,
+      tensors: slots,
+    };
+  } else if (c instanceof TurboQuantKVCache) {
     const t = c.state();
     if (t.length === 0) throw new Error("cannot persist an empty cache");
     const [kIdx, kScales, kZeros, vPacked, vScales] = t as [MlxArray, MlxArray, MlxArray, MlxArray, MlxArray];
@@ -211,6 +249,10 @@ export function cloneKvCaches(caches: Cache[]): Cache[] {
     const [B, H, , D] = a.shape as [number, number, number, number];
     return a.slice([0, 0, 0, 0], [B, H, upTo, D]);
   };
+  const mlaView = (a: MlxArray, upTo: number): MlxArray => {
+    const [B, , D] = a.shape as [number, number, number];
+    return a.slice([0, 0, 0], [B, upTo, D]);
+  };
   const tripleView = (t: ops.QuantizedTensor, upTo: number | null): ops.QuantizedTensor => ({
     packed: upTo === null ? view(t.packed) : liveView(t.packed, upTo),
     scales: upTo === null ? view(t.scales) : liveView(t.scales, upTo),
@@ -219,7 +261,24 @@ export function cloneKvCaches(caches: Cache[]): Cache[] {
   const out: Cache[] = [];
   try {
     for (const c of caches) {
-      if (c instanceof TurboQuantKVCache) {
+      if (c instanceof Glm52Cache) {
+        if (!c.latent || !c.rope || c.batchSize === null || c.offset === 0)
+          throw new Error("cannot clone an empty GLM compressed cache");
+        const n = new MLACache({
+          kvLoraRank: c.kvLoraRank,
+          ropeHeadDim: c.ropeHeadDim,
+          ...(c.dsa ? { dsa: { headDim: c.dsa.headDim } } : {}),
+          maxTokens: c.maxTokens,
+          role: c.role,
+        });
+        n.restoreCompressedState(
+          mlaView(c.latent, c.offset),
+          mlaView(c.rope, c.offset),
+          c.dsa ? mlaView(c.dsa.data!, c.offset) : null,
+          c.offset,
+        );
+        out.push(n);
+      } else if (c instanceof TurboQuantKVCache) {
         const state = c.state();
         if (state.length === 0) throw new Error("cannot clone an empty cache");
         const [kIdx, kScales, kZeros, vPacked, vScales] = state;
@@ -514,6 +573,7 @@ export interface LoadedKvCache {
 
 export interface KvLoadExpect {
   /** Reject on metadata mismatch (pass what the server is running). */
+  modelId?: string;
   configFingerprint?: string;
   tokenizerHash?: string;
   ns?: string;
@@ -522,6 +582,50 @@ export interface KvLoadExpect {
   verify?: boolean;
 }
 
+/** Reject GLM cache/layout drift before opening the tensor mmap. Generic KV
+ *  kinds may legitimately differ from makeCache() after runtime quantization;
+ *  GLM's checkpoint-native cache geometry does not. */
+function validateGlm52Prototype(
+  path: string,
+  entry: CacheHeaderEntry,
+  prototype: Cache,
+): void {
+  const glmKind = entry.kind === "mla" || entry.kind === "mla-dsa" ||
+    entry.kind === "mtp-mla";
+  if (!(prototype instanceof Glm52Cache)) {
+    if (glmKind)
+      throw new Error(`${path}: GLM cache kind ${entry.kind} does not match model cache`);
+    return;
+  }
+  const expectedKind: CacheKind = prototype.role === "mtp"
+    ? "mtp-mla"
+    : prototype.dsa ? "mla-dsa" : "mla";
+  if (entry.kind !== expectedKind) {
+    throw new Error(
+      `${path}: cache kind ${entry.kind} != model cache kind ${expectedKind}`,
+    );
+  }
+  for (const [field, actual, expected] of [
+    ["kvLoraRank", entry.kvLoraRank, prototype.kvLoraRank],
+    ["ropeHeadDim", entry.ropeHeadDim, prototype.ropeHeadDim],
+    ["dsaHeadDim", entry.dsaHeadDim, prototype.dsa?.headDim],
+    ["maxTokens", entry.maxTokens, prototype.maxTokens],
+  ] as const) {
+    if (actual !== expected) {
+      throw new Error(
+        `${path}: ${field} ${String(actual)} != model ${String(expected)}`,
+      );
+    }
+  }
+  const expectedTensors = prototype.dsa ? 3 : 2;
+  if (entry.tensors.length !== expectedTensors) {
+    throw new Error(
+      `${path}: ${entry.kind} has ${entry.tensors.length} tensors, ` +
+      `expected ${expectedTensors}`,
+    );
+  }
+}
+
 /** Copy-restore a full-attention KV tensor into STEP-rounded capacity with
  *  ≥1 token of slack (what a live mid-generation cache looks like) so the
  *  first post-restore updateAndFetch never takes the grow path — the old
@@ -569,15 +673,28 @@ export function loadKvCache(
   expect: KvLoadExpect = {},
 ): LoadedKvCache {
   const header = readKvHeader(path);
-  for (const key of ["configFingerprint", "tokenizerHash", "ns"] as const) {
+  for (const key of [
+    "modelId",
+    "configFingerprint",
+    "tokenizerHash",
+    "ns",
+  ] as const) {
     if (expect[key] !== undefined && header[key] !== expect[key])
       throw new Error(`${path}: ${key} mismatch (file ${header[key]}, expected ${expect[key]})`);
   }
   const proto = model.makeCache();
-  const cacheCount = proto.length;
-  for (const c of proto) c.dispose();
-  if (header.caches.length !== cacheCount)
-    throw new Error(`${path}: ${header.caches.length} cached layers but model has ${cacheCount}`);
+  try {
+    if (header.caches.length !== proto.length) {
+      throw new Error(
+        `${path}: ${header.caches.length} cached layers but model has ${proto.length}`,
+      );
+    }
+    for (let index = 0; index < proto.length; index++) {
+      validateGlm52Prototype(path, header.caches[index]!, proto[index]!);
+    }
+  } finally {
+    for (const c of proto) c.dispose();
+  }
 
   const mmap = MmapFile.open(path, "ro");
   const dataStart = header.dataStart;
@@ -663,6 +780,28 @@ export function loadKvCache(
           caches.push(c);
           break;
         }
+        case "mla":
+        case "mla-dsa":
+        case "mtp-mla": {
+          const hasDsa = e.kind === "mla-dsa";
+          if (e.kind === "mtp-mla" && e.dsaHeadDim !== undefined)
+            throw new Error(`${path}: native MTP cache cannot contain DSA state`);
+          const c = new MLACache({
+            kvLoraRank: e.kvLoraRank!,
+            ropeHeadDim: e.ropeHeadDim!,
+            ...(hasDsa ? { dsa: { headDim: e.dsaHeadDim! } } : {}),
+            maxTokens: e.maxTokens,
+            role: e.kind === "mtp-mla" ? "mtp" : "target",
+          });
+          c.restoreCompressedState(
+            arr(t[0]!),
+            arr(t[1]!),
+            hasDsa ? arr(t[2]!) : null,
+            e.offset,
+          );
+          caches.push(c);
+          break;
+        }
         default:
           throw new Error(`${path}: unknown cache kind ${(e as { kind: string }).kind}`);
       }

package/src/lora.ts

@@ -15,6 +15,7 @@ import { resolve } from "node:path";
 import { ptr, read } from "bun:ffi";
 import { MlxArray, cpuStream } from "./mlx/array";
 import { C } from "./mlx/ffi";
+import { transposeAxes } from "./mlx/ops";
 import { SafetensorsFile } from "./safetensors";
 import type { LoraWeights } from "./model/gemma4";
 import type { RuntimeModel } from "./model/factory";
@@ -80,21 +81,26 @@ export function loadAdapterTensors(file: string): Map<string, MlxArray> {
 }
 
 /** Adapter scale from its config: mlx-lm writes lora_parameters.scale;
- *  PEFT writes lora_alpha + r (scale = alpha / r). Reference: mount.py.
- *  `rsLora` (recorded at train time) means the effective per-layer scale is
- *  α/√rank — the caller divides by √(that layer's rank). */
+ *  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 ?? false);
+      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: r ? alpha / r : 1.0, rank: r || null, rsLora };
+      return {
+        scale: rsLora ? alpha : r ? alpha / r : 1.0,
+        rank: r || null,
+        rsLora,
+      };
     }
   }
   throw new Error(`no adapter_config.json in ${dir}`);
@@ -215,7 +221,12 @@ export class AdapterManager {
 
     // Group adapter tensors into (modulePath → {a, b}), probing both the
     // pure-LLM and VLM-wrapped prefixes like the reference _find_weight_pair.
-    const pairs = new Map<string, { a?: MlxArray; b?: MlxArray }>();
+    const pairs = new Map<string, {
+      a?: MlxArray;
+      b?: MlxArray;
+      aIsPeft?: boolean;
+      bIsPeft?: boolean;
+    }>();
     let skipped = 0;
     const ourPrefix = `${this.#model.prefixBase}.layers.`;
     const altPrefix = ourPrefix.startsWith("language_model.")
@@ -225,11 +236,23 @@ export class AdapterManager {
       const m = name.match(/^(.*)\.(lora_a|lora_A|lora_b|lora_B)(\.weight)?$/);
       if (!m) { skipped++; arr.dispose(); continue; }
       let modulePath = m[1]!;
+      // Standard PEFT state dicts wrap the model path in
+      // `base_model.model.`. The remainder is the underlying HF module path
+      // (`model.layers...` or `language_model.model.layers...`).
+      if (modulePath.startsWith("base_model.model."))
+        modulePath = modulePath.slice("base_model.model.".length);
       if (modulePath.startsWith(altPrefix))
         modulePath = ourPrefix + modulePath.slice(altPrefix.length);
       if (!targets.has(modulePath)) { skipped++; arr.dispose(); continue; }
       const slot = pairs.get(modulePath) ?? {};
-      if (m[2]!.toLowerCase() === "lora_a") slot.a = arr; else slot.b = arr;
+      const peftName = m[3] === ".weight" || m[2] === "lora_A" || m[2] === "lora_B";
+      if (m[2]!.toLowerCase() === "lora_a") {
+        slot.a = arr;
+        slot.aIsPeft = peftName;
+      } else {
+        slot.b = arr;
+        slot.bIsPeft = peftName;
+      }
       pairs.set(modulePath, slot);
     }
 
@@ -240,21 +263,61 @@ export class AdapterManager {
       for (const { a, b } of pairs.values()) { a?.dispose(); b?.dispose(); }
     };
     try {
-      for (const [modulePath, { a, b }] of pairs) {
+      for (const [modulePath, pair] of pairs) {
+        const { a, b } = pair;
         if (!a || !b)
           throw new Error(`${modulePath}: adapter has only one of lora_a/lora_b`);
         const linear = targets.get(modulePath)!;
-        const [aIn, rank] = a.shape as [number, number];
-        const [bRank, bOut] = b.shape as [number, number];
-        if (aIn !== linear.inFeatures || bRank !== rank || bOut !== linear.outFeatures)
+        if (a.shape.length !== 2 || b.shape.length !== 2)
+          throw new Error(
+            `${modulePath}: LoRA tensors must be rank-2; got ` +
+            `lora_a [${a.shape}] / lora_b [${b.shape}]`,
+          );
+
+        const internalLayout =
+          a.shape[0] === linear.inFeatures &&
+          b.shape[0] === a.shape[1] &&
+          b.shape[1] === linear.outFeatures;
+        const peftLayout =
+          a.shape[1] === linear.inFeatures &&
+          b.shape[1] === a.shape[0] &&
+          b.shape[0] === linear.outFeatures;
+        if (!internalLayout && !peftLayout)
           throw new Error(
             `${modulePath}: shape mismatch — lora_a [${a.shape}] / lora_b [${b.shape}] ` +
-            `vs base [in ${linear.inFeatures}, out ${linear.outFeatures}]; ` +
-            `was this adapter trained for a different base model?`,
+            `vs base [in ${linear.inFeatures}, out ${linear.outFeatures}]. Expected ` +
+            `mlx-lm [in, rank] + [rank, out] or PEFT [rank, in] + [out, rank].`,
           );
+
+        // Full-rank square adapters satisfy both shape predicates. In that
+        // ambiguous case the PEFT spelling is authoritative; otherwise the
+        // unique compatible layout decides.
+        const usePeftLayout =
+          peftLayout && (!internalLayout || (pair.aIsPeft === true && pair.bIsPeft === true));
+        let mountedA = a;
+        let mountedB = b;
+        if (usePeftLayout) {
+          let transposedA: MlxArray | null = null;
+          let transposedB: MlxArray | null = null;
+          try {
+            transposedA = transposeAxes(a, [1, 0]);
+            transposedB = transposeAxes(b, [1, 0]);
+          } catch (e) {
+            transposedA?.dispose();
+            transposedB?.dispose();
+            throw e;
+          }
+          mountedA = transposedA;
+          mountedB = transposedB;
+          pair.a = mountedA;
+          pair.b = mountedB;
+          a.dispose();
+          b.dispose();
+        }
+        const rank = mountedA.shape[1]!;
         // rsLoRA: effective per-layer scale is α/√rank (matches training).
         const effScale = rsLora ? scale / Math.sqrt(rank) : scale;
-        validated.push({ linear, lw: { a, b, scale: effScale, rank } });
+        validated.push({ linear, lw: { a: mountedA, b: mountedB, scale: effScale, rank } });
       }
       if (validated.length === 0)
         throw new Error(