@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/layers/compiled-normalization.ts

@@ -0,0 +1,41 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { CompiledFunction } from "@mlx-bun/mlx/compile";
+import { residualRmsNorm } from "../kernels/residual-rms-norm";
+import { gatedRmsNorm } from "../kernels/gated-rms-norm";
+
+// Shared traces take weights as inputs: they retain no model weights or cache
+// state. Shape-aware compilation retraces the custom grid when dimensions change.
+const residualFunctions = new Map<number, CompiledFunction>();
+const gatedFunctions = new Map<number, CompiledFunction>();
+
+/** Residual addition and weighted normalization, preserving both rounding steps. */
+export class CompiledResidualRMSNorm {
+  readonly #operation: CompiledFunction;
+  constructor(readonly weight: MlxArray, readonly eps: number) {
+    let operation = residualFunctions.get(eps);
+    if (!operation) {
+      operation = new CompiledFunction(inputs => residualRmsNorm(inputs[0]!, inputs[1]!, inputs[2]!, eps), false);
+      residualFunctions.set(eps, operation);
+    }
+    this.#operation = operation;
+  }
+  forward(x: MlxArray, update: MlxArray): [MlxArray, MlxArray] {
+    return this.#operation.apply([x, update, this.weight]) as [MlxArray, MlxArray];
+  }
+}
+
+/** Weighted normalization followed by Qwen's float32 SiLU gate. */
+export class CompiledGatedRMSNorm {
+  readonly #operation: CompiledFunction;
+  constructor(readonly weight: MlxArray, readonly eps: number) {
+    let operation = gatedFunctions.get(eps);
+    if (!operation) {
+      operation = new CompiledFunction(inputs => [gatedRmsNorm(inputs[0]!, inputs[1]!, inputs[2]!, eps)], false);
+      gatedFunctions.set(eps, operation);
+    }
+    this.#operation = operation;
+  }
+  forward(x: MlxArray, gate: MlxArray): MlxArray {
+    return this.#operation.apply([x, gate, this.weight])[0]!;
+  }
+}

package/src/layers/delta-heads.ts

@@ -0,0 +1,23 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { CompiledFunction } from "@mlx-bun/mlx/compile";
+import { siluRmsHeads } from "../kernels/delta/silu-rms-heads";
+
+const functions = new Map<string, CompiledFunction>();
+
+/** Shape-aware replay of the fused activation and Q/K RMS norms.
+ * Shared traces take both scale vectors as inputs and retain no layer weights. */
+export class CompiledSiluRMSHeads {
+  readonly #operation: CompiledFunction;
+  constructor(readonly keyHeads: number, readonly valueHeads: number, readonly valueWidth: number) {
+    const key = `${keyHeads}:${valueHeads}:${valueWidth}`;
+    let operation = functions.get(key);
+    if (!operation) {
+      operation = new CompiledFunction(inputs => siluRmsHeads(inputs[0]!, inputs[1]!, inputs[2]!, keyHeads, valueHeads, valueWidth), false);
+      functions.set(key, operation);
+    }
+    this.#operation = operation;
+  }
+  forward(x: MlxArray, qScale: MlxArray, kScale: MlxArray): [MlxArray, MlxArray, MlxArray] {
+    return this.#operation.apply([x, qScale, kScale]) as [MlxArray, MlxArray, MlxArray];
+  }
+}

package/src/layers/dense-embedding.ts

