Fix 2pass sdpa on < M2 (#3099)

mlx/backend/metal/kernels/sdpa_vector.h

@@ -323,6 +323,7 @@ template <typename T, int D>
     const device float* sums [[buffer(1)]],
     const device float* maxs [[buffer(2)]],
     device T* out [[buffer(3)]],
+    const constant int& blocks [[buffer(4)]],
     uint3 tid [[threadgroup_position_in_grid]],
     uint3 tpg [[threadgroups_per_grid]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
@@ -368,7 +369,7 @@ template <typename T, int D>
 
     // Update the output accumulator
     for (int i = 0; i < elem_per_thread; i++) {
-      o[i] += factor * partials[i];
+      o[i] += factor * static_cast<U>(partials[i]);
     }
     maxs += BN;
     sums += BN;

mlx/backend/metal/scaled_dot_product_attention.cpp

@@ -523,6 +523,7 @@ void sdpa_vector_2pass(
   // Get the kernel
   auto& compute_encoder = d.get_command_encoder(s.index);
   auto kernel = d.get_kernel(kname, hash_name, func_consts);
+  check_kernel_threadgroup_size(kernel, group_dims, hash_name);
 
   compute_encoder.set_compute_pipeline_state(kernel);
 
@@ -564,13 +565,8 @@ void sdpa_vector_2pass(
   kname += "_";
   kname += std::to_string(v.shape(-1));
 
-  func_consts = {
-      {&blocks, MTL::DataType::DataTypeInt, 26},
-  };
-  hash_name = kname + "_" + std::to_string(blocks);
-
   // Get the kernel
-  kernel = d.get_kernel(kname, hash_name, func_consts);
+  kernel = d.get_kernel(kname);
   compute_encoder.set_compute_pipeline_state(kernel);
 
   // Set its arguments
@@ -578,10 +574,12 @@ void sdpa_vector_2pass(
   compute_encoder.set_input_array(sums, 1);
   compute_encoder.set_input_array(maxs, 2);
   compute_encoder.set_output_array(out, 3);
+  compute_encoder.set_bytes(blocks, 4);
 
   // Launch
   group_dims = MTL::Size(1024, 1, 1);
   grid_dims = MTL::Size(q.shape(0) * q.shape(1), q.shape(2), 1);
+  check_kernel_threadgroup_size(kernel, group_dims, kname);
   compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
 }
 

mlx/backend/metal/utils.h

@@ -81,4 +81,19 @@ inline size_t ceildiv(size_t n, size_t m) {
   return (n + m - 1) / m;
 }
 
+inline void check_kernel_threadgroup_size(
+    const MTL::ComputePipelineState* kernel,
+    MTL::Size group_dims,
+    const std::string& name) {
+  auto max_size = kernel->maxTotalThreadsPerThreadgroup();
+  auto requested_size = group_dims.width * group_dims.height * group_dims.depth;
+
+  if (max_size < requested_size) {
+    std::ostringstream msg;
+    msg << "Maximum threads per threadgroup is " << max_size
+        << " but requested " << requested_size << " for kernel " << name << ".";
+    throw std::runtime_error(msg.str());
+  }
+}
+
 } // namespace mlx::core