@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/layers/rope.ts

@@ -0,0 +1,389 @@
+// Universal rope factory — faithful port of mlx_lm/models/rope_utils.py
+// (oracle venv, mlx-lm 0.31.3): initialize_rope with the default / linear /
+// llama3 (Llama3RoPE) / yarn (YarnRoPE) / longrope (SuScaledRoPE) /
+// proportional (ProportionalRoPE) variants.
+//
+// Porting discipline (three-level fidelity tree, L1): every frequency
+// table is computed with the SAME mlx float32 ops the oracle uses at
+// module init — not host-side float64 math — so the freqs bytes match the
+// oracle's bit-for-bit. Host math (Math.log/floor/ceil) is used only where
+// the oracle uses python-float math (yarn correction range, mscale, su
+// scale), which is IEEE double on both sides.
+//
+// The existing runtime paths (ops.rope default, gemma4's proportional
+// freqs) are untouched; this factory serves the universal (Tier-0) models.
+
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+
+/** Raw HF `rope_scaling` dict (values read exactly like the oracle). */
+export type RopeScalingConfig = Record<string, unknown>;
+
+const num = (v: unknown, fallback: number): number =>
+  typeof v === "number" ? v : fallback;
+
+const requireNum = (sc: RopeScalingConfig, key: string): number => {
+  const v = sc[key];
+  if (typeof v !== "number") throw new Error(`rope_scaling.${key} missing or not a number`);
+  return v;
+};
+
+/** One rope instance (shared across layers — identical params ⇒ identical
+ *  values). `apply` returns a NEW array; the caller owns disposal of both. */
+export class UniversalRope {
+  constructor(
+    readonly dims: number,
+    readonly traditional: boolean,
+    /** Base for native fast_rope; null when `freqs` supplies the table. */
+    readonly base: number | null,
+    /** fast_rope scale (1/factor for "linear", else 1). */
+    readonly scale: number,
+    /** Precomputed frequency table (llama3/yarn/longrope/proportional). */
+    readonly freqs: MlxArray | null,
+    /** Pre-rope input scale on x[..., :dims] (yarn mscale / su scale). */
+    readonly preScale: number,
+    /** Apply preScale even when it is exactly 1.0 (SuScaledRoPE does; Yarn
+     *  skips when mscale == 1.0). */
+    readonly preScaleAlways: boolean,
+  ) {}
+
+  apply(x: MlxArray, offset: number): MlxArray {
+    const input = this.#preScaled(x);
+    const out = ops.ropeScaled(
+      input ?? x, this.dims, this.traditional, this.base, this.scale, offset, this.freqs,
+    );
+    input?.dispose();
+    return out;
+  }
+
+  /** Per-row array-offset variant (batched decode under left-padding: each
+   *  row's REAL position rides an int32 [B] array — the scalar offset would
+   *  mis-position every padded row by its pad amount). Same math as
+   *  `apply`; only the offset spelling differs. */
+  applyDynamic(x: MlxArray, offsetArr: MlxArray): MlxArray {
+    const input = this.#preScaled(x);
+    const out = ops.ropeScaledDynamic(
+      input ?? x, this.dims, this.traditional, this.base, this.scale, offsetArr, this.freqs,
+    );
+    input?.dispose();
+    return out;
+  }
+
+  /** x[..., :dims] = preScale * x[..., :dims] (weak-scalar mul); null when
+   *  no pre-scale applies (caller uses x unchanged). */
+  #preScaled(x: MlxArray): MlxArray | null {
+    if (!(this.preScaleAlways || this.preScale !== 1.0)) return null;
+    const D = x.shape[x.shape.length - 1]!;
+    if (this.dims === D) return ops.mulScalar(x, this.preScale);
+    const start = x.shape.map(() => 0);
+    const stop = [...x.shape];
+    stop[stop.length - 1] = this.dims;
+    const head = x.slice(start, stop);
+    const headScaled = ops.mulScalar(head, this.preScale);
+    head.dispose();
+    const scaled = ops.sliceUpdate(x, headScaled, start, stop);
+    headScaled.dispose();
+    return scaled;
+  }
+
+  dispose(): void {
+    this.freqs?.dispose();
+  }
+}
+
+/** SmolLM3 NoPE: identity stand-in for disabled-rope layers. */
+export class NoRope extends UniversalRope {
+  constructor() {
+    super(0, false, null, 1, null, 1, false);
+  }
+  override apply(x: MlxArray, _offset: number): MlxArray {
+    // Return an owned copy-view so callers can dispose uniformly.
+    return x.slice(x.shape.map(() => 0), [...x.shape]);
+  }
+  override applyDynamic(x: MlxArray, _offsetArr: MlxArray): MlxArray {
+    return x.slice(x.shape.map(() => 0), [...x.shape]);
+  }
+}
+
+/** `base ** (arange(0, dims, 2) / dims)` with the oracle's dtypes:
+ *  int32 arange for llama3 (no explicit dtype in the source), float32
+ *  arange for yarn/longrope. Integer÷integer divides to float32 in mlx,
+ *  so both spellings compute the pow in f32. */
+function baseFreqs(dims: number, base: number, arangeDtype: Dtype): MlxArray {
+  const ar = ops.arange(0, dims, 2, arangeDtype);
+  const dimsScalar = arangeDtype === Dtype.int32
+    ? ops.fromInt32([dims], [])
+    : MlxArray.fromFloat32(new Float32Array([dims]), []);
+  const exps = ops.div(ar, dimsScalar); // → float32
+  ar.dispose();
+  dimsScalar.dispose();
+  const baseScalar = ops.scalarLike(base, exps);
+  const out = ops.pow(baseScalar, exps);
+  baseScalar.dispose();
+  exps.dispose();
+  return out; // float32 [dims/2]
+}
+
+/** Llama3RoPE frequency table (rope_utils.py:74-108, verbatim). */
+function llama3Freqs(dims: number, base: number, sc: RopeScalingConfig): MlxArray {
+  const factor = requireNum(sc, "factor");
+  const lowFreqFactor = num(sc.low_freq_factor, 1.0);
+  const highFreqFactor = num(sc.high_freq_factor, 4.0);
+  const oldContextLen = num(sc.original_max_position_embeddings, 8192);
+
+  const lowFreqWavelen = oldContextLen / lowFreqFactor;
+  const highFreqWavelen = oldContextLen / highFreqFactor;
+
+  let freqs = baseFreqs(dims, base, Dtype.int32);
+  const wavelens = ops.mulScalar(freqs, 2 * Math.PI);
+
+  // freqs = where(wavelens > low_freq_wavelen, freqs * factor, freqs)
+  const lfw = ops.scalarLike(lowFreqWavelen, wavelens);
+  const gtLow = ops.less(lfw, wavelens);
+  const freqsTimesFactor = ops.mulScalar(freqs, factor);
+  const freqs1 = ops.where(gtLow, freqsTimesFactor, freqs);
+  for (const a of [gtLow, freqsTimesFactor, freqs]) a.dispose();
+  freqs = freqs1;
+
+  // is_medium = (wavelens > high_freq_wavelen) & (wavelens < low_freq_wavelen)
+  const hfw = ops.scalarLike(highFreqWavelen, wavelens);
+  const gtHigh = ops.less(hfw, wavelens);
+  const ltLow = ops.less(wavelens, lfw);
+  const isMedium = ops.logicalAnd(gtHigh, ltLow);
+  for (const a of [hfw, lfw, gtHigh, ltLow]) a.dispose();
+
+  // smooth_factors = (old_context_len / wavelens - low) / (high - low)
+  const oldScalar = ops.scalarLike(oldContextLen, wavelens);
+  const oldOverWav = ops.div(oldScalar, wavelens);
+  oldScalar.dispose();
+  wavelens.dispose();
+  const lowScalar = ops.scalarLike(lowFreqFactor, oldOverWav);
+  const numArr = ops.sub(oldOverWav, lowScalar);
+  oldOverWav.dispose();
+  lowScalar.dispose();
+  const denScalar = ops.scalarLike(highFreqFactor - lowFreqFactor, numArr);
+  const smooth = ops.div(numArr, denScalar);
+  numArr.dispose();
+  denScalar.dispose();
+
+  // smooth_freqs = freqs / ((1 - smooth) / factor + smooth)
+  const one = ops.scalarLike(1, smooth);
+  const oneMinus = ops.sub(one, smooth);
+  one.dispose();
+  const factorScalar = ops.scalarLike(factor, oneMinus);
+  const scaledDown = ops.div(oneMinus, factorScalar);
+  oneMinus.dispose();
+  factorScalar.dispose();
+  const denom = ops.add(scaledDown, smooth);
+  scaledDown.dispose();
+  smooth.dispose();
+  const smoothFreqs = ops.div(freqs, denom);
+  denom.dispose();
+
+  const out = ops.where(isMedium, smoothFreqs, freqs);
+  for (const a of [isMedium, smoothFreqs, freqs]) a.dispose();
+  return out;
+}
+
+/** YarnRoPE (rope_utils.py:128-181, verbatim): returns [freqs, mscale]. */
+function yarnFreqs(
+  dims: number,
+  base: number,
+  scalingFactor: number,
+  originalMaxPositionEmbeddings: number,
+  betaFast: number,
+  betaSlow: number,
+  mscale: number,
+  mscaleAllDim: number,
+): [MlxArray, number] {
+  const findCorrectionDim = (numRotations: number): number =>
+    (dims * Math.log(originalMaxPositionEmbeddings / (numRotations * 2 * Math.PI))) /
+    (2 * Math.log(base));
+  let low = Math.floor(findCorrectionDim(betaFast));
+  let high = Math.ceil(findCorrectionDim(betaSlow));
+  low = Math.max(low, 0);
+  high = Math.min(high, dims - 1);
+  const getMscale = (scale: number, m: number): number =>
+    scale <= 1 ? 1.0 : 0.1 * m * Math.log(scale) + 1.0;
+  const outMscale =
+    getMscale(scalingFactor, mscale) / getMscale(scalingFactor, mscaleAllDim);
+
+  const freqExtra = baseFreqs(dims, base, Dtype.float32);
+  const freqInter = ops.mulScalar(freqExtra, scalingFactor);
+
+  // freq_mask = 1 - linear_ramp_mask(low, high, dims // 2)
+  let maxVal = high;
+  if (low === maxVal) maxVal += 0.001; // prevent singularity (oracle comment)
+  const arangeHalf = ops.arange(0, Math.floor(dims / 2), 1, Dtype.float32);
+  const minScalar = ops.scalarLike(low, arangeHalf);
+  const shifted = ops.sub(arangeHalf, minScalar);
+  arangeHalf.dispose();
+  minScalar.dispose();
+  const rangeScalar = ops.scalarLike(maxVal - low, shifted);
+  const linear = ops.div(shifted, rangeScalar);
+  shifted.dispose();
+  rangeScalar.dispose();
+  const lo = ops.scalarLike(0, linear);
+  const hi = ops.scalarLike(1, linear);
+  const ramp = ops.clip(linear, lo, hi);
+  for (const a of [linear, lo, hi]) a.dispose();
+  const one = ops.scalarLike(1, ramp);
+  const freqMask = ops.sub(one, ramp);
+  ramp.dispose();
+
+  // freqs = (inter * extra) / (inter * mask + extra * (1 - mask))
+  const numArr = ops.mul(freqInter, freqExtra);
+  const interMasked = ops.mul(freqInter, freqMask);
+  const oneMinusMask = ops.sub(one, freqMask);
+  one.dispose();
+  const extraMasked = ops.mul(freqExtra, oneMinusMask);
+  oneMinusMask.dispose();
+  const den = ops.add(interMasked, extraMasked);
+  const freqs = ops.div(numArr, den);
+  for (const a of [freqExtra, freqInter, freqMask, numArr, interMasked, extraMasked, den])
+    a.dispose();
+  return [freqs, outMscale];
+}
+
+/** SuScaledRoPE (rope_utils.py:10-71): [freqs, scale]. NOTE the oracle uses
+ *  long_factor for ALL positions (short_factor is accepted and ignored). */
+function suFreqsAndScale(
+  dims: number,
+  base: number,
+  maxPositionEmbeddings: number,
+  originalMaxPositionEmbeddings: number,
+  longFactor: number[] | number,
+  longMscale?: number,
+): [MlxArray, number] {
+  const freqs0 = baseFreqs(dims, base, Dtype.float32);
+  const lf = Array.isArray(longFactor)
+    ? MlxArray.fromFloat32(new Float32Array(longFactor), [longFactor.length])
+    : MlxArray.fromFloat32(new Float32Array([longFactor]), []);
+  const freqs = ops.mul(lf, freqs0);
+  lf.dispose();
+  freqs0.dispose();
+
+  const factor = maxPositionEmbeddings / originalMaxPositionEmbeddings;
+  const defaultScale =
+    factor <= 1.0
+      ? 1.0
+      : Math.sqrt(1 + Math.log(factor) / Math.log(originalMaxPositionEmbeddings));
+  // python: `long_mscale or default` — 0 is falsy there too.
+  const scale = longMscale ? longMscale : defaultScale;
+  return [freqs, scale];
+}
+
+/** ProportionalRoPE (rope_utils.py:199-232): rotate the first rotated_dims,
+ *  ∞-freq (identity) padding for the rest. */
+function proportionalFreqs(
+  dims: number, rotatedDims: number, base: number, factor: number,
+): MlxArray {
+  if (rotatedDims > dims) throw new Error("rotated_dims should be smaller than dims");
+  const ar = ops.arange(0, rotatedDims, 2, Dtype.float32);
+  const dimsScalar = MlxArray.fromFloat32(new Float32Array([dims]), []);
+  const exps = ops.div(ar, dimsScalar);
+  ar.dispose();
+  dimsScalar.dispose();
+  const baseScalar = ops.scalarLike(base, exps);
+  const powed = ops.pow(baseScalar, exps);
+  baseScalar.dispose();
+  exps.dispose();
+  const scaled = ops.mulScalar(powed, factor);
+  powed.dispose();
+  const padLen = Math.floor((dims - rotatedDims) / 2);
+  if (padLen === 0) return scaled;
+  const pad = MlxArray.fromFloat32(
+    new Float32Array(padLen).fill(Number.POSITIVE_INFINITY), [padLen],
+  );
+  const out = ops.concatAxis([scaled, pad], 0);
+  scaled.dispose();
+  pad.dispose();
+  return out;
+}
+
+/** What a scaling type builds its rope from. `scaling` is empty for the plain `default` type. */
+interface RopeInit {
+  dims: number;
+  base: number;
+  traditional: boolean;
+  scaling: RopeScalingConfig;
+  maxPositionEmbeddings: number | null;
+}
+
+const yarnRope = ({ dims, base, traditional, scaling: sc }: RopeInit): UniversalRope => {
+  const [freqs, mscale] = yarnFreqs(
+    dims,
+    base,
+    requireNum(sc, "factor"),
+    num(sc.original_max_position_embeddings, 4096),
+    num(sc.beta_fast, 32),
+    num(sc.beta_slow, 1),
+    num(sc.mscale, 1),
+    num(sc.mscale_all_dim, 0),
+  );
+  // YarnRoPE scales x only when mscale != 1.0 (preScaleAlways=false).
+  return new UniversalRope(dims, traditional, null, 1.0, freqs, mscale, false);
+};
+
+/** The `rope_scaling` types initialize_rope understands, keyed by the config's type string (data,
+ *  not a model family). A new scaling scheme is one entry here. */
+const ropeScalers: Readonly<Record<string, (init: RopeInit) => UniversalRope>> = {
+  default: ({ dims, base, traditional }) => new UniversalRope(dims, traditional, base, 1.0, null, 1, false),
+  linear: ({ dims, base, traditional, scaling }) =>
+    new UniversalRope(dims, traditional, base, 1 / requireNum(scaling, "factor"), null, 1, false),
+  llama3: ({ dims, base, traditional, scaling }) =>
+    new UniversalRope(dims, traditional, null, 1.0, llama3Freqs(dims, base, scaling), 1, false),
+  yarn: yarnRope,
+  deepseek_yarn: yarnRope,
+  "telechat3-yarn": yarnRope,
+  longrope: ({ dims, base, maxPositionEmbeddings, scaling: sc }) => {
+    const [freqs, scale] = suFreqsAndScale(
+      dims,
+      base,
+      maxPositionEmbeddings ?? 131072,
+      requireNum(sc, "original_max_position_embeddings"),
+      (sc.long_factor as number[] | number | undefined) ?? 1.0,
+    );
+    // SuScaledRoPE hardcodes traditional=False and ALWAYS pre-scales.
+    return new UniversalRope(dims, false, null, 1.0, freqs, scale, true);
+  },
+  proportional: ({ dims, base, traditional, scaling: sc }) => {
+    const rotatedDims = Math.floor(dims * num(sc.partial_rotary_factor, 1.0));
+    const freqs = proportionalFreqs(dims, rotatedDims, base, num(sc.factor, 1.0));
+    return new UniversalRope(dims, traditional, null, 1.0, freqs, 1, false);
+  },
+};
+
+/** Port of rope_utils.initialize_rope — the shared entry every
+ *  initialize_rope-style arch (llama, qwen2/3, olmo2, granite, …) uses. */
+export function initializeRope(
+  dims: number,
+  base: number,
+  traditional: boolean,
+  scalingConfig: RopeScalingConfig | null,
+  maxPositionEmbeddings: number | null,
+): UniversalRope {
+  const ropeType = scalingConfig
+    ? String(scalingConfig.type ?? scalingConfig.rope_type ?? "default")
+    : "default";
+  const scaler = Object.hasOwn(ropeScalers, ropeType) ? ropeScalers[ropeType] : undefined;
+  if (!scaler) throw new Error(`Unsupported RoPE type ${ropeType}`);
+  return scaler({ dims, base, traditional, scaling: scalingConfig ?? {}, maxPositionEmbeddings });
+}
+
+/** SuScaledRoPE constructed the phi3.py way (NOT via initialize_rope):
+ *  original_max_position_embeddings comes from the TOP-LEVEL config, and
+ *  rope_dim = head_dim * partial_rotary_factor. */
+export function phi3SuRope(
+  ropeDim: number,
+  base: number,
+  maxPositionEmbeddings: number,
+  originalMaxPositionEmbeddings: number,
+  longFactor: number[] | number,
+): UniversalRope {
+  const [freqs, scale] = suFreqsAndScale(
+    ropeDim, base, maxPositionEmbeddings, originalMaxPositionEmbeddings, longFactor,
+  );
+  return new UniversalRope(ropeDim, false, null, 1.0, freqs, scale, true);
+}

