fix kernel issue (#185)

Co-authored-by: Peiyuan Zhang <a1286225768@gmail.com>
This commit is contained in:
Yongqi Chen
2025-02-16 21:35:56 -08:00
committed by GitHub
co-authored by Peiyuan Zhang
parent 9ea39cee57
commit ac490399c6
@@ -526,7 +526,15 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
mem_size
);
fwd_attend_ker<128, false, false, 2, 0, 5, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
} else {
} else if (kernel_t_size == 5 && kernel_h_size == 6 && kernel_w_size == 1){
cudaFuncSetAttribute(
fwd_attend_ker<128, false, false, 2, 3, 0, 5, 6, 10>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
mem_size
);
fwd_attend_ker<128, false, false, 2, 3, 0, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
}
else {
// print error
std::cout << "Invalid kernel size" << std::endl;
//print kernel size