@mlx-bun/inference 0.0.0-stage → 1.0.0

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

@@ -0,0 +1,71 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { gateUpFactoredRow, gateUpFactoredRows } from "../kernels/trellis/gate-up-factored";
+import { mixedGateUpFactoredRows } from "../kernels/trellis/mixed-gate-up-factored";
+import { gateUpMma } from "../kernels/trellis/gate-up-mma";
+import { gateUpMmaWide } from "../kernels/trellis/gate-up-mma-wide";
+import { gateUpCodebookMma } from "../kernels/trellis/gate-up-codebook-mma";
+import { TrellisLinear } from "./trellis-linear";
+
+// Fused gate/up/SwiGLU layers over a Qwen MLP's axis-1 Trellis gate and up
+// projections. Each layer runs one kernel for one row form; the owning graph
+// picks the layer for the request's row count and the weights' bit widths.
+
+function sameAxis1(name: string, gate: TrellisLinear, up: TrellisLinear, sameWidth: boolean): 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 || (a.k === b.k) !== sameWidth)
+    throw new Error(`${name}: gate/up need packed axis-1 codes of one geometry, ${sameWidth ? "the same" : "different"} widths`);
+}
+
+/** One activation row; gate and up at the same width. */
+export class FactoredGateUpRow {
+  constructor(readonly gate: TrellisLinear, readonly up: TrellisLinear) { sameAxis1("FactoredGateUpRow", gate, up, true); }
+  forward(x: MlxArray): MlxArray { return gateUpFactoredRow(x, this.gate, this.up); }
+}
+
+/** 2..4 activation rows sharing each decoded weight; same widths. */
+export class FactoredGateUpRows {
+  constructor(readonly gate: TrellisLinear, readonly up: TrellisLinear) { sameAxis1("FactoredGateUpRows", gate, up, true); }
+  forward(x: MlxArray): MlxArray { return gateUpFactoredRows(x, this.gate, this.up); }
+}
+
+/** 1..4 activation rows; gate and up at different widths. */
+export class FactoredMixedGateUp {
+  constructor(readonly gate: TrellisLinear, readonly up: TrellisLinear) { sameAxis1("FactoredMixedGateUp", gate, up, false); }
+  forward(x: MlxArray): MlxArray { return mixedGateUpFactoredRows(x, this.gate, this.up); }
+}
+
+/** 1..8 activation rows on the simdgroup matrix unit; widths may differ. */
+export class MmaGateUp {
+  constructor(readonly gate: TrellisLinear, readonly up: TrellisLinear) {
+    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 || a.rows % 8 !== 0 || a.cols % 256 !== 0)
+      throw new Error("MmaGateUp: gate/up need packed axis-1 codes of one geometry (rows % 8, cols % 256)");
+  }
+  forward(x: MlxArray): MlxArray { return gateUpMma(x, this.gate, this.up); }
+}
+
+/** Any number of activation rows on the simdgroup matrix unit, each decoded weight fragment feeding several 8-row
+ *  blocks (a threadgroup codebook, activations staged in threadgroup memory); widths may differ. */
+export class MmaGateUpWide {
+  constructor(readonly gate: TrellisLinear, readonly up: TrellisLinear) {
+    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 || a.rows % 8 !== 0 || a.cols % 256 !== 0 || a.L > 12)
+      throw new Error("MmaGateUpWide: gate/up need packed axis-1 codes of one geometry (rows % 8, cols % 256, L <= 12)");
+  }
+  forward(x: MlxArray): MlxArray { return gateUpMmaWide(x, this.gate, this.up); }
+}
+
+/** 1..8 activation rows on the simdgroup matrix unit, each code value read from a threadgroup codebook (MmaGateUp's
+ *  bits); widths may differ. */
+export class CodebookMmaGateUp {
+  constructor(readonly gate: TrellisLinear, readonly up: TrellisLinear) {
+    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 || a.rows % 8 !== 0 || a.cols % 256 !== 0 || a.L > 12)
+      throw new Error("CodebookMmaGateUp: gate/up need packed axis-1 codes of one geometry (rows % 8, cols % 256, L <= 12)");
+  }
+  forward(x: MlxArray): MlxArray { return gateUpCodebookMma(x, this.gate, this.up); }
+}