package/src/layers/swiglu.ts

@@ -0,0 +1,23 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { CompiledFunction } from "@mlx-bun/mlx/compile";
+import * as ops from "@mlx-bun/mlx/ops";
+
+
+/** Verbatim port of mlx_lm/models/activations.py:
+ *    `@partial(mx.compile, shapeless=True) def swiglu(gate, x): return nn.silu(gate) * x`
+ *  mx.compile fuses sigmoid + mul + mul into ONE kernel (mlx-lm's decode
+ *  `CV2ISigmoid…Multiply`) instead of our three separate dispatches. Traced
+ *  once, replayed thereafter for decode and prefill. */
+let _swigluClosure: CompiledFunction | null = null;
+export function compiledSwiglu(gate: MlxArray, up: MlxArray): MlxArray {
+  if (!_swigluClosure) {
+    _swigluClosure = new CompiledFunction((inputs) => {
+      const g = inputs[0]!, u = inputs[1]!;              // nn.silu(gate) * x
+      const sig = ops.sigmoid(g);
+      const silu = ops.mul(g, sig); sig.dispose();
+      const out = ops.mul(silu, u); silu.dispose();
+      return [out];
+    });
+  }
+  return _swigluClosure.apply([gate, up])[0]!;
+}

package/src/layers/trellis-codebook-down.ts

