 #define WARP_SIZE 32
 
 template <typename T>
-__inline__ __device__ T warpReduceSum(T val) {
+__inline__ __device__ T warpReduceSum(T val, unsigned lane_mask) {
   for (int mask = HALF_WARP; mask > 0; mask >>= 1)
 #if __CUDA_ARCH__ >= 350 && CUDA_VERSION >= 9000
-    val += __shfl_xor_sync(FINAL_MASK, val, mask, warpSize);
+    val += __shfl_xor_sync(lane_mask, val, mask, warpSize);
 #else
     val += __shfl_xor(val, mask, warpSize);
 #endif

                                       bias_qk_[threadIdx.x + bias_offset]))
                 : 0.0f;
  float tmp = threadIdx.x < seq_len ? static_cast<float>(qk) : -1e20f;
  float max_val = blockReduceMax<float>(tmp);
  if (threadIdx.x == 0) s_max = max_val;
  __syncthreads();

  float qk_tmp =
      threadIdx.x < seq_len ? __expf(static_cast<float>(tmp - s_max)) : 0.0f;
  float sum_val = blockReduceSum<float>(qk_tmp);

  if (threadIdx.x == 0) {
    s_sum = sum_val + 1e-6f;

__global__ void softmax_kernel_with_eltadd(T *qk_buf_, const T *bias_qk_,
                                           const int batch_size,
                                           const int head_num,
                                           const int seq_len,
                                           const unsigned mask) {
  int seq_id = blockIdx.x % seq_len;
  int qk_offset = blockIdx.x * seq_len;
  int bias_offset = blockIdx.x % (head_num * seq_len) * seq_len;