package/src/layers/trellis-linear.ts

@@ -0,0 +1,211 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype, deviceArchitecture } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import { quantFor, type ModelConfig, type QuantSpec } from "../artifacts/config";
+import type { Weights } from "../artifacts/weights";
+import { LoraDelta, type LoraState, type LoraWeights } from "./lora";
+import { QuantizedLinear } from "./quantized-linear";
+import { wordsPerBlock } from "../kernels/trellis/codebook";
+import {
+  TRELLIS_MATVEC_MAX_M as MATVEC_MAX_M, type TrellisGeometry,
+  expandTrellis, trellisReduce, trellisScatter, fusedGateUpSwiglu as gateUpKernel,
+  tiledTrellisPrefill, tiledTrellisPrefillEligible,
+  splitKTrellisPrefill, splitKTrellisPrefillEligible,
+  wideTrellisPrefill, wideTrellisPrefillEligible,
+  qualifiedRepresentations, type TrellisRepresentations,
+} from "../kernels/trellis/index";
+export { TRELLIS_MATVEC_MAX_M, type TrellisGeometry, type MixedGateUpTail } from "../kernels/trellis/index";
+
+/** The decode variant these layers pass the kernels (semantics in
+ *  kernels/trellis/index.ts): 13, which passed the M4 Pro closeout matrix.
+ *  Benches and tests that compare variants call the kernels directly. */
+const VARIANT = 13;
+
+export function trellisGeometry(codes: MlxArray, spec: QuantSpec): TrellisGeometry {
+  const tr = spec.trellis;
+  if (spec.mode !== "trellis" || !tr) throw new Error("trellisGeometry: spec is not a trellis spec");
+  const k = spec.bits, T = spec.groupSize;
+  if (codes.ndim === 3) {
+    const [groups, rows, words] = codes.shape as [number, number, number];
+    if (codes.dtype !== Dtype.uint32 || tr.axis !== 0 || k !== 3 || T !== 256 ||
+        tr.L !== 12 || tr.code !== "1mad" || words !== 48 || groups < 1 || rows < 1)
+      throw new Error("trellis: unsupported interleaved code shape or quantization");
+    const cols = groups * 512;
+    return { k, L: tr.L, T, axis: 0, rows, cols, inFeatures: rows,
+      outFeatures: cols, blockInterleave: 2 };
+  }
+  if (codes.ndim !== 2) throw new Error("trellis: expected a 2D or interleaved 3D code tensor");
+  const [rows, words] = codes.shape as [number, number];
+  if ((words * 32) % k !== 0) throw new Error(`trellis: ${words} words not a whole number of ${k}-bit symbols`);
+  const cols = (words * 32) / k;
+  if (cols % T !== 0) throw new Error(`trellis: ${cols} coded columns not a multiple of block ${T}`);
+  wordsPerBlock(T, k);
+  return {
+    k, L: tr.L, T, axis: tr.axis, rows, cols,
+    inFeatures: tr.axis === 1 ? cols : rows,
+    outFeatures: tr.axis === 1 ? rows : cols,
+  };
+}
+
+/** Can gate and up be served by the fused kernel? Same axis-1 geometry and k. */
+export function fusedGateUpEligible(gate: TrellisLinear, up: TrellisLinear): boolean {
+  const a = gate.geometry, b = up.geometry;
+  return a.axis === 1 && b.axis === 1 &&
+    a.k === b.k && a.rows === b.rows && a.cols === b.cols && a.T === b.T && a.L === b.L;
+}
+
+/** Borrow the layers' weights; the gate's stored representations choose the decode. */
+export function fusedGateUpSwiglu(x: MlxArray, gate: TrellisLinear, up: TrellisLinear): MlxArray {
+  return gateUpKernel(x, gate, up, VARIANT, gate.representations.gateUpCodebook);
+}
+
+/** Output rows [start, stop) of an axis-1 (row-coded) Trellis projection as an
+ *  owned projection: each output row has its own code row and scale, so the
+ *  slice decodes exactly the same values. A leading-axis range of row-major
+ *  codes is already row-contiguous, so these are views of the parent's memory. */
+export function trellisRows(lin: TrellisLinear, start: number, stop: number): TrellisLinear {
+  const g = lin.geometry;
+  if (g.axis !== 1 || g.blockInterleave || lin.codes.ndim !== 2) throw new Error("trellisRows: needs row-coded (axis 1) codes");
+  const codes = lin.codes.slice([start, 0], [stop, lin.codes.shape[1]!]);
+  const scales = lin.scales.slice([start], [stop]);
+  ops.evalAll([codes, scales]);
+  return new TrellisLinear(codes, scales, lin.spec, false, lin.representations);
+}
+
+/** Stored rows [start, stop) of an axis-0 Trellis projection (input features of
+ *  a down projection): 2D codes slice to a view; block-interleaved 3D codes
+ *  [groups, rows, words] slice their middle axis (kernels copy that view). */
+export function trellisStoredRows(lin: TrellisLinear, start: number, stop: number): TrellisLinear {
+  const g = lin.geometry;
+  if (g.axis !== 0) throw new Error("trellisStoredRows: needs axis-0 codes");
+  const codes = lin.codes.ndim === 3
+    ? lin.codes.slice([0, start, 0], [lin.codes.shape[0]!, stop, lin.codes.shape[2]!])
+    : lin.codes.slice([start, 0], [stop, lin.codes.shape[1]!]);
+  const scales = lin.scales.slice([start], [stop]);
+  ops.evalAll([codes, scales]);
+  return new TrellisLinear(codes, scales, lin.spec, false, lin.representations);
+}
+
+export class TrellisLinear {
+  readonly delta: 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; }
+  readonly geometry: TrellisGeometry;
+  readonly spec: QuantSpec;
+  /** The decode variant every kernel over these weights runs (kernels/trellis/index.ts): the row-form layers
+   *  built on this projection take it at their construction. */
+  readonly variant = VARIANT;
+  /** What the kernels take from this machine's GPU family (single-row
+   *  representations, wide prefill), fixed here and passed on every call. */
+  readonly representations: TrellisRepresentations;
+
+  constructor(
+    readonly codes: MlxArray,
+    readonly scales: MlxArray,
+    spec: QuantSpec,
+    readonly useSharedScatterCodebook = false,
+    representations: TrellisRepresentations = qualifiedRepresentations(deviceArchitecture()),
+  ) {
+    this.spec = spec;
+    this.geometry = trellisGeometry(codes, spec);
+    this.delta = new LoraDelta(this.geometry.inFeatures, this.geometry.outFeatures);
+    if (this.geometry.blockInterleave && (scales.ndim !== 1 || scales.size !== this.geometry.rows))
+      throw new Error("trellis: interleaved codes require one scale per stored row");
+    this.representations = representations;
+  }
+
+  static load(weights: Weights, path: string, config: ModelConfig, useSharedScatterCodebook = false, representations?: TrellisRepresentations): TrellisLinear {
+    const spec = quantFor(config.quantization, path);
+    if (!spec || spec.mode !== "trellis")
+      throw new Error(`${path}: expected a trellis quant spec`);
+    if (!weights.has(`${path}.scales`)) throw new Error(`${path}: trellis tensor has no .scales`);
+    return new TrellisLinear(weights.tensor(`${path}.weight`), weights.tensor(`${path}.scales`), spec, useSharedScatterCodebook, representations);
+  }
+
+  static isTrellis(config: ModelConfig, path: string): boolean {
+    return quantFor(config.quantization, path)?.mode === "trellis";
+  }
+
+  get inFeatures(): number { return this.geometry.inFeatures; }
+  get outFeatures(): number { return this.geometry.outFeatures; }
+
+  /** The weight as [out, in] bf16 (decoded; transposed for axis=0). */
+  expandWeight(dtype: Dtype = Dtype.bfloat16): MlxArray {
+    const g = this.geometry;
+    const stored = expandTrellis(this.codes, this.scales, g, dtype, VARIANT);
+    if (g.axis === 1) return stored;
+    const t = ops.transposeAxes(stored, [1, 0]);
+    const w = ops.contiguous(t);
+    t.dispose(); stored.dispose();
+    return w;
+  }
+
+  /** A caller may prove row-contiguous, aligned input from an allocating op
+   * such as RMSNorm. Lazy array strides cannot establish that before eval.
+   * Rows past the tiled kernels expand the dense weight and multiply; that
+   * product is evaluated before returning, so building the next projection
+   * never holds a second expansion, unless `boundExpansion` is false: a caller
+   * that queues several projections at once (the Neural Engine split's GPU
+   * share) holds their expansions until it evaluates them. */
+  forward(x: MlxArray, inputRowContiguous = false, boundExpansion = true): MlxArray {
+    const g = this.geometry;
+    const lead = x.shape.slice(0, -1);
+    const M = lead.reduce((a, b) => a * b, 1);
+    if (x.shape[x.shape.length - 1] !== g.inFeatures)
+      throw new Error(`TrellisLinear: input dim ${x.shape[x.shape.length - 1]} != ${g.inFeatures}`);
+    const x2 = ops.reshape(x, [M, g.inFeatures]);
+    let y: MlxArray;
+    let expandedWeights = false;
+    if (M <= MATVEC_MAX_M) y = g.axis === 1 ? trellisReduce(x2, this.codes, this.scales, g, VARIANT) : trellisScatter(x2, this.codes, this.scales, g, VARIANT, this.useSharedScatterCodebook, this.representations.scatter);
+    else if (inputRowContiguous && wideTrellisPrefillEligible(g, M, x.dtype, this.representations.widePrefill))
+      y = wideTrellisPrefill(x2, this.codes, this.scales, g);
+    else if (tiledTrellisPrefillEligible(g, M, x.dtype, this.representations.widePrefill))
+      y = tiledTrellisPrefill(x2, this.codes, this.scales, g);
+    else if (splitKTrellisPrefillEligible(g, M, x.dtype))
+      y = splitKTrellisPrefill(x2, this.codes, this.scales, g);
+    else {
+      expandedWeights = true;
+      const stored = expandTrellis(this.codes, this.scales, g, x.dtype, VARIANT);
+      if (g.axis === 1) {
+        const wt = ops.transposeAxes(stored, [1, 0]);
+        y = ops.matmul(x2, wt);
+        wt.dispose();
+      } else y = ops.matmul(x2, stored);
+      stored.dispose();
+    }
+    x2.dispose();
+    const out = ops.reshape(y, [...lead, g.outFeatures]);
+    y.dispose();
+    // Bound dense weight expansions before building the next projection.
+    // Disposing `stored` above releases only its JS handle, not the lazy
+    // matmul's reference. Direct tiles and packed matvecs stay lazy: they
+    // hold activations and bounded partials, without a dense weight matrix.
+    if (expandedWeights && boundExpansion) out.eval();
+    return this.delta.forward(x, out);
+  }
+}
+
+/** Packed Trellis weights decoded once, at construction, into 8-bit g64 affine
+ *  and served through the stock QuantizedLinear (the eval-carrier numerics,
+ *  about -45 dB against the packed kernels): for a machine where the packed
+ *  kernels lose. Built explicitly by whoever selects it; holds no packed codes. */
+export class ExpandedTrellisLinear {
+  readonly carrier: QuantizedLinear;
+
+  constructor(packed: TrellisLinear) {
+    const w = packed.expandWeight(Dtype.bfloat16);
+    const q = ops.quantize(w, 64, 8, "affine");
+    ops.evalAll([q.packed, q.scales, ...(q.biases ? [q.biases] : [])]);
+    w.dispose();
+    this.carrier = new QuantizedLinear(q.packed, q.scales, q.biases, { bits: 8, groupSize: 64, mode: "affine" });
+  }
+
+  get inFeatures(): number { return this.carrier.inFeatures; }
+  get outFeatures(): number { return this.carrier.outFeatures; }
+
+  forward(x: MlxArray): MlxArray { return this.carrier.forward(x); }
+}

