@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/layers/quantized-linear.ts

@@ -0,0 +1,137 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { LoraDelta, type LoraState, type LoraWeights } from "./lora";
+import { quantFor,type ModelConfig } from "../artifacts/config";
+import type { Weights } from "../artifacts/weights";
+import { quantizedMatmulRows } from "../kernels/quantization/affine-rows";
+
+/** Replaceable quantized payload. Inference sees this as read-only through
+ *  QuantizedLinear getters; sensitivity analysis exchanges it atomically. */
+export interface QuantizedLinearState {
+  w: MlxArray;
+  scales: MlxArray;
+  biases: MlxArray | null;
+  spec: ops.QuantSpec;
+}
+
+export class QuantizedLinear {
+  readonly delta = new LoraDelta();
+  get adapters(): Map<string, LoraWeights> | null { return this.delta.adapters; }
+  set adapters(value: Map<string, LoraWeights> | null) { this.delta.adapters = value; }
+  get loraState(): LoraState | null { return this.delta.loraState; }
+  set loraState(value: LoraState | null) { this.delta.loraState = value; }
+  get dropoutId(): number { return this.delta.dropoutId; }
+  set dropoutId(value: number) { this.delta.dropoutId = value; }
+  private _w: MlxArray;
+  private _scales: MlxArray;
+  private _biases: MlxArray | null;
+  private _spec: ops.QuantSpec;
+
+  constructor(
+    w: MlxArray,
+    scales: MlxArray,
+    biases: MlxArray | null,
+    spec: ops.QuantSpec,
+    /** ADDITIVE bias term (`.bias`, mlx nn.QuantizedLinear's optional bias —
+     *  qwen2 qkv, starcoder2, …). Distinct from `biases`, the quantization
+     *  zero-points. Applied after the matmul, before any LoRA residual,
+     *  exactly like mlx's QuantizedLinear.__call__. */
+    readonly bias: MlxArray | null = null,
+  ) {
+    this._w = w;
+    this._scales = scales;
+    this._biases = biases;
+    this._spec = spec;
+  }
+
+  get w(): MlxArray { return this._w; }
+  get scales(): MlxArray { return this._scales; }
+  get biases(): MlxArray | null { return this._biases; }
+  get spec(): ops.QuantSpec { return this._spec; }
+
+  /** Atomically install a quantized payload and return the previous one.
+   *  This is the sole mutation capability used by sensitivity sweeps. */
+  exchangeQuantizedState(next: QuantizedLinearState): QuantizedLinearState {
+    const previous = {
+      w: this._w,
+      scales: this._scales,
+      biases: this._biases,
+      spec: this._spec,
+    };
+    this._w = next.w;
+    this._scales = next.scales;
+    this._biases = next.biases;
+    this._spec = next.spec;
+    return previous;
+  }
+
+  static load(weights: Weights, path: string, config: ModelConfig): QuantizedLinear {
+    if (!weights.has(`${path}.scales`))
+      throw new Error(`${path}: expected quantized linear (no .scales tensor)`);
+    const spec = quantFor(config.quantization, path);
+    if (!spec) throw new Error(`${path}: no quant spec`);
+    return new QuantizedLinear(
+      weights.tensor(`${path}.weight`),
+      weights.tensor(`${path}.scales`),
+      weights.has(`${path}.biases`) ? weights.tensor(`${path}.biases`) : null,
+      spec,
+      weights.has(`${path}.bias`) ? weights.tensor(`${path}.bias`) : null,
+    );
+  }
+
+  /** (in_features, out_features) from the quantized tensors (reference
+   *  _infer_linear_shape: scales are [out, in/group_size]). */
+  get inFeatures(): number {
+    return this._scales.shape[1]! * this._spec.groupSize;
+  }
+  get outFeatures(): number {
+    return this._scales.shape[0]!;
+  }
+
+  forward(x: MlxArray, independentRows = false): MlxArray {
+    let out = independentRows ? quantizedMatmulRows(
+      x, this._w, this._scales, this._biases, this._spec,
+    ) : ops.quantizedMatmul(
+      x,
+      this._w,
+      this._scales,
+      this._biases,
+      this._spec,
+      true,
+    );
+    // Additive bias (mlx QuantizedLinear: `x = x + self["bias"]`), before
+    // the LoRA residual (LoRALinear calls the base linear bias-inclusive).
+    if (this.bias) { const previous = out; out = ops.add(out, this.bias); previous.dispose(); }
+    return this.delta.forward(x, out);
+  }
+}
+
+/** Stock affine projection. The weights and adapter target belong to base. */
+export class QuantizedLinearPlain {
+  readonly #bias: (out: MlxArray) => MlxArray;
+  constructor(readonly base: QuantizedLinear) {
+    this.#bias = base.bias ? out => { const next = ops.add(out, base.bias!); out.dispose(); return next; } : out => out;
+  }
+  static load(weights: Weights, path: string, config: ModelConfig): QuantizedLinearPlain {
+    return new QuantizedLinearPlain(QuantizedLinear.load(weights, path, config));
+  }
+  forward(x: MlxArray): MlxArray {
+    const out = ops.quantizedMatmul(x, this.base.w, this.base.scales, this.base.biases, this.base.spec, true);
+    return this.base.delta.forward(x, this.#bias(out));
+  }
+}
+
+/** Each position uses independent one-row affine arithmetic. */
+export class QuantizedLinearRows {
+  readonly #bias: (out: MlxArray) => MlxArray;
+  constructor(readonly base: QuantizedLinear) {
+    this.#bias = base.bias ? out => { const next = ops.add(out, base.bias!); out.dispose(); return next; } : out => out;
+  }
+  static load(weights: Weights, path: string, config: ModelConfig): QuantizedLinearRows {
+    return new QuantizedLinearRows(QuantizedLinear.load(weights, path, config));
+  }
+  forward(x: MlxArray): MlxArray {
+    const out = quantizedMatmulRows(x, this.base.w, this.base.scales, this.base.biases, this.base.spec);
+    return this.base.delta.forward(x, this.#bias(out));
+  }
+}

package/src/layers/quantized-switch-linear.ts

@@ -0,0 +1,37 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { quantFor,type ModelConfig } from "../artifacts/config";
+import { expertOffloadArray } from "../artifacts/expert-offload";
+import type { Weights } from "../artifacts/weights";
+
+/** Port of switch_layers.QuantizedSwitchLinear (gather_qmm over stacked
+ *  expert weights; rhs_indices selects the expert per row). */
+export class QuantizedSwitchLinear {
+  constructor(
+    readonly w: MlxArray,
+    readonly scales: MlxArray,
+    readonly biases: MlxArray | null,
+    readonly spec: ops.QuantSpec,
+  ) {}
+
+  static load(weights: Weights, path: string, config: ModelConfig): QuantizedSwitchLinear {
+    if (!weights.has(`${path}.scales`))
+      throw new Error(`${path}: expected quantized switch linear (no .scales tensor)`);
+    const spec = quantFor(config.quantization, path);
+    if (!spec) throw new Error(`${path}: no quant spec`);
+    // Expert WEIGHT (the ~94% of expert bytes) comes from the page-aligned
+    // offload mmap when --expert-offload is active (else resident); scales/
+    // biases stay resident (small). Same bytes either way → bit-exact.
+    const wName = `${path}.weight`;
+    return new QuantizedSwitchLinear(
+      expertOffloadArray(wName) ?? weights.tensor(wName),
+      weights.tensor(`${path}.scales`),
+      weights.has(`${path}.biases`) ? weights.tensor(`${path}.biases`) : null,
+      spec,
+    );
+  }
+
+  forward(x: MlxArray, indices: MlxArray, sortedIndices: boolean): MlxArray {
+    return ops.gatherQmm(x, this.w, this.scales, this.biases, indices, this.spec, sortedIndices);
+  }
+}

package/src/layers/qwen-append.ts

@@ -0,0 +1,12 @@
+// MLX 0.32.2, mlx/backend/metal/scaled_dot_product_attention.cpp (MIT):
+// https://github.com/ml-explore/mlx/blob/v0.32.2/mlx/backend/metal/scaled_dot_product_attention.cpp
+// For the qualified Qwen GQA-6/D256 graph on applegpu_g16s, one-pass/vector
+// selection and two-pass block counts change after these inclusive KV lengths.
+// A committed span must stay within one arithmetic regime to match M=1.
+const ATTENTION_BOUNDARIES = [1023, 1024, 8192, 32768, 65536];
+
+export function qwenAppendChunkSize(cacheTokens: number): number {
+  for (const end of ATTENTION_BOUNDARIES)
+    if (cacheTokens < end) return Math.min(4, end - cacheTokens);
+  return 4;
+}

package/src/layers/qwen-mrope.ts

@@ -0,0 +1,238 @@
+// Qwen3.5/3.8 multimodal RoPE (mRoPE) — the language-side position machinery
+// for vision requests, port of mlx-vlm's qwen3_5 language rope:
+// MRoPERotaryEmbedding (style "interleaved", mrope_section [11,11,10]) +
+// get_rope_index + the rope_deltas decode continuation.
+//
+// TEXT-ONLY EQUIVALENCE: with all three position streams equal, interleaved
+// selection collapses to plain partial RoPE — which is why the text path's
+// ops.rope stays bit-exact untouched. This module only activates for
+// requests whose prompt contains vision spans; while active, EVERY
+// full-attention rope in the request (prefill AND decode) goes through the
+// manual interleaved path, mirroring the reference (which never uses the
+// fused fast-rope kernel when position_ids are supplied).
+//
+// Position semantics (get_rope_index, B=1 serial): text tokens advance all
+// three streams together; an image span of grid (t, h/merge, w/merge) gets
+// t/h/w grid positions offset by the current base; text after an image
+// resumes at max+1. delta = (max position + 1) - sequence length; decode
+// positions are offset + delta on all three streams.
+
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+import * as ops from "@mlx-bun/mlx/ops";
+
+const MROPE_SECTION = [11, 11, 10] as const;
+
+/** Interleaved frequency→stream selector (rope_utils
+ *  _interleaved_position_selector): default t(0); indices 1,4,7,… take
+ *  h(1) while < section[1]*3; indices 2,5,8,… take w(2) while < section[2]*3. */
+export function interleavedSelector(freqDim: number): Int32Array {
+  const sel = new Int32Array(freqDim);
+  for (let dim = 1; dim <= 2; dim++) {
+    const offset = dim; // h starts at 1, w at 2
+    for (let i = offset; i < Math.min(MROPE_SECTION[dim]! * 3, freqDim); i += 3)
+      sel[i] = dim;
+  }
+  return sel;
+}
+
+/** inv_freq for the language rope — computed on-device exactly like
+ *  compute_inv_freq (Metal powf is not correctly rounded; a JS emulation
+ *  drifts by 1 ulp on some entries). Returned as a live f32 device array
+ *  (the apply kernel reads it directly). */
+export function mropeInvFreq(rotaryDims: number, base: number): MlxArray {
+  const ar = ops.arange(0, rotaryDims, 2, Dtype.float32);
+  const ex = ops.mulScalar(ar, 1 / rotaryDims);
+  ar.dispose();
+  const theta = MlxArray.fromFloat32(Float32Array.from([base]), [1]);
+  const p = ops.pow(theta, ex);
+  theta.dispose();
+  ex.dispose();
+  const one = MlxArray.fromFloat32(Float32Array.from([1]), [1]);
+  const invA = ops.div(one, p);
+  one.dispose();
+  p.dispose();
+  return invA;
+}
+
+/** Effective positions [3,1,L] int32 for window [offset, offset+L) —
+ *  prompt-grid positions while inside the prompt, offset+delta beyond it
+ *  (the reference's rope_deltas continuation). */
+export function buildMropePositions(
+  state: MropeRequestState, offset: number, L: number,
+  invFreq: MlxArray, rotaryDims: number,
+): MropeForwardState {
+  return { posIds: mropePositionIds(state, offset, L), invFreq, rotaryDims };
+}
+
+/** Request-owned prompt positions; the returned native array belongs to the caller. */
+export function mropePositionIds(state: MropeRequestState, offset: number, L: number): MlxArray {
+  const P = state.positions;
+  const promptLen = P[0].length;
+  const joint = new Int32Array(3 * L);
+  for (let axis = 0; axis < 3; axis++) {
+    for (let l = 0; l < L; l++) {
+      const idx = offset + l;
+      joint[axis * L + l] = idx < promptLen ? P[axis]![idx]! : idx + state.delta;
+    }
+  }
+  return MlxArray.fromInt32(joint, [3, 1, L]);
+}
+
+// The reference language-side apply is a CUSTOM METAL KERNEL (rope_utils
+// _mrope_apply_kernel, pairing "half_split", position_ndim 3): per rotated
+// pair it computes angle = pos * inv_freq in f32, cos/sin in-kernel, and
+// writes x*c − x_pair*s / x_pair*c + x*s with ONE rounding to the tensor
+// dtype. A bf16 cos/sin manual path is ~0.4% off per rotation and moved
+// step-0 logits by ~1.0 — this verbatim port is the oracle's arithmetic.
+const kernelCache = new Map<number, MetalKernel>();
+function mropeKernel(rotaryDims: number): MetalKernel {
+  let k = kernelCache.get(rotaryDims);
+  if (k) return k;
+  const half = rotaryDims / 2;
+  const source = `
+        uint elem = thread_position_in_grid.x;
+
+        const int half_dim = ${half};
+        const int q_bsz = x_shape[0];
+        const int q_heads = x_shape[1];
+        const int q_len = x_shape[2];
+        const int q_dim = x_shape[3];
+        const int slots = half_dim + q_dim - ${rotaryDims};
+        const int work_size = q_bsz * q_heads * q_len * slots;
+
+        if (elem >= uint(work_size)) {
+            return;
+        }
+
+        int local = int(elem);
+        int slot = local % slots;
+        int tmp = local / slots;
+        int t = tmp % q_len;
+        tmp = tmp / q_len;
+        int h = tmp % q_heads;
+        int b = tmp / q_heads;
+        int base = ((b * q_heads + h) * q_len + t) * q_dim;
+
+        if (slot >= half_dim) {
+            int pass_d = ${rotaryDims} + slot - half_dim;
+            int pass_idx = base + pass_d;
+            x_out[pass_idx] = x[pass_idx];
+            return;
+        }
+
+        int freq_idx = slot;
+        int d = freq_idx;
+        int pair_d = d + half_dim;
+        int axis = int(position_selector[freq_idx]);
+        float pos = static_cast<float>(position_ids[(axis * q_bsz + b) * q_len + t]);
+        float angle = pos * static_cast<float>(inv_freq[freq_idx]);
+        float c = metal::cos(angle);
+        float s = metal::sin(angle);
+
+        int idx = base + d;
+        float xv = static_cast<float>(x[idx]);
+        float xp = static_cast<float>(x[base + pair_d]);
+        x_out[idx] = static_cast<T>(xv * c - xp * s);
+        x_out[base + pair_d] = static_cast<T>(xp * c + xv * s);
+`;
+  k = new MetalKernel({
+    name: `mrope_apply_half_split_${rotaryDims}_3d`,
+    inputNames: ["x", "position_ids", "inv_freq", "position_selector"],
+    outputNames: ["x_out"],
+    source,
+    ensureRowContiguous: true,
+  });
+  kernelCache.set(rotaryDims, k);
+  return k;
+}
+
+const selectorCache = new Map<number, MlxArray>();
+function selectorArr(half: number): MlxArray {
+  let s = selectorCache.get(half);
+  if (!s) {
+    s = MlxArray.fromInt32(interleavedSelector(half), [half]);
+    selectorCache.set(half, s);
+  }
+  return s;
+}
+
+/** Apply interleaved mRoPE to one [B, heads, L, headDim] tensor via the
+ *  reference's exact Metal kernel. `invFreq` is the f32 device array from
+ *  mropeInvFreq. */
+export function applyInterleavedRope(
+  x: MlxArray, fwd: MropeForwardState,
+): MlxArray {
+  const rotaryDims = fwd.rotaryDims;
+  const [B, H, L, D] = x.shape as [number, number, number, number];
+  const half = rotaryDims / 2;
+  const slots = half + D - rotaryDims;
+  const work = B * H * L * slots;
+  const [out] = mropeKernel(rotaryDims).apply(
+    [x, fwd.posIds, fwd.invFreq, selectorArr(half)],
+    {
+      outputs: [{ shape: [B, H, L, D], dtype: x.dtype }],
+      grid: [work, 1, 1],
+      threadGroup: [256, 1, 1],
+      templateDtypes: { T: x.dtype },
+    },
+  );
+  return out!;
+}
+
+/** get_rope_index port (B=1, no attention mask): positions per stream + the
+ *  decode delta for a prompt with vision spans. `grids` are the per-image
+ *  (t, h, w) grids in PATCH units and must appear in prompt order. */
+export function qwenRopeIndex(
+  ids: readonly number[],
+  grids: readonly [number, number, number][],
+  mergeSize: number,
+  imageTokenId: number,
+  videoTokenId: number,
+): MropeRequestState {
+  const L = ids.length;
+  const P: [Int32Array, Int32Array, Int32Array] =
+    [new Int32Array(L), new Int32Array(L), new Int32Array(L)];
+  let write = 0; // next output index (== st in the reference walk)
+  let base = 0; // st_idx: next position after everything placed so far
+  let gi = 0;
+  let i = 0;
+  while (i < L) {
+    const tok = ids[i]!;
+    if ((tok === imageTokenId || tok === videoTokenId) && gi < grids.length) {
+      const [t, h, w] = grids[gi]!;
+      gi++;
+      const gh = Math.trunc(h / mergeSize);
+      const gw = Math.trunc(w / mergeSize);
+      const span = t * gh * gw;
+      for (let k = 0; k < span; k++) {
+        const tt = Math.trunc(k / (gh * gw));
+        const rem = k % (gh * gw);
+        const hh = Math.trunc(rem / gw);
+        const ww = rem % gw;
+        P[0][write + k] = base + tt;
+        P[1][write + k] = base + hh;
+        P[2][write + k] = base + ww;
+      }
+      // Reference: next base = max position in the span + 1 = base +
+      // max(t, gh, gw) … computed as llm_pos_ids.max()+1.
+      base = base + Math.max(t, gh, gw);
+      write += span;
+      i += span;
+    } else {
+      P[0][write] = base;
+      P[1][write] = base;
+      P[2][write] = base;
+      base++;
+      write++;
+      i++;
+    }
+  }
+  if (write !== L)
+    throw new Error(`qwenRopeIndex covered ${write} of ${L} tokens (grids misaligned?)`);
+  return { positions: P, delta: base - L };
+}
+
+import { MropeForwardState,MropeRequestState } from "../contracts/mlx/positions";
+export { type MropeForwardState,type MropeRequestState } from "../contracts/mlx/positions";

package/src/layers/rope-2d.ts

@@ -0,0 +1,110 @@
+// Two-dimensional rotary embedding of the Gemma vision towers (optiq's
+// apply_multidimensional_rope): the head is split into one partition per
+// spatial axis, and each partition rotates by its own axis position. The tables
+// are built on device with the oracle's op sequence (arange -> power -> div ->
+// cos/sin -> concat) so they are bit-identical to the reference's; a host-computed
+// table rounds to bf16 differently and, applied to q and k in every layer,
+// compounds.
+
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+
+const AXES = 2;
+
+/** cos/sin for every position, [B, L, 1, headDim] (broadcast over heads). */
+export interface Rope2dTables {
+  readonly cos: MlxArray;
+  readonly sin: MlxArray;
+  dispose(): void;
+}
+
+/** Tables for `positions` [B, L, 2] (integer or float grid coordinates, x then
+ *  y), cast to the activation dtype the tables multiply. */
+export function rope2dTables(positions: MlxArray, headDim: number, theta: number, dtype: Dtype): Rope2dTables {
+  const channelsPerAxis = 2 * Math.floor(headDim / (2 * AXES));
+  if (channelsPerAxis * AXES !== headDim) throw new Error(`2D RoPE needs a head dimension divisible by ${2 * AXES}, got ${headDim}`);
+  const half = channelsPerAxis / 2;
+  const [B, L] = positions.shape as [number, number, number];
+
+  // timescale = theta ** ((2 / channelsPerAxis) * arange(half))
+  const arange = ops.arange(0, half, 1, Dtype.float32);
+  const exponents = ops.mulScalar(arange, 2 / channelsPerAxis);
+  arange.dispose();
+  const base = ops.scalarLike(theta, exponents);
+  const timescale = ops.pow(base, exponents);
+  base.dispose();
+  exponents.dispose();
+
+  const cosParts: MlxArray[] = [];
+  const sinParts: MlxArray[] = [];
+  for (let axis = 0; axis < AXES; axis++) {
+    const column = positions.slice([0, 0, axis], [B, L, axis + 1]);
+    const coordinate = column.astype(Dtype.float32);
+    column.dispose();
+    const angle = ops.div(coordinate, timescale); // [B, L, half]
+    coordinate.dispose();
+    const cos = ops.cos(angle);
+    const sin = ops.sin(angle);
+    angle.dispose();
+    cosParts.push(ops.concatAxis([cos, cos], -1)); // [B, L, channelsPerAxis]
+    sinParts.push(ops.concatAxis([sin, sin], -1));
+    cos.dispose();
+    sin.dispose();
+  }
+  timescale.dispose();
+
+  const assemble = (parts: MlxArray[]): MlxArray => {
+    const full = ops.concatAxis(parts, -1); // [B, L, headDim]
+    for (const part of parts) part.dispose();
+    const shaped = ops.reshape(full, [B, L, 1, headDim]);
+    full.dispose();
+    const cast = shaped.astype(dtype);
+    shaped.dispose();
+    return cast;
+  };
+  const cos = assemble(cosParts);
+  const sin = assemble(sinParts);
+  return { cos, sin, dispose() { cos.dispose(); sin.dispose(); } };
+}
+
+/** rotate_half applied within each axis partition of the head (not across the
+ *  whole head): [-x2, x1] per partition. x: [..., headDim]. */
+function partitionedRotateHalf(x: MlxArray): MlxArray {
+  const shape = x.shape;
+  const last = shape.length - 1;
+  const channelsPerAxis = shape[last]! / AXES;
+  const half = channelsPerAxis / 2;
+  const slice = (from: number, to: number): MlxArray => {
+    const start = shape.map(() => 0);
+    const stop = [...shape];
+    start[last] = from;
+    stop[last] = to;
+    return x.slice(start, stop);
+  };
+  const parts: MlxArray[] = [];
+  const scratch: MlxArray[] = [];
+  for (let axis = 0; axis < AXES; axis++) {
+    const origin = axis * channelsPerAxis;
+    const x1 = slice(origin, origin + half);
+    const x2 = slice(origin + half, origin + channelsPerAxis);
+    const negated = ops.neg(x2);
+    parts.push(ops.concatAxis([negated, x1], last));
+    scratch.push(x1, x2, negated);
+  }
+  const out = ops.concatAxis(parts, last);
+  for (const array of [...scratch, ...parts]) array.dispose();
+  return out;
+}
+
+/** x * cos + rotate_half(x) * sin. x: [B, L, H, headDim]; returns a new array. */
+export function applyRope2d(x: MlxArray, tables: Rope2dTables): MlxArray {
+  const rotated = partitionedRotateHalf(x);
+  const direct = ops.mul(x, tables.cos);
+  const turned = ops.mul(rotated, tables.sin);
+  rotated.dispose();
+  const out = ops.add(direct, turned);
+  direct.dispose();
+  turned.dispose();
+  return out;
+}