@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/kernels/quantization/affine-split-mma.ts

@@ -0,0 +1,144 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+
+// Matmul over MLX affine 3- or 4-bit, group-64 weights on the simdgroup matrix
+// unit: y = x · Wᵀ for any number of activation rows, in blocks of 8 (grid z;
+// the row count is read at run time, one pipeline per projection shape).
+// The A/B layout of affine{3,4}-mma: A is 8 output rows x 8 k with the k index
+// permuted so that A column i at step s is position base + 32·i + s (every lane
+// dequantizes one 64-value group of its row over a 256-value K step); B is xᵀ.
+// For the M1 Max, whose matrix products issue on the ordinary ALU pipes: each
+// simdgroup owns TILES 8-row output tiles that share every B fragment, the
+// threadgroup's SPLITS simdgroups take interleaved K steps and add their tiles
+// through threadgroup memory, and a weight dequantizes as fma(1 + q/2^BITS,
+// 2^BITS·scale, bias − 2^BITS·scale), q shifted into the mantissa of 1.0f.
+// B columns past the last activation row repeat it; those columns of C are
+// never written, so a block's cost does not depend on its rows. Summation order
+// differs from MLX's qmv and qmm (not bit-identical; gate by KL). Measured on
+// the M1 Max above eight rows (3-bit 5120 -> 10240): 9 rows 0.50 ms, 16 0.38,
+// 32 0.66, 64 1.25, against MLX's quantized matmul at 0.74, 0.75, 0.67, 1.27.
+
+const SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint fm = ((lane / 4u) & 4u) + ((lane / 2u) % 4u);
+  const uint fn = ((lane / 4u) & 2u) * 2u + (lane % 2u) * 2u;
+  const uint M = (uint)x_shape[0];
+  const uint n0 = threadgroup_position_in_grid.y * 8u * (uint)TILES;
+  const uint m0 = threadgroup_position_in_grid.z * 8u;
+  const uint groups = (uint)K / 64u;
+  const uint groupWords = 2u * (uint)BITS;
+  const device T* x0 = x + (ulong)metal::min(m0 + fn, M - 1u) * (ulong)K + fm * 32u;
+  const device T* x1 = x + (ulong)metal::min(m0 + fn + 1u, M - 1u) * (ulong)K + fm * 32u;
+  const device uint32_t* wRow[TILES];
+  metal::simdgroup_float8x8 acc[TILES];
+  #pragma clang loop unroll(full)
+  for (uint t = 0; t < (uint)TILES; ++t) {
+    wRow[t] = (const device uint32_t*)w + (ulong)(n0 + 8u * t + fm) * groups * groupWords;
+    acc[t] = metal::make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
+  }
+  for (uint step = sg; step < (uint)K / 256u; step += (uint)SPLITS) {
+    const uint base = step * 256u;
+    const uint g = step * 4u + fn / 2u;
+    uint q[TILES][2 * BITS];
+    float s2[TILES], b2[TILES];
+    #pragma clang loop unroll(full)
+    for (uint t = 0; t < (uint)TILES; ++t) {
+      const ulong at = (ulong)(n0 + 8u * t + fm) * groups + g;
+      s2[t] = float(scales[at]) * (float)(1u << (uint)BITS);
+      b2[t] = float(biases[at]) - s2[t];
+      #pragma clang loop unroll(full)
+      for (uint i = 0; i < groupWords; ++i) q[t][i] = wRow[t][g * groupWords + i];
+    }
+    const device metal::vec<T, 4>* xa = (const device metal::vec<T, 4>*)(x0 + base);
+    const device metal::vec<T, 4>* xb = (const device metal::vec<T, 4>*)(x1 + base);
+    #pragma clang loop unroll(full)
+    for (uint s4 = 0; s4 < 8u; ++s4) {
+      const float4 va = float4(xa[s4]);
+      const float4 vb = float4(xb[s4]);
+      #pragma clang loop unroll(full)
+      for (uint c = 0; c < 4u; ++c) {
+        const uint v = s4 * 4u + c;
+        metal::simdgroup_float8x8 bx;
+        bx.thread_elements()[0] = va[c];
+        bx.thread_elements()[1] = vb[c];
+        #pragma clang loop unroll(full)
+        for (uint t = 0; t < (uint)TILES; ++t) {
+          metal::simdgroup_float8x8 a;
+          a.thread_elements()[0] = metal::fma(AFFINE_UNIT(q[t], v), s2[t], b2[t]);
+          a.thread_elements()[1] = metal::fma(AFFINE_UNIT(q[t], v + 32u), s2[t], b2[t]);
+          metal::simdgroup_multiply_accumulate(acc[t], a, bx, acc[t]);
+        }
+      }
+    }
+  }
+  // Add the K splits: element (fm, fn | fn + 1) of each tile, then thread e of
+  // the threadgroup writes C element e of its tile, columns below M only.
+  threadgroup float part[SPLITS][TILES][64];
+  #pragma clang loop unroll(full)
+  for (uint t = 0; t < (uint)TILES; ++t) {
+    part[sg][t][fm * 8u + fn] = acc[t].thread_elements()[0];
+    part[sg][t][fm * 8u + fn + 1u] = acc[t].thread_elements()[1];
+  }
+  threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+  for (uint e = sg * 32u + lane; e < 64u * (uint)TILES; e += 32u * (uint)SPLITS) {
+    const uint t = e / 64u, r = (e % 64u) / 8u, m = e % 8u;
+    if (m0 + m >= M) continue;
+    float sum = 0.0f;
+    for (uint s = 0; s < (uint)SPLITS; ++s) sum += part[s][t][r * 8u + m];
+    y[(ulong)(m0 + m) * (ulong)N + n0 + 8u * t + r] = T(sum);
+  }
+`;
+
+// Value v of a 64-value group (MLX's LSB-first bit stream) as the exact float
+// 1 + q/2^B: q shifted into the top B mantissa bits of 1.0f; a value may
+// straddle into the next word. The words go by value: v is a constant of the
+// unrolled loop, so every shift and mask folds.
+const HEADER = String.raw`
+template <uint B>
+inline float affine_unit(uint lo, uint hi, uint v) {
+  const uint bit = v * B, off = bit & 31u;
+  const uint top = 23u - B;
+  uint u = off <= top ? lo << (top - off) : lo >> (off - top);
+  if (off + B > 32u) u |= hi << (top + 32u - off);
+  return as_type<float>((u & (((1u << B) - 1u) << top)) | 0x3f800000u);
+}
+#define AFFINE_UNIT(q, v) affine_unit<BITS>((q)[((v) * BITS) >> 5], (q)[metal::min((((v) * BITS) >> 5) + 1u, 2u * BITS - 1u)], (v))
+`;
+
+let kernel: MetalKernel | undefined;
+
+/** The launch for an N-output projection: output tiles per simdgroup and K splits per threadgroup. Wide outputs take
+ *  two tiles over two splits; narrow ones (the attention k/v and the DeltaNet a/b projections) one tile over eight,
+ *  for enough simdgroups in flight. */
+export interface AffineSplitMmaPlan { readonly tiles: number; readonly splits: number }
+export function affineSplitMmaPlan(N: number): AffineSplitMmaPlan {
+  return N >= 2048 && N % 16 === 0 ? { tiles: 2, splits: 2 } : { tiles: 1, splits: 8 };
+}
+
+/** Eligible: affine 3- or 4-bit, group 64, K a multiple of 256, N of 8, at least one row. */
+export function affineSplitMmaEligible(rows: number, w: MlxArray, scales: MlxArray, spec: { bits: number; groupSize: number; mode: string }): boolean {
+  if (spec.mode !== "affine" || spec.groupSize !== 64 || (spec.bits !== 3 && spec.bits !== 4)) return false;
+  const N = scales.shape[0]!, K = scales.shape[1]! * 64;
+  return rows >= 1 && K % 256 === 0 && N % 8 === 0 && w.shape.length === 2;
+}
+
+/** x [..., K] (any number of rows) · Wᵀ → [..., N] in x's dtype, for `bits`-bit weights under `plan`. Borrows
+ *  inputs. */
+export function affineSplitMma(x: MlxArray, w: MlxArray, scales: MlxArray, biases: MlxArray, bits: 3 | 4,
+  plan: AffineSplitMmaPlan): MlxArray {
+  const K = x.shape.at(-1)!, M = x.size / K, N = scales.shape[0]!;
+  if (M < 1) throw new Error("affineSplitMma: no rows");
+  kernel ??= new MetalKernel({ name: "mlx_bun_affine_split_mma", inputNames: ["x", "w", "scales", "biases"],
+    outputNames: ["y"], source: SOURCE, header: HEADER, ensureRowContiguous: true });
+  using x2 = ops.reshape(x, [M, K]);
+  const [y] = kernel.apply([x2, w, scales, biases], {
+    outputs: [{ shape: [M, N], dtype: x.dtype }],
+    grid: [32 * plan.splits, N / (8 * plan.tiles), Math.ceil(M / 8)], threadGroup: [32 * plan.splits, 1, 1],
+    templateDtypes: { T: x.dtype },
+    templateInts: { N, K, BITS: bits, TILES: plan.tiles, SPLITS: plan.splits },
+  });
+  using y2 = y!;
+  return ops.reshape(y2, [...x.shape.slice(0, -1), N]);
+}

package/src/kernels/quantization/affine3-mma.ts

@@ -0,0 +1,117 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+
+// Verify-width matmul over MLX affine 3-bit, group-64 weights on the simdgroup
+// matrix unit: y = x · Wᵀ for 1..8 activation rows (one block of 8). One
+// simdgroup owns 8 output 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
+// dequantizes 64 contiguous values of its row (exactly one quantization group,
+// MLX's LSB-first bit stream) into its A-fragment elements, and reads two
+// activation rows contiguously. Per-weight work is independent of M. Summation
+// order differs from MLX's qmv (not bit-identical; gate by KL).
+const SG_TG = 4, MT = 1;
+
+const SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint n0 = (threadgroup_position_in_grid.y * (uint)SG_TG + simdgroup_index_in_threadgroup) * 8u;
+  if (n0 >= (uint)N) return;
+  const uint fm = ((lane / 4u) & 4u) + ((lane / 2u) % 4u);
+  const uint fn = ((lane / 4u) & 2u) * 2u + (lane % 2u) * 2u;
+  const uint n = n0 + fm;
+  const uint groups = (uint)K / 64u;
+  const device uint32_t* wRow = (const device uint32_t*)((const device uint8_t*)w + (ulong)n * (ulong)K * 3u / 8u);
+  // 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];
+  metal::simdgroup_float8x8 acc[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)K + fm * 32u;
+    x1[bk] = x + (ulong)(m1[bk] ? r0 + 1u : 0u) * (ulong)K + fm * 32u;
+    acc[bk] = metal::make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
+  }
+  for (uint base = 0; base < (uint)K; base += 256u) {
+    const uint g = base / 64u + fn / 2u;
+    const float s = float(scales[(ulong)n * groups + g]);
+    const float b = float(biases[(ulong)n * groups + g]);
+    // 64 values of one group: 6 words (4-byte aligned).
+    uint q[6];
+    const device uint32_t* qp = wRow + g * 6u;
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < 6u; ++i) q[i] = qp[i];
+    #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 c = 0; c < 4u; ++c) {
+        const uint t = s4 * 4u + c;
+        metal::simdgroup_float8x8 a;
+        a.thread_elements()[0] = metal::fma(s, affine_q3(q, t), b);
+        a.thread_elements()[1] = metal::fma(s, affine_q3(q, t + 32u), b);
+        #pragma clang loop unroll(full)
+        for (uint bk = 0; bk < (uint)MT; ++bk) {
+          metal::simdgroup_float8x8 bx;
+          bx.thread_elements()[0] = va[bk][c];
+          bx.thread_elements()[1] = vb[bk][c];
+          metal::simdgroup_multiply_accumulate(acc[bk], a, bx, acc[bk]);
+        }
+      }
+    }
+  }
+  #pragma clang loop unroll(full)
+  for (uint bk = 0; bk < (uint)MT; ++bk) {
+    const uint r0 = mBase + 8u * bk + fn;
+    if (m0[bk]) y[(ulong)r0 * (ulong)N + n] = T(acc[bk].thread_elements()[0]);
+    if (m1[bk]) y[(ulong)(r0 + 1u) * (ulong)N + n] = T(acc[bk].thread_elements()[1]);
+  }
+`;
+
+// Value t of a 64-value group as an exact f32 (2^23 mantissa trick); a 3-bit
+// value may straddle two words.
+const HEADER = String.raw`
+inline float affine_q3(thread const uint* q, uint t) {
+  const uint bit = t * 3u;
+  const uint wi = bit >> 5;
+  const uint off = bit & 31u;
+  uint v = q[wi] >> off;
+  if (off + 3u > 32u) v |= q[wi + 1] << (32u - off);
+  return as_type<float>(0x4b000000u | (v & 7u)) - 8388608.0f;
+}
+`;
+
+let kernel: MetalKernel | undefined;
+
+/** Eligible: affine 3-bit, group 64, K a multiple of 256, N of 8, 1..8 rows. */
+export function affine3MmaEligible(rows: number, w: MlxArray, scales: MlxArray, spec: { bits: number; groupSize: number; mode: string }): boolean {
+  if (spec.mode !== "affine" || spec.groupSize !== 64 || spec.bits !== 3) return false;
+  const N = scales.shape[0]!, K = scales.shape[1]! * 64;
+  return rows >= 1 && rows <= 8 && K % 256 === 0 && N % 8 === 0 && w.shape.length === 2;
+}
+
+/** x [..., K] (1..8 rows) · Wᵀ → [..., N] in x's dtype. Borrows inputs. */
+export function affine3Mma(x: MlxArray, w: MlxArray, scales: MlxArray, biases: MlxArray): MlxArray {
+  const K = x.shape.at(-1)!, M = x.size / K, N = scales.shape[0]!;
+  if (M < 1 || M > 8) throw new Error(`affine3Mma: M=${M} outside 1..8`);
+  kernel ??= new MetalKernel({ name: "mlx_bun_affine3_mma", inputNames: ["x", "w", "scales", "biases"],
+    outputNames: ["y"], source: SOURCE, header: HEADER, ensureRowContiguous: true });
+  using x2 = ops.reshape(x, [M, K]);
+  const [y] = kernel.apply([x2, w, scales, biases], {
+    outputs: [{ shape: [M, N], dtype: x.dtype }],
+    grid: [32 * SG_TG, Math.ceil(N / (8 * SG_TG)), 1], threadGroup: [32 * SG_TG, 1, 1],
+    templateDtypes: { T: x.dtype },
+    templateInts: { M, N, K, SG_TG, MT },
+  });
+  const out = ops.reshape(y!, [...x.shape.slice(0, -1), N]);
+  y!.dispose();
+  return out;
+}