package/src/layers/trellis-prompt.ts

@@ -0,0 +1,102 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { expandTrellis, trellisReduce, trellisScatter, wideTrellisPrefill, tiledTrellisPrefill, splitKTrellisPrefill } from "../kernels/trellis/index";
+import { LoraOperation, LoraState } from "./lora";
+import { compiledSwiglu } from "./swiglu";
+import { fusedGateUpSwiglu, type TrellisLinear } from "./trellis-linear";
+
+export interface RowProjection { forward(x: MlxArray): MlxArray }
+
+/** The graph supplies a projection kernel. Shapes preserve the input's leading
+ * dimensions; reshaped handles are released while the result remains lazy. */
+function project(lin: TrellisLinear, x: MlxArray, kernel: (x: MlxArray) => MlxArray): MlxArray {
+  const lead = x.shape.slice(0, -1);
+  using rows = ops.reshape(x, [lead.reduce((a, b) => a * b, 1), lin.inFeatures]);
+  using out = kernel(rows);
+  return lin.delta.forward(x, ops.reshape(out, [...lead, lin.outFeatures]));
+}
+
+export class TrellisReducePrompt {
+  constructor(readonly linear: TrellisLinear) {}
+  forward(x: MlxArray): MlxArray {
+    return project(this.linear, x, rows => trellisReduce(rows, this.linear.codes, this.linear.scales, this.linear.geometry, this.linear.variant));
+  }
+}
+export class TrellisScatterPrompt {
+  constructor(readonly linear: TrellisLinear) {}
+  forward(x: MlxArray): MlxArray {
+    return project(this.linear, x, rows => trellisScatter(rows, this.linear.codes, this.linear.scales, this.linear.geometry, this.linear.variant, this.linear.useSharedScatterCodebook, this.linear.representations.scatter));
+  }
+}
+export class TrellisWidePrompt {
+  constructor(readonly linear: TrellisLinear) {}
+  forward(x: MlxArray): MlxArray {
+    return project(this.linear, x, rows => wideTrellisPrefill(rows, this.linear.codes, this.linear.scales, this.linear.geometry));
+  }
+}
+export class TrellisTiledPrompt {
+  constructor(readonly linear: TrellisLinear) {}
+  forward(x: MlxArray): MlxArray {
+    return project(this.linear, x, rows => tiledTrellisPrefill(rows, this.linear.codes, this.linear.scales, this.linear.geometry));
+  }
+}
+export class TrellisSplitKPrompt {
+  constructor(readonly linear: TrellisLinear) {}
+  forward(x: MlxArray): MlxArray {
+    return project(this.linear, x, rows => splitKTrellisPrefill(rows, this.linear.codes, this.linear.scales, this.linear.geometry));
+  }
+}
+
+/** Axis-1 dense expansion; evaluate the matmul before the next projection so
+ * the deferred graph does not retain every expanded weight. */
+export class TrellisExpandReducePrompt {
+  constructor(readonly linear: TrellisLinear) {}
+  forward(x: MlxArray): MlxArray {
+    const out = project(this.linear, x, rows => {
+      using stored = expandTrellis(this.linear.codes, this.linear.scales, this.linear.geometry, x.dtype, this.linear.variant);
+      using wt = ops.transposeAxes(stored, [1, 0]);
+      return ops.matmul(rows, wt);
+    });
+    out.eval();
+    return out;
+  }
+}
+export class TrellisExpandScatterPrompt {
+  constructor(readonly linear: TrellisLinear) {}
+  forward(x: MlxArray): MlxArray {
+    const out = project(this.linear, x, rows => {
+      using stored = expandTrellis(this.linear.codes, this.linear.scales, this.linear.geometry, x.dtype, this.linear.variant);
+      return ops.matmul(rows, stored);
+    });
+    out.eval();
+    return out;
+  }
+}
+
+/** Fused gate/up/SwiGLU followed by the literal down projection. */
+export class TrellisFusedMLP {
+  readonly #operation: LoraOperation;
+  constructor(readonly gate: TrellisLinear, readonly up: TrellisLinear, readonly down: RowProjection,
+    state?: LoraState) {
+
+
+    const adapted = new TrellisUnfusedMLP(new TrellisReducePrompt(gate), new TrellisReducePrompt(up), down);
+    this.#operation = new LoraOperation(x => {
+      using hidden = fusedGateUpSwiglu(x, gate, up);
+      return down.forward(hidden);
+    }, x => adapted.forward(x), state, [gate.delta, up.delta]);
+  }
+  forward(x: MlxArray): MlxArray { return this.#operation.forward(x); }
+  dispose(): void { this.#operation.dispose(); }
+}
+
+/** Independent projections followed by the oracle's compiled SwiGLU. */
+export class TrellisUnfusedMLP {
+  constructor(readonly gate: RowProjection, readonly up: RowProjection, readonly down: RowProjection) {}
+  forward(x: MlxArray): MlxArray {
+    using gate = this.gate.forward(x);
+    using up = this.up.forward(x);
+    using hidden = compiledSwiglu(gate, up);
+    return this.down.forward(hidden);
+  }
+}

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

@@ -0,0 +1,73 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { gateUpReduceRow, gateUpReduceRows } from "../kernels/trellis/gate-up";
+import { fusedGateUpSwigluMixed, type MixedGateUpTail } from "../kernels/trellis/mixed-gate-up";
+import type { TrellisLinear } from "./trellis-linear";
+
+// Fused gate/up/SwiGLU layers over a Qwen MLP's axis-1 Trellis gate and up
+// projections on the reduce kernels (one SIMD group per output row, the decoded
+// weight times the scale per weight; kernels/trellis/index.ts). Each layer runs
+// one kernel for one row form and takes everything that kernel varies on at
+// construction: the projections' decode variant and the representations their
+// GPU family qualifies (`TrellisLinear.representations`). The owning graph picks
+// the layer for the rows of its width class and the weights' bit widths.
+
+function sameAxis1(name: string, gate: TrellisLinear, up: TrellisLinear, sameWidth: boolean): 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 || (a.k === b.k) !== sameWidth)
+    throw new Error(`${name}: gate/up need packed axis-1 codes of one geometry, ${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 width. The threadgroup codebook decode where the gate's GPU family
+ *  qualifies it (applied by the kernel at the measured artifact shape). */
+export class ReduceGateUpRow {
+  readonly #variant: number;
+  readonly #codebook: boolean;
+
+  constructor(readonly gate: TrellisLinear, readonly up: TrellisLinear) {
+    sameAxis1("ReduceGateUpRow", gate, up, true);
+    this.#variant = gate.variant;
+    this.#codebook = gate.representations.gateUpCodebook;
+  }
+
+  /** `x` [.., in] (borrowed, one row) to an owned [.., rows]. */
+  forward(x: MlxArray): MlxArray { return gateUpReduceRow(x, this.gate, this.up, this.#variant, this.#codebook); }
+}
+
+/** 2..4 activation rows sharing each decoded weight; gate and up at the same width. */
+export class ReduceGateUpRows {
+  readonly #variant: number;
+
+  constructor(readonly gate: TrellisLinear, readonly up: TrellisLinear) {
+    sameAxis1("ReduceGateUpRows", gate, up, true);
+    this.#variant = gate.variant;
+  }
+
+  /** `x` [.., in] (borrowed, 2..4 rows) to an owned [.., rows]. */
+  forward(x: MlxArray): MlxArray { return gateUpReduceRows(x, this.gate, this.up, this.#variant); }
+}
+
+/** 1..4 activation rows; gate and up at different widths. `tail` is the SwiGLU arithmetic the kernel reproduces
+ *  (`"split"`: MLX's compiled swiglu over the two bf16 projections, bit for bit; `"fused"`: the float32 sigmoid of the
+ *  same-width kernel). The exact-bits decoder at one row where the gate's GPU family qualifies it. */
+export class ReduceMixedGateUp {
+  readonly #variant: number;
+  readonly #exactBits: boolean;
+
+  constructor(readonly gate: TrellisLinear, readonly up: TrellisLinear, readonly tail: MixedGateUpTail) {
+    sameAxis1("ReduceMixedGateUp", gate, up, false);
+    this.#variant = gate.variant;
+    this.#exactBits = gate.representations.mixedExactBits;
+  }
+
+  /** `x` [.., in] (borrowed, 1..4 rows) to an owned [.., rows]. */
+  forward(x: MlxArray): MlxArray {
+    return fusedGateUpSwigluMixed(x, this.gate, this.up, this.#variant, this.tail, this.#exactBits);
+  }
+}
+
+/** Fixed split activation for the written M1 graph's mixed-width projections. */
+export class ReduceSplitMixedGateUp extends ReduceMixedGateUp {
+  constructor(gate: TrellisLinear, up: TrellisLinear) { super(gate, up, "split"); }
+}

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

@@ -0,0 +1,98 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { balancedScatterRow, balancedScatterRows, scatterRow, scatterRows } from "../kernels/trellis/scatter";
+import type { ScatterRepresentation } from "../kernels/trellis/qualified";
+import type { TrellisLinear } from "./trellis-linear";
+
+// Packed down-projection layers over a Qwen MLP's axis-0 Trellis codes on the
+// scatter kernels (lane per code word, split-K partials folded by one sum;
+// kernels/trellis/index.ts). Each layer runs one kernel for one row form and one
+// code layout (row-major 2- or 4-bit codes, or the 3-bit block interleave, which
+// the balanced kernels decode) and takes everything that kernel varies on at
+// construction: the projection's decode variant, the single-row representations
+// its GPU family qualifies and whether its shared-row codebook was composed
+// (`TrellisLinear.useSharedScatterCodebook`). 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 lin.delta.forward(x, 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 codes. */
+export class ScatterDownRow {
+  readonly #variant: number;
+  readonly #representation: ScatterRepresentation;
+
+  constructor(readonly down: TrellisLinear) {
+    rowMajor("ScatterDownRow", down);
+    this.#variant = down.variant;
+    this.#representation = down.representations.scatter;
+  }
+
+  forward(x: MlxArray): MlxArray {
+    const { codes, scales, geometry } = this.down;
+    return project(this.down, x, x2 => scatterRow(x2, codes, scales, geometry, this.#variant, this.#representation));
+  }
+}
+
+/** 2..4 activation rows sharing each decoded weight; row-major codes. */
+export class ScatterDownRows {
+  readonly #variant: number;
+
+  constructor(readonly down: TrellisLinear) {
+    rowMajor("ScatterDownRows", down);
+    this.#variant = down.variant;
+  }
+
+  forward(x: MlxArray): MlxArray {
+    const { codes, scales, geometry } = this.down;
+    return project(this.down, x, x2 => scatterRows(x2, codes, scales, geometry, this.#variant));
+  }
+}
+
+/** One activation row; 3-bit block-interleaved codes. */
+export class ScatterDownK3iRow {
+  readonly #variant: number;
+  readonly #representation: ScatterRepresentation;
+
+  constructor(readonly down: TrellisLinear) {
+    k3Interleaved("ScatterDownK3iRow", down);
+    this.#variant = down.variant;
+    this.#representation = down.representations.scatter;
+  }
+
+  forward(x: MlxArray): MlxArray {
+    const { codes, scales, geometry } = this.down;
+    return project(this.down, x, x2 => balancedScatterRow(x2, codes, scales, geometry, this.#variant, this.#representation));
+  }
+}
+
+/** 2..4 activation rows sharing each decoded weight; 3-bit block-interleaved codes, with the shared-row codebook at
+ *  three and four rows when the projection composed it. */
+export class ScatterDownK3iRows {
+  readonly #variant: number;
+  readonly #sharedCodebook: boolean;
+
+  constructor(readonly down: TrellisLinear) {
+    k3Interleaved("ScatterDownK3iRows", down);
+    this.#variant = down.variant;
+    this.#sharedCodebook = down.useSharedScatterCodebook;
+  }
+
+  forward(x: MlxArray): MlxArray {
+    const { codes, scales, geometry } = this.down;
+    return project(this.down, x, x2 => balancedScatterRows(x2, codes, scales, geometry, this.#variant, this.#sharedCodebook));
+  }
+}

package/src/models/ane-install.ts

@@ -0,0 +1,74 @@
+// The install step for the Neural Engine: every program the shipped compositions
+// load on this machine, compiled once into the system compile cache, the way
+// Apple's own pipeline compiles a model for the Neural Engine on the device and
+// keeps it. A graph never compiles at load (kernels/ane/linear); it loads what
+// this step left in the cache, about 4 ms a program.
+//
+// The table is the loader's: each graph table entry whose served device is this
+// machine's (or any) names the programs it loads for a composition
+// (`GraphEntry.anePrograms`), a function of the composition alone, so the step
+// needs no downloaded model. Programs identical in MIL and weights are one
+// program and compile once. Compiling loads nothing and holds none of the ANE
+// driver's clients. The cache is per code-signing identity, so the step runs
+// under the executable that serves (the app's `compile-ane` verb).
+import type { Composition } from "../contracts/portable/composition";
+import { AneProgram, aneProgramKey, type AneCompileStatus, type AneProgramSpec } from "../kernels/ane/linear";
+import { machineComposition } from "./composition";
+import { MLX_GRAPH_TABLES } from "./factory";
+import type { GraphTable } from "./implementation";
+
+/** One distinct program of the install step and the graph entries that load it. */
+export interface AneInstallProgram {
+  readonly spec: AneProgramSpec;
+  readonly graphs: readonly string[];
+}
+
+/** An install step program and what compiling it did. */
+export interface AneInstallResult extends AneInstallProgram {
+  readonly status: AneCompileStatus;
+  readonly milliseconds: number;
+}
+
+export interface AneInstallOptions {
+  /** The record compiled for. Default: {@link machineComposition}, what a load with no flags resolves here. */
+  readonly composition?: Composition;
+  /** The graph tables read. Default: the loader's ({@link MLX_GRAPH_TABLES}). */
+  readonly tables?: Readonly<Record<string, GraphTable<unknown, unknown>>>;
+}
+
+/** Every distinct Neural Engine program the entries of `tables` load for `composition`, among the entries whose
+ * served device is `composition.device` or unconstrained, in table order. */
+export function aneInstallPrograms(options: AneInstallOptions = {}): AneInstallProgram[] {
+  const composition = options.composition ?? machineComposition();
+  const programs = new Map<string, { spec: AneProgramSpec; graphs: string[] }>();
+  for (const table of Object.values(options.tables ?? MLX_GRAPH_TABLES))
+    for (const entry of table.entries) {
+      if (!entry.anePrograms || (entry.serves.device !== undefined && entry.serves.device !== composition.device)) continue;
+      for (const spec of entry.anePrograms()) {
+        const key = aneProgramKey(spec), known = programs.get(key);
+        if (known) { if (!known.graphs.includes(entry.id)) known.graphs.push(entry.id); }
+        else programs.set(key, { spec, graphs: [entry.id] });
+      }
+    }
+  return [...programs.values()];
+}
+
+/** Each program {@link aneInstallPrograms} lists and whether this Mac's Neural Engine compile cache holds it for
+ * this process (the framework's `compiledModelExists`). Compiles nothing. */
+export function aneInstallStatus(options: AneInstallOptions = {}): (AneInstallProgram & { readonly cached: boolean })[] {
+  return aneInstallPrograms(options).map(program => ({ ...program, cached: AneProgram.cached(program.spec) }));
+}
+
+/** The install step: compiles every program {@link aneInstallPrograms} lists into this Mac's Neural Engine
+ * compile cache, skipping those already there, and returns what it did for each, in order (`onProgram` hears each
+ * as it finishes). Throws `AneBridgeError` when the bridge does not load or the compiler refuses a program. */
+export function compileAnePrograms(options: AneInstallOptions & { readonly onProgram?: (result: AneInstallResult) => void } = {}):
+  AneInstallResult[] {
+  return aneInstallPrograms(options).map(program => {
+    const started = performance.now();
+    const status = AneProgram.compile(program.spec);
+    const result = { ...program, status, milliseconds: performance.now() - started };
+    options.onProgram?.(result);
+    return result;
+  });
+}