@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/kernels/trellis/down-factored.ts

@@ -0,0 +1,180 @@
+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 { TrellisGeometry } from "./geometry";
+import { HEADER, wordsPerBlock } from "./codebook";
+import { TRELLIS_THREADS, TRELLIS_SG_PER_TG } from "./launch";
+
+// Axis-0 (down) projections over row-major Trellis codes [rows, cols·K/32]
+// with the input scale factored out: each input row's activation is
+// pre-multiplied by scale·(1/147.8) and weights contribute their code value y
+// (trellis_y_f32). Lane-per-word of a block with running sums per position,
+// split-K over input rows into SPLITS f32 partials that one MLX sum folds.
+// Same traversal as trellisScatter; reassociated relative to its f32
+// code·scale weights (not bit-identical; gate by KL).
+const SPLITS = 128;
+
+// One activation row.
+const ROW_SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint wpb = (uint)(BT * K / 32);
+  const uint nb = 32u / wpb;                       // blocks per SIMD group
+  const uint nBlocks = (uint)C / (uint)BT;
+  const uint groups = (nBlocks + nb - 1u) / nb;
+  const uint group = thread_position_in_grid.y * (uint)SG_TG + sg;
+  const uint split = thread_position_in_grid.z;
+  const uint bi = lane / wpb;
+  const uint wi = lane - bi * wpb;
+  const uint blk = group * nb + bi;
+  const bool active = (group < groups) && (bi < nb) && (blk < nBlocks);
+  const uint nbrLane = bi * wpb + ((wi + 1u == wpb) ? 0u : wi + 1u);
+  // This lane's positions: symbol offsets p = multiples of K inside [32wi, 32wi+32).
+  const uint pBase = 32u * wi;
+  const uint p0 = ((pBase + (uint)K - 1u) / (uint)K) * (uint)K;
+  uint offs[NP];
+  bool valid[NP];
+  bool spill[NP];
+  uint cols[NP];
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i < (uint)NP; ++i) {
+    const uint p = p0 + i * (uint)K;
+    valid[i] = active && (p < pBase + 32u);
+    offs[i] = p - pBase;
+    spill[i] = offs[i] + (uint)L > 32u;
+    const uint t = (uint)BT - 1u - p / (uint)K;
+    cols[i] = blk * (uint)BT + t;
+  }
+  const ulong rowWords = (ulong)nBlocks * wpb;
+  const uint rowsPer = ((uint)R + (uint)SPLITS - 1u) / (uint)SPLITS;
+  const uint r0 = split * rowsPer;
+  const uint r1 = metal::min((uint)R, r0 + rowsPer);
+  const device uint32_t* col0 = codes + (ulong)blk * wpb + wi;
+  const uint mask = (1u << (uint)L) - 1u;
+  float acc[NP];
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i < (uint)NP; ++i) acc[i] = 0.0f;
+  for (uint r = r0; r < r1; ++r) {
+    const uint word = active ? col0[(ulong)r * rowWords] : 0u;
+    const uint nxt = metal::simd_shuffle(word, (ushort)nbrLane);
+    const float xr = float(x[r]);
+    const float sr = float(scales[r]);
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < (uint)NP; ++i) {
+      uint win = word >> offs[i];
+      if (spill[i]) win |= nxt << (32u - offs[i]);
+      win &= mask;
+      acc[i] = metal::fma(trellis_y_f32(win), xr * (sr * TRELLIS_RCP), acc[i]);
+    }
+  }
+  const ulong outBase = (ulong)split * (ulong)C;
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i < (uint)NP; ++i)
+    if (valid[i]) partial[outBase + cols[i]] = acc[i];
+`;
+
+// M activation rows sharing each decoded weight.
+const ROWS_SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint wpb = (uint)(BT * K / 32);
+  const uint nb = 32u / wpb;                       // blocks per SIMD group
+  const uint nBlocks = (uint)C / (uint)BT;
+  const uint groups = (nBlocks + nb - 1u) / nb;
+  const uint group = thread_position_in_grid.y * (uint)SG_TG + sg;
+  const uint split = thread_position_in_grid.z;
+  const uint bi = lane / wpb;
+  const uint wi = lane - bi * wpb;
+  const uint blk = group * nb + bi;
+  const bool active = (group < groups) && (bi < nb) && (blk < nBlocks);
+  const uint nbrLane = bi * wpb + ((wi + 1u == wpb) ? 0u : wi + 1u);
+  const uint pBase = 32u * wi;
+  const uint p0 = ((pBase + (uint)K - 1u) / (uint)K) * (uint)K;
+  uint offs[NP];
+  bool valid[NP];
+  bool spill[NP];
+  uint cols[NP];
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i < (uint)NP; ++i) {
+    const uint p = p0 + i * (uint)K;
+    valid[i] = active && (p < pBase + 32u);
+    offs[i] = p - pBase;
+    spill[i] = offs[i] + (uint)L > 32u;
+    const uint t = (uint)BT - 1u - p / (uint)K;
+    cols[i] = blk * (uint)BT + t;
+  }
+  const ulong rowWords = (ulong)nBlocks * wpb;
+  const uint rowsPer = ((uint)R + (uint)SPLITS - 1u) / (uint)SPLITS;
+  const uint r0 = split * rowsPer;
+  const uint r1 = metal::min((uint)R, r0 + rowsPer);
+  const device uint32_t* col0 = codes + (ulong)blk * wpb + wi;
+  const uint mask = (1u << (uint)L) - 1u;
+  float acc[M][NP];
+  #pragma clang loop unroll(full)
+  for (uint m = 0; m < (uint)M; ++m)
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < (uint)NP; ++i) acc[m][i] = 0.0f;
+  for (uint r = r0; r < r1; ++r) {
+    const uint word = active ? col0[(ulong)r * rowWords] : 0u;
+    const uint nxt = metal::simd_shuffle(word, (ushort)nbrLane);
+    float xr[M];
+    #pragma clang loop unroll(full)
+    for (uint m = 0; m < (uint)M; ++m) xr[m] = float(x[(ulong)m * (ulong)R + r]) * (float(scales[r]) * TRELLIS_RCP);
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < (uint)NP; ++i) {
+      uint win = word >> offs[i];
+      if (spill[i]) win |= nxt << (32u - offs[i]);
+      win &= mask;
+      const float weight = trellis_y_f32(win);
+      #pragma clang loop unroll(full)
+      for (uint m = 0; m < (uint)M; ++m) acc[m][i] = metal::fma(weight, xr[m], acc[m][i]);
+    }
+  }
+  #pragma clang loop unroll(full)
+  for (uint m = 0; m < (uint)M; ++m) {
+    const ulong outBase = ((ulong)m * (ulong)SPLITS + split) * (ulong)C;
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < (uint)NP; ++i)
+      if (valid[i]) partial[outBase + cols[i]] = acc[m][i];
+  }
+`;
+
+let rowKernel: MetalKernel | undefined, rowsKernel: MetalKernel | undefined;
+
+function launch(kernel: MetalKernel, x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry, M: number): MlxArray {
+  if (g.axis !== 0 || g.blockInterleave || codes.ndim !== 2) throw new Error("downFactored: needs row-major axis-0 codes");
+  const wpb = wordsPerBlock(g.T, g.k), nb = Math.floor(32 / wpb);
+  if (nb < 1) throw new Error(`downFactored: block of ${wpb} words exceeds one SIMD group`);
+  const groups = Math.ceil((g.cols / g.T) / nb);
+  const [partial] = kernel.apply([x2, codes, scales], {
+    outputs: [{ shape: [M, SPLITS, g.cols], dtype: Dtype.float32 }],
+    grid: [TRELLIS_THREADS, Math.ceil(groups / TRELLIS_SG_PER_TG), SPLITS],
+    threadGroup: [TRELLIS_THREADS, 1, 1],
+    templateDtypes: { T: x2.dtype },
+    templateInts: { M, R: g.rows, C: g.cols, BT: g.T, K: g.k, L: g.L, SG_TG: TRELLIS_SG_PER_TG, SPLITS, NP: Math.ceil(32 / g.k) },
+  });
+  const sum = ops.sumAxis(partial!, 1, false);
+  partial!.dispose();
+  const out = sum.astype(x2.dtype);
+  sum.dispose();
+  return out;
+}
+
+/** x2 [1, inFeatures] · W → [1, outFeatures] in x2's dtype. Borrows inputs. */
+export function downFactoredRow(x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry): MlxArray {
+  if (x2.shape[0] !== 1) throw new Error("downFactoredRow: expects one activation row");
+  rowKernel ??= new MetalKernel({ name: "mlx_bun_trellis_down_factored_row", inputNames: ["x", "codes", "scales"],
+    outputNames: ["partial"], source: ROW_SOURCE, header: HEADER, ensureRowContiguous: true });
+  return launch(rowKernel, x2, codes, scales, g, 1);
+}
+
+/** x2 [M, inFeatures] · W → [M, outFeatures] for 2..4 rows sharing each
+ *  decoded weight. Borrows inputs. */
+export function downFactoredRows(x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry): MlxArray {
+  const M = x2.shape[0]!;
+  if (M < 2 || M > 4) throw new Error(`downFactoredRows: M=${M} outside 2..4`);
+  rowsKernel ??= new MetalKernel({ name: "mlx_bun_trellis_down_factored_rows", inputNames: ["x", "codes", "scales"],
+    outputNames: ["partial"], source: ROWS_SOURCE, header: HEADER, ensureRowContiguous: true });
+  return launch(rowsKernel, x2, codes, scales, g, M);
+}