package/src/kernels/quantization/affine3-rows.ts

@@ -0,0 +1,103 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+
+// Verify-width matmul over MLX affine 3-bit, group-64 weights: y = x · Wᵀ for
+// 2..8 activation rows. MLX's qmv runs each row as its own threadgroups, so
+// every weight is fetched and decoded M times. Here each simdgroup owns NR = 4
+// output rows; per 512-value K step, every lane decodes its VPT = 16 values of
+// each output row once (exact floats via the 2^23 mantissa trick; MLX packs
+// eight 3-bit values little-endian into three bytes) and reuses them for all M
+// rows held in registers.
+// SGS = 2 simdgroups per threadgroup. Same algebra as MLX:
+// Σ_g scale·Σ q·x + bias·Σ x. Summation order differs from MLX's qmv (not
+// bit-identical; gate by KL).
+const NR = 4, SGS = 2, VPT = 16;
+
+const SOURCE = 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)SGS + sg) * (uint)NR;
+  if (n0 >= (uint)N) return;
+  const ulong rowBytes = (ulong)K * 3u / 8u;
+  const uint groups = (uint)K / 64u;
+  const device uint8_t* wb = (const device uint8_t*)w;
+  float acc[NR][M];
+  #pragma clang loop unroll(full)
+  for (uint r = 0; r < (uint)NR; ++r)
+    #pragma clang loop unroll(full)
+    for (uint m = 0; m < (uint)M; ++m) acc[r][m] = 0.0f;
+  for (uint k0 = lane * (uint)VPT; k0 < (uint)K; k0 += 32u * (uint)VPT) {
+    float xr[M][VPT];
+    float xs[M];
+    #pragma clang loop unroll(full)
+    for (uint m = 0; m < (uint)M; ++m) {
+      const device metal::vec<T, 4>* xp = (const device metal::vec<T, 4>*)(x + (ulong)m * (ulong)K + k0);
+      float sum = 0.0f;
+      #pragma clang loop unroll(full)
+      for (uint v = 0; v < (uint)VPT / 4u; ++v) {
+        const float4 f = float4(xp[v]);
+        xr[m][4 * v] = f.x; xr[m][4 * v + 1] = f.y; xr[m][4 * v + 2] = f.z; xr[m][4 * v + 3] = f.w;
+        sum += (f.x + f.y) + (f.z + f.w);
+      }
+      xs[m] = sum;
+    }
+    const uint g = k0 / 64u;
+    #pragma clang loop unroll(full)
+    for (uint r = 0; r < (uint)NR; ++r) {
+      const uint n = n0 + r;
+      const device uint8_t* p = wb + (ulong)n * rowBytes + (ulong)k0 * 3u / 8u;
+      float q[VPT];
+      #pragma clang loop unroll(full)
+      for (uint c = 0; c < (uint)VPT / 8u; ++c) {
+        const uint packed = uint(p[3 * c]) | (uint(p[3 * c + 1]) << 8) | (uint(p[3 * c + 2]) << 16);
+        #pragma clang loop unroll(full)
+        for (uint j = 0; j < 8u; ++j)
+          q[8 * c + j] = as_type<float>(0x4b000000u | ((packed >> (3u * j)) & 7u)) - 8388608.0f;
+      }
+      const float s = float(scales[(ulong)n * groups + g]);
+      const float b = float(biases[(ulong)n * groups + g]);
+      #pragma clang loop unroll(full)
+      for (uint m = 0; m < (uint)M; ++m) {
+        float dot = 0.0f;
+        #pragma clang loop unroll(full)
+        for (uint j = 0; j < (uint)VPT; ++j) dot = metal::fma(q[j], xr[m][j], dot);
+        acc[r][m] = metal::fma(s, dot, metal::fma(b, xs[m], acc[r][m]));
+      }
+    }
+  }
+  #pragma clang loop unroll(full)
+  for (uint r = 0; r < (uint)NR; ++r)
+    #pragma clang loop unroll(full)
+    for (uint m = 0; m < (uint)M; ++m) {
+      const float v = metal::simd_sum(acc[r][m]);
+      if (lane == 0u) y[(ulong)m * (ulong)N + n0 + r] = T(v);
+    }
+`;
+
+let kernel: MetalKernel | undefined;
+
+/** Eligible: affine 3-bit, group 64, K a multiple of 512, N of 4, 2..8 rows. */
+export function affine3RowsEligible(rows: number, w: MlxArray, scales: MlxArray, spec: { bits: number; groupSize: number; mode: string }): boolean {
+  if (spec.mode !== "affine" || spec.groupSize !== 64 || spec.bits !== 3) return false;
+  const N = scales.shape[0]!, K = scales.shape[1]! * 64;
+  return rows >= 2 && rows <= 8 && K % (32 * VPT) === 0 && N % NR === 0 && w.shape.length === 2;
+}
+
+/** x [..., K] (2..8 rows) · Wᵀ → [..., N] in x's dtype. Borrows inputs. */
+export function affine3Rows(x: MlxArray, w: MlxArray, scales: MlxArray, biases: MlxArray): MlxArray {
+  const K = x.shape.at(-1)!, M = x.size / K, N = scales.shape[0]!;
+  if (M < 2 || M > 8) throw new Error(`affine3Rows: M=${M} outside 2..8`);
+  kernel ??= new MetalKernel({ name: "mlx_bun_affine3_rows", inputNames: ["x", "w", "scales", "biases"],
+    outputNames: ["y"], source: SOURCE, ensureRowContiguous: true });
+  using x2 = ops.reshape(x, [M, K]);
+  const [y] = kernel.apply([x2, w, scales, biases], {
+    outputs: [{ shape: [M, N], dtype: x.dtype }],
+    grid: [32 * SGS, Math.ceil(N / (NR * SGS)), 1], threadGroup: [32 * SGS, 1, 1],
+    templateDtypes: { T: x.dtype },
+    templateInts: { M, N, K, NR, SGS, VPT },
+  });
+  const out = ops.reshape(y!, [...x.shape.slice(0, -1), N]);
+  y!.dispose();
+  return out;
+}

package/src/kernels/quantization/affine4-mma.ts

@@ -0,0 +1,113 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+
+// Verify-width matmul over MLX affine 4-bit, group-64 weights on the simdgroup
+// matrix unit: y = x · Wᵀ for 1..8 activation rows (one block of 8). One
+// simdgroup owns 8 output 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
+// dequantizes 64 contiguous values of its row (exactly one quantization group,
+// MLX's LSB-first bit stream) into its A-fragment elements, and reads two
+// activation rows contiguously. Per-weight work is independent of M. Summation
+// order differs from MLX's qmv (not bit-identical; gate by KL).
+const SG_TG = 4, MT = 1;
+
+const SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint n0 = (threadgroup_position_in_grid.y * (uint)SG_TG + simdgroup_index_in_threadgroup) * 8u;
+  if (n0 >= (uint)N) return;
+  const uint fm = ((lane / 4u) & 4u) + ((lane / 2u) % 4u);
+  const uint fn = ((lane / 4u) & 2u) * 2u + (lane % 2u) * 2u;
+  const uint n = n0 + fm;
+  const uint groups = (uint)K / 64u;
+  const device uint32_t* wRow = (const device uint32_t*)((const device uint8_t*)w + (ulong)n * (ulong)K / 2u);
+  // 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];
+  metal::simdgroup_float8x8 acc[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)K + fm * 32u;
+    x1[bk] = x + (ulong)(m1[bk] ? r0 + 1u : 0u) * (ulong)K + fm * 32u;
+    acc[bk] = metal::make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
+  }
+  for (uint base = 0; base < (uint)K; base += 256u) {
+    const uint g = base / 64u + fn / 2u;
+    const float s = float(scales[(ulong)n * groups + g]);
+    const float b = float(biases[(ulong)n * groups + g]);
+    // 64 values of one group: 8 words.
+    uint q[8];
+    const device uint32_t* qp = wRow + g * 8u;
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < 8u; ++i) q[i] = qp[i];
+    #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 c = 0; c < 4u; ++c) {
+        const uint t = s4 * 4u + c;
+        metal::simdgroup_float8x8 a;
+        a.thread_elements()[0] = metal::fma(s, affine_q4(q, t), b);
+        a.thread_elements()[1] = metal::fma(s, affine_q4(q, t + 32u), b);
+        #pragma clang loop unroll(full)
+        for (uint bk = 0; bk < (uint)MT; ++bk) {
+          metal::simdgroup_float8x8 bx;
+          bx.thread_elements()[0] = va[bk][c];
+          bx.thread_elements()[1] = vb[bk][c];
+          metal::simdgroup_multiply_accumulate(acc[bk], a, bx, acc[bk]);
+        }
+      }
+    }
+  }
+  #pragma clang loop unroll(full)
+  for (uint bk = 0; bk < (uint)MT; ++bk) {
+    const uint r0 = mBase + 8u * bk + fn;
+    if (m0[bk]) y[(ulong)r0 * (ulong)N + n] = T(acc[bk].thread_elements()[0]);
+    if (m1[bk]) y[(ulong)(r0 + 1u) * (ulong)N + n] = T(acc[bk].thread_elements()[1]);
+  }
+`;
+
+// Value t of a 64-value group as an exact f32 (2^23 mantissa trick). 4-bit
+// values never straddle a word.
+const HEADER = String.raw`
+inline float affine_q4(thread const uint* q, uint t) {
+  const uint bit = t * 4u;
+  return as_type<float>(0x4b000000u | ((q[bit >> 5] >> (bit & 31u)) & 15u)) - 8388608.0f;
+}
+`;
+
+let kernel: MetalKernel | undefined;
+
+/** Eligible: affine 4-bit, group 64, K a multiple of 256, N of 8, 1..8 rows. */
+export function affine4MmaEligible(rows: number, w: MlxArray, scales: MlxArray, spec: { bits: number; groupSize: number; mode: string }): boolean {
+  if (spec.mode !== "affine" || spec.groupSize !== 64 || spec.bits !== 4) return false;
+  const N = scales.shape[0]!, K = scales.shape[1]! * 64;
+  return rows >= 1 && rows <= 8 && K % 256 === 0 && N % 8 === 0 && w.shape.length === 2;
+}
+
+/** x [..., K] (1..8 rows) · Wᵀ → [..., N] in x's dtype. Borrows inputs. */
+export function affine4Mma(x: MlxArray, w: MlxArray, scales: MlxArray, biases: MlxArray): MlxArray {
+  const K = x.shape.at(-1)!, M = x.size / K, N = scales.shape[0]!;
+  if (M < 1 || M > 8) throw new Error(`affine4Mma: M=${M} outside 1..8`);
+  kernel ??= new MetalKernel({ name: "mlx_bun_affine4_mma", inputNames: ["x", "w", "scales", "biases"],
+    outputNames: ["y"], source: SOURCE, header: HEADER, ensureRowContiguous: true });
+  using x2 = ops.reshape(x, [M, K]);
+  const [y] = kernel.apply([x2, w, scales, biases], {
+    outputs: [{ shape: [M, N], dtype: x.dtype }],
+    grid: [32 * SG_TG, Math.ceil(N / (8 * SG_TG)), 1], threadGroup: [32 * SG_TG, 1, 1],
+    templateDtypes: { T: x.dtype },
+    templateInts: { M, N, K, SG_TG, MT },
+  });
+  const out = ops.reshape(y!, [...x.shape.slice(0, -1), N]);
+  y!.dispose();
+  return out;
+}