@@ -0,0 +1,52 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { downCodebookRow, downK3InterleavedCodebookRow } from "../kernels/trellis/down-codebook";
+import type { TrellisLinear } from "./trellis-linear";
+
+// Packed down-projection layers over a Qwen MLP's axis-0 Trellis codes at one
+// row, with the scatter kernels' output bits (trellisScatter; kernels/trellis/index.ts)
+// through a threadgroup codebook. Each layer runs one kernel for one code layout
+// (row-major 2- or 4-bit codes, or the 3-bit block interleave) and checks the
+// projection at construction. forward(x) keeps x's leading dimensions.
+
+function project(lin: TrellisLinear, x: MlxArray, kernel: (x2: MlxArray) => MlxArray): MlxArray {
+  const g = lin.geometry, lead = x.shape.slice(0, -1);
+  using x2 = ops.reshape(x, [lead.reduce((a, b) => a * b, 1), g.inFeatures]);
+  using y = kernel(x2);
+  return ops.reshape(y, [...lead, g.outFeatures]);
+}
+function rowMajor(name: string, lin: TrellisLinear): void {
+  const g = lin.geometry;
+  if (g.axis !== 0 || g.blockInterleave || lin.codes.ndim !== 2 || (g.k === 3 && g.T === 256 && g.L <= 12))
+    throw new Error(`${name}: needs packed row-major axis-0 codes the balanced kernels do not decode (3-bit codes take the K3i layers)`);
+}
+function k3Interleaved(name: string, lin: TrellisLinear): void {
+  const g = lin.geometry;
+  if (g.axis !== 0 || g.blockInterleave !== 2 || g.k !== 3 || g.T !== 256 || g.L !== 12)
+    throw new Error(`${name}: needs 3-bit block-interleaved axis-0 codes`);
+}
+
+/** One activation row; row-major 2- or 4-bit codes at T 256, decoding through a threadgroup codebook with
+ *  word-aligned lanes (kernels/trellis/down-codebook.ts): trellisScatter's single-row output bit for bit. */
+export class CodebookDownRow {
+  constructor(readonly down: TrellisLinear) {
+    rowMajor("CodebookDownRow", down);
+    const g = down.geometry;
+    if ((g.k !== 2 && g.k !== 4) || g.T !== 256 || g.L > 12) throw new Error("CodebookDownRow: needs 2- or 4-bit codes at T 256, L <= 12");
+  }
+
+  forward(x: MlxArray): MlxArray {
+    const { codes, scales, geometry } = this.down;
+    return this.down.delta.forward(x, project(this.down, x, x2 => downCodebookRow(x2, codes, scales, geometry)));
+  }
+}
+
+/** One activation row; 3-bit block-interleaved codes, the same way: trellisScatter's single-row output bit for bit. */
+export class CodebookDownK3iRow {
+  constructor(readonly down: TrellisLinear) { k3Interleaved("CodebookDownK3iRow", down); }
+
+  forward(x: MlxArray): MlxArray {
+    const { codes, scales, geometry } = this.down;
+    return this.down.delta.forward(x, project(this.down, x, x2 => downK3InterleavedCodebookRow(x2, codes, scales, geometry)));
+  }
+}