package/src/kernels/trellis/down-k3-interleaved-factored.ts

@@ -0,0 +1,146 @@
+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 { TrellisGeometry } from "./geometry";
+import { HEADER } from "./codebook";
+import { TRELLIS_THREADS, TRELLIS_SG_PER_TG } from "./launch";
+
+// Axis-0 (down) projections over 3-bit Trellis codes in the block-interleaved
+// layout [cols/512, rows, 48] (T 256, L 12): every lane owns 8 consecutive
+// positions (24 bits) of one block, so all 32 lanes are busy. The input scale
+// is factored out: each input row's activation is pre-multiplied by
+// scale·(1/147.8) and weights contribute their code value y. Split-K over
+// input rows into SPLITS f32 partials folded by one MLX sum. Same traversal as
+// the balanced trellisScatter; reassociated relative to its f32 code·scale
+// weights (not bit-identical; gate by KL).
+const SPLITS = 128, K = 3, BT = 256, L = 12, WPB = BT * K / 32;
+
+// One activation row.
+const ROW_SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint block = thread_position_in_grid.y * (uint)SG_TG + sg;
+  const uint nBlocks = (uint)C / ${BT}u;
+  const uint split = thread_position_in_grid.z;
+  if (block >= nBlocks) return;
+  const ulong rowWords = 2ul * ${WPB}u;
+  const uint pBase = lane * 8u * ${K}u;
+  const uint wi = pBase >> 5u;
+  const uint off = pBase & 31u;
+  const uint next = (wi + 1u == ${WPB}u) ? 0u : wi + 1u;
+  const device uint32_t* words = codes + (ulong)(block / 2) * R * 2 * ${WPB}u + (block % 2) * ${WPB}u;
+  const uint rowsPer = ((uint)R + (uint)SPLITS - 1u) / (uint)SPLITS;
+  const uint r0 = split * rowsPer;
+  const uint r1 = metal::min((uint)R, r0 + rowsPer);
+  float acc[8];
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i < 8u; ++i) acc[i] = 0.0f;
+  for (uint r = r0; r < r1; ++r) {
+    const uint word = words[(ulong)r * rowWords + wi];
+    const uint nxt = words[(ulong)r * rowWords + next];
+    const uint lo = off == 0u ? word : (word >> off) | (nxt << (32u - off));
+    const uint hi = off == 0u ? nxt : nxt >> off;
+    const float xr = float(x[r]);
+    const float sr = float(scales[r]);
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < 8u; ++i) {
+      const uint shift = i * ${K}u;
+      uint win = lo >> shift;
+      if (shift + ${L}u > 32u) win |= hi << (32u - shift);
+      win &= (1u << ${L}u) - 1u;
+      acc[i] = metal::fma(trellis_y_f32(win), xr * (sr * TRELLIS_RCP), acc[i]);
+    }
+  }
+  const ulong outBase = (ulong)split * (ulong)C;
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i < 8u; ++i)
+    partial[outBase + block * ${BT}u + ${BT - 1}u - lane * 8u - i] = acc[i];
+`;
+
+// M activation rows sharing each decoded weight.
+const ROWS_SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint block = thread_position_in_grid.y * (uint)SG_TG + sg;
+  const uint nBlocks = (uint)C / ${BT}u;
+  const uint split = thread_position_in_grid.z;
+  if (block >= nBlocks) return;
+  const ulong rowWords = 2ul * ${WPB}u;
+  const uint pBase = lane * 8u * ${K}u;
+  const uint wi = pBase >> 5u;
+  const uint off = pBase & 31u;
+  const uint next = (wi + 1u == ${WPB}u) ? 0u : wi + 1u;
+  const device uint32_t* words = codes + (ulong)(block / 2) * R * 2 * ${WPB}u + (block % 2) * ${WPB}u;
+  const uint rowsPer = ((uint)R + (uint)SPLITS - 1u) / (uint)SPLITS;
+  const uint r0 = split * rowsPer;
+  const uint r1 = metal::min((uint)R, r0 + rowsPer);
+  float acc[M][8];
+  #pragma clang loop unroll(full)
+  for (uint m = 0; m < (uint)M; ++m)
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < 8u; ++i) acc[m][i] = 0.0f;
+  for (uint r = r0; r < r1; ++r) {
+    const uint word = words[(ulong)r * rowWords + wi];
+    const uint nxt = words[(ulong)r * rowWords + next];
+    const uint lo = off == 0u ? word : (word >> off) | (nxt << (32u - off));
+    const uint hi = off == 0u ? nxt : nxt >> off;
+    float xr[M];
+    #pragma clang loop unroll(full)
+    for (uint m = 0; m < (uint)M; ++m) xr[m] = float(x[(ulong)m * (ulong)R + r]) * (float(scales[r]) * TRELLIS_RCP);
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < 8u; ++i) {
+      const uint shift = i * ${K}u;
+      uint win = lo >> shift;
+      if (shift + ${L}u > 32u) win |= hi << (32u - shift);
+      win &= (1u << ${L}u) - 1u;
+      const float weight = trellis_y_f32(win);
+      #pragma clang loop unroll(full)
+      for (uint m = 0; m < (uint)M; ++m) acc[m][i] = metal::fma(weight, xr[m], acc[m][i]);
+    }
+  }
+  #pragma clang loop unroll(full)
+  for (uint m = 0; m < (uint)M; ++m) {
+    const ulong outBase = ((ulong)m * (ulong)SPLITS + split) * (ulong)C;
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < 8u; ++i)
+      partial[outBase + block * ${BT}u + ${BT - 1}u - lane * 8u - i] = acc[m][i];
+  }
+`;
+
+let rowKernel: MetalKernel | undefined, rowsKernel: MetalKernel | undefined;
+
+function launch(kernel: MetalKernel, x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry, M: number): MlxArray {
+  if (g.axis !== 0 || g.blockInterleave !== 2 || g.k !== K || g.T !== BT || g.L !== L)
+    throw new Error("downK3InterleavedFactored: needs 3-bit block-interleaved axis-0 codes (T 256, L 12)");
+  const [partial] = kernel.apply([x2, codes, scales], {
+    outputs: [{ shape: [M, SPLITS, g.cols], dtype: Dtype.float32 }],
+    grid: [TRELLIS_THREADS, Math.ceil(g.cols / BT / TRELLIS_SG_PER_TG), SPLITS],
+    threadGroup: [TRELLIS_THREADS, 1, 1],
+    templateDtypes: { T: x2.dtype },
+    templateInts: { M, R: g.rows, C: g.cols, SG_TG: TRELLIS_SG_PER_TG, SPLITS },
+  });
+  const sum = ops.sumAxis(partial!, 1, false);
+  partial!.dispose();
+  const out = sum.astype(x2.dtype);
+  sum.dispose();
+  return out;
+}
+
+/** x2 [1, inFeatures] · W → [1, outFeatures] in x2's dtype. Borrows inputs. */
+export function downK3InterleavedFactoredRow(x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry): MlxArray {
+  if (x2.shape[0] !== 1) throw new Error("downK3InterleavedFactoredRow: expects one activation row");
+  rowKernel ??= new MetalKernel({ name: "mlx_bun_trellis_down_k3i_factored_row", inputNames: ["x", "codes", "scales"],
+    outputNames: ["partial"], source: ROW_SOURCE, header: HEADER, ensureRowContiguous: true });
+  return launch(rowKernel, x2, codes, scales, g, 1);
+}
+
+/** x2 [M, inFeatures] · W → [M, outFeatures] for 2..4 rows sharing each
+ *  decoded weight. Borrows inputs. */
+export function downK3InterleavedFactoredRows(x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry): MlxArray {
+  const M = x2.shape[0]!;
+  if (M < 2 || M > 4) throw new Error(`downK3InterleavedFactoredRows: M=${M} outside 2..4`);
+  rowsKernel ??= new MetalKernel({ name: "mlx_bun_trellis_down_k3i_factored_rows", inputNames: ["x", "codes", "scales"],
+    outputNames: ["partial"], source: ROWS_SOURCE, header: HEADER, ensureRowContiguous: true });
+  return launch(rowsKernel, x2, codes, scales, g, M);
+}

