skinny_mm: simdgroup-matrix MMA kernel for big-weight decode matmuls (M 5–32)

.github/workflows/ci.yml

@@ -1,11 +1,18 @@
 name: CI
 
-# Build the engine and run the test suite on every PR targeting master or develop.
+# Build the engine and run the test suite on every PR targeting master or develop,
+# and on every push to master. The master runs exist for the dependency cache:
+# GitHub scopes caches saved by a PR branch to that branch alone, so without a
+# master-saved cache every new PR starts with a cold MLX build (~1.5-2h on the
+# small runner) — which is exactly what the run history showed before this.
 # mlxforge is Apple Silicon only (MLX Metal backend), so this runs on GitHub's
 # ARM64 macOS runners. The integration tests self-skip when the model isn't
 # present, so we don't download weights here — only the pure-logic units and the
 # golden-reference fixtures (committed under reference/fixtures/) are exercised.
 on:
+  push:
+    branches:
+      - master
   pull_request:
     branches:
       - master
@@ -21,11 +28,11 @@ jobs:
     name: Build & test (Apple Silicon)
     runs-on: macos-15  # ARM64 Apple Silicon runner — required for the Metal backend
     # The GitHub macos-15 arm64 runner is only 3 vCPU / 7GB, so a *cold* build
-    # (compiling MLX's Metal kernels from scratch) can take ~an hour. The cache
-    # below makes every later run fast, but the cache is only written when a run
-    # finishes — so the first run must be allowed to complete. Hence the generous
-    # timeout; subsequent cache-warm runs finish in a few minutes.
-    timeout-minutes: 120
+    # (compiling MLX's Metal kernels from scratch) takes 1.5-2h — one run hit the
+    # old 2h limit on a slow runner. The cache below makes cache-warm runs finish
+    # in minutes, but the cache is only written when a run completes, so cold
+    # runs must be allowed to finish. Hence the generous timeout.
+    timeout-minutes: 200
     steps:
       - name: Checkout
         uses: actions/checkout@v4

doc/architecture.md

@@ -165,11 +165,17 @@ In `Worker::decode_step()`:
 weight-bandwidth-bound, and MLX's tiled GEMM drops to ~1/3 of GEMV bandwidth
 the moment the batch reaches 2 rows (ml-explore/mlx#3661) — historically a
 ~2.6× per-row decode cliff between B=1 and B=2. `model/skinny_matmul` provides
-custom `fast::metal_kernel` kernels for the B∈[2,16] dense-fp16 decode shape:
-each simdgroup reads a weight row once and keeps the batch's activations as
-register accumulators (a one-column variant for B≤4 at ~GEMV bandwidth, a
-two-column variant for 5–16). Past B=16 the tiled GEMM wins and `linear()`
-falls back. Accumulation is fp32 in a different order than `mx::matmul`, so
+custom `fast::metal_kernel` kernels for the dense-fp16 decode shape, picked by
+(B, weight size): a one-column-per-simdgroup scalar variant for B≤4 (~GEMV
+bandwidth), a two-column scalar variant for B∈[5,16] on the small per-layer
+weights (barrier-free simdgroups tolerate the latency of short back-to-back
+ops best), and a simdgroup-matrix MMA variant for B∈[5,32] on big weights —
+in practice the vocab head — using hardware `simdgroup_half8x8`
+multiply-accumulates over 8-output-column tiles, with the activation chunk
+staged once per threadgroup in threadgroup memory and each weight element
+streamed from device exactly once grid-wide (~134 GB/s at B=8, still ~68 at
+B=32, vs the GEMM's flat ~55). Anywhere else the tiled GEMM wins and
+`linear()` falls back. Accumulation is fp32 in a different order than `mx::matmul`, so
 logits differ at fp16-noise scale; the gate is row-for-row token equality of a
 kernel-on batch against the stock-matmul batch
 (`tests/scheduler/worker_test.cpp`), plus a pure-kernel `allclose` grid

src/model/skinny_matmul.cpp

@@ -40,8 +40,9 @@ constexpr const char* kSourceOneCol = R"(
 
 // Two output columns per simdgroup: each activation load feeds two weight
 // rows, halving the redundant x traffic that degrades the one-column variant
-// past M ~ 8. Best for M in [5, 16]; the crossover vs the tiled GEMM is past
-// 16 (66 GB/s at M=16 vs the GEMM's flat ~56 on M1 Pro).
+// past M ~ 8. Best for M in [5, 16] on the small per-layer weights, where its
+// barrier-free independent simdgroups ride out the latency of short
+// back-to-back ops better than the MMA tile kernel below.
 constexpr const char* kSourceTwoCol = R"(
     const uint lane = thread_position_in_grid.x;   // 0..31
     const uint pair = thread_position_in_grid.y;   // output column pair
@@ -78,8 +79,82 @@ constexpr const char* kSourceTwoCol = R"(
     }
 )";
 
