mlx-bun 0.0.13 → 0.3.0

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)
 // ---------------------------------------------------------------------
@@ -38,7 +43,16 @@ const DEFAULT_MAX_BYTES = 32 * 1024 * 1024;
 /** Process-local TTL + byte-capped LRU (port of response_store.py).
  *  Does not survive restarts — sufficient for local serving; anything
  *  sturdier means external storage, out of scope upstream too. */
-export class ResponseStore {
+export interface ResponseHistory {
+  readonly ttlMs: number;
+  readonly maxBytes: number;
+  readonly size: number;
+  readonly totalBytes: number;
+  put(id: string, entry: StoredResponse): void;
+  get(id: string): StoredResponse | null;
+}
+
+export class ResponseStore implements ResponseHistory {
   #items = new Map<string, StoreEntry>(); // Map iterates in insertion order
   #bytes = 0;
 
@@ -109,6 +123,24 @@ export class ResponseStore {
   }
 }
 
+/** Conversation policy shared by direct and isolated application hosts. */
+export function resolveResponsesConversation(body: ResponsesRequest, store: Pick<ResponseHistory, "get">) {
+  const previousId = body.previous_response_id ?? null;
+  if (previousId) {
+    const prior = store.get(previousId);
+    if (!prior) throw new Error(`previous_response_id '${previousId}' not found or expired`);
+    const next = typeof body.input === "string"
+      ? body.input ? [{ type: "message", role: "user", content: body.input }] : []
+      : body.input ?? [];
+    body = { ...body, input: [...prior.input, ...outputItemsToInputItems(prior.output), ...next] as Array<Record<string, unknown>>,
+      instructions: body.instructions ?? prior.instructions ?? undefined };
+  }
+  return { body, previousId,
+    input: typeof body.input === "string" ? [{ type: "message", role: "user", content: body.input }] : [...body.input ?? []],
+    instructions: body.instructions ?? null,
+  };
+}
+
 // ---------------------------------------------------------------------
 // Request translation: Responses → Chat Completions
 // ---------------------------------------------------------------------
@@ -167,6 +199,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 +313,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 +806,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,81 @@
+import { AsyncLocalStorage } from "node:async_hooks";
+
+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);
+const executionConfig = new AsyncLocalStorage<RuntimeConfig>();
+
+/** Legacy kernels read through this port. A bound execution sees its own
+ * immutable snapshot across awaits, independently of later host settings. */
+export function withRuntimeConfig<T>(config: RuntimeConfig, run: () => T): T {
+  return executionConfig.getStore() === config ? run() : executionConfig.run(config, run);
+}
+
+export function runtimeConfig(): RuntimeConfig {
+  return executionConfig.getStore() ?? active;
+}
+
+export function runtimeValue(name: RuntimeKey): string | undefined {
+  return runtimeConfig().value(name);
+}
+
+/** On/off env flag: an explicit "1"/"0" always wins; otherwise `defaultOn`.
+ *  (Formerly src/flags.ts `flagOn` — same resolver, one module.) */
+export function flagOn(name: RuntimeKey, defaultOn: boolean): boolean {
+  return runtimeFlag(name, defaultOn);
+}
+export function runtimeFlag(name: RuntimeKey, defaultOn: boolean): boolean {
+  return runtimeConfig().flag(name, defaultOn);
+}
+
+export function runtimeNumber(name: RuntimeKey, fallback: number): number {
+  return runtimeConfig().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; };
+}