diff --git a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_128.cu b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_128.cu
index 3a673d3e3..d1b924a20 100644
--- a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_128.cu
+++ b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_128.cu
@@ -25,10 +25,13 @@
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 #define MGQA_LAUNCH_KERNEL(                                                                                            \
-    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                \
-    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);          \
+    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                                    \
+    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);                              \
     dim3   grid(params.num_heads, params.batch_size);                                                                  \
-    mmha::masked_groupedquery_attention_kernel<T,                                                                         \
+    cudaFuncSetAttribute(mmha::masked_groupedquery_attention_kernel<T,                                                 \
+        Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS>,                                          \
+        cudaFuncAttributeMaxDynamicSharedMemorySize, smem_sz);                                                         \
+    mmha::masked_groupedquery_attention_kernel<T,                                                                      \
                                             Dh,                                                                        \
                                             Dh_MAX,                                                                    \
                                             THDS_PER_KEY,                                                              \
@@ -43,7 +46,6 @@ void mgqa_launch_kernel(const KERNEL_PARAMS_TYPE& params, const cudaStream_t& st
 {
     constexpr int  THREADS_PER_VALUE  = threads_per_value_t<T, Dh_MAX>::value;
     int            tlength            = params.timestep;
-    // printf("tlength, CROSS_ATTENTION = %d, %d\n", tlength);
     if (params.cache_indir == nullptr) {
         if (tlength < 32) {
             MGQA_LAUNCH_KERNEL(T, Dh, Dh_MAX, 4, THREADS_PER_VALUE, 64, false, stream);
diff --git a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_144.cu b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_144.cu
index 7e20bdccc..827c5439e 100644
--- a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_144.cu
+++ b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_144.cu
@@ -25,10 +25,13 @@
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 #define MGQA_LAUNCH_KERNEL(                                                                                            \
-    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                \
-    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);          \
+    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                                    \
+    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);                              \
     dim3   grid(params.num_heads, params.batch_size);                                                                  \
-    mmha::masked_groupedquery_attention_kernel<T,                                                                         \
+    cudaFuncSetAttribute(mmha::masked_groupedquery_attention_kernel<T,                                                 \
+        Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS>,                                          \
+        cudaFuncAttributeMaxDynamicSharedMemorySize, smem_sz);                                                         \
+    mmha::masked_groupedquery_attention_kernel<T,                                                                      \
                                             Dh,                                                                        \
                                             Dh_MAX,                                                                    \
                                             THDS_PER_KEY,                                                              \
diff --git a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_160.cu b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_160.cu
index 57c6dd1aa..90b6940cb 100644
--- a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_160.cu
+++ b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_160.cu
@@ -25,10 +25,13 @@
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 #define MGQA_LAUNCH_KERNEL(                                                                                            \
-    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                \
-    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);          \
+    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                                    \
+    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);                              \
     dim3   grid(params.num_heads, params.batch_size);                                                                  \
-    mmha::masked_groupedquery_attention_kernel<T,                                                                         \
+    cudaFuncSetAttribute(mmha::masked_groupedquery_attention_kernel<T,                                                 \
+        Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS>,                                          \
+        cudaFuncAttributeMaxDynamicSharedMemorySize, smem_sz);                                                         \
+    mmha::masked_groupedquery_attention_kernel<T,                                                                      \
                                             Dh,                                                                        \
                                             Dh_MAX,                                                                    \
                                             THDS_PER_KEY,                                                              \
diff --git a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_192.cu b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_192.cu
index d8c349cad..4c1bc3359 100644
--- a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_192.cu
+++ b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_192.cu
@@ -25,10 +25,13 @@
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 #define MGQA_LAUNCH_KERNEL(                                                                                            \
-    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                \
-    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);          \
+    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                                    \
+    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);                              \
     dim3   grid(params.num_heads, params.batch_size);                                                                  \
-    mmha::masked_groupedquery_attention_kernel<T,                                                                         \
+    cudaFuncSetAttribute(mmha::masked_groupedquery_attention_kernel<T,                                                 \
+        Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS>,                                          \
+        cudaFuncAttributeMaxDynamicSharedMemorySize, smem_sz);                                                         \
+    mmha::masked_groupedquery_attention_kernel<T,                                                                      \
                                             Dh,                                                                        \
                                             Dh_MAX,                                                                    \
                                             THDS_PER_KEY,                                                              \
diff --git a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_224.cu b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_224.cu
index 03ff2cadd..bb962d2f7 100644
--- a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_224.cu
+++ b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_224.cu
@@ -25,10 +25,13 @@
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 #define MGQA_LAUNCH_KERNEL(                                                                                            \
-    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                \
-    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);          \
+    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                                    \
+    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);                              \
     dim3   grid(params.num_heads, params.batch_size);                                                                  \
-    mmha::masked_groupedquery_attention_kernel<T,                                                                         \
+    cudaFuncSetAttribute(mmha::masked_groupedquery_attention_kernel<T,                                                 \
+        Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS>,                                          \
+        cudaFuncAttributeMaxDynamicSharedMemorySize, smem_sz);                                                         \
+    mmha::masked_groupedquery_attention_kernel<T,                                                                      \
                                             Dh,                                                                        \
                                             Dh_MAX,                                                                    \
                                             THDS_PER_KEY,                                                              \
diff --git a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_256.cu b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_256.cu
index fe496d4a7..82577421f 100644
--- a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_256.cu
+++ b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_256.cu
@@ -25,10 +25,13 @@
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 #define MGQA_LAUNCH_KERNEL(                                                                                            \
-    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                \
-    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);          \
+    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                                    \
+    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);                              \
     dim3   grid(params.num_heads, params.batch_size);                                                                  \
