mlx-bun 0.0.7 → 0.0.9

package/src/model/universal/modules.ts

@@ -0,0 +1,163 @@
+// Universal-tier primitives: dense (unquantized) linear/embedding modules
+// (Phase 1.5 — mlx nn.Linear / nn.Embedding semantics), quantized-or-dense
+// loader helpers, norm loaders, and the load-time weight audit.
+//
+// The quantized paths delegate to the proven gemma4-base primitives; the
+// dense paths are verbatim ports of mlx.nn.Linear (`mx.addmm(bias, x, W.T)`
+// / `x @ W.T`) and mlx.nn.Embedding (`weight[ids]`, `x @ weight.T`).
+
+import type { ModelConfig } from "../../config";
+import type { Weights } from "../../weights";
+import { MlxArray } from "../../mlx/array";
+import * as ops from "../../mlx/ops";
+import { QuantizedEmbedding, QuantizedLinear, RMSNorm } from "../gemma4-base";
+
+/** Records every tensor a universal load consumes; `finish` diffs against
+ *  the shard index so a descriptor mistake is a LOAD error (unconsumed /
+ *  missing tensors named), never a silently-wrong model
+ *  (docs/design/generic-model-support.md §3.4). */
+export class WeightAudit {
+  readonly consumed = new Set<string>();
+
+  use(name: string): void {
+    this.consumed.add(name);
+  }
+
+  /** `drop` = the arch's sanitize rules (tensors mlx-lm discards on load,
+   *  e.g. rotary_emb.inv_freq, or lm_head.weight under tied embeddings). */
+  finish(weights: Weights, drop: RegExp[]): void {
+    const unconsumed = weights.tensorNames.filter(
+      (n) => !this.consumed.has(n) && !drop.some((re) => re.test(n)),
+    );
+    if (unconsumed.length > 0)
+      throw new Error(
+        `weight audit: ${unconsumed.length} tensor(s) in the checkpoint were not consumed ` +
+        `by the universal module (descriptor mismatch?): ${unconsumed.slice(0, 12).join(", ")}` +
+        (unconsumed.length > 12 ? ", …" : ""),
+      );
+  }
+}
+
+/** Load a tensor and record it in the audit. */
+function tensorUsed(weights: Weights, audit: WeightAudit, name: string): MlxArray {
+  audit.use(name);
+  return weights.tensor(name);
+}
+
+/** Plain unquantized linear — mlx nn.Linear:
+ *  `mx.addmm(bias, x, weight.T)` when bias is present, else `x @ weight.T`. */
+export class DenseLinear {
+  readonly wT: MlxArray;
+
+  constructor(readonly w: MlxArray, readonly bias: MlxArray | null) {
+    this.wT = ops.transposeAxes(w, [1, 0]); // lazy view, shared across calls
+  }
+
+  get inFeatures(): number {
+    return this.w.shape[1]!;
+  }
+  get outFeatures(): number {
+    return this.w.shape[0]!;
+  }
+
+  forward(x: MlxArray): MlxArray {
+    return this.bias ? ops.addmm(this.bias, x, this.wT) : ops.matmul(x, this.wT);
+  }
+}
+
+/** Either linear flavor behind one call surface. */
+export type AnyLinear = QuantizedLinear | DenseLinear;
+
+/** Quantized when `.scales` exists (MLX-quantized checkpoints), else dense
+ *  bf16/f16/f32 (Phase 1.5). Both flavors carry the optional ADDITIVE
+ *  `.bias` term (qwen2 qkv, starcoder2, …). */
+export function loadLinear(
+  weights: Weights, path: string, config: ModelConfig, audit: WeightAudit,
+): AnyLinear {
+  if (weights.has(`${path}.scales`)) {
+    const lin = QuantizedLinear.load(weights, path, config);
+    audit.use(`${path}.weight`);
+    audit.use(`${path}.scales`);
+    if (weights.has(`${path}.biases`)) audit.use(`${path}.biases`);
+    if (weights.has(`${path}.bias`)) audit.use(`${path}.bias`);
+    return lin;
+  }
+  if (!weights.has(`${path}.weight`))
+    throw new Error(`${path}: no .weight tensor (nor .scales) in checkpoint`);
+  return new DenseLinear(
+    tensorUsed(weights, audit, `${path}.weight`),
+    weights.has(`${path}.bias`) ? tensorUsed(weights, audit, `${path}.bias`) : null,
+  );
+}
+
+/** Plain unquantized embedding — mlx nn.Embedding (+ as_linear tied head). */
+export class DenseEmbedding {
+  readonly wT: MlxArray;
+
+  constructor(readonly w: MlxArray) {
+    this.wT = ops.transposeAxes(w, [1, 0]);
+  }
+
+  encode(ids: MlxArray): MlxArray {
+    return ops.takeAxis(this.w, ids, 0);
+  }
+
+  /** nn.Embedding.as_linear: `x @ weight.T`. */
+  asLinear(h: MlxArray): MlxArray {
+    return ops.matmul(h, this.wT);
+  }
+}
+
+export type AnyEmbedding = QuantizedEmbedding | DenseEmbedding;
+
+export function loadEmbedding(
+  weights: Weights, path: string, config: ModelConfig, audit: WeightAudit,
+): AnyEmbedding {
+  if (weights.has(`${path}.scales`)) {
+    const emb = QuantizedEmbedding.load(weights, path, config);
+    audit.use(`${path}.weight`);
+    audit.use(`${path}.scales`);
+    if (weights.has(`${path}.biases`)) audit.use(`${path}.biases`);
+    return emb;
+  }
+  return new DenseEmbedding(tensorUsed(weights, audit, `${path}.weight`));
+}
+
+/** RMSNorm loader; `plusOne` = the gemma-family `1.0 + weight` variant
+ *  (the add is folded once at load — identical value to mlx-lm's
+ *  per-call `1.0 + self.weight`). */
+export function loadRmsNorm(
+  weights: Weights, path: string, eps: number, plusOne: boolean, audit: WeightAudit,
+): RMSNorm {
+  const w = tensorUsed(weights, audit, `${path}.weight`);
+  if (!plusOne) return new RMSNorm(w, eps);
+  const one = ops.scalarLike(1, w);
+  const w1 = ops.add(one, w);
+  one.dispose();
+  return new RMSNorm(w1, eps);
+}
+
+/** mlx nn.LayerNorm (weight + bias affine — starcoder2's norm). */
+export class LayerNorm {
+  constructor(
+    readonly weight: MlxArray,
+    readonly bias: MlxArray | null,
+    readonly eps: number,
+  ) {}
+
+  forward(x: MlxArray): MlxArray {
+    return ops.layerNorm(x, this.weight, this.bias, this.eps);
+  }
+}
+
+export type AnyNorm = RMSNorm | LayerNorm;
+
+export function loadLayerNorm(
+  weights: Weights, path: string, eps: number, audit: WeightAudit,
+): LayerNorm {
+  return new LayerNorm(
+    tensorUsed(weights, audit, `${path}.weight`),
+    weights.has(`${path}.bias`) ? tensorUsed(weights, audit, `${path}.bias`) : null,
+    eps,
+  );
+}

