From ac490399c645eced0a2bfa30f04430044a6796a7 Mon Sep 17 00:00:00 2001
From: Yongqi Chen <144848849+BrianChen1129@users.noreply.github.com>
Date: Mon, 17 Feb 2025 00:35:56 -0500
Subject: [PATCH] fix kernel issue (#185)

Co-authored-by: Peiyuan Zhang <a1286225768@gmail.com>
---
 csrc/sliding_tile_attention/st_attn/st_attn_h100.cu | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/csrc/sliding_tile_attention/st_attn/st_attn_h100.cu b/csrc/sliding_tile_attention/st_attn/st_attn_h100.cu
index aaad99678..6cd633ee7 100644
--- a/csrc/sliding_tile_attention/st_attn/st_attn_h100.cu
+++ b/csrc/sliding_tile_attention/st_attn/st_attn_h100.cu
@@ -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