Speed up NAX split-K by better tuning and routing and fix NAX addmm (#3422)

I 'll merge now and comment with more benchmarks later since this also fixes two bugs so worst case we 'll do another tuning, it isn't like we won't need the functionality of this PR.

mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.h

@@ -36,12 +36,10 @@ void gemm_epilogue(
   using CFrag = typename NAXTile_t::NAXFrag_t;
   using cfrag_t = typename CFrag::template dtype_frag_t<T>;
 
-  STEEL_PRAGMA_UNROLL
-  for (short mm = 0; mm < TM; mm++) {
-    STEEL_PRAGMA_UNROLL
-    for (short nn = 0; nn < TN; nn++) {
-      const short m = mm * CFrag::kFragRows;
-      const short n = nn * CFrag::kFragCols;
+  const_for_loop<0, TM, 1>([&](auto mm) {
+    const_for_loop<0, TN, 1>([&](auto nn) {
+      auto m = mm * Int<CFrag::kFragRows>{};
+      auto n = nn * Int<CFrag::kFragCols>{};
 
       cfrag_t celems;
 
@@ -59,7 +57,7 @@ void gemm_epilogue(
             n);
       }
 
-      auto delems = Dtile.frag_at(mm, nn);
+      thread auto& delems = Dtile.template frag_at<mm, nn>();
 
       STEEL_PRAGMA_UNROLL
       for (short i = 0; i < kElemsPerFrag; i++) {
@@ -70,8 +68,8 @@ void gemm_epilogue(
           delems[i] += static_cast<V>(celems[i]);
         }
       }
-    }
-  }
+    });
+  });
 }
 
 // clang-format off

mlx/backend/metal/kernels/steel/utils/integral_constant.h

@@ -23,10 +23,6 @@ struct integral_constant {
   METAL_FUNC constexpr operator value_type() const noexcept {
     return value;
   }
-
-  // METAL_FUNC constexpr value_type operator()() const noexcept {
-  //   return value;
-  // }
 };
 
 template <bool B>
@@ -131,4 +127,4 @@ METAL_FUNC constexpr auto sum(T x, Us... us) {
 } // namespace steel
 } // namespace mlx
 
-#pragma METAL internals : disable
+#pragma METAL internals : disable

mlx/backend/metal/matmul.cpp