package/src/model/universal/rope.ts

@@ -0,0 +1,366 @@
+// Universal rope factory — faithful port of mlx_lm/models/rope_utils.py
+// (oracle venv, mlx-lm 0.31.3): initialize_rope with the default / linear /
+// llama3 (Llama3RoPE) / yarn (YarnRoPE) / longrope (SuScaledRoPE) /
+// proportional (ProportionalRoPE) variants.
+//
+// Porting discipline (three-level fidelity tree, L1): every frequency
+// table is computed with the SAME mlx float32 ops the oracle uses at
+// module init — not host-side float64 math — so the freqs bytes match the
+// oracle's bit-for-bit. Host math (Math.log/floor/ceil) is used only where
+// the oracle uses python-float math (yarn correction range, mscale, su
+// scale), which is IEEE double on both sides.
+//
+// The existing runtime paths (ops.rope default, gemma4's proportional
+// freqs) are untouched; this factory serves the universal (Tier-0) models.
+
+import { MlxArray } from "../../mlx/array";
+import { Dtype } from "../../mlx/ffi";
+import * as ops from "../../mlx/ops";
+
+/** Raw HF `rope_scaling` dict (values read exactly like the oracle). */
+export type RopeScalingConfig = Record<string, unknown>;
+
+const num = (v: unknown, fallback: number): number =>
+  typeof v === "number" ? v : fallback;
+
+const requireNum = (sc: RopeScalingConfig, key: string): number => {
+  const v = sc[key];
+  if (typeof v !== "number") throw new Error(`rope_scaling.${key} missing or not a number`);
+  return v;
+};
+
+/** One rope instance (shared across layers — identical params ⇒ identical
+ *  values). `apply` returns a NEW array; the caller owns disposal of both. */
+export class UniversalRope {
+  constructor(
+    readonly dims: number,
+    readonly traditional: boolean,
+    /** Base for native fast_rope; null when `freqs` supplies the table. */
+    readonly base: number | null,
+    /** fast_rope scale (1/factor for "linear", else 1). */
+    readonly scale: number,
+    /** Precomputed frequency table (llama3/yarn/longrope/proportional). */
+    readonly freqs: MlxArray | null,
+    /** Pre-rope input scale on x[..., :dims] (yarn mscale / su scale). */
+    readonly preScale: number,
+    /** Apply preScale even when it is exactly 1.0 (SuScaledRoPE does; Yarn
+     *  skips when mscale == 1.0). */
+    readonly preScaleAlways: boolean,
+  ) {}
+
+  apply(x: MlxArray, offset: number): MlxArray {
+    let input = x;
+    let scaled: MlxArray | null = null;
+    if (this.preScaleAlways || this.preScale !== 1.0) {
+      // x[..., :dims] = preScale * x[..., :dims] (weak-scalar mul)
+      const D = x.shape[x.shape.length - 1]!;
+      if (this.dims === D) {
+        scaled = ops.mulScalar(x, this.preScale);
+      } else {
+        const start = x.shape.map(() => 0);
+        const stop = [...x.shape];
+        stop[stop.length - 1] = this.dims;
+        const head = x.slice(start, stop);
+        const headScaled = ops.mulScalar(head, this.preScale);
+        head.dispose();
+        scaled = ops.sliceUpdate(x, headScaled, start, stop);
+        headScaled.dispose();
+      }
+      input = scaled;
+    }
+    const out = ops.ropeScaled(
+      input, this.dims, this.traditional, this.base, this.scale, offset, this.freqs,
+    );
+    scaled?.dispose();
+    return out;
+  }
+
+  dispose(): void {
+    this.freqs?.dispose();
+  }
+}
+
+/** SmolLM3 NoPE: identity stand-in for disabled-rope layers. */
+export class NoRope extends UniversalRope {
+  constructor() {
+    super(0, false, null, 1, null, 1, false);
+  }
+  override apply(x: MlxArray, _offset: number): MlxArray {
+    // Return an owned copy-view so callers can dispose uniformly.
+    return x.slice(x.shape.map(() => 0), [...x.shape]);
+  }
+}
+
+/** `base ** (arange(0, dims, 2) / dims)` with the oracle's dtypes:
+ *  int32 arange for llama3 (no explicit dtype in the source), float32
+ *  arange for yarn/longrope. Integer÷integer divides to float32 in mlx,
+ *  so both spellings compute the pow in f32. */
+function baseFreqs(dims: number, base: number, arangeDtype: Dtype): MlxArray {
+  const ar = ops.arange(0, dims, 2, arangeDtype);
+  const dimsScalar = arangeDtype === Dtype.int32
+    ? ops.fromInt32([dims], [])
+    : MlxArray.fromFloat32(new Float32Array([dims]), []);
+  const exps = ops.div(ar, dimsScalar); // → float32
+  ar.dispose();
+  dimsScalar.dispose();
+  const baseScalar = ops.scalarLike(base, exps);
+  const out = ops.pow(baseScalar, exps);
+  baseScalar.dispose();
+  exps.dispose();
+  return out; // float32 [dims/2]
+}
+
+/** Llama3RoPE frequency table (rope_utils.py:74-108, verbatim). */
+function llama3Freqs(dims: number, base: number, sc: RopeScalingConfig): MlxArray {
+  const factor = requireNum(sc, "factor");
+  const lowFreqFactor = num(sc.low_freq_factor, 1.0);
+  const highFreqFactor = num(sc.high_freq_factor, 4.0);
+  const oldContextLen = num(sc.original_max_position_embeddings, 8192);
+
+  const lowFreqWavelen = oldContextLen / lowFreqFactor;
+  const highFreqWavelen = oldContextLen / highFreqFactor;
+
+  let freqs = baseFreqs(dims, base, Dtype.int32);
+  const wavelens = ops.mulScalar(freqs, 2 * Math.PI);
+
+  // freqs = where(wavelens > low_freq_wavelen, freqs * factor, freqs)
+  const lfw = ops.scalarLike(lowFreqWavelen, wavelens);
+  const gtLow = ops.less(lfw, wavelens);
+  const freqsTimesFactor = ops.mulScalar(freqs, factor);
+  const freqs1 = ops.where(gtLow, freqsTimesFactor, freqs);
+  for (const a of [gtLow, freqsTimesFactor, freqs]) a.dispose();
+  freqs = freqs1;
+
+  // is_medium = (wavelens > high_freq_wavelen) & (wavelens < low_freq_wavelen)
+  const hfw = ops.scalarLike(highFreqWavelen, wavelens);
+  const gtHigh = ops.less(hfw, wavelens);
+  const ltLow = ops.less(wavelens, lfw);
+  const isMedium = ops.logicalAnd(gtHigh, ltLow);
+  for (const a of [hfw, lfw, gtHigh, ltLow]) a.dispose();
+
+  // smooth_factors = (old_context_len / wavelens - low) / (high - low)
+  const oldScalar = ops.scalarLike(oldContextLen, wavelens);
+  const oldOverWav = ops.div(oldScalar, wavelens);
+  oldScalar.dispose();
+  wavelens.dispose();
+  const lowScalar = ops.scalarLike(lowFreqFactor, oldOverWav);
+  const numArr = ops.sub(oldOverWav, lowScalar);
+  oldOverWav.dispose();
+  lowScalar.dispose();
+  const denScalar = ops.scalarLike(highFreqFactor - lowFreqFactor, numArr);
+  const smooth = ops.div(numArr, denScalar);
+  numArr.dispose();
+  denScalar.dispose();
+
+  // smooth_freqs = freqs / ((1 - smooth) / factor + smooth)
+  const one = ops.scalarLike(1, smooth);
+  const oneMinus = ops.sub(one, smooth);
+  one.dispose();
+  const factorScalar = ops.scalarLike(factor, oneMinus);
+  const scaledDown = ops.div(oneMinus, factorScalar);
+  oneMinus.dispose();
+  factorScalar.dispose();
+  const denom = ops.add(scaledDown, smooth);
+  scaledDown.dispose();
+  smooth.dispose();
+  const smoothFreqs = ops.div(freqs, denom);
+  denom.dispose();
+
+  const out = ops.where(isMedium, smoothFreqs, freqs);
+  for (const a of [isMedium, smoothFreqs, freqs]) a.dispose();
+  return out;
+}
+
+/** YarnRoPE (rope_utils.py:128-181, verbatim): returns [freqs, mscale]. */
+function yarnFreqs(
+  dims: number,
+  base: number,
+  scalingFactor: number,
+  originalMaxPositionEmbeddings: number,
+  betaFast: number,
+  betaSlow: number,
+  mscale: number,
+  mscaleAllDim: number,
+): [MlxArray, number] {
+  const findCorrectionDim = (numRotations: number): number =>
+    (dims * Math.log(originalMaxPositionEmbeddings / (numRotations * 2 * Math.PI))) /
+    (2 * Math.log(base));
+  let low = Math.floor(findCorrectionDim(betaFast));
+  let high = Math.ceil(findCorrectionDim(betaSlow));
+  low = Math.max(low, 0);
+  high = Math.min(high, dims - 1);
+  const getMscale = (scale: number, m: number): number =>
+    scale <= 1 ? 1.0 : 0.1 * m * Math.log(scale) + 1.0;
+  const outMscale =
+    getMscale(scalingFactor, mscale) / getMscale(scalingFactor, mscaleAllDim);
+
+  const freqExtra = baseFreqs(dims, base, Dtype.float32);
+  const freqInter = ops.mulScalar(freqExtra, scalingFactor);
+
+  // freq_mask = 1 - linear_ramp_mask(low, high, dims // 2)
+  let maxVal = high;
+  if (low === maxVal) maxVal += 0.001; // prevent singularity (oracle comment)
+  const arangeHalf = ops.arange(0, Math.floor(dims / 2), 1, Dtype.float32);
+  const minScalar = ops.scalarLike(low, arangeHalf);
+  const shifted = ops.sub(arangeHalf, minScalar);
+  arangeHalf.dispose();
+  minScalar.dispose();
+  const rangeScalar = ops.scalarLike(maxVal - low, shifted);
+  const linear = ops.div(shifted, rangeScalar);
+  shifted.dispose();
+  rangeScalar.dispose();
+  const lo = ops.scalarLike(0, linear);
+  const hi = ops.scalarLike(1, linear);
+  const ramp = ops.clip(linear, lo, hi);
+  for (const a of [linear, lo, hi]) a.dispose();
+  const one = ops.scalarLike(1, ramp);
+  const freqMask = ops.sub(one, ramp);
+  ramp.dispose();
+
+  // freqs = (inter * extra) / (inter * mask + extra * (1 - mask))
+  const numArr = ops.mul(freqInter, freqExtra);
+  const interMasked = ops.mul(freqInter, freqMask);
+  const oneMinusMask = ops.sub(one, freqMask);
+  one.dispose();
+  const extraMasked = ops.mul(freqExtra, oneMinusMask);
+  oneMinusMask.dispose();
+  const den = ops.add(interMasked, extraMasked);
+  const freqs = ops.div(numArr, den);
+  for (const a of [freqExtra, freqInter, freqMask, numArr, interMasked, extraMasked, den])
+    a.dispose();
+  return [freqs, outMscale];
+}
+
+/** SuScaledRoPE (rope_utils.py:10-71): [freqs, scale]. NOTE the oracle uses
+ *  long_factor for ALL positions (short_factor is accepted and ignored). */
+function suFreqsAndScale(
+  dims: number,
+  base: number,
+  maxPositionEmbeddings: number,
+  originalMaxPositionEmbeddings: number,
+  longFactor: number[] | number,
+  longMscale?: number,
+): [MlxArray, number] {
+  const freqs0 = baseFreqs(dims, base, Dtype.float32);
+  const lf = Array.isArray(longFactor)
+    ? MlxArray.fromFloat32(new Float32Array(longFactor), [longFactor.length])
+    : MlxArray.fromFloat32(new Float32Array([longFactor]), []);
+  const freqs = ops.mul(lf, freqs0);
+  lf.dispose();
+  freqs0.dispose();
+
+  const factor = maxPositionEmbeddings / originalMaxPositionEmbeddings;
+  const defaultScale =
+    factor <= 1.0
+      ? 1.0
+      : Math.sqrt(1 + Math.log(factor) / Math.log(originalMaxPositionEmbeddings));
+  // python: `long_mscale or default` — 0 is falsy there too.
+  const scale = longMscale ? longMscale : defaultScale;
+  return [freqs, scale];
+}
+
+/** ProportionalRoPE (rope_utils.py:199-232): rotate the first rotated_dims,
+ *  ∞-freq (identity) padding for the rest. */
+function proportionalFreqs(
+  dims: number, rotatedDims: number, base: number, factor: number,
+): MlxArray {
+  if (rotatedDims > dims) throw new Error("rotated_dims should be smaller than dims");
+  const ar = ops.arange(0, rotatedDims, 2, Dtype.float32);
+  const dimsScalar = MlxArray.fromFloat32(new Float32Array([dims]), []);
+  const exps = ops.div(ar, dimsScalar);
+  ar.dispose();
+  dimsScalar.dispose();
+  const baseScalar = ops.scalarLike(base, exps);
+  const powed = ops.pow(baseScalar, exps);
+  baseScalar.dispose();
+  exps.dispose();
+  const scaled = ops.mulScalar(powed, factor);
+  powed.dispose();
+  const padLen = Math.floor((dims - rotatedDims) / 2);
+  if (padLen === 0) return scaled;
+  const pad = MlxArray.fromFloat32(
+    new Float32Array(padLen).fill(Number.POSITIVE_INFINITY), [padLen],
+  );
+  const out = ops.concatAxis([scaled, pad], 0);
+  scaled.dispose();
+  pad.dispose();
+  return out;
+}
+
+/** Port of rope_utils.initialize_rope — the shared entry every
+ *  initialize_rope-style arch (llama, qwen2/3, olmo2, granite, …) uses. */
+export function initializeRope(
+  dims: number,
+  base: number,
+  traditional: boolean,
+  scalingConfig: RopeScalingConfig | null,
+  maxPositionEmbeddings: number | null,
+): UniversalRope {
+  const ropeType = scalingConfig
+    ? String(scalingConfig.type ?? scalingConfig.rope_type ?? "default")
+    : "default";
+
+  if (ropeType === "default" || ropeType === "linear") {
+    const scale = ropeType === "linear" ? 1 / requireNum(scalingConfig!, "factor") : 1.0;
+    return new UniversalRope(dims, traditional, base, scale, null, 1, false);
+  }
+
+  if (ropeType === "llama3") {
+    const freqs = llama3Freqs(dims, base, scalingConfig!);
+    return new UniversalRope(dims, traditional, null, 1.0, freqs, 1, false);
+  }
+
+  if (ropeType === "yarn" || ropeType === "deepseek_yarn" || ropeType === "telechat3-yarn") {
+    const sc = scalingConfig!;
+    const [freqs, mscale] = yarnFreqs(
+      dims,
+      base,
+      requireNum(sc, "factor"),
+      num(sc.original_max_position_embeddings, 4096),
+      num(sc.beta_fast, 32),
+      num(sc.beta_slow, 1),
+      num(sc.mscale, 1),
+      num(sc.mscale_all_dim, 0),
+    );
+    // YarnRoPE scales x only when mscale != 1.0 (preScaleAlways=false).
+    return new UniversalRope(dims, traditional, null, 1.0, freqs, mscale, false);
+  }
+
+  if (ropeType === "longrope") {
+    const sc = scalingConfig!;
+    const [freqs, scale] = suFreqsAndScale(
+      dims,
+      base,
+      maxPositionEmbeddings ?? 131072,
+      requireNum(sc, "original_max_position_embeddings"),
+      (sc.long_factor as number[] | number | undefined) ?? 1.0,
+    );
+    // SuScaledRoPE hardcodes traditional=False and ALWAYS pre-scales.
+    return new UniversalRope(dims, false, null, 1.0, freqs, scale, true);
+  }
+
+  if (ropeType === "proportional") {
+    const sc = scalingConfig!;
+    const rotatedDims = Math.floor(dims * num(sc.partial_rotary_factor, 1.0));
+    const freqs = proportionalFreqs(dims, rotatedDims, base, num(sc.factor, 1.0));
+    return new UniversalRope(dims, traditional, null, 1.0, freqs, 1, false);
+  }
+
+  throw new Error(`Unsupported RoPE type ${ropeType}`);
+}
+
+/** SuScaledRoPE constructed the phi3.py way (NOT via initialize_rope):
+ *  original_max_position_embeddings comes from the TOP-LEVEL config, and
+ *  rope_dim = head_dim * partial_rotary_factor. */
+export function phi3SuRope(
+  ropeDim: number,
+  base: number,
+  maxPositionEmbeddings: number,
+  originalMaxPositionEmbeddings: number,
+  longFactor: number[] | number,
+): UniversalRope {
+  const [freqs, scale] = suFreqsAndScale(
+    ropeDim, base, maxPositionEmbeddings, originalMaxPositionEmbeddings, longFactor,
+  );
+  return new UniversalRope(ropeDim, false, null, 1.0, freqs, scale, true);
+}