+// simdgroup-matrix MMA tile kernel for M in [5, 32] on *large* weights (the
+// vocab head): each simdgroup owns an 8-output-column tile across the full D,
+// accumulating y^T = w_tile * x_tile^T in hardware simdgroup_float8x8 ops, so
+// the per-element FMA/issue cost that throttles the scalar kernels past M ~ 4
+// disappears. Memory hierarchy:
+//  - weights stream from device exactly once grid-wide (plain, non-transposed
+//    simdgroup_loads; the transpose lands on x instead);
+//  - the x chunk is staged in threadgroup memory once per threadgroup and the
+//    transposed tile loads hit that on-chip copy, not device/L1 — SG=8
+//    simdgroups (64 output columns) share each staged chunk;
+//  - the staging loop zero-fills rows past M, so no host-side padding of x.
+// Tail tiles clamp o0 to O-8 and overlap-recompute (duplicate stores of
+// identical values are benign); applies() guarantees O >= 8 and CHUNK | D.
+constexpr const char* kSourceMma = R"(
+    const uint lane   = thread_position_in_threadgroup.x;  // 0..31
+    const uint sg     = thread_position_in_threadgroup.y;  // 0..SG-1
+    const uint o_tile = thread_position_in_grid.y;
+    const int  D = w_shape[1];
+    const int  O = w_shape[0];
+    const size_t o0 = (size_t)min((int)(o_tile * 8), O - 8);
+
+    threadgroup half xs[MT * 8 * CHUNK];
+    simdgroup_float8x8 acc[MT];
+    for (int t = 0; t < MT; ++t)
+        acc[t] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
+    simdgroup_half8x8 xt, wt[UN];
+    const uint tid = sg * 32 + lane;
+    for (int kc = 0; kc < D; kc += CHUNK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        for (uint i = tid; i < (uint)(MT * 8 * CHUNK / 4); i += 32 * SG) {
+            uint m = i / (CHUNK / 4), j = i % (CHUNK / 4);
+            ((threadgroup half4*)xs)[i] = (m < (uint)M)
+                ? ((const device half4*)(x + (size_t)m * D + kc))[j]
+                : half4(0.0h);
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        for (int k = 0; k < CHUNK; k += 8 * UN) {
+            for (int u = 0; u < UN; ++u)
+                simdgroup_load(wt[u], w + o0 * D + kc + k + u * 8, D);
+            for (int t = 0; t < MT; ++t) {
+                for (int u = 0; u < UN; ++u) {
+                    simdgroup_load(xt, (threadgroup half*)xs +
+                                   (size_t)(t * 8) * CHUNK + k + u * 8, CHUNK, 0, true);
+                    simdgroup_multiply_accumulate(acc[t], wt[u], xt, acc[t]);
+                }
+            }
+        }
+    }
+    threadgroup float scratch[SG][64];
+    for (int t = 0; t < MT; ++t) {
+        simdgroup_store(acc[t], scratch[sg], 8);
+        simdgroup_barrier(mem_flags::mem_threadgroup);
+        for (uint i = lane; i < 64; i += 32) {
+            int m = t * 8 + (int)(i % 8);
+            if (m < M) y[(size_t)m * O + o0 + (i / 8)] = (half)scratch[sg][i];
+        }
+        simdgroup_barrier(mem_flags::mem_threadgroup);
+    }
+)";
+
 constexpr int kOneColMaxM = 4;
