bugfix: use actual sm count for num_sm90_ctas (#762)

flashinfer-ai · Jan 29, 2025 · e5a3bef · e5a3bef
1 parent 0e25eb2
commit e5a3bef
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/include/flashinfer/attention/scheduler.cuh b/include/flashinfer/attention/scheduler.cuh
@@ -824,7 +824,11 @@ inline cudaError_t PrefillSM90Plan(void* float_buffer, size_t float_workspace_si
     cta_tile_q = 192;
   }
 
-  const int num_sm90_ctas = 132;  // for sm90, the num_ctas is fixed
+  int device = 0;
+  FLASHINFER_CUDA_CALL(cudaGetDevice(&device));
+  int num_sm90_ctas = 0;
+  FLASHINFER_CUDA_CALL(
+      cudaDeviceGetAttribute(&num_sm90_ctas, cudaDevAttrMultiProcessorCount, device));
 
   CTACostHeap cta_cost_heap(num_sm90_ctas);
   std::vector<std::vector<IdType>> cta_qo_tile_indices(num_sm90_ctas, std::vector<IdType>()),