fix: kernel launch failure due to smem overflow

punica-ai · Nov 22, 2023 · a6b6ab9 · a6b6ab9
1 parent 1495038
commit a6b6ab9
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/csrc/sgmv_flashinfer/sgmv_all.cu b/csrc/sgmv_flashinfer/sgmv_all.cu
@@ -50,9 +50,15 @@ bool sgmv_shrink(T* y, T* x, T** w, int32_t* s, void* tmp,
 
   cudaError_t status;
   if (use_cooperative) {
+    if (smem > 46 * 1024) {
+      cudaFuncSetAttribute(cooperative_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
+    }
     status = cudaLaunchCooperativeKernel((void*)cooperative_kernel, nblks,
                                          nthrs, args, smem, stream);
   } else {
+    if (smem > 46 * 1024) {
+      cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
+    }
     status = cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem, stream);
   }
   return status == cudaSuccess;