mlx-bun 0.0.12 → 0.2.0

package/src/quantize/sensitivity.ts

@@ -26,7 +26,10 @@ import { cpuStream, MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
 import { trainForward } from "../train/forward";
 import type { RuntimeModel } from "../model/factory";
-import type { QuantizedLinear } from "../model/gemma4-base";
+import type {
+  QuantizedLinear,
+  QuantizedLinearState,
+} from "../model/gemma4-base";
 
 /** Per-layer KL sensitivity at each candidate bit-width (port of
  *  optiq SensitivityResult). */
@@ -43,46 +46,23 @@ export interface AnalyzeSensitivityOptions {
 }
 
 // --------------------------------------------------------------------------
-// Minimal in-place mutation of a QuantizedLinear.
-//
-// gemma4-base.ts declares QuantizedLinear's w/scales/biases/spec as `readonly`
-// (they never change at inference time). The sensitivity sweep needs to swap a
-// layer's quantized weight transiently. Rather than edit the model file, we
-// poke the public fields here via a typed cast — a quantize-only mutation
-// helper. The fields are public, so this is a write to public state, not a
-// reach into private internals. Flagged in the port report.
+// Scoped state exchange for sensitivity probes. QuantizedLinear owns the only
+// mutation capability; inference consumers retain a read-only getter surface.
 // --------------------------------------------------------------------------
 
-/** Mutable view of the public QuantizedLinear fields we swap. */
-type MutableQuantizedLinear = {
-  w: MlxArray;
-  scales: MlxArray;
-  biases: MlxArray | null;
-  spec: ops.QuantSpec;
-};
-
 /** Snapshot of a QuantizedLinear's quantized state (for restore). */
-interface LayerState {
-  w: MlxArray;
-  scales: MlxArray;
-  biases: MlxArray | null;
-  spec: ops.QuantSpec;
-}
+type LayerState = QuantizedLinearState;
 
 function captureState(layer: QuantizedLinear): LayerState {
   return { w: layer.w, scales: layer.scales, biases: layer.biases, spec: layer.spec };
 }
 
 function restoreState(layer: QuantizedLinear, state: LayerState): void {
-  const m = layer as unknown as MutableQuantizedLinear;
+  const transient = layer.exchangeQuantizedState(state);
   // Dispose the transient swapped-in tensors before restoring the originals.
-  if (m.w !== state.w) m.w.dispose();
-  if (m.scales !== state.scales) m.scales.dispose();
-  if (m.biases && m.biases !== state.biases) m.biases.dispose();
-  m.w = state.w;
-  m.scales = state.scales;
-  m.biases = state.biases;
-  m.spec = state.spec;
+  if (transient.w !== state.w) transient.w.dispose();
+  if (transient.scales !== state.scales) transient.scales.dispose();
+  if (transient.biases && transient.biases !== state.biases) transient.biases.dispose();
 }
 
 /** Quantize `bf16Weight` at `newBits`/`groupSize` and swap it into `layer`
@@ -97,15 +77,16 @@ function mutateLayerToBits(
   original: LayerState,
 ): void {
   const q = ops.quantize(bf16Weight, groupSize, newBits, "affine", cpuStream);
-  const m = layer as unknown as MutableQuantizedLinear;
+  const previous = layer.exchangeQuantizedState({
+    w: q.packed,
+    scales: q.scales,
+    biases: q.biases,
+    spec: { bits: newBits, groupSize, mode: "affine" },
+  });
   // Free any previously swapped-in transient tensors (not the original).
-  if (m.w !== original.w) m.w.dispose();
-  if (m.scales !== original.scales) m.scales.dispose();
-  if (m.biases && m.biases !== original.biases) m.biases.dispose();
-  m.w = q.packed;
-  m.scales = q.scales;
-  m.biases = q.biases;
-  m.spec = { bits: newBits, groupSize, mode: "affine" };
+  if (previous.w !== original.w) previous.w.dispose();
+  if (previous.scales !== original.scales) previous.scales.dispose();
+  if (previous.biases && previous.biases !== original.biases) previous.biases.dispose();
   ops.evalAll([q.packed, q.scales, q.biases]);
 }
 
@@ -193,29 +174,30 @@ export function analyzeSensitivityExact(
         ops.evalAll([bf16Source]);
 
         const sensitivities: Record<number, number> = {};
-
-        for (const bits of candidateBits) {
-          if (bits === baselineBits) {
-            // KL is 0 by construction (re-quantizing to the baseline bit-width
-            // reproduces the running model). Skip the forwards.
-            sensitivities[bits] = 0.0;
-            continue;
+        try {
+          for (const bits of candidateBits) {
+            if (bits === baselineBits) {
+              // KL is 0 by construction (re-quantizing to the baseline bit-width
+              // reproduces the running model). Skip the forwards.
+              sensitivities[bits] = 0.0;
+              continue;
+            }
+            mutateLayerToBits(layer, bf16Source, bits, groupSize, original);
+
+            let klSum = 0;
+            for (let i = 0; i < calArrays.length; i++) {
+              const cur = trainForward(model, calArrays[i]!);
+              klSum += klFromRef(cur, refLogits[i]!);
+              cur.dispose();
+            }
+            sensitivities[bits] = klSum / Math.max(calArrays.length, 1);
           }
-          mutateLayerToBits(layer, bf16Source, bits, groupSize, original);
-
-          let klSum = 0;
-          for (let i = 0; i < calArrays.length; i++) {
-            const cur = trainForward(model, calArrays[i]!);
-            klSum += klFromRef(cur, refLogits[i]!);
-            cur.dispose();
-          }
-          sensitivities[bits] = klSum / Math.max(calArrays.length, 1);
+        } finally {
+          // Restore even when quantization or a calibration forward throws.
+          restoreState(layer, original);
+          bf16Source.dispose();
         }
 
-        // Restore the original quantized weight; release the bf16 source.
-        restoreState(layer, original);
-        bf16Source.dispose();
-
         // param count = product of the dequantized (full) weight shape.
         const fullShape = [layer.outFeatures, layer.inFeatures];
         const paramCount = fullShape.reduce((a, b) => a * b, 1);

package/src/quantize/weight-transform.ts

@@ -0,0 +1,292 @@
+// Quantizer-facing offline weight-transform seam. Recipe math lives in
+// rotate.ts; this module owns planning, family adaptation, and tensor ownership
+// so the artifact writer has one contract for every fold implementation.
+
+import type { ModelConfig } from "../config";
+import type { MlxArray } from "../mlx/array";
+import type { Weights } from "../weights";
+import {
+  foldLlamaWeights,
+  planQwen35Fold,
+  planQwenMtpFold,
+  QwenFoldContext,
+  type FoldGeometry,
+  type FoldOptions,
+  type QwenFoldOptions,
+  type QwenFoldPlan,
+} from "./rotate";
+
+/** Serialized alongside a quantized artifact so a folded basis is a property
+ *  of the artifact, not an out-of-band experiment-script convention. */
+export interface WeightTransformMetadata {
+  id: string;
+  seed: number;
+  family: string;
+  deviations: string[];
+}
+
+/** Pure output plan. `sourceByOutput` makes cloning/untie operations explicit:
+ *  most outputs map to themselves; a synthesized head maps to its embedding. */
+export interface WeightTransformPlan {
+  readonly id: string;
+  readonly outputNames: readonly string[];
+  readonly sourceByOutput: ReadonlyMap<string, string>;
+  readonly untieWordEmbeddings: boolean;
+  readonly metadata: WeightTransformMetadata;
+}
+
+/** Per-run executor. `apply` returns a lazy, owned array and never disposes
+ *  `source`; the quantizer owns the result and the Weights object owns source. */
+export interface WeightTransformContext {
+  apply(outputName: string, source: MlxArray): MlxArray;
+  dispose(): void;
+}
+
+/** Common seam for the Llama, Qwen3.5 trunk, and Qwen MTP fold recipes. */
+export interface WeightTransform {
+  readonly id: string;
+  /** Name/config analysis only: no mlx arrays, model load, or device work. */
+  plan(names: readonly string[], config: ModelConfig): WeightTransformPlan;
+  createContext(weights: Weights, plan: WeightTransformPlan): WeightTransformContext;
+}
+
+export interface LlamaWeightTransformPlan extends WeightTransformPlan {
+  readonly kind: "llama";
+  readonly geometry: FoldGeometry;
+  readonly options: FoldOptions;
+}
+
+export interface QwenWeightTransformPlan extends WeightTransformPlan {
+  readonly kind: "qwen3_5" | "qwen3_5_mtp";
+  readonly hiddenSize: number;
+  readonly seed: number;
+  readonly fold: QwenFoldPlan;
+}
+
+function sourceMap(names: readonly string[]): Map<string, string> {
+  return new Map(names.map((name) => [name, name]));
+}
+
+function assertPow2(n: number, what: string): void {
+  if (n < 2 || (n & (n - 1)) !== 0)
+    throw new Error(
+      `rotation fold: ${what}=${n} is not a power of two — the Kronecker ` +
+      "Hadamard path is not implemented",
+    );
+}
+
+/** Pure Llama fold plan. Array work remains inside createContext/apply. */
+export function planLlamaWeightTransform(
+  names: readonly string[],
+  config: ModelConfig,
+  options: FoldOptions,
+): LlamaWeightTransformPlan {
+  const has = new Set(names);
+  for (const required of ["model.embed_tokens.weight", "model.norm.weight"])
+    if (!has.has(required)) throw new Error(`rotation fold: missing tensor ${required}`);
+  if (has.has("lm_head.weight"))
+    throw new Error("rotation fold: llama adapter currently requires tied embeddings");
+
+  for (let i = 0; i < config.text.numHiddenLayers; i++) {
+    const layer = `model.layers.${i}`;
+    for (const required of [
+      `${layer}.input_layernorm.weight`,
+      `${layer}.post_attention_layernorm.weight`,
+      `${layer}.self_attn.q_proj.weight`,
+      `${layer}.self_attn.k_proj.weight`,
+      `${layer}.self_attn.v_proj.weight`,
+      `${layer}.self_attn.o_proj.weight`,
+      `${layer}.mlp.gate_proj.weight`,
+      `${layer}.mlp.up_proj.weight`,
+      `${layer}.mlp.down_proj.weight`,
+    ]) {
+      if (!has.has(required)) throw new Error(`rotation fold: missing tensor ${required}`);
+    }
+  }
+
+  const geometry: FoldGeometry = {
+    numLayers: config.text.numHiddenLayers,
+    hiddenSize: config.text.hiddenSize,
+    numHeads: config.text.numAttentionHeads,
+    numKvHeads: config.text.numKeyValueHeads,
+    headDim: config.text.headDim,
+  };
+  if (options.r1 ?? true) assertPow2(geometry.hiddenSize, "hidden_size");
+  if (options.r2 ?? true) assertPow2(geometry.headDim, "head_dim");
+
+  const outputNames = [...names, "lm_head.weight"];
+  const sourceByOutput = sourceMap(names);
+  sourceByOutput.set("lm_head.weight", "model.embed_tokens.weight");
+  return {
+    kind: "llama",
+    id: "rotation.llama",
+    outputNames,
+    sourceByOutput,
+    untieWordEmbeddings: true,
+    geometry,
+    options: { ...options },
+    metadata: {
+      id: "rotation.llama",
+      seed: options.seed,
+      family: "llama",
+      deviations: [
+        "no-embedding-mean-centering",
+        "no-R4-downproj-input-fold",
+        "no-R3",
+        "gamma-kept-in-module-as-ones (eps preserved)",
+        "fold-precision-f32",
+      ],
+    },
+  };
+}
+
+function createLlamaContext(
+  weights: Weights,
+  plan: LlamaWeightTransformPlan,
+): WeightTransformContext {
+  const folded = foldLlamaWeights(weights, plan.geometry, plan.options);
+  const pending = new Map(folded.tensors.map((tensor) => [tensor.name, tensor.array]));
+  return {
+    apply(outputName, source) {
+      const transformed = pending.get(outputName);
+      if (transformed) {
+        pending.delete(outputName);
+        return transformed;
+      }
+      // Preserve family-specific tensors outside the fold corridor.
+      return source.astype(source.dtype);
+    },
+    dispose() {
+      for (const array of pending.values()) array.dispose();
+      pending.clear();
+    },
+  };
+}
+
+function qwenTransformPlan(
+  names: readonly string[],
+  config: ModelConfig,
+  seed: number,
+  companion: boolean,
+  prefix?: string,
+): QwenWeightTransformPlan {
+  const fold = companion
+    ? planQwenMtpFold(names)
+    : planQwen35Fold(
+        names,
+        prefix ?? (names.includes("language_model.model.norm.weight") ? "language_model." : ""),
+      );
+  const outputNames = [...names];
+  const sourceByOutput = sourceMap(names);
+  if (fold.extraHead) {
+    outputNames.push(fold.extraHead.name);
+    sourceByOutput.set(fold.extraHead.name, fold.extraHead.from);
+  }
+  const family = companion ? "qwen3_5_mtp" : "qwen3_5";
+  return {
+    kind: family,
+    id: `rotation.${family}`,
+    outputNames,
+    sourceByOutput,
+    untieWordEmbeddings: fold.extraHead !== null,
+    hiddenSize: config.text.hiddenSize,
+    seed,
+    fold,
+    metadata: {
+      id: `rotation.${family}`,
+      seed,
+      family,
+      deviations: [...fold.deviations],
+    },
+  };
+}
+
+function createQwenContext(
+  weights: Weights,
+  plan: QwenWeightTransformPlan,
+): WeightTransformContext {
+  // Lazy mode composes into the quantizer's existing writer. The experiment
+  // script keeps QwenFoldContext's eager default because it releases shards.
+  const context = new QwenFoldContext(
+    weights,
+    plan.hiddenSize,
+    plan.seed,
+    plan.fold.gammaNames,
+    false,
+  );
+  return {
+    apply(outputName, source) {
+      const extra = plan.fold.extraHead;
+      if (extra && outputName === extra.name)
+        return context.apply({ kind: "input", gamma: extra.gamma }, source);
+      const op = plan.fold.ops.get(outputName);
+      if (!op) throw new Error(`weight transform ${plan.id}: unplanned output ${outputName}`);
+      return context.apply(op, source);
+    },
+    dispose: () => context.dispose(),
+  };
+}
+
+/** Llama-family γ+R1+R2 adapter. */
+export function llamaWeightTransform(options: FoldOptions): WeightTransform {
+  return {
+    id: "rotation.llama",
+    plan: (names, config) => planLlamaWeightTransform(names, config, options),
+    createContext(weights, plan) {
+      if (!("kind" in plan) || plan.kind !== "llama")
+        throw new Error(`weight transform rotation.llama received plan ${plan.id}`);
+      return createLlamaContext(weights, plan as LlamaWeightTransformPlan);
+    },
+  };
+}
+
+/** Qwen3.5 trunk/VL adapter (R1 only because of the attention output gate). */
+export function qwen35WeightTransform(options: QwenFoldOptions): WeightTransform {
+  return {
+    id: "rotation.qwen3_5",
+    plan: (names, config) =>
+      qwenTransformPlan(names, config, options.seed, false, options.prefix),
+    createContext(weights, plan) {
+      if (!("kind" in plan) || plan.kind !== "qwen3_5")
+        throw new Error(`weight transform rotation.qwen3_5 received plan ${plan.id}`);
+      return createQwenContext(weights, plan as QwenWeightTransformPlan);
+    },
+  };
+}
+
+/** Qwen3.5 MTP companion adapter. Must use the same seed as its trunk. */
+export function qwenMtpWeightTransform(seed: number): WeightTransform {
+  return {
+    id: "rotation.qwen3_5_mtp",
+    plan: (names, config) => qwenTransformPlan(names, config, seed, true),
+    createContext(weights, plan) {
+      if (!("kind" in plan) || plan.kind !== "qwen3_5_mtp")
+        throw new Error(`weight transform rotation.qwen3_5_mtp received plan ${plan.id}`);
+      return createQwenContext(weights, plan as QwenWeightTransformPlan);
+    },
+  };
+}
+
+/** Select one of the three adapters from the source model and tensor schema. */
+export function automaticRotationWeightTransform(
+  options: FoldOptions & { prefix?: string },
+): WeightTransform {
+  return {
+    id: "rotation.auto",
+    plan(names, config) {
+      if (names.includes("fc.weight") && names.includes("pre_fc_norm_hidden.weight"))
+        return qwenTransformPlan(names, config, options.seed, true);
+      if (config.modelType === "qwen3_5")
+        return qwenTransformPlan(names, config, options.seed, false, options.prefix);
+      if (config.modelType === "llama")
+        return planLlamaWeightTransform(names, config, options);
+      throw new Error(`rotation transform: unsupported model_type ${config.modelType}`);
+    },
+    createContext(weights, plan) {
+      if (!("kind" in plan)) throw new Error(`rotation transform: invalid plan ${plan.id}`);
+      return plan.kind === "llama"
+        ? createLlamaContext(weights, plan as LlamaWeightTransformPlan)
+        : createQwenContext(weights, plan as QwenWeightTransformPlan);
+    },
+  };
+}

package/src/responses.ts

@@ -17,6 +17,11 @@
 //   - previous_response_id is echoed in the response when the request
 //     carried one (oracle always emits null).
 
+import type {
+  CompletionEvent,
+  CompletionStreamProtocol,
+} from "./serve/completion-sink";
+
 // ---------------------------------------------------------------------
 // Response store (previous_response_id resumption)
 // ---------------------------------------------------------------------
@@ -167,6 +172,55 @@ function coerceText(content: unknown): string {
   return parts.join("");
 }
 
+/** Convert a Responses message content array to chat-shaped content: text
+ *  flattens, media TRANSLATES (`input_image` → `image_url`, incl. the
+ *  Responses-standard string `image_url` field; video symmetrically), and
+ *  unknown non-text parts REJECT loudly — a silent media drop answers
+ *  confidently about content the model never saw (the anthropic.ts audio
+ *  doctrine). Text-only arrays keep the legacy flattened-string shape. */
+function coerceContent(content: unknown): string | Record<string, unknown>[] {
+  if (content == null) return "";
+  if (typeof content === "string") return content;
+  if (!Array.isArray(content)) return String(content);
+  const parts: Record<string, unknown>[] = [];
+  let hasMedia = false;
+  for (const part of content) {
+    if (typeof part === "string") {
+      parts.push({ type: "text", text: part });
+      continue;
+    }
+    if (!part || typeof part !== "object") continue;
+    const p = part as Record<string, unknown>;
+    if (p.type === "input_image" || p.type === "image_url") {
+      const url = typeof p.image_url === "string"
+        ? p.image_url
+        : (p.image_url as { url?: string } | undefined)?.url;
+      if (typeof url !== "string")
+        throw new Error("input_image part missing image_url");
+      parts.push({ type: "image_url", image_url: { url } });
+      hasMedia = true;
+    } else if (p.type === "input_video" || p.type === "video_url") {
+      const url = typeof p.video_url === "string"
+        ? p.video_url
+        : (p.video_url as { url?: string } | undefined)?.url;
+      if (typeof url !== "string")
+        throw new Error("input_video part missing video_url");
+      parts.push({ type: "video_url", video_url: { url } });
+      hasMedia = true;
+    } else if (p.text != null) {
+      parts.push({ type: "text", text: String(p.text) });
+    } else {
+      throw new Error(
+        `unsupported Responses content part type "${String(p.type)}" — ` +
+        `use input_text, input_image, or input_video`,
+      );
+    }
+  }
+  if (!hasMedia)
+    return parts.map((p) => String((p as { text?: unknown }).text ?? "")).join("");
+  return parts;
+}
+
 /** Convert a prior response's output array into input-shaped items for
  *  previous_response_id resumption (reasoning items dropped — replaying
  *  them invites the model to repeat itself). */
@@ -232,11 +286,14 @@ export function responsesToChatBody(body: ResponsesRequest): Record<string, unkn
       }
       // regular message item (explicit type or implicit by having role)
       let role = String(item.role ?? "user");