package/src/pi-provider.ts

@@ -57,7 +57,7 @@ export interface PiProviderWiring {
 
 /**
  * Build the in-memory auth + registry and resolve the local model for a pi
- * session pointed at `baseUrl` (e.g. "http://127.0.0.1:8090/v1"). The runtime
+ * session pointed at `baseUrl` (e.g. "http://127.0.0.1:8080/v1"). The runtime
  * key satisfies pi's auth check for our zero-cost local provider; no disk
  * (no models.json, no ~/.pi cross-talk).
  */

package/src/pi-terminal.ts

@@ -77,7 +77,7 @@ Be honest about what you did and what you found, including mistakes, dead ends,
 }
 
 export interface RunEmbeddedPiOptions {
-  /** mlx-bun server base URL, e.g. "http://localhost:8090/v1". */
+  /** mlx-bun server base URL, e.g. "http://localhost:8080/v1". */
   baseUrl: string;
   /** Real served model id, used to label the model in the prompt + registry. */
   modelLabel?: string;
@@ -85,6 +85,9 @@ export interface RunEmbeddedPiOptions {
   contextWindow?: number;
   /** Whether the served model has a switchable thinking channel. */
   reasoning?: boolean;
+  /** Whether the served model accepts images (server's vision capability —
+   *  /v1/models `vision`). MUST match the real model: see PiProviderOptions. */
+  vision?: boolean;
   /** Run mode. Default: "interactive". */
   mode?: PiTerminalMode;
   /** Print mode output format ("text" final-only, "json" event stream). */
@@ -115,6 +118,7 @@ export async function runEmbeddedPi(opts: RunEmbeddedPiOptions): Promise<number>
     contextWindow: opts.contextWindow ?? DEFAULT_CONTEXT_WINDOW,
     name: opts.modelLabel ? `${opts.modelLabel} (mlx-bun local)` : undefined,
     reasoning: opts.reasoning ?? false,
+    vision: opts.vision ?? false,
   });
 
   // Shared pi tool/skill surface. Memory is included only when the user has
