 
 /* Calculate the sum of all elements in a block */
 template <typename T>
-__inline__ __device__ T blockReduceSum(T val) {
+__inline__ __device__ T blockReduceSum(T val, unsigned mask) {
   static __shared__ T shared[WARP_SIZE];
   int lane = threadIdx.x & 0x1f;
   int wid = threadIdx.x >> 5;
 
-  val = warpReduceSum<T>(val);
+  val = warpReduceSum<T>(val, mask);
 
   if (lane == 0) shared[wid] = val;
 
   __syncthreads();
 
-  val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : (T)(0.0f);
-  val = warpReduceSum<T>(val);
+  // align block_span to warpSize
+  int block_span = (blockDim.x + warpSize - 1) >> 5;
+  val = (threadIdx.x < block_span) ? shared[lane] : (T)(0.0f);
+  val = warpReduceSum<T>(val, mask);
 
   return val;
 }
 
 template <typename T>
-__inline__ __device__ T warpReduceMax(T val) {
+__inline__ __device__ T warpReduceMax(T val, unsigned lane_mask) {
   for (int mask = HALF_WARP; mask > 0; mask >>= 1)
 #if __CUDA_ARCH__ >= 350 && CUDA_VERSION >= 9000
-    val = max(val, __shfl_xor_sync(FINAL_MASK, val, mask, warpSize));
+    val = max(val, __shfl_xor_sync(lane_mask, val, mask, warpSize));
 #else
     val = max(val, __shfl_xor(val, mask, warpSize));
 #endif


/* Calculate the maximum of all elements in a block */
template <typename T>
__inline__ __device__ T blockReduceMax(T val) {
  static __shared__ T shared[WARP_SIZE];
  int lane = threadIdx.x & 0x1f;
  int wid = threadIdx.x >> 5;

  val = warpReduceMax(val);

  if (lane == 0) shared[wid] = val;

  __syncthreads();

  val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : -1e10f;
  val = warpReduceMax(val);

  return val;
}

  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>
