Add initial tuning for M5 pro and max (#3211)

Author: jagrit06

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

@@ -22,6 +22,10 @@
 
 #define instantiate_gemm_shapes_helper(iname, itype, oname, otype) \
     instantiate_gemm_transpose_helper(iname, itype, oname, otype,  64,  64, 256, 2, 2) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype,  64, 128,  64, 2, 4) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype,  64, 128, 256, 2, 4) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 128, 128,  64, 4, 4) \
+    instantiate_gemm_transpose_helper(iname, itype, oname, otype, 128, 128, 256, 4, 4) \
     instantiate_gemm_transpose_helper(iname, itype, oname, otype, 128, 128, 512, 4, 4)
 
 instantiate_gemm_shapes_helper(float16, half, float16, half);

mlx/backend/metal/matmul.cpp

@@ -204,6 +204,15 @@ void steel_matmul_regular_axpby_nax(
   int bm = 128, bn = 128, bk = 512;
   int wm = 4, wn = 4;
 
+  // Temp routing for larger devices
+  char devc = d.get_architecture().back();
+  if (devc == 's' || devc == 'c' || devc == 'd') {
+    bk = (K >= 8192 && K > (M + N)) ? 64 : 256;
+
+    bm = 64;
+    wm = 2;
+  }
+
   // Prepare kernel name
   std::ostringstream kname;
 
@@ -268,6 +277,9 @@ void steel_matmul_regular_axpby_nax(
 
   // TODO: Explore device-based tuning for swizzle
   int swizzle_log = tm <= 3 ? 0 : 1;
+  if (devc == 's' || devc == 'c' || devc == 'd') {
+    swizzle_log = 2;
+  }
 
   // Prepare steel matmul params
   GEMMParams params{/* const int M = */ M,