@@ -0,0 +1,29 @@
+// Universal-tier primitives: dense (unquantized) linear/embedding modules
+// (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 QuantizedLinear/QuantizedEmbedding
+// primitives (packages/inference/src/layers/); 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 { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+
+/** 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);
+  }
+}

package/src/layers/dense-linear.ts

@@ -0,0 +1,32 @@
+// Universal-tier primitives: dense (unquantized) linear/embedding modules
+// (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 QuantizedLinear/QuantizedEmbedding
+// primitives (packages/inference/src/layers/); 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 { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+
+/** 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);
+  }
+}

package/src/layers/dequantized-linear.ts

@@ -0,0 +1,72 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { QuantizedLinear } from "./quantized-linear";
+
+// A prompt piece's affine projection over a loaded QuantizedLinear (borrowed
+// weights, no additive bias). From its fewest rows the weight is dequantized to
+// bf16 for the call and multiplied by MLX's dense GEMM: the same bf16 weights
+// MLX's quantized matmul builds in its tiles and the same float accumulation,
+// so the output is that matmul's bit for bit, from the GEMM's larger tiles
+// (measured on the M1 Max at 3 and 4 bits, group 64, every Qwen3.8 projection
+// whose outputs are whole 64-wide tiles: bit-identical at every row count from
+// 97 to 1100 above the split-K range below, and 1.0x to 1.3x faster from 128
+// rows, 7.8 against 5.7 TFLOPS at 2048). Fewer rows, the rows where MLX's
+// dense GEMM splits K (its own accumulation order: while ceil(M/16) ·
+// ceil(N/16) <= 2048 on a Max GPU, so k_proj and v_proj's 1024 outputs up to
+// 512 rows), and the other projections (DeltaNet's in_proj_a and in_proj_b, 48
+// outputs) take MLX's quantized matmul itself. The identity is MLX's kernel
+// choice at those shapes, not a property of every shape (at 512 inputs the two
+// round differently): a graph selects this layer for the artifact it was
+// measured on. The dense weight is a transient of one call (outputs x inputs x
+// 2 bytes). Applies the request's external adapter delta after the projection.
+
+/** The fewest rows the dense GEMM takes at any projection: below this the dequantization costs about what the
+ *  GEMM saves over MLX's quantized matmul. */
+export const DEQUANTIZED_MIN_ROWS = 128;
+
+/** The most tiles of 16 x 16 outputs MLX's dense GEMM runs on its split-K kernel on a Max GPU (the M1 Max's
+ *  `applegpu_g13s`), whose accumulation order is not the quantized matmul's. */
+const SPLIT_K_TILES = 2048;
+
+/** The fewest rows a projection of `outputs` outputs takes the dense GEMM at: DEQUANTIZED_MIN_ROWS or the first row
+ *  count above the GEMM's split-K range at those outputs, whichever is more; never when they are not whole 64-wide
+ *  tiles. */
+export function dequantizedMinRows(outputs: number): number {
+  return outputs % 64 !== 0 ? Infinity : Math.max(DEQUANTIZED_MIN_ROWS, 16 * Math.floor(SPLIT_K_TILES / Math.ceil(outputs / 16)) + 1);
+}
+
+export class DequantizedLinear {
+  readonly #spec: ops.QuantSpec;
+  /** The fewest rows this projection takes the dense GEMM at (`dequantizedMinRows`). */
+  readonly #minRows: number;
+
+  constructor(readonly linear: QuantizedLinear) {
+    if (!linear.biases || linear.bias || linear.w.ndim !== 2 || linear.spec.mode !== "affine")
+      throw new Error("DequantizedLinear: unsupported projection (affine, two-dimensional, no additive bias)");
+    this.#spec = linear.spec;
+    this.#minRows = dequantizedMinRows(linear.outFeatures);
+  }
+
+  /** `x` [.., inputs] (borrowed) to an owned [.., outputs], lazy. */
+  forward(x: MlxArray): MlxArray {
+    const { w, scales, biases } = this.linear;
+    if (x.size / x.shape.at(-1)! < this.#minRows) return this.linear.delta.forward(x, ops.quantizedMatmul(x, w, scales, biases, this.#spec, true));
+    using dense = ops.dequantize(w, scales, biases, this.#spec);
+    using weightT = ops.transposeAxes(dense, [1, 0]);
+    return this.linear.delta.forward(x, ops.matmul(x, weightT));
+  }
+}
+
+/** One fixed dequantize-and-GEMM operation. The graph declares its measured row range. */
+export class DequantizedMatmulLinear {
+  constructor(readonly linear: QuantizedLinear) {
+    if (!linear.biases || linear.bias || linear.w.ndim !== 2 || linear.spec.mode !== "affine")
+      throw new Error("DequantizedMatmulLinear: unsupported projection (affine, two-dimensional, no additive bias)");
+  }
+  forward(x: MlxArray): MlxArray {
+    const {w,scales,biases,spec}=this.linear;
+    using dense=ops.dequantize(w,scales,biases,spec);
+    using weightT=ops.transposeAxes(dense,[1,0]);
+    return this.linear.delta.forward(x,ops.matmul(x,weightT));
+  }
+}

package/src/layers/folded-attention.ts

@@ -0,0 +1,42 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+
+/** The windows MLX's fused SDPA hands to its unfused fallback for want of a
+ *  vector kernel: at most 8 causal queries per row (above that its full
+ *  attention kernel serves), more than 32 query vectors per KV head (the
+ *  vector kernel's limit), no more queries than keys. q [B, H, L, D], keys
+ *  [B, KVH, N, D]. */
+export function foldedSdpaEligible(q: MlxArray, keys: MlxArray): boolean {
+  const [, H, L] = q.shape as [number, number, number, number];
+  const KVH = keys.shape[1]!;
+  return L <= 8 && H % KVH === 0 && L * (H / KVH) > 32 && L <= keys.shape[2]!;
+}
+
+/** Causal attention over bf16 keys and values with GQA folded into rows: the
+ *  rep·L query vectors of each KV head multiply its keys in one matmul and its
+ *  values in another, so the cache is read once per KV head instead of once per
+ *  query head (MLX's fallback broadcasts each KV head over its query heads).
+ *  Otherwise MLX's fallback arithmetic, op for op: the scale cast to the query
+ *  dtype, the bottom-right causal mask at -inf, the precise softmax. Bit for bit
+ *  `ops.sdpa` on the windows `foldedSdpaEligible` takes on `applegpu_g13s`, where
+ *  it was measured, and to bf16 precision elsewhere (rows (r, i): head
+ *  KVH·rep + r, query i, which sees keys j <= N - L + i). q [B, H, L, D], keys
+ *  and values [B, KVH, N, D] (any strides) → [B, H, L, D], owned by the caller. */
+export function foldedSdpa(q: MlxArray, keys: MlxArray, values: MlxArray, scale: number): MlxArray {
+  const [B, H, L, D] = q.shape as [number, number, number, number];
+  const KVH = keys.shape[1]!, N = keys.shape[2]!, rep = H / KVH;
+  using scaled = ops.mulScalar(q, scale);
+  using rows = ops.reshape(scaled, [B, KVH, rep * L, D]);
+  using keysT = ops.transposeAxes(keys, [0, 1, 3, 2]);
+  using scores = ops.matmul(rows, keysT);                          // [B, KVH, rep·L, N]
+  using cols = ops.arange(0, N, 1, Dtype.int32);
+  using colRow = ops.reshape(cols, [1, N]);
+  using limits = ops.fromInt32(Array.from({ length: rep * L }, (_, row) => N - L + (row % L)), [rep * L, 1]);
+  using visible = ops.lessEqual(colRow, limits);
+  using ninf = ops.scalarLike(-Infinity, scores);
+  using masked = ops.where(visible, scores, ninf);
+  using probs = ops.softmaxAxis(masked, -1, true);
+  using out = ops.matmul(probs, values);                            // [B, KVH, rep·L, Dv]
+  return ops.reshape(out, [B, H, L, out.shape[3]!]);
+}

package/src/layers/geglu.ts

@@ -0,0 +1,130 @@
+// GeGLU activation and feed-forward shared by the Gemma family: the compiled
+// geglu closure (mlx-lm's @mx.compile geglu), the spelled-out composition, and
+// the dense quantized MLP built on them. Which one a layer runs is its owner's
+// construction-time choice (`compiled`), never read here.
+
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+
+import { CompiledFunction } from "@mlx-bun/mlx/compile";
+import type { ModelConfig } from "../artifacts/config";
+import type { Weights } from "../artifacts/weights";
+import { isCompiledTrace } from "../runtime/compiled-trace";
+import { QuantizedLinear, QuantizedLinearPlain } from "./quantized-linear";
+
+/** Verbatim port of mlx_lm/models/gemma4_text.py:
+ *    `@partial(mx.compile, shapeless=True) def geglu(gate, x): return nn.gelu_approx(gate) * x`
+ *  The oracle @mx.compile's this and uses it for EVERY MLP shape (dense MLP
+ *  __call__ line 114, and the SwitchGLU GeGLU activation line 149-150). The
+ *  spelled-out composition runs the SAME math UNFUSED (standalone
+ *  ops.geluApprox + ops.mul), a different dispatched kernel set + the per-op
+ *  host tax the MiniCPM5 swiglu fix measured at ~5%. This closure fuses
+ *  geluApprox's ~9 element-wise ops + the final mul into one traced-once/replayed
+ *  graph, matching the oracle's kernel set. `geluApprox` is composed op-for-op
+ *  from the same primitives as the oracle's nn.gelu_approx
+ *  (0.5·x·(1+tanh(√(2/π)(x+0.044715·x³)))), so the compiled graph is the
+ *  oracle's geglu graph. Autograd-safe (verified: the vjp through this closure is
+ *  bit-identical to the plain composition). */
+let _gegluClosure: CompiledFunction | null = null;
+export function compiledGeglu(gate: MlxArray, up: MlxArray): MlxArray {
+  if (!_gegluClosure) {
+    _gegluClosure = new CompiledFunction((inputs) => {
+      const g = inputs[0]!, u = inputs[1]!; // nn.gelu_approx(gate) * x
+      const act = ops.geluApprox(g);
+      const out = ops.mul(act, u); act.dispose();
+      return [out];
+    });
+  }
+  return _gegluClosure.apply([gate, up])[0]!;
+}
+
+/** gelu_approx(gate) * up, consuming both inputs. `compiled` is the owner's
+ *  construction-time choice: the compiled closure (the oracle's @mx.compile'd
+ *  geglu, one kernel set; the default the Gemma graphs compose), or the
+ *  spelled-out operation (same L1
+ *  parity, slower). Inside a compiled-decode trace the closure records the
+ *  spelled-out ops. */
+export function geglu(gate: MlxArray, up: MlxArray, compiled: boolean): MlxArray {
+  let out: MlxArray;
+  if (compiled && !isCompiledTrace()) {
+    out = compiledGeglu(gate, up);
+  } else {
+    const act = ops.geluApprox(gate);
+    out = ops.mul(act, up);
+    act.dispose();
+  }
+  gate.dispose();
+  up.dispose();
+  return out;
+}
+
+/** Dense GeGLU feed-forward: down(geglu(gate(x), up(x))) over quantized linears,
+ *  with the activation form its owner chose at construction (`compiled`). */
+export class GegluMLP {
+  readonly gate: QuantizedLinear;
+  readonly up: QuantizedLinear;
+  readonly down: QuantizedLinear;
+
+  constructor(weights: Weights, config: ModelConfig, prefix: string, readonly compiled: boolean) {
+    this.gate = QuantizedLinear.load(weights, `${prefix}.gate_proj`, config);
+    this.up = QuantizedLinear.load(weights, `${prefix}.up_proj`, config);
+    this.down = QuantizedLinear.load(weights, `${prefix}.down_proj`, config);
+  }
+
+  forward(x: MlxArray): MlxArray {
+    const m = geglu(this.gate.forward(x), this.up.forward(x), this.compiled);
+    const out = this.down.forward(m);
+    m.dispose();
+    return out;
+  }
+}
+
+/** Plain GeGLU consumes both activation inputs. */
+export class PlainGeglu {
+  forward(gate: MlxArray, up: MlxArray): MlxArray {
+    using act = ops.geluApprox(gate);
+    const out = ops.mul(act, up);
+    gate.dispose(); up.dispose();
+    return out;
+  }
+}
+
+/** Compiled GeGLU consumes both activation inputs. */
+export class CompiledGeglu {
+  forward(gate: MlxArray, up: MlxArray): MlxArray {
+    const out = compiledGeglu(gate, up);
+    gate.dispose(); up.dispose();
+    return out;
+  }
+}
+
+export class PlainGegluMLP {
+  readonly #activation = new PlainGeglu();
+  readonly #gate: QuantizedLinearPlain;
+  readonly #up: QuantizedLinearPlain;
+  readonly #down: QuantizedLinearPlain;
+  constructor(readonly base: GegluMLP) {
+    this.#gate = new QuantizedLinearPlain(base.gate);
+    this.#up = new QuantizedLinearPlain(base.up);
+    this.#down = new QuantizedLinearPlain(base.down);
+  }
+  forward(x: MlxArray): MlxArray {
+    using hidden = this.#activation.forward(this.#gate.forward(x), this.#up.forward(x));
+    return this.#down.forward(hidden);
+  }
+}
+export class CompiledGegluMLP {
+  readonly #activation = new CompiledGeglu();
+  readonly #gate: QuantizedLinearPlain;
+  readonly #up: QuantizedLinearPlain;
+  readonly #down: QuantizedLinearPlain;
+  constructor(readonly base: GegluMLP) {
+    this.#gate = new QuantizedLinearPlain(base.gate);
+    this.#up = new QuantizedLinearPlain(base.up);
+    this.#down = new QuantizedLinearPlain(base.down);
+  }
+  forward(x: MlxArray): MlxArray {
+    using hidden = this.#activation.forward(this.#gate.forward(x), this.#up.forward(x));
+    return this.#down.forward(hidden);
+  }
+}

package/src/layers/helpers.ts

@@ -0,0 +1,7 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+
+/** Dispose `old` and return `next` — for h = disposing(h, op(h)) chains. */
+export function disposing(old: MlxArray, next: MlxArray): MlxArray {
+  old.dispose();
+  return next;
+}

package/src/layers/index.ts

@@ -0,0 +1,36 @@
+export * from "./quantized-linear";
+export * from "./quantized-embedding";
+export * from "./normalization";
+export * from "./compiled-normalization";
+export * from "./delta-heads";
+export * from "./trellis-linear";
+export * from "./quantized-switch-linear";
+export * from "./quantized-attention";
+export * from "./quantized-append-attention";
+export * from "../kernels/attention/masks";
+export * from "./helpers";
+export * from "./dense-linear";
+export * from "./dense-embedding";
+export * from "./layer-norm";
+export * from "./loading";
+export * from "./geglu";
+export * from "./swiglu";
+export * from "./qwen-mrope";
+export * from "./rope";
+export * from "./rope-2d";
+
+export * from "./trellis-prompt";
+export * from "./affine-verify-linear";
+export * from "./ane-prefill-split";
+export * from "./trellis-gate-up";
+export * from "./trellis-down";
+export * from "./trellis-reduce-gate-up";
+export * from "./trellis-scatter-down";
+export { LoraDelta, LoraOperation, LoraState } from "./lora";
+
+export * from "./dequantized-linear";
+export * from "./trellis-codebook-down";
+export * from "./trellis-codebook-gate-up";
+export * from "./folded-attention";
+
+export { AneTrellisExpandReduce, AneTrellisExpandScatter, AneTrellisExpandedMlpSplit, AneTrellisExpandedMlpSplitK3i, type AneGpuMlpRows } from "./ane-prefill-split";

package/src/layers/kv4-decode-attention.ts

@@ -0,0 +1,15 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type * as ops from "@mlx-bun/mlx/ops";
+import { kvq4DecodeEligible, kvq4DecodeSdpa } from "../kernels/attention/kvq4-decode";
+
+/** One-row decode attention over a 4-bit, group-64 quantized KV cache with
+ *  head dim 256 (kernels/attention/kvq4-decode): every packed K/V row is read
+ *  once for all query heads of its KV head, `block` tokens per threadgroup.
+ *  q [1, H, 1, 256] → [1, H, 1, 256]. Not bit-identical to quantizedSdpa
+ *  (bf16-level). */
+export function kv4DecodeAttention(q: MlxArray, keys: ops.QuantizedTensor, values: ops.QuantizedTensor, scale: number,
+  groupSize: number, bits: number, block: number): MlxArray {
+  if (!kvq4DecodeEligible(q, keys, groupSize, bits))
+    throw new Error("kv4DecodeAttention: needs one query row, head dim 256, a 4-bit group-64 cache and an even GQA factor");
+  return kvq4DecodeSdpa(q, keys, values, scale, block);
+}

package/src/layers/layer-norm.ts

@@ -0,0 +1,24 @@
+// Universal-tier primitives: dense (unquantized) linear/embedding modules
+// (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 QuantizedLinear/QuantizedEmbedding
+// primitives (packages/inference/src/layers/); 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 { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+
+/** 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);
+  }
+}

package/src/layers/loading.ts

@@ -0,0 +1,92 @@
+// Universal-tier primitives: dense (unquantized) linear/embedding modules
+// (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 QuantizedLinear/QuantizedEmbedding
+// primitives (packages/inference/src/layers/); 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 * as ops from "@mlx-bun/mlx/ops";
+import type { ModelConfig } from "../artifacts/config";
+import { WeightAudit,tensorUsed } from "../artifacts/weight-audit";
+import type { Weights } from "../artifacts/weights";
+import { DenseEmbedding } from "./dense-embedding";
+import { DenseLinear } from "./dense-linear";
+import { LayerNorm } from "./layer-norm";
+import { RMSNorm } from "./normalization";
+import { QuantizedEmbedding } from "./quantized-embedding";
+import { QuantizedLinear } from "./quantized-linear";
+import { TrellisLinear } from "./trellis-linear";
+
+/** Either linear flavor behind one call surface. */
+export type AnyLinear = QuantizedLinear | DenseLinear | TrellisLinear;
+
+/** Quantized when `.scales` exists (MLX-quantized checkpoints), else dense
+ *  bf16/f16/f32. Both flavors carry the optional ADDITIVE
+ *  `.bias` term (qwen2 qkv, starcoder2, …). */
+export function loadLinear(
+  weights: Weights, path: string, config: ModelConfig, audit: WeightAudit, representations?: import("../kernels/trellis/index").TrellisRepresentations,
+): AnyLinear {
+  if (TrellisLinear.isTrellis(config, path)) {
+    const lin = TrellisLinear.load(weights, path, config, false, representations);
+    audit.use(`${path}.weight`);
+    audit.use(`${path}.scales`);
+    return lin;
+  }
+  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,
+  );
+}
+
+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);
+}
+
+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,
+  );
+}