-constexpr int kMaxM = 16;
+constexpr int kTwoColMaxM = 16;
+constexpr int kMmaMaxM = 32;
+constexpr int kSimdgroupsPerTg = 8;  // SG: o-tiles sharing one staged x chunk
+constexpr int kUnroll = 4;           // UN: weight tiles in flight per k step
+
+// The MMA kernel only beats the alternatives on big single matmuls (the vocab
+// head, where it sustains GEMV-class bandwidth out to M=32). On the small
+// per-layer weights the engine runs as short dependent back-to-back ops, and
+// there the barrier-free scalar kernels (M <= 16) or the stock GEMM (M > 16)
+// win — measured on chained-dependent shapes, not single dispatches.
+constexpr int64_t kMmaMinWeightElems = 32 * 1024 * 1024;  // 64 MB of fp16
+
+bool is_big_weight(const mx::array& w) {
+  return (int64_t)w.shape()[0] * w.shape()[1] >= kMmaMinWeightElems;
+}
 
 }  // namespace
 
@@ -90,24 +165,47 @@ bool skinny_matmul_applies(const mx::array& x, const mx::array& w) {
   if (nd == 3 && x.shape()[1] != 1) return false;  // decode shape only, never prefill
   if (nd != 2 && nd != 3) return false;
   const int m = x.shape()[0];
-  return m >= 2 && m <= kMaxM && x.shape()[nd - 1] == w.shape()[1];
+  if (m < 2 || x.shape()[nd - 1] != w.shape()[1]) return false;
+  if (m <= kTwoColMaxM) return true;
+  // 17..32 pays off only on big weights (and the MMA kernel needs a full
+  // 8-column tile); on small weights the stock GEMM wins — fall back.
+  return m <= kMmaMaxM && w.shape()[0] >= 8 && is_big_weight(w);
 }
 
 mx::array skinny_matmul(const mx::array& x, const mx::array& w) {
   static const auto one_col = mx::fast::metal_kernel(
       "mlxforge_gemv_multirow", {"x", "w"}, {"y"}, kSourceOneCol);
   static const auto two_col = mx::fast::metal_kernel(
       "mlxforge_gemv_multirow2", {"x", "w"}, {"y"}, kSourceTwoCol);
+  static const auto mma = mx::fast::metal_kernel(
+      "mlxforge_gemv_mma", {"x", "w"}, {"y"}, kSourceMma);
 
   const int m = x.shape()[0];
   const int o = w.shape()[0];
+  const int d = w.shape()[1];
   mx::array x2 = x.ndim() == 3 ? mx::reshape(x, {m, x.shape()[2]}) : x;
-  const bool narrow = m <= kOneColMaxM;
-  std::vector<mx::array> out = (narrow ? one_col : two_col)(
-      {x2, w}, {mx::Shape{m, o}}, {mx::float16},
-      /*grid=*/{32, narrow ? o : (o + 1) / 2, 1}, /*threadgroup=*/{32, 1, 1},
-      /*template_args=*/{{"M", m}},
-      /*init_value=*/std::nullopt, /*verbose=*/false, {});
+  std::vector<mx::array> out;
+  if (m <= kOneColMaxM) {
+    out = one_col({x2, w}, {mx::Shape{m, o}}, {mx::float16},
+                  /*grid=*/{32, o, 1}, /*threadgroup=*/{32, 1, 1},
+                  /*template_args=*/{{"M", m}},
+                  /*init_value=*/std::nullopt, /*verbose=*/false, {});
+  } else if (o >= 8 && is_big_weight(w)) {
+    const int chunk = d % 256 == 0 ? 256 : 128;
+    int tiles = (o + 7) / 8;
+    tiles = (tiles + kSimdgroupsPerTg - 1) / kSimdgroupsPerTg * kSimdgroupsPerTg;
+    out = mma({x2, w}, {mx::Shape{m, o}}, {mx::float16},
+              /*grid=*/{32, tiles, 1}, /*threadgroup=*/{32, kSimdgroupsPerTg, 1},
+              /*template_args=*/
+              {{"M", m}, {"MT", (m + 7) / 8}, {"CHUNK", chunk}, {"UN", kUnroll},
+               {"SG", kSimdgroupsPerTg}},
+              /*init_value=*/std::nullopt, /*verbose=*/false, {});
+  } else {
+    out = two_col({x2, w}, {mx::Shape{m, o}}, {mx::float16},
+                  /*grid=*/{32, (o + 1) / 2, 1}, /*threadgroup=*/{32, 1, 1},
+                  /*template_args=*/{{"M", m}},
+                  /*init_value=*/std::nullopt, /*verbose=*/false, {});
+  }
   return x.ndim() == 3 ? mx::reshape(out[0], {m, 1, o}) : out[0];
 }
 