package/src/layers/trellis-codebook-gate-up.ts

@@ -0,0 +1,42 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { gateUpCodebookRow } from "../kernels/trellis/gate-up-codebook";
+import type { MixedGateUpTail } from "../kernels/trellis/mixed-gate-up";
+import type { TrellisLinear } from "./trellis-linear";
+
+// The fused gate/up/SwiGLU layer over a Qwen MLP's axis-1 Trellis gate and up
+// projections at one row, decoding through a threadgroup codebook with the reduce
+// kernels' output bits (kernels/trellis/gate-up-codebook.ts). It runs one kernel
+// and takes everything that kernel varies on at construction.
+
+/** `sameWidth` null: either. */
+function sameAxis1(name: string, gate: TrellisLinear, up: TrellisLinear, sameWidth: boolean | null): void {
+  const a = gate.geometry, b = up.geometry;
+  if (a.axis !== 1 || b.axis !== 1 || a.rows !== b.rows || a.cols !== b.cols ||
+      a.T !== b.T || a.L !== b.L || (sameWidth !== null && (a.k === b.k) !== sameWidth))
+    throw new Error(`${name}: gate/up need packed axis-1 codes of one geometry${sameWidth === null ? ""
+      : `, ${sameWidth ? "the same" : "different"} widths`}`);
+  if (gate.variant !== up.variant) throw new Error(`${name}: gate and up decode with different variants`);
+}
+
+/** One activation row; gate and up at the same or different widths, decoding through a threadgroup codebook
+ *  (kernels/trellis/gate-up-codebook.ts) with fusedGateUpSwiglu's (same widths) and fusedGateUpSwigluMixed's (different
+ *  widths) output bits. `mixedTail` is the SwiGLU arithmetic where the widths differ (fusedGateUpSwigluMixed's `tail`);
+ *  same widths take the float32 sigmoid. */
+export class CodebookGateUpRow {
+  constructor(readonly gate: TrellisLinear, readonly up: TrellisLinear, readonly mixedTail: MixedGateUpTail) {
+    sameAxis1("CodebookGateUpRow", gate, up, null);
+    if (gate.geometry.L > 12) throw new Error("CodebookGateUpRow: the codebook holds at most 2^12 states");
+  }
+
+  /** `x` [.., in] (borrowed, one row) to an owned [.., rows]. */
+  forward(x: MlxArray): MlxArray { return gateUpCodebookRow(x, this.gate, this.up, this.mixedTail); }
+}
+
+/** The mixed-width split activation tail, fixed for written two-argument bindings. */
+export class CodebookSplitGateUpRow {
+  constructor(readonly gate:TrellisLinear,readonly up:TrellisLinear) {
+    sameAxis1("CodebookSplitGateUpRow",gate,up,null);
+    if (gate.geometry.L>12) throw new Error("CodebookSplitGateUpRow: the codebook holds at most 2^12 states");
+  }
+  forward(x:MlxArray):MlxArray {return gateUpCodebookRow(x,this.gate,this.up,"split");}
+}

