@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/kernels/trellis/expand.ts

@@ -0,0 +1,42 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+import type { TrellisGeometry } from "./geometry";
+import { HEADER, lutFor, decoderVariant } from "./codebook";
+
+import { vectorTrellisExpand, vectorTrellisExpandEligible } from "./vector-expand";
+
+const EXPAND_SOURCE = String.raw`
+  const uint c = thread_position_in_grid.x;
+  const uint r = thread_position_in_grid.y;
+  if (c >= (uint)C || r >= (uint)R) return;
+  const uint wpb = (uint)(BT * K / 32);
+  const uint blk = c / (uint)BT;
+  const uint t = c - blk * (uint)BT;
+  const device uint32_t* block = codes + (INTERLEAVE
+    ? (ulong)(blk / 2) * R * 2 * wpb + (ulong)r * 2 * wpb + (blk % 2) * wpb
+    : (ulong)r * (C / BT) * wpb + blk * wpb);
+  const uint s = trellis_state(block, wpb, t, (uint)BT, (uint)K, (uint)L);
+  out[(ulong)r * (uint)C + c] = T(((VARIANT) >= 2 ? lut[s] : (VARIANT) == 1 ? trellis_val_rcp(s) : trellis_val(s)) * float(scales[r]));
+`;
+
+let kernel: MetalKernel | undefined;
+function expandKernel(): MetalKernel {
+  return kernel ??= new MetalKernel({ name: "mlx_bun_trellis_expand",
+    inputNames: ["codes", "scales", "lut"], outputNames: ["out"],
+    source: EXPAND_SOURCE, header: HEADER, ensureRowContiguous: true });
+}
+
+/** Decode the stored matrix ([rows, cols], coded along cols) to `dtype`. */
+export function expandTrellis(codes: MlxArray, scales: MlxArray, g: TrellisGeometry, dtype: Dtype, selected: number): MlxArray {
+  if (selected === 13 && vectorTrellisExpandEligible(g, dtype)) return vectorTrellisExpand(codes, scales, g);
+  const [out] = expandKernel().apply([codes, scales, lutFor(g.L)], {
+    outputs: [{ shape: [g.rows, g.cols], dtype }],
+    grid: [g.cols, g.rows, 1],
+    threadGroup: [Math.min(256, g.cols), 1, 1],
+    templateDtypes: { T: dtype },
+    templateInts: { R: g.rows, C: g.cols, BT: g.T, K: g.k, L: g.L,
+      VARIANT: decoderVariant(selected), INTERLEAVE: g.blockInterleave ?? 0 },
+  });
+  return out!;
+}

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