src/model/skinny_matmul.h

@@ -3,14 +3,27 @@
 // MLX's Metal matmul drops from ~161 GB/s (GEMV, M=1) to ~56 GB/s (tiled GEMM)
 // the moment M reaches 2, and every M in [2, 64] pays the same tile cost
 // (ml-explore/mlx#3661) — exactly the continuous-batching decode shape, where
-// each per-token step is weight-bandwidth-bound. These custom
-// fast::metal_kernel kernels read each weight row once per simdgroup and keep
-// the M activation rows as register accumulators, recovering GEMV-class
-// bandwidth: a one-column-per-simdgroup variant for M in [2, 4] (~161 GB/s)
-// and a two-column variant for M in [5, 16] (the doubled arithmetic intensity
-// halves the redundant activation reads that degrade larger M; ~125 GB/s at
-// M=8, still ahead of the tiled GEMM at M=16). Beyond 16 the fallback GEMM is
-// the right path.
+// each per-token step is weight-bandwidth-bound. Three custom
+// fast::metal_kernel kernels recover GEMV-class bandwidth, dispatched by
+// (M, weight size):
+//  - M in [2, 4]: one output column per simdgroup, scalar fp32 register
+//    accumulators (~161 GB/s — bandwidth-saturated, nothing faster exists);
+//  - M in [5, 16], small weights: two columns per simdgroup — barrier-free
+//    independent simdgroups, which ride out the latency of the short
+//    back-to-back per-layer matmuls better than anything tiled;
+//  - M in [5, 32], big weights (>= 64 MB, in practice the vocab head):
+//    simdgroup-matrix MMA tiles. The scalar approach dies past M ~ 4 on
+//    instruction issue (M loads + 4M FMAs per weight half4), so this variant
+//    moves the arithmetic into hardware simdgroup_half8x8
+//    multiply-accumulates (8 output columns x full D per simdgroup, fp32
+//    accumulator tiles), stages the x chunk in threadgroup memory once per
+//    8-simdgroup threadgroup, and streams each weight element from device
+//    exactly once grid-wide: ~134 GB/s at M=8 and still ~68 at M=32 on the
+//    vocab head, vs the two-column kernel's 106/12 and the GEMM's flat ~55.
+// The split is empirical: on *chained dependent* small matmuls (the engine's
+// per-layer regime) the scalar kernels and the stock GEMM (M > 16) win, while
+// the MMA kernel wins on big single matmuls at every M in [5, 32] — single-op
+// microbenchmarks rank these kernels differently than the engine does.
 //
 // Accumulation is fp32 but in a different order than mx::matmul, so logits can
 // differ at fp16-noise scale — the same class as the decode-vs-recompute gap.
@@ -24,10 +37,12 @@ namespace mx = mlx::core;
 
 namespace mlxforge {
 
-// True when the kernel path applies to the shapes: x is (B, 1, D) or (B, D)
-// fp16 with B in [2, 16], w is a dense fp16 (O, D) weight, and D is a multiple
-// of 128 (half4 loads across 32 lanes). Enablement is the caller's flag
-// (DecoderModel::skinny_mm_); this checks shapes only.
+// True when a kernel path applies to the shapes: x is (B, 1, D) or (B, D)
+// fp16, w is a dense fp16 (O, D) weight with D a multiple of 128 (half4
+// loads / chunked staging); B in [2, 16] always qualifies, B in [17, 32] only
+// on big weights with O >= 8 (the MMA path — small weights past 16 are the
+// GEMM's). Enablement is the caller's flag (DecoderModel::skinny_mm_); this
+// checks shapes only.
 bool skinny_matmul_applies(const mx::array& x, const mx::array& w);
 
 // x @ w.T via the multi-row GEMV kernels. Preserves x's leading shape:

tests/model/skinny_matmul_test.cpp

@@ -1,7 +1,10 @@
 // The multi-row GEMV decode kernels must agree with mx::matmul at fp16-noise
-// tolerance across the whole M range and both kernel variants (one-column for
-// M <= 4, two-column for 5..16, including an odd O exercising the tail guard),
-// and the shape gate must reject everything outside the batched-decode shape.
+// tolerance across the whole M range and all three kernel variants
+// (one-column for M <= 4, two-column for 5..16 on small weights including an
+// odd O exercising its tail guard, and the simdgroup-matrix MMA path for
+// 5..32 on a big weight — M not a multiple of 8 exercises its zero-fill row
+// guard, D 128 / 1024 cover both CHUNK template paths), and the shape gate
+// must reject everything outside the batched-decode shape.
 // Pure GPU test — no model weights needed.
 #include <doctest/doctest.h>
 
@@ -50,17 +53,42 @@ TEST_CASE("skinny_matmul matches mx::matmul across M, D, and both variants") {
   }
 }
 
