mlx-bun 0.0.4

package/src/config.ts

@@ -0,0 +1,229 @@
+// Model config + quantization metadata reader for the load path.
+// Parses HF config.json (Gemma 4 unified layout: text_config nested),
+// the OptiQ per-layer quantization map, and kv_config.json when present.
+
+export interface RopeParams {
+  ropeTheta: number;
+  ropeType: string;
+  partialRotaryFactor: number;
+  /** ProportionalRoPE frequency multiplier (rope_utils: factor, default 1). */
+  factor: number;
+}
+
+export interface TextConfig {
+  hiddenSize: number;
+  numHiddenLayers: number;
+  numAttentionHeads: number;
+  numKeyValueHeads: number;
+  headDim: number;
+  /** Global (full-attention) layers can use different KV geometry. */
+  numGlobalKeyValueHeads: number;
+  globalHeadDim: number;
+  attentionKEqV: boolean;
+  intermediateSize: number;
+  hiddenActivation: string;
+  rmsNormEps: number;
+  vocabSize: number;
+  maxPositionEmbeddings: number;
+  slidingWindow: number;
+  /** "sliding_attention" | "full_attention" per layer. */
+  layerTypes: string[];
+  /** e2b/e4b: per-layer input embedding width (0 = disabled). */
+  hiddenSizePerLayerInput: number;
+  vocabSizePerLayerInput: number;
+  /** Trailing layers that share KV with earlier same-type layers. */
+  numKvSharedLayers: number;
+  /** 26B-A4B MoE block: dense MLP + routed experts per layer. */
+  enableMoeBlock: boolean;
+  numExperts: number;
+  topKExperts: number;
+  moeIntermediateSize: number;
+  /** Qwen3.5 hybrid (gated-DeltaNet) geometry; 0/false for other families. */
+  linearNumValueHeads: number;
+  linearNumKeyHeads: number;
+  linearKeyHeadDim: number;
+  linearValueHeadDim: number;
+  linearConvKernelDim: number;
+  /** Every `fullAttentionInterval`-th layer is full-attention (Qwen3.5). */
+  fullAttentionInterval: number;
+  /** Qwen3.5 full-attention output gate: q_proj emits 2× and gates the output. */
+  attnOutputGate: boolean;
+  /** Rotary fraction of head_dim (Qwen3.5 full attention: 0.25 → 64 of 256). */
+  partialRotaryFactor: number;
+  ropeParameters: Record<string, RopeParams>;
+  finalLogitSoftcapping: number | null;
+  tieWordEmbeddings: boolean;
+  bosTokenId: number;
+  eosTokenId: number | number[];
+}
+
+export interface QuantSpec {
+  bits: number;
+  groupSize: number;
+  mode: string;
+}
+
+export interface QuantizationConfig {
+  default: QuantSpec;
+  /** Per-module overrides keyed by module path (e.g. "language_model.model.embed_tokens"). */
+  perLayer: Map<string, QuantSpec>;
+}
+
+export interface KvQuantSpec {
+  layerIdx: number;
+  bits: number;
+  groupSize: number;
+}
+
+export interface ModelConfig {
+  modelDir: string;
+  modelType: string;
+  architectures: string[];
+  dtype: string;
+  text: TextConfig;
+  quantization: QuantizationConfig | null;
+  kvQuant: KvQuantSpec[] | null;
+  hasVisionSidecar: boolean;
+  eosTokenIds: number[];
+  raw: Record<string, unknown>;
+}
+
+function parseRope(raw: Record<string, any>): Record<string, RopeParams> {
+  const out: Record<string, RopeParams> = {};
+  for (const [kind, p] of Object.entries(raw ?? {})) {
+    out[kind] = {
+      ropeTheta: p.rope_theta,
+      ropeType: p.rope_type ?? "default",
+      partialRotaryFactor: p.partial_rotary_factor ?? 1.0,
+      factor: p.factor ?? 1.0,
+    };
+  }
+  return out;
+}
+
+function parseQuantization(raw: Record<string, any> | undefined): QuantizationConfig | null {
+  if (!raw) return null;
+  const def: QuantSpec = {
+    bits: raw.bits,
+    groupSize: raw.group_size,
+    mode: raw.mode ?? "affine",
+  };
+  const perLayer = new Map<string, QuantSpec>();
+  for (const [key, value] of Object.entries(raw)) {
+    if (key === "bits" || key === "group_size" || key === "mode") continue;
+    if (typeof value === "object" && value !== null) {
+      perLayer.set(key, {
+        bits: (value as any).bits ?? def.bits,
+        groupSize: (value as any).group_size ?? def.groupSize,
+        mode: (value as any).mode ?? def.mode,
+      });
+    } else if (value === false) {
+      // mlx convention: `"layer": false` means not quantized
+      perLayer.set(key, { bits: 0, groupSize: 0, mode: "none" });
+    }
+  }
+  return { default: def, perLayer };
+}
+
+/** Quant spec for a module path, or null if the module is unquantized. */
+export function quantFor(q: QuantizationConfig | null, modulePath: string): QuantSpec | null {
+  if (!q) return null;
+  const spec = q.perLayer.get(modulePath) ?? q.default;
+  return spec.mode === "none" ? null : spec;
+}
+
+export async function loadModelConfig(modelDir: string): Promise<ModelConfig> {
+  const raw = (await Bun.file(`${modelDir}/config.json`).json()) as Record<string, any>;
+  // Gemma 4 unified nests the LM config; plain text models keep it at top level.
+  const t = (raw.text_config ?? raw) as Record<string, any>;
+  const modelType = raw.model_type;
+  const isLlama = modelType === "llama";
+  const isQwen35 = typeof modelType === "string" && modelType.startsWith("qwen3_5");
+  // Qwen3.5 rope_parameters is a flat dict ({type, rope_theta, mrope_section,
+  // partial_rotary_factor}), not the gemma per-attention-type map — and
+  // type "default" means plain partial nn.RoPE (mrope_section ignored for text).
+  const qwenRope = (t.rope_parameters ?? {}) as Record<string, any>;
+
+  const text: TextConfig = {
+    hiddenSize: t.hidden_size,
+    numHiddenLayers: t.num_hidden_layers,
+    numAttentionHeads: t.num_attention_heads,
+    numKeyValueHeads: t.num_key_value_heads,
+    headDim: t.head_dim ?? Math.floor(t.hidden_size / t.num_attention_heads),
+    numGlobalKeyValueHeads: t.num_global_key_value_heads ?? t.num_key_value_heads,
+    globalHeadDim: t.global_head_dim ?? t.head_dim ?? Math.floor(t.hidden_size / t.num_attention_heads),
+    attentionKEqV: t.attention_k_eq_v ?? false,
+    intermediateSize: t.intermediate_size,
+    hiddenActivation: t.hidden_activation ?? t.hidden_act ?? "gelu_pytorch_tanh",
+    rmsNormEps: t.rms_norm_eps,
+    vocabSize: t.vocab_size,
+    maxPositionEmbeddings: t.max_position_embeddings,
+    slidingWindow: t.sliding_window ?? 0,
+    layerTypes: t.layer_types ?? (isLlama ? Array(t.num_hidden_layers).fill("full_attention") : []),
+    hiddenSizePerLayerInput: t.hidden_size_per_layer_input ?? 0,
+    vocabSizePerLayerInput: t.vocab_size_per_layer_input ?? t.vocab_size,
+    numKvSharedLayers: t.num_kv_shared_layers ?? 0,
+    enableMoeBlock: t.enable_moe_block ?? false,
+    numExperts: t.num_experts ?? 0,
+    topKExperts: t.top_k_experts ?? 0,
+    moeIntermediateSize: t.moe_intermediate_size ?? 0,
+    linearNumValueHeads: t.linear_num_value_heads ?? 0,
+    linearNumKeyHeads: t.linear_num_key_heads ?? 0,
+    linearKeyHeadDim: t.linear_key_head_dim ?? 0,
+    linearValueHeadDim: t.linear_value_head_dim ?? 0,
+    linearConvKernelDim: t.linear_conv_kernel_dim ?? 0,
+    fullAttentionInterval: t.full_attention_interval ?? 0,
+    attnOutputGate: t.attn_output_gate ?? false,
+    partialRotaryFactor:
+      qwenRope.partial_rotary_factor ?? t.partial_rotary_factor ?? 1.0,
+    ropeParameters: isQwen35
+      ? {
+          full_attention: {
+            ropeTheta: qwenRope.rope_theta ?? t.rope_theta ?? 10000,
+            ropeType: "default",
+            partialRotaryFactor: qwenRope.partial_rotary_factor ?? 1.0,
+            factor: 1.0,
+          },
+        }
+      : t.rope_parameters
+        ? parseRope(t.rope_parameters)
+        : isLlama
+          ? {
+              full_attention: {
+                ropeTheta: t.rope_theta ?? 10000,
+                ropeType: "default",
+                partialRotaryFactor: 1.0,
+                factor: 1.0,
+              },
+            }
+          : {},
+    finalLogitSoftcapping: t.final_logit_softcapping ?? null,
+    tieWordEmbeddings: t.tie_word_embeddings ?? raw.tie_word_embeddings ?? false,
+    bosTokenId: t.bos_token_id ?? raw.bos_token_id,
+    eosTokenId: t.eos_token_id ?? raw.eos_token_id,
+  };
+
+  let kvQuant: KvQuantSpec[] | null = null;
+  if (await Bun.file(`${modelDir}/kv_config.json`).exists()) {
+    const kv = (await Bun.file(`${modelDir}/kv_config.json`).json()) as any[];
+    kvQuant = kv.map((e) => ({
+      layerIdx: e.layer_idx,
+      bits: e.bits,
+      groupSize: e.group_size,
+    }));
+  }
+
+  const eos = raw.eos_token_id ?? text.eosTokenId;
+  return {
+    modelDir,
+    modelType,
+    architectures: raw.architectures ?? [],
+    dtype: raw.dtype ?? "bfloat16",
+    text,
+    quantization: parseQuantization(raw.quantization ?? raw.quantization_config),
+    kvQuant,
+    hasVisionSidecar: await Bun.file(`${modelDir}/optiq_vision.safetensors`).exists(),
+    eosTokenIds: Array.isArray(eos) ? eos : [eos],
+    raw,
+  };
+}