@@ -0,0 +1,141 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+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";
+
+// Fused gate/up/SwiGLU for 1..8 activation rows on the simdgroup matrix unit,
+// decoding through a threadgroup codebook (the M1 Max form, where the matrix
+// products issue on the ordinary ALU pipes and the integer decode competes with
+// them). gate-up-mma's fragments, K permutation and accumulation order, so each
+// output is bit for bit gateUpMma's: a simdgroup owns an 8-output-row tile, A
+// column i at step j is position base + 32·i + j, and every lane decodes two
+// contiguous 32-position runs per projection from K+1 words. Each code value y
+// is read from a half table of the 2^L values (every y in [-510, 510] is exact
+// in half) that the threadgroup generates once, in place of the 1MAD
+// arithmetic; the A fragments are half, B and the accumulators float. A
+// window is one bit-field extract: one that straddles a word reads a virtual
+// word (the high half of one word and the low half of the next). Each lane
+// reads its two activation rows from device memory, 16 bytes at a time, with
+// no threadgroup staging or barrier inside the K loop; rows past M repeat the
+// last row and are not written. Measured on the M1 Max at the Qwen MLP shape
+// (17408 x 5120, gate and up together, real weights), 3-bit: 8 rows 0.476 ms
+// against gateUpMma's 1.17, 4 rows 0.465 against 0.795, 2 rows 0.458 against
+// gateUpFactoredRows' 0.664.
+const SG_TG = 4;
+
+const RUN_HEADER = String.raw`
+// Words W0 .. W0 + K of run 'run' of a row's codes (the last wraps to the block's first word).
+template <int K, int BT_>
+inline void codebook_mma_run(const device uint32_t* rowCodes, uint run, thread uint* w, thread uint* v) {
+  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) w[i] = bw[w0 + i];
+  w[K] = bw[t0 == 0u ? 0u : w0 + (uint)K];
+  // Virtual word i: bits [32i - 16, 32i + 16) of the run.
+  #pragma clang loop unroll(full)
+  for (uint i = 1; i <= (uint)K; ++i) v[i] = (w[i - 1] >> 16) | (w[i] << 16);
+}
+// The state of position j of a run (bits [(31 - j)·K, +L) of its words), one extract.
+template <int K, int L_>
+inline uint codebook_mma_state(thread const uint* w, thread const uint* v, uint j) {
+  const uint b = (31u - j) * (uint)K;
+  const uint off = b & 31u;
+  if (off + (uint)L_ <= 32u) return (w[b >> 5] >> off) & ((1u << (uint)L_) - 1u);
+  const uint vi = (b + 16u) >> 5;
+  return (v[vi] >> (b + 16u - 32u * vi)) & ((1u << (uint)L_) - 1u);
+}
+`;
+
+const SOURCE = String.raw`
+  const uint M = (uint)x_shape[0];
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  threadgroup half codebook[1 << L];
+  for (uint i = sg * 32u + lane; i < (1u << (uint)L); i += 32u * (uint)SG_TG) codebook[i] = half(trellis_y_f32(i));
+  threadgroup_barrier(metal::mem_flags::mem_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));
+  // Activation rows fn and fn + 1 from position 32·fm, 8 values per 16-byte load.
+  const device uint4* x0 = (const device uint4*)(x + (ulong)metal::min(fn, M - 1u) * (ulong)C + fm * 32u);
+  const device uint4* x1 = (const device uint4*)(x + (ulong)metal::min(fn + 1u, M - 1u) * (ulong)C + fm * 32u);
+  metal::simdgroup_float8x8 accG = metal::make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
+  metal::simdgroup_float8x8 accU = 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], gav[KG + 1], gb[KG + 1], gbv[KG + 1];
+    uint ua[KU + 1], uav[KU + 1], ub[KU + 1], ubv[KU + 1];
+    codebook_mma_run<KG, BT>(gRow, runA, ga, gav);
+    codebook_mma_run<KG, BT>(gRow, runA + 1u, gb, gbv);
+    codebook_mma_run<KU, BT>(uRow, runA, ua, uav);
+    codebook_mma_run<KU, BT>(uRow, runA + 1u, ub, ubv);
+    #pragma clang loop unroll(full)
+    for (uint s8 = 0; s8 < 4u; ++s8) {
+      const uint4 xa = x0[base / 8u + s8], xb = x1[base / 8u + s8];
+      #pragma clang loop unroll(full)
+      for (uint q = 0; q < 8u; ++q) {
+        const uint j = s8 * 8u + q;
+        metal::simdgroup_half8x8 aG, aU;
+        aG.thread_elements()[0] = codebook[codebook_mma_state<KG, L>(ga, gav, j)];
+        aG.thread_elements()[1] = codebook[codebook_mma_state<KG, L>(gb, gbv, j)];
+        aU.thread_elements()[0] = codebook[codebook_mma_state<KU, L>(ua, uav, j)];
+        aU.thread_elements()[1] = codebook[codebook_mma_state<KU, L>(ub, ubv, j)];
+        metal::simdgroup_float8x8 b;
+        b.thread_elements()[0] = float(as_type<T>(ushort((q & 1u) ? xa[q / 2u] >> 16 : xa[q / 2u] & 0xFFFFu)));
+        b.thread_elements()[1] = float(as_type<T>(ushort((q & 1u) ? xb[q / 2u] >> 16 : xb[q / 2u] & 0xFFFFu)));
+        metal::simdgroup_multiply_accumulate(accG, aG, b, accG);
+        metal::simdgroup_multiply_accumulate(accU, aU, b, accU);
+      }
+    }
+  }
+  const float gs = float(gscales[row]) * TRELLIS_RCP;
+  const float us = float(uscales[row]) * TRELLIS_RCP;
+  #pragma clang loop unroll(full)
+  for (uint j = 0; j < 2u; ++j) {
+    const uint m = fn + j;
+    if (m >= M) continue;
+    const T gateT = T(accG.thread_elements()[j] * gs);
+    const T upT = T(accU.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 bfloat16 or float16 activation rows, each output gateUpMma's bits. Gate and up
+ *  share their axis-1 geometry (rows a multiple of 8, columns of 256, L at most 12); their bit widths may differ.
+ *  Borrows x and the weights; returns owned lazy output with x's leading dimensions. */
+export function gateUpCodebookMma(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(`gateUpCodebookMma: M=${M} outside 1..8`);
+  if (x.dtype !== Dtype.bfloat16 && x.dtype !== Dtype.float16)
+    throw new Error("gateUpCodebookMma: needs bfloat16 or float16 activations");
+  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 || g.L > 12)
+    throw new Error("gateUpCodebookMma: unsupported gate/up geometry");
+  kernel ??= new MetalKernel({ name: "mlx_bun_trellis_gateup_codebook_mma", inputNames: ["x", "gcodes", "gscales", "ucodes", "uscales"],
+    outputNames: ["mid"], source: SOURCE, header: HEADER + RUN_HEADER, ensureRowContiguous: true });
+  using 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: { R: g.rows, C: g.cols, BT: g.T, KG: g.k, KU: u.k, L: g.L, SG_TG },
+  });
+  using y = mid!;
+  return ops.reshape(y, [...lead, g.rows]);
+}

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

@@ -0,0 +1,137 @@
+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 { CODEBOOK_WINDOW, HEADER } from "./codebook";
+import { TRELLIS_CODEBOOK_ROW_THREADS, TRELLIS_CODEBOOK_ROW_SG_PER_TG } from "./launch";
+import type { MixedGateUpTail } from "./mixed-gate-up";
+
+// Fused gate/up/SwiGLU for one activation row over axis-1 Trellis codes whose
+// gate and up widths may differ (KG, KU), decoding through a threadgroup
+// codebook: every code value y·(1/147.8) of the 2^L states, generated once per
+// threadgroup and shared by the rows its 32 SIMD groups take in turn (at most
+// THREADGROUPS threadgroups, so each table serves several rows). Same
+// lanes, runs, accumulation order and f32 code·scale weights as
+// fusedGateUpSwiglu (same widths) and fusedGateUpSwigluMixed (different
+// widths), bit for bit; TAIL selects the mixed kernel's SwiGLU arithmetic (1:
+// MLX's compiled swiglu over the bf16 projections; 0: the float32 sigmoid of
+// the same-width kernel).
+//
+// The threadgroup also stages the activation row, in 16-byte chunks ordered
+// so the 32 lanes of one load read consecutive chunks; a lane reads each
+// chunk of 8 positions just before it uses them (a threadgroup of 1024
+// threads fits an M1 Max core only below 53 registers per thread). Each window
+// is one bit-field extract (CODEBOOK_WINDOW, codebook.ts); code words are
+// addressed by 32-bit offsets.
+const SOURCE = String.raw`
+  threadgroup float codebook[1 << L];
+  threadgroup uint4 xs[C / 8];
+  constexpr uint N_RUNS = (uint)C / 32u;
+  for (uint i = thread_position_in_threadgroup.x; i < (1u << (uint)L); i += (uint)THREADS)
+    codebook[i] = trellis_unrefined_bits_y(trellis_y(i));
+  // Chunk q (positions 8q..8q+7) of run r at q·N_RUNS + r.
+  for (uint c = thread_position_in_threadgroup.x; c < (uint)C / 8u; c += (uint)THREADS)
+    xs[c] = ((const device uint4*)x)[(c % N_RUNS) * 4u + c / N_RUNS];
+  threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+  const threadgroup uchar* table = (const threadgroup uchar*)codebook;
+  const uint lane = thread_index_in_simdgroup;
+  const uint gWpb = (uint)(BT * KG / 32);
+  const uint uWpb = (uint)(BT * KU / 32);
+  const uint runsPerBlock = (uint)BT / 32u;
+  const uint rowStep = threadgroups_per_grid.y * (uint)ROWS_TG;
+  for (uint row = thread_position_in_grid.y * (uint)ROWS_TG + simdgroup_index_in_threadgroup; row < (uint)R; row += rowStep) {
+    const uint gRow = row * (uint)((C / BT) * gWpb);
+    const uint uRow = row * (uint)((C / BT) * uWpb);
+    const float gScale = float(gscales[row]);
+    const float uScale = float(uscales[row]);
+    float gAcc = 0.0f, gAcc2 = 0.0f, uAcc = 0.0f, uAcc2 = 0.0f;
+    for (uint run = lane; run < N_RUNS; run += 32u) {
+      const uint blk = run / runsPerBlock;
+      const uint t0 = (run - blk * runsPerBlock) * 32u;
+      const uint gw = gRow + blk * gWpb, 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], gv[KG + 1];
+      uint u[KU + 1], uv[KU + 1];
+      // Words W0 .. W0 + K - 1 stay in the block; the next word wraps to the block's first for its last run.
+      #pragma clang loop unroll(full)
+      for (uint i = 0; i < (uint)KG; ++i) g[i] = gcodes[gw + gW0 + i];
+      g[KG] = gcodes[gw + (t0 == 0u ? 0u : gW0 + (uint)KG)];
+      #pragma clang loop unroll(full)
+      for (uint i = 0; i < (uint)KU; ++i) u[i] = ucodes[uw + uW0 + i];
+      u[KU] = ucodes[uw + (t0 == 0u ? 0u : uW0 + (uint)KU)];
+      // Virtual words: bits [32i - 16, 32i + 16) of the run's words (only the ignored bits below word 0 are absent).
+      gv[0] = g[0] << 16;
+      uv[0] = u[0] << 16;
+      #pragma clang loop unroll(full)
+      for (uint i = 1; i <= (uint)KG; ++i) gv[i] = (g[i - 1] >> 16) | (g[i] << 16);
+      #pragma clang loop unroll(full)
+      for (uint i = 1; i <= (uint)KU; ++i) uv[i] = (u[i - 1] >> 16) | (u[i] << 16);
+      #pragma clang loop unroll(full)
+      for (uint q = 0; q < 4u; ++q) {
+        const uint4 chunk = xs[q * N_RUNS + run];
+        #pragma clang loop unroll(full)
+        for (uint p = 0; p < 8u; ++p) {
+          const uint j = q * 8u + p;
+          // The bf16 in either half of a word, as the float of its bits.
+          const float xv = as_type<float>((p & 1u) ? chunk[p >> 1] & 0xFFFF0000u : chunk[p >> 1] << 16);
+          const float gWeight = trellis_codebook_window(table, g, gv, (31u - j) * (uint)KG, (uint)L) * gScale;
+          const float uWeight = trellis_codebook_window(table, u, uv, (31u - j) * (uint)KU, (uint)L) * uScale;
+          if (j & 1u) { gAcc2 = metal::fma(gWeight, xv, gAcc2); uAcc2 = metal::fma(uWeight, xv, uAcc2); }
+          else { gAcc = metal::fma(gWeight, xv, gAcc); uAcc = metal::fma(uWeight, 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);
+      if ((TAIL) == 1) {
+        const T sigY = 1 / (1 + metal::exp(metal::abs(gateT)));
+        const T sigmoidT = (gateT < 0) ? sigY : 1 - sigY;
+        const T siluT = gateT * sigmoidT;
+        mid[row] = siluT * upT;
+      } else {
+        const T sigmoidT = T(1.0f / (1.0f + metal::precise::exp(-float(gateT))));
+        const T siluT = T(float(gateT) * float(sigmoidT));
+        mid[row] = T(float(siluT) * float(upT));
+      }
+    }
+  }
+`;
+
+/** Threadgroups at most: five per M1 Max GPU core, measured against 128 to 256 (the codebook generated 544 times per
+ *  projection at the artifact's 17408 rows otherwise). */
+const THREADGROUPS = 160;
+
+let kernel: MetalKernel | undefined;
+
+/** silu(gate(x))·up(x) for exactly one activation row over matching axis-1
+ *  geometry (L at most 12, at most 8192 bf16 inputs; the bit widths may differ). `tail` applies where the
+ *  widths differ; same widths take the float32 sigmoid. Borrows x and the
+ *  weights; returns owned lazy output with x's leading dimensions. */
+export function gateUpCodebookRow(x: MlxArray, gate: TrellisWeights, up: TrellisWeights, tail: MixedGateUpTail): MlxArray {
+  const g = gate.geometry, u = up.geometry;
+  if (g.axis !== 1 || u.axis !== 1 || g.rows !== u.rows || g.cols !== u.cols || g.T !== u.T || g.L !== u.L || g.L > 12)
+    throw new Error("gateUpCodebookRow: gate and up need matching axis-1 geometry with L <= 12");
+  if (x.size !== x.shape.at(-1)) throw new Error("gateUpCodebookRow: expects one activation row");
+  if ((4 << g.L) + 2 * g.cols > 32768) throw new Error("gateUpCodebookRow: the codebook and staged row exceed 32 KiB of threadgroup memory");
+  kernel ??= new MetalKernel({ name: "mlx_bun_trellis_gateup_codebook_row",
+    inputNames: ["x", "gcodes", "gscales", "ucodes", "uscales"], outputNames: ["mid"],
+    source: SOURCE, header: HEADER + CODEBOOK_WINDOW, ensureRowContiguous: true });
+  const lead = x.shape.slice(0, -1);
+  const x2 = ops.reshape(x, [1, g.inFeatures]);
+  const [mid] = kernel.apply([x2, gate.codes, gate.scales, up.codes, up.scales], {
+    outputs: [{ shape: [1, g.rows], dtype: x.dtype }],
+    grid: [TRELLIS_CODEBOOK_ROW_THREADS, Math.min(THREADGROUPS, Math.ceil(g.rows / TRELLIS_CODEBOOK_ROW_SG_PER_TG)), 1],
+    threadGroup: [TRELLIS_CODEBOOK_ROW_THREADS, 1, 1],
+    templateDtypes: { T: x.dtype },
+    templateInts: { R: g.rows, C: g.cols, BT: g.T, KG: g.k, KU: u.k, L: g.L, ROWS_TG: TRELLIS_CODEBOOK_ROW_SG_PER_TG,
+      THREADS: TRELLIS_CODEBOOK_ROW_THREADS, TAIL: g.k !== u.k && tail === "split" ? 1 : 0 },
+  });
+  x2.dispose();
+  const out = ops.reshape(mid!, [...lead, g.rows]);
+  mid!.dispose();
+  return out;
+}

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

@@ -0,0 +1,197 @@
+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 over same-width axis-1 Trellis codes with the row scale
+// factored out: each weight contributes its code value y (trellis_y_f32), and
+// scale·(1/147.8) multiplies each output once after the reduction. Same lanes
+// and runs as fusedGateUpSwiglu; reassociated relative to its f32 code·scale
+// weights (not bit-identical; gate by KL).
+
+// One activation row; two accumulators per projection (even/odd positions).
+const ROW_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 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 = 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>*)(x + 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_y_f32(gwin);
+        const float uv = trellis_y_f32(uwin);
+        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) * (gScale * TRELLIS_RCP);
+  const float up = metal::simd_sum(uAcc + uAcc2) * (uScale * TRELLIS_RCP);
+  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[row] = T(float(siluT) * float(upT));
+  }
+`;
+
+// M activation rows share each decoded weight; one accumulator per row and
+// projection (the rows supply the independent FMA chains).
+const ROWS_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 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], 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 * 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_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 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));
+    }
+  }
+`;
+
+const INPUTS = ["x", "gcodes", "gscales", "ucodes", "uscales"];
+let rowKernel: MetalKernel | undefined, rowsKernel: MetalKernel | undefined;
+
+function check(name: string, gate: TrellisWeights, up: TrellisWeights): void {
+  const a = gate.geometry, b = up.geometry;
+  if (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)
+    throw new Error(`${name}: gate and up need the same axis-1 geometry`);
+}
+function launch(kernel: MetalKernel, x: MlxArray, gate: TrellisWeights, up: TrellisWeights, M: number): MlxArray {
+  const g = gate.geometry, lead = x.shape.slice(0, -1);
+  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, K: g.k, L: g.L, ROWS_TG: TRELLIS_SG_PER_TG },
+  });
+  x2.dispose();
+  const out = ops.reshape(mid!, [...lead, g.rows]);
+  mid!.dispose();
+  return out;
+}
+
+/** silu(gate(x))·up(x) for exactly one activation row. Borrows x and the
+ *  weights; returns owned lazy output with x's leading dimensions. */
+export function gateUpFactoredRow(x: MlxArray, gate: TrellisWeights, up: TrellisWeights): MlxArray {
+  check("gateUpFactoredRow", gate, up);
+  if (x.size !== x.shape.at(-1)) throw new Error("gateUpFactoredRow: expects one activation row");
+  rowKernel ??= new MetalKernel({ name: "mlx_bun_trellis_gateup_factored_row", inputNames: INPUTS,
+    outputNames: ["mid"], source: ROW_SOURCE, header: HEADER, ensureRowContiguous: true });
+  return launch(rowKernel, x, gate, up, 1);
+}
+
+/** silu(gate(x))·up(x) for 2..4 activation rows sharing each decoded weight. */
+export function gateUpFactoredRows(x: MlxArray, gate: TrellisWeights, up: TrellisWeights): MlxArray {
+  check("gateUpFactoredRows", gate, up);
+  const M = x.size / x.shape.at(-1)!;
+  if (M < 2 || M > 4) throw new Error(`gateUpFactoredRows: M=${M} outside 2..4`);
+  rowsKernel ??= new MetalKernel({ name: "mlx_bun_trellis_gateup_factored_rows", inputNames: INPUTS,
+    outputNames: ["mid"], source: ROWS_SOURCE, header: HEADER, ensureRowContiguous: true });
+  return launch(rowsKernel, x, gate, up, M);
+}

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

@@ -0,0 +1,147 @@
+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";
+
+// Fused gate/up/SwiGLU on the simdgroup matrix unit for any number of
+// activation rows: gate-up-mma's fragments, decode and accumulation (each
+// output row bit for bit the same), with every decoded weight fragment feeding
+// an MMA for MT blocks of 8 rows. A threadgroup of SG_TG simdgroups (8 output
+// rows each) generates the 2^L code values y once into a threadgroup codebook
+// (half: every y in [-510, 510] is exact) and stages the activations of its
+// 8·MT rows over each 256-position K step in threadgroup memory, where every
+// simdgroup reads its B fragments; grid z walks chunks of 8·MT rows. The row
+// count is read at run time (one pipeline per geometry and MT). Measured on
+// the M1 Max at the Qwen MLP shape (17408 x 5120, 3-bit): 9 rows 1.01 ms, 16
+// 0.98, 24 1.52, 32 1.87, 64 3.68, 128 7.22 for gate and up together,
+// against 11.3, 10.9, 10.9, 9.5 (the prefill tiles), 5.3 and 8.1 (weight
+// expansion and GEMM).
+const SG_TG = 8;
+
+const RUN_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 = String.raw`
+  const uint M = (uint)x_shape[0];
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint tid = sg * 32u + lane;
+  threadgroup half codebook[1 << L];
+  threadgroup T xs[8 * MT][256];
+  for (uint i = tid; i < (1u << (uint)L); i += 32u * (uint)SG_TG) codebook[i] = half(trellis_y_f32(i));
+  const uint n0 = (threadgroup_position_in_grid.y * (uint)SG_TG + sg) * 8u;
+  const uint fm = ((lane / 4u) & 4u) + ((lane / 2u) % 4u);
+  const uint fn = ((lane / 4u) & 2u) * 2u + (lane % 2u) * 2u;
+  // A simdgroup past the last output row still stages its share and meets every barrier; it writes nothing.
+  const uint row = metal::min(n0 + fm, (uint)R - 1u);
+  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));
+  const uint mBase = threadgroup_position_in_grid.z * 8u * (uint)MT;
+  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) {
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    for (uint e = tid; e < 8u * (uint)MT * 64u; e += 32u * (uint)SG_TG) {
+      const uint r = e / 64u, c4 = e % 64u, m = mBase + r;
+      const metal::vec<T, 4> v = m < M ? ((const device metal::vec<T, 4>*)(x + (ulong)m * (ulong)C + base))[c4]
+        : metal::vec<T, 4>(0);
+      *((threadgroup metal::vec<T, 4>*)(&xs[r][c4 * 4u])) = v;
+    }
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    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 j = 0; j < 32u; ++j) {
+      metal::simdgroup_float8x8 aG, aU;
+      aG.thread_elements()[0] = float(codebook[run_state<KG, L>(ga, j)]);
+      aG.thread_elements()[1] = float(codebook[run_state<KG, L>(gb, j)]);
+      aU.thread_elements()[0] = float(codebook[run_state<KU, L>(ua, j)]);
+      aU.thread_elements()[1] = float(codebook[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] = float(xs[8u * bk + fn][fm * 32u + j]);
+        b.thread_elements()[1] = float(xs[8u * bk + fn + 1u][fm * 32u + j]);
+        metal::simdgroup_multiply_accumulate(accG[bk], aG, b, accG[bk]);
+        metal::simdgroup_multiply_accumulate(accU[bk], aU, b, accU[bk]);
+      }
+    }
+  }
+  if (n0 >= (uint)R) return;
+  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 >= 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;
+
+/** Row blocks per pass for M rows: two (16 rows), or three where three blocks of 8 cover them (17..24 rows: 1.52
+ *  against 1.64 ms for three passes of one block and 1.86 for two of two). */
+function rowBlocks(M: number): number { return Math.ceil(M / 8) === 3 ? 3 : 2; }
+
+/** silu(gate(x))·up(x) for any number of activation rows. Gate and up share their axis-1 geometry (rows a multiple
+ *  of 8, columns of 256, L at most 12); their bit widths may differ. Borrows x and the weights; returns owned lazy
+ *  output with x's leading dimensions. */
+export function gateUpMmaWide(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) throw new Error("gateUpMmaWide: no rows");
+  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 || g.L > 12)
+    throw new Error("gateUpMmaWide: unsupported gate/up geometry");
+  kernel ??= new MetalKernel({ name: "mlx_bun_trellis_gateup_mma_wide", inputNames: ["x", "gcodes", "gscales", "ucodes", "uscales"],
+    outputNames: ["mid"], source: SOURCE, header: HEADER + RUN_HEADER, ensureRowContiguous: true });
+  const MT = rowBlocks(M);
+  using 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), Math.ceil(M / (8 * MT))], threadGroup: [32 * SG_TG, 1, 1],
+    templateDtypes: { T: x.dtype },
+    templateInts: { R: g.rows, C: g.cols, BT: g.T, KG: g.k, KU: u.k, L: g.L, SG_TG, MT },
+  });
+  using y = mid!;
+  return ops.reshape(y, [...lead, g.rows]);
+}