mlx-bun 0.0.13 → 0.2.0

package/src/model/qwen3_5.ts

@@ -30,6 +30,10 @@ import {
   type Mask,
 } from "./gemma4-base";
 import { gatedDeltaUpdate, SSMCache } from "./qwen3-delta";
+import {
+  activeMrope, applyInterleavedRope, buildMropePositions, mropeInvFreq,
+  setActiveMrope, type MropeRequestState,
+} from "./qwen3-mrope";
 
 const PREFIX = "language_model";
 
@@ -146,30 +150,58 @@ export class GatedDeltaNet {
     const [B, S] = x.shape as [number, number, number];
     const convDim = this.keyDim * 2 + this.valueDim;
     const nKeep = this.convKernel - 1;
+    const prof = (globalThis as Record<string, unknown>).__deltaProf as
+      Record<string, number>
+      | undefined;
+
+    // Armed speculative verify round (serve loop, serial lane): this forward
+    // must be rewindable. Snapshot = hand the REPLACED state slots to the
+    // round instead of disposing (free — arrays are immutable); record the
+    // position-local kernel inputs (qkv/a/b) and install the prefix replay.
+    const spec = cache.specRound;
+    if (spec) {
+      if (!spec.armed)
+        throw new Error("GatedDeltaNet: spec round already recorded this round");
+      spec.armed = false;
+      spec.S = S;
+      spec.replay = (c, keep) => this.#replaySpecPrefix(c, keep);
+    }
 
+    const t0 = prof ? performance.now() : 0;
     const qkv = this.inProjQkv.forward(x); // [B,S,convDim]
     let z = this.inProjZ.forward(x);
     z = disposing(z, ops.reshape(z, [B, S, this.numVHeads, this.headVDim]));
     const b = this.inProjB.forward(x); // [B,S,numVHeads]
-    const a = this.inProjA.forward(x); // [B,S,numVHeads]
+    const a = this.inProjA.forward(x);
+    if (prof) { ops.evalAll([qkv, z, b, a]); prof.proj = (prof.proj ?? 0) + performance.now() - t0; }
+    const tc = prof ? performance.now() : 0; // [B,S,numVHeads]
 
     // Causal depthwise conv with the conv-state prefix (B=1: no ssm mask).
     const convState =
       cache.conv ?? ops.zeros([B, nKeep, convDim], x.dtype);
     const convInput = ops.concatAxis([convState, qkv], 1); // [B,S+nKeep,convDim]
     if (!cache.conv) convState.dispose();
-    qkv.dispose();
-    // New conv state = last nKeep rows (contiguous, the array is sliced).
-    const newConv = ops.contiguous(
-      convInput.slice([0, S, 0], [B, S + nKeep, convDim]),
-    );
-    cache.conv?.dispose();
+    if (spec) spec.qkv = qkv;
+    else qkv.dispose();
+    // New conv state = last nKeep rows. MUST be a TRUE copy (mx.copy):
+    // this tail slice is already row-contiguous at B=1, so ops.contiguous
+    // is a no-op VIEW sharing convInput's buffer — caching it pinned the
+    // whole [B, S+nKeep, convDim] chunk buffer (~42 MB/layer/chunk, 48
+    // layers = the measured 2.0 GB-per-2048-chunk prefill leak; 46 GB
+    // active at 32k, async-GPU-OOM on 24 GB — M4 2026-08-20).
+    const newConvView = convInput.slice([0, S, 0], [B, S + nKeep, convDim]);
+    const newConv = ops.copyOf(newConvView);
+    newConvView.dispose();
+    if (spec) spec.prevConv = cache.conv;
+    else cache.conv?.dispose();
     cache.conv = newConv;
 
     const conv = ops.conv1d(convInput, this.convWeight, 1, 0, 1, convDim);
     convInput.dispose();
     const convOut = compiledSilu(conv); // nn.silu (mx.compile), matches mlx-lm
     conv.dispose();
+    if (prof) { ops.evalAll([convOut]); prof.conv = (prof.conv ?? 0) + performance.now() - tc; }
+    const tn = prof ? performance.now() : 0;
 
     const [qFlat, kFlat, vFlat] = ops.split(
       convOut, [this.keyDim, 2 * this.keyDim], -1,
@@ -187,16 +219,26 @@ export class GatedDeltaNet {
     q = disposing(q, ops.mulScalar(q, invScale * invScale));
     k = disposing(k, ops.rmsNorm(k, null, 1e-6));
     k = disposing(k, ops.mulScalar(k, invScale));
+    if (prof) { ops.evalAll([q, k]); prof.norms = (prof.norms ?? 0) + performance.now() - tn; }
+    const tk = prof ? performance.now() : 0;
 
     const [out, newState] = gatedDeltaUpdate(
       q, k, v, a, b, this.aLog, this.dtBias, cache.recurrent,
     );
+    if (prof) { ops.evalAll([out]); prof.kernel = (prof.kernel ?? 0) + performance.now() - tk; }
+    const to = prof ? performance.now() : 0;
     q.dispose();
     k.dispose();
     v.dispose();
-    a.dispose();
-    b.dispose();
-    cache.recurrent?.dispose();
+    if (spec) {
+      spec.a = a;
+      spec.b = b;
+      spec.prevRecurrent = cache.recurrent;
+    } else {
+      a.dispose();
+      b.dispose();
+      cache.recurrent?.dispose();
+    }
     cache.recurrent = newState;
     cache.advance(S);
 
@@ -208,9 +250,80 @@ export class GatedDeltaNet {
     gated.dispose();
     const result = this.outProj.forward(merged);
     merged.dispose();
+    if (prof) { ops.evalAll([result]); prof.out = (prof.out ?? 0) + performance.now() - to; }
     return result;
   }
 
+  /** Speculative rollback replay: with the pre-round snapshot RESTORED onto
+   *  `cache`, re-advance the first `keep` window tokens from the recorded
+   *  position-local inputs. Every op mirrors forward() on the same values —
+   *  conv windows, silu, per-position norms, and the kernel's serial prefix
+   *  are all independent of the rejected tail — so the resulting conv +
+   *  recurrent state is BIT-EXACTLY what a forward over only the accepted
+   *  prefix would have produced. Output projections (z / out_proj) are state-
+   *  free and skipped: only the states matter here. */
+  #replaySpecPrefix(cache: SSMCache, keep: number): void {
+    const r = cache.specRound;
+    if (!r || !r.qkv || !r.a || !r.b)      throw new Error("GatedDeltaNet replay without recorded round inputs");
+    const B = r.qkv.shape[0]!;
+    const convDim = this.keyDim * 2 + this.valueDim;
+    const nKeep = this.convKernel - 1;
+
+    const qkvPfxView = r.qkv.slice([0, 0, 0], [B, keep, convDim]);
+    const qkvPfx = ops.contiguous(qkvPfxView);
+    qkvPfxView.dispose();
+    const convState = cache.conv ?? ops.zeros([B, nKeep, convDim], qkvPfx.dtype);
+    const convInput = ops.concatAxis([convState, qkvPfx], 1);
+    if (!cache.conv) convState.dispose();
+    qkvPfx.dispose();
+    // TRUE copy, same reason as forward(): the contiguous tail slice would
+    // otherwise be a view pinning the whole replay convInput buffer.
+    const rConvView = convInput.slice([0, keep, 0], [B, keep + nKeep, convDim]);
+    const newConv = ops.copyOf(rConvView);
+    rConvView.dispose();
+    cache.conv?.dispose();
+    cache.conv = newConv;
+
+    const conv = ops.conv1d(convInput, this.convWeight, 1, 0, 1, convDim);
+    convInput.dispose();
+    const convOut = compiledSilu(conv);
+    conv.dispose();
+    const [qFlat, kFlat, vFlat] = ops.split(
+      convOut, [this.keyDim, 2 * this.keyDim], -1,
+    ) as [MlxArray, MlxArray, MlxArray];
+    convOut.dispose();
+    let q = ops.reshape(qFlat, [B, keep, this.numKHeads, this.headKDim]);
+    qFlat.dispose();
+    let k = ops.reshape(kFlat, [B, keep, this.numKHeads, this.headKDim]);
+    kFlat.dispose();
+    const v = disposing(vFlat, ops.reshape(vFlat, [B, keep, this.numVHeads, this.headVDim]));
+    const invScale = Math.pow(this.headKDim, -0.5);
+    q = disposing(q, ops.rmsNorm(q, null, 1e-6));
+    q = disposing(q, ops.mulScalar(q, invScale * invScale));
+    k = disposing(k, ops.rmsNorm(k, null, 1e-6));
+    k = disposing(k, ops.mulScalar(k, invScale));
+
+    const aPfxView = r.a.slice([0, 0, 0], [B, keep, this.numVHeads]);
+    const aPfx = ops.contiguous(aPfxView);
+    aPfxView.dispose();
+    const bPfxView = r.b.slice([0, 0, 0], [B, keep, this.numVHeads]);
+    const bPfx = ops.contiguous(bPfxView);
+    bPfxView.dispose();
+
+    const [y, newState] = gatedDeltaUpdate(
+      q, k, v, aPfx, bPfx, this.aLog, this.dtBias, cache.recurrent,
+    );
+    q.dispose();
+    k.dispose();
+    v.dispose();
+    aPfx.dispose();
+    bPfx.dispose();
+    y.dispose(); // only the state advance matters on the rollback path
+    cache.recurrent?.dispose();
+    cache.recurrent = newState;
+    cache.advance(keep);
+  }
+
   private rmsNormGated(hidden: MlxArray, gate: MlxArray): MlxArray {
     const xn = ops.rmsNorm(hidden, this.normWeight, this.eps); // bf16
     // Oracle: _precise_swiglu(hidden, gate, xn) — ALWAYS one @mx.compile kernel
@@ -279,13 +392,23 @@ export class Qwen3Attention {
     // position as ropeOffsetArr (rows have different prompt lengths); the
     // dynamic-offset kernel is the same fast::rope, bit-exact vs the static
     // form (tests/compile.test.ts). Serial lane: scalar offset, unchanged.
+    // Vision requests (serial lane only) install activeMrope for the current
+    // forward: 3D interleaved positions via the manual apply — the reference
+    // never uses the fused fast-rope kernel when position_ids are supplied,
+    // so this IS the oracle's own arithmetic. Text-only stays on ops.rope.
+    const mr = activeMrope;
     const offArr = (cache as { ropeOffsetArr?: MlxArray }).ropeOffsetArr;
-    q = disposing(q, offArr
-      ? ops.ropeDynamic(q, this.ropeDims, this.ropeBase, offArr, null)
-      : ops.rope(q, this.ropeDims, this.ropeBase, cache.offset, null));
-    k = disposing(k, offArr
-      ? ops.ropeDynamic(k, this.ropeDims, this.ropeBase, offArr, null)
-      : ops.rope(k, this.ropeDims, this.ropeBase, cache.offset, null));
+    if (mr) {
+      q = disposing(q, applyInterleavedRope(q, mr));
+      k = disposing(k, applyInterleavedRope(k, mr));
+    } else {
+      q = disposing(q, offArr
+        ? ops.ropeDynamic(q, this.ropeDims, this.ropeBase, offArr, null)
+        : ops.rope(q, this.ropeDims, this.ropeBase, cache.offset, null));
+      k = disposing(k, offArr
+        ? ops.ropeDynamic(k, this.ropeDims, this.ropeBase, offArr, null)
+        : ops.rope(k, this.ropeDims, this.ropeBase, cache.offset, null));
+    }
 
     let attn: MlxArray;
     if (cache instanceof QuantizedKVCache) {
@@ -448,8 +571,42 @@ export class Qwen35Model {
     return this.forwardLayers(h, cache);
   }
 
-  forwardEmbeddings(_embeds: MlxArray, _cache: Cache[], _bidir: MlxArray | null): MlxArray {
-    throw new Error("qwen3_5 vision/input-embedding path is not supported");
+  /** Active vision mRoPE request state (serial lane; set by the generation
+   *  gateway around a vision request's run, null for text-only — which keeps
+   *  the bit-exact fast-rope path). While set, forwardLayers installs the
+   *  per-forward interleaved cos/sin consumed by every full-attn layer. */
+  mrope: MropeRequestState | null = null;
+  #mropeInvFreq: MlxArray | null = null;
+
+  /** Vision prefill: spliced input embeddings [1, L, H] (image features
+   *  overwriting image-token rows — the caller builds them). Caller keeps
+   *  ownership of `embeds`. bidir/ids/multimodal are the Gemma-shaped extras
+   *  and are unused: Qwen3.5 vision attends fully causally. */
+  forwardEmbeddings(
+    embeds: MlxArray, cache: Cache[], _bidir: MlxArray | null,
+    _ids: MlxArray | null = null, _multimodal: MlxArray | null = null,
+  ): MlxArray {
+    const h = ops.contiguous(embeds); // fresh handle — the layer loop consumes it
+    return this.forwardLayers(h, cache);
+  }
+
+  /** Layer-output tap (same contract as Gemma4Model.hiddenTap): the serve
+   *  loop's forwardMaybeTap sets it around prefill/verify forwards so
+   *  KV-borrowing draft sources (native Qwen MTP) can read PRE-final-norm
+   *  hiddens — layer 63's output stream, what mlx-vlm captures with
+   *  skip_final_norm=True. No-op (and graph-identical) when null. */
+  hiddenTap: { layers: Set<number>; pos?: number; captured: Map<number, MlxArray> } | null = null;
+
+  /** Store layer `i`'s residual stream ([1,L,H]) if the tap requests it.
+   *  Copies so the loop's h.dispose() can't free the capture. */
+  protected captureLayer(i: number, h: MlxArray): void {
+    const tap = this.hiddenTap;
+    if (!tap || !tap.layers.has(i)) return;
+    const H = h.shape[2]!;
+    const src = tap.pos !== undefined ? h.slice([0, tap.pos, 0], [1, tap.pos + 1, H]) : h;
+    const copy = ops.contiguous(src);
+    if (src !== h) src.dispose();
+    tap.captured.set(i, copy);
   }
 
   protected forwardLayers(h0: MlxArray, cache: Cache[]): MlxArray {
@@ -457,14 +614,49 @@ export class Qwen35Model {
     // One full-attention mask shared by all full layers (same offset); linear
     // layers see no ssm mask at B=1.
     const faMask = cache[this.faIdx]!.makeMask(L, null);
-    let h = h0;
-    for (let i = 0; i < this.layers.length; i++) {
-      const next = this.layers[i]!.forward(h, faMask, cache[i]!);
-      h.dispose();
-      h = next;
+    // Vision requests: one interleaved-mRoPE cos/sin table per forward,
+    // shared by all 12 full-attention layers (positions are layer-invariant).
+    let mropeFwd: ReturnType<typeof buildMropePositions> | null = null;
+    if (this.mrope) {
+      const t = this.config.text;
+      const ropeDims = Math.trunc(t.headDim * t.partialRotaryFactor);
+      const base = t.ropeParameters.full_attention?.ropeTheta ?? 10000;
+      this.#mropeInvFreq ??= mropeInvFreq(ropeDims, base);
+      mropeFwd = buildMropePositions(
+        this.mrope, cache[this.faIdx]!.offset, L, this.#mropeInvFreq, ropeDims,
+      );
+      setActiveMrope(mropeFwd);
+    }
+    let h: MlxArray | null = h0;
+    const prof = (globalThis as Record<string, unknown>).__deltaProf as
+      Record<string, number>
+      | undefined;
+    try {
+      for (let i = 0; i < this.layers.length; i++) {
+        const tl = prof ? performance.now() : 0;
+        const next = this.layers[i]!.forward(h, faMask, cache[i]!);
+        h.dispose();
+        h = next;
+        if (prof) {
+          ops.evalAll([h]);
+          const key = this.layers[i]!.isLinear ? "layerLin" : "layerFull";
+          prof[key] = (prof[key] ?? 0) + performance.now() - tl;
+        }
+        this.captureLayer(i, h); // native-MTP pre-final-norm tap (no-op unless set)
+      }
+      const out = disposing(h, this.finalNorm.forward(h));
+      h = null; // consumed — the finally must not double-free
+      return out;
+    } finally {
+      if (mropeFwd) {
+        setActiveMrope(null);
+        mropeFwd.posIds.dispose();
+      }
+      // A mid-loop layer throw must not strand the mask array or the
+      // in-flight [1,L,H] residual (2026-08-18 review).
+      faMask.arr?.dispose();
+      h?.dispose();
     }
-    faMask.arr?.dispose();
-    return disposing(h, this.finalNorm.forward(h));
   }
 
   logitsFromHidden(h: MlxArray): MlxArray {

package/src/native/frame_extract.swift

@@ -0,0 +1,140 @@
+// AVFoundation frame-extraction sidecar probe (PLAN 14w) — the video
+// equivalent of the audio path's afconvert doctrine: macOS-native codecs,
+// zero vendored dependencies.
+//
+//   frame-extract <video> <outDir> [fps=2] [maxFrames=768]
+//     → writes frame-%04d.png (RGB) + prints "N <count> W <w> H <h>"
+//   frame-extract --make-fixture <out.mov>
+//     → writes the deterministic 6-frame 320x240 gradient test clip
+//
+// Sampling mirrors mlx-vlm's do_sample_frames: timestamps k/fps for
+// k = 0..min(ceil(duration*fps), maxFrames)-1, precise tolerance.
+
+import AVFoundation
+import Foundation
+import ImageIO
+import UniformTypeIdentifiers
+
+func die(_ msg: String) -> Never {
+    FileHandle.standardError.write((msg + "\n").data(using: .utf8)!)
+    exit(1)
+}
+
+func makeFixture(_ path: String) {
+    let url = URL(fileURLWithPath: path)
+    try? FileManager.default.removeItem(at: url)
+    guard let writer = try? AVAssetWriter(outputURL: url, fileType: .mov) else {
+        die("AVAssetWriter init failed")
+    }
+    let W = 320, H = 240, FRAMES = 6
+    let settings: [String: Any] = [
+        AVVideoCodecKey: AVVideoCodecType.h264,
+        AVVideoWidthKey: W,
+        AVVideoHeightKey: H,
+    ]
+    let input = AVAssetWriterInput(mediaType: .video, outputSettings: settings)
+    let adaptor = AVAssetWriterInputPixelBufferAdaptor(
+        assetWriterInput: input,
+        sourcePixelBufferAttributes: [
+            kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_32ARGB,
+            kCVPixelBufferWidthKey as String: W,
+            kCVPixelBufferHeightKey as String: H,
+        ])
+    writer.add(input)
+    writer.startWriting()
+    writer.startSession(atSourceTime: .zero)
+    for f in 0..<FRAMES {
+        while !input.isReadyForMoreMediaData { Thread.sleep(forTimeInterval: 0.01) }
+        var pb: CVPixelBuffer?
+        CVPixelBufferCreate(nil, W, H, kCVPixelFormatType_32ARGB, nil, &pb)
+        guard let pixelBuffer = pb else { die("CVPixelBufferCreate failed") }
+        CVPixelBufferLockBaseAddress(pixelBuffer, [])
+        let base = CVPixelBufferGetBaseAddress(pixelBuffer)!
+        let bpr = CVPixelBufferGetBytesPerRow(pixelBuffer)
+        let p = base.assumingMemoryBound(to: UInt8.self)
+        for y in 0..<H {
+            for x in 0..<W {
+                let o = y * bpr + x * 4
+                p[o] = 255
+                p[o + 1] = UInt8((x * 255 / W + f * 20) % 256)
+                p[o + 2] = UInt8((y * 255 / H + f * 10) % 256)
+                p[o + 3] = UInt8(((x + y) * 255 / (W + H) + f * 30) % 256)
+            }
+        }
+        CVPixelBufferUnlockBaseAddress(pixelBuffer, [])
+        adaptor.append(pixelBuffer, withPresentationTime: CMTime(value: CMTimeValue(f), timescale: 2))
+    }
+    input.markAsFinished()
+    let sem = DispatchSemaphore(value: 0)
+    writer.finishWriting { sem.signal() }
+    sem.wait()
+    guard writer.status == .completed else { die("finishWriting status \(writer.status.rawValue)") }
+    print("fixture \(path)")
+}
+
+let args = CommandLine.arguments
+if args.count >= 3 && args[1] == "--make-fixture" {
+    makeFixture(args[2])
+    exit(0)
+}
+guard args.count >= 3 else { die("usage: frame-extract <video> <outDir> [fps] [maxFrames]") }
+let videoPath = args[1]
+let outDir = args[2]
+let fps = args.count > 3 ? Double(args[3]) ?? 2.0 : 2.0
+let maxFrames = args.count > 4 ? Int(args[4]) ?? 768 : 768
+
+let asset = AVURLAsset(url: URL(fileURLWithPath: videoPath))
+let sem = DispatchSemaphore(value: 0)
+var duration = CMTime.zero
+var loadError: Error?
+Task {
+    do { duration = try await asset.load(.duration) } catch { loadError = error }
+    sem.signal()
+}
+sem.wait()
+if let e = loadError { die("asset load failed: \(e.localizedDescription)") }
+// Fragmented/malformed containers can load with an indefinite duration —
+// CMTimeGetSeconds yields NaN and Int(NaN) is a Swift runtime TRAP. Fall
+// back to the frame cap; the extraction loop below already stops at the
+// first undecodable timestamp.
+let seconds = CMTimeGetSeconds(duration)
+let count = seconds.isFinite
+    ? max(1, min(Int((seconds * fps).rounded(.up)), maxFrames))
+    : maxFrames
+
+let gen = AVAssetImageGenerator(asset: asset)
+gen.appliesPreferredTrackTransform = true
+gen.requestedTimeToleranceBefore = .zero
+gen.requestedTimeToleranceAfter = .zero
+// Bound DECODED memory at the source: the language pipeline's video pixel
+// budget always downscales frames far below 1024px anyway (max_pixels
+// 32·32·768 across t·h·w), so capping the longest edge here loses nothing
+// it would have kept while keeping a 4K clip from materializing ~24 MB per
+// frame (768 frames ≈ 17.8 GiB) before preprocessing.
+gen.maximumSize = CGSize(width: 1024, height: 1024)
+
+try? FileManager.default.createDirectory(
+    atPath: outDir, withIntermediateDirectories: true)
+var written = 0
+var lastW = 0, lastH = 0
+for k in 0..<count {
+    let t = CMTime(seconds: Double(k) / fps, preferredTimescale: 600)
+    let img: CGImage
+    do {
+        img = try gen.copyCGImage(at: t, actualTime: nil)
+    } catch {
+        // Past-end sampling on short clips: stop at the last decodable frame.
+        break
+    }
+    lastW = img.width
+    lastH = img.height
+    let out = URL(fileURLWithPath: String(format: "%@/frame-%04d.png", outDir, k))
+    guard let dest = CGImageDestinationCreateWithURL(
+        out as CFURL, UTType.png.identifier as CFString, 1, nil)
+    else { die("png destination failed") }
+    CGImageDestinationAddImage(dest, img, nil)
+    guard CGImageDestinationFinalize(dest) else { die("png write failed") }
+    written += 1
+}
+guard written > 0 else { die("no frames decodable from \(videoPath)") }
+print("N \(written) W \(lastW) H \(lastH)")

package/src/native-pack.ts

@@ -15,18 +15,24 @@ import { existsSync, mkdirSync, renameSync, rmSync } from "node:fs";
 import { homedir, tmpdir } from "node:os";
 import { dirname, join } from "node:path";
 import { downloadOne } from "./download";
+import { runtimeValue } from "./runtime-config";
 
-export const NATIVE_PACK_VERSION = "0.2.0";
+export const NATIVE_PACK_VERSION = "0.3.0";
 export const NATIVE_PACK_FILES = [
   "libmlxc.dylib", "libmlx.dylib", "libjaccl.dylib", "mlx.metallib",
-  "libmlx_bun_expert_io.dylib",
+  "libmlx_bun_expert_io.dylib", "mlx-bun-frame-extract",
 ] as const;
 
+// RELEASE-TIME BAKE: scripts/build-native-pack.sh prints the sha256/size for
+// the freshly built tarball — publish `native-v0.3.0` with both assets and
+// paste the values here BEFORE tagging the package release (the download
+// 404s until the GitHub release exists; dev trees resolve via homebrew/env
+// and never hit this path).
 const SHA256: Record<string, string> = {
-  arm64: "9bd3795c5ea8f52b18413501f2d68c32c264a20751302300a08dc04cd67df97c",
+  arm64: "b485fed858971ea6f072f349a0234c0e2d5030d3ef2ddfa942c8975eca2d3b2e",
 };
 const SIZE: Record<string, number> = {
-  arm64: 52_307_647,
+  arm64: 52299265,
 };
 
 export function nativePackName(arch = process.arch): string {
@@ -37,7 +43,7 @@ export function nativePackName(arch = process.arch): string {
  *  the repo is private the download needs a token (GITHUB_TOKEN or
  *  `gh auth token`); once public it works anonymously. */
 export function nativePackUrl(arch = process.arch): string {
-  return process.env.MLX_BUN_NATIVE_PACK_URL ??
+  return runtimeValue("MLX_BUN_NATIVE_PACK_URL") ??
     `https://github.com/joshuarossi/mlx-bun/releases/download/native-v${NATIVE_PACK_VERSION}/${nativePackName(arch)}`;
 }
 
@@ -48,7 +54,8 @@ export function nativePackDir(arch = process.arch): string {
 /** Where the native runtime currently resolves from, or null if absent
  *  everywhere (= a fresh machine that needs the pack). */
 export function nativeRuntimeDir(): string | null {
-  if (process.env.MLX_BUN_LIBMLXC) return dirname(process.env.MLX_BUN_LIBMLXC);
+  const explicit = runtimeValue("MLX_BUN_LIBMLXC");
+  if (explicit) return dirname(explicit);
   const candidates = [
     dirname(process.execPath),
     nativePackDir(),

package/src/pi-web.ts

@@ -27,6 +27,7 @@ import { createHash } from "node:crypto";
 import { mkdirSync, rmSync } from "node:fs";
 import { homedir } from "node:os";
 import { join, resolve } from "node:path";
+import { runtimeValue } from "./runtime-config";
 import type { ServerWebSocket, WebSocketHandler } from "bun";
 import {
   createAgentSessionFromServices,
@@ -1394,7 +1395,7 @@ class PiWebSession {
           (d) => d.state === "active" && d.repoId !== this.opts.modelId,
         )?.repoId ?? null,
       }, { hasTools: WELCOME_TOOLS.length > 0 }) + surface.memoryHint;
-      if (process.env.MLX_BUN_PI_DEBUG) {
+      if (runtimeValue("MLX_BUN_PI_DEBUG")) {
         console.error(`[pi-web] prompt ${WEB_CHAT_PROMPT_VERSION} sha=${webChatPromptFingerprint(webPrompt)} memory=${surface.memoryEnabled ? "on" : "off"} codingTools=${this.codingToolsActive ? "on" : "off"}`);
       }
       const services = await createAgentSessionServices({
@@ -1826,7 +1827,7 @@ class PiWebSession {
 
     pi.on("tool_call", (event: ToolCallEvent) => {
       const block = decideBeforeToolCall(this.loopHygiene, event.toolName, event.input);
-      if (block && process.env.MLX_BUN_PI_DEBUG) {
+      if (block && runtimeValue("MLX_BUN_PI_DEBUG")) {
         console.error(`[pi-web] loop-hygiene blocked ${event.toolName}: ${block.reason.split("\n")[0]}`);
       }
       return block ? { block: true, reason: block.reason } : undefined;
@@ -1852,7 +1853,7 @@ class PiWebSession {
   private installApprovalGate(pi: ExtensionAPI): void {
     pi.on("tool_call", async (event: ToolCallEvent) => {
       const tool = event.toolName;
-      if (process.env.MLX_BUN_PI_DEBUG) {
+      if (runtimeValue("MLX_BUN_PI_DEBUG")) {
         console.error(`[pi-web] tool_call ${tool} args=${JSON.stringify(event.input)}`);
       }
 
@@ -1945,7 +1946,7 @@ class PiWebSession {
   }
 
   private onSessionEvent(event: AgentSessionEvent): void {
-    if (process.env.MLX_BUN_PI_DEBUG) {
+    if (runtimeValue("MLX_BUN_PI_DEBUG")) {
       const extra = event.type === "message_update"
         ? `/${(event as { assistantMessageEvent?: { type?: string } }).assistantMessageEvent?.type}` : "";
       console.error(`[pi-event] ${event.type}${extra}`);
@@ -2025,7 +2026,7 @@ class PiWebSession {
       // isStreaming atomically inside prompt(), so this read can't strand a
       // message at the turn boundary.
       case "prompt": {
-        if (process.env.MLX_BUN_PI_DEBUG) {
+        if (runtimeValue("MLX_BUN_PI_DEBUG")) {
           console.error(`[pi-web] prompt text=${JSON.stringify(msg.text.slice(0, 500))} images=${msg.images?.length ?? 0} streaming=${session.isStreaming} adapter=${this.selectedAdapter ?? "base"}`);
         }
         const images = toPiImages(msg.images);

package/src/prompt-cache.ts

@@ -24,7 +24,7 @@
 // reachable); onPut lets the server schedule its debounced write-behind
 // snapshot for both lanes.
 
-import { TurboQuantKVCache, type Cache } from "./model/gemma4";
+import type { Cache } from "./model/gemma4";
 import { cloneKvCaches } from "./kv-store";
 
 /** Reference-counted release: wraps an entry's `retain` (e.g. an mmap
@@ -112,7 +112,7 @@ export function cacheBytes(caches: Cache[]): number {
     // views per call (kv-store.ts snapshotCache/cloneKvCaches contract) —
     // those must be disposed here or they leak (see generate.ts
     // evalCacheState for the same hazard on the prefill path).
-    if (c instanceof TurboQuantKVCache) for (const a of state) a.dispose();
+    if (c.stateNeedsDispose) for (const a of state) a.dispose();
   }
   return total;
 }

package/src/quantize/atomic-output.ts

@@ -0,0 +1,38 @@
+import {
+  existsSync,
+  mkdirSync,
+  mkdtempSync,
+  readdirSync,
+  renameSync,
+  rmdirSync,
+  rmSync,
+  statSync,
+} from "node:fs";
+import { basename, dirname, join, resolve } from "node:path";
+
+/** Build a complete directory beside its destination, then publish it with a
+ * single same-filesystem rename. Failed writers leave no partial destination. */
+export async function writeAtomicDirectory<T>(
+  outDir: string,
+  write: (stagingDir: string) => Promise<T>,
+): Promise<T> {
+  const destination = resolve(outDir);
+  const replaceEmptyDestination = existsSync(destination);
+  if (
+    replaceEmptyDestination &&
+    (!statSync(destination).isDirectory() || readdirSync(destination).length !== 0)
+  ) throw new Error(`output directory already exists: ${outDir}`);
+
+  const parent = dirname(destination);
+  mkdirSync(parent, { recursive: true });
+  const staging = mkdtempSync(join(parent, `.${basename(destination)}.tmp-`));
+  try {
+    const result = await write(staging);
+    if (replaceEmptyDestination) rmdirSync(destination);
+    renameSync(staging, destination);
+    return result;
+  } catch (error) {
+    rmSync(staging, { recursive: true, force: true });
+    throw error;
+  }
+}

package/src/quantize/calibration.ts

@@ -20,6 +20,7 @@
 
 import { existsSync } from "node:fs";
 import type { LoadedTokenizer } from "../tokenizer";
+import { runtimeValue } from "../runtime-config";
 
 /** One JSONL calibration sample (subset of the OptIQ schema we can consume). */
 interface CalibrationSample {
@@ -34,7 +35,7 @@ function bundledMixCandidates(): string[] {
   const home = process.env.HOME ?? "";
   const rel = "lib/python3.14/site-packages/optiq/calibration/data/optiq.jsonl";
   return [
-    process.env.MLX_BUN_CALIBRATION_JSONL,
+    runtimeValue("MLX_BUN_CALIBRATION_JSONL"),
     `${home}/Code/mlx-lm-example/.venv/${rel}`,
     `${home}/Code/mlx-lm/.venv/${rel}`,
   ].filter((p): p is string => !!p);

package/src/quantize/config-writer.ts

@@ -76,6 +76,12 @@ export interface OptiqMetadata {
   group_size: number;
   achieved_bpw: number;
   per_layer_count: number;
+  weight_transforms?: Array<{
+    id: string;
+    seed: number;
+    family: string;
+    deviations: string[];
+  }>;
 }
 
 /**