Better support consumer CUDA GPUs (#3056)

mlx/backend/cuda/device.cpp

@@ -210,6 +210,10 @@ std::pair<int, int> get_graph_limits(Device& d) {
       ops = 50;
       mb = 500;
       break;
+    case 1200: // Consumer Blackwell
+      ops = 100;
+      mb = 1000;
+      break;
     case 1210: // DGX Spark
       ops = 20;
       mb = 25;

mlx/backend/cuda/device.h

@@ -119,7 +119,7 @@ class CommandEncoder {
   CudaStream stream_;
   CudaGraph graph_;
   Worker worker_;
-  char node_count_{0};
+  int node_count_{0};
   bool in_concurrent_{false};
   std::vector<cudaGraphNode_t> from_nodes_;
   std::vector<cudaGraphNode_t> to_nodes_;

mlx/backend/cuda/jit_module.cpp

@@ -278,7 +278,7 @@ void compile(
   std::vector<const char*> args;
   bool use_sass = compiler_supports_device_sass(device);
   auto cc = device.compute_capability_major();
-  std::string arch_tag = (cc == 90 || cc == 100 || cc == 121) ? "a" : "";
+  std::string arch_tag = (cc >= 9) ? "a" : "";
   std::string compute = fmt::format(
       "--gpu-architecture={}_{}{}{}",
       use_sass ? "sm" : "compute",