mlx-bun 0.0.10 → 0.0.11

package/src/config.ts

@@ -38,6 +38,13 @@ export interface TextConfig {
   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;
@@ -77,6 +84,37 @@ export interface KvQuantSpec {
   groupSize: number;
 }
 
+/** TurboQuant scheme (docs/design/turboquant-kv.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
+ *  src/generate.ts maybeQuantizeKv). Valid kBits: {2,4,5,8}; vBits:
+ *  {2,3,4,5,8} (docs/design/turboquant-kv.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;
@@ -93,6 +131,11 @@ export interface ModelConfig {
 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",
@@ -103,7 +146,7 @@ function parseRope(raw: Record<string, any>): Record<string, RopeParams> {
   return out;
 }
 
-function parseQuantization(raw: Record<string, any> | undefined): QuantizationConfig | null {
+export function parseQuantization(raw: Record<string, any> | undefined): QuantizationConfig | null {
   if (!raw) return null;
   const def: QuantSpec = {
     bits: raw.bits,
@@ -215,8 +258,12 @@ export async function loadModelConfig(modelDir: string): Promise<ModelConfig> {
     numKvSharedLayers: t.num_kv_shared_layers ?? 0,
     enableMoeBlock: t.enable_moe_block ?? false,
     numExperts: t.num_experts ?? 0,
-    topKExperts: t.top_k_experts ?? 0,
+    // gemma4 uses `top_k_experts`; qwen3_moe uses `num_experts_per_tok`.
+    topKExperts: t.top_k_experts ?? t.num_experts_per_tok ?? 0,
     moeIntermediateSize: t.moe_intermediate_size ?? 0,
+    decoderSparseStep: t.decoder_sparse_step ?? 1,
+    mlpOnlyLayers: t.mlp_only_layers ?? [],
+    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,

package/src/download.ts

@@ -114,6 +114,28 @@ function acquireLock(blobPath: string, name: string): () => void {
   throw new Error(`could not acquire ${lockPath}`);
 }
 
+/** Reject a repo file path that could escape the snapshot directory once
+ *  joined onto it. `rfilename` comes verbatim from the remote HF API's
+ *  `siblings[].rfilename` (listRepoFiles) — for a repo the caller doesn't
+ *  control (attacker-owned/attacker-referenced), that string is untrusted
+ *  input reaching a path-join + symlink write (downloadModel's
+ *  `join(snapDir, f.rfilename)` + `symlinkSync`). node:path's `join()`
+ *  normalizes `..` segments, so an rfilename like
+ *  `../../../../home/x/.ssh/authorized_keys` would otherwise resolve
+ *  OUTSIDE the intended snapshots dir and write a symlink at an arbitrary
+ *  filesystem path. Reachable over HTTP via POST /api/hub/download
+ *  (src/hub-rest.ts), which passes any caller-supplied `repo` straight
+ *  through with no allowlist — this check is the actual gate, not that
+ *  route. A leading `/` (absolute path) is equally rejected: join() would
+ *  otherwise treat it as relative to snapDir and still normalize safely,
+ *  but there's no legitimate reason a repo file entry needs one, and
+ *  rejecting it removes any doubt. */
+export function isSafeRepoFilename(rfilename: string): boolean {
+  if (!rfilename || rfilename.startsWith("/") || rfilename.includes("\0")) return false;
+  const parts = rfilename.split("/");
+  return parts.every((p) => p !== "" && p !== "." && p !== "..");
+}
+
 /** git blob identity: sha1("blob <size>\0" + content) — what blobId is. */
 export function gitBlobSha1(bytes: Uint8Array): string {
   const hasher = new Bun.CryptoHasher("sha1");
@@ -294,6 +316,14 @@ export function downloadsSnapshot(): DownloadStatus[] {
   return downloadLog.slice(-5);
 }
 
+/** True when a download for this exact repoId is currently "active" in this
+ *  process's tracker (used by the Hub's POST /api/hub/download to refuse a
+ *  duplicate kick-off — same repoId, still running — rather than starting a
+ *  second writer that would race the first on the same blob lockfile). */
+export function isDownloadActive(repoId: string): boolean {
+  return downloadLog.some((d) => d.repoId === repoId && d.state === "active");
+}
+
 export async function downloadModel(
   repoId: string, opts: DownloadOptions = {},
 ): Promise<string> {
@@ -303,6 +333,16 @@ export async function downloadModel(
   const hub = opts.cacheDir ?? DEFAULT_HUB;
 
   const listing = await listRepoFiles(repoId, { ...opts, token });
+  // Fail the whole download closed (not skip-and-continue) if the repo's
+  // file listing contains a path-traversal rfilename — see
+  // isSafeRepoFilename's doc comment. This gates BOTH the CLI and the web
+  // Hub's POST /api/hub/download (src/hub-rest.ts), which accepts any
+  // caller-supplied repo string with no allowlist of its own.
+  for (const f of listing.files) {
+    if (!isSafeRepoFilename(f.rfilename)) {
+      throw new Error(`${repoId}: refusing unsafe file path in repo listing: ${JSON.stringify(f.rfilename)}`);
+    }
+  }
   const repoDir = join(hub, `models--${repoId.replaceAll("/", "--")}`);
   const blobsDir = join(repoDir, "blobs");
   const snapDir = join(repoDir, "snapshots", listing.sha);

package/src/eval/README.md

@@ -36,8 +36,8 @@ repo rule: free-running greedy "measures chaos").
 # drift gate on an existing non-bit-exact lever (validates the harness today):
 bun scripts/eval.ts kl --candidate e4b
 
-# once a fused kernel is wired for e4b, its drift vs compat:
-bun scripts/eval.ts kl --candidate e4b --self MLX_BUN_PERF_KERNEL --ref-value 0 --cand-value 1
+# self-arm A/B on an output-affecting lever (fused-sdpa on vs off):
+bun scripts/eval.ts kl --candidate e4b --self MLX_BUN_NO_FUSED_SDPA --ref-value 1 --cand-value 0
 
 # optiq-style absolute drift vs a reference model:
 bun scripts/eval.ts kl --candidate e4b --reference <bf16-or-4bit-id>
@@ -55,9 +55,10 @@ the throughput `runs` table). Each row records the active perf levers
 
 - **M0 (done):** KL drift gate + capability score + `quality_runs` DB + CLI.
 - **M0b (done):** *serving-path* KL (`evaluateKlServingDecode` in `kl.ts`,
-  `bun scripts/eval.ts kl --decode`). Validated: identical arms → KL 0;
-  `MLX_BUN_FUSED_DECODE` tiled-decode → KL 0.27 (real drift on the quantized
-  path). Motivation below — the compat KL couldn't see it:
+  `bun scripts/eval.ts kl --decode`). Validated at the time: identical arms →
+  KL 0; the since-deleted `MLX_BUN_FUSED_DECODE` tiled-decode → KL 0.27 (real
+  drift on the quantized path — the lever was removed 2026-07-05, but the
+  validation stands). Motivation below — the compat KL couldn't see it:
   The compat `forward()` path with `makeCache()` uses **plain bf16 caches**. That is wrong
   twice over: (1) e4b's KV is **per-layer mixed 4-bit/8-bit** (group 64, from
   `kv_config.json` → `config.kvQuant`; 8-bit layers 3,4,9,10,11,14, 4-bit
@@ -67,7 +68,7 @@ the throughput `runs` table). Each row records the active perf levers
   caches, then `cache[i] = c.toQuantized(e.groupSize, e.bits)` per
   `config.kvQuant` (the logic already at `generate.ts:68-86`), so the
   teacher-forced prefill runs the real mixed-4/8-bit quantized path and the