package/src/kernels/trellis/down-mma-wide.ts

@@ -0,0 +1,195 @@
+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 { HEADER } from "./codebook";
+import type { TrellisGeometry } from "./geometry";
+
+// Down projection on the simdgroup matrix unit for any number of activation
+// rows: down-mma's splits and accumulation (each output row bit for bit the
+// same), with every decoded B fragment feeding an MMA for MT blocks of 8 rows;
+// grid z walks chunks of 8·MT rows. A simdgroup's S sub-tiles and MT follow
+// the layout and the rows (`plan`): neither changes a sum. Each code value y
+// is read from a half table of the 2^L values (exact in half) the threadgroup
+// generates once, in place of the 1MAD arithmetic (the M1 Max form, where the
+// matrix products issue on the ordinary ALU pipes and the integer decode
+// competes with them); 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) and the B fragments are half. Input rows come in whole steps of 8 (the
+// rows a multiple of 8, so is every split's); activation rows past M repeat
+// the last row and are not written. The row count is read at run time (one
+// pipeline per geometry and launch). SPLITS f32 partials [M, SPLITS, C] fold
+// in one MLX sum. Two kernels, one per code layout: row-major [rows, cols·K/32]
+// and the 3-bit block interleave [cols/512, rows, 48]. Measured on the M1 Max at
+// the Qwen down shape (17408 x 5120, real weights) at 8, 16, 24, 32 and 64 rows,
+// 3-bit interleaved: 0.24, 0.44, 0.64, 0.84 and 1.65 ms, against 0.66, 0.67,
+// 0.87, 1.24 and 2.39 decoding with the 1MAD arithmetic (downMma at 8 rows
+// 0.45); 2-bit row-major: 0.24, 0.40, 0.58, 0.76 and 1.49, against 0.62, 0.64,
+// 0.88, 1.08 and 2.13 (downMma 0.43).
+const SG_TG = 4, SPLITS = 8;
+
+/** One pass of the kernel: S sub-tiles per simdgroup, MT blocks of 8 rows per decoded fragment. */
+export interface DownPass { readonly S: number; readonly MT: number }
+
+/** One block of 8 rows per pass, sixteen sub-tiles: the cheapest pass for 1..8 rows in either layout. */
+export const DOWN_ONE_BLOCK: DownPass = Object.freeze({ S: 16, MT: 1 });
+
+/** Each layout's passes and their relative cost, measured per pass at 8 to 128 rows: the 3-bit interleave runs one
+ *  block per pass at every row count; row-major codes take sixteen sub-tiles for one block and eight for two or
+ *  three. */
+const PLANS: Readonly<Record<"rowMajor" | "interleaved", readonly (DownPass & { readonly cost: number })[]>> = {
+  interleaved: [{ ...DOWN_ONE_BLOCK, cost: 1 }],
+  rowMajor: [{ ...DOWN_ONE_BLOCK, cost: 0.6 }, { S: 8, MT: 2, cost: 1 }, { S: 8, MT: 3, cost: 1.44 }],
+};
+
+/** The pass for M rows: the one whose passes cost least in all (the more blocks on a tie). */
+function plan(layout: keyof typeof PLANS, M: number): DownPass {
+  let best = { S: 0, MT: 0, cost: Infinity };
+  for (const { S, MT, cost } of PLANS[layout]) {
+    const total = Math.ceil(M / (8 * MT)) * cost;
+    if (total <= best.cost) best = { S, MT, cost: total };
+  }
+  return { S: best.S, MT: best.MT };
+}
+
+const SOURCE = String.raw`
+  const uint M = (uint)x_shape[0];
+  const uint lane = thread_index_in_simdgroup;
+  threadgroup half codebook[1 << L];
+  for (uint i = simdgroup_index_in_threadgroup * 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 id = threadgroup_position_in_grid.y * (uint)SG_TG + simdgroup_index_in_threadgroup;
+  const uint groupsPerBlock = 32u / (uint)S;
+  const uint nGroups = ((uint)C / 256u) * groupsPerBlock;
+  if (id >= nGroups * (uint)SPLITS) return;
+  const uint cg = id % nGroups, split = id / nGroups;
+  const uint block = cg / groupsPerBlock, pg = (cg % groupsPerBlock) * 8u * (uint)S;
+  const uint fm = ((lane / 4u) & 4u) + ((lane / 2u) % 4u);
+  const uint fn = ((lane / 4u) & 2u) * 2u + (lane % 2u) * 2u;
+  const uint wpb = 8u * (uint)K;
+  const ulong rowWords = (ulong)(INTERLEAVE ? 2u : (uint)C / 256u) * wpb;
+  const device uint32_t* words = codes + (INTERLEAVE
+    ? (ulong)(block / 2u) * (ulong)R * 2u * wpb + (block % 2u) * wpb : (ulong)block * wpb);
+  // This lane's 2S positions start at p0; realign its bit window to bit 0.
+  const uint p0 = pg + (uint)S * fn;
+  const uint bit0 = p0 * (uint)K, w0 = bit0 >> 5, off0 = bit0 & 31u;
+  const uint mBase = threadgroup_position_in_grid.z * 8u * (uint)MT;
+  const uint rowsPer = (((uint)R + (uint)SPLITS - 1u) / (uint)SPLITS + 7u) & ~7u;
+  const uint r0 = split * rowsPer;
+  const uint r1 = metal::min((uint)R, r0 + rowsPer);
+  const device T* xm[MT];
+  metal::simdgroup_float8x8 acc[S][MT];
+  #pragma clang loop unroll(full)
+  for (uint bk = 0; bk < (uint)MT; ++bk) {
+    xm[bk] = x + (ulong)metal::min(mBase + 8u * bk + fm, M - 1u) * (ulong)R;
+    #pragma clang loop unroll(full)
+    for (uint s = 0; s < (uint)S; ++s) acc[s][bk] = metal::make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
+  }
+  // The window's words in a row's block (the last ones wrap to its first).
+  uint wordAt[NW + 1];
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i <= (uint)NW; ++i) wordAt[i] = (w0 + i) % wpb;
+  const device uint32_t* rowCodes = words + (ulong)(r0 + fm) * rowWords;
+  for (uint r = r0; r < r1; r += 8u, rowCodes += 8u * rowWords) {
+    uint raw[NW + 1];
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i <= (uint)NW; ++i) raw[i] = rowCodes[wordAt[i]];
+    uint w[NW], v[NW];
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < (uint)NW; ++i) w[i] = off0 ? (raw[i] >> off0) | (raw[i + 1] << (32u - off0)) : raw[i];
+    // Virtual word i: bits [32i - 16, 32i + 16) of the realigned window.
+    #pragma clang loop unroll(full)
+    for (uint i = 1; i < (uint)NW; ++i) v[i] = (w[i - 1] >> 16) | (w[i] << 16);
+    const uint ra = r + fn;
+    const float sa = float(scales[ra]) * TRELLIS_RCP, sb = float(scales[ra + 1u]) * TRELLIS_RCP;
+    metal::simdgroup_float8x8 a[MT];
+    #pragma clang loop unroll(full)
+    for (uint bk = 0; bk < (uint)MT; ++bk) {
+      a[bk].thread_elements()[0] = float(xm[bk][ra]) * sa;
+      a[bk].thread_elements()[1] = float(xm[bk][ra + 1u]) * sb;
+    }
+    #pragma clang loop unroll(full)
+    for (uint s = 0; s < (uint)S; ++s) {
+      metal::simdgroup_half8x8 b;
+      b.thread_elements()[0] = codebook[block_state<K, L>(w, v, s)];
+      b.thread_elements()[1] = codebook[block_state<K, L>(w, v, s + (uint)S)];
+      #pragma clang loop unroll(full)
+      for (uint bk = 0; bk < (uint)MT; ++bk)
+        metal::simdgroup_multiply_accumulate(acc[s][bk], a[bk], b, acc[s][bk]);
+    }
+  }
+  #pragma clang loop unroll(full)
+  for (uint bk = 0; bk < (uint)MT; ++bk) {
+    const uint m = mBase + 8u * bk + fm;
+    if (m >= M) continue;
+    device float* out = partial + ((ulong)m * (ulong)SPLITS + split) * (ulong)C + (ulong)block * 256u + 255u - p0;
+    #pragma clang loop unroll(full)
+    for (uint s = 0; s < (uint)S; ++s) {
+      out[-(int)s] = acc[s][bk].thread_elements()[0];
+      out[-(int)((uint)S + s)] = acc[s][bk].thread_elements()[1];
+    }
+  }
+`;
+
+// State at position t of a realigned window (bits [t·K, +L) of w), one extract (virtual words v past a word edge).
+const STATE_HEADER = String.raw`
+template <int K, int L_>
+inline uint block_state(thread const uint* w, thread const uint* v, uint t) {
+  const uint b = t * (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);
+}
+`;
+
+function launch(kernel: MetalKernel, x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry, interleave: 0 | 1,
+  { S, MT }: DownPass): MlxArray {
+  const M = x2.shape[0]!;
+  if (x2.ndim !== 2 || M < 1) throw new Error("downMmaWide: expects [rows, inputs]");
+  if (g.axis !== 0 || g.T !== 256 || g.cols % 256 !== 0 || g.rows % 8 !== 0 || g.L > 12)
+    throw new Error("downMmaWide: unsupported geometry");
+  const nGroups = (g.cols / 256) * (32 / S);
+  // Words covering 2S positions plus the L-bit window, after realignment.
+  const NW = Math.ceil((2 * S * g.k + g.L - 1) / 32) + 1;
+  const [partial] = kernel.apply([x2, codes, scales], {
+    outputs: [{ shape: [M, SPLITS, g.cols], dtype: Dtype.float32 }],
+    grid: [32 * SG_TG, Math.ceil((nGroups * SPLITS) / SG_TG), Math.ceil(M / (8 * MT))], threadGroup: [32 * SG_TG, 1, 1],
+    templateDtypes: { T: x2.dtype },
+    templateInts: { R: g.rows, C: g.cols, K: g.k, L: g.L, SPLITS, SG_TG, INTERLEAVE: interleave, MT, S, NW },
+  });
+  using sum = ops.sumAxis(partial!, 1, false);
+  partial!.dispose();
+  return sum.astype(x2.dtype);
+}
+
+let rowMajor: MetalKernel | undefined, interleaved: MetalKernel | undefined;
+
+/** x2 [M, inFeatures] · W → [M, outFeatures] in x2's dtype for row-major axis-0 codes (T 256, L at most 12, input
+ *  features a multiple of 8), any number of rows, in passes of the given shape. Borrows inputs. */
+export function downMmaWidePass(x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry, pass: DownPass): MlxArray {
+  if (g.blockInterleave || codes.ndim !== 2) throw new Error("downMmaWide: needs row-major codes");
+  rowMajor ??= new MetalKernel({ name: "mlx_bun_trellis_down_mma_wide", inputNames: ["x", "codes", "scales"],
+    outputNames: ["partial"], source: SOURCE, header: HEADER + STATE_HEADER, ensureRowContiguous: true });
+  return launch(rowMajor, x2, codes, scales, g, 0, pass);
+}
+
+/** downMmaWidePass for 3-bit block-interleaved codes [cols/512, rows, 48]. */
+export function downK3InterleavedMmaWidePass(x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry,
+  pass: DownPass): MlxArray {
+  if (g.blockInterleave !== 2 || g.k !== 3) throw new Error("downK3InterleavedMmaWide: needs 3-bit block-interleaved codes");
+  interleaved ??= new MetalKernel({ name: "mlx_bun_trellis_down_k3i_mma_wide", inputNames: ["x", "codes", "scales"],
+    outputNames: ["partial"], source: SOURCE, header: HEADER + STATE_HEADER, ensureRowContiguous: true });
+  return launch(interleaved, x2, codes, scales, g, 1, pass);
+}
+
+/** downMmaWidePass in the pass the row count plans. */
+export function downMmaWide(x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry): MlxArray {
+  return downMmaWidePass(x2, codes, scales, g, plan("rowMajor", x2.shape[0]!));
+}
+
+/** downK3InterleavedMmaWidePass in the pass the row count plans. */
+export function downK3InterleavedMmaWide(x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry): MlxArray {
+  return downK3InterleavedMmaWidePass(x2, codes, scales, g, plan("interleaved", x2.shape[0]!));
+}