package/src/layers/trellis-down.ts

@@ -0,0 +1,117 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { TrellisGeometry } from "../kernels/trellis/geometry";
+import { downFactoredRow, downFactoredRows } from "../kernels/trellis/down-factored";
+import { downK3InterleavedFactoredRow, downK3InterleavedFactoredRows } from "../kernels/trellis/down-k3-interleaved-factored";
+import { downK3InterleavedMma, downMma } from "../kernels/trellis/down-mma";
+import {
+  DOWN_ONE_BLOCK, downK3InterleavedMmaWide, downK3InterleavedMmaWidePass, downMmaWide, downMmaWidePass,
+} from "../kernels/trellis/down-mma-wide";
+import { TrellisLinear } from "./trellis-linear";
+
+// Down-projection layers over a Qwen MLP's axis-0 Trellis codes. Each layer
+// runs one kernel for one row form and one code layout (row-major, or the
+// 3-bit block interleave); the owning graph picks the layer for the request's
+// row count and the weights' layout. forward(x) keeps x's leading dimensions.
+
+type Kernel = (x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry) => MlxArray;
+
+function project(lin: TrellisLinear, kernel: Kernel, x: MlxArray): MlxArray {
+  const g = lin.geometry, lead = x.shape.slice(0, -1);
+  using x2 = ops.reshape(x, [lead.reduce((a, b) => a * b, 1), g.inFeatures]);
+  using y = kernel(x2, lin.codes, lin.scales, g);
+  return lin.delta.forward(x, ops.reshape(y, [...lead, g.outFeatures]));
+}
+// The any-row codebook down kernel in one block of 8 rows per pass, fixed for the 1..8-row layers.
+const oneBlockDown: Kernel = (x2, codes, scales, g) => downMmaWidePass(x2, codes, scales, g, DOWN_ONE_BLOCK);
+const oneBlockDownK3i: Kernel = (x2, codes, scales, g) => downK3InterleavedMmaWidePass(x2, codes, scales, g, DOWN_ONE_BLOCK);
+
+function rowMajor(name: string, lin: TrellisLinear): void {
+  const g = lin.geometry;
+  if (g.axis !== 0 || g.blockInterleave || lin.codes.ndim !== 2) throw new Error(`${name}: needs packed row-major axis-0 codes`);
+}
+function k3Interleaved(name: string, lin: TrellisLinear): void {
+  const g = lin.geometry;
+  if (g.axis !== 0 || g.blockInterleave !== 2 || g.k !== 3 || g.T !== 256 || g.L !== 12)
+    throw new Error(`${name}: needs 3-bit block-interleaved axis-0 codes`);
+}
+
+/** One activation row; row-major codes. */
+export class FactoredDownRow {
+  constructor(readonly down: TrellisLinear) { rowMajor("FactoredDownRow", down); }
+  forward(x: MlxArray): MlxArray { return project(this.down, downFactoredRow, x); }
+}
+
+/** 2..4 activation rows; row-major codes. */
+export class FactoredDownRows {
+  constructor(readonly down: TrellisLinear) { rowMajor("FactoredDownRows", down); }
+  forward(x: MlxArray): MlxArray { return project(this.down, downFactoredRows, x); }
+}
+
+/** One activation row; 3-bit block-interleaved codes. */
+export class FactoredDownK3iRow {
+  constructor(readonly down: TrellisLinear) { k3Interleaved("FactoredDownK3iRow", down); }
+  forward(x: MlxArray): MlxArray { return project(this.down, downK3InterleavedFactoredRow, x); }
+}
+
+/** 2..4 activation rows; 3-bit block-interleaved codes. */
+export class FactoredDownK3iRows {
+  constructor(readonly down: TrellisLinear) { k3Interleaved("FactoredDownK3iRows", down); }
+  forward(x: MlxArray): MlxArray { return project(this.down, downK3InterleavedFactoredRows, x); }
+}
+
+/** 1..8 activation rows on the simdgroup matrix unit; row-major codes. */
+export class MmaDown {
+  constructor(readonly down: TrellisLinear) {
+    rowMajor("MmaDown", down);
+    if (down.geometry.T !== 256 || down.geometry.L > 12) throw new Error("MmaDown: needs T 256, L <= 12");
+  }
+  forward(x: MlxArray): MlxArray { return project(this.down, downMma, x); }
+}
+
+/** 1..8 activation rows on the simdgroup matrix unit; 3-bit block-interleaved codes. */
+export class MmaDownK3i {
+  constructor(readonly down: TrellisLinear) { k3Interleaved("MmaDownK3i", down); }
+  forward(x: MlxArray): MlxArray { return project(this.down, downK3InterleavedMma, x); }
+}
+
+/** Any number of activation rows on the simdgroup matrix unit through a threadgroup codebook, each decoded fragment
+ *  feeding up to three 8-row blocks (MmaDown's bits per block); row-major codes. */
+export class MmaDownWide {
+  constructor(readonly down: TrellisLinear) {
+    rowMajor("MmaDownWide", down);
+    const g = down.geometry;
+    if (g.T !== 256 || g.L > 12 || g.rows % 8 !== 0) throw new Error("MmaDownWide: needs T 256, L <= 12, rows % 8");
+  }
+  forward(x: MlxArray): MlxArray { return project(this.down, downMmaWide, x); }
+}
+
+/** Any number of activation rows on the simdgroup matrix unit through a threadgroup codebook (MmaDownK3i's bits per
+ *  block of 8); 3-bit block-interleaved codes. */
+export class MmaDownWideK3i {
+  constructor(readonly down: TrellisLinear) {
+    k3Interleaved("MmaDownWideK3i", down);
+    if (down.geometry.rows % 8 !== 0) throw new Error("MmaDownWideK3i: needs rows % 8");
+  }
+  forward(x: MlxArray): MlxArray { return project(this.down, downK3InterleavedMmaWide, x); }
+}
+
+/** 1..8 activation rows on the simdgroup matrix unit through a threadgroup codebook, one block of 8 rows per pass (the
+ *  any-row kernel's one-block pass; MmaDown's bits); row-major codes. */
+export class CodebookMmaDown {
+  constructor(readonly down: TrellisLinear) {
+    rowMajor("CodebookMmaDown", down);
+    const g = down.geometry;
+    if (g.T !== 256 || g.L > 12 || g.rows % 8 !== 0) throw new Error("CodebookMmaDown: needs T 256, L <= 12, rows % 8");
+  }
+  forward(x: MlxArray): MlxArray { return project(this.down, oneBlockDown, x); }
+}
+
+/** CodebookMmaDown for 3-bit block-interleaved codes (MmaDownK3i's bits). */
+export class CodebookMmaDownK3i {
+  constructor(readonly down: TrellisLinear) {
+    k3Interleaved("CodebookMmaDownK3i", down);
+    if (down.geometry.rows % 8 !== 0) throw new Error("CodebookMmaDownK3i: needs rows % 8");
+  }
+  forward(x: MlxArray): MlxArray { return project(this.down, oneBlockDownK3i, x); }
+}