-    mmha::masked_groupedquery_attention_kernel<T,                                                                         \
+    cudaFuncSetAttribute(mmha::masked_groupedquery_attention_kernel<T,                                                 \
+        Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS>,                                          \
+        cudaFuncAttributeMaxDynamicSharedMemorySize, smem_sz);                                                         \
+    mmha::masked_groupedquery_attention_kernel<T,                                                                      \
                                             Dh,                                                                        \
                                             Dh_MAX,                                                                    \
                                             THDS_PER_KEY,                                                              \
diff --git a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_32.cu b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_32.cu
index ceeb96484..319890ba4 100644
--- a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_32.cu
+++ b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_32.cu
@@ -25,10 +25,13 @@
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 #define MGQA_LAUNCH_KERNEL(                                                                                            \
-    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                \
-    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);          \
+    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                                    \
+    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);                              \
     dim3   grid(params.num_heads, params.batch_size);                                                                  \
-    mmha::masked_groupedquery_attention_kernel<T,                                                                         \
+    cudaFuncSetAttribute(mmha::masked_groupedquery_attention_kernel<T,                                                 \
+        Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS>,                                          \
+        cudaFuncAttributeMaxDynamicSharedMemorySize, smem_sz);                                                         \
+    mmha::masked_groupedquery_attention_kernel<T,                                                                      \
                                             Dh,                                                                        \
                                             Dh_MAX,                                                                    \
                                             THDS_PER_KEY,                                                              \
diff --git a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_48.cu b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_48.cu
index f225bef82..5eaf6ca73 100644
--- a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_48.cu
+++ b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_48.cu
@@ -25,10 +25,13 @@
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 #define MGQA_LAUNCH_KERNEL(                                                                                            \
-    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                \
-    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);          \
+    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                                    \
+    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);                              \
     dim3   grid(params.num_heads, params.batch_size);                                                                  \
-    mmha::masked_groupedquery_attention_kernel<T,                                                                         \
+    cudaFuncSetAttribute(mmha::masked_groupedquery_attention_kernel<T,                                                 \
+        Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS>,                                          \
+        cudaFuncAttributeMaxDynamicSharedMemorySize, smem_sz);                                                         \
+    mmha::masked_groupedquery_attention_kernel<T,                                                                      \
                                             Dh,                                                                        \
                                             Dh_MAX,                                                                    \
                                             THDS_PER_KEY,                                                              \
diff --git a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_64.cu b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_64.cu
index 7a9679952..4c7b2d6fa 100644
--- a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_64.cu
+++ b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_64.cu
@@ -25,10 +25,13 @@
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 #define MGQA_LAUNCH_KERNEL(                                                                                            \
-    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                \
-    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);          \
+    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                                    \
+    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);                              \
     dim3   grid(params.num_heads, params.batch_size);                                                                  \
-    mmha::masked_groupedquery_attention_kernel<T,                                                                         \
+    cudaFuncSetAttribute(mmha::masked_groupedquery_attention_kernel<T,                                                 \
+        Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS>,                                          \
+        cudaFuncAttributeMaxDynamicSharedMemorySize, smem_sz);                                                         \
+    mmha::masked_groupedquery_attention_kernel<T,                                                                      \
                                             Dh,                                                                        \
                                             Dh_MAX,                                                                    \
                                             THDS_PER_KEY,                                                              \
diff --git a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_80.cu b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_80.cu
index 8af12155f..9069915e6 100644
--- a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_80.cu
+++ b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_80.cu
@@ -25,10 +25,13 @@
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 #define MGQA_LAUNCH_KERNEL(                                                                                            \
-    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                \
-    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);          \
+    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                                    \
+    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);                              \
     dim3   grid(params.num_heads, params.batch_size);                                                                  \
-    mmha::masked_groupedquery_attention_kernel<T,                                                                         \
+    cudaFuncSetAttribute(mmha::masked_groupedquery_attention_kernel<T,                                                 \
+        Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS>,                                          \
+        cudaFuncAttributeMaxDynamicSharedMemorySize, smem_sz);                                                         \
+    mmha::masked_groupedquery_attention_kernel<T,                                                                      \
                                             Dh,                                                                        \
                                             Dh_MAX,                                                                    \
                                             THDS_PER_KEY,                                                              \
diff --git a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_96.cu b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_96.cu
index 9160b43c1..fb57189f9 100644
--- a/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_96.cu
+++ b/src/fastertransformer/kernels/llama/decoder_masked_groupedquery_attention/decoder_masked_groupedquery_attention_96.cu
@@ -25,10 +25,13 @@
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 #define MGQA_LAUNCH_KERNEL(                                                                                            \
-    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                \
-    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);          \
+    T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS, stream)                                    \
+    size_t smem_sz = mmha::smem_size_in_bytes<T>(params, THDS_PER_VALUE, THDS_PER_BLOCK);                              \
     dim3   grid(params.num_heads, params.batch_size);                                                                  \
-    mmha::masked_groupedquery_attention_kernel<T,                                                                         \
+    cudaFuncSetAttribute(mmha::masked_groupedquery_attention_kernel<T,                                                 \
+        Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE, THDS_PER_BLOCK, HAS_BEAMS>,                                          \
+        cudaFuncAttributeMaxDynamicSharedMemorySize, smem_sz);                                                         \
+    mmha::masked_groupedquery_attention_kernel<T,                                                                      \
                                             Dh,                                                                        \
                                             Dh_MAX,                                                                    \
                                             THDS_PER_KEY,                                                              \