@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/kernels/trellis/gate-up-mma.ts

@@ -0,0 +1,148 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+import type { TrellisWeights } from "./geometry";
+import { HEADER } from "./codebook";
+
+// Verify-width fused gate/up/SwiGLU on the simdgroup matrix unit: one
+// simdgroup owns an 8-output-row tile for up to 8 activation rows. The K index
+// inside each 8x8 step is permuted so that A column i at step s is position
+// base + 32·i + s: every lane decodes two contiguous 32-position runs per
+// projection from K+1 words (the matvec's word window, unrolled) straight into
+// its A fragment and reads its two activation rows contiguously, so per-weight
+// decode work matches the one-row matvec. Raw code values y accumulate in f32;
+// scale·(1/147.8) applies once per output. Gate and up may differ in bit width.
+// Not bit-identical to the per-lane matvec reduction (different summation
+// order); gate it by KL. Fragment coordinates follow MLX steel's Metal layout.
+const SG_TG = 4, MT = 1;
+
+const RUN3_HEADER = String.raw`
+template <int K, int BT_>
+inline void load_run(const device uint32_t* rowCodes, uint run, thread uint* w) {
+  const uint wpb = (uint)(BT_ * K / 32);
+  const uint runsPerBlock = (uint)BT_ / 32u;
+  const uint blk = run / runsPerBlock;
+  const uint t0 = (run - blk * runsPerBlock) * 32u;
+  const device uint32_t* bw = rowCodes + blk * wpb;
+  const uint w0 = ((uint)BT_ - 32u - t0) * (uint)K / 32u;
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i <= (uint)K; ++i) {
+    const uint wi = w0 + i;
+    w[i] = bw[wi >= wpb ? wi - wpb : wi];
+  }
+}
+template <int K, int L_>
+inline uint run_state(thread const uint* w, uint j) {
+  const uint b = (31u - j) * (uint)K;
+  const uint wi = b >> 5;
+  const uint off = b & 31u;
+  uint win = w[wi] >> off;
+  if (off + (uint)L_ > 32u) win |= w[wi + 1] << (32u - off);
+  return win & ((1u << (uint)L_) - 1u);
+}
+`;
+const SOURCE_V3 = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint n0 = (threadgroup_position_in_grid.y * (uint)SG_TG + sg) * 8u;
+  if (n0 >= (uint)R) return;
+  const uint fm = ((lane / 4u) & 4u) + ((lane / 2u) % 4u);
+  const uint fn = ((lane / 4u) & 2u) * 2u + (lane % 2u) * 2u;
+  const uint row = n0 + fm;
+  const device uint32_t* gRow = gcodes + (ulong)row * (ulong)((C / BT) * (BT * KG / 32));
+  const device uint32_t* uRow = ucodes + (ulong)row * (ulong)((C / BT) * (BT * KU / 32));
+  // MT activation blocks of 8 rows; grid z walks chunks of 8·MT rows.
+  const uint mBase = threadgroup_position_in_grid.z * 8u * (uint)MT;
+  bool m0[MT], m1[MT];
+  const device T* x0[MT];
+  const device T* x1[MT];
+  #pragma clang loop unroll(full)
+  for (uint bk = 0; bk < (uint)MT; ++bk) {
+    const uint r0 = mBase + 8u * bk + fn;
+    m0[bk] = r0 < (uint)M; m1[bk] = r0 + 1u < (uint)M;
+    x0[bk] = x + (ulong)(m0[bk] ? r0 : 0u) * (ulong)C + fm * 32u;
+    x1[bk] = x + (ulong)(m1[bk] ? r0 + 1u : 0u) * (ulong)C + fm * 32u;
+  }
+  metal::simdgroup_float8x8 accG[MT], accU[MT];
+  #pragma clang loop unroll(full)
+  for (uint bk = 0; bk < (uint)MT; ++bk) {
+    accG[bk] = metal::make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
+    accU[bk] = metal::make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
+  }
+  for (uint base = 0; base < (uint)C; base += 256u) {
+    const uint runA = base / 32u + fn;
+    uint ga[KG + 1], gb[KG + 1], ua[KU + 1], ub[KU + 1];
+    load_run<KG, BT>(gRow, runA, ga);
+    load_run<KG, BT>(gRow, runA + 1u, gb);
+    load_run<KU, BT>(uRow, runA, ua);
+    load_run<KU, BT>(uRow, runA + 1u, ub);
+    #pragma clang loop unroll(full)
+    for (uint s4 = 0; s4 < 8u; ++s4) {
+      float4 va[MT], vb[MT];
+      #pragma clang loop unroll(full)
+      for (uint bk = 0; bk < (uint)MT; ++bk) {
+        va[bk] = m0[bk] ? float4(((const device metal::vec<T, 4>*)(x0[bk] + base))[s4]) : float4(0.0f);
+        vb[bk] = m1[bk] ? float4(((const device metal::vec<T, 4>*)(x1[bk] + base))[s4]) : float4(0.0f);
+      }
+      #pragma clang loop unroll(full)
+      for (uint q = 0; q < 4u; ++q) {
+        const uint j = s4 * 4u + q;
+        metal::simdgroup_float8x8 aG, aU;
+        aG.thread_elements()[0] = trellis_y_f32((run_state<KG, L>(ga, j)));
+        aG.thread_elements()[1] = trellis_y_f32((run_state<KG, L>(gb, j)));
+        aU.thread_elements()[0] = trellis_y_f32((run_state<KU, L>(ua, j)));
+        aU.thread_elements()[1] = trellis_y_f32((run_state<KU, L>(ub, j)));
+        #pragma clang loop unroll(full)
+        for (uint bk = 0; bk < (uint)MT; ++bk) {
+          metal::simdgroup_float8x8 b;
+          b.thread_elements()[0] = va[bk][q];
+          b.thread_elements()[1] = vb[bk][q];
+          metal::simdgroup_multiply_accumulate(accG[bk], aG, b, accG[bk]);
+          metal::simdgroup_multiply_accumulate(accU[bk], aU, b, accU[bk]);
+        }
+      }
+    }
+  }
+  const float gs = float(gscales[row]) * TRELLIS_RCP;
+  const float us = float(uscales[row]) * TRELLIS_RCP;
+  #pragma clang loop unroll(full)
+  for (uint bk = 0; bk < (uint)MT; ++bk)
+    #pragma clang loop unroll(full)
+    for (uint j = 0; j < 2u; ++j) {
+      const uint m = mBase + 8u * bk + fn + j;
+      if (m >= (uint)M) continue;
+      const T gateT = T(accG[bk].thread_elements()[j] * gs);
+      const T upT = T(accU[bk].thread_elements()[j] * us);
+      const T sigmoidT = T(1.0f / (1.0f + metal::precise::exp(-float(gateT))));
+      const T siluT = T(float(gateT) * float(sigmoidT));
+      mid[(ulong)m * (ulong)R + row] = T(float(siluT) * float(upT));
+    }
+`;
+
+let kernel: MetalKernel | undefined;
+
+/** silu(gate(x))·up(x) for 1..8 activation rows. Gate and up share their
+ *  axis-1 geometry (rows a multiple of 8, columns of 256); their bit widths may
+ *  differ. Borrows x and the weights; returns owned lazy output with x's
+ *  leading dimensions. */
+export function gateUpMma(x: MlxArray, gate: TrellisWeights, up: TrellisWeights): MlxArray {
+  const g = gate.geometry, u = up.geometry;
+  const lead = x.shape.slice(0, -1), M = lead.reduce((a, b) => a * b, 1);
+  if (M < 1 || M > 8) throw new Error(`gateUpMma: M=${M} outside 1..8`);
+  if (g.axis !== 1 || u.axis !== 1 || g.rows !== u.rows || g.cols !== u.cols || g.T !== u.T || g.L !== u.L ||
+      g.rows % 8 !== 0 || g.cols % 256 !== 0 || g.T % 32 !== 0)
+    throw new Error("gateUpMma: unsupported gate/up geometry");
+  kernel ??= new MetalKernel({ name: "mlx_bun_trellis_gateup_mma", inputNames: ["x", "gcodes", "gscales", "ucodes", "uscales"],
+    outputNames: ["mid"], source: SOURCE_V3, header: HEADER + RUN3_HEADER, ensureRowContiguous: true });
+  const x2 = ops.reshape(x, [M, g.inFeatures]);
+  const [mid] = kernel.apply([x2, gate.codes, gate.scales, up.codes, up.scales], {
+    outputs: [{ shape: [M, g.rows], dtype: x.dtype }],
+    grid: [32 * SG_TG, Math.ceil(g.rows / 8 / SG_TG), 1], threadGroup: [32 * SG_TG, 1, 1],
+    templateDtypes: { T: x.dtype },
+    templateInts: { M, R: g.rows, C: g.cols, BT: g.T, KG: g.k, KU: u.k, L: g.L, SG_TG, MT },
+  });
+  x2.dispose();
+  const out = ops.reshape(mid!, [...lead, g.rows]);
+  mid!.dispose();
+  return out;
+}

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

@@ -0,0 +1,248 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import type { TrellisGeometry, TrellisWeights } from "./geometry";
+import { HEADER, lutFor, decoderVariant } from "./codebook";
+import { TRELLIS_THREADS, TRELLIS_SG_PER_TG, TRELLIS_CODEBOOK_THREADS, TRELLIS_CODEBOOK_SG_PER_TG } from "./launch";
+import { qualifiedRepresentations } from "./qualified";
+
+import { TRELLIS_MATVEC_MAX_M as MATVEC_MAX_M } from "./reduce";
+
+const GATEUP_SOURCE = String.raw`
+  threadgroup float lutTG[4096];
+  if ((VARIANT) == 2 || CODEBOOK) {
+    for (uint i = thread_position_in_threadgroup.x; i < 4096u; i += (uint)THREADS)
+      lutTG[i] = CODEBOOK ? trellis_unrefined_y(trellis_y(i)) : lut[i];
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+  }
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint row = thread_position_in_grid.y * (uint)ROWS_TG + sg;
+  const uint sample = thread_position_in_grid.z;
+  if (row >= (uint)R || sample >= (uint)M) return;
+  const uint wpb = (uint)(BT * K / 32);
+  const uint runsPerBlock = (uint)BT / 32u;
+  const uint nRuns = (uint)C / 32u;
+  const ulong rowWords = (ulong)((C / BT) * wpb);
+  const device uint32_t* gRow = gcodes + (ulong)row * rowWords;
+  const device uint32_t* uRow = ucodes + (ulong)row * rowWords;
+  const device T* xs = x + (ulong)sample * (ulong)C;
+  const float gScale = float(gscales[row]);
+  const float uScale = float(uscales[row]);
+  float gAcc = 0.0f, gAcc2 = 0.0f;
+  float uAcc = 0.0f, uAcc2 = 0.0f;
+  for (uint run = lane; run < nRuns; run += 32u) {
+    const uint blk = run / runsPerBlock;
+    const uint t0 = (run - blk * runsPerBlock) * 32u;
+    const device uint32_t* gw = gRow + blk * wpb;
+    const device uint32_t* uw = uRow + blk * wpb;
+    const uint w0 = ((uint)BT - 32u - t0) * (uint)K / 32u;
+    uint g[K + 1];
+    uint u[K + 1];
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i <= (uint)K; ++i) {
+      const uint wi = w0 + i;
+      const uint idx = wi >= wpb ? wi - wpb : wi;
+      g[i] = gw[idx];
+      u[i] = uw[idx];
+    }
+    const uint c0 = blk * (uint)BT + t0;
+    const device metal::vec<T, 4>* x4 = (const device metal::vec<T, 4>*)(xs + c0);
+    #pragma clang loop unroll(full)
+    for (uint j4 = 0; j4 < 8u; ++j4) {
+      const float4 xv4 = float4(x4[j4]);
+      #pragma clang loop unroll(full)
+      for (uint q = 0; q < 4u; ++q) {
+        const uint j = j4 * 4u + q;
+        const uint b = (31u - j) * (uint)K;
+        const uint wi = b >> 5;
+        const uint off = b & 31u;
+        uint gwin = g[wi] >> off;
+        uint uwin = u[wi] >> off;
+        if (off + (uint)L > 32u) { gwin |= g[wi + 1] << (32u - off); uwin |= u[wi + 1] << (32u - off); }
+        gwin &= (1u << (uint)L) - 1u;
+        uwin &= (1u << (uint)L) - 1u;
+        const float xv = xv4[q];
+        const float gv = TRELLIS_ROUND((CODEBOOK ? lutTG[gwin] : TRELLIS_DECODE(gwin, lutTG, lut)) * gScale);
+        const float uv = TRELLIS_ROUND((CODEBOOK ? lutTG[uwin] : TRELLIS_DECODE(uwin, lutTG, lut)) * uScale);
+        if (q & 1u) { gAcc2 = metal::fma(gv, xv, gAcc2); uAcc2 = metal::fma(uv, xv, uAcc2); }
+        else { gAcc = metal::fma(gv, xv, gAcc); uAcc = metal::fma(uv, xv, uAcc); }
+      }
+    }
+  }
+  const float gate = metal::simd_sum(gAcc + gAcc2);
+  const float up = metal::simd_sum(uAcc + uAcc2);
+  if (lane == 0u) {
+    const T gateT = T(gate);
+    const T upT = T(up);
+    const T sigmoidT = T(1.0f / (1.0f + metal::precise::exp(-float(gateT))));
+    const T siluT = T(float(gateT) * float(sigmoidT));
+    mid[(ulong)sample * (ulong)R + row] = T(float(siluT) * float(upT));
+  }
+`;
+
+const GATEUP_SHARED_M_SOURCE = String.raw`
+  threadgroup float lutTG[4096];
+  if ((VARIANT) == 2) {
+    for (uint i = thread_position_in_threadgroup.x; i < 4096u; i += ${TRELLIS_THREADS}u) lutTG[i] = lut[i];
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+  }
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint row = thread_position_in_grid.y * (uint)ROWS_TG + sg;
+  const uint sample = 0u;
+  if (row >= (uint)R || sample >= (uint)M) return;
+  const uint wpb = (uint)(BT * K / 32);
+  const uint runsPerBlock = (uint)BT / 32u;
+  const uint nRuns = (uint)C / 32u;
+  const ulong rowWords = (ulong)((C / BT) * wpb);
+  const device uint32_t* gRow = gcodes + (ulong)row * rowWords;
+  const device uint32_t* uRow = ucodes + (ulong)row * rowWords;
+  const float gScale = float(gscales[row]);
+  const float uScale = float(uscales[row]);
+  float gAcc[M], gAcc2[M], uAcc[M], uAcc2[M];
+  #pragma clang loop unroll(full)
+  for (uint m = 0; m < (uint)M; ++m) {
+    gAcc[m] = 0.0f; gAcc2[m] = 0.0f; uAcc[m] = 0.0f; uAcc2[m] = 0.0f;
+  }
+  for (uint run = lane; run < nRuns; run += 32u) {
+    const uint blk = run / runsPerBlock;
+    const uint t0 = (run - blk * runsPerBlock) * 32u;
+    const device uint32_t* gw = gRow + blk * wpb;
+    const device uint32_t* uw = uRow + blk * wpb;
+    const uint w0 = ((uint)BT - 32u - t0) * (uint)K / 32u;
+    uint g[K + 1];
+    uint u[K + 1];
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i <= (uint)K; ++i) {
+      const uint wi = w0 + i;
+      const uint idx = wi >= wpb ? wi - wpb : wi;
+      g[i] = gw[idx];
+      u[i] = uw[idx];
+    }
+    const uint c0 = blk * (uint)BT + t0;
+    #pragma clang loop unroll(full)
+    for (uint j4 = 0; j4 < 8u; ++j4) {
+      float4 xv4[M];
+      #pragma clang loop unroll(full)
+      for (uint m = 0; m < (uint)M; ++m)
+        xv4[m] = float4(((const device metal::vec<T, 4>*)(x + (ulong)m * (ulong)C + c0))[j4]);
+      #pragma clang loop unroll(full)
+      for (uint q = 0; q < 4u; ++q) {
+        const uint j = j4 * 4u + q;
+        const uint b = (31u - j) * (uint)K;
+        const uint wi = b >> 5;
+        const uint off = b & 31u;
+        uint gwin = g[wi] >> off;
+        uint uwin = u[wi] >> off;
+        if (off + (uint)L > 32u) { gwin |= g[wi + 1] << (32u - off); uwin |= u[wi + 1] << (32u - off); }
+        gwin &= (1u << (uint)L) - 1u;
+        uwin &= (1u << (uint)L) - 1u;
+        const float gv = TRELLIS_ROUND(TRELLIS_DECODE(gwin, lutTG, lut) * gScale);
+        const float uv = TRELLIS_ROUND(TRELLIS_DECODE(uwin, lutTG, lut) * uScale);
+        #pragma clang loop unroll(full)
+        for (uint m = 0; m < (uint)M; ++m) {
+          const float xv = xv4[m][q];
+          if (q & 1u) { gAcc2[m] = metal::fma(gv, xv, gAcc2[m]); uAcc2[m] = metal::fma(uv, xv, uAcc2[m]); }
+          else { gAcc[m] = metal::fma(gv, xv, gAcc[m]); uAcc[m] = metal::fma(uv, xv, uAcc[m]); }
+        }
+      }
+    }
+  }
+  #pragma clang loop unroll(full)
+  for (uint m = 0; m < (uint)M; ++m) {
+    const float gate = metal::simd_sum(gAcc[m] + gAcc2[m]);
+    const float up = metal::simd_sum(uAcc[m] + uAcc2[m]);
+    if (lane == 0u) {
+      const T gateT = T(gate);
+      const T upT = T(up);
+      const T sigmoidT = T(1.0f / (1.0f + metal::precise::exp(-float(gateT))));
+      const T siluT = T(float(gateT) * float(sigmoidT));
+      mid[(ulong)m * (ulong)R + row] = T(float(siluT) * float(upT));
+    }
+  }
+`;
+
+let base: MetalKernel | undefined, shared: MetalKernel | undefined;
+function gateUpKernel(): MetalKernel {
+  return base ??= new MetalKernel({ name: "mlx_bun_trellis_gateup_swiglu",
+    inputNames: ["x", "gcodes", "gscales", "ucodes", "uscales", "lut"], outputNames: ["mid"],
+    source: GATEUP_SOURCE, header: HEADER, ensureRowContiguous: true });
+}
+function sharedGateUpKernel(): MetalKernel {
+  return shared ??= new MetalKernel({ name: "mlx_bun_trellis_gateup_shared_m",
+    inputNames: ["x", "gcodes", "gscales", "ucodes", "uscales", "lut"], outputNames: ["mid"],
+    source: GATEUP_SHARED_M_SOURCE, header: HEADER, ensureRowContiguous: true });
+}
+
+/** The measured artifact geometry, row count, dtype and numerical variant of the codebook decode. */
+function codebookShape(g: TrellisGeometry, m: number, dtype: Dtype, selected: number): boolean {
+  return selected === 13 && m === 1 && dtype === Dtype.bfloat16 &&
+    g.axis === 1 && g.T === 256 && g.L === 12 && [2, 3, 4].includes(g.k) &&
+    g.inFeatures === 5120 && g.outFeatures === 17408 && g.cols === 5120 && g.rows === 17408 && !g.blockInterleave;
+}
+
+/** Qualification is deliberately limited to the measured artifact geometry and
+ * GPU family. Other shapes, row counts, dtypes and numerical variants retain
+ * inline decode. The table uses variant 6 arithmetic, not the host's rounded LUT. */
+export function gateUpCodebookEligible(g: TrellisGeometry, m: number, dtype: Dtype, selected: number,
+  architecture: string): boolean {
+  return qualifiedRepresentations(architecture).gateUpCodebook && codebookShape(g, m, dtype, selected);
+}
+
+/** Variants 7..13 share each decoded weight across rows (`GATEUP_SHARED_M_SOURCE`). */
+const sharesRows = (selected: number) => selected >= 7 && selected <= 13;
+
+/** silu(gate(x))·up(x) for exactly one activation row, gate and up at the same
+ * width: the per-row kernel. `codebook` applies at the measured shape; layers
+ * pass the choice they resolved at construction (`qualifiedRepresentations`).
+ * Borrows x and the weights; returns owned lazy output with x's leading dimensions. */
+export function gateUpReduceRow(x: MlxArray, gate: TrellisWeights, up: TrellisWeights, selected: number,
+  codebook: boolean): MlxArray {
+  if (x.size !== x.shape.at(-1)) throw new Error("gateUpReduceRow: expects one activation row");
+  return launch(gateUpKernel(), x, gate, up, selected, codebook, 1, false);
+}
+
+/** silu(gate(x))·up(x) for 2..4 activation rows sharing each decoded weight
+ * (variants 7..13), gate and up at the same width. Borrows x and the weights;
+ * returns owned lazy output with x's leading dimensions. */
+export function gateUpReduceRows(x: MlxArray, gate: TrellisWeights, up: TrellisWeights, selected: number): MlxArray {
+  const M = x.size / x.shape.at(-1)!;
+  if (M < 2 || M > MATVEC_MAX_M) throw new Error(`gateUpReduceRows: M=${M} outside 2..${MATVEC_MAX_M}`);
+  if (!sharesRows(selected)) throw new Error(`gateUpReduceRows: variant ${selected} does not share rows (7..13 do)`);
+  return launch(sharedGateUpKernel(), x, gate, up, selected, false, M, true);
+}
+
+/** Borrow matching axis-1 gate/up weights with the same bit width; M must be 1..4.
+ * Return owned lazy SwiGLU output, retaining the input's leading dimensions.
+ * `codebook` applies at the measured shape: layers pass the choice they resolved
+ * at construction (`qualifiedRepresentations`), benches and tests the one they compare. */
+export function fusedGateUpSwiglu(x: MlxArray, gate: TrellisWeights, up: TrellisWeights, selected: number,
+  codebook: boolean): MlxArray {
+  const M = x.shape.slice(0, -1).reduce((a, b) => a * b, 1);
+  if (M > MATVEC_MAX_M) throw new Error(`fusedGateUpSwiglu: M=${M} > ${MATVEC_MAX_M}`);
+  const shared = sharesRows(selected) && M > 1;
+  return launch(shared ? sharedGateUpKernel() : gateUpKernel(), x, gate, up, selected, codebook, M, shared);
+}
+
+function launch(kernel: MetalKernel, x: MlxArray, gate: TrellisWeights, up: TrellisWeights, selected: number,
+  codebook: boolean, M: number, shared: boolean): MlxArray {
+  const g = gate.geometry;
+  const lead = x.shape.slice(0, -1);
+  const useCodebook = codebook && codebookShape(g, M, x.dtype, selected);
+  const threads = useCodebook ? TRELLIS_CODEBOOK_THREADS : TRELLIS_THREADS;
+  const rowsPerGroup = useCodebook ? TRELLIS_CODEBOOK_SG_PER_TG : TRELLIS_SG_PER_TG;
+  const x2 = ops.reshape(x, [M, g.inFeatures]);
+  const [mid] = kernel.apply([x2, gate.codes, gate.scales, up.codes, up.scales, lutFor(g.L)], {
+    outputs: [{ shape: [M, g.rows], dtype: x.dtype }],
+    grid: [threads, Math.ceil(g.rows / rowsPerGroup), shared ? 1 : M],
+    threadGroup: [threads, 1, 1],
+    templateDtypes: { T: x.dtype },
+    templateInts: { M, R: g.rows, C: g.cols, BT: g.T, K: g.k, L: g.L, ROWS_TG: rowsPerGroup, VARIANT: decoderVariant(selected), THREADS: threads, CODEBOOK: Number(useCodebook) },
+  });
+  x2.dispose();
+  const out = ops.reshape(mid!, [...lead, g.rows]);
+  mid!.dispose();
+  return out;
+}

package/src/kernels/trellis/geometry.ts

@@ -0,0 +1,22 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+
+export interface TrellisGeometry {
+  k: number;
+  L: number;
+  T: number;
+  axis: 0 | 1;
+  /** Stored-matrix rows / coded columns. */
+  rows: number;
+  cols: number;
+  inFeatures: number;
+  outFeatures: number;
+  /** Two coded blocks interleaved across rows; absent for row-major codes. */
+  blockInterleave?: 2;
+}
+
+/** Borrowed packed weights and their caller-supplied geometry. */
+export interface TrellisWeights {
+  codes: MlxArray;
+  scales: MlxArray;
+  geometry: TrellisGeometry;
+}

package/src/kernels/trellis/index.ts

@@ -0,0 +1,103 @@
+// Packed Trellis projection kernels (QTIP 1MAD code) and the decode-variant
+// semantics their callers pass (layers/trellis-linear.ts passes 13).
+//
+// Format: `.weight` = uint32 bitstream, `.scales` = fp16 [rows], config entry
+// `{mode:"trellis", bits:k, group_size:T, trellis:{L, code:"1mad", axis}}`.
+// axis=1 (gate/up) is coded along the INPUT dim, stored [out, in*k/32]; axis=0
+// (down) is coded along the OUTPUT dim, stored [in, out*k/32] (the stored matrix
+// is W^T, coded along its last axis). The k3/axis0 layout [cols/512, rows, 48]
+// interleaves two coded blocks across rows (`blockInterleave`); the tensor
+// shape identifies it and no second code copy is kept.
+//
+// One decode primitive: state_t is the L-bit window at bit offset (T-1-t)*k of
+// the block (wrapping), so any weight decodes in O(1) with the 1MAD code
+// computed inline (codebook.ts HEADER), no LUT and no threadgroup memory.
+//   - reduce   (M <= 4, axis=1): one SIMD group per output row; each lane owns
+//               runs of 32 consecutive positions held in K+1 registers (every
+//               window shift a compile-time constant); simd_sum.
+//   - scatter  (M <= 4, axis=0): lane-per-WORD of a block, running sums per
+//               position in the word, split-K over input rows (partials folded
+//               by one mlx sum). One kernel each: scatterRow / balancedScatterRow
+//               (one row; balanced = 3-bit codes at T 256) and scatterRows /
+//               balancedScatterRows (2..4 rows sharing each decoded weight);
+//               trellisScatter chooses among them by row count and geometry.
+//   - expand   (M > 4 fallback): decode the whole tensor to bf16, then a stock
+//               matmul; the transient is one tensor.
+//   - gate-up / mixed-gate-up: fused gate+up+SwiGLU for M <= 4 rows
+//               (gateUpReduceRow one row, gateUpReduceRows 2..4 shared rows;
+//               fusedGateUpSwiglu chooses between them by row count). The mixed
+//               kernel serves layers whose gate and up bit widths differ (one
+//               word layout and bit offset each); TAIL "split" reproduces MLX's
+//               compiled swiglu over two bf16 projections bit for bit (no output
+//               changes vs the two-kernel path), "fused" the float32 sigmoid of
+//               the same-width kernel.
+//
+// Decode VARIANT, an explicit argument of each kernel. TrellisLinear passes
+// 13; benches and tests pass the others to compare them. 0-6 select the decode
+// arithmetic (see HEADER in codebook.ts): 0-3 reconstruct bf16(lut[state]*scale);
+// 6 is f32 code*scale in the packed kernels (expansion still stores bf16).
+// Variants 7-13 keep variant 6's decoded values (`decoderVariant`) and add
+// kernel selection, each including the ones below it:
+//    7  weight decoding shared across M=2..4 in reduce and gate-up
+//    8  balanced 3-bit scatter (k=3, T=256, L<=12)
+//    9  same kernels as 8; it once also deferred prefill projection evaluation
+//       to the Qwen layer barrier, which no longer exists
+//   10  shared-M scatter across M=2..4
+//   11  tiled axis-1 prefill at M=5..32 on the measured MLP shape; wide prefill
+//       when the input is proven row-contiguous
+//   12  tiled axis-0 prefill at M=5..8 with MLX's split-K reduction
+//   13  shared packed-word loads in the remaining bf16 expansion (vector
+//       expand), and the shared scatter codebook for M=3..4 bf16 on the
+//       interleaved k3 layout
+//
+// Threadgroup-codebook one-row kernels (every code value of the 2^L states
+// generated once per threadgroup), the variant-13 output of the kernels above
+// bit for bit:
+//   - gate-up-codebook: gate/up/SwiGLU, the same or different widths
+//   - down-codebook: lanes owning word-aligned runs of positions (every window
+//     shift a constant), row-major 2/4-bit and 3-bit block-interleaved codes
+//
+// Factored-scale kernels: weights contribute their code value y and the scale
+// (with 1/147.8) applies once per output (gate/up) or per input row (down);
+// one kernel per row form and code layout:
+//   - gate-up-factored: same-width gate/up, one row / 2..4 shared rows
+//   - mixed-gate-up-factored: different-width gate/up, 1..4 rows
+//   - down-factored: row-major axis-0 codes, one row / 2..4 rows
+//   - down-k3-interleaved-factored: 3-bit block-interleaved axis-0 codes
+//   - gate-up-mma / down-mma: 1..8 rows on the simdgroup matrix unit
+//   - gate-up-mma-wide / down-mma-wide: the same fragments for any number of
+//     rows, each decoded fragment feeding the MMAs of several 8-row blocks,
+//     each code value read from a threadgroup codebook; down-mma-wide in
+//     passes the rows plan or the caller fixes (one block for 1..8 rows)
+//   - gate-up-codebook-mma: gate-up-mma's fragments and bits for 1..8 rows,
+//     each code value read from a threadgroup codebook (the M1 Max form)
+//
+// ExpandedTrellisLinear (layers/trellis-linear.ts) decodes a trellis tensor at
+// construction into 8-bit g64 affine (the eval-carrier numerics, about -45 dB)
+// and serves it through the stock QuantizedLinear, for a machine where the
+// kernels lose.
+export type { TrellisGeometry, TrellisWeights } from "./geometry";
+export { vectorTrellisExpand, vectorTrellisExpandEligible } from "./vector-expand";
+export { expandTrellis } from "./expand";
+export { trellisReduce, TRELLIS_MATVEC_MAX_M } from "./reduce";
+export { trellisScatter, scatterRow, scatterRows, balancedScatterRow, balancedScatterRows } from "./scatter";
+export { qualifiedRepresentations, type ScatterRepresentation, type TrellisRepresentations } from "./qualified";
+export { fusedGateUpSwiglu, gateUpReduceRow, gateUpReduceRows } from "./gate-up";
+export { gateUpCodebookRow } from "./gate-up-codebook";
+export { downCodebookRow, downK3InterleavedCodebookRow } from "./down-codebook";
+export { fusedGateUpSwigluMixed, type MixedGateUpTail } from "./mixed-gate-up";
+export { tiledTrellisPrefill, tiledTrellisPrefillEligible } from "./tiled-prefill";
+export { splitKTrellisPrefill, splitKTrellisPrefillEligible } from "./splitk-prefill";
+export { wideTrellisPrefill, wideTrellisPrefillEligible, nativeTrellisWidePrefill } from "./wide-prefill";
+export { lut1mad, wordsPerBlock } from "./codebook";
+export { gateUpFactoredRow, gateUpFactoredRows } from "./gate-up-factored";
+export { mixedGateUpFactoredRows } from "./mixed-gate-up-factored";
+export { downFactoredRow, downFactoredRows } from "./down-factored";
+export { downK3InterleavedFactoredRow, downK3InterleavedFactoredRows } from "./down-k3-interleaved-factored";
+export { gateUpMma } from "./gate-up-mma";
+export { downMma, downK3InterleavedMma } from "./down-mma";
+export { gateUpMmaWide } from "./gate-up-mma-wide";
+export {
+  downMmaWide, downK3InterleavedMmaWide, downMmaWidePass, downK3InterleavedMmaWidePass, DOWN_ONE_BLOCK, type DownPass,
+} from "./down-mma-wide";
+export { gateUpCodebookMma } from "./gate-up-codebook-mma";

package/src/kernels/trellis/launch.ts

@@ -0,0 +1,12 @@
+/** Default Trellis matvec launch: four SIMD groups of 32 lanes per threadgroup,
+ * one output row per SIMD group. Qualified codebook decode uses wider groups. */
+export const TRELLIS_THREADS = 128;
+export const TRELLIS_SG_PER_TG = 4;
+/** Amortize the 16 KiB codebook over sixteen rows in the qualified M1 decode. */
+export const TRELLIS_CODEBOOK_THREADS = 512;
+export const TRELLIS_CODEBOOK_SG_PER_TG = 16;
+/** The one-row gate/up codebook kernel: one threadgroup per M1 Max core, each table serving 32 rows at a time. */
+export const TRELLIS_CODEBOOK_ROW_THREADS = 1024;
+export const TRELLIS_CODEBOOK_ROW_SG_PER_TG = 32;
+/** The one-row down codebook kernel: eight SIMD groups per threadgroup (each split's lanes own 16 positions). */
+export const TRELLIS_CODEBOOK_DOWN_THREADS = 256;

package/src/kernels/trellis/mixed-gate-up-factored.ts

@@ -0,0 +1,119 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+import type { TrellisWeights } from "./geometry";
+import { HEADER } from "./codebook";
+import { TRELLIS_THREADS, TRELLIS_SG_PER_TG } from "./launch";
+
+// Fused gate/up/SwiGLU for layers whose gate and up were coded at different
+// bit widths (KG, KU), 1..4 activation rows sharing each decoded weight. Row
+// scales are factored out of the reduction as in gate-up-factored.ts. The
+// SwiGLU tail is MLX's compiled swiglu over bf16 gate and up (sigmoid via
+// exp(|g|) in T), the arithmetic of two separate projections plus swiglu.
+const SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint row = thread_position_in_grid.y * (uint)ROWS_TG + sg;
+  if (row >= (uint)R) return;
+  const uint gWpb = (uint)(BT * KG / 32);
+  const uint uWpb = (uint)(BT * KU / 32);
+  const uint runsPerBlock = (uint)BT / 32u;
+  const uint nRuns = (uint)C / 32u;
+  const device uint32_t* gRow = gcodes + (ulong)row * (ulong)((C / BT) * gWpb);
+  const device uint32_t* uRow = ucodes + (ulong)row * (ulong)((C / BT) * uWpb);
+  const float gScale = float(gscales[row]);
+  const float uScale = float(uscales[row]);
+  float gAcc[M], uAcc[M];
+  #pragma clang loop unroll(full)
+  for (uint m = 0; m < (uint)M; ++m) { gAcc[m] = 0.0f; uAcc[m] = 0.0f; }
+  for (uint run = lane; run < nRuns; run += 32u) {
+    const uint blk = run / runsPerBlock;
+    const uint t0 = (run - blk * runsPerBlock) * 32u;
+    const device uint32_t* gw = gRow + blk * gWpb;
+    const device uint32_t* uw = uRow + blk * uWpb;
+    const uint gW0 = ((uint)BT - 32u - t0) * (uint)KG / 32u;
+    const uint uW0 = ((uint)BT - 32u - t0) * (uint)KU / 32u;
+    uint g[KG + 1];
+    uint u[KU + 1];
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i <= (uint)KG; ++i) {
+      const uint wi = gW0 + i;
+      g[i] = gw[wi >= gWpb ? wi - gWpb : wi];
+    }
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i <= (uint)KU; ++i) {
+      const uint wi = uW0 + i;
+      u[i] = uw[wi >= uWpb ? wi - uWpb : wi];
+    }
+    const uint c0 = blk * (uint)BT + t0;
+    #pragma clang loop unroll(full)
+    for (uint j4 = 0; j4 < 8u; ++j4) {
+      float4 xv4[M];
+      #pragma clang loop unroll(full)
+      for (uint m = 0; m < (uint)M; ++m)
+        xv4[m] = float4(((const device metal::vec<T, 4>*)(x + (ulong)m * (ulong)C + c0))[j4]);
+      #pragma clang loop unroll(full)
+      for (uint q = 0; q < 4u; ++q) {
+        const uint j = j4 * 4u + q;
+        const uint gb = (31u - j) * (uint)KG;
+        const uint ub = (31u - j) * (uint)KU;
+        const uint gwi = gb >> 5, goff = gb & 31u;
+        const uint uwi = ub >> 5, uoff = ub & 31u;
+        uint gwin = g[gwi] >> goff;
+        uint uwin = u[uwi] >> uoff;
+        if (goff + (uint)L > 32u) gwin |= g[gwi + 1] << (32u - goff);
+        if (uoff + (uint)L > 32u) uwin |= u[uwi + 1] << (32u - uoff);
+        gwin &= (1u << (uint)L) - 1u;
+        uwin &= (1u << (uint)L) - 1u;
+        const float gv = trellis_y_f32(gwin);
+        const float uv = trellis_y_f32(uwin);
+        #pragma clang loop unroll(full)
+        for (uint m = 0; m < (uint)M; ++m) {
+          const float xv = xv4[m][q];
+          gAcc[m] = metal::fma(gv, xv, gAcc[m]); uAcc[m] = metal::fma(uv, xv, uAcc[m]);
+        }
+      }
+    }
+  }
+  #pragma clang loop unroll(full)
+  for (uint m = 0; m < (uint)M; ++m) {
+    const float gate = metal::simd_sum(gAcc[m]) * (gScale * TRELLIS_RCP);
+    const float up = metal::simd_sum(uAcc[m]) * (uScale * TRELLIS_RCP);
+    if (lane == 0u) {
+      const T gateT = T(gate);
+      const T upT = T(up);
+      const T sigY = 1 / (1 + metal::exp(metal::abs(gateT)));
+      const T sigmoidT = (gateT < 0) ? sigY : 1 - sigY;
+      const T siluT = gateT * sigmoidT;
+      mid[(ulong)m * (ulong)R + row] = siluT * upT;
+    }
+  }
+`;
+
+let kernel: MetalKernel | undefined;
+
+/** silu(gate(x))·up(x) for 1..4 activation rows when gate and up share their
+ *  axis-1 geometry but not their bit width. Borrows x and the weights; returns
+ *  owned lazy output with x's leading dimensions. */
+export function mixedGateUpFactoredRows(x: MlxArray, gate: TrellisWeights, up: TrellisWeights): MlxArray {
+  const g = gate.geometry, u = up.geometry;
+  if (g.axis !== 1 || u.axis !== 1 || g.k === u.k || g.rows !== u.rows || g.cols !== u.cols || g.T !== u.T || g.L !== u.L)
+    throw new Error("mixedGateUpFactoredRows: gate and up need the same axis-1 geometry at different widths");
+  const lead = x.shape.slice(0, -1);
+  const M = lead.reduce((a, b) => a * b, 1);
+  if (M < 1 || M > 4) throw new Error(`mixedGateUpFactoredRows: M=${M} outside 1..4`);
+  kernel ??= new MetalKernel({ name: "mlx_bun_trellis_gateup_mixed_factored", inputNames: ["x", "gcodes", "gscales", "ucodes", "uscales"],
+    outputNames: ["mid"], source: SOURCE, header: HEADER, ensureRowContiguous: true });
+  const x2 = ops.reshape(x, [M, g.inFeatures]);
+  const [mid] = kernel.apply([x2, gate.codes, gate.scales, up.codes, up.scales], {
+    outputs: [{ shape: [M, g.rows], dtype: x.dtype }],
+    grid: [TRELLIS_THREADS, Math.ceil(g.rows / TRELLIS_SG_PER_TG), 1],
+    threadGroup: [TRELLIS_THREADS, 1, 1],
+    templateDtypes: { T: x.dtype },
+    templateInts: { M, R: g.rows, C: g.cols, BT: g.T, KG: g.k, KU: u.k, L: g.L, ROWS_TG: TRELLIS_SG_PER_TG },
+  });
+  x2.dispose();
+  const out = ops.reshape(mid!, [...lead, g.rows]);
+  mid!.dispose();
+  return out;
+}