   T max_val = phi::funcs::blockReduceMax<T>(stride_max, mask);
 
   T stride_sum = 0.f;
-  for (int i = 0; i < seq_len; i += blockDim.x) {
+  for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
     stride_sum += __expf(qk_buf[threadIdx.x + i + qk_offset] +
                          bias_qk[threadIdx.x + i + qk_offset] - max_val);
   }
   T sum_val = phi::funcs::blockReduceSum<T>(stride_sum, mask);
 
-  for (int i = 0; i < seq_len; i += blockDim.x) {
+  for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
     qk_buf[threadIdx.x + i + qk_offset] =
         (T)(__expf(qk_buf[threadIdx.x + i + qk_offset] +
                    bias_qk[threadIdx.x + i + qk_offset] - max_val) /

      phi::funcs::blockReduceMax<float>(max(stride_max.x, stride_max.y), mask);

  float2 stride_sum = make_float2(0.f, 0.f);
  for (int i = 0; i < seq_len; i += blockDim.x) {
    float2 cur =
        phi::funcs::ToFloat2<half2>(qk_buf_[threadIdx.x + i + qk_offset] +
                                    bias_qk_[threadIdx.x + i + qk_offset]);

  assert(blockDim.x % 32 == 0);

  float2 stride_max = make_float2(-1e20f, -1e20f);
  for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
    float2 cur =
        phi::funcs::ToFloat2<half2>(qk_buf_[threadIdx.x + i + qk_offset] +
                                    bias_qk_[threadIdx.x + i + qk_offset]);
