mlx-bun 0.0.13 → 0.2.0

package/src/cli.ts

@@ -13,7 +13,8 @@ import { existsSync } from "node:fs";
 import { basename, join, resolve } from "node:path";
 import { Registry } from "./registry";
 import { loadModelConfig, parseTurboQuantScheme, type TurboQuantScheme } from "./config";
-import { fit, skuMatrix, thisMachine } from "./fit";
+import { fit, skuMatrix, thisMachine, type FitKvScheme } from "./fit";
+import { resolveKvScheme } from "./kv-scheme";
 import { EvalDB } from "./evaldb";
 import pkg from "../package.json" with { type: "json" };
 import { renderHelp } from "./tui";
@@ -27,6 +28,11 @@ import {
   planGlm52MemoryForArtifact,
 } from "./model/glm52-memory";
 import type { JobEvent } from "./jobs/types";
+import {
+  configureRuntime,
+  runtimeValue,
+  type RuntimeKey,
+} from "./runtime-config";
 
 const argv = process.argv.slice(2);
 // The appliance path: naked `mlx-bun` (or only options, e.g.
@@ -118,10 +124,10 @@ const SERVER_FLAGS = `Server options:
   --unix <path>             (internal) listen on a unix socket instead of a
                             TCP port — the engine half of --isolate.
   --batch <n>               Max concurrent requests decoded together
-                            [default: 8]. A LONE request runs the exact
-                            serial engine (same bits, same speed); the cap
-                            engages only when concurrent requests arrive
-                            (agent fan-out). 1 pins strict serial.
+                            [default: 8]. A lone scheduler request uses the
+                            parity- and latency-gated B=1 fast path; concurrent
+                            requests can form B=N (agent fan-out). 1 pins the
+                            preserved strict serial executor.
                             --decode-concurrency is accepted as the
                             mlx_lm.server alias (semantics differ; see
                             docs/reference/server-config.md)
@@ -152,9 +158,13 @@ Model & quality:
                             Serial lane only — with --batch N a mounted draft
                             routes every request serial, like mlx_lm.server.
   --draft-kind <kind>       Override draft-artifact detection:
-                            two-model | assistant | dspark | deepspec | ngram
+                            two-model | assistant | dspark | deepspec | mtp | ngram
                             (deepspec = DeepSeek's released Gemma4DSparkModel
-                            drafters; ngram = MODEL-FREE prompt lookup — no
+                            drafters; mtp = a NATIVE multi-token-prediction
+                            head split from the target's own release, e.g.
+                            mlx-community/Qwen3.8-27B-MTP-bf16 — auto-detected
+                            from model_type *_mtp;
+                            ngram = MODEL-FREE prompt lookup — no
                             --draft-model, drafts copied from the request's
                             own context, lossless)  [default: auto]
   --num-draft-tokens <n>    Drafts per verify round  [default: 3; ngram: 10;
@@ -343,6 +353,8 @@ Usage: mlx-bun fit <query> [options]
 
 Options:
   --ctx <tokens>       Context size to assess  [default: 8192]
+  --kv-quant <mode>    Bill the KV cache quantized: 4 | 8 | config | off
+                       (mirrors serve --kv-quant; default bf16)
   --skus               Also print the Apple silicon SKU matrix
 
 Solves weights + KV growth + prefill transient against wired memory and
@@ -576,6 +588,9 @@ knapsack per-layer bit allocation; implies quantization, no -q needed):
   --candidate-bits <l> Comma list the knapsack may pick from [default: 4,8]
   --calibration-mix <m> "optiq" or a JSONL path  [default: optiq]
   --n-calibration <n>  Calibration samples  [default: 2]
+  --rotate-weights     Fold the model's offline TurboQuant rotation before
+                       quantization (auto-detects Llama/Qwen3.5/Qwen MTP)
+  --rotation-seed <n>  Deterministic rotation seed  [default: 42]
 
 Not supported (mlx_lm.convert flags we don't implement — the command
 exits with an error rather than guessing): --dtype, -d/--dequantize,
@@ -805,11 +820,11 @@ function openChatUi(url: string, hostPort: string): void {
 }
 
 /** Server/runtime flags shared by every mode that loads a model
- *  (serve, pi). Env levers are set here so they're in place before the
+ *  (serve, pi). Runtime overrides are installed here before the
  *  generate/compiled-decode modules read them. */
 /** Resolve the decode ROUTE: a tier alias (--l1/--l2) sets the whole route,
  *  and an explicit per-fork flag (--kv-quant/--fused-sdpa/…) overrides the
- *  alias. Sets the decode env levers and returns the kv-quant mode. See
+ *  alias. Installs decode runtime options and returns the kv-quant mode. See
  *  docs/design/parity-tier-dag.md + unified-engine-frontier-plan.md.
  *  Each tier is a GUARANTEE about which reference you reproduce bit-for-bit:
  *    --l1  bit-for-bit IDENTICAL to mlx-lm    — drop-in replacement for mlx-lm
@@ -863,8 +878,9 @@ function applyDecodeRoute(): { kvQuant?: "off" | "config" | number; turboQuant?:
   const pick = (name: string, base: boolean | undefined): boolean | null => {
     const ex = onOff(name); return ex !== null ? ex : (base ?? null); // explicit flag wins, else tier preset
   };
-  const set = (env: string, val: boolean | null, invert = false) => {
-    if (val !== null) process.env[env] = (invert ? !val : val) ? "1" : "0";
+  const set = (name: RuntimeKey, val: boolean | null, invert = false) => {
+    if (val !== null)
+      configureRuntime({ [name]: (invert ? !val : val) ? "1" : "0" });
   };
   // Explicit --kv-quant picks the composition its ORACLE uses: config (the
   // optiq mixed scheme) runs optiq's fused SDPA (the L2 golden composition);
@@ -874,7 +890,7 @@ function applyDecodeRoute(): { kvQuant?: "off" | "config" | number; turboQuant?:
   set("MLX_BUN_COMPILED_DECODE", pick("compiled-decode", p.compiled));
   set("MLX_BUN_NO_FUSED_SDPA", pick("fused-sdpa", fusedSdpaBase), true); // inverted env
   // Compiled activations (the faithful geglu/swiglu — mlx-lm's @mx.compile kernel).
-  // One fork drives both env vars; qwen3/qwen3.5/universal compile unconditionally,
+  // One fork drives both runtime options; qwen3/qwen3.5/universal compile unconditionally,
   // so this toggles gemma geglu + minicpm5 swiglu (the only unfused-capable sites).
   set("MLX_BUN_COMPILED_GEGLU", pick("compiled-activations", p.compiledAct));
   set("MLX_BUN_COMPILED_SWIGLU", pick("compiled-activations", p.compiledAct));
@@ -909,11 +925,12 @@ function serverRuntimeFlags(): { port: number; serverOptions: import("./server")
     process.exit(1);
   };
   const route = applyDecodeRoute(); // --l1/--l2 tier alias, with per-fork flags overriding
-  if (flag("force-wire")) process.env.MLX_BUN_FORCE_WIRE = "1";
+  if (flag("force-wire")) configureRuntime({ MLX_BUN_FORCE_WIRE: "1" });
   // Remote image_url/audio_url media fetches block private/loopback/link-
   // local destinations by default (SSRF — src/media-fetch.ts); this is the
   // LAN-hosts escape hatch.
-  if (flag("allow-private-media")) process.env.MLX_BUN_ALLOW_PRIVATE_MEDIA = "1";
+  if (flag("allow-private-media"))
+    configureRuntime({ MLX_BUN_ALLOW_PRIVATE_MEDIA: "1" });
 
   const serverOptions: import("./server").ServerOptions = {};
   const budgetGB = Number(opt("memory-budget", "0"));
@@ -945,7 +962,7 @@ function serverRuntimeFlags(): { port: number; serverOptions: import("./server")
       console.warn("--ssd-cache-verify has no effect without --ssd-cache — ignored");
   }
   // --batch N: max concurrent requests batched through the mlx-lm-parity
-  // engine (N=1 = today's serial path). --decode-concurrency is accepted for
+  // engine (N=1 pins the strict serial executor). --decode-concurrency is accepted for
   // drop-in compatibility with mlx_lm.server, but the semantics differ: there
   // it caps per-BatchGenerator decode parallelism (default 32); in mlx-bun it
   // caps continuous batching (default 8; `--batch 1` pins the serial path —
@@ -965,7 +982,7 @@ function serverRuntimeFlags(): { port: number; serverOptions: import("./server")
   // storage, default off (docs/design/paged-kv-cache.md). Serial-only in v1:
   // with no explicit --batch, pin --batch 1 (the default is 8); an explicit
   // --batch N>1 is refused loudly by createServer rather than downgraded.
-  if (flag("paged-kv") || process.env.MLX_BUN_PAGED_KV === "1") {
+  if (flag("paged-kv") || runtimeValue("MLX_BUN_PAGED_KV") === "1") {
     const bsRaw = opt("paged-kv-block-size");
     serverOptions.pagedKv = bsRaw !== null ? { blockSize: Number(bsRaw) } : {};
     if (batchRaw === null) {
@@ -1086,7 +1103,7 @@ async function mountStartupAdapter(
 function runtimeSummary(o: import("./server").ServerOptions): string {
   const kv = o.turboQuant ? `turbo-k${o.turboQuant.kBits}v${o.turboQuant.vBits}`
     : o.kvQuant === "off" ? "off" : typeof o.kvQuant === "number" ? `kv${o.kvQuant}` : "config";
-  const lever = (env: string, dflt: string) => process.env[env] ?? dflt;
+  const lever = (name: RuntimeKey, dflt: string) => runtimeValue(name) ?? dflt;
   return `kv-quant ${kv} · compiled-decode ${lever("MLX_BUN_COMPILED_DECODE", "1") === "1" ? "on" : "off"}` +
     (lever("MLX_BUN_COMPILED_GEGLU", "1") === "0" ? " · compiled-activations off" : "") +
     (o.batch && o.batch > 1 ? ` · batch ${o.batch}` : "") +
@@ -1376,7 +1393,7 @@ switch (cmd) {
 
   case "fit": {
     const query = positional(0);
-    if (!query) { console.error("usage: mlx-bun fit <query> [--ctx N] [--skus]"); process.exit(1); }
+    if (!query) { console.error("usage: mlx-bun fit <query> [--ctx N] [--kv-quant 4|8|config] [--skus]"); process.exit(1); }
     const reg = new Registry();
     if (reg.list().length === 0) await reg.scan();
     const m = reg.resolve(query);
@@ -1443,9 +1460,32 @@ switch (cmd) {
       console.log();
       break;
     }
-    const r = fit(config, m.sizeBytes, ctx, thisMachine(), undefined, m.expertsBytes);
+    // --kv-quant mirrors serve: bill the quantized cache's true bytes so
+    // the reported window matches what serving with the same flag admits.
+    const kvQuantOpt = opt("kv-quant");
+    let fitKvScheme: FitKvScheme | undefined;
+    if (kvQuantOpt === "4" || kvQuantOpt === "8") {
+      fitKvScheme = resolveKvScheme({ override: Number(kvQuantOpt) }).fitOptions;
+    } else if (kvQuantOpt === "config") {
+      if (!config.kvQuant?.length) {
+        console.error(`${m.repoId} ships no per-layer kv-quant config — use --kv-quant 4|8`);
+        process.exit(1);
+      }
+      fitKvScheme = resolveKvScheme({
+        override: "config",
+        config: config.kvQuant,
+        missingConfig: "error",
+      }).fitOptions;
+    } else if (kvQuantOpt != null && kvQuantOpt !== "off") {
+      console.error(`--kv-quant ${kvQuantOpt}: expected 4, 8, config, or off`);
+      process.exit(1);
+    }
+    const r = fit(config, m.sizeBytes, ctx, thisMachine(), undefined, m.expertsBytes, undefined, fitKvScheme);
     h1("will it fit?");
-    console.log(`  ${style.bold(m.repoId)} ${style.dim(`@ ${ctx.toLocaleString()} context · this machine`)}`);
+    const kvNote = fitKvScheme
+      ? style.dim(` · kv-quant ${fitKvScheme.kvBits ? `${fitKvScheme.kvBits}-bit` : "config"}`)
+      : "";
+    console.log(`  ${style.bold(m.repoId)} ${style.dim(`@ ${ctx.toLocaleString()} context · this machine`)}${kvNote}`);
     console.log();
     const expertsNote = config.text.enableMoeBlock && m.expertsBytes > 0
       ? style.dim(`  (experts ${gb(m.expertsBytes)}; top ${config.text.topKExperts}/${config.text.numExperts} read per token)`)
@@ -1470,7 +1510,7 @@ switch (cmd) {
           { header: "fits", paint: (c) => (c.includes("fits") ? style.green(c) : style.dim(c)) },
           { header: "max context", align: "right" }, { header: "decode", align: "right" },
         ],
-        skuMatrix(config, m.sizeBytes, ctx, m.expertsBytes).map((row) => [
+        skuMatrix(config, m.sizeBytes, ctx, m.expertsBytes, fitKvScheme).map((row) => [
           row.sku, `${row.ramGB} GB`, row.fits ? "fits" : "—",
           row.fits ? row.maxContext.toLocaleString() : "—",
           row.fits ? `~${row.decodeTps.toFixed(0)} tok/s` : "—",
@@ -1573,10 +1613,29 @@ switch (cmd) {
     const config = await loadModelConfig(m.path);
     const report = fit(config, m.sizeBytes, 8192, undefined, undefined, m.expertsBytes);
     sFit.done(`${style.bold(m.repoId)} ${style.dim(`· ${gb(m.sizeBytes)} · ~${report.predictedDecodeTps.toFixed(0)} tok/s predicted`)}${picked ? style.dim(" · auto-picked (override: mlx-bun serve <query>)") : ""}`);
+    // Near-ceiling advisory (admission doctrine: advise, never refuse). A
+    // model near the DEFAULT Metal wired ceiling (~75% of RAM) dies with an
+    // UNCATCHABLE async-GPU-OOM on long prefills (M4 24 GB × 17 GB model,
+    // 2026-08-20) — surface the standard remedy up front instead.
+    {
+      const { maxRecommendedWorkingSetSize } = await import("./mlx/ffi");
+      const ceiling = maxRecommendedWorkingSetSize();
+      const sysctlRaw = Bun.spawnSync(["sysctl", "-n", "iogpu.wired_limit_mb"]).stdout.toString().trim();
+      const limitIsDefault = sysctlRaw === "0" || sysctlRaw === "";
+      if (limitIsDefault && ceiling > 0 && m.sizeBytes > 0.8 * ceiling) {
+        const ramGb = Number(Bun.spawnSync(["sysctl", "-n", "hw.memsize"]).stdout.toString().trim()) / 2 ** 30;
+        const suggestMb = Math.floor((ramGb - 2.5) * 1024);
+        console.error(style.dim(
+          `  ! ${gb(m.sizeBytes)} model vs ~${gb(ceiling)} default GPU ceiling — long prompts can\n` +
+          `    hit an uncatchable Metal OOM. If that happens, raise the limit first:\n` +
+          `      sudo sysctl iogpu.wired_limit_mb=${suggestMb}   (resets on reboot)`,
+        ));
+      }
+    }
     const sNative = step("native runtime");
     await ensureNative(sNative);
     sNative.done("native runtime ready");
-    const { createServer, loadContext } = await import("./server");
+    const { createServer, loadContext, shutdownServer } = await import("./server");
     if (flag("expert-offload")) {
       if (m.expertsBytes === 0) {
         console.error(style.dim("--expert-offload ignored: this model has no experts (dense)"));
@@ -1609,9 +1668,9 @@ switch (cmd) {
     }
     const draftKindRaw = opt("draft-kind");
     const draftKind = (draftKindRaw ?? undefined) as
-      | "dspark" | "deepspec" | "assistant" | "two-model" | "ngram" | undefined;
-    if (draftKind !== undefined && !["dspark", "deepspec", "assistant", "two-model", "ngram"].includes(draftKind)) {
-      console.error(`--draft-kind expects two-model|assistant|dspark|deepspec|ngram (got "${draftKindRaw}")`);
+      | "dspark" | "deepspec" | "assistant" | "two-model" | "ngram" | "mtp" | undefined;
+    if (draftKind !== undefined && !["dspark", "deepspec", "assistant", "two-model", "ngram", "mtp"].includes(draftKind)) {
+      console.error(`--draft-kind expects two-model|assistant|dspark|deepspec|mtp|ngram (got "${draftKindRaw}")`);
       process.exit(1);
     }
     // Prompt-lookup window bounds (ngram kind only; validated here so a typo
@@ -1639,7 +1698,10 @@ switch (cmd) {
     const ctx = await loadContext(m.path, m.repoId, {
       memoryBudgetBytes: rt.serverOptions.memoryBudgetBytes,
       glm: glm52RuntimeFlags(rt.serverOptions),
-      ...(draftModelDir || draftKind === "ngram"
+      // ngram is model-free; mtp resolves the bundled <model>/mtp/ companion
+      // in loadContext when no --draft-model is given (cli.md contract —
+      // this gate silently dropped that case until 2026-08-20).
+      ...(draftModelDir || draftKind === "ngram" || draftKind === "mtp"
         ? {
             ...(draftModelDir ? { draftModelDir } : {}),
             numDraftTokens,
@@ -1652,6 +1714,30 @@ switch (cmd) {
     sLoad.done(`weights loaded ${style.dim(`in ${(performance.now() - t0).toFixed(0)} ms`)}${draftModelDir ? style.dim(` · draft: ${draftModelDir.split("/").filter(Boolean).at(-1)}`) : draftKind === "ngram" ? style.dim(" · draft: ngram (prompt lookup)") : ""}`);
     await mountStartupAdapter(ctx, rt.serverOptions);
     const server = createServer(ctx, rt.port, { ...rt.serverOptions, owner: "serve" });
+    let shutdownStarted = false;
+    const onServeSignal = (signal: "SIGINT" | "SIGTERM"): void => {
+      if (shutdownStarted) return;
+      shutdownStarted = true;
+      const rawTimeout = Number(runtimeValue("MLX_BUN_SHUTDOWN_TIMEOUT_MS"));
+      const timeoutMs = Number.isFinite(rawTimeout) && rawTimeout > 0 ? rawTimeout : 120_000;
+      console.log(`\n[server] ${signal}: flushing prompt-cache durability`);
+      void shutdownServer(server, timeoutMs).then((result) => {
+        if (!result.durability.durable) {
+          const d = result.durability;
+          console.warn(
+            `[server] cache flush incomplete after ${Math.round(d.elapsedMs)} ms: ` +
+              `${d.pendingSnapshots} snapshots, ${d.pendingSpills} spills, ` +
+              `${d.droppedSpills} dropped, ${d.failedSpills} failed`,
+          );
+        }
+        process.exit(0);
+      }).catch((error) => {
+        console.warn(`[server] shutdown flush failed: ${(error as Error).message}`);
+        process.exit(0);
+      });
+    };
+    process.once("SIGINT", () => onServeSignal("SIGINT"));
+    process.once("SIGTERM", () => onServeSignal("SIGTERM"));
     const shownHost = rt.serverOptions.hostname ?? "localhost";
     console.log();
     box([
@@ -1682,7 +1768,7 @@ switch (cmd) {
     // three entry points (raw / OpenAI API / chat UI) differ only in how params
     // are POPULATED; here every param is explicit. Decode-path levers mirror
     // serve so you can pin the route (mlx-lm compat = --l1).
-    // --l1/--l2 + per-fork overrides set the decode env levers; the
+    // --l1/--l2 + per-fork overrides set the decode runtime options; the
     // returned KV scheme is applied below so `generate --l2` runs the same
     // quantized-KV route (and produces the same tokens) as `serve --l2`.
     const route = applyDecodeRoute();
@@ -1706,13 +1792,12 @@ switch (cmd) {
     // bf16, same silent fallback as serve), N = uniform bits from decode
     // start. Unset (no tier, no --kv-quant) keeps generate's historical
     // default: bf16 bit-exact greedy (generateText's parity path).
-    const kvScheme =
-      route.turboQuant ? { turboQuant: route.turboQuant, quantizedKvStart: 0 }
-      : route.kvQuant === "off" ? {}
-      : route.kvQuant === "config"
-        ? (tm.config.kvQuant?.length ? { kvConfig: tm.config.kvQuant } : {})
-      : typeof route.kvQuant === "number"
-        ? { kvBits: route.kvQuant, quantizedKvStart: 0 }
+    const kvScheme = route.turboQuant || route.kvQuant !== undefined
+      ? resolveKvScheme({
+          override: route.kvQuant,
+          turboQuant: route.turboQuant,
+          config: tm.config.kvQuant,
+        }).generationOptions
       : undefined;
     const num = (n: string): number | undefined => { const v = opt(n); return v == null ? undefined : Number(v); };
     const text = await generateText(tm, prompt, {
@@ -1794,7 +1879,7 @@ switch (cmd) {
   case "benchmark": {
     const { banner, step, box, style } = await import("./tui");
     banner(pkg.version);
-    serverRuntimeFlags(); // env levers (--compiled-decode, --fused-sdpa, ...) apply to the run
+    serverRuntimeFlags(); // decode options (--compiled-decode, --fused-sdpa, ...) apply to the run
     const { m, picked } = await resolveModelAuto(opt("model") ?? positional(0) ?? opt("query"));
     const tokens = Number(opt("tokens", "256"));
     const runs = Number(opt("runs", "3"));
@@ -2446,6 +2531,12 @@ switch (cmd) {
       console.error(`--candidate-bits expects a comma list of integers in [2, 8] (got "${opt("candidate-bits")}")`);
       process.exit(1);
     }
+    const rotateWeights = flag("rotate-weights");
+    const rotationSeed = Number(opt("rotation-seed", "42"));
+    if (!Number.isInteger(rotationSeed)) {
+      console.error(`--rotation-seed expects an integer (got "${opt("rotation-seed")}")`);
+      process.exit(1);
+    }
     const mlxPath = opt("mlx-path", "mlx_model")!;
     const { existsSync } = await import("node:fs");
     if (existsSync(mlxPath)) {
@@ -2487,7 +2578,7 @@ switch (cmd) {
         ? `quantizing (mixed, target ${targetBpw} bpw — sensitivity sweep, ~minutes)`
         : `quantizing (${qBits}-bit, group ${qGroup})`,
     );
-    const { quantizeModelDir } = await import("./quantize");
+    const { quantizeModelDir, automaticRotationWeightTransform } = await import("./quantize");
     try {
       const r = await quantizeModelDir(
         srcDir,
@@ -2500,6 +2591,9 @@ switch (cmd) {
           ...(candidateBits ? { candidateBits } : {}),
           ...(opt("calibration-mix") ? { calibrationMix: opt("calibration-mix")! } : {}),
           ...(opt("n-calibration") ? { nCalibration: Number(opt("n-calibration")) } : {}),
+          ...(rotateWeights
+            ? { weightTransform: automaticRotationWeightTransform({ seed: rotationSeed }) }
+            : {}),
         },
         (e) => sQ.update(e.message),
       );
@@ -2511,6 +2605,7 @@ switch (cmd) {
         `source    ${style.dim(srcDir)}`,
         `model     ${style.bold(r.outDir)} ${style.dim(`· ${gb(r.write.totalSize)}`)}`,
         `quant     ${style.dim(targetBpw !== undefined ? `mixed ${r.achievedBpw.toFixed(2)} bpw (target ${targetBpw})` : `${qBits}-bit g${qGroup} affine`)}`,
+        ...(rotateWeights ? [`transform ${style.dim(`TurboQuant rotation seed ${rotationSeed}`)}`] : []),
         "",
         `serve it   ${style.accent(`mlx-bun serve ${r.outDir}`)}`,
       ]);

package/src/eval/kl.ts

@@ -9,7 +9,7 @@
 //
 // Two reference modes:
 //   * self-flag (DEFAULT, the kernel-drift gate): one model, one weight
-//     load; forward each prompt twice with an env lever set to its
+//     load; forward each prompt twice with a runtime lever set to its
 //     reference value vs its candidate value. Directly measures the drift
 //     a perf lever introduces. Fits in RAM trivially (no second model).
 //   * two-model (optiq-style absolute quality): a separate reference
@@ -30,6 +30,10 @@ import { loadModelConfig, type ModelConfig, type TurboQuantScheme } from "../con
 import { createModel, type RuntimeModel } from "../model/factory";
 import { loadTokenizer, type LoadedTokenizer } from "../tokenizer";
 import { maybeQuantizeKv } from "../generate";
+import {
+  configureRuntime,
+  runtimeKey,
+} from "../runtime-config";
 
 export interface KLResult {
   nPrompts: number;
@@ -135,23 +139,28 @@ export async function evaluateKlSelfFlag(opts: {
   const { model, tokenizer } = await loadRunnable(opts.candidate);
   const tokenized = prepPrompts(opts.prompts, tokenizer, seqLen, n);
 
-  const prev = process.env[opts.flag];
+  const key = runtimeKey(opts.flag);
   const perPrompt: Float32Array[] = [];
-  try {
-    for (const ids of tokenized) {
-      process.env[opts.flag] = opts.refValue;
-      const p = forwardLogits(model, ids);
-      process.env[opts.flag] = opts.candValue;
-      const q = forwardLogits(model, ids);
-      const T = Math.min(p.shape[1]!, q.shape[1]!);
-      const pT = p.slice([0, 0, 0], [1, T, p.shape[2]!]);
-      const qT = q.slice([0, 0, 0], [1, T, q.shape[2]!]);
-      perPrompt.push(klPerToken(pT, qT));
-      for (const a of [p, q, pT, qT]) a.dispose();
+  for (const ids of tokenized) {
+    let p: MlxArray;
+    const restoreRef = configureRuntime({ [key]: opts.refValue });
+    try {
+      p = forwardLogits(model, ids);
+    } finally {
+      restoreRef();
     }
-  } finally {
-    if (prev === undefined) delete process.env[opts.flag];
-    else process.env[opts.flag] = prev;
+    let q: MlxArray;
+    const restoreCandidate = configureRuntime({ [key]: opts.candValue });
+    try {
+      q = forwardLogits(model, ids);
+    } finally {
+      restoreCandidate();
+    }
+    const T = Math.min(p.shape[1]!, q.shape[1]!);
+    const pT = p.slice([0, 0, 0], [1, T, p.shape[2]!]);
+    const qT = q.slice([0, 0, 0], [1, T, q.shape[2]!]);
+    perPrompt.push(klPerToken(pT, qT));
+    for (const a of [p, q, pT, qT]) a.dispose();
   }
 
   const agg = aggregate(perPrompt);
@@ -232,8 +241,7 @@ function decodeArm(
   flag: string, value: string,
   kvOverride?: Parameters<typeof maybeQuantizeKv>[1],
 ): Float32Array[] {
-  const prev = process.env[flag];
-  process.env[flag] = value;
+  const restore = configureRuntime({ [runtimeKey(flag)]: value });
   try {
     const cache = model.makeCache();
     try {
@@ -256,8 +264,7 @@ function decodeArm(
       for (const c of cache) c.dispose();
     }
   } finally {
-    if (prev === undefined) delete process.env[flag];
-    else process.env[flag] = prev;
+    restore();
   }
 }
 

package/src/eval/runner.ts

@@ -6,9 +6,11 @@
 import { homedir } from "node:os";
 import { readFileSync } from "node:fs";
 import { generate, type GenerateOptions } from "../generate";
+import { clearCache } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
 import { ChatTemplate } from "../chat-template";
 import { loadModelConfig, type ModelConfig } from "../config";
+import { runtimeValue } from "../runtime-config";
 import {
   openModel,
   type Glm52RuntimeOpenOptions,
@@ -19,7 +21,7 @@ import { DiffusionGemmaModel } from "../model/diffusion-gemma";
 import { loadTokenizer, type LoadedTokenizer } from "../tokenizer";
 import { resolveModelDir } from "./kl";
 
-export const EVAL_DATA_DIR = process.env.MLX_BUN_EVAL_DATA ?? `${homedir()}/.cache/mlx-bun/eval-data`;
+export const EVAL_DATA_DIR = runtimeValue("MLX_BUN_EVAL_DATA") ?? `${homedir()}/.cache/mlx-bun/eval-data`;
 
 /** Read a jsonl dataset exported by scripts/eval/export-datasets.py. */
 export function loadJsonl<T = Record<string, unknown>>(name: string): T[] {
@@ -124,7 +126,23 @@ export function greedyDecodeBitExact(tm: TaskModel, ids: number[], maxTokens: nu
   if (tm.activeAdapters) tm.model.loraState.active = tm.activeAdapters; // apply the mounted adapter
   const cache = tm.model.makeCache();
   try {
-    if (ids.length > 1) tm.model.forward(ids.slice(0, -1), cache).dispose(); // prefill
+    // Chunked, LM-head-free prefill. The old single-shot `forward(prompt)`
+    // materialized [1, L, vocab] logits just to discard them — at a 248k
+    // vocab that's ~1 GB per 1k prompt tokens, and with no allocator-cache
+    // clears the Metal buffer cache retained every giant freed tensor
+    // (the 27B capability-suite 73 GB swap-thrash). Prefill only needs the
+    // cache state, which forwardHidden advances identically (the LM head
+    // never touches the cache) — bit-exact by construction. Chunk size and
+    // clear cadence mirror mlx-lm generate_step / our generate().
+    const PREFILL_CHUNK = 2048;
+    const prompt = ids.slice(0, -1);
+    for (let s = 0; s < prompt.length; s += PREFILL_CHUNK) {
+      const chunk = prompt.slice(s, s + PREFILL_CHUNK);
+      const cIds = ops.fromInt32(chunk, [1, chunk.length]);
+      tm.model.forwardHidden(cIds, cache).dispose();
+      cIds.dispose();
+      clearCache();
+    }
     let last = ids[ids.length - 1]!;
     const eos = new Set(tm.config.eosTokenIds);
     const out: number[] = [];
@@ -134,6 +152,7 @@ export function greedyDecodeBitExact(tm: TaskModel, ids: number[], maxTokens: nu
       lg.dispose();
       const best = ops.itemUint32(am);
       am.dispose();
+      if (i % 256 === 0) clearCache(); // mlx-lm decode cadence
       if (eos.has(best)) break; // mlx-lm halts on EOS and excludes it from the output
       out.push(best);
       last = best;
@@ -148,7 +167,7 @@ export function greedyDecodeBitExact(tm: TaskModel, ids: number[], maxTokens: nu
  *  pass `opts.sampler` to drive temperature/HLG/etc. */
 export async function generateText(tm: TaskModel, body: string, opts: GenOpts = {}): Promise<string> {
   const maxTokens = opts.maxTokens ?? 256;
-  const enableThinking = opts.enableThinking ?? process.env.MLX_BUN_EVAL_THINK === "1";
+  const enableThinking = opts.enableThinking ?? runtimeValue("MLX_BUN_EVAL_THINK") === "1";
   const templated = opts.useChat !== false && tm.template !== null;
   const text = templated
     ? tm.template!.render([{ role: "user", content: body }], { addGenerationPrompt: true, enableThinking })
@@ -165,7 +184,7 @@ export async function generateText(tm: TaskModel, body: string, opts: GenOpts =
   // 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)
+    (runtimeValue("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 || kvScheme.turboQuant));

package/src/eval/tasks/bfcl.ts

@@ -16,6 +16,7 @@
 
 import { generateText, loadJsonl, sampleIndices, type TaskModel } from "../runner";
 import { type ToolDefinition } from "../../chat-template";
+import { runtimeValue } from "../../runtime-config";
 
 // ---------------------------------------------------------------------------
 // Data shapes (the merged jsonl: each line = {query, answer}).
@@ -465,7 +466,7 @@ export async function evaluateBfcl(
 
   // Optiq-parity mode (DEFAULT): optiq's EXACT 200 BFCL-simple questions, scored
   // through OUR pipeline. MLX_BUN_BFCL_FROZEN=0 reverts to our copy + sampling.
-  const useFrozen = opts.frozen ?? (process.env.MLX_BUN_BFCL_FROZEN !== "0");
+  const useFrozen = opts.frozen ?? (runtimeValue("MLX_BUN_BFCL_FROZEN") !== "0");
   let items: BfclItem[];
   if (useFrozen) {
     items = loadJsonl<BfclFrozenRow>("bfcl_optiq_frozen").map((r) => ({

package/src/eval/tasks/gsm8k.ts

@@ -2,6 +2,7 @@
 // 3-shot CoT, greedy decode, extract the numeric answer, compare.
 
 import { generateText, loadJsonl, sampleIndices, type TaskModel } from "../runner";
+import { runtimeValue } from "../../runtime-config";
 
 // 3-shot exemplars (verbatim from gsm8k.py).
 const FEW_SHOT: { question: string; answer: string }[] = [
@@ -73,7 +74,7 @@ export async function evaluateGsm8k(
   // his published GSM8K used) but through OUR OWN pipeline — our few-shot builder,
   // our chat template, our scorer. If our number reproduces his, our runtime is
   // faithful. MLX_BUN_GSM8K_FROZEN=0 falls back to our own sampling of the full set.
-  const useFrozen = opts.frozen ?? (process.env.MLX_BUN_GSM8K_FROZEN !== "0");
+  const useFrozen = opts.frozen ?? (runtimeValue("MLX_BUN_GSM8K_FROZEN") !== "0");
   let rows: Gsm8kRow[];
   if (useFrozen) {
     rows = loadJsonl<Gsm8kRow>("gsm8k_optiq_frozen"); // optiq's exact draw; we use {question, answer}

package/src/eval/tasks/hashhop.ts

@@ -10,6 +10,7 @@
 // fused-attention kernel (opportunity A) has the most to prove.
 
 import { generateText, loadJsonl, type TaskModel } from "../runner";
+import { runtimeValue } from "../../runtime-config";
 
 const INSTR =
   "You will be given a list of hash assignments and a starting hash. " +
@@ -96,7 +97,7 @@ export async function evaluateHashhop(
   // RECONSTRUCTION producing DIFFERENT problems, so it cannot reproduce optiq's
   // number — the frozen problems are the only way to match. MLX_BUN_HASHHOP_FROZEN=0
   // uses our generator.
-  const useFrozen = opts.frozen ?? (process.env.MLX_BUN_HASHHOP_FROZEN !== "0");
+  const useFrozen = opts.frozen ?? (runtimeValue("MLX_BUN_HASHHOP_FROZEN") !== "0");
   if (useFrozen) {
     const rows = loadJsonl<HashhopFrozen>("hashhop_optiq_frozen");
     const hopCorrect: Record<number, number> = {};

package/src/eval/tasks/humaneval.ts

@@ -16,11 +16,12 @@ import { mkdtempSync, rmSync, writeFileSync } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { generateText, loadJsonl, sampleIndices, type TaskModel } from "../runner";
+import { runtimeValue } from "../../runtime-config";
 
 // Oracle venv python — the same interpreter the Python reference uses.
 // Machine-specific: overridable via MLX_BUN_ORACLE_PYTHON (the committed default is
 // the reference box; other laptops set the env, never re-commit the path).
-const ORACLE_PYTHON = process.env.MLX_BUN_ORACLE_PYTHON ?? "/Users/joshrossi/Code/mlx-lm/.venv/bin/python";
+const ORACLE_PYTHON = runtimeValue("MLX_BUN_ORACLE_PYTHON") ?? "/Users/joshrossi/Code/mlx-lm/.venv/bin/python";
 
 interface HumanevalRow {
   task_id: string;
@@ -185,7 +186,7 @@ function runProgram(program: string, timeoutSec: number, memoryMb: number): RunR
       env,
       timeout: timeoutSec * 1000, // hard wall-clock kill
       killSignal: "SIGKILL",
-      stdio: ["ignore", "ignore", process.env.MLX_BUN_EVAL_DEBUG === "1" ? "pipe" : "ignore"],
+      stdio: ["ignore", "ignore", runtimeValue("MLX_BUN_EVAL_DEBUG") === "1" ? "pipe" : "ignore"],
       encoding: "utf8",
     });
 
@@ -219,7 +220,7 @@ export async function evaluateHumaneval(
   // Optiq-parity mode (DEFAULT): optiq's EXACT 164-problem set (he runs the full
   // HumanEval test split), scored through OUR pipeline. MLX_BUN_HUMANEVAL_FROZEN=0
   // reverts to our own copy + optional sampling.
-  const useFrozen = opts.frozen ?? (process.env.MLX_BUN_HUMANEVAL_FROZEN !== "0");
+  const useFrozen = opts.frozen ?? (runtimeValue("MLX_BUN_HUMANEVAL_FROZEN") !== "0");
   const rows = loadJsonl<HumanevalRow>(useFrozen ? "humaneval_optiq_frozen" : "humaneval");
   const idx = useFrozen
     ? Array.from({ length: rows.length }, (_, i) => i)
@@ -252,7 +253,7 @@ export async function evaluateHumaneval(
 
     const rr = runProgram(program, timeoutSec, memoryMb);
     if (rr.ok) nPass++;
-    else if (process.env.MLX_BUN_EVAL_DEBUG === "1") {
+    else if (runtimeValue("MLX_BUN_EVAL_DEBUG") === "1") {
       const err = rr.stderr ?? "";
       const reason = rr.timedOut ? "TIMEOUT"
         : completion.trim() === "" ? "EMPTY"

package/src/eval/tasks/ifeval.ts

@@ -21,6 +21,7 @@
 // default in both implementations, so behavior is identical.
 
 import { generateText, loadJsonl, sampleIndices, type TaskModel } from "../runner";
+import { runtimeValue } from "../../runtime-config";
 
 export interface IfevalInstance {
   prompt: string;
@@ -472,7 +473,7 @@ export async function evaluateIfeval(
   // Optiq-parity mode (DEFAULT): score optiq's EXACT prompt set (he runs the full
   // IFEval split — 541 prompts) through OUR verifiers + OUR chat template.
   // MLX_BUN_IFEVAL_FROZEN=0 reverts to our own copy + optional sampling.
-  const useFrozen = opts.frozen ?? (process.env.MLX_BUN_IFEVAL_FROZEN !== "0");
+  const useFrozen = opts.frozen ?? (runtimeValue("MLX_BUN_IFEVAL_FROZEN") !== "0");
   const rows = loadJsonl<IfevalRow>(useFrozen ? "ifeval_optiq_frozen" : "ifeval");
   const nSamples = opts.nSamples ?? rows.length;
   const idx = (!useFrozen && nSamples < rows.length)

package/src/eval/tasks/mmlu.ts

@@ -10,6 +10,7 @@
 // table. MLX_BUN_MMLU_FROZEN=0 falls back to our own sampling of the full 14k set.
 
 import { loadJsonl, sampleIndices, type TaskModel } from "../runner";
+import { runtimeValue } from "../../runtime-config";
 
 interface MmluRow {
   question: string;
@@ -85,7 +86,7 @@ export async function evaluateMmlu(
   // parity, not sampling. Set MLX_BUN_MMLU_FROZEN=0 (or frozen:false) to fall back to
   // our own stratified sampling of the full 14k test set (a DIFFERENT subset: JS RNG
   // ≠ numpy RandomState, so it does not match optiq's questions even at the same seed).
-  const useFrozen = opts.frozen ?? (process.env.MLX_BUN_MMLU_FROZEN !== "0");
+  const useFrozen = opts.frozen ?? (runtimeValue("MLX_BUN_MMLU_FROZEN") !== "0");
 
   const testRows = loadJsonl<MmluRow>(useFrozen ? "mmlu_optiq_frozen" : "mmlu_test");
   const devRows = loadJsonl<MmluRow>(useFrozen ? "mmlu_optiq_dev" : "mmlu_dev");