mlx-bun 0.0.8 → 0.0.10

package/src/config.ts

@@ -186,6 +186,9 @@ export async function loadModelConfig(modelDir: string): Promise<ModelConfig> {
   const t = isDiffusion ? { ...diffusionGemmaRawDefaults(), ...baseT } : baseT;
   const isLlama = modelType === "llama";
   const isQwen35 = typeof modelType === "string" && modelType.startsWith("qwen3_5");
+  // Plain Qwen3 (Qwen3ForCausalLM, e.g. Qwen3-Embedding): a flat HF config with
+  // a scalar rope_theta and no rope_parameters map — handled like llama below.
+  const isQwen3 = modelType === "qwen3";
   // Qwen3.5 rope_parameters is a flat dict ({type, rope_theta, mrope_section,
   // partial_rotary_factor}), not the gemma per-attention-type map — and
   // type "default" means plain partial nn.RoPE (mrope_section ignored for text).
@@ -234,7 +237,7 @@ export async function loadModelConfig(modelDir: string): Promise<ModelConfig> {
         }
       : t.rope_parameters
         ? parseRope(t.rope_parameters)
-        : isLlama
+        : isLlama || isQwen3
           ? {
               full_attention: {
                 ropeTheta: t.rope_theta ?? 10000,

package/src/download.ts

@@ -15,7 +15,10 @@
 // Files download sequentially (decode is bandwidth-bound, the network
 // rarely is; resumability matters more than parallelism here).
 
-import { mkdirSync, existsSync, statSync, createWriteStream, renameSync, rmSync, symlinkSync, readFileSync } from "node:fs";
+import {
+  mkdirSync, existsSync, statSync, createWriteStream, renameSync, rmSync,
+  symlinkSync, readFileSync, openSync, closeSync, readSync, writeSync,
+} from "node:fs";
 import { dirname, join } from "node:path";
 import { DEFAULT_HUB } from "./registry";
 
@@ -54,6 +57,63 @@ export function hfToken(): string | null {
   }
 }
 
+/** 401/403 almost always means a gated/private repo, not a broken network —
+ *  say so instead of leaving the user staring at a bare status code. */
+function authHint(status: number): string {
+  return status === 401 || status === 403
+    ? " — gated repo? run `hf auth login` or set HF_TOKEN"
+    : "";
+}
+
+// --- download lockfile -----------------------------------------------------
+// Two writers on one blob (a foreground `mlx-bun get` racing the server's
+// background auto-download) both append to the same .incomplete file and
+// corrupt it. An O_EXCL .lock beside the .incomplete serializes them; a lock
+// older than LOCK_STALE_MS (crashed owner) or whose pid is gone is stolen.
+
+const LOCK_STALE_MS = 60 * 60 * 1000; // ~1 h: longer than any sane blob fetch
+
+function lockIsStale(lockPath: string): boolean {
+  try {
+    const st = statSync(lockPath);
+    if (Date.now() - st.mtimeMs > LOCK_STALE_MS) return true;
+    const pid = Number(readFileSync(lockPath, "utf8").split("\n")[0]);
+    if (Number.isFinite(pid) && pid > 0) {
+      try { process.kill(pid, 0); } catch { return true; } // owner is gone
+    }
+    return false;
+  } catch {
+    return true; // unreadable/vanished — treat as stale, O_EXCL still arbitrates
+  }
+}
+
+/** Take `<blobPath>.lock` (O_EXCL). Returns a release fn, or throws the
+ *  friendly "another download in progress" error when live-held. */
+function acquireLock(blobPath: string, name: string): () => void {
+  const lockPath = `${blobPath}.lock`;
+  for (let attempt = 0; attempt < 2; attempt++) {
+    try {
+      const fd = openSync(lockPath, "wx");
+      writeSync(fd, `${process.pid}\n${new Date().toISOString()}\n`);
+      closeSync(fd);
+      return () => { try { rmSync(lockPath, { force: true }); } catch {} };
+    } catch (e) {
+      if ((e as NodeJS.ErrnoException).code !== "EEXIST") throw e;
+      if (lockIsStale(lockPath)) {
+        try { rmSync(lockPath, { force: true }); } catch {}
+        continue; // one retry after stealing
+      }
+      let holder = "";
+      try { holder = ` (pid ${readFileSync(lockPath, "utf8").split("\n")[0]})`; } catch {}
+      throw new Error(
+        `another download of ${name} is in progress${holder} — ` +
+        `wait for it to finish, or delete ${lockPath} if you're sure it isn't`,
+      );
+    }
+  }
+  throw new Error(`could not acquire ${lockPath}`);
+}
+
 /** git blob identity: sha1("blob <size>\0" + content) — what blobId is. */
 export function gitBlobSha1(bytes: Uint8Array): string {
   const hasher = new Bun.CryptoHasher("sha1");
@@ -73,7 +133,10 @@ export async function listRepoFiles(
     headers: token ? { authorization: `Bearer ${token}` } : {},
   });
   if (!res.ok)
-    throw new Error(`HF API ${res.status} for ${repoId}@${revision}: ${(await res.text()).slice(0, 200)}`);
+    throw new Error(
+      `HF API ${res.status} for ${repoId}@${revision}: ${(await res.text()).slice(0, 200)}` +
+      authHint(res.status),
+    );
   const body = (await res.json()) as {
     sha: string;
     siblings: { rfilename: string; size?: number; blobId?: string; lfs?: { oid: string; size: number } }[];
@@ -115,6 +178,21 @@ export async function downloadOne(
   url: string, token: string | null, blobPath: string,
   expected: { size: number; sha256?: string; sha1?: string; name: string },
   onProgress?: DownloadOptions["onProgress"],
+): Promise<void> {
+  // Serialize concurrent writers (foreground get vs background auto-download)
+  // on this blob — both appending to one .incomplete corrupts it.
+  const release = acquireLock(blobPath, expected.name);
+  try {
+    await downloadOneLocked(url, token, blobPath, expected, onProgress);
+  } finally {
+    release();
+  }
+}
+
+async function downloadOneLocked(
+  url: string, token: string | null, blobPath: string,
+  expected: { size: number; sha256?: string; sha1?: string; name: string },
+  onProgress?: DownloadOptions["onProgress"],
 ): Promise<void> {
   const partPath = `${blobPath}.incomplete`;
   let offset = existsSync(partPath) ? statSync(partPath).size : 0;
@@ -124,14 +202,27 @@ export async function downloadOne(
   }
 
   // Streaming verification — resume re-hashes the existing prefix so the
-  // final digest always covers every byte on disk.
+  // final digest always covers every byte on disk. Chunked reads: a resumed
+  // multi-GB blob must not be materialized as one allocation.
   const sha256 = new Bun.CryptoHasher("sha256");
   const sha1 = new Bun.CryptoHasher("sha1");
   sha1.update(`blob ${expected.size}\0`);
   if (offset > 0) {
-    const existing = await Bun.file(partPath).arrayBuffer();
-    sha256.update(existing);
-    sha1.update(existing);
+    const fd = openSync(partPath, "r");
+    try {
+      const buf = Buffer.alloc(8 * 2 ** 20);
+      let pos = 0;
+      while (pos < offset) {
+        const n = readSync(fd, buf, 0, Math.min(buf.length, offset - pos), pos);
+        if (n <= 0) break;
+        const chunk = buf.subarray(0, n);
+        sha256.update(chunk);
+        sha1.update(chunk);
+        pos += n;
+      }
+    } finally {
+      closeSync(fd);
+    }
   }
 
   if (offset < expected.size) {
@@ -139,10 +230,10 @@ export async function downloadOne(
     if (res.status === 200 && offset > 0) {
       // server ignored Range — start the file (and hashes) over
       rmSync(partPath);
-      return downloadOne(url, token, blobPath, expected, onProgress);
+      return downloadOneLocked(url, token, blobPath, expected, onProgress);
     }
     if (res.status !== 200 && res.status !== 206)
-      throw new Error(`download failed (${res.status}) for ${expected.name}`);
+      throw new Error(`download failed (${res.status}) for ${expected.name}${authHint(res.status)}`);
     const out = createWriteStream(partPath, { flags: offset > 0 ? "a" : "w" });
     try {
       for await (const chunk of res.body as ReadableStream<Uint8Array>) {

package/src/embed.ts

@@ -0,0 +1,77 @@
+// Text-embedding entry point. Wraps Qwen3Model.embedPooled (last-token hidden →
+// L2-normalized vector) with the Qwen3-Embedding input convention so the CLI
+// (`mlx-bun embed`), the server (`/v1/embeddings`), and in-repo experiments all
+// produce the SAME vectors — the ones verified bit-exact vs mlx-lm in
+// tests/qwen3-embed-parity.test.ts.
+
+import { Qwen3Model } from "./model/qwen3";
+import type { RuntimeModel } from "./model/factory";
+import type { LoadedTokenizer } from "./tokenizer";
+import * as ops from "./mlx/ops";
+
+/** <|endoftext|> — the token Qwen3-Embedding terminates every input with and
+ *  pools the last position of. */
+export const EMBED_EOD = 151643;
+
+/** Only plain-Qwen3 (the Qwen3-Embedding backbone) exposes a pooled-embedding
+ *  path today. Narrow the served RuntimeModel before embedding. */
+export function isEmbeddingModel(model: RuntimeModel): model is Qwen3Model {
+  return model instanceof Qwen3Model;
+}
+
+/** Qwen3-Embedding query format: a task instruction steers WHICH similarity axis
+ *  the geometry reflects ("represent this for topic routing" vs "...for
+ *  sentiment"). Documents are embedded raw; only queries get an instruction. */
+export function withInstruction(text: string, instruction?: string): string {
+  return instruction ? `Instruct: ${instruction}\nQuery:${text}` : text;
+}
+
+/** Tripwire: counts embedOne invocations (embedMany accrues via embedOne) so
+ *  callers can assert the pipeline embeds exactly as many times as expected. */
+let embedCallCount = 0;
+
+/** Reset the embed tripwire counter to 0. */
+export function resetEmbedCounter(): void {
+  embedCallCount = 0;
+}
+
+/** Read the embed tripwire counter (total embedOne calls since last reset). */
+export function getEmbedCounter(): number {
+  return embedCallCount;
+}
+
+export interface EmbedResult {
+  /** L2-normalized embedding (length = model hidden size). */
+  vector: Float32Array;
+  /** Token count of the embedded input (incl. the EOD pooling token). */
+  tokens: number;
+}
+
+/** Embed one text. `instruction` (optional) applies the query format. */
+export function embedOne(
+  model: Qwen3Model,
+  tok: LoadedTokenizer,
+  text: string,
+  instruction?: string,
+): EmbedResult {
+  embedCallCount++;
+  // addSpecialTokens=false matches the bit-exact parity path (Qwen3 has no BOS;
+  // we append the EOD pooling token ourselves).
+  const ids = [...tok.encode(withInstruction(text, instruction), false), EMBED_EOD];
+  const idArr = ops.fromInt32(ids, [1, ids.length]);
+  const vec = model.embedPooled(idArr);
+  idArr.dispose();
+  const out = vec.toFloat32();
+  vec.dispose();
+  return { vector: out, tokens: ids.length };
+}
+
+/** Embed many texts (one forward each — the runtime is single-sequence). */
+export function embedMany(
+  model: Qwen3Model,
+  tok: LoadedTokenizer,
+  texts: string[],
+  instruction?: string,
+): EmbedResult[] {
+  return texts.map((t) => embedOne(model, tok, t, instruction));
+}

package/src/eval/perplexity.ts

@@ -0,0 +1,164 @@
+// Perplexity over a local text/JSONL dataset — mlx_lm.perplexity methodology.
+//
+// The reference (mlx_lm/perplexity.py, mlx-lm 0.31.3) does:
+//   1. tokenize dataset samples, visit them in a seeded random permutation,
+//   2. concatenate tokens into one stream and cut it into NON-OVERLAPPING
+//      rows of `sequence_length` (no sliding-window stride, no BOS re-seeding
+//      per row), keeping at most `num_samples` rows,
+//   3. per batch: logits = model(batch[:, :-1]) in f32, per-token
+//      cross-entropy vs batch[:, 1:] — EVERY position counts,
+//   4. ppl = exp(mean loss); SE via the delta method:
+//      se_ppl = ppl · std(losses, ddof=1) / √N.
+//
+// We reproduce 2-4 exactly. Step 1 differs only in the data source: the
+// reference pulls a Hugging Face dataset (network); we take a LOCAL .txt or
+// .jsonl file — never a download. The permutation RNG is a seeded
+// deterministic shuffle (not NumPy's), so row selection is reproducible per
+// seed but not bit-matched to the Python tool's sampling.
+//
+// The forward reuses the existing full-sequence machinery
+// (src/train/forward.ts trainForward → model.forwardHidden + logitsFromHidden).
+
+import { clearCache, Dtype } from "../mlx/ffi";
+import * as ops from "../mlx/ops";
+import { trainForward } from "../train/forward";
+import type { RuntimeModel } from "../model/factory";
+
+export interface PplDataOptions {
+  sequenceLength: number;
+  /** Max rows to keep (-1 / 0 = all available). */
+  numSamples: number;
+  /** Seed for the sample-order shuffle (reference default: 123). */
+  seed: number;
+}
+
+export interface PplResult {
+  ppl: number;
+  /** Delta-method standard error of the perplexity. */
+  standardError: number;
+  meanLoss: number;
+  tokens: number;
+  rows: number;
+}
+
+/** Deterministic PRNG (mulberry32) for the seeded sample shuffle. */
+function mulberry32(seed: number): () => number {
+  let s = seed >>> 0;
+  return () => {
+    s = (s + 0x6d2b79f5) >>> 0;
+    let t = s;
+    t = Math.imul(t ^ (t >>> 15), t | 1);
+    t ^= t + Math.imul(t ^ (t >>> 7), t | 61);
+    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
+  };
+}
+
+/** Split a raw dataset file into text samples: `.jsonl` → one sample per
+ *  line from the `text` field; anything else → the whole file is one sample. */
+export function parseSamples(raw: string, path: string): string[] {
+  if (path.endsWith(".jsonl") || path.endsWith(".ndjson")) {
+    const samples: string[] = [];
+    for (const [i, line] of raw.split("\n").entries()) {
+      if (!line.trim()) continue;
+      let row: unknown;
+      try {
+        row = JSON.parse(line);
+      } catch {
+        throw new Error(`${path}:${i + 1}: not valid JSON`);
+      }
+      const text = (row as { text?: unknown }).text;
+      if (typeof text !== "string")
+        throw new Error(`${path}:${i + 1}: expected a {"text": "..."} row`);
+      samples.push(text);
+    }
+    return samples;
+  }
+  return raw.length > 0 ? [raw] : [];
+}
+
+/** Tokenize samples (seeded shuffle), concatenate, and pack into
+ *  non-overlapping rows of `sequenceLength` tokens (reference load_data). */
+export function packRows(
+  sampleTokens: number[][],
+  opts: PplDataOptions,
+): Int32Array[] {
+  const { sequenceLength: L, numSamples, seed } = opts;
+  const order = sampleTokens.map((_, i) => i);
+  const rnd = mulberry32(seed);
+  for (let i = order.length - 1; i > 0; i--) {
+    const j = Math.floor(rnd() * (i + 1));
+    [order[i], order[j]] = [order[j]!, order[i]!];
+  }
+
+  const wanted = numSamples > 0 ? L * numSamples : Infinity;
+  const stream: number[] = [];
+  for (const idx of order) {
+    if (stream.length >= wanted) break;
+    stream.push(...sampleTokens[idx]!);
+  }
+
+  const nRows = Math.min(
+    Math.floor(stream.length / L),
+    numSamples > 0 ? numSamples : Infinity,
+  );
+  const rows: Int32Array[] = [];
+  for (let r = 0; r < nRows; r++)
+    rows.push(Int32Array.from(stream.slice(r * L, (r + 1) * L)));
+  return rows;
+}
+
+/** Evaluate perplexity over packed rows. Per batch (all rows are the same
+ *  length, so the stock unpadded forward applies): f32 logits over
+ *  batch[:, :-1], per-token CE vs batch[:, 1:], every position counted. */
+export function evalPpl(
+  model: RuntimeModel,
+  rows: Int32Array[],
+  batchSize: number,
+  onBatch?: (done: number, total: number) => void,
+): PplResult {
+  const L = rows[0]!.length;
+  const allLosses: number[] = [];
+  const numBatches = Math.ceil(rows.length / batchSize);
+
+  for (let s = 0, bi = 0; s < rows.length; s += batchSize, bi++) {
+    const batch = rows.slice(s, s + batchSize);
+    const B = batch.length;
+    const inputs: number[] = [];
+    const targets: number[] = [];
+    for (const row of batch) {
+      for (let i = 0; i < L - 1; i++) inputs.push(row[i]!);
+      for (let i = 1; i < L; i++) targets.push(row[i]!);
+    }
+    const ids = ops.fromInt32(inputs, [B, L - 1]);
+    const tgt = ops.fromInt32(targets, [B, L - 1, 1]);
+    const logits = trainForward(model, ids);
+    const f32 = logits.astype(Dtype.float32);
+    logits.dispose();
+    // CE[t] = logsumexp(logits[t]) − logits[t, target]  (f32, like nn.losses)
+    const lse = ops.logsumexpAxis(f32, 2, false); // [B, L-1]
+    const picked = ops.takeAlongAxis(f32, tgt, 2); // [B, L-1, 1]
+    f32.dispose();
+    const picked2d = ops.reshape(picked, [B, L - 1]);
+    picked.dispose();
+    const losses = ops.sub(lse, picked2d);
+    lse.dispose();
+    picked2d.dispose();
+    const vals = losses.toFloat32();
+    losses.dispose();
+    ids.dispose();
+    tgt.dispose();
+    for (const v of vals) allLosses.push(v);
+    clearCache();
+    onBatch?.(bi + 1, numBatches);
+  }
+
+  const n = allLosses.length;
+  const mean = allLosses.reduce((a, b) => a + b, 0) / n;
+  const ppl = Math.exp(mean);
+  // Sample std (ddof=1) → delta-method SE of the perplexity.
+  const varSum = allLosses.reduce((a, b) => a + (b - mean) ** 2, 0);
+  const std = n > 1 ? Math.sqrt(varSum / (n - 1)) : 0;
+  const standardError = ppl * (std / Math.sqrt(n));
+
+  return { ppl, standardError, meanLoss: mean, tokens: n, rows: rows.length };
+}

package/src/eval/runner.ts

@@ -99,6 +99,11 @@ export interface GenOpts {
    *  (its capability numbers are non-thinking; e.g. MiniCPM5 IFEval 64.7 is
    *  non-thinking). Set MLX_BUN_EVAL_THINK=1 or pass true to evaluate thinking mode. */
   enableThinking?: boolean;
+  /** Explicit KV-quant scheme (cli `generate --l2/--l3/--kv-quant`, resolved with
+   *  serve's semantics). Overrides the MLX_BUN_EVAL_KV_QUANT env default; an empty
+   *  object forces bf16. When it names quantized KV, generation runs the product
+   *  generate() path (same tokens as `serve` with the same scheme). */
+  kvScheme?: Pick<GenerateOptions, "kvBits" | "kvConfig" | "quantizedKvStart">;
 }
 
 /** Bit-exact greedy decode: raw `model.forward` + argmax loop, which matches mlx-lm
@@ -106,7 +111,7 @@ export interface GenOpts {
  *  (forwardHidden/logitsFromHidden + pipelined sampler) diverges from the raw forward
  *  on near-ties past ~32 tokens, so the EVAL decodes bit-exactly itself rather than
  *  alter the serving path. Greedy + full-precision KV only. */
-function greedyDecodeBitExact(tm: TaskModel, ids: number[], maxTokens: number): string {
+export function greedyDecodeBitExact(tm: TaskModel, ids: number[], maxTokens: number): string {
   if (tm.activeAdapters) tm.model.loraState.active = tm.activeAdapters; // apply the mounted adapter
   const cache = tm.model.makeCache();
   try {
@@ -148,23 +153,27 @@ export async function generateText(tm: TaskModel, body: string, opts: GenOpts =
   // plain mlx_lm.generate (kv-quant lives in the serving runtime, not the eval). So
   // for like-for-like parity we default the eval to UNQUANTIZED KV; MMLU's argmax
   // path was already unquantized, which is why it matched while generation didn't.
-  // MLX_BUN_EVAL_KV_QUANT=1 generates through the model's quantized KV (serving config).
-  const kv = (process.env.MLX_BUN_EVAL_KV_QUANT === "1" && tm.config.kvQuant?.length)
-    ? tm.config.kvQuant : undefined;
+  // MLX_BUN_EVAL_KV_QUANT=1 generates through the model's quantized KV (serving
+  // config); an explicit opts.kvScheme (cli generate's resolved tier route) wins.
+  const envKv: GenOpts["kvScheme"] | undefined =
+    (process.env.MLX_BUN_EVAL_KV_QUANT === "1" && tm.config.kvQuant?.length)
+      ? { kvConfig: tm.config.kvQuant, quantizedKvStart: 0 } : undefined;
+  const kvScheme = opts.kvScheme ?? envKv;
+  const kvActive = !!(kvScheme && (kvScheme.kvBits || kvScheme.kvConfig?.length));
 
   // Parity default — greedy + full-precision KV + no sampler arm: decode bit-exactly
   // via the raw forward (matches mlx-lm token-for-token). Sampler arms / kv-quant fall
   // through to the product generate(). DiffusionGemma is non-autoregressive (no AR
   // forward) — always go through generate(), which routes it to the denoising engine.
   const isDiffusion = tm.model instanceof DiffusionGemmaModel;
-  if (!isDiffusion && !kv && !opts.sampler && !tm.samplerOverride) {
+  if (!isDiffusion && !kvActive && !opts.sampler && !tm.samplerOverride) {
     return greedyDecodeBitExact(tm, ids, maxTokens);
   }
   const gen = generate(tm.model, ids, {
     maxTokens,
     temperature: 0, // greedy default — deterministic head-to-head arms
     adapters: tm.activeAdapters ?? [], // per-request activation of the mounted eval-adapter
-    ...(kv ? { kvConfig: kv, quantizedKvStart: 0 } : {}),
+    ...(kvActive ? kvScheme : {}),
     ...(opts.sampler ?? {}), // overrides temperature when provided
     ...(tm.samplerOverride ?? {}), // arm override wins over the task's own sampler
   });

package/src/eval/tasks/bfcl.ts

@@ -15,7 +15,7 @@
 // of the task — is ported exactly. (See summary for this one note.)
 
 import { generateText, loadJsonl, sampleIndices, type TaskModel } from "../runner";
-import { type ToolDefinition } from "../chat-template";
+import { type ToolDefinition } from "../../chat-template";
 
 // ---------------------------------------------------------------------------
 // Data shapes (the merged jsonl: each line = {query, answer}).