@@ -202,6 +206,11 @@ export interface ParsedPiArgs {
   printFormat: "text" | "json";
   /** Joined free-text message (the trailing words / `-p <text>`). */
   message?: string;
+  /** pi's verbose startup banner (interactive only). */
+  verbose: boolean;
+  /** Unknown flags that were dropped (with their values) — callers should
+   *  warn so `--resume`-style pi flags don't fail silently. */
+  ignored: string[];
 }
 
 /**
@@ -222,32 +231,44 @@ export function parsePiArgs(passthrough: string[], isStdinPiped: boolean): Parse
   let print = false;
   let rpc = false;
   let json = false;
+  let verbose = false;
   const words: string[] = [];
+  const ignored: string[] = [];
   for (let i = 0; i < passthrough.length; i++) {
     const a = passthrough[i]!;
     if (a === "-p" || a === "--print") {
       print = true;
     } else if (a === "--json") {
       json = true;
+    } else if (a === "--verbose") {
+      verbose = true;
     } else if (a === "--mode") {
       const m = passthrough[++i];
       if (m === "rpc") rpc = true;
       else if (m === "json") { print = true; json = true; }
       else if (m === "text") print = true;
     } else if (a.startsWith("-")) {
-      // Unknown flag for the built-in agent: ignore it (the full pi flag
-      // surface lives in the user's own pi via `mlx-bun harness pi`). Don't
-      // treat it as message text.
-      continue;
+      // Unknown flag for the built-in agent: drop it AND its value (the full
+      // pi flag surface lives in the user's own pi via `mlx-bun harness pi`).
+      // Without the value-swallow, `mlx-bun pi --resume abc123` started a
+      // chat with the message "abc123". A following token that looks like
+      // another flag is not a value, so it isn't consumed.
+      const next = passthrough[i + 1];
+      if (next !== undefined && !next.startsWith("-")) {
+        ignored.push(`${a} ${next}`);
+        i++;
+      } else {
+        ignored.push(a);
+      }
     } else {
       words.push(a);
     }
   }
   const message = words.length > 0 ? words.join(" ") : undefined;
   const printFormat = json ? "json" : "text";
-  if (rpc) return { mode: "rpc", printFormat: "text", message };
+  if (rpc) return { mode: "rpc", printFormat: "text", message, verbose, ignored };
   // `-p`/`--mode` request print; piped stdin makes a bare invocation one-shot
   // too, matching pi's own CLI (it reads stdin → switches to print mode).
-  if (print || isStdinPiped) return { mode: "print", printFormat, message };
-  return { mode: "interactive", printFormat, message };
+  if (print || isStdinPiped) return { mode: "print", printFormat, message, verbose, ignored };
+  return { mode: "interactive", printFormat, message, verbose, ignored };
 }

package/src/pi-web.ts

@@ -67,6 +67,21 @@ const GATED_TOOLS = new Set(["bash", "edit", "write"]);
  */
 const WELCOME_TOOLS = ["read", "web_search"] as const;
 
