From b56782be52f3f6be10b3bbf3d60528f8e9982c81 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Mon, 2 Feb 2026 08:45:09 -0800 Subject: [PATCH] [Metal] Tune splitk gemm dispatch conditions and partition sizes (#3087) --- mlx/backend/metal/matmul.cpp | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/mlx/backend/metal/matmul.cpp b/mlx/backend/metal/matmul.cpp index c03f6e86..98d7b668 100644 --- a/mlx/backend/metal/matmul.cpp +++ b/mlx/backend/metal/matmul.cpp @@ -535,6 +535,8 @@ void steel_gemm_splitk_axpby( float beta = 0.0f) { using namespace mlx::steel; + int _tm = (M + 32 - 1) / 32; + int _tn = (N + 32 - 1) / 32; int _tk = K / 16; int bm = M < 40 ? 16 : 32; @@ -542,7 +544,9 @@ void steel_gemm_splitk_axpby( int bk = 16; int wm = 2, wn = 2; - int split_k_partitions = _tk < 16 ? 2 : (_tk < 32 ? 4 : (_tk < 64 ? 8 : 16)); + // As _tk grows use more partitions, as _tm * _tn grow use fewer partitions + int split_k_partitions = + std::min(std::max(2, next_power_of_2(_tk / (_tm * _tn))), 32); int split_k_partition_stride = M * N; int gemm_k_iterations = (K / bk) / split_k_partitions; int split_k_partition_size = gemm_k_iterations * bk; @@ -909,12 +913,16 @@ void steel_matmul_axpby( ///////////////////////////////////////////////////////////////////////////// // Split K specialization - int _tm = M / 16; - int _tn = N / 16; + int _tm = (M + 16 - 1) / 16; + int _tn = (N + 16 - 1) / 16; int _tk = K / 16; // Case 1: Small M×N with large K, use SIMD split-K - if (batch_size_out == 1 && (_tm * _tn) <= 32 && _tk >= 8) { + char devc = d.get_architecture().back(); + // Max and Ultra dispatch larger sizes to splitk + int min_tmn_threshold = (devc == 's' || devc == 'd') ? 2048 : 1024; + if (batch_size_out == 1 && (_tm * _tn) <= min_tmn_threshold && _tk >= 8 && + K >= std::max(M, N)) { return steel_gemm_splitk_axpby( /* const Stream& s = */ s, /* metal::Device& d = */ d,