diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 62a9ea3..dfa63b1 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/doc/architecture.md b/doc/architecture.md index 4b345d0..b1d8c5b 100644 --- a/doc/architecture.md +++ b/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 diff --git a/src/model/skinny_matmul.cpp b/src/model/skinny_matmul.cpp index 9d3eef4..c76ecdb 100644 --- a/src/model/skinny_matmul.cpp +++ b/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(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,7 +165,11 @@ 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) { @@ -98,16 +177,35 @@ mx::array skinny_matmul(const mx::array& x, const mx::array& w) { "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 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 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]; } diff --git a/src/model/skinny_matmul.h b/src/model/skinny_matmul.h index dcb000e..b485697 100644 --- a/src/model/skinny_matmul.h +++ b/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: diff --git a/tests/model/skinny_matmul_test.cpp b/tests/model/skinny_matmul_test.cpp index 182d2f9..333f813 100644 --- a/tests/model/skinny_matmul_test.cpp +++ b/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 @@ -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)); } diff --git a/tests/scheduler/worker_test.cpp b/tests/scheduler/worker_test.cpp index b431fd2..7540df1 100644 --- a/tests/scheduler/worker_test.cpp +++ b/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).