L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Metal] Tune splitk gemm dispatch conditions and partition sizes (#3087)

Browse Source
This commit is contained in:
Awni Hannun
2026-02-02 08:45:09 -08:00
committed by GitHub
parent 8ef539522c
commit b56782be52
1 changed files with 12 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx/backend/metal/matmul.cpp
+12 -4
View File
@@ -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<CHECK_AB>(
/* const Stream& s = */ s,
/* metal::Device& d = */ d,