Update unfused_attention_kernels.cu

fix bug of softmax kernel
NVIDIA · Apr 20, 2023 · d7ccf83 · d7ccf83
1 parent adb21c3
commit d7ccf83
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/fastertransformer/kernels/unfused_attention_kernels.cu b/src/fastertransformer/kernels/unfused_attention_kernels.cu
@@ -602,6 +602,7 @@ __global__ void softmax_kernel_h2_v2(T*        attn_score,
 
 #define LAUNCH_MAKSED_SOFTMAX_(T_, ITEMS_PER_THREAD)                                                                   \
     block.x /= ITEMS_PER_THREAD;                                                                                       \
+    block.x = (block.x + 31) / 32 * 32;                                                                                \
     assert(block.x <= 1024);                                                                                           \
     if (is_half2) {                                                                                                    \
         if (grid.x % 4 == 0) {                                                                                         \