-      const content = coerceText(item.content);
       if (role === "system" || role === "developer") {
-        if (content) systemParts.push(content);
+        // System prompts are text by contract — flatten (media in a system
+        // item has no defined semantics on any surface we mirror).
+        const text = coerceText(item.content);
+        if (text) systemParts.push(text);
         continue;
       }
+      const content = coerceContent(item.content);
       if (!["user", "assistant", "tool"].includes(role)) role = "user";
       messages.push({ role, content });
     }
@@ -722,6 +779,49 @@ export class ResponsesStreamTranslator {
   }
 }
 
+/** Direct semantic-event adapter used by the server. `onComplete` receives
+ * the same response.completed payload stored for previous_response_id. */
+export function createResponsesStreamProtocol(
+  model: string,
+  previousResponseId: string | null,
+  onComplete: (final: Record<string, unknown>) => void,
+): CompletionStreamProtocol {
+  const translator = new ResponsesStreamTranslator(model, previousResponseId);
+  const chunk = (delta: Record<string, unknown>) => ({ choices: [{ delta }] });
+  return {
+    start: () => translator.addChunk(chunk({ role: "assistant", content: "" })),
+    addEvents(events: CompletionEvent[]) {
+      return events.flatMap((event) => {
+        if (event.type === "reasoning") {
+          return translator.addChunk(chunk({ reasoning: event.text }));
+        }
+        if (event.type === "content") {
+          return translator.addChunk(chunk({ content: event.text }));
+        }
+        return translator.addChunk(chunk({
+          tool_calls: event.calls.map((call, index) => ({ index, ...call })),
+        }));
+      });
+    },
+    finish(reason, usage) {
+      const frames = [
+        ...translator.addChunk({ choices: [{ delta: {}, finish_reason: reason }], usage }),
+        ...translator.finalize(),
+      ];
+      onComplete(translator.finalResponse());
+      return frames;
+    },
+    error(message) {
+      return [sse("error", {
+        type: "error",
+        code: "server_error",
+        message,
+        param: null,
+      })];
+    },
+  };
+}
+
 /** Wrap our OpenAI SSE byte stream as a Responses SSE byte stream.
  *  onComplete fires after response.completed with the final payload so
  *  the route can store it for previous_response_id resumption. */