package/src/curve-sampler.ts

@@ -0,0 +1,122 @@
+// curve-sampler.ts — the v2 log-prob transfer-curve sampler. Replaces the whole
+// temperature+softmax stage with a drawn MONOTONE CUBIC curve in log-probability
+// space:
+//
+//   p_out_i = exp(curve(log p_in_i)) / Σ_j exp(curve(log p_in_j))
+//
+// `lp` reaching the sampler is already log p_in (toLogprobs = logits − logsumexp),
+// so the curve maps it directly and an identity curve reproduces temperature 1.
+// The curve is a Fritsch–Carlson monotone cubic Hermite (PCHIP) through N movable
+// control points in nat log-prob space (DaVinci-curves style — the endpoints start
+// at the corners but are draggable too). Smooth, never overshoots, monotone. A
+// control point dragged below the diagonal suppresses those tokens; the gate is
+// emergent from the curve shape, not a separate threshold.
+//
+// Reading the curve: at a given p_in, ABOVE the diagonal raises those tokens,
+// BELOW lowers them (renorm makes the literal number an intent); the SLOPE across
+// a region sets how evenly the tokens inside it spread (flat<1 even, steep>1 separated).
+
+import { MlxArray } from "./mlx/array";
+import * as ops from "./mlx/ops";
+
+export interface CurvePoint { x_pct: number; y_pct: number; }
+export interface CurveParams {
+  space?: "logprob";
+  points: CurvePoint[]; // ≥2 control points, all movable; first/last are the endpoints
+  monotonic?: boolean;
+}
+
+/** Bottom-left default anchor / plot floor (p = 1e-6). */
+export const CURVE_UMIN = Math.log(1e-6);
+const L = (pt: CurvePoint): [number, number] => [Math.log(Math.max(pt.x_pct, 1e-9) / 100), Math.log(Math.max(pt.y_pct, 1e-9) / 100)];
+
+/** Control points in nat log-prob space, sorted ascending in x. */
+function knotArrays(p: CurveParams): { X: number[]; Y: number[] } {
+  const k = p.points.map(L).sort((a, b) => a[0] - b[0]);
+  return { X: k.map((q) => q[0]), Y: k.map((q) => q[1]) };
+}
+
+/** Consecutive-point secants, top→bottom (region slopes; what the editor readout shows). */
+export function curveSecants(p: CurveParams): number[] {
+  const { X, Y } = knotArrays(p);
+  const s: number[] = [];
+  for (let i = X.length - 1; i > 0; i--) s.push((Y[i]! - Y[i - 1]!) / (X[i]! - X[i - 1]!));
+  return s;
+}
+/** True iff every consecutive-point secant is ≥ 0 (monotone non-decreasing). */
+export function isMonotone(p: CurveParams): boolean {
+  return p.points.length >= 2 && curveSecants(p).every((s) => s >= -1e-12);
+}
+
+interface Seg { A: number; B: number; C: number; D: number } // cubic in u: A+Bu+Cu²+Du³
+export interface Spline { X: number[]; Y: number[]; M: number[]; segs: Seg[] }
+
+function cubicInU(x0: number, h: number, y0: number, y1: number, m0: number, m1: number): Seg {
+  const c0 = y0, c1 = h * m0;
+  const c2 = -3 * y0 - 2 * h * m0 + 3 * y1 - h * m1;
+  const c3 = 2 * y0 + h * m0 - 2 * y1 + h * m1;
+  const s = 1 / h, s2 = s * s, s3 = s2 * s;
+  let A = c0, B = 0, C = 0, D = 0;
+  B += c1 * s; A += -c1 * s * x0;
+  C += c2 * s2; B += -2 * x0 * c2 * s2; A += x0 * x0 * c2 * s2;
+  D += c3 * s3; C += -3 * x0 * c3 * s3; B += 3 * x0 * x0 * c3 * s3; A += -x0 * x0 * x0 * c3 * s3;
+  return { A, B, C, D };
+}
+
+/** Monotone cubic (Fritsch–Carlson) over the N control points. */
+export function buildSpline(p: CurveParams): Spline {
+  const { X, Y } = knotArrays(p);
+  const n = X.length;
+  const Dsec: number[] = [];
+  for (let i = 0; i < n - 1; i++) Dsec.push((Y[i + 1]! - Y[i]!) / (X[i + 1]! - X[i]!));
+  const M = new Array<number>(n);
+  M[0] = Dsec[0]!; M[n - 1] = Dsec[n - 2]!;
+  for (let i = 1; i < n - 1; i++) {
+    if (Dsec[i - 1]! * Dsec[i]! <= 0) { M[i] = 0; continue; }
+    const h1 = X[i]! - X[i - 1]!, h2 = X[i + 1]! - X[i]!;
+    const w1 = 2 * h2 + h1, w2 = h2 + 2 * h1;
+    M[i] = (w1 + w2) / (w1 / Dsec[i - 1]! + w2 / Dsec[i]!);
+  }
+  const segs: Seg[] = [];
+  for (let i = 0; i < n - 1; i++) segs.push(cubicInU(X[i]!, X[i + 1]! - X[i]!, Y[i]!, Y[i + 1]!, M[i]!, M[i + 1]!));
+  return { X, Y, M, segs };
+}
+
+/** Scalar eval (editor preview / tests). Extrapolates linearly past both endpoints. */
+export function evalSpline(sp: Spline, u: number): number {
+  const n = sp.X.length;
+  if (u < sp.X[0]!) return sp.Y[0]! + Math.max(sp.M[0]!, 1) * (u - sp.X[0]!); // deep tail: slope ≥1, stays gated
+  if (u > sp.X[n - 1]!) return sp.Y[n - 1]! + sp.M[n - 1]! * (u - sp.X[n - 1]!);
+  let i = sp.segs.length - 1;
+  while (i > 0 && u < sp.X[i]!) i--;
+  const s = sp.segs[i]!;
+  return s.A + s.B * u + s.C * u * u + s.D * u * u * u;
+}
+export function curveAt(p: CurveParams, u: number): number { return evalSpline(buildSpline(p), u); }
+
+/** Apply the monotone-cubic transfer to normalized logprobs `lp` [.., V] (on-device). */
+export function applyCurve(lp: MlxArray, p: CurveParams): MlxArray {
+  const sp = buildSpline(p);
+  const n = sp.X.length;
+  const tmp: MlxArray[] = [];
+  const k = <T extends MlxArray>(a: T): T => { tmp.push(a); return a; };
+  const C = (v: number): MlxArray => k(ops.scalarLike(v, lp));
+  const lp2 = k(ops.mul(lp, lp));
+  const lp3 = k(ops.mul(lp2, lp));
+  const seg = (s: Seg): MlxArray => k(ops.add(
+    k(ops.add(C(s.A), k(ops.mulScalar(lp, s.B)))),
+    k(ops.add(k(ops.mulScalar(lp2, s.C)), k(ops.mulScalar(lp3, s.D)))),
+  ));
+  // below the bottom point: slope ≥ 1 (deep tail stays gated)
+  const lowSlope = Math.max(sp.M[0]!, 1);
+  let out: MlxArray = k(ops.add(C(sp.Y[0]! - lowSlope * sp.X[0]!), k(ops.mulScalar(lp, lowSlope))));
+  for (let i = 0; i < sp.segs.length; i++) {
+    const next = ops.where(k(ops.greaterEqual(lp, C(sp.X[i]!))), seg(sp.segs[i]!), out);
+    tmp.push(next); out = next;
+  }
+  // above the top point: linear extrapolation with the top tangent (if it was dragged left of 100%)
+  const aboveExt = k(ops.add(C(sp.Y[n - 1]! - sp.M[n - 1]! * sp.X[n - 1]!), k(ops.mulScalar(lp, sp.M[n - 1]!))));
+  const result = ops.where(k(ops.greaterEqual(lp, C(sp.X[n - 1]!))), aboveExt, out);
+  for (const a of tmp) a.dispose();
+  return result;
+}