-  perf levers (NO_FUSED_SDPA, perf-kernel) actually bite.
+  perf levers (NO_FUSED_SDPA) actually bite.
 - **M1 (done):** all six capability tasks ported + runtime-validated, plus
   `capability` aggregation and the `smoketest` (KL + GSM8K-50) gate. Datasets
   exported by `scripts/eval/export-datasets.py` → `~/.cache/mlx-bun/eval-data/`.

package/src/eval/kl.ts

@@ -26,7 +26,7 @@ import { Dtype } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
 import { Registry } from "../registry";
 import { Weights } from "../weights";
-import { loadModelConfig, type ModelConfig } from "../config";
+import { loadModelConfig, type ModelConfig, type TurboQuantScheme } from "../config";
 import { createModel, type RuntimeModel } from "../model/factory";
 import { loadTokenizer, type LoadedTokenizer } from "../tokenizer";
 import { maybeQuantizeKv } from "../generate";
@@ -223,10 +223,14 @@ function klScalar(p: Float32Array, q: Float32Array): number {
 }
 
 /** Prefill (bf16) → quantize → teacher-forced decode; return per-step logit
- *  distributions for the last `tokens.length - prefillLen` positions. */
+ *  distributions for the last `tokens.length - prefillLen` positions.
+ *  `kvOverride` replaces the default kvConfig-from-model-config scheme when
+ *  given — the seam evaluateKlKvArm uses to run a TurboQuant (or bf16-null)
+ *  arm instead of the env-flag self-flip. */
 function decodeArm(
   model: RuntimeModel, config: ModelConfig, tokens: number[], prefillLen: number,
   flag: string, value: string,
+  kvOverride?: Parameters<typeof maybeQuantizeKv>[1],
 ): Float32Array[] {
   const prev = process.env[flag];
   process.env[flag] = value;
@@ -237,7 +241,10 @@ function decodeArm(
       lp.dispose();
       // quantize the populated caches (mixed 4/8 per kv_config) — exactly
       // what generate() does between prefill and decode.
-      maybeQuantizeKv(cache, { kvConfig: config.kvQuant ?? undefined, quantizedKvStart: 0 });
+      maybeQuantizeKv(
+        cache,
+        kvOverride ?? { kvConfig: config.kvQuant ?? undefined, quantizedKvStart: 0 },
+      );
       const out: Float32Array[] = [];
       for (let i = prefillLen; i < tokens.length; i++) {
         const logits = model.forward([tokens[i]!], cache); // [1,1,V] quantized decode
@@ -292,3 +299,52 @@ export async function evaluateKlServingDecode(opts: {
     refLabel: `serving-decode self:${opts.flag}=${opts.refValue}→${opts.candValue} (${decodeSteps} steps)`,
   };
 }
+
+/** KV-scheme A/B on the teacher-forced serving-decode path: bf16 baseline
+ *  (no kvOverride) vs an explicit KV scheme (e.g. TurboQuant), same model /
+ *  same prompts / same weight load — the quality-vs-bpw curve gate
+ *  (scripts/eval-turboquant-curve.ts). Unlike evaluateKlServingDecode (env-flag
+ *  self-flip), the two arms differ by an actual GenerateOptions kv scheme
+ *  object, so this is the seam for schemes that aren't env-lever-shaped. */
+export async function evaluateKlKvArm(opts: {
+  candidate: string;
+  /** Undefined ⇒ bf16 (no quantization). */
+  candidateScheme?: { turboQuant?: TurboQuantScheme; kvBits?: number; kvConfig?: ModelConfig["kvQuant"] };
+  prompts: string[];
+  nPrompts?: number;
+  seqLen?: number;
+  decodeSteps?: number;
+}): Promise<KLResult> {
+  const t0 = Date.now();
+  const n = opts.nPrompts ?? 64;
+  const seqLen = opts.seqLen ?? 256;
+  const decodeSteps = opts.decodeSteps ?? 32;
+  const { model, tokenizer, config } = await loadRunnable(opts.candidate);
+  const tokenized = prepPrompts(opts.prompts, tokenizer, seqLen, n);
+  const noopFlag = "MLX_BUN_EVAL_KV_ARM_NOOP"; // decodeArm always flips a flag; this one is inert
+  const refKv = {}; // bf16: maybeQuantizeKv no-ops with no kvBits/kvConfig/turboQuant
+  const candKv = opts.candidateScheme
+    ? { ...opts.candidateScheme, quantizedKvStart: 0, kvConfig: opts.candidateScheme.kvConfig ?? undefined }
+    : {};
+
+  const perPrompt: Float32Array[] = [];
+  for (const tokens of tokenized) {
+    if (tokens.length < 2) continue;
+    const prefillLen = Math.max(1, tokens.length - decodeSteps);
+    const refSteps = decodeArm(model, config, tokens, prefillLen, noopFlag, "0", refKv);
+    const candSteps = decodeArm(model, config, tokens, prefillLen, noopFlag, "0", candKv);
+    const m = Math.min(refSteps.length, candSteps.length);
+    const kls = new Float32Array(m);
+    for (let s = 0; s < m; s++) kls[s] = klScalar(refSteps[s]!, candSteps[s]!);
+    perPrompt.push(kls);
+  }
+
+  const agg = aggregate(perPrompt);
+  return {
+    nPrompts: perPrompt.length, seqLen, meanKl: agg.mean, medianKl: agg.median, p95Kl: agg.p95,
+    elapsedSec: (Date.now() - t0) / 1000,
+    refLabel: opts.candidateScheme?.turboQuant
+      ? `bf16 vs turbo k${opts.candidateScheme.turboQuant.kBits}v${opts.candidateScheme.turboQuant.vBits} (${decodeSteps} steps)`
+      : `bf16 vs kv-scheme (${decodeSteps} steps)`,
+  };
+}

package/src/eval/quality-db.ts

@@ -5,7 +5,7 @@
 //
 // One row = one (model, config) measured on one task (or the aggregate
 // "capability" / "kl" rows). `config_json` records the perf levers that
-// were active (compile / fused-decode / perf-kernel) so a row is
+// were active (compile / fused-sdpa / kv-quant) so a row is
 // self-describing: which arm of the head-to-head produced this number.
 
 import { Database } from "bun:sqlite";
@@ -19,7 +19,7 @@ CREATE TABLE IF NOT EXISTS quality_runs (
   model_path TEXT NOT NULL,
   commit_sha TEXT,
   task TEXT NOT NULL,            -- 'kl' | 'gsm8k' | 'mmlu' | ... | 'capability'
-  config_json TEXT NOT NULL,     -- {compile,fusedDecode,perfKernel,...} active levers
+  config_json TEXT NOT NULL,     -- {compile,fusedSdpa,kvQuant,...} active levers
   n_samples INTEGER NOT NULL,
   -- task accuracy (percent 0-100); null for kl-only rows
   pct REAL,

package/src/eval/runner.ts

@@ -99,11 +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
+  /** Explicit KV-quant scheme (cli `generate --l2/--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">;
+  kvScheme?: Pick<GenerateOptions, "kvBits" | "kvConfig" | "quantizedKvStart" | "turboQuant">;
 }
 
 /** Bit-exact greedy decode: raw `model.forward` + argmax loop, which matches mlx-lm
@@ -159,7 +159,7 @@ export async function generateText(tm: TaskModel, body: string, opts: GenOpts =
     (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));
+  const kvActive = !!(kvScheme && (kvScheme.kvBits || kvScheme.kvConfig?.length || kvScheme.turboQuant));
 
   // 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

package/src/expert-offload.ts

@@ -50,12 +50,25 @@ export function isExpertOffload(): boolean {
 /** 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. */
+ *  required for the GPU to gather from it.
+ *
+ *  Lifetime: these arrays are created ONCE at model load
+ *  (QuantizedSwitchLinear.load, gemma4-base.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;
-  return MlxArray.fromView(mm.view(e.offset, e.length), e.shape, SAFETENSORS_TO_MLX[e.dtype]);
+  // 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]);
 }
 
 if (process.env.MLX_BUN_EXPERT_OFFLOAD) activateExpertOffload(process.env.MLX_BUN_EXPERT_OFFLOAD);

package/src/flags.ts

@@ -0,0 +1,13 @@
+// Generic on/off resolver for the per-fork decode flags (--compiled-decode,
+// --compiled-activations, --fused-sdpa, …). An explicit "1"/"0" in the env
+// always wins; otherwise the caller's `defaultOn` applies. The tier aliases
+// (--l1/--l2, applyDecodeRoute in cli.ts) set the env vars; model code
+// reads them through this helper so a hand-set flag and a tier preset resolve
+// identically.
+
+/** Resolve an on/off env flag. An explicit "1"/"0" always wins; otherwise the
+ *  flag takes `defaultOn`. */
+export function flagOn(name: string, defaultOn: boolean): boolean {
+  const v = process.env[name];
+  return v === "1" ? true : v === "0" ? false : defaultOn;
+}