package/src/kernels/quantization/affine4-rows.ts

@@ -0,0 +1,102 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+
+// Verify-width matmul over MLX affine 4-bit, group-64 weights: y = x · Wᵀ for
+// 2..8 activation rows. MLX's qmv runs each row as its own threadgroups, so
+// every weight is fetched and decoded M times. Here each simdgroup owns NR = 4
+// output rows; per 512-value K step, every lane decodes its VPT = 16 values of
+// each output row once (exact floats via the 2^23 mantissa trick, MLX's
+// little-endian packing) and reuses them for all M rows held in registers.
+// SGS = 2 simdgroups per threadgroup. Same algebra as MLX:
+// Σ_g scale·Σ q·x + bias·Σ x. Summation order differs from MLX's qmv (not
+// bit-identical; gate by KL).
+const NR = 4, SGS = 2, VPT = 16;
+
+const SOURCE = 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)SGS + sg) * (uint)NR;
+  if (n0 >= (uint)N) return;
+  const ulong rowBytes = (ulong)K / 2u;
+  const uint groups = (uint)K / 64u;
+  const device uint8_t* wb = (const device uint8_t*)w;
+  float acc[NR][M];
+  #pragma clang loop unroll(full)
+  for (uint r = 0; r < (uint)NR; ++r)
+    #pragma clang loop unroll(full)
+    for (uint m = 0; m < (uint)M; ++m) acc[r][m] = 0.0f;
+  for (uint k0 = lane * (uint)VPT; k0 < (uint)K; k0 += 32u * (uint)VPT) {
+    float xr[M][VPT];
+    float xs[M];
+    #pragma clang loop unroll(full)
+    for (uint m = 0; m < (uint)M; ++m) {
+      const device metal::vec<T, 4>* xp = (const device metal::vec<T, 4>*)(x + (ulong)m * (ulong)K + k0);
+      float sum = 0.0f;
+      #pragma clang loop unroll(full)
+      for (uint v = 0; v < (uint)VPT / 4u; ++v) {
+        const float4 f = float4(xp[v]);
+        xr[m][4 * v] = f.x; xr[m][4 * v + 1] = f.y; xr[m][4 * v + 2] = f.z; xr[m][4 * v + 3] = f.w;
+        sum += (f.x + f.y) + (f.z + f.w);
+      }
+      xs[m] = sum;
+    }
+    const uint g = k0 / 64u;
+    #pragma clang loop unroll(full)
+    for (uint r = 0; r < (uint)NR; ++r) {
+      const uint n = n0 + r;
+      const device uint8_t* p = wb + (ulong)n * rowBytes + (ulong)k0 / 2u;
+      float q[VPT];
+      #pragma clang loop unroll(full)
+      for (uint c = 0; c < (uint)VPT / 8u; ++c) {
+        const uint packed = ((const device uint32_t*)p)[c];
+        #pragma clang loop unroll(full)
+        for (uint j = 0; j < 8u; ++j)
+          q[8 * c + j] = as_type<float>(0x4b000000u | ((packed >> (4u * j)) & 15u)) - 8388608.0f;
+      }
+      const float s = float(scales[(ulong)n * groups + g]);
+      const float b = float(biases[(ulong)n * groups + g]);
+      #pragma clang loop unroll(full)
+      for (uint m = 0; m < (uint)M; ++m) {
+        float dot = 0.0f;
+        #pragma clang loop unroll(full)
+        for (uint j = 0; j < (uint)VPT; ++j) dot = metal::fma(q[j], xr[m][j], dot);
+        acc[r][m] = metal::fma(s, dot, metal::fma(b, xs[m], acc[r][m]));
+      }
+    }
+  }
+  #pragma clang loop unroll(full)
+  for (uint r = 0; r < (uint)NR; ++r)
+    #pragma clang loop unroll(full)
+    for (uint m = 0; m < (uint)M; ++m) {
+      const float v = metal::simd_sum(acc[r][m]);
+      if (lane == 0u) y[(ulong)m * (ulong)N + n0 + r] = T(v);
+    }
+`;
+
+let kernel: MetalKernel | undefined;
+
+/** Eligible: affine 4-bit, group 64, K a multiple of 512, N of 4, 2..8 rows. */
+export function affine4RowsEligible(rows: number, w: MlxArray, scales: MlxArray, spec: { bits: number; groupSize: number; mode: string }): boolean {
+  if (spec.mode !== "affine" || spec.groupSize !== 64 || spec.bits !== 4) return false;
+  const N = scales.shape[0]!, K = scales.shape[1]! * 64;
+  return rows >= 2 && rows <= 8 && K % (32 * VPT) === 0 && N % NR === 0 && w.shape.length === 2;
+}
+
+/** x [..., K] (2..8 rows) · Wᵀ → [..., N] in x's dtype. Borrows inputs. */
+export function affine4Rows(x: MlxArray, w: MlxArray, scales: MlxArray, biases: MlxArray): MlxArray {
+  const K = x.shape.at(-1)!, M = x.size / K, N = scales.shape[0]!;
+  if (M < 2 || M > 8) throw new Error(`affine4Rows: M=${M} outside 2..8`);
+  kernel ??= new MetalKernel({ name: "mlx_bun_affine4_rows", inputNames: ["x", "w", "scales", "biases"],
+    outputNames: ["y"], source: SOURCE, ensureRowContiguous: true });
+  using x2 = ops.reshape(x, [M, K]);
+  const [y] = kernel.apply([x2, w, scales, biases], {
+    outputs: [{ shape: [M, N], dtype: x.dtype }],
+    grid: [32 * SGS, Math.ceil(N / (NR * SGS)), 1], threadGroup: [32 * SGS, 1, 1],
+    templateDtypes: { T: x.dtype },
+    templateInts: { M, N, K, NR, SGS, VPT },
+  });
+  const out = ops.reshape(y!, [...x.shape.slice(0, -1), N]);
+  y!.dispose();
+  return out;
+}

package/src/kernels/residual-rms-norm.ts

@@ -0,0 +1,65 @@
+// Copyright © 2024 Apple Inc. RMS reduction adapted under the MIT license;
+// see THIRD_PARTY_NOTICES.md.
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { MetalKernel } from "@mlx-bun/mlx/metal-kernel";
+
+// Reduction and rounding follow MLX 0.32.2's rms_norm.metal (Apple Inc., MIT).
+// Round the residual sum before squaring, and the normalized value before
+// multiplying by the weight, just as add followed by fast::rms_norm does.
+const SOURCE = String.raw`
+  const uint row = threadgroup_position_in_grid.x;
+  const uint tid = thread_position_in_threadgroup.x;
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  threadgroup float sums[32];
+  threadgroup float inv[1];
+  float acc = 0.0f;
+  for (uint r = 0; r < uint(D); r += uint(THREADS) * 4u) {
+    for (uint i = 0; i < 4u; ++i) {
+      const uint col = r + tid * 4u + i;
+      if (col < uint(D)) {
+        const uint index = row * uint(D) + col;
+        const T value = T(float(x[index]) + float(update[index]));
+        residual[index] = value;
+        const float v = float(value);
+        acc += v * v;
+      }
+    }
+  }
+  acc = metal::simd_sum(acc);
+  if (sg == 0) sums[lane] = 0.0f;
+  threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+  if (lane == 0) sums[sg] = acc;
+  threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+  if (sg == 0) {
+    acc = metal::simd_sum(sums[lane]);
+    if (lane == 0) inv[0] = metal::precise::rsqrt(acc / float(D) + as_type<float>(uint(EPS)));
+  }
+  threadgroup_barrier(metal::mem_flags::mem_threadgroup | metal::mem_flags::mem_device);
+  for (uint r = 0; r < uint(D); r += uint(THREADS) * 4u) {
+    for (uint i = 0; i < 4u; ++i) {
+      const uint col = r + tid * 4u + i;
+      if (col < uint(D)) {
+        const uint index = row * uint(D) + col;
+        normalized[index] = weight[col] * T(float(residual[index]) * inv[0]);
+      }
+    }
+  }
+`;
+
+let kernel: MetalKernel | undefined;
+
+/** Equal-shaped, row-contiguous inputs and a same-dtype vector weight. Returns
+ * two owned arrays: the rounded residual and its weighted RMS normalization. */
+export function residualRmsNorm(x: MlxArray, update: MlxArray, weight: MlxArray, eps: number): [MlxArray, MlxArray] {
+  const shape = x.shape, width = shape.at(-1)!;
+  const threads = Math.min(1024, Math.ceil(width / 128) * 32);
+  const epsBits = new Uint32Array(new Float32Array([eps]).buffer)[0]!;
+  kernel ??= new MetalKernel({ name: "mlx_bun_residual_rms_norm", inputNames: ["x", "update", "weight"],
+    outputNames: ["residual", "normalized"], source: SOURCE });
+  return kernel.apply([x, update, weight], {
+    outputs: [{ shape, dtype: x.dtype }, { shape, dtype: x.dtype }],
+    grid: [x.size / width * threads, 1, 1], threadGroup: [threads, 1, 1],
+    templateDtypes: { T: x.dtype }, templateInts: { D: width, THREADS: threads, EPS: epsBits },
+  }) as [MlxArray, MlxArray];
+}

package/src/kernels/sampling/index.ts

@@ -0,0 +1,2 @@
+export * from "./normalized-argmax";
+export * from "./token-bitmask";