+/**
+ * The web chat session's tool allowlist: the two welcome tools, plus the
+ * memory/reference read tools whenever the surface has memory enabled. The
+ * system prompt (surface.memoryHint) and the bundled memory skill instruct
+ * the model to call memory_resolve/memory_read/… — pi treats `tools` as an
+ * allowlist, so leaving them out makes every such call fail. All additions
+ * are in READ_ONLY_TOOLS (no approval round-trip).
+ */
+export function webChatToolAllowlist(memoryEnabled: boolean): string[] {
+  return [
+    ...WELCOME_TOOLS,
+    ...(memoryEnabled ? [...MEMORY_TOOL_NAMES, ...REFERENCE_TOOL_NAMES] : []),
+  ];
+}
+
 /** Auto-deny a pending approval after this long with no browser decision. */
 const APPROVAL_TIMEOUT_MS = 120_000;
 
@@ -564,8 +579,10 @@ class PiWebSession {
           // Allowlist (not surface.tools): `read` is a pi built-in enabled by
           // name; `web_search` is the custom tool from surface.customTools. pi
           // exposes only names in this list, so web_fetch/weather stay defined
-          // but hidden.
-          tools: [...WELCOME_TOOLS],
+          // but hidden. Memory/reference tools ride along when the surface has
+          // memory enabled — the prompt + memory skill advertise them, so they
+          // must be callable (webChatToolAllowlist).
+          tools: webChatToolAllowlist(surface.memoryEnabled),
           customTools: surface.customTools,
         })),
         services,

package/src/quantize/config-writer.ts

@@ -109,7 +109,7 @@ export async function writeQuantizedConfig(
 }
 
 /** Copy known aux files (and any `*.model`) from src → out when present. */
-async function copyAuxFiles(srcDir: string, outDir: string): Promise<void> {
+export async function copyAuxFiles(srcDir: string, outDir: string): Promise<void> {
   const seen = new Set<string>();
   for (const f of AUX_FILES) {
     const src = join(srcDir, f);