 __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 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;

                                       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;

  int grid = m;
  int block = k;

  unsigned mask = block < 32 ? (((unsigned)1 << block) - 1) : FINAL_MASK;
  softmax_kernel_with_eltadd<T><<<grid, block, 0, stream>>>(
      qk_buf_, bias_qk, batch_size, head_num, seq_len, mask);
}

template <typename T>