@@ -363,35 +363,6 @@ void steel_matmul_regular_axpby(
     int64_t C_batch_stride /* = 0*/,
     float alpha /* = 1.0f */,
     float beta /* = 0.0f */) {
-  if (metal::is_nax_available() && !issubdtype(a.dtype(), complexfloating) &&
-      (env::enable_tf32() || a.dtype() != float32)) {
-    return steel_matmul_regular_axpby_nax<CHECK_AB>(
-        /* const Stream& s = */ s,
-        /* metal::Device& d = */ d,
-        /* const array& a = */ a,
-        /* const array& b = */ b,
-        /* const array& c = */ c,
-        /* array& out = */ out,
-        /* int M = */ M,
-        /* int N = */ N,
-        /* int K = */ K,
-        /* int batch_size_out = */ batch_size_out,
-        /* int lda = */ lda,
-        /* int ldb = */ ldb,
-        /* int ldd = */ ldd,
-        /* bool transpose_a = */ transpose_a,
-        /* bool transpose_b = */ transpose_b,
-        /* std::vector<array>& copies = */ copies,
-        /* Shape batch_shape = */ batch_shape,
-        /* Strides batch_strides = */ batch_strides,
-        /* int64_t A_batch_stride = */ A_batch_stride,
-        /* int64_t B_batch_stride = */ B_batch_stride,
-        /* int64_t matrix_stride_out = */ matrix_stride_out,
-        /* int64_t C_batch_stride = */ C_batch_stride,
-        /* float alpha = */ alpha,
-        /* float beta = */ beta);
-  }
-
   using namespace mlx::steel;
 
   // Determine dispatch kernel
@@ -704,14 +675,26 @@ void steel_gemm_splitk_axpby_nax(
     float beta = 0.0f) {
   using namespace mlx::steel;
 
-  constexpr int bm = 128, bn = 128, bk = 512;
-  constexpr int wm = 4, wn = 4;
+  int bm = 128, bn = 128, bk = 512;
+  int wm = 4, wn = 4;
+  int split_k_partition_size = 4096;
+
+  if ((M + N) / 2 < 512 || K <= 4096) {
+    bm = bn = 64;
+    bk = 256;
+    wm = wn = 2;
+  }
+  if (K <= 1024) {
+    split_k_partition_size = K / 2;
+  } else if (K <= 2048) {
+    split_k_partition_size = 1024;
+  } else if (K <= 4096) {
+    split_k_partition_size = 2048;
+  }
 
   // Determine how many partitions to split K into
-  constexpr int split_k_partition_size = 3072;
   int split_k_partitions =
       (K + split_k_partition_size - 1) / split_k_partition_size;
-
   const int bk_iters_per_partition = split_k_partition_size / bk;
   const int split_k_partition_stride = M * N;
 
@@ -928,13 +911,17 @@ void steel_matmul_axpby(
   int _tm = (M + 16 - 1) / 16;
   int _tn = (N + 16 - 1) / 16;
   int _tk = K / 16;
+  int64_t matrix_size = static_cast<int64_t>(M) * N;
+  bool use_nax = metal::is_nax_available() &&
+      !issubdtype(a.dtype(), complexfloating) &&
+      (env::enable_tf32() || a.dtype() != float32);
+  char devc = d.get_architecture().back();
+  int min_tmn_threshold = (devc == 's' || devc == 'd') ? 2048 : 1024;
 
   // Case 1: Small M×N with large K, use SIMD split-K
-  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)) {
+  if (!use_nax && 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,
@@ -956,14 +943,9 @@ void steel_matmul_axpby(
   }
 
   // Case 2: Large K with sufficient M, N, and NAX is available, use NAX split-K
-  // TODO: Add device-specific tuning for more NAX GPUs in the future
-  constexpr int min_mn_threshold = 2048 * 2048;
-  constexpr int min_k_threshold = 10240;
-  if (batch_size_out == 1 && metal::is_nax_available() &&
-      !issubdtype(a.dtype(), complexfloating) &&
-      (env::enable_tf32() || a.dtype() != float32) &&
-      int64_t(M) * N >= min_mn_threshold && K >= min_k_threshold &&
-      K >= (3 * std::max(M, N))) {
+  if (use_nax && batch_size_out == 1 &&
+      (K >= 3 * std::max(M, N) ||
+       (std::max(M, N) <= 1024 && K > 2 * std::max(M, N)))) {
     return steel_gemm_splitk_axpby_nax<CHECK_AB>(
         /* const Stream& s = */ s,
         /* metal::Device& d = */ d,
@@ -998,6 +980,34 @@ void steel_matmul_axpby(
   int64_t B_batch_stride_ = B_batch_stride.empty() ? 0 : B_batch_stride.back();
   int64_t C_batch_stride_ = C_batch_stride.empty() ? 0 : C_batch_stride.back();
 
+  if (use_nax) {
+    return steel_matmul_regular_axpby_nax<CHECK_AB>(
+        /* const Stream& s = */ s,
+        /* metal::Device& d = */ d,
+        /* const array& a = */ a,
+        /* const array& b = */ b,
+        /* const array& c = */ c,
+        /* array& out = */ out,
+        /* int M = */ M,
+        /* int N = */ N,
+        /* int K = */ K,
+        /* int batch_size_out = */ batch_size_out,
+        /* int lda = */ lda,
+        /* int ldb = */ ldb,
+        /* int ldd = */ N,
+        /* bool transpose_a = */ transpose_a,
+        /* bool transpose_b = */ transpose_b,
+        /* std::vector<array>& copies = */ copies,
+        /* Shape batch_shape = */ std::move(batch_shape),
+        /* Strides batch_strides = */ std::move(batch_strides),
+        /* int64_t A_batch_stride = */ A_batch_stride_,
+        /* int64_t B_batch_stride = */ B_batch_stride_,
+        /* int64_t matrix_stride_out = */ int64_t(M) * N,
+        /* int64_t C_batch_stride = */ C_batch_stride_,
+        /* float alpha = */ alpha,
+        /* float beta = */ beta);
+  }
+
   return steel_matmul_regular_axpby<CHECK_AB>(
       /* const Stream& s = */ s,
       /* metal::Device& d = */ d,