package/src/kernels/trellis/down-mma.ts

@@ -0,0 +1,145 @@
+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 { HEADER } from "./codebook";
+import type { TrellisGeometry } from "./geometry";
+
+// Verify-width down projection on the simdgroup matrix unit, 1..8 activation
+// rows (one block of 8). A simdgroup owns 8·S consecutive positions of one
+// 256-position coded block and a split of the input rows. Per 8-row step every
+// lane decodes 2S contiguous positions of one input row (its bit window
+// realigned once) into the B fragments of S output sub-tiles; the A fragments
+// hold x·scale·(1/147.8) and each feeds S multiplies. Sub-tile s, fragment
+// column i is position P_g + S·i + s (output column 255 - that, per block).
+// SPLITS f32 partials [M, SPLITS, C] fold in one MLX sum. Two kernels, one per
+// code layout: row-major [rows, cols·K/32] and the 3-bit block interleave
+// [cols/512, rows, 48]. Not bit-identical to the matvec kernels; gate by KL.
+const SG_TG = 4, MT = 1, S = 16, SPLITS = 8;
+
+const SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint id = threadgroup_position_in_grid.y * (uint)SG_TG + simdgroup_index_in_threadgroup;
+  const uint groupsPerBlock = 32u / (uint)S;
+  const uint nGroups = ((uint)C / 256u) * groupsPerBlock;
+  if (id >= nGroups * (uint)SPLITS) return;
+  const uint cg = id % nGroups, split = id / nGroups;
+  const uint block = cg / groupsPerBlock, pg = (cg % groupsPerBlock) * 8u * (uint)S;
+  const uint fm = ((lane / 4u) & 4u) + ((lane / 2u) % 4u);
+  const uint fn = ((lane / 4u) & 2u) * 2u + (lane % 2u) * 2u;
+  const uint wpb = 8u * (uint)K;
+  const ulong rowWords = (ulong)(INTERLEAVE ? 2u : (uint)C / 256u) * wpb;
+  const device uint32_t* words = codes + (INTERLEAVE
+    ? (ulong)(block / 2u) * (ulong)R * 2u * wpb + (block % 2u) * wpb : (ulong)block * wpb);
+  // This lane's 2S positions start at p0; realign its bit window to bit 0.
+  const uint p0 = pg + (uint)S * fn;
+  const uint bit0 = p0 * (uint)K, w0 = bit0 >> 5, off0 = bit0 & 31u;
+  const uint mBase = threadgroup_position_in_grid.z * 8u * (uint)MT;
+  const uint rowsPer = (((uint)R + (uint)SPLITS - 1u) / (uint)SPLITS + 7u) & ~7u;
+  const uint r0 = split * rowsPer;
+  const uint r1 = metal::min((uint)R, r0 + rowsPer);
+  metal::simdgroup_float8x8 acc[S][MT];
+  #pragma clang loop unroll(full)
+  for (uint s = 0; s < (uint)S; ++s)
+    #pragma clang loop unroll(full)
+    for (uint bk = 0; bk < (uint)MT; ++bk) acc[s][bk] = metal::make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
+  for (uint r = r0; r < r1; r += 8u) {
+    const uint rw = r + fm;
+    const device uint32_t* rowCodes = words + (ulong)metal::min(rw, (uint)R - 1u) * rowWords;
+    uint raw[NW + 1];
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i <= (uint)NW; ++i) {
+      const uint wi = w0 + i;
+      raw[i] = rowCodes[wi % wpb];
+    }
+    uint w[NW];
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < (uint)NW; ++i) w[i] = off0 ? (raw[i] >> off0) | (raw[i + 1] << (32u - off0)) : raw[i];
+    metal::simdgroup_float8x8 a[MT];
+    const uint ra = r + fn;
+    const float sa = ra < r1 ? float(scales[ra]) * TRELLIS_RCP : 0.0f;
+    const float sb = ra + 1u < r1 ? float(scales[ra + 1u]) * TRELLIS_RCP : 0.0f;
+    #pragma clang loop unroll(full)
+    for (uint bk = 0; bk < (uint)MT; ++bk) {
+      const uint m = mBase + 8u * bk + fm;
+      const bool ok = m < (uint)M;
+      const device T* xm = x + (ulong)(ok ? m : 0u) * (ulong)R;
+      a[bk].thread_elements()[0] = ok && ra < r1 ? float(xm[ra]) * sa : 0.0f;
+      a[bk].thread_elements()[1] = ok && ra + 1u < r1 ? float(xm[ra + 1u]) * sb : 0.0f;
+    }
+    const bool rowOk = rw < r1;
+    #pragma clang loop unroll(full)
+    for (uint s = 0; s < (uint)S; ++s) {
+      metal::simdgroup_float8x8 b;
+      b.thread_elements()[0] = rowOk ? trellis_y_f32((block_state<K, L>(w, s))) : 0.0f;
+      b.thread_elements()[1] = rowOk ? trellis_y_f32((block_state<K, L>(w, s + (uint)S))) : 0.0f;
+      #pragma clang loop unroll(full)
+      for (uint bk = 0; bk < (uint)MT; ++bk)
+        metal::simdgroup_multiply_accumulate(acc[s][bk], a[bk], b, acc[s][bk]);
+    }
+  }
+  #pragma clang loop unroll(full)
+  for (uint bk = 0; bk < (uint)MT; ++bk) {
+    const uint m = mBase + 8u * bk + fm;
+    if (m >= (uint)M) continue;
+    device float* out = partial + ((ulong)m * (ulong)SPLITS + split) * (ulong)C + (ulong)block * 256u + 255u - p0;
+    #pragma clang loop unroll(full)
+    for (uint s = 0; s < (uint)S; ++s) {
+      out[-(int)s] = acc[s][bk].thread_elements()[0];
+      out[-(int)((uint)S + s)] = acc[s][bk].thread_elements()[1];
+    }
+  }
+`;
+
+// State at position t of a realigned window (bit t·K of w, LSB first).
+const STATE_HEADER = String.raw`
+template <int K, int L_>
+inline uint block_state(thread const uint* w, uint t) {
+  const uint b = t * (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);
+}
+`;
+
+
+function launch(kernel: MetalKernel, x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry, interleave: 0 | 1): MlxArray {
+  const M = x2.shape[0]!;
+  if (M < 1 || M > 8) throw new Error(`downMma: M=${M} outside 1..8`);
+  if (g.axis !== 0 || g.T !== 256 || g.cols % 256 !== 0 || g.L > 12) throw new Error("downMma: unsupported geometry");
+  const nGroups = (g.cols / 256) * (32 / S);
+  // Words covering 2S positions plus the L-bit window, after realignment.
+  const NW = Math.ceil((2 * S * g.k + g.L - 1) / 32) + 1;
+  const [partial] = kernel.apply([x2, codes, scales], {
+    outputs: [{ shape: [M, SPLITS, g.cols], dtype: Dtype.float32 }],
+    grid: [32 * SG_TG, Math.ceil((nGroups * SPLITS) / SG_TG), 1], threadGroup: [32 * SG_TG, 1, 1],
+    templateDtypes: { T: x2.dtype },
+    templateInts: { M, R: g.rows, C: g.cols, K: g.k, L: g.L, SPLITS, SG_TG, INTERLEAVE: interleave, MT, S, NW },
+  });
+  const sum = ops.sumAxis(partial!, 1, false);
+  partial!.dispose();
+  const out = sum.astype(x2.dtype);
+  sum.dispose();
+  return out;
+}
+
+let rowMajor: MetalKernel | undefined, interleaved: MetalKernel | undefined;
+
+/** x2 [M, inFeatures] · W → [M, outFeatures] in x2's dtype for row-major
+ *  axis-0 codes (T 256). Borrows inputs. */
+export function downMma(x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry): MlxArray {
+  if (g.blockInterleave || codes.ndim !== 2) throw new Error("downMma: needs row-major codes");
+  rowMajor ??= new MetalKernel({ name: "mlx_bun_trellis_down_mma", inputNames: ["x", "codes", "scales"],
+    outputNames: ["partial"], source: SOURCE, header: HEADER + STATE_HEADER, ensureRowContiguous: true });
+  return launch(rowMajor, x2, codes, scales, g, 0);
+}
+
+/** downMma for 3-bit block-interleaved codes [cols/512, rows, 48]. */
+export function downK3InterleavedMma(x2: MlxArray, codes: MlxArray, scales: MlxArray, g: TrellisGeometry): MlxArray {
+  if (g.blockInterleave !== 2 || g.k !== 3) throw new Error("downK3InterleavedMma: needs 3-bit block-interleaved codes");
+  interleaved ??= new MetalKernel({ name: "mlx_bun_trellis_down_k3i_mma", inputNames: ["x", "codes", "scales"],
+    outputNames: ["partial"], source: SOURCE, header: HEADER + STATE_HEADER, ensureRowContiguous: true });
+  return launch(interleaved, x2, codes, scales, g, 1);
+}