+TEST_CASE("skinny_matmul MMA path matches mx::matmul on a big weight, M 5..32") {
+  // 32768 x 1024 crosses the big-weight threshold (>= 32M elements), routing
+  // M in [5, 32] through the simdgroup-matrix kernel. 32775 columns make the
+  // last tile partial, exercising the clamped overlap-recompute path.
+  for (int o : {32775, 32768}) {
+    mx::array w = mx::astype(mx::random::normal({o, 1024}), mx::float16);
+    for (int m : {5, 8, 16, 23, 32}) {
+      CAPTURE(m);
+      CAPTURE(o);
+      mx::array x = mx::astype(
+          mx::multiply(mx::random::normal({m, 1024}), mx::array(0.05f)), mx::float16);
+      REQUIRE(mlxforge::skinny_matmul_applies(x, w));
+      mx::array ref = mx::matmul(x, mx::transpose(w));
+      mx::array got = mlxforge::skinny_matmul(x, w);
+      CHECK(got.shape() == ref.shape());
+      CHECK(max_abs_diff(got, ref) < 5e-3f);
+    }
+  }
+}
+
 TEST_CASE("skinny_matmul_applies rejects everything outside the decode shape") {
   mx::array w = mx::astype(mx::random::normal({64, 1024}), mx::float16);
   auto x = [&](mx::Shape s, mx::Dtype t = mx::float16) {
     return mx::astype(mx::random::normal(std::move(s)), t);
   };
   CHECK_FALSE(mlxforge::skinny_matmul_applies(x({1, 1024}), w));     // M=1: GEMV is faster
-  CHECK_FALSE(mlxforge::skinny_matmul_applies(x({17, 1024}), w));    // past the GEMM crossover
+  CHECK_FALSE(mlxforge::skinny_matmul_applies(x({17, 1024}), w));    // small w: GEMM past 16
   CHECK(mlxforge::skinny_matmul_applies(x({16, 1024}), w));
   CHECK_FALSE(mlxforge::skinny_matmul_applies(x({4, 2, 1024}), w));  // prefill (L > 1)
   CHECK(mlxforge::skinny_matmul_applies(x({4, 1, 1024}), w));
   CHECK_FALSE(mlxforge::skinny_matmul_applies(x({4, 1024}, mx::float32), w));  // dtype
   mx::array w_odd = mx::astype(mx::random::normal({64, 1000}), mx::float16);
   CHECK_FALSE(mlxforge::skinny_matmul_applies(x({4, 1000}), w_odd));  // D % 128 != 0
+  // Big weights extend the range to 32 via the MMA path — but no further.
+  mx::array w_big = mx::astype(mx::random::normal({32768, 1024}), mx::float16);
+  CHECK(mlxforge::skinny_matmul_applies(x({17, 1024}), w_big));
+  CHECK(mlxforge::skinny_matmul_applies(x({32, 1024}), w_big));
+  CHECK_FALSE(mlxforge::skinny_matmul_applies(x({33, 1024}), w_big));
 }

tests/scheduler/worker_test.cpp

@@ -120,7 +120,7 @@ TEST_CASE("skinny_mm decode kernels reproduce the stock-matmul greedy stream") {
     prompt.insert(prompt.end(), ids.begin(), ids.end());
   }
   const int kMax = 16;
-  const int kBatch = 4;  // decode at B=4 routes every linear through the kernels
+  const int kBatch = 6;  // B > 4 routes decode linears through the MMA tile kernel
 
   // Reuse may only change speed, never tokens: the kernel-on batch must match
   // the kernel-off batch row for row (both greedy on identical prompts).