Skip to content

Commit

Permalink
fix: kernel launch failure due to smem overflow
Browse files Browse the repository at this point in the history
  • Loading branch information
JCao authored and abcdabcd987 committed Nov 22, 2023
1 parent 1495038 commit a6b6ab9
Showing 1 changed file with 6 additions and 0 deletions.
6 changes: 6 additions & 0 deletions csrc/sgmv_flashinfer/sgmv_all.cu
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down

0 comments on commit a6b6ab9

Please sign in to comment.