package/src/runtime-config.ts

@@ -0,0 +1,67 @@
+export type RuntimeKey = `MLX_BUN_${string}`;
+export type RuntimeOverrides = Readonly<Partial<Record<RuntimeKey, string | undefined>>>;
+
+export interface RuntimeConfig {
+  readonly values: Readonly<Record<string, string | undefined>>;
+  value(name: RuntimeKey): string | undefined;
+  flag(name: RuntimeKey, defaultOn: boolean): boolean;
+  number(name: RuntimeKey, fallback: number): number;
+}
+
+/** Build one immutable runtime snapshot. Only mlx-bun-owned keys cross this
+ *  boundary; feature code never reads or mutates process.env. */
+export function createRuntimeConfig(
+  source: Readonly<Record<string, string | undefined>>,
+): RuntimeConfig {
+  const values = Object.freeze(Object.fromEntries(
+    Object.entries(source).filter(([key, value]) =>
+      key.startsWith("MLX_BUN_") && value !== undefined),
+  ));
+  return Object.freeze({
+    values,
+    value: (name: RuntimeKey) => values[name],
+    flag: (name: RuntimeKey, defaultOn: boolean) => {
+      const value = values[name];
+      return value === "1" ? true : value === "0" ? false : defaultOn;
+    },
+    number: (name: RuntimeKey, fallback: number) => {
+      const value = values[name];
+      if (value === undefined) return fallback;
+      const parsed = Number(value);
+      return Number.isFinite(parsed) ? parsed : fallback;
+    },
+  });
+}
+
+let active = createRuntimeConfig(process.env);
+
+export function runtimeConfig(): RuntimeConfig {
+  return active;
+}
+
+export function runtimeValue(name: RuntimeKey): string | undefined {
+  return active.value(name);
+}
+
+export function runtimeFlag(name: RuntimeKey, defaultOn: boolean): boolean {
+  return active.flag(name, defaultOn);
+}
+
+export function runtimeNumber(name: RuntimeKey, fallback: number): number {
+  return active.number(name, fallback);
+}
+
+export function runtimeKey(name: string): RuntimeKey {
+  if (!name.startsWith("MLX_BUN_"))
+    throw new Error(`runtime flag must start with MLX_BUN_ (got ${name})`);
+  return name as RuntimeKey;
+}
+
+/** Replace selected values by installing a new frozen snapshot. The CLI uses
+ *  this once while resolving flags; tests may use the returned restore hook.
+ *  No environment mutation is involved. */
+export function configureRuntime(overrides: RuntimeOverrides): () => void {
+  const previous = active;
+  active = createRuntimeConfig({ ...previous.values, ...overrides });
+  return () => { active = previous; };
+}